VAE-Var: Variational Autoencoder-Enhanced Variational Methods for Data Assimilation in Meteorology
Authors: Yi Xiao
Qilong Jia
Kun Chen
LEI BAI
Wei Xue
Venue: ICLR 2025 Poster (2025)
Keywords: Data assimilation
Variational Autoencoder
Weather Forecasting
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Data assimilation (DA) is an essential statistical technique for generating accurate estimates of a physical system's states by combining prior model predictions with observational data, especially in the realm of weather forecasting. Effectively modeling the prior distribution while adapting to diverse observational sources presents significant challenges for both traditional and neural network-based DA algorithms. This paper introduces VAE-Var, a novel neural network-based data assimilation algorithm aimed at 1) enhancing accuracy by capturing the non-Gaussian characteristics of the conditional background distribution $p(\mathbf{x}|\mathbf{x}_b)$, and 2) efficiently assimilating real-world observational data. VAE-Var utilizes a variational autoencoder to learn the background error distribution, with its decoder creating a variational cost function to optimize the analysis states. The advantages of VAE-Var include: 1) it maintains the framework of traditional variational assimilation, enabling it to accommodate various observation operators, particularly irregular observations; 2) it lessens the dependence on expert knowledge for constructing the background distribution, allowing for improved modeling of non-Gaussian structures; and 3) experimental results indicate that, when applied to the FengWu weather forecasting model, VAE-Var outperforms DiffDA and two traditional algorithms (interpolation and 3DVar) in terms of assimilation accuracy in sparse observational contexts, and is capable of assimilating real-world GDAS prepbufr observations over a year.

Published as a conference paper at ICLR 2025
VAE-VAR: VARIATIONAL JUTOENCODER-EEHANNEE VARIATIONAL METHODS FOR DATA ASSIMILATION IN METEOROLOGY
Yi Xiao!.2 Qilong Jia' Kun Chen
Lei Bai
Wei Xue
l Department of Computer Science and Technology, Tsinghua University
Shanghai Artificial Intelligence Laboratory (y-xiaoe ao22, jq123 iaa.ls.. www.aalsssssinnaaaaa.a. edu. en kunc3301@163.com baisanshi @ @mmall com
xuewei etareittttingua.euu..cr en
ABSTRACT
Data assimilation (DA) an essential statistical technique for generating accurate estimates of a physical system's 's states by combining prior model predictions with observational data, especially in the realm of weather forecasting. Effectively modeling the prior distribution while adapting to diverse observational sources presents significant challenges for both traditional and neural network-based DA algorithms. This paper introduces VAE-Var, novel neural network-based data assimilation algorithm aimed at 1) enhancing accuracy by capturing the non- Gaussian characteristics of the conditional background distribution p(x(x,), and 2) efficienctly assimilating real-world observational data. VAE-Va Var utilizes vari- ational autoencoder to learn the background error distribution, with its decoder creating a variational cost function to optimize the analysis states. The advan- tages of VAE-Var include: 1) it maintains the framework of traditional variational assimilation, enabling it to accommodate various observation operators, partic- ularly irregular observations;
it lessens the dependence on expert knowledge for constructing the background distribution, allowing for improved modeling of Gaussian structures; and 3) experimental results indicate that, when applied to the Feng Wu weather forecasting model, VAE-Var outperforms DiffDA and two traditional algorithms (interpolation and 3DVar) terms of assimilation accuracy in sparse observational contexts, and is capable of assimilating real-world GDAS prepbufr observations over a year.
INTRODUCTION
Data assimilation (DA) is a statistical approach aimed at modeling the posterior likelihood distribu- tion of physical system's states, represented as x[x|xx,,)) This distribution is conditioned prior predictions, known as background states Xs, and observational data y. ultimately yielding an accu- rate estimate, referred as analysis states x., through sampling from p(x]x,, y) or by maximizing the likelihood. Data assimilation is critical for deriving initial states in various fields, particularly in
numerical weather forecasting (Bauer et al., 2015) Kalnay 2003
Carrassi al. 2018).
In data assimilation for weather forecasting, background states X₁, are essentially predictions pro- duced by an imprecise forecasting model, with the conditional likelihood p(x[x;, often remaining unknown. Due to the high dimensionality of weather systems, modeling (/X(x6) presents significant challenges (Strogatz) 2018). Since meteorological observations originate from different sources, the types of y are highly diverse. Typically, the observation operator H. is used to represent the mapping from physical space to observational space. An effective data assimilation algorithm must not only deliver accurate estimates of the conditional background distribution p( p(x[x;, but also proficiently manage various observation operators, preferably in a zero-shot manner.
Traditional algorithms and neural network- based approaches offer distinct advantages in address- ing the data assimilation problem. Traditional data assimilation methods, particularly variational assimilation algorithms, operate under the assumption that the background emror, Xg, follows

---

Published as a conference paper at ICLR 2025
Gaussian distribution that is independent of x (Barker et al. 2004 Bannister 2008 Descombes et al. 2015 Trémolet 2006). Under this assumption, the posterior likelihood p(x x6. y) can be calculated explicitly and estimation can be achieved by maximizing the likelihood function. The advantage of variational methods lies in its robustness to various observation operators; regardless of how H is formulated, maximization can always be performed using techniques such gradi- ent descent. This approach has proven successful for over two decades (Rabier et al., 1998] 2000). However, relies heavily on expert knowledge for constructing the correlations of background er- ror structures (Huang et al. 2209);; additionally, the assumption of Gaussian error distribution may constrain the potential accuracy of this approach.
In recent years, generative neural networks, particularly diffusion models, have been increasingly utilized in data assimilation to model the background conditional distribution. Two notable exam- ples are SDA (Rozet & Louppe, 2023) and DifIDA (Huang et al. (024). In both studies, diffu- sion model is first employed represent the background distribution p(x[xx)) Subsequently, SDA leverages diffusion posterior sampling techniques (Chung et al. (022)) to incorporate observational information y, while DifIDA utilizes the repaint technique (Lugmayr et al. 2022) to integrate ob- servations and sample from p(x X6. y) during inference. The advantage of these approaches over traditional algorithms lies in their ability to model more complex structures for p(x[xs). potentially enhancing the accuracy of the analysis states Xa- However, these methods frequently show lim- itations in dealing with different observation operators. For example, DiffDA's repaint technique assumes that observations are aligned with a grid, whereas in real-world scenarios, observations are typically very sparse and likely fall outside the grid. Similarly, SDA relies on certain assumptions apply diffusion posterior sampling, and experimental results indicate that these assumptions can undermine assimilation accuracy, particularly in cases where observations are sparse or noisy.
In this paper, we aim to develop a neural network- based data ussimilation algorithm that not only cap- tures the non-Gaussian characteristics of the condi- tional background distribution for enhanced accu- racy but also effectively assimilates data under real- world observations (sparse and outside of the grid). We introduce VAE- Var, a novel data assimilation al- gorithm in which variational autoencoder is first employed to learn the conditional background dis- tribution and then the decoder component is utilized to construct a variational cost function, which, when optimized, yields the analysis states. The advantages of our approach are outlined as follows:
Table 1:
Comparison among traditional DA method (Trad-DA), diffusion-basee DA (NN-DA) and VAE-Var
Structure of p(x[x;)
Observation
Operator
Trad-DA NN-DA VAE- Var
Gaussian non-Gaussian
arbitrary
limited
non-Gaussian
arbitrary
This algorithm inherits the framework of traditional variational assimilation by explicitly modeling the posterior probability function p(x[x>, y) and maximizing it to derive the anal- ysis states. As result, compared to other neural network data assimilation methods such as SDA and DifIDA, VAE-Var can better handle different types of observation operators, particularly irregular observations that do not fall on the grid points of the physical field.
Unlike traditional variational assimilation algorithms, VAE- Var alleviates the dependence on expert knowledge for constructing the conditional background distribution, enabling the model to effectively capture non-Gaussian structures.
Experimental results on the Feng Wu weather forecasting system of 0.25% resolution demonstrate that VAE-Var achieves higher assimilation accuracy than both DifIDA and two traditional algorithms for most variables in the sparse observational settings. Further- more, when integrated with Feng Wu, it can reliably assimilate real-worke GDAS prepbufr observations over one-year period.
2 PRELIMINARIES
Variational Assimilation Variational assimilation seeks to determine the probability density func- tion p(x[xx.y) of the physical state's distribution at specific time, given known observational conditions y and the background state Xs, and calculate the analysis state x, by maximizing this probability density function, that is, x, arg maxx p(x[x,,y).
2

---

Published as a conference paper at ICLR 2025
Since X; is derived from weather forecasting model and y from observatories and satellites, it is reasonable to assume that X5, and y are independent. Under this assumption, the following equation holds, with the derivation provided in the Appendix:
arg max p(x[Xb: y) = arg max p(x|x)((xxxx))
(1)
We define the variational cost as:
(xx) =
log logply/x)p(xxxx) = - ply x) - log p(x(
(2)
Maximizing the probability density function is equivalent minimizing the variational cost. This cost comprises two terms: an observation term, (x. log p(y[x) and a background term,
(x Xx) log p(x(xx).
Observation Term In this paper, only consider observations where the observed variables consistent with the physical field variables, and exclude more complex observations like brightness temperature. In this case, the enror of the observation operator be simplified distribution with diagonal covariance matrix (Kalnay) 2003), that is, y/x ~ (H((x, R), where R corresponds to the variance matrix. The observation term then be formulated as follows:
Gaussian
1 Lx y) = 2 (x H(x)) (((x)))R=-(( H(x)).
Background Term As previously mentioned, the background state X₁, is derived by integrating a forecasting model from previous state. In most cases, the forecasting model cannot perfectly cap- ture all dynamical information and is subject to inherent errors. Traditional variational assimilation algorithms assume that these errors are independent of x; and they follow Gaussian distribution, that is, x X2 ~ (O, where B is the background error covariance matrix. The background term then be expressed as:
(x(x,x)) = _(x (x- = (x xx)).
(4)
In global weather systems, the dimensionality of physical states can be extremely large. For in- stance, in the FengWu forecasting model with 37 vertical layers (Chen et al. 2023), the dimen- sion physical fields is 189 x 721 x 1440 xe 2 x 10a, making the dimension of B approx- imately 4 x 1016. Explicitly constructing such a large matrix is generally impractical. To ad- dress this, a common approach is to decompose B as B = UU and apply variable trans- formation, x Uz + Xs. With this transformation, the background term can be rewritten as: Ls(2) Ls Xb: Xb) fizzz,, and the observation term becomes: L.(z) Lo X y). By minimizing L(z) = (z) + (z)
1/ZT2 + (U + x),y), it eliminates the need to ex- plicitly calculate the inverse of B. In traditional 3DVar, the sparse matrix U is derived by utilizing expert knowledge to statistically analyze background error samples, which are constructed using the NMC method (Descombes et 2015).
VAE-VAR
3.1 BACKGROUND ERROR ESTIMATION BASED ON VAE
Reviewing the process of traditional variational algorithms, we can identify two key factors for their successful implementation: First, it assumes that the error of the background states x X; is independent of the background states Xs, which enables the explicit construction of the expression p(x[x;). Second, it assumes that the background states' error follows Gaussian distribution and constructs linear mapping U through expert knowledge, which can map variable z from the standard normal distribution to the background states' error distribution (U (U-1(x- - X1) ~ (O I)).
In our study, we choose to maintain the first assumption. This decision is based on the fact that, in real-world scenarios, the dimensionality is often very high while the available training dataset is relatively small. If we were to eliminate this assumption, there is significant risk that the net- work would learn specious correlations between x and x;, a common issue in data assimilation. Conversely, opt to discard the second assumption, the high non-linearity of the forecasting
3

---

Published as a conference paper at ICLR 2025
model is likely to result in Gaussann error distribution. An illustrative example is provided in the Appendix for further clarification. If the error distribution is non- nan-Gassiinn then a non- non-linear mapping (denoted as F) is needed to map the standard normal distribution the background states' error distribution (J-1(xx Xb) ~ NO, I)).
There three primary types of candidate network structures suitable for this task. The first option the normalizing flow (Rezende & Mohamed| 2013 Kingma & Dhariwal 2018). which estab- lishes deterministic mapping between the standard normal distribution and target distribution. However, its high computational cost and large parameter size make it impractical for learning high- dimensional global weather fields. The second option the diffusion model. which has strong fitting capabilities, but introduces multiple rounds of sampling during inference, and the transformation be- tween the Gaussian and target space is not straightforward track. The third and preferred choice the variational autoencoder (VAE) (Kingma & Welling, 2013). VAE exhibits strong fitting abilities and has been successful on high-dimensional data (Doersch 2016). While VAE creates stochastic mapping from latent Gaussian variable to the target space, the core of the mapping (i.e., the de- coder) is deterministic, making it easier to formulate the variational cost. Thus, we adopt the VAE to perform the transformation.
A VAE typically comprises two components: an encoder & and a decoder D. The encoder maps the input data into mean and variance, from which sample Gaussian distribution. This sampled output is then fed into the decoder. The goal of the VAE is to learn the structure of the data space. Once the VAE well-trained,
simply sample from standard normal distribution, pass it through the decoder, and the output will resemble a sample from the original dataset. Essentially, the decoder the VAE good choice serve the non-linear mapping J for transforming the latent space to the data space.
3.2 GENERAL FORMULATION OF VAE-VAR VAR
Training
Errrr samples
forecast for time periods of 2r (lefi) and (right)
Reanalysis fields at time (4: - -) (left) and {4) (right)
then obtain the difference
Hidden states
Reconstructed samples
VAE Encoder
Mean
sampling
VAE Decoder
Std
Assimilation
Background field
Analysis field
Hidden state (Initialized as zeros)
Analysis increment
Observations
Obs Loss
VAE Decoder
BG Loss
calculate gradient with auto-differentiation mww.oeeeeoeeeee.................
Loss
update the hidden state with L-BFGS
Figure 1: VAE- Var Framework. The upper half panel demonstrates the training phase; the lower half panel demonstrates the assimilation phase. Please refer to the main text for detailed explanation.
Here, we propose our VAE- Var framework for data assimilation. As shown Figure it comprises two phases: training and assimilation.
4

---

Published as a conference paper at ICLR 2025
Training Set Construction During the training phase, adopt the traditional NMC (National Metcorological Center) method to generate historical error samples, as outlined in Algorithm ! The reanalysis data are treated the ground truth of states. We select historical reanalysis states at intervals of т, then integrate them for a time period of r and 27 respectively, ensuring that both are predictions for the same timestamp. By subtracting these states, generate error sample. Repeating this process across a long period of reanalysis data allows us to construct a training set.
Algorithm Training Set Construction
Algorithm VAE-Var Assimilation
Require: Forecasting model M, sample number N, time gap r for i from 1 to N do
Require: : Trained decoder D, background state Xt.
observation y
2 5 0
(z(z) + f
Randomly pick reanalysis data with
an interval of r: Xo, X
[z(2) + L. (D(z) +b,
X),
X1 f MM-22-(x, (x-) X2 4 M Ma (Xo)
Minimize L(z)
(z(z)
+ (z( with L-
BFGS (Liu & Nocedal 1989
Paszke et al. 2017)
Add & &2 to the training set
and get the minimum point
end for
X, D(z") + Xb
Loss Function We employ the standard approach for training the VAE, where the loss function consists of two components: the reconstruction loss (Lossree) and the Kullback-eeillrr (KL) di- vergence (L)sski.). The reconstruction loss quantifies how accurately the decoder can regenerate the original input from the latent space, which, in our case, is defined as the mean squared error be- tween the input and its reconstruction. The KL divergence ensures that the learned latent distribution aligns closely with a standard normal distribution. These two terms together guide the optimization process, enabling the VAE to capture the underlying data distribution effectively. Additionally, a hyperparameter is introduced to balance the reconstruction loss and the KI. divergence, and the total loss is expressed as Loss J Lossyec + LOssKL.
its decoder implement VAE TAE-Var assimilation, as de- scribed in Algorithm 2 Similar to traditional algorithms, VAE-Var aims to optimize variables in the latent space, denoted by The key difference that the transformation from the latent space the physical space is non-linear, i.e., x = D(z) + Xs. where D(z)) represents the decoder output. As a result, the observation term is modified accordingly, expressed as Cz(z) = L. (D(z) y). It important to note that, from a rigorous mathematical perspective, when the mapping between the latent space and the physical space is nonlinear, the transformation of the probability density func- tion introduces an additional non-constant determinant term in the background field, which is almost computationally intractable. To account for this, we empirically scale the original background term 1 T by a positive parameter À for proper compensation. We provide a further explanation on the formulation of the background term in the Appendix. For any given observation y. as long as the observation operator H is constructed in an auto-differentiable way, which will be explained in Sec- tion then all components of [(z) are differentiable, and we can leverage anto-differentiation to ad(z),
Assimilation Once VAE is trained,
perform back-propagation and compute the gradient 8z with the parameters of D fixed. Follow-
ing this, the L-BFGS algorithm can be used to find the minimum of [(z).
Algorithm Cyclic Forecasting and Assimilation with VAE-Var
Require: Initial state x,"
forecasting model M, observations, steps L, assimilation cycle T. Initialize the time stamp
1: --00 0
2: for slep from 0 to L do
3:
(t) t arg min, C 2 x
(t) yr(t)
D Obtain the analysis control variable
4: (1) D(z( (1) +
Recover the analysis states. D Obtain the background at the next time step.
[t+T) + Me-Mr(xxx!
5:
x,
6:
++
7:
end for
5

---

Published as a conference paper at ICLR 2025
Cyclic Forecasting and Assimilation VAE-Vir and the weather forecasting model can be coupled to realize the cyclic forecasting and assimilation, similar to Feng Wa (Xiao et al.) 2024). Denoting E (-|xp. y) the objective function with respect to observations y and the background states X6, and using the superseript represent time steps, the coupling of VAE-Var with model forecasts implemented as shown in Algorithm 3
4 RESULTS
4.1 EXPERIMENTAL SETUP
Forecasting Model Our experiments utilize FengWu (Chen et al. 2023), prominent data-driven global medium-range weather forecasting model. Feng Wu simulates five atmospheric variables (each across 13 pressure levels the version we use) and four surface variables, leading to a total of 69 predictands. In this study, the atmospheric variables (q). zonal wind component (u), meridional wind component (v), and air temperature (t); the 13 sub-variables across various vertical levels are denoted using abbreviations of their short names and pressure levels (e.g.. z500 represents geopotential height at 500 hPa). The four surface variables 2-meter temperature (12m), 10-meter zonal wind component (u10), 10-meter meridional wind component (y10), and mean sea level pressure (mslp). We test spatial resolution of 721 x 1440. Following the convention of the original paper (Chen al. 2023), the six-hour forecasting model is trained using the ERAS dataset from 1979 to 2015.
geopotential (z). specific humidity
VAE Structure and Training Our approach allows considerable flexibility in selecting the neural network structure for VAE. Rather than designing new architecture from the ground up, opted to adapt the well-established FengWu structure, which primarily consists of multiple "Swin Trans- former v2" (Liu et al. 2022) networks and is proved to be capable of capturing complex spatial and temporal relationships. In our experiments, we employ two Feng Wu networks to construct the encoder and decoder, shown in Figure 2 As for the encoder, the channel number of "Mean_i" and "Std_i" is set to (the original channel number is 13) for 1 < S. 5 and the channel number of "Mean 0" and "Std 0" is set (the original channel number is 4); this guarantees that the latent space has smaller dimensionality than the original space. The decoder is almost symmetric to the encoder, with the key difference being that the input layer channel number half of the encoder's output layer channel number. We use ERA5 reanalysis data from 1979 to 2015 to train the VAE model, and the loss weight set to 2.0.
Encoder
Decoder
Geopotential
SF-v2
SFav2
Mean I
Latent Var
SF.v2
SF-x2
Geopatential
Swin Transforerr v2 Mocia Fuser)
I
Swin Tramsformer (Croes Modal Fuxee)
Humidity
SF-x2
SF-x2
Meun
Latent
SF-v2
SF-v2
Humidity
Std 2
Mean
SF-v2
Var
V_ Wind
SF.v2
SF.v2
Mean
Latent Var
SF-v2
SF-v2
V. Wind
Sid
Temperature
Mean
Latent
SF-v2
Temperature
Std
Surface Variables
SF-x2
SF-v2
Mean 0
Latent Var. 0
SF42
SF-v2
Surface Variahles
Input
Merge
Split
Sampling
Merge
Reconsiruction
Figure 2: VAE Structure for the Feng Wu forecasting system. "SF-v2" is abbreviated for "Swin Transformer v2". Both the encoder and the decoder adopt the Feng Wu backbone structure.
Observation Setup Two types of observations are considered our experiments. The first type involves simulated observations from the ERA5 (Hersbach et al. 2020) dataset. Following Huang ct al. (2024), we apply sparse grid column masks to mimic the sparse distribution real-world observations. For instance, if we have 1000 observation columns, only 1000 points from the 721 x

---

Published as a conference paper at ICLR 2025
1440 grid are observed, and when grid point is observed, we have data for all variables at all height levels. The second type uses observations from the Global Data Assimilation System (GDAS) prepbufr dataset (Rodell et al.) 2004), which provides data from various observing systems and instruments. In the second experiment, the observations may fall outside the grid points and are typically sparse across height levels, which closely resembles operational scenario.
Cyclic Forecasting and Assimilation Setup The initial state for our cyclic forecasting is derived from the ERA5 dataset. In our experiments, the forecasting system begins at 00:00 on January 1. To generate the background states x we begin with the ERA5 reanalysis field at 00:00 on December 30 of the previous year, run the 6-hour forecasting model for eight steps, and use the resulting fields of variables the starting point for the cyclic forecasting process. The assimilation cycle T equals six hours, consistent with the pre-trained six-hour Fengwu forecasting model.
Baselines VAE- Var is compared against three baselines: DifIDA, interpolation and a traditional variational method (3DVar). For DifIDA, draw from the results presented in the original pa- per (Huang et al.] 2024); the interpolation method employs bilinear interpolation for station-based data; the 3Dvr method implemented following the approach outlined in Descombees al. (2015)
4.2 ERAS-SIMULATED OBSERVATIONS
In the ERA5-simulates observational setting, all observations are positioned on the grid. Thus, the observation operator 7 defined as a mask operator, where the values at the observation locations set one, and the values locations without observations with element-wise matrix multiplication in PyTorch, and it is clearly differentiable. The observation covariance matrix R is assumed be diagonal, with the square root of each entry set to 0.1 times the standard deviation of the respective variable. The parameter À is set to 4.0.
set zero. We implement H
Fixed Observations First, we examine the scenario where the observation positions remain con- stant over time. To align with DiffDA, initiate our forecasting system January 2022, sim- ulate it for 15 days, and conduct experiments with four different observation amounts: 1000, 2000, 4000, and 8000 columns. We calculate the root mean square errors (RMSE) of the analysis states at various time steps to assess assimilation accuracy. The results are illustrated in Figure B which indicate that VAE-Va Var consistently outperforms all three baseline methods, across all three demon- strated variables and the four observation amounts. Results for other important variables are left in the Appendix. Notably, when the observation amount is relatively low (1000 columns), VAE Var demonstrates marked advantage over the other two algorithms, highlighting its strong capability in leveraging background information. As the observation amount increases, VAE- Var maintains superior performance, albeit with a reduced margin, since the significance of effectively utilizing background information diminishes with more observations.
Unfixed Observations Next, we conduct experiments where the observation positions vary over time, while keeping the observation amount constant. The results, presented in Figure are consis- tent with the "fixed observation" case. VAE-Va Var continues to outperform DifIDA, the interpolation method and 3DVar across all demonstrated variables and all observational settings.
4.3 GDAS PREPBUFR OBSERVATIONS
In this experiment, we aim to assimilate real-world observations. The observation dataset we choose is the GDAS prepbufr dataset. This dataset consists of a global collection of surface and upper air reports gathered by the National Centers for Environmental Prediction (NCEP). It includes obser- vations from land and marine surfaces, radiosondes, pibals, aircraft, and Global Telecommunica- tions System (GTS) reports. also incorporates data from wind profilers, U.S. derived winds, SSM/I oceanic winds and TCW retrievals, as well as satellite wind data provided by the National En- vironmental Satellite Data and Information Service (NESDIS). We utilize this observational dataset because it is comprehensive and directly applied in today's operational weather forecasting systems.
7

---

Published as a conference paper at ICLR 2025
obs 1000 (0.10%)
obs
2000 (0.19%)
obs
4000 (0.39%)
obs amount (0.77%)
500
VAE-Var
201
Inurp-lation
1DV
200
I&SII RMSB
2.-
1.5
10-
0.5
RMNE
:00
200
100
200
100
200
300
100
200
500
Simslated time (hours)
Simulated time (hours)
Simudated time (hours)
Simulated time (hours)
Figure 3: Results for fixed observation positions (15 days). The system is simulated an auto- regressive manner for 15 days, starting from January 2022. Following DifIDA, we demonstrate the RMSES of the analysis field for three variables (z500, t850, 12m) three different rows.
obs amount 1000 (0.10%) obs amount 2000 (0.19%)
obs amount 4000 (0.39%) obs amount 8000 (0. 77%
(mi 1x)
/
Intep-laation
3D
102
RMSE
20
1.0
300
200
10c
200
100
300
100
200
503
12m RMSE
200
206
300
10c
200
300
10c
300
300
10c
200
303
Simnlaate time (hours)
Simulated time
Simulated time (hours)
Simulated time (bours)
Figure 4: Results for unfixed observation positions (15 days). The system is simulated in auto- regressive manner for 15 days, starting from January 2022. Following DiffDA, we demonstrate the RMSEs of the analysis field for three variables (z500, t850, 12m) in three different rows.
Data Processing In the GDAS prepbufr dataset, each message represents single measurement taken by an observational instrument at a specific time and location. Since our assimilation al- gorithm focuses on specific time point, we only select observations within half an hour of the assimilation time for processing. For instance, if the current time is 3 PM, we choose observations taken between 2:30 PM and 3:30 PM. Moreover, the spatial locations of observations are often not located on the grid of Feng Wu. To address this, we construct new, finer grid for the observations and map the observations to the nearest grid points. Theoretically, this grid can be fine enough to minimize the error introduced by nearest-neighhor approximation. As a trade-off between accuracy and computational cost, the grid is designed to have the same latitude and longitude resolution as the FengWu physical grid, but with 40 vertical layers, which is much denser than FengWu. Also, applicd quality control after gridding the data, filtering out observations that deviate substantially from the corresponding ERAS values.
8

---

Published as a conference paper at ICLR 2025
Observation Operator After data processing, all observations are placed grid with finer resolution than FengWu. The observation operator H is then defined as the interpolation from the Feng Wu grid to the observation grid, which can be implemented using differentiable 'Linear' layer in py Torch. Details of this implementation are provided in the Appendix.
5500 ground truth (m /s)
5500 background (m is')
z500 obs mask, Hobs-2009 2009 (0.194%)
50000
$5000
50000
55000
Background error (RMSE-T) 74.3 m//s)
z500 analysis increment (m 1s²)
Analysis error (RMSE- 16MMSE-66..1 56.9m'15 153
500
-500 500
-500
Figure 5: Visualization of z500 assimilation results at time 2017-01-01 00:00:00. ERAS ground truth, the background field and the observation mask are demonstrated in the first row; the error of the background field (background field minus ground truth), the analysis increment (analysis field minus background field) and the error of the analysis field (analysis field minus ground truth) shown in the second row. Strictly, the figure titled "z500 obs mask" correspond to observations at 501hPa because the observations use different height axis. Please refer to the Appendix for details.
Analysis Field Visualization We select the year 2017 for conducting the assimilation experiment because the completeness of the observational data highest for that year. In Figure 5 visualize the assimilation results at time 2017-01-01 00:00. At this time, there are a total of 375,371 obser- vations, which are reduced to 142,647 (approximately 0.2% of the total grid points) after quality control. From the visualization of the observation mask, it is evident that the observations are pri- marily concentrated over Eurasia, North America, and the South Atlantic Ocean. The observations over the first two regions are mainly derived from radiosondes launched from the ground, while those over the South Atlantic come from retrievals based on geostationary satellite data. The anal- ysis increment closely aligns with the observations, and in areas where the observations are more concentrated, the analysis increment is more pronounced. By comparing the errors between the analysis field and the background field, observe reduction in error after assimilation, which highlights the effectiveness of our method.
Cyclic Forecasting and Assimilation Results 6 the RMSE and Bias of the analy- sis fields at various time steps when VAE-Var is coupled with FengWu for a year-long simulation. In the GDAS prepbufr observational settings, it common to encounter positions where only a subset of variables is observed. Since the DifIDA is currently unable to address such cases, we compare VAE- Var with the interpolation and 3DVar baselines this section. The results show that for all eight reported variables, VAE- WAr-Var outperforms both methods in RMSE and Bias, demonstrating that our approach effectively reduces both error and systematic bias in the cyclic forecasting system.
4.4 COMPUTATIONAL COST
In VAE-Var, the linear transformation U in 3DVar is replaced with nonlinear D. The number of floating-poit operations in D is no samller than in U, so from a floating-pont computation perspec- tive, the computational complexity of VAE Var is not lower than that of 3DVar. However, VAE-Var Var's main advantage its easy implementation on GPUs, which significantly improves overall assimila- tion efficiency. For example, on single 1100 GPU, one cycle of assimilation takes approximately

---

Published as a conference paper at ICLR 2025
2500 (m')s') '
1500 (K)
u500 (m/s)
M (m/s)
1000
lexppelazin
3010
100
200
100
1i0
20m
san
100
30
100
102
Sinmllated (days)
Simmlaaed (days)
Simulated time (days)
Stumllated time (days)
u10 (m/ss
v10 (m/i))
12m (K)
mesp (Pu)
www.mmm
4003
200
60
62
260
110
200
sio
160
200
102
Sanulated (days)
Simulated (dayys)
Simmlated (days)
Stumllted (days)
Figure 6: Results for GDAS prepbufi observations (one year). The system is simulated in an auto-regressive manner for one year, starting from January 1, 2017. The y-axis represents RMSE and Bias with respect to the ERA5 ground truth. Due to the varying amounts of observations over time, the RMSE for the interpolation method experiences significant fluctuations. The curves have been smoothed to provide better visualization.
18 seconds. In contrast, 3DVar is typically implemented on CPUs, where one cycle of assimilation takes several minutes (Smith et al. 2014). Therefore, in practice, VAE- Var can easily leverage GPU capabilities to achieve assimilation one order of magnitude faster than traditional 3DVar.
RELATED WORK
The concept behind our work is similar to "latent space data assimilation".. For instance, in Amen- dola et al. (2020). an LSTM used to map physical states into latent space, where observed data are assimilated using a Kalman Filter. Similarly, Peyron et al. (2221) and Fan et al. (2025) employ autoencoders construct the latent space and have achieved promising results. In Meline & Zaplot- nik (2023), a VAE is trained for learning the map, with experiments conducted on the ERA5 dataset. However, the observational settings and forecasting models these works are overly simplified and the effectiveness on more complicated systems remains to be evaluated.
While these approaches focus on finding the latent space of the physical fields, the critical aspect in data assimilation is not the physical fields themselves but the uncertainty (or error) of these fields. In contrast, the VAE- Var algorithm we propose directly learns the latent space of the error fields, offering straightforward yet effective method for data assimilation. Moreover, our algorithm is proved effective on a global-scale AI forecasting model with real-world observations.
CONCLUSION
In this paper, we introduce a novel data assimilation algorithm, VAE Var, which leverages varia- tional autoencoder to model the non-Gaussian structures of the background enror distribution while utilizing the traditional variational methods for effective assimilation real-world irregularly sam- pled observations. Experimental results ERAS-simulated observations show that our method sur- passes both an AI data assimilation algorithm, DifIDA, and two traditional methods (interpolation conducted experiments using real-world observations, demonstrating
and 3DVar). Additionally, that VAE Var capable of performing cyclic forecasting and assimilation over the span of one year. VAE- Var is also shown have high computational efficiency GPU devices.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work supported by the National Natural Science Foundation of China (NO.U2242210), This work is also supported by the Shanghai Science and Technology Commission Project (23DZ1204704) and Shanghai Artificial Intelligence Laboratory.
We sincerely acknowledge the European Centre for Medium-Range Weather Forecasts (ECMWF) for providing the ERA5 dataset and the National Center for Environmental Prediction (NCEP) for providing the GDAS dataset, which are instrumental in this study. Their efforts in data collection, archiving, and dissemination are greatly appreciated.
We also grateful the Research Support, IT, and Infrastructure team at the Shanghai Al Labo- ratory for their provision of computational resources and network support.
Additionally, extend our heartfelt thanks Dr. Ben Fei from the Chinese University of Hong Kong and Hang Fan from Shanghai AI Laboratory for their valuable assistance and insightful dis- cussions throughout this research. Their contributions have greatly enriched this work.
REPRODUCIBILITY STATEMENT
All experiments in this paper are fully reproducible. The code of VAE-Var is available Var. For the baseline algorithms, three baselines are considered: DifIDA, 3DVar, and interpolation. Regarding DifID,, we directly adopt the results from the original paper. For 3DVar, the construction details of the background error covariance matrix are provided in the Appendix. As for the interpolation method, we clarify that it is based on the bilinear inter- polation algorithm. Additionally, all the datasets we use, including ERA5 and GDAS prepbufr, are available online.
http://gituub.com/iioiilll
REFERENCES
Maddalena Amendola, Rossella Arcucci, Laetitia Mottet, Cesar Quilodran Casas, Shiwei Fan, Christopher Pain. Paul Linden, and Yi-Ke Guo. Data assimilation in the latent space of neural network. arXiv preprint arXiv:2012. 22../2200 2056, 2020.
Ross N Bannister. A review of forecast error covariance statistics in atmospheric variational data assimilation. Characteristics and measurements of forecast error covariances. Quarterly Journal of the Royal Meteorological Society: A journal of the almospheric sciences, applied meteorology and physical oceanography, 134(637):1951-1977 2008.
Dale M Barker, Wei Huang, Yong-Run Run Guo, AJ Bourgeois, and QN Xiao. A three-dimensional vari- ational data assimilation system for mm5: Implementation and initial results. Monthly Weather Review, 132(4):897-91 2004.
Peter Bauer, Alan Thorpe, and Gilbert Brunet. The quiet revolution of numerical weather prediction. Nature, 525(7567):4755, 2015.
Alberto Carrassi, Mare Bocquet, Laurent Bertino, and Geir Evensen. Data assimilation in the geo- sciences: An overview of methods, issues, and perspectives. Wiley Interdisciplinary Reviews: Climate Change, (5))e533, 2018.
Kang Chen, Tao Han, Junchao Gong. Lei Bai, Fenghua Ling, Jing-Jia Luo, Xi Chen, Leiming Ma, Tianning Zhang, Rui Su, et al. Fengwu: Pushing the skillful global medium-range weather forecast beyond 10 days lead. arXiv preprint arXiv:23044 02948, 2023.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L. Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 14687, 2022.
Gael Descombes, T Auligné, Francois Vandenberghe, DM Barker, and Jerome Barre. Generalized background error covariance matrix model (gen_be v2. 0). Geoscientific Model Development, (3):669-696, 2015.
Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05990 2016.
11

---

Published as a conference paper at ICLR 2025
Hang Fan, Ben Fei, Pierre Gentine, Yi Xiao, Kun Chen, Yubao Liu, Yongquan Qu, Fenghua Ling, and Lei Bai. Physically consistent global atmospheric data assimilation with machine learning in latent space. arXiv preprint 250.v:2502002888 0282.02884. 2025.
Mariano Giaquinta and Giuseppe Modica. Mathematical analysis: An introduction to functions of several variables. Springer Science & Business Media, 2010.
Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoo--Saater,, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, et al. The era5 global reanalysis. Quarterly Journal of the Royal Meteorological Society, (197733):19992299) 2020.
Langwen Huang, Lukas Gianinazzi, Yuejiang Yu, Peter D Dueben, and Torsten Hoefler. Diffda: diffusion model for weather-scale data assimilation. arXiv preprint WxXiv:2401.05922 2024.
Xiang- Yu Huang. Qingmong Xiao, Dale M Barker, Xin Zhang, John Michalakes, Wei Huang, Tom Henderson, John Bray, Yongsheng Chen, Zaizhong Ma, et al. Four-dimensional variational data assimilation for wrf: Formulation and preliminary results. Monthly Weather Review, 137(1): 299-314, 2009.
Eugenia Kalnay. Atmospheric modeling, data assimilation and predictability. Cambridge university press, 2003.
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational bayes. arXiv preprint
Durk P Kingma and Prafulla Dhariwal. Glow:
Generative flow with invertible 1xl convolutions. Advances in neural information processing systems, 31, 2018.
Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, 45(1):503-528, 1989.
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12009-12019, 2022.
Edward N Lorenz. Deterministic nonperiodic flow. Journal of atmospheric sciences, 20(2):130-141 1963.
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Lue Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11461-11471, 2022.
Boštjan Meline and Ziga Zaplotiik. Neural- data assimilation using variational autoencoder.
arXiv preprint arXiv:2308. /6073, 2023.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Mathis Peyron, Anthony Fillion, Selime Gürrl,, Victor Marchais, Serge Gratton, Pierre Boudier, and Gael Goret. Latent space data assimilation by using deep learning. Quarterly Journal tfthe the Royal Meteorological Society, 147(740):3759-377, 2021.
Florence Rabier, Jean-Noel Thépaut, and Philippe Courtier. Extended assimilation and forecast experiments with a four-dimensional variational assimilation system. Quarterly Journal of the Royal Meteorological Society, 1854550::166-1110
1998.
Florence Rabier, Heikki Järvinen, E Klinker, J-F Mahfouf, and A Simmons. The ecmwf opera- tional implementation of four-dimensional variational assimilation. Experimental results with simplified physics. Quarterly Journal of the Royal Meteorological Society, 1243564) 2000.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. Interna- fional conference machine learning, pp. 1530-1538. PMLR, 2015.
12

---

Published as a conference paper at ICLR 2025
Matthew Rodell, PR Houser, UEA Jambor, Gottschalck, Kieran Mitchell, C-J Meng, Kristi Arse- nault, B Cosgrove, J Radakovich, M Bosilovich, et al. The global land data assimilation system. Bulletin of the American Meteorological sociery, 85(3):381-394, 2004.
François Rozet and Gilles Louppe. Score-based data assimilation for two-layer quasi-geostrophic model. arXiv preprini arXiv:2310.01853, 2023.
Travis M Smith, Jidong Gao, Kristin M Calhoun, David J Stensrud, Kevin L Manross, Kiel L Or- tega, Chenghao Fu, Darrel Kingfield, Kimberly L Elmore, Valliappa Lakshmanan, et Ex- amination of a real-time 3dvar analysis system in the hazardous weather testbed. Weather and forecasting, 29(1):3377,, 2014.
Steven H Strogatz. Nonlinear dynamics and chaos: with applications to physics, biology, chemistry, and engineering. CRC press, 2018.
Yannick Trémolet. Accounting for an imperfect model in 4d-var. Quarterly Journal of the Royal Meteorological Society: A journal of the atmospheric sciences, applied meteorology and physical 2006.
oceanography, 132(621):2483-2555)
Yi Xiao, Lei Bai, Wei Xue, Hao Chen, Kun Chen, Tao Han, Wanli Ouyang, al. Towards a self- contained data-driven global weather forecasting framework. In Forty-first International Confer- ence on Machine Learning, 2024.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 DERIVATION OF EQUATION
We aim maximize p(x Xb. y), which can rewritten as:
p(x6, x(x)p(x) p(Xb; y)
arg max X y) arg max
arg max p(x., (
Here, p(x.. y) is independent of x and thus be omitted from the optimization.
Now, assuming that X; and y are independent conditional on x, we can decompose p(x6; y (x) as p(x; ( x). Substituting this into the equation gives:
arg max p(x| y) arg max =argmaat((xxx/)xxxx/////x
Next, applying Bayes' rule to o(x/|x), we have p(x)) x) = p(x(x))u(xx)/p(x) Substituting this back, and noting that p(xx) constant with respect x, obtain:
arg max p(x|Xx, arg max p(xxxx))(x)xx))
This is the right-hand side of Equation
which follows logically under the given assumptions.
A.2 NON-GAUSSIAN FEATURES OF THE BACKGROUND ERROR
We use the Lorenz 63 Lorenz 1963 system to provide an illustrative demonstration of the background error's non non- Gaussian features. This system involves three parameters: 0, B, as shown in Equation 5 We create dynamical models using two different sets of parameter values:
10, p 28, 3 (for ground truth model) and 10, P 29, 3 (for prediction model). By numerically integrating these models from randomly-chosen identical initial states and calculat- ing the difference between their outputs, generate set of error samples, depicted with the The error distribution is observed to be distinctly non-Gaussian and exhibits a
blue dots in Figure 7 non-convex structure.
02
dX o(V X)
dt
-0.
dy dt
-X(p 2) Y (5)
dz
XY - 3Z
dl
0.5
-02
-0.6
Figure 7: Distribution of the background error for the Lorenz 63 system.
A.3 MATHEMATICAL EXPLANATION OF VARIATIONAL COST
Although our method is generally heuristic, mathematical explanation can be provided under some assumptions. Suppose that the VAE is well leamt and the decoder is bijection map from the low-dimensional latent space to a manifold of the high-dimensional physical space. In this case, according the change of variables formula in Theorem (2.80) of Giaquinta & Modical (2010)
14

---

Published as a conference paper at ICLR 2025
(Chapter 2.5.2), supposing that the physical space has dimensionality of dz. and the latent space has a dimensionality of d2, then for any measurable set A _ RR the following equation holds:
J, px px(((z)) (DD) (z) dz -/o. px(x)H² (AnD (nx))((()) dH (x),
where Px correspond the probability density function of x, D D- ((x}) represents the preimage of the set under the mapping D, Ht corresponds to the Hausdorff s-dimensional measure, and (x) is defined as:
(aʃ(x) 0x
(as(x) 0x
((D))xx - Vdet
Intuitively, the same event should have the same probability measured in two different space. De- noting Ps the probability density function of z, the following "equation" holds:
(x(x) d7( (x) Pz(2) dz.
(8)
Combining these two equations above, we have:
J Px Px (DD) (z) dz / Pz( pz(z) (An {z)) dz.
is apparent that for z E A, 1tth (An {z}} holds; thus,
J Xx (x( (DD) (z) dz / Pz(z) dz.
(10)
Since this equation holds for any measurable set A. we have
(x((())) (DD) (z) = Pz(z).
(11)
Denoting 5 = Xb. Ps the probability density function of 5; then
0x
Po( (x)) det
Px(x).
(12)
06
Further, we have:
ps(5) = Px(z) (DD) (x))-²
[OD(x)
(8D(z)
-1/2
(13)
(z(z)
det
Oz
Oz
Since, the latent space of VAE is assumed to follow the standard Gaussian distribution, therefore, p(z) I). Then,
(aD(z)
@D(z))
[x(2) logg((((x)) (8)
log Pz(2) log det
Oz
Oz
(14)
@D(z)
@0D(z)
2 2. +
log det
2
Oz
Oz
Compared to the formula used in the paper, the strict expression includes an additional determinant term, which represents the scaling factor of the volume element when mapping from the latent space to the physical space.
Since Transformer-basee neural networks are highly complex, nonlinear, and high- dimensional systems, this scaling factor can fluctuate unpredictably over time. Including this term would make the assimilation objective function extremely difficult to optimize. For this reason, empirically omitted the determinant term from our formula. As compensatory measure, introduced an adjustable parameter in the background term of the assimilation objective function, that is, (z( 1 Tz. Since the observation deviation term L, (D(z) Xb, y) is calculated in
Simplification
15

---

Published as a conference paper at ICLR 2025
the physical space, and the background error term Lo(z) 2 is calculated in the latent space, the former would intuitively be d=
2
times larger than the latter. To balance this dimensionality
difference, the parameter À should be approximately equal to the physical space has dimensions dæ 69 721 x 1440 and the latent space has dimensions d₂ = 32 x 721 x 1440. This gives À = 4.64, and we round this value to 4 for simplicity.
dx d=
In our experimental setup,
We admit that such a simplification not mathematically rigorous, but it greatly reduces the opti- mization difficulty and is experimentally proven to be effective.
Rigorous Consideration of Jacobian Determinant In Bayesian modeling, one class of ap- proaches rigorously addressing the Jacobian determinant of the mapping from a Gaussian prior to the target dataset distribution is normalizing flow (Rezende & Mohamed] (015). These methods constrain the structure of "D to ensure the Jacobian determinant computationally tractable. Typical convolution layers, single-variable activation functions, in-
constraints include the use of vertible transformations with simple Jacobian structures. For instance, Glow (Kingma Dhariwal 2018 is well-known example in this family of methods.
While normalizing flow approaches retain the universal approximation property, they face significant challenges when applied to high-dimensional systems, such as weather modeling. Specifically, the structural constraints lead to substantial computational complexity, making training time-intensive and sesource-emmanding. To knowledge, few success has been achieved in modeling high- dimensional systems using normalizing flow.
In future work, will explore incorporating the advantages of normalizing flows into VAE-Var, with a focus on addressing the scalability and efficiency challenges in high-dimensional meteoro- logical systems.
A.4 GDAS PREBUFR OBSERVATION HEIGHTS
The potential height of the observation coordinate axis is evenly distributed logarithmically from 50 hPa to 1000 hPa. Specifically, the coordinate axis includes the following 40 potential height levels:
50 hPa 58 hPa, 63 hPa, 68 hPa, 73 hPa, 79 hPa, 86 hPa, hPa, 100hPa, 108 hPa, 116 126 hPa, 136 hPa, 147 hPa, 158 hPa, 171 hPa, 185 hPa, 199 hPa, 215 hPa, 232 hPa, 251 hPa, 271 hPa, 293 hPa, hPa, 341 hPa, 368 hPa, 398 hPa, 430 hPa, 464 hPa, 501 hPa, 541 hPa, 584 hPa, 631 hPa, 681 hPa, 735 hPa, 794 hPa, 858 hPa, 926 hPa, 1000 hPa
In contrast, the potential height axis of the physical space coordinate axis of FengWu includes the following 13 levels:
50 hPa, 100hPa, 150 hPa, 200 hPa, 250 hPa, 300 hPa, 400 hPa, 500 hPa, 600 hPa, 700 hPa, 850 hPa, 925 hPa, 1000 hPa
A.5 IMPLEMENTATION THE 3DVAR BASELINE
The implementation of the 3DVar baseline follows the GEN BE method from Descombes et al. (015) and we reproduce it with Python. Generally speaking, GEN BE consists of three steps. First, 1000ERAS fields are sampled between 1979 and 2015; then the NMC method is applied to calculate the resulting background error samples; last, calculate the statistics of U. In GEN BE, U is assumed to decomposed into four operators,
U,SU Un, where Up handles the transformation of physical variables, S represents the standard deviation, U, corresponds to vertical transformation, and U, handles horizontal transformation.
The details for constructing the four operators are follows:
Implementation of Up Up is the only transformation that does not require parameters derived from the NMC samples. It is used to convert the stream function and potential function (ø. t)) into wind components (u, v) in the latinueellongiude directions, while keeping other variables un-
16

---

Published as a conference paper at ICLR 2025
changed. Specifically, it is computed using the following equations:
Oup
Ob
dy
Ox
Ox
dy
Implementation of S S represents the error variance for each variable layer. To calculate it, we first apply the inverse transformation of Up all error samples (i.e., converting winds from the latinude-onggiude direction back to stream functions and potential functions). This is done using spectral methods, which are supported by the 'torch-harmonic library in Py PyTorch for solving the differential equations. After transforming the fields, we compute the sample variance at each pixel. The resulting variances are averaged in the latitude-longiude direction, yielding 69 values corre- sponding to the error variance at each vertical layer.
Implementation of U, U, is computed using eigenvalue decomposition. Specifically, for each upper-air variable, first average the variable in the latitude-longitue direction obtain 1000 69-1ength vectors. The covariance matrix A of these 1000 vectors is then computed. We perform eigen-decomposition on such that A PAP-1, and U, is then defined as U, PA where P is the eigenvector matrix and A is the diagonal matrix of eigenvalues.
Implementation of U Un is essentially a horizontal recursive filter applied to each of the 69 layers. The length scale L for the filter is given by:
8 Variance (Ø) 1/4
L
Variance V2
This filter is applied to each layer individually to perform horizontal smoothing.
To summarize, Up is a fixed, parameter-free variable transformation operator that converts stream and potential functions into wind components; S represents
variance with 69 parameters cor- responding to each vertical layer; U, corresponds to eigenvalue decomposition applied to the 1000 samples, yielding 69*69 + 69 parameters; U, is horizontal recursive filter applied to each of the 69 layers, with 69 parameters.
A.6 IMPLEMENTATION OF THE OBSERVATION OPERATOR
As for the real-world observations, the observation operator aims vertical interpolation from the physical grids to the observational grids. Since the geopotential height changes approximately exponentially with physical height, logarithmic linear interpolation required when performing vertical layer interpolation. The computation of the interpolation matrix P E 40 is shown in Algorithm 4 where X represents the scale of the geopotential height at the physical grid points and Y represents that the observational grid points. After calculating the interpolation matrix, the PyTorch construction of the observation operator is shown in Algorithm
Algorithm Logarithmic Interpolation Matrix Construction
Require: Vectors X (scale the vertical axis of physical space) and Y (scale of the vertical axis of the observational space)
P + zeros(len(Y), len(X)
for i 0 to (Y)
do
for jo to len X) - do
if Y X [j] then
Pi.jj] <-
else if [
X and Y < X then
P[i.j] < Tog(X -o
log(X
P 1] + box
end if
end for
end for
return P
17

---

Published as a conference paper at ICLR 2025
Algorithm 5 Observation Operator Implementation in PyTorch
Require: Physical layer æ E 1000721x11440 interpolation matrix P
2o - I.unsqueeze(())
f-
y yaa 4]) for i = to 5 do
mat zoli, + 13 4- 13
cetuunttioaal.inaa((nne. http:///.iiiemmmmmm.rooossss//// transposee1 P) ttranspooe(1,.) (mat)) end for torch.cat(1)) 1) < y-quuezee() return y
mat +
A.7 RELATIONSHIP WITH SPATIO-TEMPORAL DATA ASSIMILATION METHODS
The framework of VAE-Var is compatible with traditional spatio-temporal data assimilation meth- ods, such as 4DVar. In 4DVar, the observation term can be expressed as
N-1 Σ yy Mo R-- (y: -M-ssx
Co (x, (y))
(15)
where {yr are observations within an assimilation window and Mo are forecasting models for representing the flow dependency. As long as the forecasting models are auto-differentiable (e.g.. implemented by neural networks). VAE- Var can be easily extended to VAE- Var by substituting the original L. in Algorithm with the expression above. Experiments have not been conducted due to the limitation of computational resources and detailed analysis on VAE-Var's ability of doing spatio-temporal assimilation will be conducted in our future work.
A.8 ABLATION STUDY
Different Initial States To investigate how the change of initial state x would affect the perfor-
mance of VAE-Var, we construct seven initial fields to conduct the ablation study:
1. Start from ERAS A5 24 hours earlier and use four 6-hour FengWu forecasts as the initial field.
2. Start from ERA5 48 hours earlier and eight 6-hour Feng forecasts as the initial field.
3. Start from ERA5 192 hours earlier and use thirty-two 6-hour FengWu forecasts as the initial field.
4. Directly use ERA5 fields from 24 hours earlier as the initial field. Directly use ERAS fields from hours earlier as the initial field.
6. Directly use ERAS fields from 192 hours earlier as the initial field. Use ERA5 fields from half a year earlier as the initial field.
These seven settings are denoted as "init- 1", "init-2", .... "init-7", respectively. All the other experi- mental settings consistent with the unfixed ERAS-simulated observation experiments in the main paper. The cyclic forecasting and assimilation is performed for 25 days and the results of RMSE are demonstrated in Figure It can be found that if the initial field is not derived from Feng forecast (init-4, init-5, init-6, init-7), the assimilation process can still continuously reduce errors over cycles, eventually achieving comparable performance. Even when the initial field significantly different from the ground truth, e.g., with a half-year phase (init-7), z500 still converges after about 20 days of assimilation. This ablation study proves that VAE- Var is robust to initial conditions.
Different Observation Quantities To investigate the impact of how observation quantity affects the performance of VAE-Var, added experiments with 100, 200, 400, 800, 10,000, 20,000, and 40,000 observation columns (the largest density corresponding to approximately one observation per grid point global 1.25-degree latitude-longitude grid).
18

---

Published as a conference paper at ICLR 2025
x500 (m 16
1500 (K)
u500 (15'5)
v500 (m/s)
1000
10
in2-1
ini-2
830
300
0s2-4
ink-5
25-
15-
10-
200
0-
0.0
460
200
200
460
Simmllated time (houss)
Simulated time (hours)
Simulated time (hours)
Simmiated time (hours)
ull (m/s)
v10 (m/s)
12m (K)
mslp
10
2000
500
83155
200
400
600
200
400
600
200
400
620
200
400
500
Simulated time (hours)
Simulated time (hours)
Simulated time (hours)
Simulated time (houm)
Figure 8: Results for different initial states. The observations are unfixed in different time steps. The system is simulated for 25 days, starting from January 1, 2022.
The experiments use a 15-day assimilation cycle from January 1 to January 15, 2022, with differ- ent observation locations each time step but consistent number of observations. The results for observation quantities from 100 to 8000 are demonstrated in Figure 9 and those from 8000 to 40000 are demonstrated in Figure [O It can be found that when the number of observations is low (below 8000), VAE Var shows a clear improvement in analysis accuracy as the observation quan- tity increases. However, when the observation density exceeds 8000, the improvement in accuracy becomes less pronounced. In contrast, 3DVar demonstrates more consistent improvement as the observation count increases. When the number of observations is below 10,000, VAE- Var generally outperforms 3DV, with the advantage becoming more pronounced as the observations are sparser. However, when the number of observations exceeds 10,000, the performance of VAE-Var slightly lags behind 3Dv
Regarding why the improvement in analysis accuracy with higher observation density is less pro- nounced for VAE- Var, we think this is due to the fact that the increments in the analysis field are constrained within the latent space defined by the decoder output, D(z). During training. VAE may not fully generalize to the complete analysis increment space, meaning that certain finer details of the analysis may not be captured effectively. When observation density is high, achieving further improvement requires fine- grained adjustments to the analysis increments, which may not be well- represented in the latent space. As result, the analysis accuracy may plateau once the background field error is already small. Currently, the RMSE for the z500 field has lower bound of approx- imately 26 m²/s², which represents good level of accuracy, but further improvements may be limited due to these modeling constraints.
19

---

Published as a conference paper at ICLR 2025
RMSE aar /8")
2500 RMSB of 3Da (m)
Rots 100
19bs 1500
-400
Robs 800
wubs 9200
300
500
100
100-
50
100
150
200
250
303
350
50
100
150
200
250
300
350
Simulated time (hours)
Simulated time (howrs)
2500 (m (s') Robs 100
2500 (m (6"), Hobs 200
2500 (ai 1/s) Hobe- 400
2500 (mi 16)) ), Hobs 800
403
SDWa VAll
150-
100
500
75
1200
50-
100
3D15r
SDM
25
3DVr
300
100
200
100
Simsulate time (hours)
Simulated time (hours)
time (hours)
Simulated time (hours)
z500 (m is", Hobs- 1000
z500 (m is"), Nooss 2000
2500 (m is") Hobs 4000
>500 (m /s" Nobs 8000
100
X0
VAT: Ve
50-
XMSE
50-
40
SDWb
1Dio
20
SD3
20-
VAE
100
100
100
Simulated time (hours)
Simulated time (hours)
Sinu Isted time (hours)
Simulated
(hours)
Figure 9: Results for different observation quantities (100 < #obs < 8000). The observations are unfixed in different time steps. The system is simulated for 15 days, starting from January 1, 2022.
A.9 EXPERIMENTS EVALUATED BASED ON GDAS OBSERVATIONS
Another important evaluation method in the field of data assimilation to treat the obser- vation sites ground truth for assessment. In this method, all observations are first divided into parts: one part is used for assimilation, and the other part is reserved as ground truth evaluate the quality of the assimilated anal- ysis field. In this section, we also adopt this method to assess the assimilation performance of VAE- Var. Specifically, we select the obser- vation sites within the region bounded by 50"E to 150'E longitude and 1 to 40"N latitude (with the main area covering East Asia). Dur- ing the assimilation process, do not assim- ilate these observations, but use them as truth for evaluation.
Figure 11: Division of observations. The region marked by the red dashed box represents the ob- servation area reserved for evaluation.
20

---

Published as a conference paper at ICLR 2025
2500 RMSE of VAB (m
2500 RMSE of 3D (m
Trdoo 8000
Vodo 13000
Vodoo +2000
10
20-
20-
10
Simolated time (hours)
(hours)
x500 7s"",, Hobs 8000
2500 (mi 1s")
(mi /s")
20000
(m is')
40000
30V
3UV
3DMVrr
60
60
62
60-
40
20-
20
100
201
300
100
300
500
100
200
300
100
200
300
Sinulaated time (bours)
Simmlated time (Bours)
Simulated time (houss)
Simullted time (hours)
Figure 10: Results for different observation quantities (8000 #obs 40000). The observations are unfixed in different time steps. The system is simulated for 15 days, starting from January 1, 2022.
Starting from January 1, 2017, the cyclic forecasting and assimilation system simulated for one year and the results are shown in Figure 12 observations, VAE-Va Var still consistently outperforms 3DVar on most variables.
13
14
It is shown that under the evaluation of GDAS
u10 (m/as
v10 (mis)
12m (K)
mslp (Pa)
VAE Nor
1000
110)
400
Difforesce betwoon RMSES
10
10
5-
-5-
-5-
-5
-500
100
10
12
300
Simullated time (Days))
Stamallated time (Days)
Sinnulated time (Dayys)
Simallted time (Days)
Figure 12: Evaluation results based GDAS observations. Surface variables (u10, v10, 12m, mslp) are reported. The system is simulated in an auto-regressive manner for one year, starting from January 1, 2017. The upper row demonstrates the RMSEs between the analysis fields and the GDAS observations. The RMSEs are averaged among all the observations available at a given layer. For example, calculate the RMSE of z501 at 00:00, January 2017, we first interpolate the analysis field to the height of 501hPa; then, calculate the squared error between the value an observation the height 501hPa value and the value of the analysis field at that point; by averaging all these observations and taking the square root, we can obtain the RMSE. The lower row shows the value of the RMSE of 3DVar minus that of VAE Var. A value larger than zero indicates that VAE- Var is better than Var
21

---

Published as a conference paper at ICLR 2025
2501 (m )/5)
1501 (K)
uS01 (miss
v501 (m/ss
1300
3DNu
VAB Ver
10.0
12.5
12.5
7.5-
7.5
5.0
500
0.5
0.0
1000
5-
Diffocanes
-5
-100
-10
-10
-10-
102
200
300
300
Simolated
(Tayss
Simuuatted (Days)
(Dags)
Simollated (Days)
Figure 13: Evaluation results based on GDAS observations. z501, t501, u501 and v501 are reported. Please refer the caption of Figure for details.
(K)
(m/s)
1000
12
12
3DH
800
600
403
200
1000
10
20
500
-500
-5-
-5
-2000
-10
300
100 Simmlated
Simulated time (Days)
Simulsted time (Days))
Simulated time (Days)
(Drys)
Figure 14: Evaluation results based on GDAS observations. z858, t858, u858 and v858 are reported. Please refer to the caption of Figure 12 for details.
A.
ADDITIONAL. RESULTS FOR EXPERIMENTS IN THE MAIN PAPER
Results for other important variables in the fixed- position
observational settings
demonstrated in Figure 15 17 and those in the fifxxed-position are demonstrated Figure [8]
20 The fields of other important variables for the GDAS prepbufr observations are visualized in
Figure 27
22

---

Published as a conference paper at ICLR 2025
obs amount 1000 (0%10%)
obs amount 2000 (0.19%)
obs amount = 4000 (0.39%)
obs amount 8000 (0.77%)
ROOS (c02))
2.5
VAL Var DEDDA
2.0
Interpolation
50%
1.5
2001
100
S10)
+10 RMSE (mis)
2.5
2.0
1.0
100
200
300
100
200
300
100
200
300
100
200
300
300
malp RMSE (Pa)
100
200
300
100
200
300
200 300
200
300
Simulated Lime (hours)
Simulated time (hours)
Sinulated time (hours)
Simulated lame (bours)
Figure 15: Results for fixed observation positions (15 days). The system is simulated in an auto- regressive manner for 15 days, starting from January 1, 2022. The RMSEs of the analysis field for three variables (ul0, v10, mslp) are demonstrated in three different rows.
obs amount 1000 (0.10%)
obs amount 2000 (0.19%)
obs amount = 4000 (0.39%)
obs amount 8000 (0.77%)
VAB-Var
2850 RMSE Emi
200
DEDDA
150
Interpolation
30V
100
200
300
100
200
100
200
500
100
300
6850 RMSE (m35)
100
200
300
100
200
300
100
230
300
100
200
300
v150 RMSE (1955)
100
200
300
100
200
300
100
200
100 Simulated time (hours)
Simukted time (hours)
Simulated time (hours)
Simulated time (hours)
Figure 16: Results for fixed observation positions (15 days). The system is simulated in an auto- regressive manner for 15 days, starting from January 1, 2022. The RMSEs of the analysis field for three variables (z850, u850, v850) are demonstrated three different rows.
23

---

Published as a conference paper at ICLR 2025
obs amount 1000 (0.10%)
obs amount 2000 (0.19%)
obs amount 4000 (0.39%)
obs amount 8000 (0.77%)
0,000 ROOER (mal))
VAL Var DEDDA
Interpolttion
50%
100
200
300
100
200
300
100
200
300
100
200
300
v500 RMSE (mok)
100
200
100
200
300
100
200
300
100
200
300
1500 RMSE (K)
0.5
100
200
100
200
200
100
200
300
100
200
300
Simulated lime (hours)
Simulated time (houms)
Simulated time (hours)
Simulated time (hour))
Figure 17: Results for fixed observation positions (15 days). The system is simulated in an auto- regressive manner for 15 days, starting from January 1, 2022. The RMSEs of the analysis field for three variables (u500, v500, 1500) are demonstrated in three different rows.
obs amount 1000 (0. 10%)
obs amount 2000 (0.19%)
obs amount = 4000 (0.39%)
obs amount 8000 (0.77%
RMSE (anis)
2.5
VAB-Var
DEDDA
Interpoltiion
30V
1.5
200
300
100
300
200
300
300
v10 RMSE mm))
2.5
M
2.0
1.0
100
200
300
100
200
300
100
200
300
130
200
300
300
mspp RMSE (Pa)
Ww
100
200
300
100
200
500
200 300
200
309
Simulated time (hours)
Simulared time (hours)
Simulated time (hours)
Simulated time (hours)
Figure 18: Results for unfixed observation positions (15 days). The system is simulated in an auto-regressive manner for 15 days, starting from January 1, 2022. The RMSEs of the analysis field for three variables (u10, v10, mslp) are demonstrated in three different rows.
24

---

Published as a conference paper at ICLR 2025
obs amount 1000 (0.10%)
obs amount 2000 (0.19%)
obs amount = 4000 (0.39%)
obs amount 8000 (0.77%)
250
4850 RMSE (m 1s
MM /
VAL Var
200
DEDDA
150
Interpolttion
50%
200
300
100
2001
100
200
500
100
U850 RMSE (m/s)
100
200
300
100
200
300
100
250
300
100
250
300
v850 RMSE (108)
100
200
300
300
200
300
100
200
500
100
300
Simubtedd lime (bours)
Simulated lime (bours)
Simulated time (hours)
Simulated Lime (hours)
Figure 19: Results for unfixed observation positions (15 days). The system is simulated in an auto-regressive manner for 15 days, starting from January 1, 2022. The RMSEs of the analysis field for three variables (2850, u850, v850) are demonstrated in three different rows.
obs amount- 1000 (0.10%)
obs amount 2000 (0.19%)
obs amount 4000 (0.39%)
obs amount 8000 (0.77%)
0,500 RMME (mii))
VAB-Var
DEDDA
Interpobtiion
30V
100
200
300
100
200
300
100
200
300
130
200
300
v500 ROOSE (msk))
100
200
100
200
300
100
200
300
100
200
300
1500 RMSE (3C)
0.5
100
100
200
200
100
200
300
200
300
Simulated time (hours)
Simulated time (hours)
Simulated time (hours)
Simulared time (hous)
Figure 20: Results for unfixed observation positions (15 days). The system is simulated in an auto-regressive manner for 15 days, starting from January 1, 2022. The RMSEs of the analysis field for three variables (u500, v500, 1500) are demonstrated three different rows.
25

---

Published as a conference paper at ICLR 2025
z850 ground truth (m (m'1s")
z850 background 1s²)
z850 obs mask, #obs 3022 (0.291%)
10000
12500
15000
10000
12500
15000
Background error (RMSE-58..) 58.7--58.mm/ 1s")
2850 unalysis increment (m )
Analysis enor (RMSE-48ml 115)
-500
500
-500
500
500
500
Figure 21: Visualization of z850 assimilation results at time 2017-01-01 00:00:00. The figure titled "2850 obs mask" correspond to observations at 858hPa because the observations different height axis.
u500 ground truth (m/s)
u500 backgrround (m/s)
u,000 ubs mask, Hubs- 1obs-1190 1190 (0.115%)
50
Background (RMSE-2.22 m/s)
u500 anulysis increment (m/s)
Anulysis (RMSE-2. 1.0SS--2.14 2.14 m/s)
-10
0
10
100
10
-10
10
Figure 22: Visualization of u500 assimilation results at time 2017-01-01 00:00:00. The figure titled "u500 obs mask" correspond observations at 501hPa because the observations use a different height axis.
26

---

Published as a conference paper at ICLR 2025
u850 ground truth (m/s)
u850 background (m/s)
u850 ohs mask, #obs-1300 (0.126%)
Background emrrr (RMSE-1 1.71 m/s)
u850 analysis increment (mis)
Anal Analysis error (RMSF- 1.63 m/s)
-10
10
110
-10
10
Figure 23: Visualization of u850 assimilation results at time 2017-01-01 00:00:00. The figure
titled "u850 obs mask" correspond to observations at 858hPa because the observations a different height axis.
v500 ground truth (m/s)
v500 background (m/s)
v500 mask, #obs- 1048 (0.101%) %)
-50
Background error (RMSE-2.16m(s)
v500 analysis increment (mis)
Analysis crror (RMSE=2.077 m/s)
-10
10
-10
10
Figure 24: Visualization of v500 assimilation results at time 2017-01-01 00:00:00. The figure
titled "v500 obs mask" correspond to observations at 501hPa because the observations a different height axis.
27

---

Published as a conference paper at ICLR 2025
v850 ground truth (m/s)
v850 background (m/s)
v850 obs mask, #obs-1 194 (0.115%)
Background errir (RMSE-1 1.63 m/s)
v850 analysis increment (mis)
Analysis error (RMSF- 1 1.57 m/s)
-10
10
110
-10
10
Figure 25: Visualization of v850 assimilation results at time 2017-01-01 00:00:00. The figure
titled "v850 obs mask" correspond to observations at 858hPa because the observations a different height axis.
ground truth (K)
1500 background (K)
1500 obs mask, Hobs=483 (0.0465%)
240
260
260
Background error (RMSE=0.639K)
1500 analysis increment (K)
Analysis error (RMSF-0.60) K)
Figure 26: Visualization of t500 assimilation results at time 2017-01-01 00:00:00. The figure titled "t500 obs mask" correspond to observations at 50thPa because the observations a different height axis.
28

---

Published as a conference paper at ICLR 2025
1850 ground truth (K)
1850 background (K)
(850 obs mask, Nobs- 769 (0.0741%)
250
300
Background emorr (RMSE-0.88 K)
1850 analysis increment (K)
Analysis error (RMSE 00822 K)
Figure 27: Visualization of 1850 assimilation results at time 2017-01-01 00:00:00. The figure titled "1850 obs mask" correspond to observations at 858hPa because the observations a different height axis.
29


Reviews:
[{"review_id": "4cIaNZ39C1", "reviewer": "Reviewer_D5gS", "summary": "The paper presents VAE-Var a new data assimilation algorithm based on variational assimilation in order to capture non-Gaussian characteristics of the conditional background distribution p(x|x_b) and thus enhance the accuracy of the method. The algorithm also able to incorporate real measurements coming form different systems.\n\nThe theory is clear. The architecture is well justified. The experiments too. The results of the main paper are very interesting, the one in the supplementary material bring some nuances.\n\nThe experiments rely on FengWu, an adaptation is done to construct the encoder and the decoder. ERA5 reanalysis data from 1979 to 2015 are used as training set. Two types of observations are used : simulated one from ERA5 to mimic real-world observations, and GDAS prepbufr with various system and instruments. The initial setting is derived from ERA5 dataset. For ERA5 simulations, fixed or unfixed observations are used. For fixed observations VAE-Var outperforms the other methods  in the main paper. In the supplementary material figure 8, sometimes 3DVar perform better than VAE-Var as for mslpRSME with 8000 observations. Graphically speaking, VAE-Var performs closely to 3DVar. It is worth noting that VAE-Var performs better most of the time with 1000 observations.\nFor GDAS, the RMSE is always smaller for VAE-Var. Bias is always better, except for z500.\n\nAll in all, this is a good paper.", "questions": "I am not completely convinced by the right part of the equation 2, could you give more details or a reference?\n\nLine 129, why is the matrix R of the observation error is assumed to be diagonal?\n\nWhy do the authors use ERA5 reanalysis data from 1979 to 2015?\n\nIn the supplementary material A.2., could the authors give a reference where this term is rigorously considered? So, further work to make a link between the complexity and the improvements of VAE-Var approach can be supported.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "In section 3.1.the use of the hypothesis of independence between the background states x-x_ b and background state x_b is well supported. The same goes for the non-Gaussian error distribution assumption.\n\nThe justification of the use VAE is clear. The generation of the training set construction is well formulated, the same goes for the loss function.\n\nThe authors are honest regarding their results, especially regarding the limitations in supplementary material A.2.\n\nVAE-Var focuses learning the latent space of the error fields. This is an AI forecasting model effective at global scale with real-world measurements.", "weaknesses": "Limitations :\n\nThe performance presented in the main paper show very interesting results regarding VAE-Var. It must be nuanced with the supplementary material figure for broader considerations.\n\nThe authors mention in the supplementary material A.2. that they do not consider the determinant term in the equation 8 because it would make the assimilation of the objective function extremely hard to optimize. Thus, they introduce a scaling factor in the background term of the assimilation objective function. This is honesty from the authors. Maybe some tricks can be found later to make it rigorous.\n\nMinor comments that do not impact the score:\n\nLine 243-244, in Chen et al., the dimensionality of the problem is 189 * 721 * 1440. What does the 69 correspond to? Is it a confusion with what is used in the present study section 4.1?\n\nLine 262-263, the authors claim that all the components of L are differentiable. Why H is differentiable?\n\nThere is no 2.5.3. in Giaquinta & Modica 2012, only 2.5.a-d regarding convexity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kSTd8g81tf", "reviewer": "Reviewer_2mtD", "summary": "This paper proposes a new Variational Data Assimilation method based on Variational Autoencoders, called VAE-Var. The method is separated into two stages, the training and the assimilation stage. During the training stage, a dataset of error samples is constructed by comparing two forecasts made at different time intervals. These samples are then used to train the VAE model, which learns the error distribution of the background states. During the assimilation stage, the VAE decoder is used to map latent vectors back to the physical state, that when combined to the background field represents the analysis states after data assimilation. The proposed methodology is compared against DiffDA, 1DVAR, and interpolation for the FengWu forecasting model and show superior performance across benchmarks", "questions": "- In my understanding, Diffusion models can capture more complex distributions but require a lot of data. Could the authors provide more information on how fair the comparison between the VAE and the Diffusion model is? Meaning that if the distribution that they model is not  complex and the number of available data points is not large enough, isn't Diffusion models expected to perform worse?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written, and clearly communicates the goal. The authors also explain the shortcomings of the other methods and how they propose to overcome them. The methods section is clear and the results show that VAE-Var out-performs all baselines.", "weaknesses": "There is not an obvious weakness to the paper. However I am not an expert in meteorology so maybe I missed something.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a4ylRrzNHJ", "reviewer": "Reviewer_bhb7", "summary": "The study introduces VAE-Var, a neural network-based data assimilation method that combines VAE with traditional variational DA techniques. The authors claim that VAE-Var is designed to enhance DA by (i) Capturing non-Gaussian structures in background error distributions. (ii) Allowing effective assimilation of real-world, often irregularly distributed observations. Tests conducted using the FengWu weather model demonstrate that VAE-Var outperforms DiffDA and 3D-Var. Although the research topic is of interest, from the reviewer’s point of view, the novelty of the proposed approach is limited since (varitional-) autoencoder based data assimilation methods are already well established.", "questions": "The authors should also compare their approach to ensemble-based data assimilation algorithms, such as EnKF, which are known to yield superior results for chaotic dynamical systems.\n\nThe role of the parameter lambda needs further explanation and why lambda = 4 is chosen in this paper?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The VAE-Var approach can integrate partial and irregular observation points, giving it certain advantages over purely supervised machine learning surrogates for data assimilation. Numerical results show that the VAE-Var approach consistently outperforms both DiffDA and the conventional 3D-Var approach in terms of accuracy and efficiency.", "weaknesses": "The combination of neural network based AE (including VAE) and variational data assimilation is a well known approach. There are more than 10 publications regarding this idea in dynamical systems, for ex,\n- Peyron, M., Fillion, A., Gürol, S., Marchais, V., Gratton, S., Boudier, P. and Goret, G., 2021. Latent space data assimilation by using deep learning. Quarterly Journal of the Royal Meteorological Society, 147(740), pp.3759-3777.\n-Melinc, B. and Zaplotnik, Ž., 2024. 3D‐Var data assimilation using a variational autoencoder. Quarterly Journal of the Royal Meteorological Society, 150(761), pp.2273-2295.  \n- Cheng, S., Chen, J., Anastasiou, C., Angeli, P., Matar, O.K., Guo, Y.K., Pain, C.C. and Arcucci, R., 2023. Generalised latent assimilation in heterogeneous reduced spaces with machine learning surrogate models. Journal of Scientific Computing, 94(1), p.11.\n- Mohd Razak, S., Jahandideh, A., Djuraev, U. and Jafarpour, B., 2022. Deep learning for latent space data assimilation in subsurface flow systems. SPE Journal, 27(05), pp.2820-2840.\n- Zhong, C., Cheng, S., Kasoar, M. and Arcucci, R., 2023. Reduced-order digital twin and latent data assimilation for global wildfire prediction. Natural hazards and earth system sciences, 23(5), pp.1755-1768.\n…..\n\nThe proposed approach uses ground truth samples for background error specification. This might be a valid approach in twin experiments. However, ground truth is out of reach in most applications of DA methods, such as weather prediction. In other words, we never know the historical ‘true’ weather, as observations also contain uncertainties. I would suggest focusing on the innovation quantity (i.e., y-H(xb) or y-H(xa)) for a propre error covariance specification/analysis, see\n-Tandeo, P., Ailliot, P., Bocquet, M., Carrassi, A., Miyoshi, T., Pulido, M. and Zhen, Y., 2020. A review of innovation-based methods to jointly estimate model and observation error covariance matrices in ensemble data assimilation. Monthly Weather Review, 148(10), pp.3973-3994.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CcZcRMJWpU", "reviewer": "Reviewer_jauN", "summary": "The paper introduces VAE-Var, a data assimilation method that combines variational autoencoders (VAE) with a traditional variational data assimilation technique to enhance weather forecasting accuracy. By using a VAE, VAE-Var relaxes the traditional Gaussian background error assumption, enabling it to better model complex, non-Gaussian background error distributions. Through experiments with both synthesized observations and real-world observations, VAE-Var demonstrates improved or comparable accuracy over established methods, including 3DVar and DiffDA. The approach’s adaptability to diverse data types and its demonstrated year-long forecasting stability showcases VAE-Var’s promise for real-world weather forecasting applications.", "questions": "1. One major question that need to be addressed: VAE-Var improves minimally or does not improve with an increasing amount of available observations, whereas other approaches, both ML-based and traditional, show noticeable gains under similar conditions (Figures 3, 4, 8, 9, 10, 11, 12, 13). This behavior is unusual for data assimilation algorithms, as more observations typically enhance analysis accuracy. A thorough discussion or ablation study may be needed to investigate this.\n\nSome suggestions from the review feedback system : )\n- Conduct an ablation study to investigate why VAE-Var's performance doesn't improve significantly with more observations\n- Discuss potential reasons for this behavior, such as possible saturation effects or limitations in the model's capacity to incorporate additional information\n- Analyze the implications of this behavior for the algorithm's practical use in different observational settings\n- If possible, propose modifications to the algorithm that might address this limitation\n\n2. Formulating  $p(y∣x)$ as a Gaussian distribution is indeed an approximation of the observation errors. Although it is widely used in operational numerical weather prediction (NWP), it is an approximation instead of a well-defined characteristics. This assumption does not always hold. For instance, quality-controlled observations of brightness temperature from the High Resolution Infrared Sounder exhibit distributions significantly different from Gaussian. Incorrectly assuming Gaussian statistics in such cases can lead to substantial biases in the estimated state, as shown in studies like[6], which diagnose and assess the impacts of non-Gaussianity in innovations on data assimilation accuracy. For rigor, the paper could benefit from framing the Gaussian distribution of $p(y∣x)$ as a well-accepted but approximate assumption, rather than misleadingly stating that Gaussian error is a well-defined characteristic of observation errors. \n\n[6]: Pires, Carlos A., Olivier Talagrand, and Marc Bocquet. \"Diagnosis and impacts of non-Gaussianity of innovations in data assimilation.\" *Physica D: Nonlinear Phenomena* 239.17 (2010): 1701-1717.\n\n3. Minor issues:\n- Line 202-2-3 and Algorithm 1 and Figure 1. Is $\\tau$  the number of time steps or the time gap between two snapshots? It is a bit confusing here.\n- Inconsistency between the lower half panel of Figure 1 and Algorithm 2: How is the hidden state initialized? 0 or a random field?\n- Algorithm 3: the model time step T is not explicitly defined in the algorithm.\n- The last term in Equation (8) should be $\\frac{1}{2}\\log \\det \\left( \\frac{\\partial \\mathcal{D}(\\textbf{z})}{\\partial \\textbf{z}} \\right)^T \\left( \\frac{\\partial \\mathcal{D}(\\textbf{z})}{\\partial \\textbf{z}} \\right)$", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper communicates its concepts clearly. The use of VAEs to model non-Gaussian background error distributions is an original contribution, addressing a major assumption in background error modeling, with promising results. The experiment on real-world observations moves this work beyond a proof-of-concept. The topic is highly relevant: data assimilation is essential for numerical weather prediction, and developing efficient DA algorithms for AI-based weather prediction is necessary to improve their capability. Additionally, the integration of machine learning, specifically generative modeling, into DA represents a forward-looking approach that could influence further research in hybrid models for weather and climate prediction.", "weaknesses": "### 1. Regarding ERA5 as the 'ground truth'.\nERA5, as a reanalysis product, can serve as a useful reference, but it should not be treated as \"ground truth,\" especially for evaluating data assimilation algorithms. ERA5 is known to have a lot of issues that can be attributed to terrain effects, poor coverage of assimilated observations, and the numerical model used. For example, ERA5 exhibits underestimated temperature in some regions of Europe[1], an overly strong equatorial mesospheric jet, mismatches in the analyzed near-surface wind/temperature/humidity[2], significant precipitation/temperature biases in some regions of North America[3], etc.\n\nTherefore, merely evaluating data assimilation algorithms on ERA5 data overlooks the actual 'ground truth' beyond ERA5's inherent biases and errors. The authors could benefit from reviewing some non-ML DA papers that work on real observations, to see how their DA approaches are evaluated. For example, to assess whether a data assimilation algorithm truly works, it is important to compare its performance directly with real observations. A potential test could be holding out actual observations at the same positions at each time step and using these to evaluate the analyses from all implemented DA algorithms. This experiment would require no retraining and would allow for direct comparison and diagnostics of the DA algorithm developed. Paper [4], which uses the root-mean-square departures from actual observation to validate their approach, serves as a reference, though there are likely many more (for example, [5]).       \n\nRegarding ERA5 as ground truth and only evaluating algorithms on ERA5 are common issues in AI for weather research. Recognizing that ERA5 is not a definitive ground truth would provide the community with a clearer perspective and encourage more comprehensive benchmarking with different data sources, thus advancing the field.        \n[1] https://www.mdpi.com/2073-4441/14/4/543          \n[2] https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation#ERA5:datadocumentation-Knownissues          \n[3] https://hess.copernicus.org/articles/24/2527/2020/   \n[4] https://journals.ametsoc.org/view/journals/mwre/151/7/MWR-D-22-0260.1.xml         \n[5] https://rmets.onlinelibrary.wiley.com/doi/full/10.1002/qj.3917            \n\n### 2. Lack of reproducibility and missing details weaken the results and claims.\n- As the code is not provided, additional details on the configuration of algorithms and baselines are necessary. However, they are not provided.\n- There are no details on how the observation operator is constructed, particularly for real observations. \n- 3DVar details are missing, such as how $\\text{U}_p$, $S$, $\\text{U}_v$, and $\\text{U}_h$ are determined?\n- Configuration details for DiffDA are also absent, leaving readers uncertain about how this baseline was set up and tuned for the experiments.\n- Providing the above details would enhance the reliability and transparency of the reported results, allowing readers to better assess and understand the performance claims of VAE-Var relative to these baselines.\n\n\n### 3. Simplification\n- As the authors mention, the simplification in the VAE-Var objective function is not mathematically rigorous. More importantly, it also weaken the claim that VAE-Var could better leverage the non-Gaussian statistics. By omitting the exact term in the background error component that accounts for the nonlinear transformation from a Gaussian to a non-Gaussian distribution, the model loses a critical factor needed to rigorously capture the intended non-Gaussian effects.  Unfortunately, addressing this limitation may be challenging, as the most straightforward solution would involve making the original, more complex optimization problem tractable.\n- A minor issue related to appendix A.2: the reference for equation (7), *Giaquinta & Modica (2012): Mathematical analysis: functions of one variable. Springer Science & Business Media*, does not have Chapter 2.5.3 and equation (7).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Nigjm44uXH", "reviewer": "Reviewer_bsmw", "summary": "The authors offer a way of conducting sparse data assimilation using a minimization problem by estimating the error term of a forecasting model as compared to the state directly, and combining it with variational data assimilation. They achieve convincing results on weather forecasting.", "questions": "When generating, why is it integrated 0->2t and t->2t instead of just 0->t vs the ground truth at t?\n\nHow is the speed in terms of assimilation versus previous data assimilation approaches?\n\nSince you're modeling the error; suppose that the initial state is very ill-estimated. Is the data assimilation approach still able to resolve the problem, or does it rely on the initial state estimate being within some stability region?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper offers a method of data assimilation for sparse meteorological problems which surpass current baselines like DiffDA, and give very promising results. The method is clear, and the problem is particularly applicable to real-time weather forecasting.", "weaknesses": "Currently, the model seems method-specific, limiting its generalizability as a data assimilation approach since any new improvements or retraining of the original forecasting model would require retraining of this model. Is there a way to perhaps make this more robust? However, in the grand scheme of things, due to its superior performance, this isn't that big of an issue in my opinion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bvtQYmjho4", "reviewer": "Reviewer_Uc8k", "summary": "This paper studies the data assimilation problem in the context of numerical weather forecast, and proposes a deep learning model for this problem based on a variational auto-encoder architecture. This approach captures the non-Gaussian background error distribution relative to the forecasting model within a variational data assimilation framework, by training a variational auto-encoder to learn the prior distribution. This deep learning model is evaluated on high-dimensional operational weather data and shows competitive performances compared to a recent diffusion based data assimilation algorithm.", "questions": "In Figure 3, how would the authors explain that 3DVar outperforms DiffDA? Should the latter not be better than the former, as it can capture non-Gaussian statistics?\n\nIn line 261, the part about rescaling the background term is not very clear (although it may be explained in more details in the Appendix).\n\nUsing the term \"interpolation\" as a name for a nearest-neighbor method might be confusing, as \"interpolation\" already has a precise meaning in the data assimilation literature [3].\n\nA paramount aspect in data assimilation is uncertainty quantification and measuring the various possible reconstruction states, under chaotic dynamics. This is briefly mentioned in the related work section, but it would be very relevant to discuss how the probabilistic framework underlying variational auto-encoders allows the learning model to sample different predictions for a given observation.\n\n### References\n\n[1] Huang, L., Gianinazzi, L., Yu, Y., Dueben, P. D., & Hoefler, T. (2024). Diffda: a diffusion model for weather-scale data assimilation. arXiv preprint arXiv:2401.05932.\n\n[2] Lorenc, A. C., & Rawlins, F. (2005). Why does 4D‐Var beat 3D‐Var?. Quarterly Journal of the Royal Meteorological Society: A journal of the atmospheric sciences, applied meteorology and physical oceanography, 131(613), 3247-3257.\n\n[3] Asch, M., Bocquet, M., & Nodet, M. (2016). Data assimilation: methods, algorithms, and applications. Society for Industrial and Applied Mathematics.\n\n[4] Fablet, R., Chapron, B., Drumetz, L., Mémin, E., Pannekoucke, O., & Rousseau, F. (2021). Learning variational data assimilation models and solvers. Journal of Advances in Modeling Earth Systems, 13(10), e2021MS002572.\n\n[5] Blanke, M., Fablet, R., & Lelarge, M. (2024). Neural Incremental Data Assimilation. arXiv preprint arXiv:2406.15076.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The presented methodology is sound, and the proposed algorithm is accurately described. The paper is clearly written overall.\n\nThe proposed approach is evaluated on real-world data, on which it is shown to outperform the baselines. This suggests that considerable work has been invested to ensure the scalability of this method.\n\nThe VAE-Var algorithm is able to assimilate data from more general observation operators than DiffDA or other learning-based algorithms.\n\nThe authors propose a cyclic forecasting and assimilation setup where data assimilation and forecasting are coupled, with the forecast model, FengWu, being itself a large deep learning model. This method, already proposed in [1], is appealing for the future of numerical weather forecast as it does not rely on computational intensive weather models anymore once the neural networks are trained. \n\nAlthough not emphasized in the paper, a key potential strength of this approach is that the variational problem is formulated in latent space, rather than in state space. This may be a considerable advantage if the neural network efficiently compresses the state variable. Indeed, the data assimilation problem is typically of very high dimension, which is a bottleneck in practice, and reducing the dimension of the computations involved is key.", "weaknesses": "One key issue in the paper is the assumption that the state-of-the-art is limited to spatial-only data assimilation, specifically 3DVar-like algorithms, where the observations are assimilated at each new time step, as in Algorithm 3. There is no mention of the 4DVar algorithm, which is the state of the art for operational data assimilation, and which differs significantly from the description proposed in section 2. Instead of modeling error by a Gaussian distribution at a given time step relative to the prediction at the previous time, 4DVar models the error distribution at a multitude of observation times, within a so-called assimilation window, thanks to a physical model that couples the different instants. By enforcing temporal dependencies in the observations via physical laws, the reconstructed state becomes more accurate, and 4DVar is known to outperform 3DVar. As opposed to the framework described in Section 2, the error distribution in 4DVar is not Gaussian, as there is a strongly nonlinear physical model in the equation. See [2] and [3] for more details.\n\nTo properly study the impact of the proposed method, I believe that it should be compared, at least on a toy system, to 4DVar, and potentially to learning-based counterparts of 4DVar that model the temporal dimension of the signal to reconstruct it, including diffusion models [4] and end-to-end reconstruction approaches [5, 6].\nFurthermore, the upsides of using the presented variational auto-encoder framework rather than a diffusion model as in DiffDA is not so clear: the authors mention that their algorithm is more versatile in terms of observation operators, but this difference is not thoroughly explained in Section 3.\n\n\nAdditionally, I found that the paper lacks clarity in some parts. For example, (too) much space is dedicated to the Bayesian framework of variational data assimilation, whereas the training loss function of the proposed method, which relies on the VAE framework, is not explicitly state. In the experiment section, I think that enumerating the baselines in a common paragraph would improve the clarity.\n\nFinally, the paper does not mention the computational cost of data assimilation and does not compare the computational times involved. This aspect is of significant importance, and computational savings are one of the main reasons underlying the deep learning approach for weather forecast problems. Therefore, I believe that it would be beneficial for the paper to discuss this topic and to compare the computational times between the different approaches, traditional and learning-based.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)