stat.MLMar 22, 2026

Closed-form conditional diffusion models for data assimilation

Authors: Brianna Binder, Agnimitra Dasgupta, Assad Oberai

Organizations: Department of Aerospace & Mechanical Engineering, University of Southern California, 3650 McClintock Ave, Los Angeles, 90089, California, USA. · Optimization & Uncertainty Quantification, Sandia National Laboratories, 1450 Innovation Pkwy SE, Albuquerque, 87123, New Mexico, USA.

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We propose closed-form conditional diffusion models for data assimilation. Diffusion models use data to learn the score function (defined as the gradient of the log-probability density of a data distribution), allowing them to generate new samples from the data distribution by reversing a noise injection process. While it is common to train neural networks to approximate the score function, we leverage the analytical tractability of the score function to assimilate the states of a system with measurements. To enable the efficient evaluation of the score function, we use kernel density estimation to model the joint distribution of the states and their corresponding measurements. The proposed approach also inherits the capability of conditional diffusion models of operating in black-box settings, i.e., the proposed data assimilation approach can accommodate systems and measurement processes without their explicit knowledge. The ability to accommodate black-box systems combined with the superior capabilities of diffusion models in approximating complex, non-Gaussian probability distributions means that the proposed approach offers advantages over many widely used filtering methods. We evaluate the proposed method on nonlinear data assimilation problems based on the Lorenz-63 and Lorenz-96 systems of moderate dimensionality and nonlinear measurement models. Results show the proposed approach outperforms the widely used ensemble Kalman and particle filters when small to moderate ensemble sizes are used.

Figures & tables

CategoryValues
System Parametersσ=10,  ρ=28,  β=8/3\sigma=10,\;\rho=28,\;\beta=8/3
Assimilation ScheduleK=100,  Δt=0.1K=100,\;\Delta t=0.1
Process Integration SchemeForward Euler ( dt=0.01dt=0.01 )
Process Noiseϵ∼N(0,0.012I3)\bm{\epsilon}\sim\mathcal{N}(\bm{0},0.01^{2}\mathbb{I}_{3})
Observation Noiseη∼N(0,0.52)\eta\sim\mathcal{N}(0,0.5^{2})
Table 1: Lorenz-63 Experiment Settings
Ensemble size NOptimal BandwidthRange of sampling stepsAverage EW2\mathcal{E}_{W_{2}}
σx\sigma_{x}σy\sigma_{y}DiffusionEnKFSIR
200.2000.509 – 1212.80914.76817.440
500.1000.5012 – 169.77413.77317.867
1000.1000.2510 – 138.47413.76817.400
2500.0500.2511 – 156.55312.01716.882
5000.0250.2512 – 166.23311.86515.332
Table 2: For each ensemble size NN , the optimal kernel bandwidth parameters (σx,σy)(\sigma_{x},\sigma_{y}) selected for the conditional diffusion-based filter, the range of sampling steps used by the adaptive numerical integration scheme (see Section 3.3 ), and the corresponding time-averaged Wasserstein error EW2\mathcal{E}_{W_{2}} are reported for the Lorenz–63 system. The time-averaged Wasserstein errors for the EnKF and SIR filters are also reported. All errors are also averaged over S=10S=10 independent simulations. We indicate the best performing filter using bold fonts.
Figure 1: Evolution of the filtering distribution of x1x_{1} (vertical axis) for the Lorenz-63 system over K=100K=100 assimilation steps (horizontal axis). Each column corresponds to a different ensemble size and simulation. Within each column, the rows show (top to bottom): the reference distribution obtained using SIR with Ntrue=100,000N_{\rm true}=100{,}000 particles (Truth), the filtering distribution obtained with ensemble size NN using three different filters. The color intensity corresponds to the estimated probability density, with darker regions indicating higher density.
Figure 2: Predicted and updated states at assimilation step k=60k=60 and for N=500N=500 with three different filters. Red particles represent the states after the prediction step, blue particles represent the states after the update step, and the black line indicates the observation.
CategoryValues
System ParametersF=8F=8
Assimilation ScheduleK=500,  Δt=0.1K=500,\;\Delta t=0.1
Process Integration schemeRK4 ( dt=0.01dt=0.01 )
Process Noiseϵ∼N(0,0.012Id)\bm{\epsilon}\sim\mathcal{N}(\bm{0},0.01^{2}\mathbb{I}_{d})
Observation Noiseη∼N(0,0.5Id)\bm{\eta}\sim\mathcal{N}(\bm{0},0.5\mathbb{I}_{d})
Table 3: Lorenz-96 Experiment Settings
Ensemble size NOptimal BandwidthRange of sampling stepsAverage ERMSE\mathcal{E}_{RMSE}
σx\sigma_{x}σy\sigma_{y}DiffusionEnKFSIR
200.201.0013 –173.0734.9044.806
500.200.7514 –172.0494.6644.803
1000.200.5016 –191.6884.0234.675
2500.100.5016 –191.1442.1954.742
5000.100.5017 –191.0760.7294.732
Table 4: Optimal kernel bandwidths (σx,σy)(\sigma_{x},\sigma_{y}) , the range of sampling steps used by the adaptive numerical integration scheme, and the corresponding time-averaged RMSE error ERMSE\mathcal{E}_{RMSE} of the conditional diffusion model-based filter with varying ensemble sizes NN for the 10-dimensional Lorenz–96 system. The time-averaged RMSE errors for the EnKF and SIR filters are also reported. All errors are also averaged over S=10S=10 independent simulations. We indicate the best performing filter using bold fonts.
Figure 3: Visualization the performance of different filters with ensemble size N=100N=100 for one test trajectory of the 10-dimensional Lorenz-96 system over data assimilation steps 300-400. In each plot, the vertical axis corresponds to the states, and the horizontal axis corresponds to the data assimilation step. Rows 1 through 5 correspond to the observations, true states, estimated mean of the assimilated states, absolute error of mean with respect to the ground truth, and estimated ensemble spread (one standard deviation about the estimated mean), respectively. Columns 1 through 3 correspond to the conditional diffusion model-based, EnKF, and SIR filters, respectively.
Figure 4: Performance comparison of three filtering methods for a single degree of freedom, x1x_{1} , of the 10-dimensional Lorenz-96 system with ensemble size N=100N=100 over data assimilation steps 300-400, using the same test trajectory as in Fig. 3 . The top, middle, and bottom plots correspond to the conditional diffusion model-based filter, EnKF, and SIR filter, respectively. In each plot, the vertical axis corresponds to the value of x1x_{1} , and the horizontal axis corresponds to the data assimilation steps. The solid blue line represents the ensemble mean of the assimilated state, the dashed red line denotes the true state trajectory, the black crosses indicate the observations, and the shaded gray region shows the ensemble spread (one standard deviation about the estimated mean).
Ensemble size NOptimal BandwidthRange of sampling stepsAverage ERMSE\mathcal{E}_{RMSE}
σx\sigma_{x}σy\sigma_{y}DiffusionEnKFSIR
200.201.5013–173.5505.2515.008
500.151.0015–172.9045.0714.882
1000.150.7515–172.4564.8425.000
2500.100.7516–192.0743.9104.978
5000.100.7517–191.7712.1964.948
Table 5: Optimal kernel bandwidths (σx,σy)(\sigma_{x},\sigma_{y}) , the range of sampling steps used by the adaptive numerical integration scheme, and the corresponding time-averaged RMSE error ERMSE\mathcal{E}_{RMSE} of the conditional diffusion model-based filter with varying ensemble sizes NN for the 20-dimensional Lorenz–96 system. The time-averaged RMSE errors for the EnKF and SIR filters are also reported. All errors are also averaged over S=10S=10 independent simulations. We indicate the best performing filter using bold fonts.
Figure 5: Visualization the performance of different filters with ensemble size N=100N=100 for one test trajectory of the 20-dimensional Lorenz-96 system over data assimilation steps 300-400. In each plot, the vertical axis corresponds to the states, and the horizontal axis corresponds to the data assimilation step. Rows 1 through 5 correspond to the observations, true states, estimated mean of the assimilated states, absolute error of mean with respect to the ground truth, and estimated ensemble spread (one standard deviation about the estimated mean), respectively. Columns 1 through 3 correspond to the conditional diffusion model-based, EnKF, and SIR filters, respectively.
Figure 6: Performance comparison of three filtering methods for a single degree of freedom, x1x_{1} , of the 20-dimensional Lorenz-96 system with ensemble size N=100N=100 over data assimilation steps 300-400, using the same test trajectory as in Fig. 5 . The top, middle, and bottom plots correspond to the conditional diffusion model-based filter, EnKF, and SIR filter, respectively. In each plot, the vertical axis corresponds to the value of x1x_{1} , and the horizontal axis corresponds to the data assimilation steps. The solid blue line represents the ensemble mean of the assimilated state, the dashed red line denotes the true state trajectory, the black crosses indicate the observations, and the shaded gray region shows the ensemble spread (one standard deviation about the estimated mean).
More in this topic

Data Assimilation

Data assimilation combines observations with dynamical models to estimate hidden system states and their uncertainty over time.

See all 58 papers

Explore similar work

CardsList
Apr 8, 2026

Amortized Filtering and Smoothing with Conditional Normalizing Flows

Bayesian filtering and smoothing are central to data assimilation in nonlinear dynamical systems. Recent advances in deep generative models provide flexible approximations of the associated non-Gaussian posterior distributions. However, several existing approaches require the score to be re-estimated or a transport map to be constructed at each assimilation step. We propose an amortized framework for filtering and smoothing that reuses trained conditional models across observation sequences and assimilation times. The framework jointly learns a shared recurrent summary network and two conditional normalizing flows from simulated state and observation trajectories. The recurrent network represents each observation history by a fixed-dimensional summary that conditions the filtering approximation and, together with the next state, the backward transition approximation. An information-theoretic analysis shows that, under the Markov assumption, a summary sufficient for filtering is also sufficient for the backward kernel. Combining the terminal filtering approximation with the backward kernels yields an approximate joint smoothing distribution. Numerical experiments on advection-diffusion, Burgers, and Lorenz systems demonstrate the accuracy of the proposed filtering and smoothing approximations and characterize the evolution of their errors beyond the training horizon.

Tiangang Cui, Xiaodong Feng, Chenlong Pei +2

School of Mathematics and Statistics, The University of Sydney, Camperdown, NSW 2006, Australia · Institute of Advanced Study, Beijing Normal-Hong Kong Baptist University, Zhuhai 519087, China · SKLMS & Institute of Computational Mathematics and Scientific/Engineering Computing, Academy of Mathematics and Systems Science, Chinese Academy of Sciences +1

Bayesian FilteringData AssimilationBayesian InferenceAmortized Inference
AskPDF
May 18, 2026

SURGE: Approximation and Training Free Particle Filter for Diffusion Surrogate

Data assimilation (DA) addresses the problem of sequentially estimating the state of a dynamical system from noisy and incomplete observations. In this work, we employ a diffusion model as a world model to simulate and predict the system's dynamics. Recently, score-based diffusion models have learned global diffusion priors that effectively model (stochastic) dynamics, revealing strong potential for data assimilation. In this paper, we investigate how information from noisy observations can be incorporated to enable continuous correction and refinement of the predicted system state when using a diffusion prior. Motivated by particle filtering methods, we represent the posterior distribution using a set of particles. After receiving noisy observations, the diffusion model is guided using the observation likelihood to steer the generation process toward observation-consistent states. Nevertheless, such guidance does not guarantee sampling from the true posterior. We therefore employ a Sequential Monte Carlo approach over the diffusion trajectory, viewed as a path measure, to reweight and resample particles, thereby correcting the generation process and ensuring convergence toward the desired posterior distribution. This leads to an unbiased particle filtering method that rigorously fuses observational data with diffusion model simulations.

Lifu Wei, Yinuo Ren, Naichen Shi +1

Department of Mechanical Engineering, Northwestern University, Evanston, IL, United States · Institute for Computational & Mathematical Engineering, Stanford University, Stanford, CA, United States · Department of Industrial Engineering & Management Sciences, Northwestern University, Evanston, IL, United States

Particle FilteringBayesian FilteringData AssimilationDiffusion Models
AskPDF
May 27, 2026

Diffusion Model-Based Data Assimilation for Real-World Energy Consumption Forecasting

Accurate estimation and forecasting of energy consumption are important for power-system operation, planning, and demand-side management. In practice, however, complete and timely measurements may not always be available, and the observed data can be partial, noisy, or delayed. This motivates the use of learned forecasting models for predicting the evolving consumption state, together with data assimilation methods for sequential forecast correction. In this work, we study a high-dimensional data assimilation problem for real energy-consumption data. \modeltext{The forward prediction is supplied by a pretrained black-box spatio-temporal forecasting model, which is treated as the state propagator in the filtering procedure.} We employ the Ensemble Score Filter (EnSF) to assimilate partial and noisy observations and to correct the forecast trajectory over time. The EnSF uses score-based diffusion models to approximate filtering distributions and avoids retraining neural-network score models during assimilation by using a closed-form score representation and Monte Carlo approximation. Numerical experiments demonstrate that open-loop propagation of the learned forecasting model can become unreliable over long horizons, while EnSF-based correction substantially improves state estimation. Comparisons with the Ensemble Kalman Filter (EnKF) further show that EnSF provides stronger correction under the nonlinear observation setting considered in this work.

Ruoyu Hu, Dahai Yu, Feng Bao +2

Department of Mathematics, Florida State University, Tallahassee, 32306, Florida, USA · Department of Computer Science, Florida State University, Tallahassee, 32306, Florida, USA · Computer Science and Mathematics Division, Oak Ridge National Laboratory, Oak Ridge, 37831, Tennessee, USA

Time Series ForecastingData AssimilationElectricity Load Forecasting
AskPDF