Latent-EnSF: A Latent Ensemble Score Filter for High-Dimensional Data Assimilation with Sparse Observation Data

## Abstract

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.

## 1 Introduction

Many complex physical systems are traditionally modeled by partial differential equations (PDEs). First-principle PDE-based modeling and simulation have been demonstrated to be powerful in making predictions of complex systems in various scientific and engineering fields. However, in practical applications, significant discrepancy of the simulation-based prediction and the reality may arise from different sources, e.g., model inadequacy, uncertainties in model parameters, boundary and/or initial conditions, exteral forcing/loading terms, numerical approximation errors, etc. Meanwhile, data-driven machine learning (ML) models have been significantly developed and employed to make system predictions in the last few years, which have been shown to be extremely efficient in the case of large dynamical systems such as weather forecasting, a problem that requires the supercomputer of hours to run for the European Centre for Medium-Range Weather Forecasts (ECMWF), but one that can be solved by neural network surrogates such as FourCastNet (Pathak et al., )
in a matter of seconds. These properties also make them particularly suitable for ensemble forecasting and uncertainty quantification. However, these models are also often subject to the key limitation of divergence from reality in the long run, resulting in significant accumulated errors for long-term predictions because of various uncertainties and the autoregressive structure of the ML models.

To mitigate the discrepancy and accumulated errors and make more accurate predictions, data assimilation plays an essential role in
incorporating additional observation data of the reality into the current PDE/ML-based prediction models. By adapting the PDE/ML state to match the observation data, the accuracy of future state predictions can be markedly improved. Data assimilation (Sanz-Alonso et al., ; Asch et al., ) approaches such as Kalman filter (KF) (Kalman, ), particle filters (Künsch, ), and their variants have been widely applied to problems such as weather prediction, geophysical modeling, robotics, and many more areas.

The basis of many current data assimilation methods used in practice is the KF (Kalman, ). It parameterizes its state with the mean and the covariance, and presents the optimal solution to the data assimilation problem given that certain linearity properties of the transition and observation function are fulfilled. However, in practice, KFs end up being computationally expensive for high-dimensional states and biased for nonlinear problems.
An alternative approach is particle filter, which represents the density by an ensemble of particles. They relax the assumption of distributional form, e.g., Gaussian distribution, made for distribution-based approaches such as the KF. The EnKF (Kalman, ) is a particular particle filter that originates from the KF. Instead of keeping the covariance matrix, it computes the sample covariance of the ensemble of particles. Due to its robust performance, it has become a widely used method for applications today.
However, most particle filters, including EnKF, suffer from the curse of dimensionality of requiring an exponentially large number of particles to accurately describe the distribution in high dimensions.
To address this challenge, various extensions to EnKF have been proposed, such as the Localized Ensemble Transform Kalman Filter (LETKF) (Hunt et al., ), where covariance localization techniques as well as advancements from the Ensemble Transform Kalman Filter (Bishop et al., ) are applied.

Nonetheless, a few major roadblocks hinder efficient data assimilation for high dimensional systems. Particularly, standard data assimilation approaches have quite limited capabilities in high-dimensional environments, resorting to assumptions such as localization in order to make them more applicable. High dimensional systems and complex models can also limit the capability and efficiency of models such as 4D-Var (Rabier & Liu, ), which is extremely inefficient since it requires propagating through the physical model many times. A recent approach named DiffDA (Huang et al., ) uses diffusion models and masked interpolations of the weather data to assimilate sparse observations, but the approach falls behind pure interpolation after a few steps. Bao et al. () proposed a novel method called Ensemble Score Filter (EnSF) which has achieved great results on high-dimensional data assimilation problems by leveraging the power of score-based diffusion models to sample from the posterior distribution efficiently. However, sparse observations prove to be very challenging for EnSF, limiting its applicability to real-world scenarios where observation data are limited. Meanwhile, approaches such as Generalized Latent Assimilation (GLA) (Cheng et al., ) attempt to use dimension-reduction techniques to conduct data assimilation in latent space, but are still subject to the constraints of the standard EnKF for complex high-dimensional systems.

In this work, we propose an efficient data assimilation method called Latent-EnSF that builds on Ensemble Score Filters (EnSF) (Bao et al., ) to address the joint challenges of high dimensionality and data sparsity. See Figure 1 for a workflow of the Latent-EnSF. We propose to construct coupled Variational Autoencoders (VAE) (Kingma & Welling, ) to encode the sparse observations into a latent space in which we conduct data assimilation by EnSF using diffusion-based sampling. The ensemble of samples of the assimilated latent state is then decoded to an ensemble of samples of the full state. Importantly, our paper addresses the key weakness of the EnSF when dealing with sparse observations, and retains its ability to efficiently assimilate observations for high-dimensional nonlinear systems. We demonstrate this performance for both PDE and ML based prediction models of some challenging examples, experimenting with data that is sparse in both space and time.

Figure: Figure 1: Flow of the Latent-EnSF. An ensemble of prior states $x_{t}|y_{1:t-1}$ is assimilated with sparse observation $y_{t}$ in the latent autoencoder space to obtain samples $x_{t}|y_{1:t}$ from the posterior.
Refer to caption: extracted/5848781/images/ga_ensf.png

We will begin by formally introducing data assimilation problems with a review of the EnSF in Section 2, follow up with how we integrate this with coupled variational autoencoders in Section 3 to create our Latent-EnSF, and finally conclude with experiments in Section 4.

## 2 Dynamical Systems and Data Assimilation

To model a dynamical physical system, we denote the state at a certain time $t$ as $x_{t}\in\mathbb{R}^{d}$, where $d$ is typically very high for complex models. For a system with discrete time steps $t=1,2,\dots$, given initial condition $x_{0}$, we then model the evolution of the state from the time $t-1$ to time $t$ as

$$ $x_{t}=M(x_{t-1},\varepsilon_{t}),$ (1) $$

where $\varepsilon_{t}$ is a process noise term coming from a known distribution; this term accounts for hidden interactions and numerical errors. The corresponding observation $y_{t}\in\mathbb{R}^{m}$ at time $t$ is given by

$$ $y_{t}=H(x_{t})+\gamma_{t},$ (2) $$

where $\gamma_{t}$ (also coming from a known distribution, typically Gaussian) represents the observation noise because of instrumental inaccuracies, etc. Both the model map $M:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ and the observation map $H:\mathbb{R}^{d}\to\mathbb{R}^{m}$ can be nonlinear.

Now, the model of the dynamical system in Equation 1 may lead to inaccurate representation of the ground truth with increasing discrepancy in time because of, e.g., model inadequacy or input uncertainty. The goal of data assimilation is to assimilate the observation data
of the true state
to the model of the dynamical system and recover an accurate representation of the true state.
Note that in practical applications the observation data
can be sparsely distributed in both space and time,
which makes the recovery of the true state very challenging, especially for high dimensional problems with the state dimension $d\gg 1$, and the observation dimension $m$ much smaller than $d$.

### 2.1 Bayesian Filtering

From a Bayesian filtering perspective (Dore et al., ), a data assimilation problem at each time $t$ 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 $x_{t-1}$ given observation data $y_{1:t-1}=(y_{1},\dots,y_{t-1})$, denoted as $P(x_{t-1}|y_{1:t-1})$, is available at time $t-1$, with $P(x_{0}|y_{1:0})=P(x_{0})$ given for time $t-1=0$, then the prediction step provides the density of $x_{t}$ as

$$ $\textbf{Prediction: }P(x_{t}|y_{1:t-1})=\int P(x_{t}|x_{t-1})P(x_{t-1}|y_{1:t- 1})dx_{t-1},$ (3) $$

where $P(x_{t}|x_{t-1})$ represents the transition probability governed by the dynamical system in Equation 1. Note that $x_{t}$ depends on $y_{1:t-1}$ only through $x_{t-1}$.
Let $P(y_{t}|x_{t})$ denote the likelihood function of the data $y_{t}$ given state $x_{t}$ from Equation 2. Then, the update step provides the updated posterior density $P(x_{t}|y_{1:t})$ of the state $x_{t}$ given the new observation data $y_{t}$ by Bayes’ rule as

$$ $\textbf{Update: }P(x_{t}|y_{1:t})=\frac{P(y_{t}|x_{t})P(x_{t}|y_{1:t-1})}{P(y_ {t}|y_{1:t-1})},$ (4) $$

where the normalization constant or model evidence term $P(y_{t}|y_{1:t-1})$ is given by $P(y_{t}|y_{1:t-1})=\int P(y_{t}|x_{t})P(x_{t}|y_{1:t-1})dx_{t}$, which is typically intractable to compute.

### 2.2 Diffusion Models and the Ensemble Score Filter

A recent development for a new particle filter is the EnSF, where the ensemble of samples can be drawn from the posterior distribution through a diffusion process. It builds on the recent advancement in score-based generative modeling (Ho et al., ) in generating high-fidelity samples from a distribution by first computing a score function $\nabla_{x}P(x)$ and then sampling with it by solving stochastic differential equations (SDE) (Song et al., ). This circumvents knowing the true density 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

$$ $dx=f(x,\tau)d\tau+g(\tau)dw,$ (5) $$

which transforms data from an arbitrary distribution to an isotropic Gaussian for $\tau\in\mathcal{T}=[0,T]$. Here, $f(x,\tau)$ is a drift term, $g(\tau)$ is a diffusion term, and $w$ is a $d$-dimensional Wiener process. It has the corresponding reverse-time SDE that runs backward in time from $\tau=T$ to 0 0 as

$$ $dx=[f(x,\tau)-g^{2}(\tau)\nabla_{x}\log P_{\tau}(x)]d\tau+g(\tau)d\bar{w},$ (6) $$

where $\bar{w}$ 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 $\nabla_{x}\log P(x_{\tau})$ is known, e.g., by the Euler–Maruyama scheme at the discrete time steps $0=\tau_{0}<\tau_{2}<...<\tau_{k}=T$.

In the following part, we use $x_{t,\tau}$ to indicate the state at physical time $t$ and diffusion time $\tau$.
The EnSF method (Bao et al., ) uses $T=1$ and the following drift and diffusion terms

$$ $f(x_{t,\tau},\tau)=\frac{d\log\alpha_{\tau}}{d\tau}x_{t,\tau}\text{ and }g^{2} (\tau)=\frac{d\beta_{\tau}^{2}}{d\tau}-2\frac{d\log\alpha_{\tau}}{d\tau}\beta_ {\tau}^{2},$ (7) $$

with $\alpha_{\tau}=1-\tau(1-\epsilon_{\alpha})$ and $\beta^{2}_{\tau}=\epsilon_{\beta}+\tau(1-\epsilon_{\beta})$ for two small positive hyperparameters $\epsilon_{\alpha}$ and $\epsilon_{\beta}$ which aid in avoiding the collapse of the generated samples.
This choice leads to the distribution $x_{t,\tau}\sim N(\alpha_{\tau}x_{t,0},\beta^{2}_{\tau}I)$ conditioned on $x_{t,0}=x_{t}$, which results in the prior score function

$$ $\begin{split}\nabla_{x}\log P(x_{t,\tau}|y_{1:t-1})&=\nabla_{x}\log\int P(x_{t ,\tau}|x_{t})P(x_{t}|y_{1:t-1})dx_{t}\\ &=\int-\frac{x_{t,\tau}-\alpha_{\tau}x_{t}}{\beta_{\tau}^{2}}\omega(x_{t,\tau} ,x_{t})P(x_{t}|y_{1:t-1})dx_{t},\end{split}$ (8) $$

where the gradient $\nabla_{x}$ is taken with respect to $x_{t,\tau}$, and the weight $\omega(x_{t,\tau},x_{t})$ is given by

$$ $\omega(x_{t,\tau},x_{t})=\frac{P(x_{t,\tau}|x_{t})}{\int P(x_{t,\tau}|x_{t}^{ \prime})P(x_{t}^{\prime}|y_{1:t-1})dx_{t}^{\prime}}.$ (9) $$

Both the score function in Equation 8 and the weight in Equation 9 can be evaluated by sample average approximation with random samples from the distribution $P(x_{t}|y_{1:t-1})$ obtained in the prediction step in Equation 3. In the update step in Equation 4, the posterior score function in the EnSF method (Bao et al., ) is formulated as

$$ $\nabla_{x}\log P(x_{t,\tau}|y_{1:t})=\nabla_{x}\log P(x_{t,\tau}|y_{1:t-1})+h( \tau)\nabla_{x}\log P(y_{t}|x_{t,\tau}),$ (10) $$

with the prior score function in the first term given in Equation 8, the likelihood function $P(y_{t}|x_{t,\tau})$ in the second term given explicitly from the observation map in Equation 2, and a damping function $h(\tau)$ that is monotonically decreasing with $h(0)=1$ and $h(1)=0$, e.g., $h(\tau)=1-\tau$. Note that the posterior score function is consistent with the Bayes’ rule for the posterior in Equation 4 at $\tau=0$.

To this end, a brief description of one step of the EnSF algorithm is presented in Algorithm 1.

Figure: Algorithm 1 One Step of EnSF

There are a number of advantages that EnSF provides over EnKF, particularly in high dimensions, where a much larger number of ensemble members are typically required for EnKF to have a good estimate of the approximate Gaussian density, while EnSF leverages the explicit likelihood function and the diffusion process to generate samples. In addition, EnSF does not assume approximate linearity of the dynamical system, making it highly applicable to high-dimensional and nonlinear filtering problems, e.g., a million-dimensional Lorenz-96 system (Bao et al., ) and a surface quasigeostrophic model (Bao et al., ). However, the key limitation of EnSF to filtering with sparse observation arises also from the explicit use of the likelihood function, as explained below.

### 2.3 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_{t})=x_{t}[S]$, which only make observations of $x_{t}$ in the dimensions from the subset $S\subset\{1,...,d\}$, where $S$ may have a much smaller cardinality $|S|\ll d$ than the full dimension. In this case, the gradient of the log-likelihood $\nabla_{x}\log P(y_{t}|x_{t})$ vanishes in the dimensions in $\{1,...,d\}/S$, as shown in Figure 2, where we calculate $\nabla_{x}P(y_{t}|x_{t})$ for one of our experiments in Section 4.1 with 1/16 of the total dimensions used for the observations. Note that, the gradient is nonzero only at the points which are observed as seen in Figure 2. Though the ramifications are minimal for sufficiently dense observations, the vanishing gradients for higher sparsity scenarios significantly limits the effectiveness of the standard EnSF, which we demonstrate empirically in section 4.1.1.

Figure: Figure 2: The gradient of the log-likelihood function $\nabla_{x}P(y_{t}|x_{t})$ (right) vanishes at the points where the sparse observation data $y_{t}$ (middle) do not have any information of the state $x_{t}$ (left).
Refer to caption: extracted/5848781/images/score_x.png

To mitigate this weakness of the EnSF,
rather than conducting data assimilation in the full state space with sparse observations, we propose to transform the data assimilation into a learned latent space with sufficiently preserved information and significantly reduced dimension.

## 3 Latent ensemble score filter

To address the key limitation of EnSF in the case of sparse observations due to the vanishing gradient of the log-likelihood function, we propose a latent representation of the sparse observations by a 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 allows us to not only make use of the benefits of standard latent score-based generative models (Vahdat et al., ; Rombach et al., ) which offer optimized sampling speed and expressivity, but also create a expressive map from a sparse observation space to a full-dimensional latent space. Additionally, the successful application of latent-diffusion models to videos (Blattmann et al., ) offers a potential application of exploring latent-assimilation coupled with latent dynamics.

### 3.1 Compression by variational autoencoder

The Variational Autoencoder (VAE) (Kingma & Welling, ) provides a compressed representation of a high-dimensional distribution by means of bottlenecking.
Note that since this is a family of models, it admits many specific instances such as VQ-VAE (Van Den Oord et al., ), InfoVAE (Zhao et al., ), etc., which can promote desirable properties such as independence and adaptable priors in the latent space. VAEs have two components, an encoder $\mathcal{E}$ and a decoder $\mathcal{D}$. The encoder casts the original state $x\in\mathbb{R}^{d}$ into a low dimensional representation $\mathcal{E}(x)\in\mathbb{R}^{2r}$ such that $2r\ll d$. Then, the model splits $\mathcal{E}(x)$ into mean $\mu\in{\mathbb{R}^{r}}$ and variance $\sigma^{2}\in{\mathbb{R}^{r}}$, and samples from the normal distribution $N(\mu,\Sigma)$ with $\Sigma=\text{diag}(\sigma_{1}^{2},\dots,\sigma_{r}^{2})$ by applying a reparameterization trick $z=\mu+\sigma\cdot\epsilon$, where $\epsilon\in N(0,I_{r})$. The decoder then constructs an approximation of the state as $\tilde{x}=\mathcal{D}(z)$.

VAEs are trained by optimizing for the evidence lower bound (ELBO), which is composed of a reconstruction term in the state space (e.g. the MSE $||x-\tilde{x}||_{2}^{2}$) and a Kullback–Leibler divergence (KLD) regularization term in the latent space,
which is given by

$$ $\text{KLD}(\mu,\Sigma)=\sum_{i=1}^{r}-\log(\sigma_{i})+\frac{\sigma_{i}^{2}+ \mu_{i}^{2}}{2}-\frac{1}{2},$ (11) $$

which skews $N(\mu,\Sigma)$ to be close to $N(0,I_{r})$.

### 3.2 Coupled VAEs with Latent-Space Matching

We use VAEs to compress both the sparse observations and the full states to the latent space. To achieve consistent latent representations of sparse observations and full states, we couple the VAEs in the latent space as shown in Figure 3.

Figure: Figure 3: A coupled VAE for consistent latent representations of sparse observations and full states.
Refer to caption: extracted/5848781/images/vae_obsmatching.png

Specifically, we formulate a coupled VAE with two encoders, one encoder $\mathcal{E}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{2r}$ that encodes the full state $x_{t}$ to a latent variable $z_{t}$ with distribution $N(\mu_{t},\Sigma_{t})$, where $(\mu_{t},\sigma^{2}_{t})=\mathcal{E}(x_{t})$, and $z_{t}=\mu_{t}+\sigma_{t}\cdot\epsilon_{t}$, and the other encoder $\mathcal{E}_{\text{obs}}:\mathbb{R}^{m}\rightarrow\mathbb{R}^{2r}$ that encodes the sparse observation $y_{t}$ to another latent variable $z_{t}^{\text{obs}}$ with distribution $N(\mu^{\text{obs}}_{t},\Sigma^{\text{obs}}_{t})$, where $(\mu^{\text{obs}}_{t},(\sigma_{t}^{\text{obs}})^{2})=\mathcal{E}_{\text{obs}}(
y_{t})$ and $z_{t}^{\text{obs}}=\mu^{\text{obs}}_{t}+\sigma^{\text{obs}}_{t}\cdot\epsilon_{t}$.
To assimilate the encoded latent data to the encoded latent state, we seek to match the two latent representations for their means and variances. To build a consistent reconstruction of the full state from the latent representation, we use the same decoder $\mathcal{D}:\mathbb{R}^{r}\to\mathbb{R}^{d}$. To this end, we formulate the following loss function to account for all three contributions:

$$ $\displaystyle\ell_{t}(\theta)$ $\displaystyle=||x_{t}-\mathcal{D}(z_{t})||_{2}^{2}+||x_{t}-\mathcal{D}(z^{ \text{obs}}_{t})||_{2}^{2}$ (Reconstruction Term) (12) $\displaystyle+\text{KLD}(\mu_{t},\Sigma_{t})+\text{KLD}(\mu^{\text{obs}}_{t}, \Sigma^{\text{obs}}_{t})$ (Regularization Term) (13) $\displaystyle+||\mu_{t}-\mu^{\text{obs}}_{t}||_{2}^{2}+||\Sigma_{t}-\Sigma^{ \text{obs}}_{t}||_{2}^{2}$ $\displaystyle\text{ (Latent Matching Term)},$ (14) $$

where $\theta$ represents the parameters of the encoders and decoder, for which we use convolutional neural networks (CNNs) in this work. Then for the training of the parameter $\theta$ in the coupled VAE, we minimize an empirical loss function as a sum of the above loss function over the training trajectories of the dynamical system, and over all the time steps $t$ in each trajectory.

Note that both of the latent representations use the same decoder for reconstruction of the full state. By training a coupled VAE which minimizes the error for the reconstruction of the true state conditioned on the observations, we can strive to obtain a latent representation of the sparse observations consistent to that of the full states.
This approach is similar to the Generalized Latent Assimilation (GLA) representation proposed by Cheng et al. (), but with the addition of matching the means and variances of the state and observation encoders, as illustrated in Figure 3. Experimentally, training the encoders without matching the latent representations can lead to a solution where the latent observations are segregated from the latent states, yet result in sufficient reconstructions. Using these two separate encoders, we can match the latent representations such that we are able to apply EnSF in a full-dimensional manner, even with extremely sparse observations.

Some early papers have utilized autoencoder architectures to conduct data assimilation in various ways. GLA proposed in Cheng et al. ()
uses the EnKF to assimilate observations in the latent space of the VAE. However, for latent spaces with high dimensions, EnKFs are computationally expensive and also subject to the curse of dimensionality. Meanwhile, overly constricting the latent dimension to simplify the computation would lead to magnified reconstruction errors, so the GLA algorithm needs to strike a balance between efficiency and accuracy. Additionally, depending on the type of architecture, local structures may be lost when encoding the state vector into the latent vector, so localized methods such as the LETKF (Hunt et al., ) can lose their advantage.

### 3.3 Latent ensemble score filter

For simplicity, let the latent state $\xi_{t}$ and latent data $\zeta_{t}$ be defined as the encoded state and data as

$$ $\xi_{t}=(\mu_{t},\sigma_{t}^{2})=\mathcal{E}(x_{t})\text{ and }\zeta_{t}=(\mu_ {t}^{\text{obs}},(\sigma_{t}^{\text{obs}})^{2})=\mathcal{E}_{\text{obs}}(y_{t}),$ (15) $$

respectively. We conduct data assimilation in the latent space by assimilating the latent data $\zeta_{t}$ to the latent state $\xi_{t}$. Mathematically, the latent state $\xi_{t}$ follows approximately (up to the VAE reconstruction error) the latent dynamical system

$$ $\xi_{t}\approx\mathcal{E}(M(\mathcal{D}(z_{t-1}(\xi_{t-1})),\varepsilon_{t}))$ (16) $$

as a result of the full dynamical system in Equation 1 and the VAE with the reparametrization $z_{t-1}(\xi_{t-1})=\mu_{t-1}+\sigma_{t-1}\cdot\epsilon_{t-1}$ at time $t-1$. The latent data $\zeta_{t}$ approximately satisfies

$$ $\zeta_{t}=\mathcal{E}_{\text{obs}}(H(\mathcal{D}(z_{t}(\xi_{t}))+\gamma_{t})$ (17) $$

as a result of the observation map in Equation 2 and the VAE with the reparametrization $z_{t}(\xi_{t})=\mu_{t}+\sigma_{t}\cdot\epsilon_{t}$ at time $t$. When calculating the likelihood of the observation in the latent space, we can simply Equation 17 and use an additive latent observation noise $\hat{\gamma}_{t}$ estimated from the full space observation noise $\gamma_{t}$ through the observation encoder. As the coupled encoders seek to match the full states and the sparse observations in their latent representations, we can further approximate the latent observation map in Equation 17 as an identity map, i.e., $\zeta_{t}\approx\xi_{t}+\hat{\gamma}_{t}$, which simplifies and accelerates the computation of the gradient of the log-likelihood function by avoiding automatic differentiation through the map in Equation 17 with respect to the latent state.

With the latent state following the latent dynamics, and the approximate latent observations, we consider the Bayesian filtering problem in the latent space by formulating the prediction step as

$$ $\textbf{Prediction: }P(\xi_{t}|\zeta_{1:t-1})=\int P(\xi_{t}|\xi_{t-1})P(\xi_{ t-1}|\zeta_{1:t-1})d\xi_{t-1},$ (18) $$

with the transition probability $P(\xi_{t}|\xi_{t-1})$ detemined by the latent dynamical system in Equation 16. The update step for the posterior of the latent state $\xi_{t}$ given the latent data $\zeta_{t}$ is defined as

$$ $\textbf{Update: }P(\xi_{t}|\zeta_{1:t})=\frac{P(\zeta_{t}|\xi_{t})P(\xi_{t}| \zeta_{1:t-1})}{P(\zeta_{t}|\zeta_{1:t-1})},$ (19) $$

where $P(\zeta_{t}|\xi_{t})$ is the likelihood of the latent data given latent state $\xi_{t}$, which can be computed by the latent observation map in Equation 17. The latent normalization term $P(\zeta_{t}|\zeta_{1:t-1})=\int P(\zeta_{t}|\xi_{t})P(\xi_{t}|\zeta_{1:t-1})d
\xi_{t}$ is also generally intractable.

We employ the diffusion-based ensemble score filter in Section 2.2 for the above Bayesian filtering problem in the latent space by replacing the state $x_{t}$ with the latent state $\xi_{t}$, and the observation $y_{t}$ with the latent observation $\zeta_{t}$. Note that the latent state $\xi_{t}$ and the latent observation $\zeta_{t}$ have the same dimension $2r\ll d$. We present one step of the Latent-EnSF in Algorithm 2.

Figure: Algorithm 2 One Step of Latent EnSF

Because the sparse observations have been mapped to the same dimension as the state in the latent space, we are able to circumvent the weakness of the EnSF with vanishing gradient of the log-likelihood function when dealing with the sparse observation problems. In addition, since the encoder is a flexible neural-network representation, the model is able to restrict the latent observation map $H_{\text{latent}}$
in the latent space to be the identity function $H_{\text{latent}}(\xi_{t})=\xi_{t}$, allowing for analytical solutions for the score function without the need for automatic differentiation. This means that the model can leverage the full computational speed-up of the EnSF in the latent space. Finally, the regularization done by the VAE such that the latent vector tends to follow a $N(0,1)$ distribution simplifies hyperparameter search of the appropriate noise $\gamma_{t}$ for real world problems.

### 3.4 Addressing EnSF’s numerical instability on small scales

In practice, when applying the EnSF to problems where the values of the states and their variances are very small, specifying the corresponding observation noise such that $\gamma_{t}\ll 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 7.

This numerical instability issue is also present in the VAE latent space where the components are regularized such that $\mu_{t}\approx 0$ and $\sigma_{t}\approx 1$. We address this issue by multiplying the latent representations with a scalar constant $\psi_{\text{latent}}$, conducting the EnSF for the rescaled latent representation, and subsequently scaling the posterior samples back to the original latent representation. We conducted a small grid search experiment when working with the shallow water wave propogation problem in Section 4.1, where $\psi_{\text{latent}}=500$ seemed to work reasonably.

Remarkably, this $\psi_{\text{latent}}$ applies well in all our experiments, which implies that this scaling hyperparameter stays fairly steady across all problems. Intuitively, this is consistent; the latent regularization conducted by the VAE constrains the latent space to be in a similar range. In contrast, when EnSF is conducted on the full state space, such advantages are lost and a different scaling parameter $\psi$ is needed for different problems with different scales of the state and observation noise.

## 4 Experiments

In this section, we illustrate the limitation of using the EnSF for sparse observations, and demonstrate 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 demonstrate the merits of Latent-EnSF for medium-range weather forecasting with real-world data and a ML-based dynamical system.

### 4.1 Shallow water wave propagation

We consider the propagation of shallow water waves described by shallow-water equations (Vreugdenhil, ), which are a system of hyperbolic PDEs that describe fluid flow of free surface with the conservation of momentum and mass, with three state variables including the water height and two components 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, we consider water wave propagation with the initial water displacement modeled by a local Gaussian bump perturbation from the flat surface in a square domain of size $L\times L$ with $L=10^{6}$ m in each direction and a constant depth $h=100$ m for the topography of the floor. For the simulation, we use a finite difference code adapted from jostbr (), with the domain discretized into a uniform grid of $150\times 150$.
The simulation is then run for 2000 time steps using a upwind scheme with a time step $\Delta t=0.1\frac{\Delta x}{\sqrt{gh}}$ with $\Delta x=L/150$ to satisfy the Courant–Friedrichs–Lewy condition (Courant et al., ), see Figure 5 for the evolution of the height and its sparse observations with 4x ($75\times 75$) and 225x ($10\times 10$) data reduction.

For the data assimilation, we start with a mis-estimated initial state where the Gaussian bump perturbation is shifted away from the true initial state. 100 ensemble members are then taken by adding a negligible amount of noise to the mis-estimated state. The data assimilation is conducted with observations every 20 time steps unless otherwise stated, for a total of 100 assimilation steps, with an observation noise variance of 1. When conducting the EnSF, we take 100 forward Euler steps to solve the reverse-time SDE, and use hyperparameters $\epsilon_{\alpha}=0.05$, $\epsilon_{\beta}=0$.

#### 4.1.1 EnSF and Sparsity

In Section 2.3, we illustrated the vanishing 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 at different levels of sparsity.

Figure: Figure 4: Evolution of states and sparse observations.
Refer to caption: extracted/5848781/images/samples_grid_with_perturbed.png

In Figure 5, we plot the performance in terms of relative root mean square error (RMSE) of the assimilated state as the observation space becomes increasingly sparse. The relative error of the perturbed trajectory and the original trajectory is also plotted as a reference. We can observe that when using the EnSF, increasing the sparsity of the observation data leads to increasing assimilation errors and slower convergence to the true state. In particular, for extreme cases such as when the observation data are reduced by a factor of 225 times, there is almost no discernible difference between the assimilated error of EnSF compared to a non-assimilated trajectory, or EnSF does not bring any improvement in the data assimilation. In contrast, the Latent-EnSF achieves small relative errors with fast convergence to the true state even in the case of extreme sparsity, with the errors at the same level achieved by EnSF with the observation of the full state. These findings align with Bao et al. () where the sparsity in their experiments for EnSF is limited to a 2x reduction.

#### 4.1.2 Latent-EnSF Results

We train the coupled VAE as shown in Figure 3 using convolutional layers with input dimension $3\times 150\times 150$ and latent dimension $4\times 10\times 10=400$ with 4 channels for the experiments other than that in Figure 6, which is for the evaluation of the performance with respect to the latent dimension size. The following experiments use a $10\times 10$ observation grid, corresponding to the 225x sparse entry in Figure 5. We compare the Latent-EnSF and the Latent-EnKF, which shares the same VAE weights as the Latent-EnSF. For the approximate $\gamma_{t}$ observation noise terms in the latent space, we adopt a heuristic by taking the mean standard deviation values given by encoded latent states for all three approaches. In Figure 6, we can see that the Latent-EnKF struggles with higher dimensional latent spaces, whereas the Latent-EnSF gets more accurate as a result of a decrease in reconstruction error for the VAE, even though they perform similarly for the 100-dimensional case.

In addition to the Latent-EnKF, we test against the Latent-LETKF in high-dimensional latent space, as shown in the right part of Figure 6. This utilizes the LETKF which in practice handles high-dimensional data assimilation much better than the EnKF. However, we can see that though it indeed beats the Latent-EnKF benchmark, especially with higher latent dimensions, our approach is still much better in terms of assimilation speed, accuracy, and efficiency (see Table 1) for all cases.

Figure: Figure 6: Relative RMSE of Latent-EnSF compared to Latent-EnKF (left) and Latent-LETKF (right) with varying latent dimensions, with the error of the perturbed state shown as a reference.
Refer to caption: extracted/5848781/images/dims_rmse.png

Figure 7 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). We can observe that even though the perturbed dynamics is very far from the ground truth, the Latent-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-EnSF performs the best among the data assimilation methods in the same latent space.

Figure: Figure 7: Comparison of different methods at time step 1 (top two 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.
Refer to caption: extracted/5848781/images/swe_1.png

In Figure 8, 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. In contrast, the Latent-LETKF is very sensitive to the temporal sparsity, with much slower convergence for larger sparsity, while Latent-EnSF is robust to sparse observations not only in space but also in time in this experiment.

Figure: Figure 8: Relative RMSE of Latent-EnSF vs Latent-LETKF for varying observation sparsity in time, with the assimilation frequency at every 5, 20, and 50 time steps, respectively.
Refer to caption: extracted/5848781/images/sparse_time_ensf.png

### 4.2 Medium-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 commonly used dataset for conducting these weather forecasting experiments is the ERA5 (Hersbach et al., ), coming from the European Centre for Medium-Range Weather Forecasts (EMCWF). The main variables at play in this dataset include the wind speeds, temperature, geopotential, and humidity, just to name a few, all of which are available at any of the 137 different pressure levels. This dataset comes with 40 years of reanalyzed data, and can support various degrees of fidelity. Currently, ML-based prediction approaches such as FourCastNet (Pathak et al., ), Pangu Weather (Bi et al., ), and GraphCast (Lam et al., ) have shown to be extremely efficient when applied to weather forecasting compared to standard PDE-based approaches. We adopt the 21 variables used in the FourCastNet paper which include wind velocities and geopotential at 3 different pressure levels, along with relative humidity and temperature at a few key pressure levels and the total column water vapor at the surface. For forecasting, we train a FourCastNet model on our coarse subset of the dataset. When training the VAE for data assimilation, we create a ResNet-style architecture similar to what we constructed in the shallow-water experiments. The encoder downsamples by a factor of 2 twice between convolution layers, resulting in a latent dimension size of $32\times 36\times 18=20736$. In between the layers we use Leaky-ReLU layers with a coefficient of 0.2. We use this as a demonstrative task to show the feasibility of applying Latent-EnSF to complex systems, so we train and test on 10 years of data with $2.5^{\circ}$ resolution, with a grid of size $144\times 72$, as a proof of concept. Evaluations are done for the year 2015, where we average the metrics across nine 41-day windows. Data assimilation is done once a day, incorporating observations at a 64x sparsity rate, with 8x reduction in each dimension, or a grid of $18\times 9$ observations.

As shown in Figure 9, the forecasting by the FourCastNet-based ML model without data assimilation becomes increasingly inaccurate because of the accumulated errors, which limits its relatively reliable prediction for the first 14 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-LETKF leads to increasing 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 we assume the dynamics of the state does not evolve, i.e., $x_{t+1}=M(x_{t})=x_{t}$. From the comparison of Latent-EnSF 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.

Figure: Figure 9: Relative RMSE for the assimilation of Z500 (geopotential at 500hPA) and U500 (u-component of wind at 500hPA) on the ERA5 with respect to the number of days of forecasting.
Refer to caption: extracted/5848781/images/era5_relative_rmse.png

We observe in Figure 9 that the assimilation for the state Z500 is much more accurate than that of U500, which can also be observed in Figure 10 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\times 9$ observations for more accurate data assimilation.

Figure: Figure 10: ERA5 Z500 (top two rows) and U500 (bottom two rows) 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 a baseline FourCastNet model and Latent-LETKF.
Refer to caption: extracted/5848781/images/fourcastnet_plot_z500.png

To examine the computational efficiency, we report the mean and standard deviation of the time it takes the same machine (a combination of an Nvidia A6000 GPU with an AMD 7543 CPU) for one step of the assimilation by each of the methods in Table 1 for the two applications. The latent dimension $2r$ is applicable to all methods except EnSF with dimension $d$ in the full space.
We can see that even though the Latent-LETKF leads to higher accuracy than Latent-EnKF, it takes much more time than all the other methods, due to the need to iterate through the dimensions to construct each individual localization matrix (this could be ameliorated by parallel computing in many processors). However, the computational cost increases approximately linearly with respect to the dimensions compared to the much faster increase of the Latent-EnKF. In comparison, the 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 a 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 4 times for Latent-EnSF, which implies the much better scalability of the Latent-EnSF compared to the other methods.
Additionally, note that even with the added complexity of doing a forward pass through the encoder and decoder, the assimilation by Latent-EnSF is still faster than that by EnSF in the full space, making Latent-EnSF highly conducive for real-time data assimilation.

**Table 1: Mean and standard deviation of wall-clock time in seconds of different assimilation approaches for both our experiments, obtained using Nvidia A6000 GPU + AMD 7543 CPU.**
| Dataset | Latent Dim | Latent-EnKF | Latent-LETKF | Latent-EnSF | EnSF |
| --- | --- | --- | --- | --- | --- |
| SWE | 400 | $0.056\pm 0.002$ | $6.567\pm 0.069$ | $0.051\pm 0.001$ | $0.223\pm 0.002$ |
| ERA5 | 20736 | $41.161\pm 0.640$ | $450.213\pm 3.512$ | $0.198\pm 0.008$ | $0.501\pm 0.001$ |

## 5 Conclusion

In this paper, we introduced Latent-EnSF, a novel data assimilation method in addressing the joint challenges of high dimensionality in the state and high sparsity in the observation for nonlinear Bayesian filtering problems, which have been the key and common challenges for the existing approaches. We proposed a coupled VAE for information compression to the latent space, which is trained for a consistent match of the full states and sparse observations for their latent representations and with latent regularization, as well as for a consistent reconstruction of the full states. We also presented the corresponding approximate latent dynamical system and the simplified latent observation map in accelerating the sampling process in the latent space. We demonstrated the high accuracy, fast convergence, and high efficiency of the proposed method compared to other approaches for a challenging data assimilation problem with PDE-based shallow water wave propagation model and synthetic data that is sparse in both space and time. We also applied and demonstrated the advantage of the proposed method to a ML-based model with real-world data for data assimilation in medium-range weather forecasting.

For further avenues of interest, extending Latent-EnSF to incorporate latent dynamics could further enhance the computational efficiency. Additionally, exploring more complex architectures could allow us to conduct data assimilation for unstructured observation points and continuous-time data where extremely sparse observations arrive in a stochastic manner. Overall, Latent-EnSF presents a promising direction for advancing data assimilation techniques and bringing an efficient way of utilizing machine learning models to simulate real-world systems with improved accuracy.

## Acknowledgement

This work is partially supported by NSF grant # 2325631, # 2245111, and # 2245674. We acknowledge helpful discussions with Prof. Felix Herrmann, Dr. Jinwoo Go, and Pengpeng Xiao.

## Peer Reviews

confidence: 3
contribution: 1
initial_rating: 3
presentation: 2
questions: What initial state distributions are used for the experiments? It should be explicitly given.
Regarding the perturbed trajectory, how is the perturbation applied? It should also be explicitly given.
soundness: 2
strengths: This paper's strength lies in demonstrating that data assimilation in the latent space can effectively solve the problem of vanishing gradients.
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.
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.

Moreover, 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.

Regarding 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.

Additionally, 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.

# LETKF is included in the experiments thanks to the authors' feedback. However, how much the initial conditions affect the performance is still unclear.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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.

 Instead of only sparse observations, it would also be interesting to compare latent-ensf against existing methods for partial and irregular observations.
soundness: 3
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.
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.
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.

As 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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: ## What is the actual deficiency of EnSF?

The empirical results are great.
What 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.
I've expanded on this question in the _weaknesses_ section, but I invite you to answer it here.

As 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. 

## What is the training distribution?

What 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

## Section 4.1.1
I'm confused about both figures and the explanation. Can you answer for me the following?

1. 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?
2. Figure 3:  

   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
   2. I'm not even sure which model is being plotted in these figures. It looks like a fluid PDE?
  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.

## What sparsity patterns are allowed?

The 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?

## section 4.1

>However, though the Latent-LETKF beats the Latent-
>EnKF benchmark, especially with higher latent dimensions, our approach is still much better in
>terms of assimilation speed, accuracy, and efficiency (see Table 1) for all cases.

Some 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.
Moreover, 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? 
These 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?

Why 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?

Figure 8: The paper talks about $\nabla_x \log P(y|x)$ and yet graph shows $\nabla_x P(y|x)$. Why?

Figure 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?
soundness: 3
strengths: This problem is high value, and the approach unifies several interesting ideas.

Empirical results are impressive.
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.
weaknesses: Minor cosmetic and presentational issues. These are not crucial, but I will attempt to note when I see them

More 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.

Firstly, 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.

The 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.

The explanation for what problem is being solved here and why the solution works seems incorrect or at least incomplete.
If 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).

As 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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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.
- 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.
- 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"?
- 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.
- 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?
- 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.
- 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?
soundness: 4
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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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"?

* 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.

* Line 314: can you comment on the computational cost of this step because equation (16) involves both encoder and decoder.
soundness: 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.
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.
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.

2. 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.

3. 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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
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.

2. The presentation of Section 3.2 is somewhat confusing, and a few points could use further clarification:

   - 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?
   - 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?

3. Minor Issues:

 - Consider moving Figure 1 to Section 3, as it would better illustrate the proposed method in context.

 - In Section 3.1, certain notations are unclear, such as $ \Sigma_t $ and $ \sigma_t $. Clearer definitions would improve readability.

 - 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? 

 - 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),

      - Chen Y, Sanz-Alonso D, Willett R. "Reduced-order Autodifferentiable Ensemble Kalman Filters," *Inverse Problems*, 2023; 39(12):124001.

     I believe including the discussion of more related work could make the proposed method's positioning more comprehensive.
soundness: 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.
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.
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.

Meta-review: The reviewers recommend acceptance (5-6-6-8-6-6). The paper proposes Latent Ensemble Score Filter for improving data assimilation in high-dimensional and nonlinear systems. The approach is sound and the results are good. The author-reviewer discussion has been very constructive and has led to a number of clarifications and improvements that have convinced the reviewers. Given the overall positive reviews, I recommend acceptance. I request the authors to implement the many suggestions made during the author-reviewer discussion period in the final version of the paper.

Decision: Accept (Poster)