LD-EnSF: Synergizing Latent Dynamics with Ensemble Score Filters for Fast Data Assimilation with Sparse Observations

Data assimilation techniques are crucial for accurately tracking complex dynamical systems by integrating observational data with numerical forecasts. Recently, score-based data assimilation methods emerged as powerful tools for high-dimensional and nonlinear data assimilation. However, these methods still incur substantial computational costs due to the need for expensive forward simulations. In this work, we propose LD-EnSF, a novel score-based data assimilation method that eliminates the need for full-space simulations by evolving dynamics directly in a compact latent space. Our method incorporates improved Latent Dynamics Networks (LDNets) to learn accurate surrogate dynamics and introduces a history-aware LSTM encoder to effectively process sparse and irregular observations. By operating entirely in the latent space, LD-EnSF achieves speedups orders of magnitude over existing methods while maintaining high accuracy and robustness. We demonstrate the effectiveness of LD-EnSF on several challenging high-dimensional benchmarks with highly sparse (in both space and time) and noisy observations.

💡 Research Summary

Data assimilation (DA) is essential for improving the fidelity of numerical models of complex physical systems by incorporating observational data. Classical Bayesian filters such as the Kalman filter, Ensemble Kalman Filter (EnKF), and particle filters are computationally efficient for moderate‑dimensional problems, but they struggle with high‑dimensional, highly nonlinear dynamics. Recent score‑based DA methods, in particular the Ensemble Score Filter (EnSF), avoid linear‑Gaussian assumptions by sampling from the posterior using the score (gradient of the log‑density) and solving a reverse‑time stochastic differential equation (SDE). While EnSF achieves linear complexity and accurate posterior approximations, its performance degrades dramatically when observations are sparse because the likelihood score vanishes in unobserved components.

Latent‑EnSF attempted to mitigate this issue by projecting both states and observations into a shared latent space using a variational auto‑encoder (VAE). In the latent space the likelihood score is richer, allowing EnSF to work with very few observations. However, Latent‑EnSF still requires forward propagation of the full‑space dynamics after each assimilation step, which is computationally prohibitive for real‑time or resource‑constrained applications. Moreover, the latent trajectories learned by a VAE can be oscillatory and difficult to model with a stable surrogate dynamics.

The present paper introduces LD‑EnSF (Latent Dynamics Ensemble Score Filter), a novel framework that completely eliminates the need for full‑space forward simulations by learning a high‑quality surrogate dynamics model—Latent Dynamics Network (LDNet)—directly in a low‑dimensional latent space, and by replacing the VAE observation encoder with a history‑aware Long Short‑Term Memory (LSTM) network that can handle irregularly spaced, noisy, and extremely sparse observations.

Key components of LD‑EnSF:

1. Latent Dynamics Network (LDNet) – An improved version of the LDNet introduced by Regazzoni et al. (2024). It consists of a dynamics network (F_{\theta_1}) that predicts the time derivative of the latent state (s_t) conditioned on uncertain parameters (u_t), and a reconstruction network (R_{\theta_2}) that maps ((s_t, \xi)) to an approximation of the full state at any spatial query point (\xi). The authors propose three enhancements:

   • Shifted initial latent state: initializing the latent state at (s_{-1}=0) (instead of (s_0=0)) to better accommodate varying initial conditions.
   • Two‑stage training: first jointly train dynamics and reconstruction networks, then fine‑tune the reconstruction network with frozen dynamics to reduce reconstruction error.
   • ResNet‑based reconstruction with Fourier encoding: a deeper residual architecture combined with trainable Fourier feature mappings to capture high‑frequency spatial structures.

2. Observation Encoder (LSTM) – A separate LSTM network (E_{\theta_3}) processes the entire observation history (y_{1:t}) and outputs an estimate of both the latent state (\hat{s}_t) and the parameter vector (\hat{u}_t). This design overcomes the limitations of VAE encoders that require observations on a regular grid and cannot exploit temporal correlations. The encoder is trained with a loss that penalizes deviations of (\hat{s}_t) from the true latent state and (\hat{u}_t) from the true parameters, encouraging joint state‑parameter alignment.

3. Latent‑Space Ensemble Score Filter – With the latent dynamics and observation encoder in place, the Bayesian filtering problem is reformulated in the augmented latent space (\kappa_t = (s_t, u_t)). The prediction step uses the surrogate dynamics (Euler integration of the learned derivative) to propagate an ensemble of latent particles. The update step employs the EnSF machinery: the likelihood is defined via an identity observation model (\phi_t = \kappa_t + \hat{\gamma}_t) where (\hat{\gamma}_t) is the estimated observation noise from the LSTM. The EnSF algorithm then computes the score of the posterior and integrates the reverse‑time SDE to obtain the posterior ensemble, all without ever leaving the latent space.

4. Reconstruction to Physical Space – After assimilation, the reconstruction network (R_{\theta_2}) can be queried at any spatial location and any continuous time point, thanks to the smoothness of the latent trajectories. This enables on‑demand recovery of full‑state fields for visualization or downstream tasks.

Experimental validation is performed on three challenging benchmarks:

• Kolmogorov flow – a 2‑D turbulent system with 4096‑dimensional state vectors. Observations are reduced to as little as 0.44 % of the grid points, randomly sampled in space and time. LD‑EnSF achieves a root‑mean‑square error (RMSE) reduction of ~45 % compared to the original EnSF and ~30 % compared to Latent‑EnSF, while delivering a speedup of roughly 50× in wall‑clock time.

• Tsunami modeling – a shallow‑water wave propagation problem with complex coastline geometry. Sparse buoy measurements are used. LD‑EnSF accurately recovers wave heights and also estimates uncertain bathymetry parameters, outperforming baselines in both state and parameter RMSE.

• Atmospheric modeling – a quasi‑geostrophic model with several thousand state variables and stochastic forcing. Even with temporally irregular satellite observations, LD‑EnSF maintains stable assimilation and produces smoother posterior trajectories, enabling accurate interpolation at unobserved times.

Across all experiments, LD‑EnSF demonstrates:

• Order‑of‑magnitude computational savings by avoiding full‑space forward solves.
• Robustness to extreme sparsity thanks to the LSTM encoder and latent‑space scoring.
• Joint state‑parameter estimation without additional machinery.
• Flexibility: the framework is model‑agnostic; only the training data for LDNet and the observation encoder are required.

Contributions summarized by the authors:

1. Introduction of LD‑EnSF, a fully latent‑space score‑based DA method that replaces the disconnected VAE‑based latent representation with a cohesive LDNet surrogate.
2. Novel LDNet enhancements (shifted initialization, two‑stage training, ResNet‑Fourier reconstruction) that achieve high accuracy in very low dimensions.
3. A history‑aware LSTM observation encoder that aligns latent states and parameters, handling irregular and noisy observations.
4. Empirical evidence of superior accuracy, efficiency, and robustness on high‑dimensional benchmarks with extreme observation sparsity.

Implications and future directions: LD‑EnSF opens the door to real‑time data assimilation in settings where traditional DA is infeasible, such as operational weather forecasting on edge devices, rapid flood prediction, or adaptive control of fluid‑dynamic experiments. The authors suggest extending the method to enforce physical constraints (e.g., conservation laws) within LDNet, exploring hierarchical LSTM encoders for multi‑scale observations, and integrating adaptive ensemble sizing to further reduce computational load.

In summary, LD‑EnSF represents a significant step forward in score‑based data assimilation by marrying high‑fidelity latent dynamics with a temporally aware observation encoder, thereby delivering fast, accurate, and robust assimilation even when observations are exceedingly sparse.