Latent Autoencoder Ensemble Kalman Filter for Nonlinear Data assimilation

Xin T. Tong; Yanyan Wang; Liang Yan

cs.LGMar 6, 2026

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

Abstract

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

\mathscr{E}

that maps physical states to a latent space, a decoder

\mathscr{D}

that reconstructs states from latent representations, a linear transition operator

\mathbf{A}

that governs latent dynamics, an observation encoder

\mathscr{E}_{\mathrm{obs}}

that embeds observations into the latent space, and a latent observation operator

\mathbf{H}

, which may be specified a priori or learned.

\mathbf{P}_{zy}=\frac{1}{N_{e}-1}\!\sum\limits_{j=1}^{N_{e}}(\widehat{\mathbf{z}}_{k|k-1}^{(j)}-\overline{\mathbf{z}}_{k|k-1})(\widehat{\mathbf{y}}_{k}^{(j)}-\overline{\mathbf{y}}_{k})^{\top},

Algorithm 1 Latent Autoencoder Ensemble Kalman Filter (LAE-EnKF)

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

(z_{i},z_{j})

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.

$\boldsymbol{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

(x_{i},x_{j})

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.

Panel of Figure 7 — 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

\frac{\|\mathscr{D}\mathbf{A}^{k}\mathscr{E}(\mathbf{x}_{0})-\mathbf{x}_{k}\|_{2}}{\|\mathbf{x}_{k}\|_{2}}

. Right: relative latent consistency error

\frac{\|\mathbf{A}^{k}\mathscr{E}(\mathbf{x}_{0})-\mathscr{E}(\mathbf{x}_{k})\|_{2}}{\|\mathscr{E}(\mathbf{x}_{k})\|_{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

\mathrm{E}_{\text{Rel},1:K}

is reported in each subplot.

$\boldsymbol{\sigma}$	EnKF	AE-EnKF	DAE-EnKF	LAE-EnKF
$\sigma=0.10$	0.7987	0.3371	0.4369	0.1058
$\sigma=0.05$	0.7553	0.2855	0.3400	0.0780
$\sigma=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.

$\Delta t$	$\boldsymbol{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

\mathrm{E}_{\mathrm{Rel},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 $\Delta 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

\Delta 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 $\Delta 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

\Delta 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

\mathrm{E}_{\text{Rel},1:K}

is reported above each panel.

Data Assimilation

Data assimilation combines observations with dynamical models to estimate hidden system states and their uncertainty over time.

See all 58 papers

LD-EnFF: Latent-Dynamics Ensemble Flow Filtering for Data Assimilation with Sparse Observations
Oct 2, 2026Ziyu Tian, Kaichen Shen, Wenbo Hao +3Data Assimilation Bayesian Filtering
School of Mathematics, Georgia Institute of Technology · School of Computational Science and Engineering, Georgia Institute of Technology
LEVDA: Latent Ensemble Variational Data Assimilation via Differentiable Dynamics
Feb 23, 2026Phillip Si, Peng ChenData Assimilation Latent Space Optimization
School of Computational Sciences and Engineering, Georgia Institute of Technology
Feature-preserving Latent-EnKF for Data Assimilation of Flows with Shocks
Jun 10, 2026Hemanth Chandravamsi, Hangchuan Hu, Ponkrshnan Thiagarajan +1Kalman Filtering Data Assimilation
Department of Mechanical Engineering, Johns Hopkins University, Baltimore, MD 21218, USA

Oct 2, 2026J/K move · Enter open · S save

LD-EnFF: Latent-Dynamics Ensemble Flow Filtering for Data Assimilation with Sparse Observations

Ziyu Tian, Kaichen Shen, Wenbo Hao, Phillip Si+2

School of Mathematics, Georgia Institute of Technology · School of Computational Science and Engineering, Georgia Institute of Technology

Data Assimilation Bayesian Filtering Latent Dynamics Modeling

Open paper AskPDF arXiv

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.

Abstract

Data Assimilation

Explore similar work

LD-EnFF: Latent-Dynamics Ensemble Flow Filtering for Data Assimilation with Sparse Observations

LEVDA: Latent Ensemble Variational Data Assimilation via Differentiable Dynamics

Feature-preserving Latent-EnKF for Data Assimilation of Flows with Shocks