Latent Autoencoder Ensemble Kalman Filter for Nonlinear Data assimilation
Authors: Xin T. Tong, Yanyan Wang, Liang Yan
Organizations: Department of Mathematics, National University of Singapore, 119076, Singapore · School of Mathematics, Southeast University, Nanjing 210096, China
The ensemble Kalman filter (EnKF) is widely used for data assimilation in high-dimensional systems, but its performance often deteriorates for strongly nonlinear dynamics due to the structural mismatch between the Kalman update and the underlying system behavior. In this work, we propose a latent autoencoder ensemble Kalman filter (LAE-EnKF) that addresses this limitation by reformulating the assimilation problem in a learned latent space with linear and stable dynamics. The proposed method learns a nonlinear encoder--decoder together with a stable linear latent evolution operator and a consistent latent observation mapping, yielding a closed linear state-space model in the latent coordinates. This construction restores compatibility with the Kalman filtering framework and allows both forecast and analysis steps to be carried out entirely in the latent space. Compared with existing autoencoder-based and latent assimilation approaches that rely on unconstrained nonlinear latent dynamics, the proposed formulation emphasizes structural consistency, stability, and interpretability. We provide a theoretical analysis of learning linear dynamics on low-dimensional manifolds and establish generalization error bounds for the proposed latent model. Numerical experiments on representative nonlinear and chaotic systems demonstrate that the LAE-EnKF yields more accurate and stable assimilation than the standard EnKF and related latent-space methods, while maintaining comparable computational cost and data-driven.
Figures & tables
Figure 1: The architecture of latent autoencoder framework. It consists of a state encoder E that maps physical states to a latent space, a decoder D that reconstructs states from latent representations, a linear transition operator A that governs latent dynamics, an observation encoder Eobs that embeds observations into the latent space, and a latent observation operator H , which may be specified a priori or learned.
Figure 2: Learned latent representations for latent dimensions n=2,3,4 . Top: latent variables learned by the autoencoder component of DAE-EnKF (denoted as DAE), where no structural constraint is imposed on the latent dynamics. Bottom: latent variables learned by the proposed LAE, which enforces linear evolution in the latent space. For n>2 , the latent variables are projected onto their first two principal components using PCA for visualization.
Figure 3: Pairwise latent coordinate plots (zi,zj) for n=3 learned by the LAE.
Figure 4: Long-term prediction relative error for latent dimensions n=2,3,4 . Predictions are obtained by recursively propagating the learned latent dynamics over time steps and decoding back to the physical space.
n
EnKF
AE-EnKF
DAE-EnKF
LAE-EnKF
–
1.9757E-01
–
–
–
n=2
–
1.6412E-01
2.3312E-01
3.8947E-02
n=3
–
1.8667E-01
2.1106E-01
1.4244E-02
n=4
–
1.9632E-01
4.0424E-01
1.5682E-02
Table 1: Relative error \mathrm{E}_{\text{Rel},1:{\color[rgb]{0,0,0}K}} for different methods with various values of n .
Figure 5: Evolution of the relative RMSE over discrete time steps for different methods with latent dimension n=2 . Solid curves show the mean over 10 independent runs, and shaded regions indicate the corresponding 95% confidence intervals.
Figure 6: Reconstructed state trajectories in two representative coordinate pairs (xi,xj) from the D=100 dimensional physical state for different methods with n=2 , compared with the true trajectory. Red stars indicate the initial reconstructed state, and red dots denote the initial true state.
Figure 7: True solution fields of the advection–diffusion–reaction equation at the initial time (Left) and final time T=3 (Right). White circles indicate the observation locations.
Figure 8: Ablation study of the loss function ( 8 ) with latent dimension n=20 . Left: relative state prediction error ∥xk∥2∥DAkE(x0)−xk∥2 . Right: relative latent consistency error ∥E(xk)∥2∥AkE(x0)−E(xk)∥2 .
Figure 9: Time evolution of the relative RMSE for Example 5.2 under latent dimensions n=16,20,24 . Solid curves indicate the mean over 10 independent runs, and the shaded regions represent the corresponding 95% confidence intervals.
Figure 10: Reconstructed fields (Top) and corresponding pointwise absolute errors (Bottom) at the final time T=3 by LAE-EnKF for different latent dimensions n=16 , n=20 , and n=24 .
Figure 11: Comparison of reconstructed fields (Top) and corresponding pointwise absolute errors (Bottom) at the final time T=3 for latent dimension n=20 .
Figure 12: Time evolution of the relative RMSE for Example 5.2 under different latent observation operators: identity operator (Left), randomly operator (Middle), and learned operator (Right). Solid curves denote the mean over 10 independent runs, and shaded regions indicate the corresponding 95% confidence intervals. The global relative error ERel,1:K is reported in each subplot.
σ
EnKF
AE-EnKF
DAE-EnKF
LAE-EnKF
σ=0.10
0.7987
0.3371
0.4369
0.1058
σ=0.05
0.7553
0.2855
0.3400
0.0780
σ=0.01
0.7516
0.2344
0.2860
0.0633
Time (s)
2837.29
6.75
0.10
0.05
Table 2: Relative error \mathrm{E}_{\text{Rel},1:{\color[rgb]{0,0,0}K}} and online computation time for Example 2.
Δt
n
EnKF
AE-EnKF
DAE-EnKF
LAE-EnKF
–
0.5977 (0.1590)
–
–
–
64
–
0.4504 (0.2027)
0.8276 (0.8195)
0.1475 (0.1480)
0.1
80
–
0.3793 (0.2026)
0.8202 (0.8697)
0.1509 (0.1512)
100
–
0.3563 (0.2080)
0.6696 (0.6804)
0.1426 (0.1431)
–
0.3223 (0.1689)
–
–
–
64
–
0.4376 (0.2416)
1.1994 (1.0393)
0.1862 (0.1864)
Table 3: Relative RMSE ERel,1:K for different methods with varying latent dimension n , reported without (with) covariance localization, under two observation intervals.
Figure 13: Reconstructed spatial state fields (Left) and corresponding error (Right) with the observation interval Δt=0.1 and latent dimension n=64 . From top to bottom, the results show the localization EnKF, localization AE-EnKF , DAE-EnKF and the LAE-EnKF.
Figure 14: Time evolution of a representative observed variable (Left) and an unobserved variable (Right) reconstructed with observation interval Δt=0.1 and latent dimension n=64 , compared with the corresponding true trajectories. From top to bottom, the results show the localization EnKF, localization AE-EnKF , DAE-EnKF and the LAE-EnKF.
Figure 15: Reconstructed spatial state fields (Left) and corresponding error (Right) with the observation interval Δt=0.2 and latent dimension n=64 . From top to bottom, the results show the localization EnKF, localization AE-EnKF , DAE-EnKF and the LAE-EnKF.
Figure 16: Time evolution of a representative observed variable (Left) and an unobserved variable (Right) reconstructed by LAE-EnKF with observation interval Δt=0.2 for different latent dimension n=64,80,100 .
Figure 17: Reconstruction of representative state components under nonlinear observations. Comparison between localized EnKF and LAE-EnKF under different latent observation operators, including the identity, random, and learned operators. The global RMSE ERel,1:K is reported above each panel.
Data assimilation combines model forecasts with noisy, incomplete observations to estimate the evolving state of a dynamical system. Existing methods face two compounding challenges: high-dimensional nonlinear dynamics make repeated forward simulation computationally expensive, while sparse observations provide limited direct information about the full state. To address these challenges, we propose the Latent-Dynamics Ensemble Flow Filter (LD-EnFF), a sequential Bayesian filtering framework that performs both forecast propagation and filtering updates in a compact latent space. LD-EnFF combines a latent dynamics surrogate for ensemble propagation with a variational autoencoder (VAE)-based observation model that evaluates a state-dependent observation likelihood in latent space. At each assimilation step, an ensemble filtering update based on flow matching uses the forecast ensemble and this likelihood to generate posterior samples, jointly updating latent states and uncertain parameters. This design avoids repeated full-state simulation during forecasting and full-field reconstruction during likelihood evaluation. LD-EnFF substantially outperforms a broad range of data assimilation algorithms on benchmarks spanning Kolmogorov flow, tsunami propagation, and atmospheric modeling, all featuring complex dynamics and sparse, noisy observations.
Ziyu Tian, Kaichen Shen, Wenbo Hao +3
School of Mathematics, Georgia Institute of Technology · School of Computational Science and Engineering, Georgia Institute of Technology
Long-range geophysical forecasts are fundamentally limited by chaotic dynamics and numerical errors. While data assimilation can mitigate these issues, classical variational smoothers require computationally expensive tangent-linear and adjoint models. Conversely, recent efficient latent filtering methods often enforce weak trajectory-level constraints and assume fixed observation grids. To bridge this gap, we propose Latent Ensemble Variational Data Assimilation (LEVDA), an ensemble-space variational smoother that operates in the low-dimensional latent space of a pretrained differentiable neural dynamics surrogate. By performing four-dimensional ensemble-variational (4DEnVar) optimization within an ensemble subspace, LEVDA jointly assimilates states and unknown parameters without the need for adjoint code or auxiliary observation-to-latent encoders. Leveraging the fully differentiable, continuous-in-time-and-space nature of the surrogate, LEVDA naturally accommodates highly irregular sampling at arbitrary spatiotemporal locations. Across three challenging geophysical benchmarks, LEVDA matches or outperforms state-of-the-art latent filtering baselines under severe observational sparsity while providing more reliable uncertainty quantification. Simultaneously, it achieves substantially improved assimilation accuracy and computational efficiency compared to full-state 4DEnVar.
Phillip Si, Peng Chen
School of Computational Sciences and Engineering, Georgia Institute of Technology
The ensemble Kalman filter (EnKF) is widely adopted for sequential data assimilation, but fails for solutions with discontinuities, such as shocks in compressible flows. For ensembles with uncertainty in discontinuity location, the EnKF analysis can lose physical consistency, which can be viewed through three lenses: (i) bimodal pointwise statistics in the forecast ensemble, (ii) mixed-sign analysis weights in the affine superposition of forecast members, and (iii) spatially distributed Kalman-gain profiles that modify the state around a discontinuity rather than translating it coherently. In order to preserve flow features in the analysis state and, as such, maintain physical consistency, a feature-preserving latent-EnKF is introduced. The proposed method performs the ensemble update in a low-dimensional latent space learned by a coordinate-conditioned auto-decoder. The latent representation retains near-Gaussian ensemble statistics and provides coordinates in which flow features vary coherently across the forecast ensemble. The EnKF updates performed on the latent codes are decoded into coherent changes in the represented flow features, translating shocks through coordinate morphing rather than superposing them at displaced locations. The corresponding Kalman-gain profiles generated by the latent-EnKF exhibit an approximately step-like variation across the discontinuities, allowing the fronts to translate while preserving sharp features and monotonicity. Sod shock-tube, Mach~2 shock-cylinder, and immersed-body shape-inference tests with sparse, noisy observations demonstrate accurate recovery of shocks and contact discontinuities without spurious oscillations.