LO-SDA: Latent Optimization for Score-based Atmospheric Data Assimilation

## Abstract

Abstract Data assimilation (DA) plays a pivotal role in numerical weather prediction by systematically integrating sparse observations with model forecasts to estimate optimal atmospheric initial condition for forthcoming forecasts.
Traditional Bayesian DA methods adopt a Gaussian background prior as a practical compromise for the curse of dimensionality in atmospheric systems, that simplifies the nonlinear nature of atmospheric dynamics and can result in biased estimates.
To address this limitation, we propose a novel generative DA method, LO-SDA. First, a variational autoencoder is trained to learn compact latent representations that disentangle complex atmospheric correlations. Within this latent space, a background-conditioned diffusion model is employed to directly learn the conditional distribution from data, thereby generalizing and removing assumptions in the Gaussian prior in traditional DA methods. Most importantly, we introduce latent optimization during the reverse process of the diffusion model to ensure strict consistency between the generated states and sparse observations.
Idealized experiments demonstrate that LO-SDA not only outperforms score-based DA methods based on diffusion posterior sampling but also surpasses traditional DA approaches.
To our knowledge, this is the first time that a diffusion-based DA method demonstrates the potential to outperform traditional approaches on high-dimensional global atmospheric systems.
These findings suggest that long-standing reliance on Gaussian priors—a foundational assumption in operational atmospheric DA—may no longer be necessary in light of advances in generative modeling.

## 1 Introduction

In numerical weather prediction, data assimilation (DA) is essential for generating accurate initial conditions that directly determine forecast skill lorenc1986analysis; gustafsson2018survey; asch2016varda. Modern DA methods estimate the optimal atmospheric state $\boldsymbol{x}$ within a Bayesian framework by combining sparse observations $\boldsymbol{y}$ with model forecasts $\boldsymbol{x}_{b}$ (also known as background fields) asch2016varda; rabier2003varda; carrassi2018varda; le1986variational.
Specifically, DA aims to estimate the Bayesian posterior distribution $p(\boldsymbol{x}|\boldsymbol{x}_{b},\boldsymbol{y})$.
Given that forecasts and observations are typically conditionally independent, the posterior simplifies to $p(\boldsymbol{x}|\boldsymbol{x}_{b},\boldsymbol{y})\propto p(\boldsymbol{y}|\boldsymbol{x})p(\boldsymbol{x}|\boldsymbol{x}_{b})$.

Traditional DA methods typically assume both the prior $p(\boldsymbol{x}\mid\boldsymbol{x}_{b})$ and the likelihood $p(\boldsymbol{y}|\boldsymbol{x})$ follow Gaussian distributions to simplify the inference process bannister2017review.
While this assumption is relatively reasonable for observation errors, it breaks down for background uncertainty, which often becomes non-Gaussian after undergoing nonlinear model evolution.
Furthermore, this assumption makes traditional DA methods rely on the background error covariance matrix $\mathbf{B}$ to define the solution space for assimilation bannister2008review1. Nevertheless, $\mathbf{B}$ often spans more than $10^{12}$ degrees of freedom in high-resolution systems, making it extremely challenging to estimate and potentially introducing significant additional error into the assimilation process bannister2017review; bannister2008review2. These limitations have spurred generative DA frameworks.

Generative DA models perform posterior inference using score functions, offering a promising alternative to traditional approaches by relaxing the need for Gaussian assumptions rozet2023sda; rozet2023mlpsda; qu2024qusda; huang2024diffda; manshausen2025sda. However, existing approaches face notable limitations, both in practical implementation and theoretical understanding.
For instance, Huang et al. (DiffDA) huang2024diffda condition the diffusion model on the background and incorporate observations through a repainting strategy, but their method underperforms in sparse observation settings and cannot effectively handle nonlinear observation operators such as satellite radiative transfer.
Qu et al. qu2024qusda encode background and multi-modal observations into a unified guidance signal, though their reliance on specific observation distribution assumptions restricts generalization to complex DA scenarios.
Moreover, Rozet et al. rozet2023sda; rozet2023mlpsda and Manshausen et al. manshausen2025sda treat observations as guidance during the reverse process, ignoring the background prior. While this works well when observations are dense and clean, but often fails under sparse or noisy conditions, where background information becomes essential.
These limitations underscore the necessity of a unified framework that jointly leverages both background information and observational guidance in generative DA.

To this end, we propose the Latent Optimization Score-based Data Assimilation (LO-SDA) framework, which seeks to offer a more principled and reliable formulation of generative DA.
First, we train a variational autoencoder (VAE) to learn a compact latent representation of the high-dimensional atmospheric states, capturing nonlinear dependencies among variables and enabling more efficient probabilistic modeling.
Second, we train a score-based model to model the background conditioned prior in latent space $p(\boldsymbol{z}|\boldsymbol{z}_{b})$, where $\boldsymbol{z}$ represents the latent representation of model state $\boldsymbol{x}$.
Third, we introduce an alternating latent optimization scheme song2024LO that iteratively enforces observational constraints during guided diffusion sampling.
In our framework, the diffusion-estimated prior yields a less biased analysis compared to traditional approaches (Figure [1](https://arxiv.org/html/2510.22562v1#S2.F1) (a)), while latent optimization significantly enhances analysis-observation consistency (Figure [1](https://arxiv.org/html/2510.22562v1#S2.F1) (b)).
Our framework behaves similarly to multiple posterior likelihood maximization during guidance sampling, distinguishing it from single-step gradient descent in DPS. Meanwhile, our framework has the shared optimization-based strategy employed by variational DA. This mechanism offers a plausible explanation for our framework’s superior performance.

Our contributions are outlined as follows:

- •
We identify a theoretical connection between latent optimization and variational DA methods. Building upon this theoretical analogy, we develop a novel framework that integrates observational information into the background conditional prior through latent optimization techniques. Our alternating latent optimization scheme provably achieves multiple maximizations of the posterior likelihood $p(\bm{y}|\bm{z})$ during guided sampling, guaranteeing progressive refinement.
- •
To the best of our knowledge, this is the first work to demonstrate that removing Gaussian assumptions via diffusion enables score-based DA to outperform traditional methods in high-dimensional global atmospheric settings.
- •
By incorporating latent optimization into score-based DA, we iteratively enforce observational constraints during sampling, resulting in more accurate and observation-consistent analyses than those produced by existing approaches.

## 2 Related work

The variational assimilation methods.
Variational assimilation is a representative class of traditional DA methods and is widely used in operational numerical weather prediction systems.
In the three-dimensional case, it seeks to maximize the posterior likelihood asch2016varda; rabier2003varda; carrassi2018varda:

$$ $\displaystyle\boldsymbol{x}_{a}=\operatorname*{arg\,max}_{\boldsymbol{x}}p(\boldsymbol{x}|\boldsymbol{x}_{b},\boldsymbol{y})=\operatorname*{arg\,max}_{\boldsymbol{x}}p(\boldsymbol{x}|\boldsymbol{x}_{b})p(\boldsymbol{y}|\boldsymbol{x}),$ (1) $$

where the assumption of independence between observation errors and background errors is applied.
By assuming that the prior distribution $p(\boldsymbol{x}|\boldsymbol{x}_{b})$ and the observation likelihood $p(\boldsymbol{y}|\boldsymbol{x})$ follow Gaussian distributions, three-dimensional variational DA (3DVar) is equivalent to minimizing the following cost function:

$$ $\displaystyle J(\boldsymbol{x})=\frac{1}{2}(\boldsymbol{x}-\boldsymbol{x}_{b})^{T}\textbf{B}^{-1}(\boldsymbol{x}-\boldsymbol{x}_{b})+\frac{1}{2}(\boldsymbol{y}-\mathcal{H}(\boldsymbol{x}))^{T}\textbf{R}^{-1}(\boldsymbol{y}-\mathcal{H}(\boldsymbol{x}))\ .$ (2) $$

where $\mathbf{B}$ and $\mathbf{R}$ denote the covariance matrices of background and observation errors, respectively, and $\mathcal{H}$ is the observation operator that maps model states to observation space.
As noted by Bannister bannister2008review1; bannister2017review, $\mathbf{B}$ plays a central role in variational DA by shaping the feasible solution space and promoting physical consistency in the resulting analysis.
In practice, the high-dimensional $\mathbf{B}$ is commonly simplified via a control variable transformation that approximately diagonalizes it descombesGeneralizedBackgroundError2015. Although this facilitates its inversion, the simplified $\mathbf{B}$ may fail to capture the evolving physical consistency of atmospheric states, leading to suboptimal assimilation outcomes.

Figure: Figure 1: Comparison between LOSDA and other DA approaches. (a) Prior estimation: The true background conditional prior $p(\boldsymbol{x}|\boldsymbol{x}_{b})$ (blue dashed) is approximated as Gaussian in traditional DA (green), while LOSDA directly estimates it through diffusion modeling (red). By incorporating observation likelihood $p(\boldsymbol{y}|\boldsymbol{x})$, LOSDA achieves posterior estimation $p(\boldsymbol{x}|\boldsymbol{x}_{b},\boldsymbol{y})$ closer to the ground truth. (b) Observation integration methods: Top - Diffusion Posterior Sampling (DPS) updates denoised $\boldsymbol{x}_{t}$ via observation error gradient guidance (single-step consistency). Bottom - LOSDA’s optimization approach directly minimizes observation error for optimal denoised $\boldsymbol{x}_{t}$ (strict multi-step consistency). Our framework enforces tighter observation constraints than gradient-based DPS.
Refer to caption: Figures/dps_laop.png

The latent assimilation methods.
Recently, latent data assimilation (LDA) cheng2024multilda; melinc20243dlda; peyron2021latentlda; amendola2021datalda; zhengGeneratingUnseenNonlinear; fan2025physicallylda; fanNovelLatentSpace2025; fanNovelLatentSpace2025a has been proposed to apply traditional DA methods with Gaussian priors in a compact latent space learned via autoencoders.
For example, the latent formulation of the widely used 3DVar—referred to as L3DVar—optimizes the following loss function:

$$ $\displaystyle J(\boldsymbol{z})=\frac{1}{2}(\boldsymbol{z}-\boldsymbol{z}_{b})^{T}\textbf{B}_{z}^{-1}(\boldsymbol{z}-\boldsymbol{z}_{b})+\frac{1}{2}(\boldsymbol{y}-\mathcal{H}(D(\boldsymbol{z}))^{T}\textbf{R}^{-1}(\boldsymbol{y}-\mathcal{H}(D(\boldsymbol{z}))).$ (3) $$

where $\boldsymbol{z}$ and $\mathbf{B}_{z}$ denote the latent state and the background error covariance matrix in the latent space, respectively. Several studies have found that $\mathbf{B}_{z}$ is inherently near-diagonal, as the latent space effectively captures correlations among atmospheric variables. Consequently, LDA can adopt a diagonal $\mathbf{B}_{z}$, greatly simplifying its implementation melinc20243dlda; zhengGeneratingUnseenNonlinear; fan2025physicallylda. Fan et al. fan2025physicallylda further showed that performing variational assimilation in latent space can outperform its model-space counterpart. Nevertheless, most latent DA methods remains constrained by the Gaussian prior assumption.
To overcome this limitation, our work also leverages latent representations of high-dimensional atmospheric states, but replaces the Gaussian prior with a more expressive, data-driven distribution modeled by a latent score-based model.

## 3 Method

### 3.1 Preliminary

Score-based model.
The score-based model, a promising class of the generative models BeatGANdhariwal2021diffusion; BeatGANho2022cascaded, offering high-quality generation and excellent model convergence Conversong2021maximum; Converhuang2021variational; Converkingma2021variational. It comprises a forward process and a reverse process Diffho2020denoising; Diffsohl2015deep; Diffsong2020score. In the forward process, the original data distribution is transformed into a known prior, by gradually injecting noise. Such a process is governed by a stochastic differential equation (SDE) and a corresponding reverse-time SDE Diffsong2020score,

$$ $\displaystyle d\boldsymbol{x}$ $\displaystyle={\bf f}(\boldsymbol{x},t)dt+g(t)d\boldsymbol{w}$ (4) $\displaystyle d\boldsymbol{x}$ $\displaystyle=[{\bf f}(\boldsymbol{x},t)-g(t)^{2}\nabla_{\boldsymbol{x}}\log p_{t}(\boldsymbol{x})]dt+g(t)d\bar{\boldsymbol{w}},$ (5) $$

where the reverse SDE transforms the prior distribution back into the data distribution by gradually removing the noise. Here, $\boldsymbol{w}$ and $\bar{\boldsymbol{w}}$ both represent the standard Wiener processes (Gaussian white noise), with ${\bf f}(\boldsymbol{x},t)$ the drift coefficient and $g(t)$ the diffusion coefficient of $\boldsymbol{x}(t)$. Accordingly, the perturbation kernel from $\boldsymbol{x}_{0}$ to $\boldsymbol{x}_{t}$ takes form $p(\boldsymbol{x}_{t}|\boldsymbol{x})\sim\mathcal{N}(\mu(t),\sigma^{2}(t)\boldsymbol{I})$, where $\mu(t),\sigma^{2}(t)$ can be determined by the ${\bf f}(\boldsymbol{x},t)$ and $g(t)$. In this work, we take the widely used variance-preserving SDE and the cosine schedule for $\mu(t)$ rozet2023sda. In the generative diffusion model, the score function $\nabla_{\boldsymbol{x}}\log p_{t}(\boldsymbol{x})$ can be estimated by a neural network with parameter $\boldsymbol{\theta}$ via minimizing the denoising score matching loss $\mathcal{L}_{t}\equiv\mathbb{E}_{p(\boldsymbol{x}_{t})}||\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x},t)-\nabla_{\boldsymbol{x}}\log p_{t}(\boldsymbol{x}|\boldsymbol{x}_{0})||^{2}$, which theoretically guarantees $\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x},t)\approx\nabla_{\boldsymbol{x}}\log p_{t}(\boldsymbol{x})$ Diffsong2020score. Once we have a trained $s_{\theta}(\boldsymbol{x},t)$, the trajectory from the prior distribution to the real data distribution can be determined following Equation [5](https://arxiv.org/html/2510.22562v1#S3.E5).

Score-based data assimilation.
Data assimilation under this framework reformulates the Bayesian posterior as a composite scoring process:

$$ $\displaystyle\nabla_{\boldsymbol{x}_{t}}\log p(\bm{x}_{t}|\bm{x}_{b},\bm{y})={\nabla_{\boldsymbol{x}_{t}}\log p(\bm{x}|\bm{x}_{b})}+{\nabla_{\boldsymbol{x}_{\tau}}\log p(\bm{y}|\bm{x}_{t})}$ (6) $$

The prior term leverages the diffusion model’s capacity to capture complex spatial correlations, bypassing the oversimplified Gaussian assumptions in the conventional DA methods. The constraint term enforces observation consistency through conditional guidance. In the DPS paradigm chung2023diffusionDPS, the observation term is supposed to follow Gaussian distribution,

$$ $\displaystyle p(\bm{y}|\bm{x}_{t})\sim\mathcal{N}\left(\mathcal{H}(\tilde{\boldsymbol{x}}_{0}(\boldsymbol{x}_{t})),\boldsymbol{R}\right)$ (7) $$

where the posterior mean $\tilde{\boldsymbol{x}}_{0}$ derives from Tweedie’s formula Diffho2020denoising; Diffsohl2015deep:

$$ $\displaystyle\tilde{\boldsymbol{x}}_{0}(\boldsymbol{x}_{t})=\frac{\boldsymbol{x}_{t}+\sigma(t)^{2}\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t},\boldsymbol{x}_{b})}{\mu(t)}.$ (8) $$

Figure: Algorithm 1 Comparison of DPS Guidance and Latent Optimization for Score-Based Data Assimilation

### 3.2 Latent score-based data assimilation

Due to the computational challenges in high-dimensional systems, LDA cheng2024multilda; melinc20243dlda; peyron2021latentlda; amendola2021datalda; fan2025physicallylda is proposed to leverage VAEs for compressing physical fields into low-dimensional manifolds kingma2013autovae; doersch2016tutorialvae and performing efficient optimization in this reduced space.
Specifically, the traditional 3DVar formulation (Equation [1](https://arxiv.org/html/2510.22562v1#S2.E1)) is adapted to the latent space with cost function described by Equation [3](https://arxiv.org/html/2510.22562v1#S2.E3).
The gradient descent iteratively optimizes the latent representation of the analysis field.
The optimized latent is then decoded to reconstruct the assimilated state.
Although LDA alleviates non-linearity challenges, its retention of Gaussian assumptions for latent background distributions imposes theoretical limitations as above-discussed, particularly in capturing the multiscale complexity characteristic of real atmospheric states fan2025physicallylda.
In this work, we train a score-based model in the latent space to model the background conditional distribution.
Additionally, we integrate observations through guidance sampling in the latent space. Mathematically, the latent score modeling for DA can be expressed as:

$$ $\displaystyle\nabla_{\boldsymbol{z}_{t}}\log p(\bm{z}_{t}|\bm{z}_{b},\bm{y})$ $\displaystyle=\nabla_{\boldsymbol{z}_{t}}\log p(\bm{z}|\bm{z}_{b})+\nabla_{\boldsymbol{z}_{t}}\log p(\bm{y}|\bm{z}_{t})$ $\displaystyle=\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{z}_{t},\boldsymbol{z}_{b})+\nabla_{\boldsymbol{z}_{t}}\log p(\bm{y}|\bm{z}_{t}).$ (11) $$

For the guidance term, we implemented the latent counterpart of DPS guidance where the observation term preserves Gaussian distributions $p(\bm{y}|\bm{z}_{t})\sim\mathcal{N}(\mathcal{H}(D(\hat{\boldsymbol{z}}_{0}(\boldsymbol{z}_{t})),\boldsymbol{R})$.
Similar to Equation [8](https://arxiv.org/html/2510.22562v1#S3.E8), $\hat{\boldsymbol{z}}_{0}$ is the posterior mean. $D(\cdot)$ denotes the decoder of VAE. While Algorithm [1](https://arxiv.org/html/2510.22562v1#alg1) outlines the sampling process using DPS guidance, its single-step gradient update mechanism may provide insufficient constraint enforcement, potentially compromising observation consistency in high-dimensional scenarios.

### 3.3 Latent optimization techniques

To perform strict observation consistency, we aim to integrate variational optimization (Equation [2](https://arxiv.org/html/2510.22562v1#S2.E2) and Equation [3](https://arxiv.org/html/2510.22562v1#S2.E3)) used in traditional DA. Inspired by inverse problem solving techniques song2024LO within diffusion models, a two-stage latent optimization strategy is proposed: Hard-Constrained Optimization: (1) Solving $\hat{\boldsymbol{z}}_{0}(\boldsymbol{y})=\operatorname*{arg\,min}_{\bm{z}}(\bm{y}-\mathcal{H}(D(\bm{z})))^{T}{\boldsymbol{R}^{-1}}(\bm{y}-\mathcal{H}(D(\bm{z})))$ to ensure strict observation consistency, (2) Projecting the optimized latent back to the noisy data manifold using the reverse process. Since the $\hat{\boldsymbol{z}}_{0}(\boldsymbol{y})=(\boldsymbol{z}_{t}-\sigma(t)\varepsilon)/\mu(t)$ can be viewed as the estimated mean of latent $\boldsymbol{z}_{t}$ with the observation single $\boldsymbol{y}$, one have that

$$ $\displaystyle p(\boldsymbol{z}_{t}|\hat{\boldsymbol{z}}_{0}(\boldsymbol{y}),\boldsymbol{y})\sim\mathcal{N}\left(\mu(t)\hat{\boldsymbol{z}}_{0}(\boldsymbol{y}),\sigma^{2}(t)\boldsymbol{I}\right),$ (12) $$

from the forward process.
When we map the $\hat{\boldsymbol{z}}_{0}(\boldsymbol{y})$ back to noise data manifold, we need the distributions $p({\bm{z}}_{t}|\tilde{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$. By Bayesian formula,
$p({\bm{z}}_{t}|\tilde{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})\propto p(\tilde{\bm{z}}_{t}|{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})p(\boldsymbol{z}_{t}|\hat{\boldsymbol{z}}_{0}(\boldsymbol{y}),\boldsymbol{y})$. The posterior distribution $p(\tilde{\bm{z}}_{t}|{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ is assumed as Gaussian distribution with variance $\lambda_{t}^{2}$ and the $p(\boldsymbol{z}_{t}|\hat{\boldsymbol{z}}_{0}(\boldsymbol{y}),\boldsymbol{y})$ is supposed to provide the prior of its mean.
Thus, it is accordingly derived (see Appendix):

$$ $\displaystyle p({\bm{z}}_{t}|\tilde{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})=\mathcal{N}\left(\frac{\lambda_{t}^{2}\mu(t)\hat{\bm{z}}_{0}(\bm{y})+\sigma^{2}(t)\tilde{\bm{z}}_{t}}{\lambda_{t}^{2}+\sigma^{2}(t)},\frac{\lambda_{t}^{2}\sigma^{2}(t)}{\lambda_{t}^{2}+\sigma^{2}(t)}\bm{I}\right).$ (13) $$

Following song2024LO, we choose the variance $\lambda_{t}^{2}$ schedule as $\lambda_{t}^{2}=\lambda\left(\frac{1-\mu^{2}(t-\Delta t)}{\mu^{2}(t)}\right)\left(1-\frac{\mu^{2}(t)}{\mu^{2}(t-\Delta t)}\right)$ with a hyperparameter $\lambda$. This approach integrates variational optimization within the diffusion sampling framework (Algorithm [1](https://arxiv.org/html/2510.22562v1#alg1)), where latent variables are iteratively refined at multiple diffusion steps, preserving highly observation consistency.

We identify a potential theoretical connection between the latent optimization technique and the variational DA methods. In variational DA, the analysis field is obtained by minimizing a cost function (Equation [3](https://arxiv.org/html/2510.22562v1#S2.E3)) that balances the background term and observation posterior likelihood. Analogously, as demonstrated in Algorithm [1](https://arxiv.org/html/2510.22562v1#alg1), each latent optimization step during reverse sampling updates the latent state through strict observation consistency (corresponding to the observation term in Equation [3](https://arxiv.org/html/2510.22562v1#S2.E3)). The optimized latent is subsequently fed into the background conditional diffusion kernel
$p(\boldsymbol{z}_{t}|\boldsymbol{z}_{t+1})$ to correct the diffusion trajectory.
This optimization-based similarity not only motivates us to apply latent optimization to integrate observational information into the background conditional prior but also provides a plausible explanation of our framework’s outstanding performance.

Figure: Figure 2: Comparative visualization of t850 analysis fields across assimilation methods under 1% idealized observation (valid at 2019-01-03 00:00 UTC). Top row (left to right): ERA5 ground truth, background field, and background error. Middle row: Assimilation results from (a) proposed LO-SDA method, (b) DPS framework, and (c) Repaint approach. Bottom row: Corresponding absolute error fields relative to ERA5 truth. The reduced error magnitude (lighter hues) in LO-SDA results demonstrates our method’s superior error reduction capability compared to alternative approaches.
Refer to caption: Figures/vis_t850.png

## 4 Experiments

### 4.1 Experimental Settings and Evaluations

Dataset and metrics. We conduct our experiments on the ERA5 reanalysis dataset hersbach2020era5, a global atmospheric data product maintained by the European Centre for Medium-Range Weather Forecasts (ECMWF). Our study utilizes 5 upper-air atmospheric variables (geopotential, temperature, specific humidity, zonal wind, and meridional wind) across 13 pressure levels (50hPa, 100hPa, 150hPa, 200hPa, 250hPa, 300hPa, 400hPa, 500hPa, 600hPa, 700hPa, 850hPa, 925hPa, and 1000hPa), combined with 4 surface variables (10-meter zonal component of wind (u10), 10-meter meridional component of wind (v10), 2-meter temperature (msl) and mean sea level pressure (msl)), forming a total of 69 meteorological variables. The pressure-level variables follow the standardized ERA5 naming convention (e.g., t850 denotes temperature at 850 hPa). We use a subset spanning 1979-2018 for training and evaluations.
For evaluations, the assimilation quality is assessed by direct comparison with ERA5 reference fields. Three metrics quantifying performance are overall mean square error (MSE), mean absolute error (MAE), and the latitude-weighted root mean square error (WRMSE) (see Appendix), which is a statistical metric widely used in geospatial analysis and atmospheric science rasp2020weatherbench; rasp2024weatherbench2. The validation procedure conducts assimilation cycles at 00:00 UTC for each day throughout 2019. For each test case, we calculate the above three metrics. Final performance scores represent the annual average of these daily metrics, ensuring statistically significant results across all seasons and synoptic conditions.

Experimental setting. The Fengwu AI forecasting model chen2023fengwuAIfore (6-hour temporal resolution) is integrated into our DA framework to produce the background field. These fields are generated through an 8-step autoregressive forecasting procedure, initialized with ERA5 conditions from 48 hours prior to the target assimilation lead time. To simulate realistic observing system characteristics, we create synthetic observations by randomly masking the ERA5 truth data at two sparsity levels (95% and 99%), mimicking typical satellite coverage constraints.
The 1.40625° ($128\times 256$ grid) spatial resolution is employed, yielding input arrays of size $69\times 128\times 256$.

**Table 1: Quantitative performance comparison of different methods under 1% and 5% observations.**
| Ratio | Model | MSE | MAE | WRMSE |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| msl | u10 | u700 | v500 | z500 | t850 |  |  |  |  |
| 48h background | 0.0505 | 0.1178 | 98.7265 | 1.2727 | 1.9953 | 2.4217 | 89.2752 | 0.9310 |  |
| 1% observation | 3DVAR | 0.0483 | 0.1138 | 81.6384 | 1.2235 | 1.9850 | 2.4298 | 62.9377 | 0.8797 |
| L3DVAR | 0.0474 | 0.1105 | 62.1054 | 1.1862 | 1.9797 | 2.3392 | 53.0902 | 0.8975 |  |
| Repaint | 0.0592 | 0.1311 | 114.3672 | 1.4167 | 2.1059 | 2.5664 | 104.1351 | 1.0363 |  |
| DPS | 0.0545 | 0.1247 | 95.9850 | 1.3286 | 2.0220 | 2.3981 | 85.5673 | 1.0269 |  |
| LO-SDA(ours) | 0.0472 | 0.1101 | 62.4505 | 1.1836 | 1.8981 | 2.2439 | 53.6468 | 0.9243 |  |
| 5% observation | 3DVAR | 0.0430 | 0.0982 | 63.2562 | 1.1661 | 1.8765 | 2.2365 | 45.8574 | 0.7849 |
| L3DVAR | 0.0315 | 0.0903 | 45.8037 | 0.9350 | 1.7166 | 1.9400 | 38.6411 | 0.8024 |  |
| Repaint | 0.0496 | 0.1219 | 106.3938 | 1.2934 | 1.9889 | 2.3622 | 95.9167 | 0.9856 |  |
| DPS | 0.0486 | 0.1199 | 93.247 | 1.2673 | 1.9585 | 2.3271 | 85.3329 | 0.9891 |  |
| LO-SDA(ours) | 0.0309 | 0.0851 | 42.3498 | 0.8992 | 1.5873 | 1.7894 | 32.8990 | 0.8094 |  |

The background conditional diffusion model.
We present a unified framework for conditional physics field modeling through variational autoencoding and latent diffusion. Our architecture begins with a window-attention transformer VAE han2024cra5 that compresses high-dimensional fields ($69\times 128\times 256$) to compact latent representations ($69\times 32\times 64$). Trained for 80 epochs using AdamW loshchilov2017fixingadam with batch size 32, the VAE employs a hybrid learning rate schedule: linear warmup to $2\times 10^{-4}$ over 10,000 iterations followed by cosine decay, achieving 0.0067 reconstruction MSE as detailed in Appendix. The latent diffusion process then learns conditional distributions $p(\boldsymbol{z}|\boldsymbol{z}_{b})$ through a 28-layer transformer backbone peebles2023scalabledit with 1152-dimensional hidden states, (2,2) patch embedding, and 16-head cross-attention for background latent $\boldsymbol{z}_{b}$ conditioning. For diffusion setting, the variance-preserving SDE Diffsong2020score with cosine noise scheduling. Optimized via AdamW loshchilov2017fixingadam at constant $1\times 10^{-4}$ learning rate (batch size 32), the model converges stably over 100k training steps. Sampling employs a modified Predictor-Corrector scheme combining 128-step prediction with 2 iterations of Langevin correction Diffsong2020score. See the Appendix for more comprehensive resource usage.

Baselines. For diffusion-based experiments, we incorporate observations through two baseline methods: a latent-space implementation of the repainting technique from DiffDA huang2024diffda and the latent version of DPS described in Algorithm [1](https://arxiv.org/html/2510.22562v1#alg1). The latent repaint implementation follows:

$$ $\displaystyle\boldsymbol{z}^{obs}_{t}\sim\mathcal{N}(\mu(t)$ $\displaystyle E(\boldsymbol{x}^{*}),\sigma^{2}(t)\boldsymbol{I}),\quad\tilde{\boldsymbol{z}}^{t}\leftarrow\text{SolutionAtTime}(t)$ (14) $\displaystyle\boldsymbol{z}_{t-1}$ $\displaystyle=E\left(m\odot D(\boldsymbol{z}^{obs}_{t})+(1-m)\odot D(\tilde{\boldsymbol{z}}^{t})\right)$ (15) $$

where $z^{obs}_{t}$ is noised latent with $E(\boldsymbol{x}^{*})$ is the encoded ERA5 ground truth latent. $\boldsymbol{z}_{t}$ is the sampled prior latent at diffusion time $t$. We finally combine the decoded observed and prior latents in model space using a masking matrix $m$, with $\odot$ denoting element-wise multiplication.

For the DPS assimilation, we apply guidance during the final third of the reverse sampling process with a scale factor of 800. Our LO-SDA implementation employs a hyperparameter setting of $\lambda=100$ to balance observation constraints and prior knowledge. In non-diffusion experiments, we implement both conventional 3DVar and its latent-space counterpart (L3DVar) via optimizing cost function in Equation [2](https://arxiv.org/html/2510.22562v1#S2.E2) and Equation [3](https://arxiv.org/html/2510.22562v1#S2.E3), respectively. This dual approach allows us to systematically evaluate the performance improvements offered by latent-space assimilation techniques across different methodological frameworks.

### 4.2 Results and Discussions

Table [1](https://arxiv.org/html/2510.22562v1#S4.T1) provides a quantitative evaluation of analysis errors by comparing LO-SDA with the baseline methods.
The error of the background field is presented in the first row.
Among diffusion-based approaches (Repaint, DPS, and LO-SDA), LO-SDA with latent optimization demonstrates superior performance over both repainting and DPS guidance across comprehensive metrics (MSE, MAE, and WRMSE), particularly for most prognostic variables.
For example, LO-SDA achieves 13.39% and 20.27% improvement over DPS and Repaint in overall MSE metrics, aligning with our theoretical expectation in Algorithm [1](https://arxiv.org/html/2510.22562v1#alg1).
These results confirm that latent optimization enforces rigorous data consistency.
When compared to traditional 3DVAR data assimilation, LO-SDA exhibits significantly improved accuracy.
Notably, LO-SDA performs comparably to the ML-enhanced variational method, L3DVAR.
Moreover, LO-SDA demonstrates superior scalability with increased observational density - at 5% observation coverage, it surpasses L3DVAR’s performance, highlighting the effectiveness of latent optimization in leveraging observational constraints.

To further elucidate these findings,
we quantify the percentage improvement relative to 3DVAR and L3DVAR. Under 1% observation conditions, LO-SDA achieves a 14.76% improvement over 3DVAR in z500 WRMSE, while maintaining comparable performance ($-1.05\%$) with L3DVAR. This performance gap substantially widens at 5% observation density, where LO-SDA delivers 28.25% and 14.86% improvements over 3DVAR and L3DVAR, respectively, for z500 WRMSE.
These results underscore two critical findings: (1) The latent optimization framework effectively assimilates observational information, with its advantage becoming more pronounced as observational density increases; (2) LO-SDA exhibits strong scalability, suggesting its suitability for operational implementation with realistic observational networks.

**Table 2: Quantitative performance comparison under a 1% observation setting, with varying observation errors modeled as Gaussian noise. The standard deviations are set to 0.02, 0.05, and 0.10 relative to the ERA5 climatological standard deviation.**
| Ratio | Model | MSE | MAE | WRMSE |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| msl | u10 | u700 | v500 | z500 | t850 |  |  |  |  |
| 48h background | 0.0505 | 0.1178 | 98.7265 | 1.2727 | 1.9953 | 2.4217 | 89.2752 | 0.9310 |  |
| std = 0.02 | 3DVAR | 0.0484 | 0.1141 | 82.4252 | 1.2243 | 1.9805 | 2.4202 | 67.3522 | 0.8679 |
| L3DVAR | 0.0475 | 0.1109 | 63.0360 | 1.1935 | 1.9900 | 2.3496 | 53.6535 | 0.9054 |  |
| LO-SDA(ours) | 0.0470 | 0.1109 | 64.7142 | 1.1940 | 2.0618 | 2.2888 | 55.3858 | 0.9354 |  |
| std = 0.05 | 3DVAR | 0.0485 | 0.1158 | 85.4325 | 1.2246 | 1.9643 | 2.3893 | 77.8867 | 0.8845 |
| L3DVAR | 0.0481 | 0.1132 | 72.2672 | 1.1852 | 1.9515 | 2.3199 | 63.6825 | 0.9143 |  |
| LO-SDA(ours) | 0.0479 | 0.1127 | 68.0784 | 1.1989 | 1.9104 | 2.2641 | 58.1329 | 0.9542 |  |
| std = 0.10 | 3DVAR | 0.0495 | 0.1173 | 91.2331 | 1.2291 | 1.9493 | 2.3636 | 84.4765 | 0.9101 |
| L3DVAR | 0.0489 | 0.1147 | 83.1325 | 1.1954 | 1.9448 | 2.3275 | 75.3260 | 0.9341 |  |
| LO-SDA(ours) | 0.0498 | 0.1188 | 82.0772 | 1.2408 | 1.9602 | 2.3194 | 69.4214 | 1.0308 |  |

The success of LO-SDA demonstrates that generative modeling can enable atmospheric DA transcend traditional Gaussian assumptions.
By learning non-parametric conditional distributions $p(\boldsymbol{x}|\boldsymbol{x}_{b})$, our framework achieves more accurate analysis fields than both variational methods and recent diffusion-based DA approaches.
This advancement stems from two key innovations: (1) a background-conditioned diffusion model that replaces restrictive Gaussian priors, and (2) latent optimization during the reverse process to enforce hard observation constraints - a novel hybrid approach that marries the flexibility of generative modeling with the observation consistency requirements of operational DA.

### 4.3 Real-world observations

To evaluate our framework under real-world conditions, we employ the Global Data Assimilation System (GDAS) prepbufr dataset, which incorporates multi-source observations.
For this study, only surface and radiosonde observations are utilized. These observations are first interpolated onto the model state grid, and any multiple observations at a single grid point are averaged. High-elevation surface observations are vertically interpolated and reclassified as upper-air data.
A quality control procedure is further applied to remove observations with large deviations. Observations are dropped if their deviation from the ERA5 reference exceeds 0.05 of the ERA5 climatological standard deviation. We perform data assimilation daily at 00:00 UTC throughout 2017, using a 48-hour background field. As shown in Table [3](https://arxiv.org/html/2510.22562v1#S4.T3), the results indicate that LO-SDA achieves performance comparable to L3DVAR, and slightly outperforms the traditional 3DVAR when using real observations.

**Table 3: Assimilation performance on real-world observation across various methods.**
|  | MSE | MAE | WRMSE |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | msl | u10 | u700 | v500 | z500 | t850 |  |  |
| 48h background | 0.0475 | 0.1158 | 98.7910 | 1.2588 | 1.9692 | 2.4061 | 89.2800 | 0.9232 |
| 3DVAR | 0.0472 | 0.1150 | 87.1536 | 1.2532 | 1.9646 | 2.3950 | 83.8241 | 0.9068 |
| L3DVAR | 0.0467 | 0.1143 | 84.8751 | 1.2376 | 1.9597 | 2.3778 | 78.1842 | 0.8962 |
| LO-SDA(ours) | 0.0469 | 0.1140 | 81.4013 | 1.2128 | 1.9891 | 2.3657 | 77.1881 | 0.9914 |

### 4.4 Ablation studies

Observation Error Robustness.
To evaluate the robustness of LO-SDA under realistic observational conditions, we conduct experiments with simulated observation errors by injecting additive Gaussian noise with standard deviations of 2%, 5%, and 10% of the ERA5 climatological standard deviation.
As evidenced by the quantitative results in Table [2](https://arxiv.org/html/2510.22562v1#S4.T2), our framework maintains consistent performance across various noise levels, demonstrating remarkable error tolerance.

**Table 4: Comparison of latent optimization frequencies (skip=2, 4, 8) in reverse diffusion sampling under sparse observation settings (1% and 5%).**
| Ratio | Frequency | MSE | MAE | WRMSE |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| msl | u10 | u700 | v500 | z500 | t850 |  |  |  |  |
| 48h background | 0.0505 | 0.1178 | 98.7265 | 1.2727 | 1.9953 | 2.4217 | 89.2752 | 0.9310 |  |
| 1% observation | skip=2 | 0.0472 | 0.1101 | 62.4505 | 1.1836 | 1.8981 | 2.2439 | 53.6468 | 0.9243 |
| skip=4 | 0.0518 | 0.1162 | 70.2005 | 1.2648 | 1.9677 | 2.33571 | 59.6587 | 0.9611 |  |
| skip=8 | 0.0549 | 0.1205 | 74.2434 | 1.3258 | 2.0143 | 2.4029 | 62.7748 | 0.9840 |  |
| 5% observation | skip=2 | 0.0309 | 0.0851 | 42.3498 | 0.8992 | 1.5873 | 1.7894 | 32.8990 | 0.8094 |
| skip=4 | 0.0370 | 0.0939 | 49.5966 | 0.9892 | 1.6697 | 1.8892 | 41.5379 | 0.8538 |  |
| skip=8 | 0.0415 | 0.1001 | 54.9766 | 1.0740 | 1.7544 | 2.0039 | 45.8147 | 0.8890 |  |

Latent Optimization Frequency Analysis.
To empirically validate the theoretical connection (see Section [3.3](https://arxiv.org/html/2510.22562v1#S3.SS3)) and demonstrate its impact on our framework’s effectiveness, we conduct an ablation study on the latent optimization frequency by adjusting the skip interval parameter.
Table [4](https://arxiv.org/html/2510.22562v1#S4.T4) reveals a systematic performance degradation as the skip interval increases (i.e., fewer optimization steps).
Specifically, under sparse 5% observation conditions, the overall MSE increases by 19.74% and 34.30% for skip intervals of 4 and 8, respectively, compared to the baseline configuration with skip=2.
This degradation aligns with our theoretical insight: frequent latent optimization ensures proper integration of observations into the diffusion process, analogous to how iterative refinement in variational DA minimizes the analysis cost function.

## 5 Conclusion and Discussion

Conclusion.
The proposed LO-SDA framework presents a significant advancement in data assimilation by introducing a generative approach that effectively overcomes the limitations of traditional Gaussian assumptions. Building upon the plausible theoretical connection between latent optimization and variational DA methods, we develop a novel framework that integrates observational information into the background conditional prior through latent optimization techniques. Through a combination of latent-conditioned diffusion modeling and such optimization-based observation constraints, LO-SDA provably achieves multiple posterior likelihood optimization during guidance sampling, guaranteeing progressive refinement while demonstrating superior performance compared to both classical variational methods and recent diffusion-based techniques.
Experimental results show that the method achieves substantial improvements in assimilation accuracy, outperforming 3DVAR by 28.25% and L3DVAR by 14.86% in z500 WRMSE at 5% observation coverage. The framework’s robustness is further validated under noisy observation conditions, maintaining consistent performance across varying error levels and latent optimization frequencies ablations.

Limitations and Impacts.
The current implementation operates under idealized background conditions, relying on 48-hour forecast backgrounds, which are not compatible with the cyclic DA system. The computational demands of frequent latent optimization also pose challenges for real-time operational use.
While the architecture naturally extends to 4D DA applications by incorporating model dynamics into the latent optimization process, further research is needed to evaluate its performance under non-Gaussian background error assumptions.

Despite these limitations, LO-SDA marks a pivotal shift in atmospheric data assimilation, demonstrating that generative models can effectively replace traditional Gaussian assumptions in high-dimensional systems. The framework’s success opens new possibilities for nonparametric DA in applications such as climate reanalysis and ensemble forecasting. By merging the flexibility of deep generative models with the rigorous constraints required in operational DA, LO-SDA represents a critical step toward next-generation data assimilation systems. Future work should focus on improving computational efficiency and expanding the framework’s applicability to more diverse and realistic atmospheric conditions.

## Appendix A WRMSE

The latitude-weighted root mean square error (WRMSE) is a statistical metric widely used in geospatial analysis and atmospheric science. Given the estimate $\hat{x}_{h,w,c}$ and the truth $x_{h,w,c}$, the WRMSE is defined as,

$$ $\mathrm{WRMSE}(c)=\sqrt{\frac{1}{H\cdot W}\sum_{h,w}H\frac{\cos(\alpha_{h,w})}{\sum_{h^{\prime}=1}^{H}\cos(\alpha_{h^{\prime},w})}(x_{h,w,c}-\hat{x}_{h,w,c})^{2}}\ .$ (16) $$

Here $H$ and $W$ represent the number of grid points in the longitudinal and latitudinal directions,
respectively, and $\alpha_{h,w}$ is the latitude of point $(h,w)$.

## Appendix B The VAE training and results

Model structure and training We utilize a transformer-based variational autoencoder framework (VAEformer) to effectively reduce the dimensionality of atmospheric data, mapping high-dimensional fields to a compact latent representation  han2024cra5. The architecture incorporates window-based attention mechanisms  liu2021swin to efficiently model atmospheric circulation patterns. Following the "vit_large" design paradigm, our implementation features identical encoder and decoder structures employing 4×4 patch embeddings with matching stride, a 1024-dimensional latent space, and a 24-layer transformer network utilizing window attention. The model was trained on ERA5 reanalysis data spanning 1979-2016, with the subsequent two-year period (2016-2018) serving as validation, over the course of 60 training epochs.

Results Our trained VAE achieves 0.0067 overall MSE and 0.0486 overall MAE. The varibles WRMSE are presented in Table [5](https://arxiv.org/html/2510.22562v1#A2.T5).

**Table 5: The VAE training results on WRMSE**
| u10 | v10 | t2m | msl | z50 | z100 | z150 | z200 | z250 | z300 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 0.54832 | 0.50501 | 0.82944 | 34.002 | 75.529 | 55.645 | 42.436 | 38.426 | 35.963 | 35.008 |
| z400 | z500 | z600 | z700 | z850 | z925 | z1000 | q50 | q100 | q150 |
| 31.623 | 28.4 | 25.948 | 24.563 | 23.415 | 24.37 | 27.266 | 9.64E-09 | 6.35E-08 | 4.90E-07 |
| q200 | q250 | q300 | q400 | q500 | q600 | q700 | q850 | q925 | q1000 |
| 3.04E-06 | 1.02E-05 | 2.47E-05 | 7.81E-05 | 1.68E-04 | 2.73E-04 | 4.01E-04 | 6.02E-04 | 5.95E-04 | 4.69E-04 |
| u50 | u100 | u150 | u200 | u250 | u300 | u400 | u500 | u600 | u700 |
| 0.91052 | 1.1085 | 1.3769 | 1.5108 | 1.5418 | 1.5148 | 1.3712 | 1.2184 | 1.1193 | 1.0552 |
| u850 | u925 | u1000 | v50 | v100 | v150 | v200 | v250 | v300 | v400 |
| 0.95107 | 0.78308 | 0.60413 | 0.80967 | 0.91698 | 1.1081 | 1.2589 | 1.3588 | 1.3544 | 1.2148 |
| v500 | v600 | v700 | v850 | v925 | v1000 | t50 | t100 | t150 | t200 |
| 1.0672 | 0.96791 | 0.90445 | 0.84409 | 0.71597 | 0.55351 | 0.59292 | 0.64698 | 0.47627 | 0.39086 |
| t250 | t300 | t400 | t500 | t600 | t700 | t850 | t925 | t1000 |  |
| 0.39115 | 0.41718 | 0.47806 | 0.48918 | 0.50497 | 0.5381 | 0.64627 | 0.61494 | 0.66865 |  |

## Appendix C Resampling

Here we provide a derivation of Equation [13](https://arxiv.org/html/2510.22562v1#S3.E13). Assume we have two independent Gaussian distributions:

$$ $\displaystyle p_{a}(x)$ $\displaystyle=\mathcal{N}(x;\mu_{a},\sigma_{a}^{2})$ (17) $\displaystyle p_{b}(x)$ $\displaystyle=\mathcal{N}(x;\mu_{b},\sigma_{b}^{2})$ (18) $$

The product distribution $p_{c}(x)=p_{a}(x)p_{b}(x)$ is also Gaussian with parameters:

$$ $\displaystyle\mu_{c}$ $\displaystyle=\frac{\mu_{a}/\sigma_{a}^{2}+\mu_{b}/\sigma_{b}^{2}}{1/\sigma_{a}^{2}+1/\sigma_{b}^{2}}$ (19) $\displaystyle\sigma_{c}^{2}$ $\displaystyle=\frac{1}{1/\sigma_{a}^{2}+1/\sigma_{b}^{2}}$ (20) $$

Proof:

$$ $\displaystyle p_{c}(x)$ $\displaystyle\propto\exp\left(-\frac{(x-\mu_{a})^{2}}{2\sigma_{a}^{2}}\right)\exp\left(-\frac{(x-\mu_{b})^{2}}{2\sigma_{b}^{2}}\right)$ $\displaystyle=\exp\left(-\frac{1}{2}\left(\frac{1}{\sigma_{a}^{2}}+\frac{1}{\sigma_{b}^{2}}\right)x^{2}+x\left(\frac{\mu_{a}}{\sigma_{a}^{2}}+\frac{\mu_{b}}{\sigma_{b}^{2}}\right)+C\right)$ (21) $$

where $C$ contains terms independent of $x$. Completing the square, we obtain:

$$ $\displaystyle p_{c}(x)$ $\displaystyle\propto\exp\left(-\frac{(x-\mu_{c})^{2}}{2\sigma_{c}^{2}}\right)$ (22) $$

with $\mu_{c}$ and $\sigma_{c}^{2}$ as defined above.

Now consider the conditional distribution in Equation [13](https://arxiv.org/html/2510.22562v1#S3.E13), where:

$$ $\displaystyle p(\bm{z}_{t}|\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ $\displaystyle\sim\mathcal{N}\left(\mu(t)\hat{\bm{z}}_{0}(\bm{y}),\sigma^{2}(t)\bm{I}\right)$ (23) $\displaystyle p(\tilde{\bm{z}}_{t}|\bm{z}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ $\displaystyle\sim\mathcal{N}(\bm{z}_{t},\lambda_{t}^{2}\bm{I})$ (24) $$

The posterior distribution is given by:

$$ $\displaystyle p(\bm{z}_{t}=\bm{a}|\tilde{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ $\displaystyle\propto p(\tilde{\bm{z}}_{t}|\bm{z}_{t}=\bm{a})p(\bm{z}_{t}=\bm{a}|\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ $\displaystyle\propto\exp\left(-\frac{\|\bm{a}-\tilde{\bm{z}}_{t}\|^{2}}{2\lambda_{t}^{2}}\right)\exp\left(-\frac{\|\bm{a}-\mu(t)\hat{\bm{z}}_{0}(\bm{y})\|^{2}}{2\sigma^{2}(t)}\right)$ $\displaystyle\propto\exp\left(-\frac{1}{2}\left(\frac{1}{\lambda_{t}^{2}}+\frac{1}{\sigma^{2}(t)}\right)\|\bm{a}\|^{2}+\left\langle\bm{a},\frac{\tilde{\bm{z}}_{t}}{\lambda_{t}^{2}}+\frac{\mu(t)\hat{\bm{z}}_{0}(\bm{y})}{\sigma^{2}(t)}\right\rangle\right)$ (25) $$

Applying the product formula for Gaussians, we obtain:

$$ $\displaystyle p(\bm{z}_{t}|\tilde{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ $\displaystyle\sim\mathcal{N}\left(\frac{\lambda_{t}^{2}\mu(t)\hat{\bm{z}}_{0}(\bm{y})+\sigma^{2}(t)\tilde{\bm{z}}_{t}}{\lambda_{t}^{2}+\sigma^{2}(t)},\frac{\lambda_{t}^{2}\sigma^{2}(t)}{\lambda_{t}^{2}+\sigma^{2}(t)}\bm{I}\right)$ (26) $$

## Appendix D More visualization.

We provide additional visualization results comparing different assimilation methods under 5% observation. In all appendix figures, the top row (left to right) displays the ERA5 ground truth, background field, and background error. The middle row shows assimilation results from (a) our proposed LO-SDA method, (b) the DPS framework, and (c) the Repaint approach, while the bottom row presents the corresponding absolute error fields relative to ERA5 truth. The significantly lighter error magnitudes in the LO-SDA results highlight our method’s superior error reduction capability compared to alternative approaches.

Figure: Figure 3: Visulaization of u500 at a 2019-08-26-06:00 UTC.
Refer to caption: Figures/u_vis.png

Figure: Figure 4: Visulaization of z500 at a 2019-05-18-06:00 UTC.
Refer to caption: z_vis.png

Figure: Figure 5: Visulaization of q700 at a 2019-02-02-06:00 UTC.
Refer to caption: Figures/q_vis.png

Figure: Figure 6: Visulaization of msl at a 2019-04-07-06:00 UTC.
Refer to caption: Figures/msl_vis.png

## References

- [1]
Andrew C Lorenc.
Analysis methods for numerical weather prediction.
Quarterly Journal of the Royal Meteorological Society,
112(474):1177–1194, 1986.
- [2]
Nils Gustafsson, Tijana Janjić, Christoph Schraff, Daniel Leuenberger,
Martin Weissmann, Hendrik Reich, Pierre Brousseau, Thibaut Montmerle, Eric
Wattrelot, Antonín Bučánek, et al.
Survey of data assimilation methods for convective-scale numerical
weather prediction at operational centres.
Quarterly Journal of the Royal Meteorological Society,
144(713):1218–1256, 2018.
- [3]
Mark Asch, Marc Bocquet, and Maëlle Nodet.
Data assimilation: methods, algorithms, and applications.
SIAM, 2016.
- [4]
Florence Rabier and Zhiquan Liu.
Variational data assimilation: theory and overview.
In Proc. ECMWF Seminar on Recent Developments in Data
Assimilation for Atmosphere and Ocean, Reading, UK, pages 29–43, 2003.
- [5]
Alberto Carrassi, Marc Bocquet, Laurent Bertino, and Geir Evensen.
Data assimilation in the geosciences: An overview of methods, issues,
and perspectives.
Wiley Interdisciplinary Reviews: Climate Change, 9(5):e535,
2018.
- [6]
François-Xavier Le Dimet and Olivier Talagrand.
Variational algorithms for analysis and assimilation of
meteorological observations: theoretical aspects.
Tellus A: Dynamic Meteorology and Oceanography, 38(2):97–110,
1986.
- [7]
Ross N Bannister.
A review of operational methods of variational and
ensemble-variational data assimilation.
Quarterly Journal of the Royal Meteorological Society,
143(703):607–633, 2017.
- [8]
Ross N Bannister.
A review of forecast error covariance statistics in atmospheric
variational data assimilation. i: Characteristics and measurements of
forecast error covariances.
Quarterly Journal of the Royal Meteorological Society: A journal
of the atmospheric sciences, applied meteorology and physical oceanography,
134(637):1951–1970, 2008.
- [9]
Ross N Bannister.
A review of forecast error covariance statistics in atmospheric
variational data assimilation. ii: Modelling the forecast error covariance
statistics.
Quarterly Journal of the Royal Meteorological Society: A journal
of the atmospheric sciences, applied meteorology and physical oceanography,
134(637):1971–1996, 2008.
- [10]
François Rozet and Gilles Louppe.
Score-based data assimilation.
In Thirty-seventh Conference on Neural Information Processing
Systems, 2023.
- [11]
François Rozet and Gilles Louppe.
Score-based data assimilation for a two-layer quasi-geostrophic
model.
2023.
- [12]
Yongquan Qu, Juan Nathaniel, Shuolin Li, and Pierre Gentine.
Deep generative data assimilation in multimodal setting.
In Proceedings of the IEEE/CVF Conference on Computer Vision and
Pattern Recognition, pages 449–459, 2024.
- [13]
Langwen Huang, Lukas Gianinazzi, Yuejiang Yu, Peter Dominik Dueben, and Torsten
Hoefler.
DiffDA: a diffusion model for weather-scale data assimilation.
In Forty-first International Conference on Machine Learning,
2024.
- [14]
Peter Manshausen, Yair Cohen, Peter Harrington, Jaideep Pathak, Mike Pritchard,
Piyush Garg, Morteza Mardani, Karthik Kashinath, Simon Byrne, and Noah
Brenowitz.
Generative data assimilation of sparse weather station observations
at kilometer scales, 2025.
- [15]
Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, and Liyue Shen.
Solving inverse problems with latent diffusion models via hard data
consistency.
In The Twelfth International Conference on Learning
Representations, 2024.
- [16]
G. Descombes, T. Auligné, F. Vandenberghe, D. M. Barker, and J. Barré.
Generalized background error covariance matrix model (GEN_BE
v2.0).
Geoscientific Model Development, 8(3):669–696, March 2015.
- [17]
Sibo Cheng, Yilin Zhuang, Lyes Kahouadji, Che Liu, Jianhua Chen, Omar K Matar,
and Rossella Arcucci.
Multi-domain encoder–decoder neural networks for latent data
assimilation in dynamical systems.
Computer Methods in Applied Mechanics and Engineering,
430:117201, 2024.
- [18]
Boštjan Melinc and Žiga Zaplotnik.
3d-var data assimilation using a variational autoencoder.
Quarterly Journal of the Royal Meteorological Society,
150(761):2273–2295, 2024.
- [19]
Mathis Peyron, Anthony Fillion, Selime Gürol, Victor Marchais, Serge
Gratton, Pierre Boudier, and Gael Goret.
Latent space data assimilation by using deep learning.
Quarterly Journal of the Royal Meteorological Society,
147(740):3759–3777, 2021.
- [20]
Maddalena Amendola, Rossella Arcucci, Laetitia Mottet, César Quilodrán
Casas, Shiwei Fan, Christopher Pain, Paul Linden, and Yi-Ke Guo.
Data assimilation in the latent space of a convolutional autoencoder.
In International Conference on Computational Science, pages
373–386. Springer, 2021.
- [21]
Qingyu Zheng, Guijun Han, Wei Li, Lige Cao, Gongfu Zhou, Haowen Wu, Qi Shao,
Ru Wang, Xiaobo Wu, Xudong Cui, Hong Li, and Xuan Wang.
Generating Unseen Nonlinear Evolution in Sea Surface
Temperature Using a Deep Learning-Based Latent Space Data Assimilation
Framework.
- [22]
Hang Fan, Ben Fei, Pierre Gentine, Yi Xiao, Kun Chen, Yubao Liu, Yongquan Qu,
Fenghua Ling, and Lei Bai.
Physically consistent global atmospheric data assimilation with
machine learning in a latent space.
arXiv preprint arXiv:2502.02884, 2025.
- [23]
Hang Fan, Yubao Liu, Zhaoyang Huo, Yuewei Liu, Yueqin Shi, and Yang Li.
A Novel Latent Space Data Assimilation Framework with
Autoencoder-Observation to Latent Space (AE-O2L) Network.
Part I: The Observation-Only Analysis Method.
Monthly Weather Review, 153(8):1335–1348, August 2025.
- [24]
Hang Fan, Yubao Liu, Yuewei Liu, Zhaoyang Huo, Baojun Chen, and Yu Qin.
A Novel Latent Space Data Assimilation Framework with
Autoencoder-Observation to Latent Space (AE-O2L) Network.
Part II: Observation and Background Assimilation with
Interpretability.
Monthly Weather Review, 153(8):1349–1363, August 2025.
- [25]
Prafulla Dhariwal and Alex Nichol.
Diffusion models beat gans on image synthesis.
In Proceedings of the 35th International Conference on Neural
Information Processing Systems, NIPS ’21, Red Hook, NY, USA, 2024. Curran
Associates Inc.
- [26]
Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi,
and Tim Salimans.
Cascaded diffusion models for high fidelity image generation.
Journal of Machine Learning Research, 23(47):1–33, 2022.
- [27]
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon.
Maximum likelihood training of score-based diffusion models.
Advances in neural information processing systems,
34:1415–1428, 2021.
- [28]
Chin-Wei Huang, Jae Hyun Lim, and Aaron C Courville.
A variational perspective on diffusion-based generative models and
score matching.
Advances in Neural Information Processing Systems,
34:22863–22876, 2021.
- [29]
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho.
Variational diffusion models.
Advances in neural information processing systems,
34:21696–21707, 2021.
- [30]
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
In Proceedings of the 34th International Conference on Neural
Information Processing Systems, NIPS ’20, Red Hook, NY, USA, 2020. Curran
Associates Inc.
- [31]
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
Deep unsupervised learning using nonequilibrium thermodynamics.
In International conference on machine learning, pages
2256–2265. PMLR, 2015.
- [32]
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano
Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential
equations.
In International Conference on Learning Representations, 2021.
- [33]
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and
Jong Chul Ye.
Diffusion posterior sampling for general noisy inverse problems.
In The Eleventh International Conference on Learning
Representations, 2023.
- [34]
Diederik P Kingma, Max Welling, et al.
Auto-encoding variational bayes, 2013.
- [35]
Carl Doersch.
Tutorial on variational autoencoders.
arXiv preprint arXiv:1606.05908, 2016.
- [36]
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, et al.
The era5 global reanalysis.
Quarterly journal of the royal meteorological society,
146(730):1999–2049, 2020.
- [37]
Stephan Rasp, Peter D Dueben, Sebastian Scher, Jonathan A Weyn, Soukayna
Mouatadid, and Nils Thuerey.
Weatherbench: a benchmark data set for data-driven weather
forecasting.
Journal of Advances in Modeling Earth Systems,
12(11):e2020MS002203, 2020.
- [38]
Stephan Rasp, Stephan Hoyer, Alexander Merose, Ian Langmore, Peter Battaglia,
Tyler Russell, Alvaro Sanchez-Gonzalez, Vivian Yang, Rob Carver, Shreya
Agrawal, et al.
Weatherbench 2: A benchmark for the next generation of data-driven
global weather models.
Journal of Advances in Modeling Earth Systems,
16(6):e2023MS004019, 2024.
- [39]
Kang Chen, Tao Han, Junchao Gong, Lei Bai, Fenghua Ling, Jing-Jia Luo, Xi Chen,
Leiming Ma, Tianning Zhang, Rui Su, et al.
Fengwu: Pushing the skillful global medium-range weather forecast
beyond 10 days lead.
arXiv preprint arXiv:2304.02948, 2023.
- [40]
Tao Han, Zhenghao Chen, Song Guo, Wanghan Xu, and Lei Bai.
Cra5: Extreme compression of era5 for portable global climate and
weather research via an efficient variational transformer.
arXiv preprint arXiv:2405.03376, 2024.
- [41]
Ilya Loshchilov, Frank Hutter, et al.
Fixing weight decay regularization in adam.
arXiv preprint arXiv:1711.05101, 5:5, 2017.
- [42]
William Peebles and Saining Xie.
Scalable diffusion models with transformers.
In Proceedings of the IEEE/CVF international conference on
computer vision, pages 4195–4205, 2023.
- [43]
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and
Baining Guo.
Swin transformer: Hierarchical vision transformer using shifted
windows.
In Proceedings of the IEEE/CVF international conference on
computer vision, pages 10012–10022, 2021.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Does the optimization always converge, and how many iterations are typically needed?
- How does LO-SDA handle strongly nonlinear observation operators such as satellite radiances?
soundness: 3
strengths: Clear motivation. The work targets the Gaussian-prior limitation of classical DA and builds a principled bridge between score-based and variational formulations.

Methodological coherence. The latent optimisation scheme is well justified and neatly integrated into the diffusion framework. Algorithm 1 clearly contrasts it with DPS guidance.

Presentation quality. Figures are clear and the connection to variational DA is articulated.
summary: The paper proposes LO-SDA, a score-based data assimilation (DA) framework that integrates latent optimisation with background-conditioned diffusion models. A variational autoencoder (VAE) is used to learn low-dimensional latent representations of atmospheric states; within this space, a diffusion model learns the conditional prior. During reverse diffusion, the authors introduce latent optimisation steps that iteratively enforce observation consistency, which they interpret as a generative analogue of variational cost-function minimisation. Experiments on ERA5 global reanalysis show that LO-SDA achieves performance comparable to or slightly better than variational DA (L3DVAR), outperforming single-step diffusion-based methods such as DPS and Repaint under sparse-observation settings.
weaknesses: - Incremental novelty. The paper’s central idea—enforcing observation consistency within score-based DA—is conceptually close to existing approaches such as score-based DA, DiffDA, and Manshausen et al., 2025. LO-SDA differs mainly by performing optimization in latent space rather than directly in model space and by framing the repeated guidance as “latent optimization.” This is a useful but evolutionary step rather than a paradigm shift.

- Overstated claims. The manuscript repeatedly suggests LO-SDA is the first score-based method to outperform traditional DA, but the margin over L3DVAR is modest and depends on selected metrics and observation density. The contribution would be stronger if positioned as a refinement of existing score-based methods rather than a breakthrough.

- Computational cost. Iterative latent optimisation requires ~5 minutes per assimilation versus seconds for variational baselines; scalability to operational use remains unclear.

- Limited theoretical depth. The link between latent optimisation and variational minimisation is described heuristically without formal guarantees or convergence analysis.

- Related-work positioning. Discussion could better situate LO-SDA among other generative and variational DA frameworks, e.g., Tensor-Var (2025) or DiffDA (2024), highlighting specific differences in assumptions and efficiency instead of focusing primarily on 3DVAR/L3DVAR comparisons.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Can the computational cost be reduced in the futher?
2. The paper notes a significant computational burden. Could the authors clarify which parts of the process dominate runtime (diffusion sampling, latent optimization, or decoding)?
3. The ERA5 experiments are convincing, but how does LO-SDA scale to operational forecast cycles or real-time assimilation scenarios?
soundness: 3
strengths: Establishes a theoretical bridge between variational DA and score-based diffusion modeling.
summary: This paper introduces LO-SDA, a generative data assimilation (DA) framework that replaces traditional Gaussian priors with a latent diffusion model conditioned on background atmospheric states. By incorporating an iterative latent optimization step during the diffusion reverse process, the method ensures strict consistency with sparse observations. Experiments on global ERA5 datasets show that LO-SDA outperforms both traditional variational DA (3DVAR) and its latent variant (L3DVAR), particularly under sparse observation settings.
weaknesses: - Computational cost for latent optimization process is expensive
- The work mainly integrates known elements (VAE compression + score-based DA + optimization) rather than proposing fundamentally new generative principles.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: Related Work: How does LO-SDA compare to existing latent-space assimilation methods like Latent-EnSF and LD-EnSF? These approaches also target high-dimensional systems with sparse observations using VAEs and score-based updates. It would be helpful if the authors could explain the differences and why LO-SDA is expected to perform better. Was there a reason these methods were omitted from discussion? A clear positioning relative to them would strengthen the paper’s originality claim.

Sequential Filtering: The proposed method handles one assimilation step via diffusion + optimization, but how is the temporal sequence of assimilation cycles handled? Is LO-SDA applied independently at each cycle with a newly encoded background, or is there any mechanism to propagate information forward (e.g., using the posterior as the next background)? Essentially, is LO-SDA intended as a one-step analysis technique or part of a full filtering framework? If the latter, do the authors observe any issues with filter consistency or error accumulation over long sequences?

Relatedly, how does the method perform when a single observation time is insufficient to constrain the latent state, such as in scenarios with high observation noise or highly underdetermined inverse problems? Are there any experiments that test LO-SDA's robustness in these noisy or ill-posed conditions, where temporal information or smoothing would typically be necessary to resolve ambiguity?

Performance at Extreme Sparsity: The results show only a slight edge (or parity) over L3DVar at 1% observation coverage. Do the authors have insight into why the improvement is small in this regime? Is the diffusion model struggling due to very weak observational constraints, or could it be that L3DVar is already very effective at low obs densities? Understanding this would help assess LO-SDA’s usefulness in the scarcest-data scenarios. Also, if possible, what modifications might improve performance when observations are extremely sparse? (e.g., more latent optimization steps, better priors, etc.)

Computational Cost: What is the runtime of LO-SDA for a single analysis, and how does it compare to L3DVar or other baselines? The paper mentions latent 3DVar takes on the order of seconds per assimilation. If LO-SDA is significantly slower due to diffusion sampling, it might be a concern for operational use. Any details on optimization steps or potential to speed it up (like using fewer diffusion steps or a more efficient sampler) would be helpful to know.
soundness: 2
strengths: Originality: The paper proposes a creative combination of deep generative modeling with data assimilation. Using a conditional diffusion model in a VAE latent space to perform Bayesian inference is a novel approach in the context of atmospheric data assimilation. The introduction of a latent optimization step to enforce observation consistency is an innovative twist that addresses a known issue in score-based filters (i.e., ensuring the analysis state actually fits the sparse observations).

Quality: The methodology builds on recent advances in both data assimilation and generative modeling. The authors validate the approach on a high-dimensional global atmosphere simulation, which lends credibility that the method can scale to realistic scenarios. The experiments compare LO-SDA against appropriate baselines: a traditional 3D-Var and a learned latent 3D-Var (L3DVar) method, as well as other score-based sampling approaches (e.g., diffusion posterior sampling and the Repaint technique). The results show LO-SDA consistently performs well, particularly as observation density increases, and it outperforms both the traditional variational method and prior score-based methods in most settings. The paper also provides statistical significance analysis to support the improvements, adding rigor to the empirical findings.

Clarity: The paper is generally well-written and structured. It motivates the problem (limitations of Gaussian priors in DA) clearly and provides sufficient background on data assimilation and score-based models for readability. The algorithmic description of LO-SDA is detailed, and the roles of the VAE and diffusion model are explained with appropriate equations. Figures (if any, e.g. illustrating the model architecture or results) are clear. Overall, the authors do a good job making a complex method accessible to readers from both the machine learning and data assimilation communities.

Significance: If the claims hold, this work is quite significant for data assimilation research. It challenges the long-standing reliance on Gaussian assumptions by showing that a learned generative prior can surpass a state-of-the-art variational DA method on a large-scale problem. This suggests a potential paradigm shift in how operational weather forecasting systems could incorporate machine learning. Moreover, LO-SDA’s strong performance under extremely sparse observations is important—sparse observational data is a common scenario (e.g., in parts of the globe or for certain variables), and methods that can handle such settings are highly valuable. By demonstrating an approach that remains accurate with as low as 1% observation coverage, the paper addresses a critical challenge in high-dimensional Bayesian filtering. The techniques introduced here (latent-space assimilation with diffusion models) could inspire further research and development of new hybrid DA algorithms in both the geosciences and broader sequential data domains.
summary: This paper introduces LO-SDA, a novel generative data assimilation method designed for numerical weather prediction. The key idea is to relax the Gaussian prior assumptions of traditional data assimilation by leveraging score-based generative modeling in a learned latent space. First, a variational autoencoder (VAE) is trained to encode high-dimensional atmospheric states into a compact latent representation. Then a diffusion model (score-based generator) conditioned on the background state is used to sample from an approximate posterior in this latent space. During the reverse diffusion process, the method applies latent optimization to enforce consistency with sparse observations, effectively nudging the generated latent state to agree with observed values.

The contributions include: (1) a new latent-space diffusion assimilation framework that can handle highly sparse observations by combining VAE compression, conditional diffusion sampling, and a novel observation-constrained optimization step; and (2) an experimental demonstration on an idealized global atmospheric dataset showing that this generative approach can match or surpass the accuracy of state-of-the-art variational methods. Notably, under extremely sparse observation coverage (only 1% of state variables observed), LO-SDA achieves comparable performance to a leading latent variational baseline (L3DVar). As observation density increases (e.g. 5% coverage), LO-SDA decisively outperforms the baseline, achieving about a 15% reduction in error (WRMSE). This is presented as the first instance where a diffusion-based data assimilation method demonstrates the potential to outperform traditional Gaussian-based methods at a global scale. The results suggest that generative modeling techniques, when properly constrained by observations, can overcome some limitations of classical methods and improve assimilation accuracy in challenging, high-dimensional settings.
weaknesses: Baseline Comparisons: A notable weakness is the absence of comparison with some established DA methods like the Ensemble Kalman Filter (e.g., LETKF). LETKF and similar ensemble Kalman techniques are widely used operationally and are strong benchmarks for assimilation performance. The paper claims to outperform “traditional approaches,” but this is only demonstrated against 3DVar/L3DVar. Without experiments against an ensemble Kalman filter, it is unclear how LO-SDA stands relative to the true state-of-practice. This missing baseline makes the empirical evaluation feel incomplete and may overstate the improvements of LO-SDA.

Improvement Magnitude: While LO-SDA does outperform the latent 3DVar (L3DVar) at higher observation densities, the gains under very sparse observations are quite modest. In the toughest scenario (1% observations), LO-SDA’s error is essentially on par with L3DVar. This suggests that the new method, despite its sophisticated generative approach, does not dramatically outshine the variational baseline when data are extremely scarce. The paper should be careful not to overclaim superiority; the results indicate LO-SDA’s advantage emerges mainly when there are a bit more observations (5% or above), whereas in the most data-starved case its performance difference is negligible. This limits the practical advantage in scenarios of extreme sparsity, which is supposedly a key motivation.

Related Work and Novelty: The authors do not discuss or cite several closely related approaches in latent-space and generative data assimilation. For example, Latent-EnSF
arxiv.org
 and LD-EnSF
arxiv.org
 are recent methods that also perform assimilation in a learned latent space to handle high-dimensional states with sparse observations (using ensemble score-based filtering). These works similarly leverage VAEs to encode states and observations and have demonstrated improved accuracy and efficiency in scenarios with very sparse data. The submission’s technical approach (using a VAE for state compression and doing assimilation in latent space) is quite aligned with those, though using a diffusion model instead of an ensemble score matching. Failing to compare with or even acknowledge Latent-EnSF(https://arxiv.org/abs/2409.00127) / LD-EnSF( https://arxiv.org/abs/2411.19305 ) is a significant oversight. It raises questions about the novelty: the idea of latent-space data assimilation for sparse observations is not entirely new, and the paper should clarify how LO-SDA advances the state of the art relative to these methods. Similarly, there has been work on Deep Bayesian Filters (DBF https://arxiv.org/abs/2405.18674) for data assimilation that introduce learned latent dynamics and aim to perform Bayes-optimal sequential inference (e.g., by constructing latent variables and using neural networks for filtering). The authors do not compare to such approaches either, leaving a gap in understanding how LO-SDA compares to other modern ML-based filtering methods beyond the variational ones.

Theoretical Limitations: From a Bayesian perspective, LO-SDA has some theoretical limitations that the paper does not fully address. Notably, the method does not sample the true posterior distribution of the state given observations; it uses a diffusion model trained on data and then applies a deterministic optimization to fit observations. This procedure may break the theoretical guarantees of Bayesian inference – the result is an analysis that is a single point estimate consistent with observations, but not necessarily a representative sample from the posterior uncertainty. As a consequence, the approach might underestimate uncertainty or fail to explore multiple possible solutions when the observations are sparse/ambiguous. Additionally, the method’s design raises concerns about temporal consistency in sequential assimilation: each assimilation cycle involves running the diffusion backward with an optimization, which does not obviously ensure consistency over time steps. There is no explicit mechanism to enforce that the analysis at one time is statistically consistent with the background distribution at the next time. In contrast, standard Bayes filtering or sequential Monte Carlo methods propagate uncertainty forward; here it seems each cycle is solved somewhat independently in latent space. This could potentially lead to filter inconsistency or drift over long sequences. The paper would benefit from a discussion on these theoretical aspects – for example, what are the implications of using this one-shot latent optimization per cycle, and could it be integrated into a more principled sequential Bayesian framework?

Scope of Experiments: The experiments, while promising, are conducted on an idealized global atmospheric model setup. It’s not entirely clear how complex or realistic the model is (e.g., resolution, model physics) and whether the performance would hold in a fully operational NWP setting. Also, the computational cost of LO-SDA is not thoroughly discussed; diffusion models can be expensive, and although latent space reduces dimension, one wonders if the method is efficient enough for real-time forecasting. Some clarification on runtime or complexity compared to baselines would strengthen the evaluation. These are minor issues, but addressing them would give a more complete picture of the method’s practical viability.

Pretraining Assumptions and Generalization: A further limitation lies in the pretraining assumptions behind the VAE encoder. LO-SDA requires a VAE trained on fully observed high-dimensional atmospheric states to learn the latent space. However, in many real-world applications, such full-state ground truth is unavailable or prohibitively expensive to collect. It is unclear how well LO-SDA would generalize to settings where only partial or noisy state estimates are available during training, which is often the case in operational meteorology. Moreover, the method’s effectiveness hinges on the quality of the learned latent representation. If the VAE is trained on idealized or unrealistic data, the resulting latent space may not be suitable for use in sparse, noisy, or changing environments. Addressing how LO-SDA can be adapted to such partially observed training regimes, or evaluating its robustness to domain shifts in latent modeling, would strengthen its practical relevance.

---

confidence: 4
contribution: 1
initial_rating: 2
presentation: 2
questions: I'm not actually sure whether the DA is performed repeatedly or only once after 48h after calculating the background? If so, how does the background distribution changing due to the data assimilation process itself actually affect the DA procedure later down the line? If not, then this would additionally need to be tested.

What differentiates this model from previous score-based data assimilation techniques (papers listed above)? Additionally, have the authors have tried any of the previous latent diffusion data assimilation models as baselines?
soundness: 2
strengths: The model improves over other methods such as Repaint, DPS, and 3DVar-esque techniques. Seems to have a respectable improvement under 5% observations, but for 1% observations, the improvements over other methods seem marginal at best (given the error bars). Thank you for updating the manuscript with some uncertainty/error bars so we can more closely look at the trends!

They test against both 1% and 5% sparse observations conditions, demonstrating that the model works for both. However, the trend seems as though it might not scale well with higher sparsity scenarios (<1%) relative to the other methods.

Due to optimizing in a latent space, the model should be more computationally efficient at assimilation for high-dimensional data. However, though this claim is made, there is no comparison of runtime with respect to traditional DA methods.
summary: First of all, I would like to state that I have previously been randomly assigned to this paper for a prior conference as well. If there are some conflicts of interests here, I humbly request the area chairs to take this into account and/or remove the review if needed. I may reference some of my previous points if they were left unaddressed in this updated version of the manuscript.


The authors create a background-conditioned latent diffusion prior inside a VAE, such that they minimize using a variational method to optimize for the observations. During the reverse pass, they apply alternating latent optimisation steps that explicitly minimise the observation-error cost.
weaknesses: The failure mode of repaint (compared to the background) is still not explored in this revised edition of the paper. Perhaps a small discussion would help strengthen the results by establishing the baseline in a more proper form.

As mentioned before, there is some limited novelty of the paper compared to recent diffusion-based data assimilation approaches such as APPA https://arxiv.org/abs/2504.18720, latent diffusion DA https://arxiv.org/abs/2406.14815, etc. Perhaps some discussion on the particular impacts of this work relative to some of the aforementioned papers would strengthen this paper.

The model seems to have a lot more variability due to noise uncertainty, as portrayed in Table 2, showing that it might not extrapolate to higher noise regimes. Is there a specific explanation as for such?

Also, I would still recommend adding a small note for the red and yellow highlights for first/second best.


Decision: Reject