Amortized Filtering and Smoothing with Conditional Normalizing Flows

Tiangang Cui; Xiaodong Feng; Chenlong Pei; Xiaoliang Wan; Tao Zhou

stat.MLApr 8, 2026

Amortized Filtering and Smoothing with Conditional Normalizing Flows

Authors: Tiangang Cui, Xiaodong Feng, Chenlong Pei, Xiaoliang Wan, Tao Zhou

Organizations: 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 · Department of Mathematics and Center for Computation and Technology, Louisiana State University, Baton Rouge 70803, USA

Abstract

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.

Figures & tables

Figure 1: Shared recurrent summary network and conditional flows. The summary network maps the observation history

y_{1:t}

to a fixed-dimensional summary statistic

s_{t}

. Conditioned on this summary, the forward flow approximates

p(u_{t}\mid y_{1:t})

and the backward flow approximates

p(u_{k}\mid u_{k+1},y_{1:k})

. The terminal forward density and the backward kernels define an approximate joint smoothing distribution.

			$N_{\mathrm{sample}}$
Problem	$N_{b}$	$N_{e}$	Filtering	One-step	Smoothing
Problem	$N_{b}$	$N_{e}$	Filtering	backward	Smoothing
Linear advection: state-dimension study	$128$	$100n$	$500$	$500$	$500$
Advection-diffusion: mesh-refinement study	$128$	$100n$	$500$	$500$	$500$
Burgers’ equation	$128$	$20000$	$500$	$500$	$500$
Single-scale Lorenz	$128$	$10000$	$500$	$500$	$500$

Table 1: Training and posterior sampling settings for AFSF.

	Filtering				One-step backward				Smoothing
	KL	RMSE	$\mathrm{MMD}^{2}$	CRPS	KL	RMSE	$\mathrm{MMD}^{2}$	CRPS	RMSE	$\mathrm{MMD}^{2}$	CRPS
$n=10$	0.0191	0.1457	0.0514	0.0822	0.0146	0.0949	0.0222	0.0536	0.1284	0.0401	0.0725
$n=20$	0.0443	0.1345	0.0860	0.0759	0.0262	0.0956	0.0446	0.0540	0.1207	0.0704	0.0683
$n=30$	0.0450	0.1288	0.1164	0.0727	0.0358	0.0960	0.0667	0.0542	0.1172	0.0979	0.0663
$n=40$	0.0524	0.1252	0.1444	0.0706	0.0434	0.0962	0.0883	0.0543	0.1149	0.1237	0.0650
$n=50$	0.0597	0.1228	0.1710	0.0693	0.0534	0.0965	0.1097	0.0544	0.1134	0.1484	0.0641

Table 2: Filtering, one-step backward, and smoothing results for the linear advection problem at different state dimensions.

Figure 2: Filtering (left) and one-step backward (right) KL divergences at

n=50

. The exact Gaussian conditionals are used as the first arguments. The vertical dashed line marks the end of the training interval.

Figure 3: Exact Gaussian reference mean, AFSF posterior mean, and pointwise two-standard-deviation bands for filtering (left) and smoothing (right) at three representative state components with

n=50

. The vertical line marks the end of the training interval.

	Filtering				One-step backward				Smoothing
	KL	RMSE	$\mathrm{MMD}^{2}$	CRPS	KL	RMSE	$\mathrm{MMD}^{2}$	CRPS	RMSE	$\mathrm{MMD}^{2}$	CRPS
$n=16$	0.1355	0.0558	0.0123	0.0314	0.0499	0.0236	0.0022	0.0133	0.0548	0.0102	0.0281
$n=32$	0.1363	0.0360	0.0103	0.0203	0.0653	0.0162	0.0021	0.0091	0.0360	0.0082	0.0182
$n=48$	0.1135	0.0267	0.0085	0.0150	0.0587	0.0121	0.0017	0.0068	0.0237	0.0067	0.0134
$n=64$	0.1268	0.0214	0.0073	0.0121	0.0592	0.0095	0.0014	0.0054	0.0190	0.0058	0.0107

Table 3: Filtering, one-step backward, and smoothing results for the advection-diffusion problem under mesh refinement.

Figure 4: Exact Gaussian reference mean, AFSF posterior mean, and pointwise two-standard-deviation bands for filtering (left) and smoothing (right) at three representative state components with

n=48

. The vertical line marks the end of the training interval.

Figure 5: Filtering (left) and one-step backward (right) KL divergences for the mesh-refinement experiment at

n=48

. The exact Gaussian conditionals are used as the first arguments. The vertical dashed line marks the end of the training interval.

	RMSE		$\mathrm{MMD}^{2}$		CRPS
	AFSF	FBF	AFSF	FBF	AFSF	FBF
$r^{2}=0.01$	0.0751	0.0752	0.0683	0.0681	0.0411	0.0410
$r^{2}=0.04$	0.0898	0.0917	0.0961	0.0994	0.0497	0.0506
$r^{2}=0.09$	0.1003	0.1015	0.1184	0.1203	0.0556	0.0561
$r^{2}=0.16$	0.1084	0.1113	0.1369	0.1426	0.0601	0.0615
$r^{2}=0.25$	0.1149	0.1179	0.1523	0.1584	0.0637	0.0650

Table 4: Filtering metrics for Burgers’ equation with

\nu=0.05

at different observation-noise variances.

	One-step backward			Smoothing
	RMSE	$\mathrm{MMD}^{2}$	CRPS	RMSE	$\mathrm{MMD}^{2}$	CRPS
$r^{2}=0.01$	0.0525	0.0340	0.0290	0.0710	0.0570	0.0373
$r^{2}=0.04$	0.0562	0.0387	0.0311	0.0879	0.0795	0.0458
$r^{2}=0.09$	0.0574	0.0405	0.0318	0.1056	0.0950	0.0502
$r^{2}=0.16$	0.0581	0.0414	0.0322	0.0981	0.1088	0.0531
$r^{2}=0.25$	0.0584	0.0418	0.0324	0.1185	0.1236	0.0581

Table 5: One-step backward and smoothing metrics for Burgers’ equation with

\nu=0.05

at different observation-noise variances.

Panel of Figure 6 — Figure 6: Filtering (top) and smoothing (bottom) for Burgers’ equation with $\nu=0.05$ and $r^{2}=0.25$ . The columns show the realized field, posterior mean, absolute error of the posterior mean, and marginal standard deviation.

	RMSE		$\mathrm{MMD}^{2}$		CRPS
	AFSF	FBF	AFSF	FBF	AFSF	FBF
$K=10$	0.1632	0.2044	0.0647	0.0967	0.0738	0.1015
$K=20$	0.1945	0.2439	0.1703	0.2530	0.0902	0.1272
$K=30$	0.2081	0.2604	0.2675	0.3987	0.0976	0.1393
$K=40$	0.2255	0.2742	0.3844	0.5295	0.1131	0.1489
$K=50$	0.2605	0.3106	0.5508	0.6918	0.1366	0.1668

Table 6: Filtering metrics for the single-scale Lorenz system at different state dimensions.

Figure 7: Time evolution of filtering RMSE,

\mathrm{MMD}^{2}

, and CRPS for AFSF and FBF in the single-scale Lorenz system at

K=50

. The vertical dashed line marks the end of the training interval.

	One-step backward			Smoothing
	RMSE	$\mathrm{MMD}^{2}$	CRPS	RMSE	$\mathrm{MMD}^{2}$	CRPS
$K=10$	0.1105	0.0301	0.0527	0.1298	0.0428	0.0598
$K=20$	0.1240	0.0738	0.0630	0.1525	0.1108	0.0742
$K=30$	0.1347	0.1253	0.0697	0.1680	0.1896	0.0838
$K=40$	0.1499	0.1979	0.0810	0.1881	0.2950	0.0991
$K=50$	0.1695	0.2966	0.0935	0.5423	0.4594	0.1844

Table 7: One-step backward and smoothing metrics for the single-scale Lorenz system.

Figure 8: Realized path, posterior mean, and pointwise two-standard-deviation bands for filtering (left) and smoothing (right) at three representative state components of the single-scale Lorenz system with

K=50

and

N_{\mathrm{train}}=3000

. The vertical line marks the end of the training interval.

	RMSE		$\mathrm{MMD}^{2}$		CRPS
	AFSF	FBF	AFSF	FBF	AFSF	FBF
$F=5$	0.4634	0.5356	0.5464	0.6065	0.2169	0.2450
$F=8$	0.4565	0.7756	0.5157	0.7799	0.2105	0.3472
$F=16$	0.4872	2.4911	0.5030	1.0136	0.2258	1.0967

Table 8: Filtering metrics for the two-scale Lorenz system at different forcing values.

	One-step backward			Smoothing
	RMSE	$\mathrm{MMD}^{2}$	CRPS	RMSE	$\mathrm{MMD}^{2}$	CRPS
$F=5$	0.0859	0.0290	0.0453	0.4232	0.4600	0.1909
$F=8$	0.0884	0.0308	0.0471	0.4121	0.4286	0.1839
$F=16$	0.1024	0.0407	0.0564	0.5141	0.4454	0.2249

Table 9: One-step backward and smoothing metrics for the two-scale Lorenz system.

Figure 9: Realized slow-state path, posterior mean, and pointwise two-standard-deviation bands for filtering (left) and smoothing based on the one-step backward approximation (right) for three observed slow variables at

F=5

. The vertical lines mark the end of the training interval.

Linear advection, $n=50$
	Filtering		One-step backward		Smoothing
Summary	RMSE	CRPS	RMSE	CRPS	RMSE	CRPS
Shared	0.1228	0.0693	0.0965	0.0544	0.1134	0.0641
Separate	0.1232	0.0696	0.0967	0.0547	0.1139	0.0644
Single-scale Lorenz, $K=50$
Shared	0.1945	0.0902	0.1240	0.0630	0.1525	0.0742

Table 10: Filtering, one-step backward, and smoothing scores with shared and separate observation summaries at

T=500

. The two-scale smoothing scores refer to the first-order closure.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

	RMSE		$\mathrm{MMD}^{2}$		CRPS
	AFSF	FBF	AFSF	FBF	AFSF	FBF
$r^{2}=0.01$	0.1140	0.1152	0.1501	0.1521	0.0588	0.0594
$r^{2}=0.04$	0.1270	0.1301	0.1830	0.1897	0.0686	0.0699
$r^{2}=0.09$	0.1380	0.1403	0.2120	0.2169	0.0754	0.0766
$r^{2}=0.16$	0.1467	0.1507	0.2360	0.2454	0.0807	0.0827
$r^{2}=0.25$	0.1546	0.1640	0.2581	0.2833	0.0852	0.0900

Appendix

Table C.1: Filtering metrics for Burgers’ equation with

\nu=0.01

Figure C.1: Realized states, posterior means, and pointwise two-standard-deviation bands for Burgers’ equation with

r^{2}=0.25

. The first two columns show filtering and smoothing, respectively, at three spatial locations for

\nu=0.05

. The last two columns show filtering for

\nu=0.01

, with temporal evolution at three spatial locations and spatial profiles at three times, respectively.

Bayesian Filtering

Bayesian filtering recursively estimates hidden states and their uncertainty from noisy, sequential observations in changing systems.

See all 56 papers

Closed-form conditional diffusion models for data assimilation
Mar 22, 2026Brianna Binder, Agnimitra Dasgupta, Assad OberaiData Assimilation Conditional Diffusion Models
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.
Generative Model Proposal based Particle Filtering for Data Assimilation
Jul 1, 2026Chandni Nagda, Mayank Shrivastava, Gudrun Thorkelsdottir +3Particle Filtering Data Assimilation
University of Illinois at Urbana-Champaign · 2NVIDIA
DAWIS: Data Assimilation with Windowed Inverse Sampling via Multitask Interpolants
Oct 2, 2026Erik Wikingsson, Martin Andrae, Tomas Landelius +1Data Assimilation Flow-Based Generative Modeling
Linköping University, Sweden · Swedish Meteorological and Hydrological Institute

Mar 22, 2026J/K move · Enter open · S save

Closed-form conditional diffusion models for data assimilation

Brianna Binder, Agnimitra Dasgupta, Assad Oberai

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.

Data Assimilation Conditional Diffusion Models Conditional Distribution Estimation

Open paper AskPDF arXiv

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.

Abstract

Bayesian Filtering

Explore similar work

Closed-form conditional diffusion models for data assimilation

Generative Model Proposal based Particle Filtering for Data Assimilation

DAWIS: Data Assimilation with Windowed Inverse Sampling via Multitask Interpolants