VAE-Var: Variational-Autoencoder-Enhanced Variational Assimilation

TL;DR

This work tackles the Gaussian-background-error limitation in variational data assimilation by introducing VAE-Var, which learns a non-Gaussian background-error distribution via a variational autoencoder. The authors derive a tractable variational cost in latent space that includes a regularization term and a Jacobian-determinant term to account for the nonlinear mapping between latent and physical spaces. They demonstrate consistent accuracy gains over traditional 3DVar and 4DVar on low-dimensional chaotic systems (Lorenz 63 and Lorenz 96) under both linear and nonlinear observation operators, with ablations confirming the importance of each cost component. The approach shows promise for improved assimilation accuracy, though scalability to high-dimensional systems remains a key limitation for future work.

Abstract

Data assimilation refers to a set of algorithms designed to compute the optimal estimate of a system's state by refining the prior prediction (known as background states) using observed data. Variational assimilation methods rely on the maximum likelihood approach to formulate a variational cost, with the optimal state estimate derived by minimizing this cost. Although traditional variational methods have achieved great success and have been widely used in many numerical weather prediction centers, they generally assume Gaussian errors in the background states, which limits the accuracy of these algorithms due to the inherent inaccuracies of this assumption. In this paper, we introduce VAE-Var, a novel variational algorithm that leverages a variational autoencoder (VAE) to model a non-Gaussian estimate of the background error distribution. We theoretically derive the variational cost under the VAE estimation and present the general formulation of VAE-Var; we implement VAE-Var on low-dimensional chaotic systems and demonstrate through experimental results that VAE-Var consistently outperforms traditional variational assimilation methods in terms of accuracy across various observational settings.

Figures (15)

• Figure 1: Comparison between the traditional 3DVar algorithm (upper part) and our proposed VAE-3DVar algorithm (lower part). The blue dots correspond to the distribution of the background state (in both the latent space and the physical space); the green dots correspond to the observation distribution; the purple dots correspond to the training samples generated with the NMC method.
• Figure 2: Results on the Lorenz 63 system with both linear (the first and the second rows) and nonlinear (the third and the last rows) observation operators under 3DVar observational settings. Different panels correspond to different observational settings. For example, the title "obs var=X,Y (i.d.) / (abs.)" means that variables $X$ and $Y$ are observed with an "identity" or "absolute" function. The x-axis represents the standard deviation of the observation noise $\epsilon_{noise}$; the y-axis corresponds to two evaluation metrics: RMSE and $\mathrm{Imp}$. $\mathrm{Imp}$ is evaluated between the VAE-3DVar ($\mathcal{L}_o+\mathcal{L}_{reg}+\mathcal{L}_{det}$) and the traditional 3DVar. The experiments are repeated 10 times using different random noise, and we report the one-sigma error bars. The label "Naive Substitute" corresponds to the algorithm where the observed variables of the background state are replaced with their observed values.
• Figure 3: Results on the Lorenz 96 system with both linear (the first and the second rows) and nonlinear (the last row) observation operators under 3DVar observational settings. In the title, "sat." is abbreviated for the saturated observation operator. The $\mathrm{Imp}$ metric is not demonstrated for nonlinear settings because $R_{bg} - R_{trad}$ can be very small, making the value of $\mathrm{Imp}$ larger than $10^4$.
• Figure 4: Results on the Lorenz 63 (the first column) and the Lorenz 96 (the second and the third columns) systems with linear observation operators under 4DVar observational setting. $\mathrm{Imp}$ is evaluated between VAE-4DVar ($\mathcal{L}_o+\mathcal{L}_{reg}+\mathcal{L}_{det}$) and traditional 4DVar.
• Figure 5: Structure of the variational autoencoder.