Closed-form conditional diffusion models for data assimilation

Brianna Binder; Agnimitra Dasgupta; Assad Oberai

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.

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

Category	Values
System Parameters	$\sigma=10,\;\rho=28,\;\beta=8/3$
Assimilation Schedule	$K=100,\;\Delta t=0.1$
Process Integration Scheme	Forward Euler ( $dt=0.01$ )
Process Noise	$\bm{\epsilon}\sim\mathcal{N}(\bm{0},0.01^{2}\mathbb{I}_{3})$
Observation Noise	$\eta\sim\mathcal{N}(0,0.5^{2})$

Table 1: Lorenz-63 Experiment Settings

Ensemble size N	Optimal Bandwidth		Range of sampling steps	Average $\mathcal{E}_{W_{2}}$
Ensemble size N	$\sigma_{x}$	$\sigma_{y}$	Range of sampling steps	Diffusion	EnKF	SIR
20	0.200	0.50	9 – 12	12.809	14.768	17.440
50	0.100	0.50	12 – 16	9.774	13.773	17.867
100	0.100	0.25	10 – 13	8.474	13.768	17.400
250	0.050	0.25	11 – 15	6.553	12.017	16.882
500	0.025	0.25	12 – 16	6.233	11.865	15.332

Table 2: For each ensemble size

N

, the optimal kernel bandwidth parameters

(\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

\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=10

independent simulations. We indicate the best performing filter using bold fonts.

Panel of Figure 1 — Figure 1: Evolution of the filtering distribution of $x_{1}$ (vertical axis) for the Lorenz-63 system over $K=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 $N_{\rm true}=100{,}000$ particles (Truth), the filtering distribution obtained with ensemble size $N$ 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=60$ and for $N=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.

Category	Values
System Parameters	$F=8$
Assimilation Schedule	$K=500,\;\Delta t=0.1$
Process Integration scheme	RK4 ( $dt=0.01$ )
Process Noise	$\bm{\epsilon}\sim\mathcal{N}(\bm{0},0.01^{2}\mathbb{I}_{d})$
Observation Noise	$\bm{\eta}\sim\mathcal{N}(\bm{0},0.5\mathbb{I}_{d})$

Table 3: Lorenz-96 Experiment Settings

Ensemble size N	Optimal Bandwidth		Range of sampling steps	Average $\mathcal{E}_{RMSE}$
Ensemble size N	$\sigma_{x}$	$\sigma_{y}$	Range of sampling steps	Diffusion	EnKF	SIR
20	0.20	1.00	13 –17	3.073	4.904	4.806
50	0.20	0.75	14 –17	2.049	4.664	4.803
100	0.20	0.50	16 –19	1.688	4.023	4.675
250	0.10	0.50	16 –19	1.144	2.195	4.742
500	0.10	0.50	17 –19	1.076	0.729	4.732

Table 4: Optimal kernel bandwidths

(\sigma_{x},\sigma_{y})

, the range of sampling steps used by the adaptive numerical integration scheme, and the corresponding time-averaged RMSE error

\mathcal{E}_{RMSE}

of the conditional diffusion model-based filter with varying ensemble sizes

N

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=10

independent simulations. We indicate the best performing filter using bold fonts.

Figure 3: Visualization the performance of different filters with ensemble size $N=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, $x_{1}$ , of the 10-dimensional Lorenz-96 system with ensemble size $N=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 $x_{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 N	Optimal Bandwidth		Range of sampling steps	Average $\mathcal{E}_{RMSE}$
Ensemble size N	$\sigma_{x}$	$\sigma_{y}$	Range of sampling steps	Diffusion	EnKF	SIR
20	0.20	1.50	13–17	3.550	5.251	5.008
50	0.15	1.00	15–17	2.904	5.071	4.882
100	0.15	0.75	15–17	2.456	4.842	5.000
250	0.10	0.75	16–19	2.074	3.910	4.978
500	0.10	0.75	17–19	1.771	2.196	4.948

Table 5: Optimal kernel bandwidths

(\sigma_{x},\sigma_{y})

, the range of sampling steps used by the adaptive numerical integration scheme, and the corresponding time-averaged RMSE error

\mathcal{E}_{RMSE}

of the conditional diffusion model-based filter with varying ensemble sizes

N

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=10

independent simulations. We indicate the best performing filter using bold fonts.

Figure 5: Visualization the performance of different filters with ensemble size $N=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, $x_{1}$ , of the 20-dimensional Lorenz-96 system with ensemble size $N=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 $x_{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).

Data Assimilation

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

See all 58 papers

Amortized Filtering and Smoothing with Conditional Normalizing Flows
Apr 8, 2026Tiangang Cui, Xiaodong Feng, Chenlong Pei +2Bayesian Filtering Data Assimilation
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
SURGE: Approximation and Training Free Particle Filter for Diffusion Surrogate
May 18, 2026Lifu Wei, Yinuo Ren, Naichen Shi +1Particle Filtering Bayesian Filtering
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
Diffusion Model-Based Data Assimilation for Real-World Energy Consumption Forecasting
May 27, 2026Ruoyu Hu, Dahai Yu, Feng Bao +2Time Series Forecasting Data Assimilation
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

Apr 8, 2026J/K move · Enter open · S save

Amortized Filtering and Smoothing with Conditional Normalizing Flows

Tiangang Cui, Xiaodong Feng, Chenlong Pei, Xiaoliang Wan+1

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 Filtering Data Assimilation Bayesian Inference Amortized Inference

Open paper AskPDF arXiv

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.

Abstract

Data Assimilation

Explore similar work

Amortized Filtering and Smoothing with Conditional Normalizing Flows

SURGE: Approximation and Training Free Particle Filter for Diffusion Surrogate

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