Latent-EnSF: A Latent Ensemble Score Filter for High-Dimensional Data Assimilation with Sparse Observation Data
Authors: Phillip Si
Peng Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Data Assimilation
Score Based Models
Diffusion Models
Weather Forecasting
Primary area: learning on time series and dynamical systems

Abstract: Accurate modeling and prediction of complex physical systems often rely on data assimilation techniques to correct errors inherent in model simulations. Traditional methods like the Ensemble Kalman Filter (EnKF) and its variants as well as the recently developed Ensemble Score Filters (EnSF) face significant challenges when dealing with high-dimensional and nonlinear Bayesian filtering problems with sparse observations, which are ubiquitous in real-world applications. In this paper, we propose a novel data assimilation method, Latent-EnSF, which leverages EnSF with efficient and consistent latent representations of the full states and sparse observations to address the joint challenges of high dimensionlity in states and high sparsity in observations for nonlinear Bayesian filtering. We introduce a coupled 
Variational Autoencoder (VAE) with two encoders to encode the full states and sparse observations in a consistent way guaranteed by a latent distribution matching and regularization as well as a consistent state reconstruction. With comparison to several methods, we demonstrate the higher accuracy, faster convergence, and higher efficiency of Latent-EnSF for two challenging applications with complex models in shallow water wave propagation and medium-range weather forecasting, for highly sparse observations in both space and time.

Published as a conference paper at ICLR 2025
LATENT-ENSF: A LATENT ENSEMBLE SCORE FILTER FOR HIGH-DIMENSIONAL DATA ASSIMILATION WITH SPARSE OBSERVATION DATA
Phillip Si, Peng Chen
Department of Computational Sciences and Engineering, Georgia Institute of Technology {psi6, pchen402) @attcch.edu
ABSTRACT
Accurate modeling and prediction of complex physical systems often rely on data assimilation techniques to correct errors inherent model simulations. Tradi- tional methods like the Ensemble Kalman Filter (EnKF) and its variants well the recently developed Ensemble Score Filters (EnSF) face significant challenges when dealing with high-dimensional and nonlinear Bayesian filtering problems with sparse observations, which are ubiquitous in real-world applications. In this paper, we propose a novel data assimilation method, Latent-EnSF, which lever- ages EnSF with efficient and consistent latent representations of the full states and sparse observations to address the joint challenges of high dimensionality in states and high sparsity in observations for nonlinear Bayesian filtering. We introduce a coupled Variational Autoencoder (VAE) with two encoders encode the full states and sparse observations in consistent way guaranteed by a latent distri- bution matching and regularization as well as consistent state reconstruction. With comparison several methods, we demonstrate the higher accuracy, faster convergence, and higher efficiency of Latent-Enn EnSF for two challenging applica- tions with complex models in shallow water wave propagation and medium-range weather forecasting, for highly sparse observations both space and time.
INTRODUCTION
Many complex physical systems are traditionally modeled by partial differential equations (PDEs). First-principle PDE-based modeling and simulation have proven to be powerful in predicting com- plex systems across various scientific and engineering fields. However, in practical applications, sig- nificant discrepancies between simulation- based predictions and reality may arise from sources such model inadequacy, uncertainties in model parameters, boundary and/or initial conditions, external forcing/loading terms, numerical approximation errors, and more. Meanwhile, data-driven machine learning (ML.) models like FourCastNet (Pathak al., (022) have been significantly developed in recent years, demonstrating extreme efficiency for system predictions as in weather forecasting. The computational efficiency also makes these ML models particularly well-suited for ensemble forecasting and uncertainty quantification. However, these models are often limited by a divergence from reality over the long run, resulting in significant accumulated errors in long-term predictions due to various uncertainties and the autoregressive structure of the ML models.
To mitigate discrepancies and accumulated errors and achieve more accurate predictions, data assim- ilation plays an essential role by incorporating additional observational data into current PDE/ML- based prediction models. By adjusting the PDE/ML state to align with the observational data, the accuracy of future state predictions ctt al. 2023 Asch et al. 2016) approaches such as the Kalman filter (KF) (Evensen) 1994), particle filters (Künsch, 2013), and their variants have been widely applied to problems including weather prediction, geophysical modeling, robotics, and many other areas.
be markedly improved. Data assimilation (Sanz-Alonse
The basis of many current data assimilation methods used in practice is the KF (Evensen 1994). It parameterizes its state with the mean and covariance, and provides the optimal solution to the data assimilation problem when certain linearity properties of the transition and observation functions

---

Published as a conference paper at ICLR 2025
met. However, in practice, KFss become computationally expensive for high sigh-dimensional states and can be biased for nonlinear problems. An alternative the particle filter, which represents the density by ensemble of particles. This approach relaxes the assumption of specific distribu- tional form, such as a Gaussian, that is inherent in distribution-based approaches like the KF. The EnKF (Kalman 1960) is particular partiele filter that originates from the KF; instead of updating the covariance matrix, it computes the sample covariance of the ensemble of particles. Due to its robust performance, has become widely used in applications today. However, most partiele fil- ters, including EnKF, suffer from the curse of dimensionality by requiring an exponentially large number of particles accurately describe the distribution in high dimensions. To address this chal- lenge, various extensions to EnKF have been proposed, such as the Localized Ensemble Transform Kalman Filter (LETKF) (Hunt al. 2007), which applies covariance localization techniques along with advancements from the Ensemble Transform Kalman Filter (Bishop et al., 200T).
Nonetheless, several major roadblocks hinder efficient data assimilation for high- High-dimensional sys- tems. Standard data assimilation approaches have limited capabilities in these environments, often relying on assumptions such as localization to remain applicable. High-dimensional systems and complex models also constrain the capability and efficiency of methods like 4D-Va Var (Rabier & Liu 2003). which requires repeatedly propagating through the physical model. Recently, DifIDA (Huang al. 2024) employs diffusion models and masked interpolations of weather data similate sparse observations, though the approach falls behind pure interpolation after a few steps. Meanwhile, Rozet Louppe (2024) utilizes conditional score-based generative models to sample from the trajectory space conditioned on observations; however, this addresses a fundamentally dif- ferent problem-sampling an entire trajectory conditioned observations rather than performing the real-time filtering needed for practical applications. |ao et al. (2024b) proposed novel method called Ensemble Score Filter (EnSF), which has achieved great results in nigh-dimensional data assimilation by leveraging score-based diffusion models to efficiently sample from the posterior distribution. Nonetheless, sparse observations remain very challenging for EnSF, limiting its appli- cability in real-world scenarios where observational data scarce. Similarly, ROAD (Chen al. 2023) develops latent dynamics through decoder but relies on dense observations.
In this work, propose efficient data assimilation method called Latent-EnSF (shown in Figure T that builds on Ensemble Score Filters (EnSF) (Bao et al. 2024b) to address the joint challenges of high dimensionality and data sparsity. We construct coupled Variational Autoencoders (VAE) (Kingma & Welling 2014) to encode the sparse observations into latent space, where we perform data assimilation using EnSF with diffusion-based sampling. Importantly, our approach addresses the key weakness of EnSF when dealing with sparse observations while retaining its ability to ef- ficiently assimilate observations for high-dimensional nonlinear systems. We demonstrate this per- formance for both PDE and ML-based prediction models on challenging examples, experimenting with data that is sparse in both space and time.
H
Edee
Holes Fdde
sparse
Reverse SUL Semping
Eyste
1yys
Figure 1: Flow of the Latent-EnSF. An ensemble of prior states 2; 191.11111111 is assimilated with sparse observation Yt in the latent autoencoder space to obtain samples ( from the posterior.
We will begin by formally introducing data assimilation problems along with review of the EnSF in Section 2 then explain how we integrate this with coupled variational autoencoders in Section to form our Latent-EnSE and finally conclude with experiments in Section
2

---

Published as a conference paper at ICLR 2025
2 DYNAMICAL SYSTEMS AND DATA ASSIMILATION
To model a dynamical physical system, we denote the state at a certain time as 22 E IRd, where typically very high for complex models. For system with discrete time steps 1, 2,.... given initial condition To, we then model the evolution of the state from the time t - 1 to time t as
It M (x(_1,E)),
(1)
where E, is process noise term coming from a known distribution; this term accounts for hidden interactions and numerical errors. The corresponding observation E 1 at time f is given by
yt = (z,) + 76.
(2)
where "i (also coming from known distribution, typically Gaussian) represents the observation noise because of instrumental inaccuracies, etc. Both the model map M R x Rd Rd and the
observation map H Rd I can be nonlinear.
The dynamical system model in Equation I may yield progressively inaccurate representations over time due factors such as model inadequacy or input uncertainty. The goal of data assimilation is to integrate the observation data into the dynamical system model to recover accurate representation of the true state. practical applications, the observation data can be sparsely distributed both space and time, making the recovery of the true state particularly challenging, especially for high- dimensional problems with a state dimension >> 1, and the observation dimension m <
2.1 BAYESIAN FILTERING
From a Bayesian filtering perspective (Dore et al. 2009). data assimilation problem at each time can be divided into two steps: a prediction step that advances the dynamical system and an update step that assimilates the data. Assuming that the posterior density of the state 2;-1 given obser- vation data f1---1 (B1...., 1t-1), denoted as P(rr-1lBh:t-1) is available at time t 1, with P(roly/1-0 = P(ro) given for time
0, then the prediction step provides the density of a, as
/
Prediction: P(
where P(zz(===1) represents the transition probability governed by the dynamical system in Equa- tion [ Here, 2t depends on V1:-11 only through 24-1- Let P(yelar) denote the likelihood function Then, the update step provides the updated posterior
of the data Y/ given state from Equation 2 density P(x, W1:c) of the state H, given the new observation data ye by Bayes' rule as
Update: P(x)()1 ((y)(x))) ((rrl ((xslyie-1)
(4)
where the normalization constant or model evidence term ((yelVra-1) is given by P(yolvin-1) f P(r.ll -1)d2, which is typically intractable to compute.
2.2 DIFFUSION MODELS AND THE ENSEMBLE SCORE FILTER
recent development for new particle filter is the EnSF, where the ensemble of samples can be drawn from the posterior distribution through diffusion process. It builds the recent advance- ment in score-hased generative modeling (Ho et al.] 2020 generating high-fidelity samples from distribution by first computing stochastic differential equations (SDE) (Song et al./ 2021). This circumvents knowing the true den- sity as it only needs to generate samples from the posterior distribution. Notably, the diffusion model takes the form of a noisy forward stochastic differential equation process
function
P(x) and then sampling with it by solving
dx = fzz r)dr + g(T)
(5)
which transforms data from an arbitrary distribution an isotropic Gaussian for E T = |0, T] Here, f(z, r) is drift term, g(T) is a diffusion term, and w is d-dimensional Wiener process. It has the corresponding reverse-time SDE that runs backward in time from 7
dx (f, (r) log (=())d g())dw

---

Published as a conference paper at ICLR 2025
where is another Wiener process independent of w. This reverse-time SDE can be solved to sample from the distribution P(x) if the score function V, log P(x,)
) is known, e.g., by the Euler- Maruyama scheme at the discrete time steps 0 To < TI ... TT
Let æt,r denote the state at physical time and diffusion time r. Bao et al. (2024b) propose the posterior score of the update step in Equation 4 as
V log ((seerll1:)) Va log P( M(T) V. log ((ya(or,r))
with the prior score function in the first term given an integral, see Appendix A that is computed by Monte Carlo estimation, the likelihood function P(ye (+x,>) in the second term given explicitly from the observation map in Equation 2 and a monotonically decreasing damping function ((T)
with (()) = 1 and ((1) = 0, c.g., h(T) = 1 - T. Note that the posterior score function is consistent with Bayes' rule in Equation att T O. To this end, one step of the EnSF algorithm is presented in Algorithm Din in Appendix A along with the details on the prior and posterior score functions.
EnSF offers several advantages over EnKF, particularly high dimensions, where EnKF typically requires many more ensemble members to accurately estimate the density, while EnSF leverages the explicit likelihood function and diffusion process to generate samples. Moreover, EnSF does not assume the approximate linearity of the dynamical system, making it highly applicable to high- dimensional and nonlinear filtering problems, such as a dilliionndimesionall Lorenz-96 system (Bao ct al. 2024b) and surface quasigeostrophic model (Bao et al. 2024a).
2.3 CHALLENGE OF ENSF FOR SPARSE OBSERVATIONS
A key limitation of EnSF in high-dimensional nonlinear filtering problems occurs when the observation data are sparse, which is the case for most practical applications. For example, suppose that the observation map H (x+) are [S]. which only makes observations of xy in the dimensions corresponding to a small subset S C {I,..., d} with |S| << d. In this case, the gradient the log-likelihooo v log P(y (xt) vanishes in the dimensions outside of S, shown in figure pointwise observations of shallow water wave height with [S] d/100.
Figure 2: The gradient of the log-likelihood func- tion Vz log P(ye (.s.) (right) vanishes at the points where the sparse observation data Yt (middle) do not have any information of the state x, (left).
As the likelihood function in Equation 7 does not account for the spatial correlation, the information from the sparse observations cannot be ade- quately assimilated to the unobserved states. The vanishing gradients significantly limit the effec- tiveness of the standard EnSF, which we demonstrate empirically in Section 4.1.1
3 LATENT ENSEMBLE SCORE FILTER
To address the key limitation of EnSF in the case of sparse observations due the vanishing gradient of the lig-liieli function, we propose a latent representation of the sparse observations by variational autoencoder (VAE) and match this latent representation to the encoded full state for the sake of consistent data assimilation in the latent space, for which we employ EnSF with the additional advantage of VAE regularization of the latent variables.
The advantages are twofold: it enables not only to exploit the benefits of standard latent score- based generative models (Vahdat et al. 2021 Rombach et al. 2022), which offer optimized sampling speed and expressivity, but also to create expressive mapping from the sparse observation space to the latent space, with the output being a full-dimensional observation in the latent space that circumvents the vanishing gradient issue of the likelikld function. Additionally, the successful application of diffnt-diffusion models to videos (Blattmann et al. (023) presents a potential avenue for exploring latent-assimllation coupled with latent dynamics.
4

---

For use the same decoder Published as a conference paper at ICLR 2025
3.1 COUPLED VARIATIONAL AUTOENCODERS WITH LATENT-SPACE MATCHING
The Variational Autoencoder (VAE) Kingma & Welling 2014). well its variants such as VQ-VAE (Van Den Oord et al., 2017 and Info VAE (Zhao et al. 2019), provides a compressed rep- resentation of high-dimensional distribution through bottlenecking. We employ VAEs compress both the sparse observations and the full states into consistent latent representation.
Observation Encoder
Mean Matching
Decoder
Varianca Matching
Encoder
Figure 3: A coupled VAE for consistent latent representations of sparse observations and full states.
Specifically, as shown in Figure S we formulate a coupled VAE with two encoders. The state encoder, & [d + R²r, encodes the full state I; into latent variable 2 with distribution is N (Mr, diag(a?)), where (He, of) E(x²) and 2 = μt + Eg. The observation encoder, Eabs Rm 2 encodes the sparse observation Yt into another latent variable zohs with distribution 1/th = N(u diag((o (orhx)²)) 2²), where (ug (orh))2) = Eebs (Vt) and 2/th obs + o consistent reconstruction of the full state from the latent representation, D Rd Rdd. To assimilate the encoded latent data to the encoded latent state, we seek to match the two latent representations the consistent input. In addition, we include a regularization of the latent distribution. To this end, we formulate the following loss function
(0) =11x; D(=)) 1|. - 21
(Reconstruction Term)
(8)
+ -118(x)) - Eatr (ve)))))1
(Latent Matching Term) (9)
+ (M + ADKI. ADK (1
(Regularization Term)
(10)
where 0 represents the parameters of the encoders & and Eats and the decoder D, for which we use convolutional neural networks (CNNs) this work. In practice, we found that latent matching for the mean is sufficient. The regularization by the Kullback-Leible divergence between the latent distribution and the reference (0, 4.) of multivariate standard normal distribution
1 r
(kL
Σ (or)) (ue)² log((o.)))
(11)
in Equation 10 is a weight parameter. We set À = 10-5 as in Rombach et al. 2022 for the of the above
encoder of the latent diffusion model. We minimize empirical loss function loss function over all time steps in each of the training trajectories of the dynamical system.
Both latent representations use the same decoder reconstruct the full state. By training a coupled VAE that minimizes the reconstruction eтror of the true state conditioned on the observations, we aim obtain a latent representation of the sparse observations that consistent with that of the full states. This approach is similar to the Generalized Latent Assimilation (GLA) representation proposed by [Cheng et al. (2023) for EnKF, but with the added matching of the latent representations of the states and sparse observations, as illustrated in Figure Experimentally, training the encoders without matching the latent representations can lead to a solution in which the latent observations segregated from the latent states, rendering the data assimilation inaccurate. Using these two separate encoders, we can align the latent representations so that we are able to apply EnSF a full-dimensional manner, with extremely sparse observations.
5

---

and the observation yi- Published as a conference paper at ICLR 2025
3.2 LATENT ENSEMBLE SCORE FILTER
For simplicity, let the latent state & and latent data Çe be defined as the encoded state and data as
& (ne; of) El (2re) and &. (nett, (o) Eabs (y)e)
(12)
We conduct data assimilation in the latent space by assimilating the latent data Ça into the latent state &t. Mathematically, the latent state & approximately follows (up to the VAE reconstruction error) the latent dynamical system
& = ((M(((2z-)((
(13)
as a result of the full dynamical system Equation [ and the VAE with the reparametrizattoon
24 (E-- = 1t-1 + 0--1 6 at time t - 1. Similarly, the latent data & approximately satisfies
= Ears (2((((E))) + 71)
(14)
result of the observation map in Equation 2 and the VAE with the reparametrization (5x) (4 -o at time f. When calculating the likelihood of the observation in the latent space, we simply apply Equation [4] and use an additive latent observation noise t. estimated from the full space observation noise V/ through the observation encoder. Since the coupled encoders aim to match the full states and the sparse observations in their latent representations, we can further approximate the latent observation map in Equation [4 as an identity map, i.e., Se R &r + 't, which simplifies and accelerates the computation of the gradient of the oog-likelhhood function by avoiding automatic differentiation through the map in Equation 14 with respect to the latent state.
Given that the latent state follows the latent dynamics and using the approximate latent observations, consider the Bayesian filtering problem in the latent space by defining the prediction step as
Prediction: P( 11=-11) P(s
(15)
with the transition probability P(& (e-1) determined by the latent dynamical system in Equation The update step for the posterior of the latent state & given the latent data &e is defined as
13
P(E P(silS1a-1)
Update: P(EelSi=)
(16)
where P(G(E))) is the likelihood of the latent data given latent state &,, computed via the latent observation map in Equation 14 The latent normalization term P(Ge(Sia-1) P((ks-n)) also generally intractable.
We employ the diffusion- beasued ensemble score filter described in Section 2.2 to address the Bayesian filtering problem in the latent space, replacing the state x, with the latent state & and the observation Ye with the latent observation G. Note that both the latent state & and the latent observation Çe have the same dimension, 2r < d. One step of the Latent-EnSS is presented in Algorithm
Algorithm 1 One Step of Latent EnSF
Input: Ensemble of the states (x--1} from distribution
State encoder E, observation encoder Eots. and decoder D.
Output: Ensemble of the states } from the posterior distribution P(re(Mia))
Encode the state & = E(x--1) and the observation Ç = Eabs (Vx).
Run the (stochastic) dynamical system model in Equation 13 from {{--1} to obtain samples {
for T = ....... To do
Estimate the prior score Ve log P(E (1a-1) ((Et,1((a-1)) Estimate the posterior score g log P(&,_ (G.)) using Equation in the latent space.
using Equation 18 and 19 in the latent space.
Solve the reverse-time SDE in Equation 6 generate the ensemble of latent states }. end for
Sample 2 = Mt + a 6 from the ensemble of the latent states & (Mu, ox) and samples of e,. Decode the ensemble of latent variables {2,} to the ensemble of the full states by = D(2.).
Because the sparse observations have been mapped to the same dimension as the state in the latent
space, we can circumvent the EnSF's weakness associated with the vanishing gradient of the likeli- hood. In addition, since the encoder is a flexible neural-netwond representation, the model can con- strain the latent observation map Hame in the latent space to be the identity function H &

---

Published as a conference paper at ICLR 2025
which allows for analytical solutions for the score function without needing automatic differentia- tion. This enables us to fully leverage the computational speed-up of the EnSF. Finally, the VAE's regularization, which encourages the latent vector to follow N (0. I), simplifies the hyperparameter search of the appropriate noise "/- We also make some adjustments to the EnSF to handle numerical instability issues for small-scaled variables, detailed in Appendix C
EXPERIMENTS
In this section, we illustrate the limitations of using the EnSF for sparse observations and demon- strate the improved accuracy and fast convergence of the Latent-EnSF compared to several data assimilation methods for a synthetic complex physical system modeled by shallow water equations with sparse observations in both space and time. We also highlight the merits of Latent- EnSF for medium-range weather forecasting using real-world data and an ML-based dynamical system.
4.1 SHALLOW WATER WAVE PROPAGATION
We consider the propagation of shallow water waves described by the shallow-water equations [Vreugdenhil 1994), system of hyperbolic PDEs that capture free-surface fluid flow through the conservation of momentum and mass, with three state variables: the water height and the two compo- nents of the water velocity field. Such models find practical applications in hydrology for predicting flood waves and in oceanography for modeling tsunami wave propagation.
In the experiment, consider water wave propagation with the initial water displacement modeled as local Gaussian bump perturbation from a flat surface in a square domain discretized on a uniform 150x 150 grid, adapted from|josthy 20199. We run the simulation for 2000 steps; specific parameters are provided in Appendix B Figure 4 displays the evolution of the water height and its pointwise sparse observations with data reduction factors of 4x (75 75) and 225x (10 10).
For the data assimilation, we start with an initial state for which the Gaussian bump perturbation is shifted away from the true initial state, as shown in the top of Figure We then generate an ensemble of samples (100 by default) by adding a small amount of noise to this initial state. Data assimilation performed with observations every (20 by default) simulation steps, unless otherwise stated. We use 100 forward Euler steps to solve the reverse-time SDE,see more details in Appendix A
4.1.1 ENSF AND SPARSITY
In Section 2.3
we illustrated the vamishing gradient of the log-likelihood function for the dimensions in which observations are not available. Here we empirically demonstrate this effect on EnSF for data assimilation with sparse observations different levels of sparsity.
In Figure [ we plot the performance in terms of relative root mean square error (RMSE) of the assimilated state compared to the true state as the observation space becomes increasingly sparse. We can observe that increasing the sparsity (from no reduction, to 4x and 225x reduction) of the observation data leads increasing assimilation errors and slower convergence EnSF with 20 samples in the ensemble. When the observation data are reduced by a factor of 225 times, the RMSE of EnSF does not decrease in time, even though we assimilate densely in time (every simulation time step) contrasted to every 20 time steps. In contrast, the Latent- EnSF achieves small relative errors with fast convergence even in the case of extreme sparsity, with the errors the same level achieved by EnSF with the observation of the full state. These findings align with [Bo]] jet al. 2024a) where the sparsity in their experiments for EnSF is limited to 2x data reduction.
the true state by
When we increase the number of samples from 20 to 1000 for EnSF with 225x data reduction, we can observe decreasing assimilation errors in time from Figure 5 This is likely due to that the prior score plays more important role compared to the likelihood function with more sparse observations, and increasing number of samples lead to more accurate estimates of the prior score.
4.1.2 LATENT-ENSF RESULTS
We train the coupled VAE as shown in Figure 3 using convolutional layers with imput dimension
3 x 150 x 150 and latent dimension 4 x 10 x 10 400 with 4 channels. The following experiments
7

---

Published as a conference paper at ICLR 2025
Peteter
Potater
and Sparsity Comparison
223 ti 23 Samplo
22ne brdr. 52 Sumplo
33t4 Endr Semplee
wayple
22
Spane
Figure 5: Relative RMSE of EnSF and Latent-EnSF for observations with differ- ent sparsity levels and number of samples.
Figure 4: Evolution of states and sparse observations. 225x sparsity markers are enlarged for ease view.
use 10 observation grid, corresponding to the 225x sparsity in Figure We compare the Latent-EnSF to the Latent-LETKF and the Latent-EnKF which share the same VAE weights as the Latent-EnSF (Figure For the approximate " observation noise terms in the latent space, we adopt heuristic by taking the mean standard deviation values given by encoded latent states (more details in Appendix C). We additionally compare against Vanilla LETKF which has desirable properties to handle sparse observations. Our approach is much better in terms of assimilation speed, accuracy, and efficiency (see Table [ for all cases, see higher observation noise in Appendix
SWE Results for Noise 10.0%
Latent-EnSF vs Lacant-LETKS, Sparse in Time and Space
lareeecier time interval
0.5
Petturbed
LETK+
0.5
400D Latan-LETKr
Balettve
400D Laten-tudr
4000 stmm
07-
000
Figure 6: Relative RMSE of Latent-EnSF com- pared baselines with 10% observation noise.
Figure 7: Relative RMSE of Latent-EnSF vs
Latent-LETKF for varying temporal sparsity.
In Figure 7 we report the comparison of Latent-EnSF and Latent-LETKF for their performances with different temporal sparsity in the observations, which come in every 5, 20, and 50 time steps with increasing sparsity. As we can observe, the assimilation accuracy of Latent-EnSF is not affected by the increased sparsity: only the burn-in time is slightly increased. contrast, the Latent-LETKF is very sensitive to the temporal sparsity, with much slower convergence for larger sparsity, while Latent-EnSFF is robust to sparse observations not only in space but also time in this experiment.
Figure 8 shows the comparison of the assimilated results of sparse observations using different methods at the first time step, time step 1000, and 2000 (the final time step), and we provide an extended version in Appendix B We can observe that even though the perturbed dynamics is very far from the ground truth, the Latent-Ent EnSF is able to push the perturbed dynamics to match the ground truth very well with small error, while EnSF fails to do so with its error very close to the difference between the perturbed and true states. Moreover, Latent-EnSh performs the best among the data assimilation methods in the same latent space.
8

---

Published as a conference paper at ICLR 2025
Ground
Pertrbed
Latent-LETXP
Latent-EuSt
Eusp
Figure 8: Comparison of different methods at time step 2000, with the ground truth, perturbed states, and assimilated states the top and their corresponding errors from the ground truth the bottom.
4.2 REDIUM-RANGE WEATHER FORECASTING
Sparse data assimilation is particularly important for accurate medium-range weather forecasting due to the complexity of the problem which can skew results of machine learning models. The com- monly used dataset for conducting these weather forecasting experiments is the ERA5 (Hersbach) al., 2020), coming from the European Centre for Medium-Range Weather Forecasts (EMCWF). This dataset comes with 40 years of reanalyzed data, and can support various degrees of fidelity. Cur- rently, ML-based prediction approaches such as FourCastNet (Pathak et al. 2022). Pangu Weather (Bi et al. 2023), and GraphCast (Lam et al. 2023) have shown to be extremely efficient when ap- plied to weather forecasting compared to standard PDE-based approaches. We adopt the 21 variables used in the FourCastNet paper and train FourCastNet model
coarse subset of the dataset for forecasting. When training the VAE for data assimilation, we create a ResNet-style architecture sim- ilar to what constructed in the shallow-water experiments. The encoder downsamples by factor of 4, resulting in latent dimension size of 32 x 36 x 18 20736. We use this as a demonstrative task show the feasibility of applying Latent-EnSF to complex systems, so we train and test on 10 years of data with 2.5" resolution, with a grid of size 144 x 72, as a proof of concept. Evaluations conducted once a
are by averaging the metrics across nine 41yday windows. Data assimilation day, incorporating observations at 64x sparsity rate, i.e. grid of 18 x observations.
As shown in Figure 9 the forecasting by the FourCastNet- ML model without data assimila- tion becomes increasingly inaccurate because of the accumulated errors, which limits its relatively reliable prediction for the first days. The data assimilation errors by EnSF follow closely those of FourCastNet forecasting without data assimilation, which implies that EnSF fails with the data assimilation with sparse observations. In contrast, the Latent- EnSF approach quickly levels off the error in the first couple of days and preserves the accuracy for all the following days, despite the increasing inaccuracy of the FourCastNet forecasting. Note that the Latent-LETKe leads to increas- ing assimilation errors, much higher than Latent-EnSF, with increasingly inaccurate FourCastNet forecasting. To demonstrate that the FourCastNet is helpful in the data assimilation, we also run the Latent-EnSF without a known dynamical model, in which case assume the dynamics of the state does not evolve, i.e., Tz+1 (x,) dt- From the comparison of Latent-EnSy No Model and Latent-EnSF (with FourCastNet),, we observe that the ML-based FourCastNet forecasting does help for the data assimilation, even though it becomes increasingly inaccurate in the long run.
We see Figure 2 that the assimilation for the state Z500 is much more accurate than that of U500, which can also be observed in Appendix E after 41 days of forecasting by FourCastNet and by data assimilation. This is due to the much higher state complexity with rich local characteristics and sharp changes in the u-component of the wind velocity U500, which would demand more data (with less sparsity) beyond the very sparse 18 9 observations for more accurate data assimilation.
4.3 COMPUTATIONAL EFFICIENCY
To examine the computational efficiency, we report the mean and standard deviation of the time it takes the same machine (Nvidia A6000 AMD 7543) for step of the assimilation by each

---

Published as a conference paper at ICLR 2025
ERAD 23CO Relative RMSSE
ERAS USDD Reletive RMSE
LaternErnS Model
0.6
Roletive T8/S2
0.4
FourCastier (Ne DAI
R155
Latert-En5F
Lasone E15F No Hodd
Lanene Ub!
Figure 9:
Relative RMSE for the assimilation of Z500 (geopotential at 500hPA) and U500 (u- component of wind at 500hPA) on the FRAS with respect to the number of days of forecasting.
of the methods in Table T for the two applications. The latent dimension 2r is applicable to all methods except EnSF with dimension d in the full space. The computational cost of the Latent- LETKF (parallelized with 64 cores) increases approximately linearly with respect to the dimensions compared to the much larger increase of the Latent-EnKe In comparison, the Latent- Latent-EnSF achieves the highest accuracy, fastest convergence, and still uses the smallest time. By comparison of the two applications, we see that although the latent dimensions differ by scale of 50, the Latent-EnKF takes over 700 times as long for assimilation in the higher latent dimension, compared to 68 times for Latent-LETKF and only times for Latent-EnSF, which implies the much better scalability of the Latent-EnSF compared to the other methods. Remarkably, even with the added complexity of doing forward pass through the encoder and decoder, the assimilation by LatenttEnSSp is still faster than that by EnSF in the full space, making Latent-EnSF suitable for real-time data assimilation.
Table I: Mean and standard deviation of wall-clock time in seconds of different assimilation ap- proaches for both our experiments, obtained using Nvidia A6000 GPU + AMD 7543 CPU.
Dataset SWE
Latent Dim
Latent-EnKF
Latent-LETKF
Lafent-EnSF
EnSF 0.223 0.002 0.501 0.001
400 20736
0.056 + 0.002
0.103 = 0.001 7.034 = 0.055
0.051 = 0.001 0.198 = 0.008
ERA5
41.161 0.640
CONCLUSION
In this paper, introduced Latent-EnSSF, novel data assimilation method that addresses the joint challenges high state dimensionality and high observation sparsity in nonlinear Bayesian filter- ing problems-kkyy key issues for existing approaches. We proposed coupled VAE for information compression into the latent space, which is trained to consistently match the full states and sparse observations in their latent representations while enforcing latent regularization and ensuring con- sistent reconstruction of the full states. We also presented the corresponding approximate latent dynamical system and a simplified latent observation map to accelerate the sampling process in the latent space. We demonstrated the high accuracy, fast convergence, and efficiency of the pro- posed method compared to other approaches for challenging data assimilation problems using both synthetic data and real-world data that sparse in both space and time.
For further avenues of interest, extending Latent-EnSF to incorporate latent dynamics (iia et al. 2024) could further enhance computational efficiency. Additionally, exploring more complex archi- tectures could enable data assimilation for unstructured observation points and continuous time data where extremely sparse observations occur in stochastic manner. Overall, Latent-EnSF presents promising direction for advancing data assimilation techniques and offers an efficient way to utilize machine learning models to simulate real-world systems with improved accuracy.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work is partially supported by NSF grant # 2325631, # 2245111, and # 2245674. We acknowl- edge helpful discussions with Prof. Felix Herrmann, Dr. Jinwoo Go, and Pengpeng Xiao.
REFERENCES
Mark Asch, Mare Bocquet, and Maëlle Nodet. Data Assimilation. Society for Industrial and Ap- plied Mathematics, Philadelphia, PA, 2016. doi: 10.1137711998811119944 URL https:
//epubs. si aaa. wwg/ dwn/ abs 10. 111.111.11111111111
Feng Bao, Hristo G. Chipilski, Siming Liang, Guannan Zhang, and Jeffrey S. Whitaker. Nonlinear ensemble filtering with diffusion models: Application to the surface quasi-geostrophic dynamics, 2024a. URL https::////. /eattt://araae /ttp:v/aavvvvvv.v http://awvvv.oo//////////////
absa 2404. 00844
Feng Bao, Zezhong Zhang, and Guannan Zhang. An ensemble score filter for tracking high- dimensional nonlinear dynamical systems, 2024b.
00983
Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Accurate medium- range global weather forecasting with 3d neural networks. Nature, 619(7970::533-538, 2023.
Craig H. Bishop, Brian J. Etherton, and Sharanya J. Majumdar. Adaptive sampling with the ensemble transform kalman filter. part i: Theoretical aspects. Monthly Weather Re- view, 129(3):420 436, 2001. doi: 10.1 (TL.T///S22000022000190000000000 CO:2. URL Dettps://journale amet soc. wtw.eettoooorr//.oooooaoo 1520-0493. 2001 _001_122________1 129__0220______0 _tw_004000__________________ 2.0.......... xm1
Andreas Blattmann, Robin Rombach, Huan Ling, Tim Dockhorn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis. Align your latents: High-resolution video synthesis with latent diffusion mod- cls. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp.
June 2023.
Yuming Chen, Daniel Sanz-Alonso, and Rebecca Willett. Reduced-orde autodifferentiable ensem- ble kalman filters. Inverse Problems, (0(12):22400) 2023.
Sibo Cheng, Jianhua Chen, Charitos Anastasiou, Panagiota Angeli, Omar K Matar, Yi-Ke Guo, Christopher C Pain, and Rossella Arcucci. Generalised latent assimilation in heterogeneous duced spaces with machine learning surrogate models. Journal of Scientific Computing, 94(1):11, 2023.
R. Courant, K. Friedrichs, and H. Lewy. On the partial difference equations of mathematical physics. IBM J. Res. Dev., 112):215-234, mar 1967. ISSN 0018-8646. doi: 112.022.. URL
112. 0215
Alessio Dore, Matteo Pinasco, and Carlo S. Regazzoni. Chapter 9 multi-modal data fusion tech- niques and applications. In Hamid Aghajan and Andrea Cavallaro (eds.), Multi-Camera Net- works, pp. 113-237. Academic Press, Oxford, 2009. ISBN 978-0-12-374633-7
doi: https:
URL https: htwp::::/wwww.. http///wwwwcccicceerreee....t...
com/ pctp://cccaoooe/rrrraaae///////p
Geir Evensen. Sequential data assimilation with a nonlinear quasi-geostrophic model using monte carlo methods forecast error statistics. Journal of Geophysical Research: Oceans, 99(C5): 10143-10162, 1994.
Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz- Sabater, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, Adrian Simmons, Cor- nel Soci, Saleh Abdalla, Xavier Abellan, Gianpaolo Balsamo, Peter Bechtold, Gionata Biavati, Jean Bidlot, Massimo Bonavita, Giovanna De Chiara, Per Dahlgren, Dick Dee, Michail Dia- mantakis, Rossana Dragani, Johannes Flemming. Richard Forbes, Manuel Fuentes, Alan Geer, Leo Haimberger, Sean Healy, Robin J. Hogan, Elías H6lm. Marta Janisková, Sarah Keeley, Patrick Laloyaux, Philippe Lopez, Cristina Lupu, Gabor Radnoti, Patricia de Rosnay, Iryna

---

Published as a conference paper at ICLR 2025
Rozum, Freja Vamborg, Sebastien Villaume, and Jean-Noël Thépaut. The era5 global reanal-
ysis. Quarterly Journal of the Royal Meteorological Society, 146(730):1999-2066 2049, 2020. doi:
https://doi.org/100 10.1000/ URL
com/
doi/abs 0. 1002/0000000 3803
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in
Neural Information Processing Systems, 33:6840-6857 2020.
Langwen Huang, Lukas Gianinazzi, Yuejiang Yu, Peter D. Dueben, and Torsten Hoefler. Diffda: diffusion model for weather-scale data assimilation, 2024. RwcLhttpps://rxii.ooo
2401. 05932
Brian R. Hunt, Eric J. Kostelich, and Istvan Szunyogh. Efficient data assimilation for spatiotemporal
chaos: A local ensemble transform kalman filter. Physica D): Nonlinear Phenomena, 230(():112- 126, 2007. ISSN 0167-2789. doi: http:///o..ogg//...bb..,.....pee
URL https://
www. wtw.::/wwcccocicccoooooooooo.aaa
Data
Assimilation.
jostbr. water,
2019.
R. E. Kalman. A New Approach Linear Filtering and Prediction Problems. Journal of Basic
Engineering, 82(1):35 03 1960. ISSN 0021-9223. doi: 10.1115/1.366252.. URL https:
7/doi. wrg. o/doi.org/10.1111 1ww./1.11111111111
Diederik Kingma and Max Welling, Auto-encoding variational bayes. International Conference on Learning Representations, 2014.
Hans R. Künsch. Particle filters, Bernouli,, 19(4), September 2013. ISSN 1350-7265, doi: 10.3150/
12-bejsp07.
hwtp::::/xxxoii..cooooo/o 0tt///11..11/11111 BEJSPO7
Remi Lam, Alvaro Sanchez-Gonzale,, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Fer- ran Alet, Suman Ravuri, Timo Ewalds, Zach Faton-Rosent Weihua Hu, et al. Learming skillful medium-range global weather forecasting. Science, 382(6677):1416-4422 2023.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, Pedram Hassanzadeh. Karthik Kashinath, and Animashree Anandkumar. Fourcastnet: A global data- driven high-resolution weather model using adaptive fourier neural operators, 2022. URL
https: //exxiv owg/ abs /2002 11214
Florence Rabier and Zhiquan Liu. Variational data assimilation: theory and overview. In Proc. ECMWF Seminar on Recent Developments in Data Assimilation for Almosphere and Ocean, Reading, UK, September 8-12, pp. 29-43, 2003.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Patlern Recognition (CVPR), pp. 10684-10695, June 2022.
François Rozet and Gilles Louppe. Score-based data assimilation. In Proceedings of the 37th International Conference Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2024. Curran Associates Inc.
Daniel Alnz-Alonso, Andrew Stuart, and Armeen Taeb. Inverse problems and data assimilation, volume 107. Cambridge University Press, 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations, 2021. URL
httpss://axxe brg. abs /0t./211ll... 13456
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 11287-11302. Curran Associates, Inc., 2021. URL http:s///prcceeeee......aroo.... http://apppppeeeee/eepaaaeapee files/
12

---

Published as a conference paper at ICLR 2025
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
C.B. Vreugdenhil. Numerical Methods for Shallow-Water Water Flow. Water Science and Technology Li- brary. Springer Netherlands, 1994. ISBN 9780792331643. URL https://books google.
com/ books? iid= SuchwO1mgaBs
Pengpeng Xiao, Phillip Si, and Peng Chen. LD-EnSF: Synergizing latent dynamics with ensemble score filters for fast data assimilation with sparse observations. arXiv preprint 1.Niv:2411119300 2024.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Balancing learning and inference in variational autoencoders. In Proceedings of the aaai conference on artificial intelligence, vol- ume 33, pp. 2019.
13

---

Published as a conference paper at ICLR 2025
A
ENSF AND ALGORITHM FOR ONE STEP ASSIMILATION
The EnSF method
[ao et al. 2024b) uses 7 and the following drift and diffusion terms
dloog a dr ttr and gz)
ds log a
[J 7)
(17)
dr
dr
with a, = - 1 (a) (a) and 83 = < + 7(1 6g) for two small positive hyperparameters ca and 6g, to ca 0.05 and Eg in our experiment, which aid in avoiding the collapse of the generated
samples. This choice leads the distribution Fe,r Naa,T1,, B.I) conditioned 21,0 Zt, which results in the prior score function
/ Pz,t,
Ve log (V)a-1) = V log
(18)
/ Dt.r
a,
w (1
x ((x
8
where the gradient V taken with respect to Tt, and the weight (ar,r 24)) ) is given by
P( (1e)
wx xe) f P(xt, (x4) P(
(19)
Both the score function in Equation 18 and the weight in Equation can can be evaluated by sample average approximation with random samples from the distribution P(x obtained in the prediction step in Equation
Algorithm One Step of EnSF
Input: Ensemble of the states { from distribution P(xx--1lv::-- and the observation yt-
Output: Ensemble of the states { from the posterior distribution P(re))1=))
Run the (stochastic) dynamical system model in Equation from {x--1] } to obtain samples {xz)-
for .... To
Estimate the prior score V. log P(xtt.lMMiaa using Equation 18 and 19
Estimate posterior score V, log P(x ()) using Equation
Solve the reverse-time SDE in Equation 6 to generate the ensemble {
end for
B ADDITIONAL DETAILS ON SHALLOW WATER EQUATIONS
We consider size L x L with L = 106 in each direction and constant depth h 100 for the topography of the floor for the shallow water wave propagation. The simulation is run for 2000 time steps using upwind scheme with time step At 0.1 Air with Ar L/150 to satisfy the FriraattFriedrichs--Leee height and its sparse observations with 4x (75 x 75) and 225x (10 x 10) data reduction.
Vah condition (Courant et al. 1967), see Figure for the evolution of the
C PINNING DOWN THE LATENT NOISE
When applying the Latent-EnSE we do not have ground-truth values for the latent observation noise. To obtain an approximate estimate, we calculated the corresponding noise of the latent state by encoding noisy full states. In Figures [1 and 12] we calculate the distribution of the latent error between the true latent state (given by encoding the non-noised state) and the latent encoding the noisy observations for two different levels of noise. Note that the standard deviation of the latent noise (calculated across ten trajectories) is stable across all time steps, and the mean is approximately
In practice when applying the EnSF to problems where the observation noise is small, specifying the corresponding observation noise such that % < 0.05 can result in numerical instability issues when sampling the posterior distribution by reverse time SDE with the drift and diffusion terms given in Equation 17
14

---

Published as a conference paper at ICLR 2025
Ground Treuth
Perambed
Intent-Enkir
LaterL-LIFTRY
1abenk-8::57
TriSp
Observadion
Pecturbed Brrec
Lateat EakE Btroe
Latent LETD7 Brewr
Latent 3:57 Error
BusE Exroc
Trush
Pettrrba
Observatioe
Perturbed Eroor
Latent-ExXP Erroe
Lakens4ETKP Error
Latest-Ense Emor
EnöF Errer
Ground Truth
Petturbed
Lateni-TaK
Latenk-CTKT
Ietent-EnSt
Obserration
Perturbed
SXKP
Latent LETKF
Latert: EnSP Error
Pv57
Figure 10: Comparison of different methods time step (top rows), 1000 (middle two rows), and 2000 (bottom two rows), with the ground truth, perturbed states, and assimilated states on the top and their corresponding errors from the ground truth at the bottom of each two rows.
This numerical instability issue is also present in the VAE latent space where the representations become compacted as regularization parameter À increases. We address this issue by multiplying the latent representations with scalar constant vitttets conducting the EnSF for the rescaled latent representation, and subsequently scaling the posterior samples back the original latent represen- tation. We conducted small grid search experiment when working with the shallow water wave propogation problem Section 4.1 where thatent = 20 seemed to work reasonably.
Remarkably, this Vaatent applies well in all our experiments, which implies that this scaling hyperpa- rameter stays fairly steady across all problems. Intuitively, this is consistent; the latent regularization conducted by the VAE constrains the latent space to be in similar range. In contrast, when EnSF conducted on the full state space, such advantages are lost and a different scaling parameter e is needed for different problems with different scales of the state and observation noise.
15

---

Published as a conference paper at ICLR 2025
Latent Differenee, No se 10.0%
Latent Difference, No 20.0%
Lateed Cafference
0.
Newn Latent Cafheeence
3.400
Risatvv Latcet 02 10100000
Noceve www.eettooooooooooooooooooo OPROCIO
2.25
0.0
10.025
0.03
-0.05
Figure 11: Difference of latent state encoded from full state and latent observation encoded from sparse observation with 10% observation noise.
Difference of latent state encoded from full state and latent observation encoded from sparse observation with 20% observation noise.
Figure 12:
D NOISY RESULTS FOR SHALLOW WATER EQUATIONS
SWE Results for Nôlse 0.0%
SWE Resuits for Nolse 20.0%
1.
Labert
20
LETIKE
4000 Lacks--Cr
0.6
2.2
02-
Figure 13: Relative RMSE LLatent-EnSE com-
Figure 14: Relative RMSE Latent-Ens com-
pared to baselines with observation noise.
pared to baselines with 20% observation noise.
16

---

Published as a conference paper at ICLR 2025
E SAMPLES FOR WEATHER FORECASTING
Chaurration
LaTE Broo
Lats Eime
Figure 15: ERA5 Z500 medium-ranged weather forecasting samples after 41 days, along with the errors. Data assimilation is conducted once a day with 64x sparse observations. We compare against baseline FourCastNet model and Latent-LETKF.
ProrCamNee No Erer
LEI Beer
Loiesi-2250 Enve
Figure 16: ERA5 U500 medium-ranged weather forecasting samples after 41 days, along with the errors. Data assimilation is conducted once a day with 64x sparse observations. We compare against baseline FourCastNet model and Latent-LETKF.
17


Reviews:
[{"review_id": "dNQEdTuXDP", "reviewer": "Reviewer_nywr", "summary": "The Ensemble Kalman Filter (EnKF) enables data assimilation for Nonlinear State-Space Models, but it encounters the problem of vanishing gradients in tasks where observations are sparse in high-dimensional state spaces. This paper proposes a solution to this issue by proposing the Latent Ensemble Score Filter (Latent-EnSF), which embeds states and observations into a latent space via a Variational Autoencoder (VAE). The proposed method claims to be effective for various data assimilation tasks with sparse observations in high-dimensional state spaces.", "questions": "What initial state distributions are used for the experiments? It should be explicitly given.\nRegarding the perturbed trajectory, how is the perturbation applied? It should also be explicitly given.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "This paper's strength lies in demonstrating that data assimilation in the latent space can effectively solve the problem of vanishing gradients.", "weaknesses": "One issue with this paper is that it unfairly compares the LETKF under unfavorable conditions. The Ensemble Kalman Filter does not inherently suffer from the vanishing gradient problem, so embedding into the latent space is not necessarily required. If the goal is to propose a good method for data assimilation in high-dimensional nonlinear state spaces, it should be directly compared with the Localized Ensemble Transform Kalman Filter (LETKF) but not with Latent-LETKF. However, this comparison is not provided.\n\nMoreover, physical laws are often described by differential equations with spatially local interactions, leading to strong local correlations. The LETKF leverages this locality to scale to data assimilation problems with large state spaces by performing local computations in parallel. Comparing the computation speed without considering this parallelization is also misleading.\n\nRegarding Equations (11)-(13), it is mentioned that the VAE's loss function minimizes the empirical loss function, but isn't the state ( $x_t$ ) unobserved? It is not clearly stated how learning is performed when the state ( $x_t$ ) is unobserved.\n\nAdditionally, it implicitly assumes the state can be estimated from observations at a single time point as shown in Equation (11) ($x_t$ can be recovered either from $D(z_t)$ or $D(z_t^{obs})$). Thus, it is unclear whether this method applies to nonlinear observations.\n\n# LETKF is included in the experiments thanks to the authors' feedback. However, how much the initial conditions affect the performance is still unclear.", "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": "c5HeUWlNd5", "reviewer": "Reviewer_76az", "summary": "The paper presents a novel data assimilation method called Latent-EnSF aimed at improving accuracy and efficiency in modeling high-dimensional and nonlinear systems with sparse observations. The authors propose using a coupled VAE with two encoders, one for the full state and another for sparse observations. This ensures consistent encoding of both types into a shared latent space, allowing efficient assimilation with fewer dimensions. Overall, the research topic is of interest in the community of data assimilation. However, from a methodology point of view, the novelty of the proposed approach might be limited to the combination of existing ENSF and latent data assimilation algorithm.", "questions": "A lot of latent data assimilation algorithms make use of the variational DA framework instead of enkf, for ex, Peyron et al., 2021 and Melinc et al, 2023. The authors should also compare 4d-var based latent DA approaches against latent-ENSF.\n\n Instead of only sparse observations, it would also be interesting to compare latent-ensf against existing methods for partial and irregular observations.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The Latent-EnSF method presented in this paper addresses key challenges in data assimilation for high-dimensional systems with sparse observations by leveraging a coupled VAE to compress both state and observation data into a shared latent space (However this idea is not new). This approach not only reduces computational complexity but also enhances the robustness of data assimilation when observation data is sparse, avoiding the gradient vanishing issues that hinder ENSF. Experimental results demonstrate Latent-EnSF's superior accuracy, faster convergence, and computational efficiency compared to existing methods.", "weaknesses": "The Latent-EnSF builds upon EnSF by integrating latent representations that avoid the computational challenges of sparse data. While the idea is sound and logical, performing data assimilation in the latent space is a well-established approach (see, for example, Peyron et al., 2021, 'Latent Space Data Assimilation Using Deep Learning'). Additionally, conducting data assimilation with a shared latent space for both the state vector and observations is also a known concept (see, for example, Cheng et al., 2024, 'Multi-domain Encoder–Decoder Neural Networks for Latent Data Assimilation in Dynamical Systems'). Therefore, from the reviewer’s perspective, the novelty of this paper primarily lies in combining existing latent data assimilation techniques with ENSF.\n\nAs the authors mentioned, since the state and observation vectors are encoded into a common latent space, the transformation function is restricted to an identity mapping. However, this may pose a problem regarding the specification of the observation error covariance matrix in the latent space; specifically, how to transform uncertainties from the full physical space to the latent space remains an important question. In the numerical experiments, the authors should also include noisy observations and demonstrate how latent-ENSF could account for observation uncertainty.", "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": "A9P8pai7q1", "reviewer": "Reviewer_GvWu", "summary": "The paper identified a shortcoming in the recently developed Ensemble Score Filtering method, which is that observations at sparse sites lead to a collapse in the model gradient, by introducing a low-dimensional latent variable encoding, which recovers improved inference in this sparse setting.", "questions": "## What is the actual deficiency of EnSF?\n\nThe empirical results are great.\nWhat is actually happening here theoretically? Since Latent-EnSF and EnSF are both targeting the same likelihood, and the VAE-based Latent-EnSF is making a looser approximation to the posterior than the original model, then it is not sufficient IMO to claim it is simply addressing a problem of missing gradient; Latent-EnSF needs to be jsutified in terms of how it draws better inference from the same information.\nI've expanded on this question in the _weaknesses_ section, but I invite you to answer it here.\n\nAs part of that I also suggest it will be good to clarify which data (training data, inference time observation) informs the final prediction. Questions about that in the next section. \n\n## What is the training distribution?\n\nWhat is the training process used for the NN? What distribution of states is in that dataset? Does an encoder trained for (however many trajectories are in this dataset) generalize badly to data assimilation problems drawn from a different dataset (say, different initial conditions) in a way that EnSF does not? That would be evidence that this method learns a good spatial smoothing, and would also suggest a deficiency in the original EnSF, if it has difficulty in spatial smoothing\n\n## Section 4.1.1\nI'm confused about both figures and the explanation. Can you answer for me the following?\n\n1. Figure 2: I'm not sure what $\\eta$ is. Is it the reconstructed state? If so, in the 4x sparse row we see that the reconstructed solution has a banded structure. Why? I know that the _observations_ updates are sparse, but when we propagate from those back through our diffusion steps, should we not recover a densely sampled field at whatever resolution our simulator $M$ runs?\n2. Figure 3:  \n\n   1. it's not clear what \"perturbed\" means here. From cross checking I think this is the error that a simulator would accrue from failing  to update at all. Why is it constant? For interesting models (e.g. with lyapunov exponent ) I would expect this to grow over time. This leads me to wonder if the systems that are being modeled here are not very interesting, or maybe that RMSE is a bad measure, or that the duaration over which we observe the system is simply not interesting. For the kind of challenging systems mentioned in the paper (e.g. the high dimensional lorenz) I would expect diverging loss metrics under perturbation\n   2. I'm not even sure which model is being plotted in these figures. It looks like a fluid PDE?\n  3. There are no error bars here. I would expect for such a central plot that we would make some attempt to plot error bars under random perturbations and ideally also random initial conditions.\n\n## What sparsity patterns are allowed?\n\nThe proposed method appears to have a curious limitation, which is that the spatially-sparse observations must occur at fixed sites $S$ for all timesteps. This is not required in the EnKF and, although i know it less well, should not be required in the EnSF. Is it required for this model also, or can that be relaxed?\n\n## section 4.1\n\n>However, though the Latent-LETKF beats the Latent-\n>EnKF benchmark, especially with higher latent dimensions, our approach is still much better in\n>terms of assimilation speed, accuracy, and efficiency (see Table 1) for all cases.\n\nSome confusing phrasing. *Is* that table 1 shows? It seems to me that it shows on two benchmarks, Latent-EnSF is the fastest in terms of wall-clock inference time. (Side point: it's ok in my opinion to use minimal wall clock time for operations with a deterministic number of FLOPs, means and standard deviations are not ideal for timing execution which is an asymmetric right-skewed distribution who randomness is mostly due to external factors) However the graph shows that on (I think on some different model?) that Latent-EnSF is more accurate than Latent-LETKF.\nMoreover, some of these methods (the VAE-based ones) require us to train a neural network, and others (EnKF) do not. How much does the NN training time add to these? When I've trained my encoder networks, does that generalise to unseen initial conditions, or do I need to train a new encoder? When do I need to train a new encoder? How costly is that? \nThese is not to say that I think anything is incorrect, but rather that the statements and supporting evidence are messy, and comparisons do not appear to be \"apples-to-apples\". Can you make the statements more precise? or the example more strictly comparable?\n\nWhy LETKF? I'm not advocating an exhaustive pairwise comparison between every possible model in this domain, but if the primary comparison is with LETKF then that method might need to be introduced more. note also that EnKF can handle this case as well, so would have been a logical method to include on the Figrues 4 and 5. If you are not going to do that (and I am sympathetic we we all get asked to do too many comparisons in these conference submissions) then can you make the justification explicit for which comparisons you wish to make?\n\nFigure 8: The paper talks about $\\nabla_x \\log P(y|x)$ and yet graph shows $\\nabla_x P(y|x)$. Why?\n\nFigure 8, figure 2: There is not color bar on the plots ot help us interpret the field values. I assume white means zero valued and the colors are bipolar? If so, I have a further question: The plots show the score function evaluated at sparse spatial locations by showing fewer pixels non-zero with values in a grid of fixed size, implying that gradient of the score function at these locations is also zero. To my mind, would be more correct to say that the score is not sampled at that location and has no value, rather than being zero. We would more conventionally show that by displaying a lower-resolution grid, i.e. with fewer, larger pixels. Moreover, we could imagine in-filling the grid, now? $\\nabla_x \\log P(y_2|x_1)$ will still convey influence, and thus a gradient between the sampling sites of  $x_1$ and $y_2$ if they are spatially near each other. Or: Am I missing something? Do I not understand the conditioning operation correctly in the score function?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This problem is high value, and the approach unifies several interesting ideas.\n\nEmpirical results are impressive.", "weaknesses": "Minor cosmetic and presentational issues. These are not crucial, but I will attempt to note when I see them\n\nMore crucially, theoretical explanation for the empirical success does not seem complete, and details are missing which might allow the reader to deduce this for themselves.\n\nFirstly, there are a number of small errors and typos, e.g. equation 14 introduces the KL \"Divergence\" as  function of the parameters of the Gaussian distribution $\\operatorname{KLD}(\\mu,\\Sigma)$ .That is not a divergence. It looks like it is maybe the conventional KL between a given diagonal Gaussian and a standard diagonal Gaussian? There are also some in the section 4.1.1, which I will address in the Questions. There are a couple more I noticed but did not note down; I'll leave these cosmetic issues aside for the moment so we can get to the substantive weakness and return if there is time.\n\nThe paper's argument on an important assertion, that sparse observations are problematic EnSF updates (l202). However, this part of the argument is not made sufficiently clearly for such a central part of the paper. The first supporting evidence  (figure 8 in appendix b) is confusing - AFAICT it simply shows what a sparse sampling grid looks like, which... doesn't really show anything. That is just a graphic depiction of the form of the observation operator $H$, I think. Section 4.1.1 does heavier lifting, which is that attempts to show that empirically, as the updates get sparser in the sense of sampling a coarse spatial grid, the quality of the EnSF estimates degrades. However, I am not convinced it does show that; It is unclear to me if this is because the experimental evidence itself is lacking; it could be the labeling in the figures. See questions for clarifications on that point.\n\nThe explanation for what problem is being solved here and why the solution works seems incorrect or at least incomplete.\nIf the baseline is EnSF, which uses (up to some discretization error and ensemble approximation etc) the \"true\" model $M$ and the \"true\" observation operator $H$, then it should in principle recover the \"best\" estimate of the system dynamics. the proposed method Latent-EnSF claims to get superior performance by introducing a further layer of approximation, mapping the observations and dynamics into a VAE. Since we have introduces a further approximation then we must assume that all else being equal, the VAE-based method should be worse in accuracy (it might be faster at inference time or something like that).\n\nAs such, if Latent-EnSF  it is doing better in accuracy terms than EnSF, it seems that it should be because it gets to exploit more information at inference time than does EnSF. Presumably this \"extra\" information is in the form of better spatial smoothing encoded in the encoder/decoder networks of the VAE, which can presumably learn the which spatial fields are in the training data. If so, we expect that the information enters the method via the training losses (11,12,13) and also by the training data set. The paper does not seem to interrogate the training losses and also I cannot easily work out what the actual training procedure is for the neural nets, so it is not clear to the reader what smoothing the neural nets are actually learning.", "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": "9SoTIKDmf1", "reviewer": "Reviewer_oZ4F", "summary": "The paper aims to address the problem of assimilating sparse observation data effectively using the ensemble score filter framework for data assimilation. Their method proceeds by encoding both state and observations to a latent space of the same dimension before applying the score-based ensemble filter. This allows one to avoid the issue of having mismatched size of the observation and state dimensions, preventing vanishing gradient of the likelihood. Experimental results demonstrate the efficacy of the approach showing further promise of score-based methods for data assimilation.", "questions": "- The authors should clarify that the Kullback-Leibler divergence in (14) is taken with respect to the standard normal measure. When talking about statistical divergences, one should specify which two measures the divergence are taken with respect to and in which order.\n- Is the loss in (11)-(13) correct? The reconstruction term should be marginalised over the latent variable $z_t$ if this were to be trained like a VAE.\n- Related to the above question, can the loss (11)-(13) be derived from first principles via variational Bayes? Or is the loss more just \"VAE-inspired\"?\n- Do we train for the encoder/decoder at every timestep $t$? Do the authors take into account its training cost if so? It seems quite expensive to achieve in practice.\n- At the moment, it seems like the method only works in the case where the observations are placed on a regular grid. Is there a way to extend this to a more realistic scenario where the observations are less regularly sampled?\n- Could the authors please elaborate more on how they compute the latent observation noise $\\hat{\\gamma}_t$? The authors only give a passing remark but this does not seem so obvious to me.\n- How do the results in the experiments compare to more traditional data assimilation methods like the EnKF or 4DVar? Does the latent EnSF still outperform these in the sparse observation setting?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper addresses a critical issue with many of the score-based filters being proposed recently, namely, that the observations being assimilated are too idealistic. The methodology proposed in this work brings it closer to more realistic scenarios, where observations are sparsely distributed. Results in the experiments are encouraging, showing that the latent EnSF can significantly outperform standard EnSF or LETKF in latent space. The writing is generally clear and sufficient details are provided to implement the proposed model.", "weaknesses": "One key weakness of this work is that baseline comparisons against traditional DA methods are missing. For example, how does it compare to standard ensemble Kalman filter (LETKF) or 4DVar in physical space? The authors consider LETKF in latent space but this approach is not actually used in practice. If we want to really demonstrate efficacy of the approach, a comparison with these traditional method should be provided. These methods are likely to perform quite well on the presented examples so it will help us see how much extra benefit we can gain by replacing them with the proposed latent EnSF framework. Furthermore, some details of the algorithm were unclear to me, which you can find below in __Questions__. Finally, plots to show the sparse observations can be improved by increasing the marker size. At the moment, it is impossible to see where the observations are without zooming in significantly.", "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": "2SoXh0uncJ", "reviewer": "Reviewer_se1f", "summary": "This paper develops the recently proposed Ensemble Score Filters (EnSF) by incorporating Variational Autoencoder (VAE) to define EnKS in the latent space derived by VAE. The resulted algorithm, hence named latent-EnSF, inherits the effectiveness of EnSF in high-dimensional data assimilation, but also solves the issued of EnSF when dealing with the sparse observations. Despite some possible typos, the proposed method is clearly explained, supported with thorough comparison with various alternative algorithms including vanilla EnSF, latent-EnKF, latent LETKF, FourCastNet using two models in water wave propagation and medium-range weather forecasting.", "questions": "* Line 222: What do you mean by \"the output as a full-dimensional observation in the latent space\"? Is it \"in the full space\"?\n\n* Equation (14), KL divergence is defined between two distributions. What are they in your equation (14)? Did you mean to say $KL(N(\\mu_t, \\Sigma_t)|| N(\\mu_t^\\text{obs}, \\Sigma_t^\\text{obs}) )? Then pay attention to the signs.\n\n* Line 314: can you comment on the computational cost of this step because equation (16) involves both encoder and decoder.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Most of the implementation in the latent space of EnSF is straightforward. The added benefit of handling sparse observations by running EnSF in the latent space is impressive and has potential impact in data assimilation in high dimensions with sparse observations. I certainly appreciate the thorough comparison offered in this paper.", "weaknesses": "1. EnSF has already achieved great results in high-dimensional data assimilation problems. VAE is used as a dimension reduction method. The mysterious capability of handling sparse observations by combining them remains largely un-explained. Though there is some on line 323, the paper may benefit in clarity from more elaboration in this aspect.\n\n2. The most difficult part to understand is the latent matching term (13) on line 244. Two encoders, the state encoder and the data encoder, compress two different things. Why do they need to match the codes (mean and variance) in the latent space? I see the authors report segregated latent representations without such term, but I still find it unnatural and lack some explanation on its contribution to the capability of handling sparse observations.\n\n3. The precision of solutions to inverse problems naturally depends on the data information available. While I am impressed by the comparison between EnSF and latent-EnSF in terms of handling the sparse observations, some comments on the limitation of the proposed latent-EnSF is also appreciated.", "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": "d18weFZNJV", "reviewer": "Reviewer_KEVJ", "summary": "This paper introduces Latent-EnSF, a data assimilation method designed to address the challenges of high-dimensional state spaces and sparse observations in nonlinear Bayesian filtering problems. By extending the Ensemble Score Filter (EnSF) with a latent space representation using coupled variational autoencoders (VAEs), the method transforms both full states and sparse observations into a shared encoded latent space. Then, it leverages EnSF, a score-based filtering method, to filter in the encoded latent space.  The paper claims that leveraging the proposed method can effectively overcome the limitations posed by sparse observations. The numerical results show the improved accuracy, faster convergence, and enhanced computational efficiency in high-dimensional nonlinear systems, with applications in shallow water wave propagation and medium-range weather forecasting.", "questions": "1. The primary question I have regarding this paper is about the training of the VAE. Specifically, it is unclear whether the encoders and decoder are trained at each time step $t$ or if they are pre-trained. Based on Equations (11–13) and Figure 1, it appears that the input data $ x_t$ for the encoder is derived iteratively during the data assimilation process (from predictive distribution), suggesting that the VAE might need to be retrained at each step. If so, this would significantly slow down the filtering process, raising questions about the efficiency claims presented in the experimental results. Additional clarification on the training procedure of the VAE is necessary to understand how the reported computational efficiency is achieved.\n\n2. The presentation of Section 3.2 is somewhat confusing, and a few points could use further clarification:\n\n   - Could you please elaborate on the need for introducing $ \\xi_t $ and $ \\zeta_t $? Are these variables different from the $ z_t $ and $ z_t^{\\text{obs}} $ used in Section 3.1, and if so, how?\n   - In lines 320–321 on page 6 and in Algorithm 1, why is the reparameterization trick needed again to sample the latent variable $ z_t $? My understanding is that after solving the reverse-time SDE in Equation 6, you already have the ensemble of latent states as samples. Could you clarify this step?\n\n3. Minor Issues:\n\n - Consider moving Figure 1 to Section 3, as it would better illustrate the proposed method in context.\n\n - In Section 3.1, certain notations are unclear, such as $ \\Sigma_t $ and $ \\sigma_t $. Clearer definitions would improve readability.\n\n - On page 5, you note that the GLA proposed by Cheng et al. (2023) can face computational issues when the latent space is high-dimensional. While I agree, typically, as shown in this paper, the latent dimension $ 2r \\ll d $, so GLA’s efficiency may not necessarily be compromised in this case. Did you consider including a comparison with GLA in the experiments? \n\n - While the method is based on the EnSF, it seems that the discussion on the existing data assimilation methods for high-dimensional state spaces and sparse observations is missing e.g. (and more),\n\n      - Chen Y, Sanz-Alonso D, Willett R. \"Reduced-order Autodifferentiable Ensemble Kalman Filters,\" *Inverse Problems*, 2023; 39(12):124001.\n\n     I believe including the discussion of more related work could make the proposed method's positioning more comprehensive.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper presents a clear and well-motivated rationale for addressing high-dimensional data assimilation with sparse observations in EnSF. The proposed method is intuitively and logically structured, and the experimental results validate its effectiveness and improvements over some of the existing approaches.", "weaknesses": "A potential weakness of this paper is the absence of certain important details regarding the developed method, which are noted in the discussion section.", "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)