LO-SDA: Latent Optimization for Score-based Atmospheric Data Assimilation
Jing-An Sun, Hang Fan, Junchao Gong, Ben Fei, Kun Chen, Fenghua Ling, Wenlong Zhang, Wanghan Xu, Li Yan, Pierre Gentine, Lei Bai
TL;DR
LO-SDA introduces a latent diffusion framework for data assimilation that replaces Gaussian priors with a background-conditioned diffusion model in latent space, coupled with an alternating latent optimization that enforces strict consistency with sparse observations. By training a VAE to capture nonlinear atmospheric correlations and modeling $p(\mathbf{z}|\mathbf{z}_b)$ with a score-based diffusion model, LO-SDA achieves improved analyses over traditional 3DVAR and recent diffusion-based methods, particularly as observational density increases. The method demonstrates strong performance on ERA5-based experiments and maintains robustness under observation noise, with real-world GDAS data showing competitive results. This approach suggests that nonparametric, generative modeling can effectively supplant Gaussian priors in high-dimensional atmospheric data assimilation, offering potential gains for climate reanalysis and ensemble forecasting, albeit with computational considerations for operational deployment.
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.