FREESIA: Covariance-Aware Posterior Transport for Expressive and Scalable Data Assimilation

Shiwei Ni; Yangwen Zhang; Hang Qi; Xiaofei Guan; Lili Ju

stat.MLSep 19, 2026

FREESIA: Covariance-Aware Posterior Transport for Expressive and Scalable Data Assimilation

Authors: Shiwei Ni, Yangwen Zhang, Hang Qi, Xiaofei Guan, Lili Ju

Organizations: School of Mathematical Sciences, Tongji University, Shanghai, China · Department of Mathematics, University of South Carolina, Columbia, SC 29208, USA

Abstract

Data assimilation aims to infer the state of complex dynamical systems based on observational data. However, accurate inference of the multimodal posteriors induced by nonlinear or non-injective observation operators remains a key challenge under high-dimensional and sparse observation conditions. Ensemble filters scale to high dimensions but are confined by restrictive distributional assumptions, while training-free generative filters (e.g., EnSF, EnFF) alleviate this limitation but may introduce structural errors and hinder information propagation under sparse observations. To address these issues, we propose a training-free, asymptotically exact posterior transport method. Firstly, a covariance-aware posterior transport scheme is designed, which embeds the forecast cross-covariance into flow-based transport and accurately recovers unobserved states while preserving the non-Gaussian posterior structure. Furthermore, the method combines a tractable observation-adaptive proposal with posterior correction, ensuring accurate approximation of the nonlinear posterior distribution. Finally, we establish the corresponding posterior flow theory, from which the asymptotic exactness of the proposed method relative to finite-ensemble surrogates and the Wasserstein error bound are derived. Experiments on Double-Well, Lorenz-96, and Kolmogorov flow show that the proposed method captures complex posterior structure and remains accurate under sparse, nonlinear, and non-injective observations. In the sparse non-injective setting, it reduces RMSE by 56% relative to the best baseline.

Figures & tables

Figure 1: Overview of the FREESIA posterior-flow update. Given $z_{t}$ , Step 1 samples an observation-adaptive GMM component, Step 2 draws a covariance-aware conditional endpoint, and Step 3 applies nonlinear likelihood-ratio correction to estimate the endpoint mean and advance the flow.

Method	JS $\downarrow$	$W_{1}$ $\downarrow$	Mode err. $\downarrow$	RMSE $\downarrow$	CRPS $\downarrow$
FREESIA (ours)	0.031	0.059	0.035	0.144	0.099
EnKF-MDA	0.185	0.320	0.259	0.433	0.374
EnKF	0.184	0.233	0.158	0.325	0.257
EnSF	0.315	0.666	0.359	0.699	0.379
EnFF-OT	0.550	0.202	0.130	0.263	0.254
EnFF-F2P	0.358	0.152	0.085	0.232	0.170

Table 1: Double-well posterior approximation errors averaged over experiments and analysis times. Lower is better; best and second-best results are bold and underlined.

	Dense $\|\arctan\|$		Stride-2 $\arctan$				Stride-4 $\arctan$
Method	RMSE $\downarrow$	CRPS $\downarrow$	RMSE $\downarrow$	Obs. $\downarrow$	Unobs. $\downarrow$	CRPS $\downarrow$	RMSE $\downarrow$	Obs. $\downarrow$	Unobs. $\downarrow$	CRPS $\downarrow$
FREESIA (ours)	0.490	0.174	0.615	0.417	0.757	0.241	1.575	0.373	1.805	0.584
EnKF-MDA	1.384	0.586	0.987	0.673	1.211	0.232	2.056	0.366	2.358	0.955
EnKF	3.053	1.712	1.380	1.120	1.573	0.422	2.684	1.689	2.928	1.356
LETKF	6.161	3.921	1.533	1.243	1.753	0.494	3.185	1.853	3.513	1.648
EnSF	3.901	2.340	2.027	1.531	2.422	0.976	2.472	1.143	2.777	1.242

Table 2: Lorenz–96 filtering under three nonlinear observation regimes. RMSE and CRPS are averaged across repeated experiments and all 101 assimilation cycles. Lower is better; best and second-best values are bold and underlined.

	1% random $\arctan$				2% random $\|\arctan\|$
Method	RMSE $\downarrow$	Obs. $\downarrow$	Unobs. $\downarrow$	CRPS $\downarrow$	RMSE $\downarrow$	Obs. $\downarrow$	Unobs. $\downarrow$	CRPS $\downarrow$
FREESIA (ours)	0.095	0.076	0.095	0.052	0.422	0.410	0.422	0.209
EnKF-MDA	0.148	0.115	0.148	0.081	1.121	1.115	1.121	0.797
EnKF	0.377	0.372	0.377	0.200	0.987	0.984	0.987	0.681
LETKF	0.218	0.177	0.218	0.117	0.970	1.088	0.967	0.572
EnSF	0.942	0.568	0.945	0.547	0.974	0.967	0.975	0.569

Table 3: Kolmogorov-flow filtering under sparse nonlinear observations. Metrics average ten trajectories and all 51 analyses; lower is better, with best and second-best values bold and underlined.

Figure 2: Final-time reconstruction for the 2% absolute-arctangent Kolmogorov-flow experiment. Top: truth, FREESIA, and EnKF vorticity, with analysis-mean velocity RMSE $0.031$ vs. $1.399$ . Middle: 328 observed spatial locations, corresponding to 656 observed velocity components (2% of the state), and vorticity residuals. Bottom: observed-point velocity directions; FREESIA/EnKF achieve $100\%/55\%$ within $90^{\circ}$ , with mean angular errors $0.9^{\circ}/82.7^{\circ}$ .

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Double-well posterior densities at analysis steps 40, 50, and 60. Rows show FREESIA, EnFF-F2P, and EnKF-MDA against the grid Bayes posterior.

Figure 4: Double-well hyperparameter search with $N_{\rm run}=3$ per setting. Left: all five reported errors versus the GMM separation parameter $\gamma$ , with a logarithmic vertical axis; the rightmost point is $\gamma=0.99$ . Center and right: RMSE and CRPS over flow steps and endpoint importance samples $N_{\mathrm{IS}}$ . Lower is better throughout; values average all search runs and analysis times.

Figure 5: Complete double-well hyperparameter search over flow steps and endpoint importance samples $N_{\mathrm{IS}}$ . The panels report RMSE, CRPS, Jensen–Shannon divergence, Wasserstein-1 distance, and mode-mass error; lower is better for every metric. Cell values are averaged over the hyperparameter-search repetitions and analysis times.

Method	Dense abs-atan		Stride-2 atan		Stride-4 atan
	$r_{\mathrm{loc}}$	$\lambda_{\mathrm{infl}}$	$r_{\mathrm{loc}}$	$\lambda_{\mathrm{infl}}$	$r_{\mathrm{loc}}$	$\lambda_{\mathrm{infl}}$
FREESIA	4	1.8	4	1.2	2	1.0
EnKF	2	1.1	4	1.2	3	1.1
EnKF-MDA	3	1.6	3	1.1	2	1.2
LETKF	2	1.2	2	1.5	2	1.3

Appendix

Table 4: Localization radius

r_{\mathrm{loc}}

and inflation parameter

\lambda_{\mathrm{infl}}

in the three Lorenz–96 experiments. GC denotes cyclic Gaspari–Cohn localization.

Figure 6: Total analysis RMSE for dense absolute-arctangent, stride-2 arctangent, and stride-4 arctangent observations. Each panel shows the four methods with the lowest time-averaged RMSE in that regime; the shared legend lists the union of methods shown across panels.

Figure 7: Observed- and unobserved-coordinate RMSE for stride-2 (top) and stride-4 (bottom) observations. The right column directly evaluates information transfer to unobserved state coordinates.

Figure 8: Representative trajectories for the first four Lorenz–96 state coordinates under dense absolute-arctangent observations. The panels show the truth and analysis trajectories for the methods selected by the corresponding RMSE comparison.

Figure 9: Representative trajectories for the first four Lorenz–96 state coordinates under stride-2 arctangent observations. The panels show the truth and analysis trajectories for the methods selected by the corresponding RMSE comparison.

Figure 10: Representative trajectories for the first four Lorenz–96 state coordinates under stride-4 arctangent observations. The panels show the truth and analysis trajectories for the methods selected by the corresponding RMSE comparison.

Figure 11: Lorenz–96 stride-4 ablation. Curves show mean analysis RMSE over 10 matched runs, and shading denotes one standard deviation across runs. Diagonal covariance removes cross-coordinate covariance terms. The no-likelihood-correction variant replaces the nonlinear likelihood-ratio weights with uniform endpoint weights.

	$\sigma=0.01$				$\sigma=0.05$
Method	RMSE	Obs.	Unobs.	CRPS	RMSE	Obs.	Unobs.	CRPS
FREESIA	1.575	0.373	1.805	0.584	2.076	0.938	2.335	1.013
EnKF-MDA	2.056	0.366	2.358	0.955	2.544	0.894	2.891	1.342
EnKF	2.684	1.689	2.928	1.356	2.642	1.603	2.903	1.407
EnSF	2.472	1.143	2.777	1.242	2.865	1.937	3.113	1.582
EnFF-F2P	3.960	4.412	3.796	2.243	2.893	2.263	3.074	1.565

Appendix

Table 5: Time-averaged total, observed-coordinate, and unobserved-coordinate RMSE and CRPS for Lorenz–96 stride-4 observations under two noise levels. Lower is better.

Figure 12: Accuracy–cost trade-off for Lorenz–96 stride-4 assimilation. FREESIA fixes

N_{\mathrm{IS}}=5

and varies only the number of flow steps

T

; error bars denote one standard deviation across repeated runs. Baseline points use the configurations reported in the main comparison.

Figure 13: Importance-weight diagnostics for the Lorenz–96 stride-4 experiment. Left: normalized effective sample size (nESS) as a function of flow time, aggregated over filtering states and flow particles. Right: nESS across the filtering sequence at selected flow times. Importance weights are strongly concentrated over the early and intermediate portions of the flow, while nESS rises rapidly near the endpoint as the conditional bridge contracts. The concentration is primarily flow-time dependent rather than showing systematic deterioration across the filtering sequence. The higher-resolution configuration is used only to resolve the near-endpoint weight behavior.

Figure 14: Analysis RMSE for 1% arctangent observations (left) and 2% absolute-arctangent observations (right). Each panel shows the four methods with the lowest time-averaged RMSE across the ten trajectories.

Figure 15: Accuracy–cost trade-off on 2-D Kolmogorov flow with 1% random arctangent observations. The sweep fixes

N_{\mathrm{IS}}=1

and

r_{\mathrm{loc}}=16

, and varies only the number of flow steps

T

. The rightmost reference point is the default benchmark configuration

(T=20,N_{\mathrm{IS}}=5)

, with the same localization radius. Error bars denote one standard deviation across ten trajectories. Baseline points use their reported configurations.

	1% atan		2% abs-atan
Method	$r_{\mathrm{loc}}$	$\lambda_{\mathrm{infl}}$	$r_{\mathrm{loc}}$	$\lambda_{\mathrm{infl}}$
FREESIA	16	1.1	16	1.1
EnKF	8	1.1	6	1.1
EnKF-MDA	10	1.2	6	1.1
LETKF	3	1.1	2	1.1

Appendix

Table 6: Localization radius

r_{\mathrm{loc}}

and covariance inflation

\lambda_{\mathrm{infl}}

for the Kolmogorov-flow comparisons.

Method	JS	$W_{1}$	Mode error	RMSE	CRPS
FREESIA	$0.03134\pm 0.00047$	$0.05936\pm 0.00053$	$0.03503\pm 0.00009$	$0.14417\pm 0.00035$	$0.09915\pm 0.00002$
EnKF-MDA	$0.18521\pm 0.00001$	$0.31992\pm 0.00004$	$0.25878\pm 0.00001$	$0.43262\pm 0.00003$	$0.37445\pm 0.00001$
EnKF	$0.18365\pm 0.00035$	$0.23281\pm 0.00050$	$0.15807\pm 0.00010$	$0.32489\pm 0.00051$	$0.25673\pm 0.00017$
EnSF	$0.31497\pm 0.00161$	$0.66648\pm 0.00404$	$0.35924\pm 0.00398$	$0.69883\pm 0.00261$	$0.37946\pm 0.00186$
EnFF-OT	$0.54998\pm 0.00007$	$0.20219\pm 0.00003$	$0.12996\pm 0.00001$	$0.26335\pm 0.00004$	$0.25411\pm 0.00004$
EnFF-F2P	$0.35776\pm 0.00115$	$0.15151\pm 0.00243$	$0.08547\pm 0.00072$	$0.23215\pm 0.00183$	$0.16993\pm 0.00080$

Appendix

Table 7: Double-Well metrics (mean

\pm

standard deviation over 10 runs).

	Dense abs-atan		Stride-2 atan				Stride-4 atan
Method	RMSE	CRPS	RMSE	Obs.	Unobs.	CRPS	RMSE	Obs.	Unobs.	CRPS
FREESIA	$0.490\pm 0.038$	$0.174\pm 0.007$	$0.615\pm 0.060$	$0.417\pm 0.025$	$0.757\pm 0.082$	$0.241\pm 0.014$	$1.575\pm 0.085$	$0.373\pm 0.017$	$1.805\pm 0.098$	$0.584\pm 0.041$
EnKF-MDA	$1.384\pm 0.255$	$0.586\pm 0.107$	$0.987\pm 0.243$	$0.673\pm 0.153$	$1.211\pm 0.307$	$0.232\pm 0.006$	$2.056\pm 0.045$	$0.366\pm 0.021$	$2.358\pm 0.051$	$0.955\pm 0.024$
EnKF	$3.053\pm 0.155$	$1.712\pm 0.148$	$1.380\pm 0.153$	$1.120\pm 0.141$	$1.573\pm 0.168$	$0.422\pm 0.026$	$2.684\pm 0.089$	$1.689\pm 0.130$	$2.928\pm 0.078$	$1.356\pm 0.020$
LETKF	$6.161\pm 1.161$	$3.921\pm 0.329$	$1.533\pm 0.225$	$1.243\pm 0.192$	$1.753\pm 0.253$	$0.494\pm 0.051$	$3.185\pm 0.052$	$1.853\pm 0.069$	$3.513\pm 0.052$	$1.648\pm 0.022$
EnSF	$3.901\pm 0.012$	$2.340\pm 0.008$	$2.027\pm 0.051$	$1.531\pm 0.036$	$2.422\pm 0.064$	$0.976\pm 0.035$	$2.472\pm 0.031$	$1.143\pm 0.017$	$2.777\pm 0.035$	$1.242\pm 0.020$

Appendix

Table 8: Lorenz–96 RMSE and CRPS (mean

\pm

standard deviation over 10 runs). For the dense observation regime, all coordinates are observed, so only total RMSE and CRPS are reported. For the sparse regimes, total, observed-coordinate, and unobserved-coordinate RMSE are reported together with CRPS.

	1% atan				2% abs-atan
Method	All	Obs.	Unobs.	CRPS	All	Obs.	Unobs.	CRPS
FREESIA	$0.095\pm 0.005$	$0.076\pm 0.007$	$0.095\pm 0.005$	$0.052\pm 0.003$	$0.422\pm 0.101$	$0.410\pm 0.105$	$0.422\pm 0.101$	$0.209\pm 0.036$
EnKF-MDA	$0.148\pm 0.007$	$0.115\pm 0.012$	$0.148\pm 0.007$	$0.081\pm 0.005$	$1.121\pm 0.088$	$1.115\pm 0.105$	$1.121\pm 0.088$	$0.797\pm 0.070$
EnKF	$0.377\pm 0.063$	$0.372\pm 0.070$	$0.377\pm 0.063$	$0.200\pm 0.051$	$0.987\pm 0.123$	$0.984\pm 0.123$	$0.987\pm 0.123$	$0.681\pm 0.100$
LETKF	$0.218\pm 0.008$	$0.177\pm 0.012$	$0.218\pm 0.008$	$0.117\pm 0.006$	$0.970\pm 0.062$	$1.088\pm 0.072$	$0.967\pm 0.062$	$0.572\pm 0.031$
EnSF	$0.942\pm 0.060$	$0.568\pm 0.067$	$0.945\pm 0.060$	$0.547\pm 0.035$	$0.974\pm 0.053$	$0.967\pm 0.085$	$0.975\pm 0.052$	$0.569\pm 0.030$

Appendix

Table 9: Kolmogorov-flow total, observed-coordinate, and unobserved-coordinate RMSE and CRPS (mean

\pm

standard deviation over ten trajectories).

Method	L96 dense	L96 stride-2	L96 stride-4	KF $\arctan$	KF $\|\arctan\|$
FREESIA (ours)	1.377	0.478	0.251	2.704	3.397
EnKF-MDA	0.216	0.083	0.040	0.232	0.394
EnKF	0.012	0.005	0.003	0.082	0.194
LETKF	0.009	0.021	0.025	0.208	0.200
EnSF	0.343	0.301	0.286	0.369	0.385
EnFF-OT	0.021	0.022	0.020	0.105	0.102

Appendix

Table 10: Mean analysis time per assimilation step in seconds.

Method	Computation	Memory
EnSF (diagonal score)	$O(MTd)$	$O(Md)$
EnFF-F2P	$O(M^{2}Td)$	$O(Md+M^{2})$
FREESIA (structured)	$O(TM^{2}C_{\mathrm{cond}}+{}$ $TMN_{\mathrm{IS}}(C_{\mathrm{end}}+C_{\mathrm{like}}))$	$O(Md+S_{\Sigma}+m^{2})$
FREESIA-Diag	$O(MTd)$	$O(Md)$

Appendix

Table 11: Leading analysis complexity of the scaling-study implementations.

M

T

d

, and

m

denote ensemble size, generative steps, state dimension, and observation dimension, respectively. For structured FREESIA,

C_{\mathrm{cond}}

C_{\mathrm{end}}

, and

C_{\mathrm{like}}

are defined in Appendix A.10 ;

S_{\Sigma}

denotes structured-covariance storage.

Figure 16: Fully observed Lorenz–96 scaling with

M=20

. Left: time-averaged analysis RMSE. Right: mean analysis time per assimilation step. Each point averages 51 analyses over three independent runs; error bars denote one standard deviation across run means.

Data Assimilation

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

See all 58 papers

Coupling-Informed Transport Maps for Bayesian Filtering in Nonlinear Dynamical Systems
May 13, 2026Dengfei Zeng, Lijian Jiang, Shuyu Sun +1Dynamical Systems Simulation-Based Inference
School of Mathematical Sciences, Tongji University, Shanghai, 200092 China · School of Mathematical Sciences, Tongji University, Shanghai, 200092 China, Key Laboratory of Intelligent Computing and Applications (Ministry of Education), Tongji University, Shanghai, 200092 China
Flow Matching for Efficient and Scalable Data Assimilation
Aug 18, 2025Taos Transue, Bohan Chen, So Takao +1Flow Matching Kalman Filtering
Department of Mathematics and Scientific Computing and Imaging Institute, University of Utah, Salt Lake City, UT, 84112 · The Computing and Mathematical Sciences Department, California Institute of Technology, 1200 E California Blvd, Pasadena, CA 91125
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

May 13, 2026J/K move · Enter open · S save

Coupling-Informed Transport Maps for Bayesian Filtering in Nonlinear Dynamical Systems

Dengfei Zeng, Lijian Jiang, Shuyu Sun, Dunhui Xiao

School of Mathematical Sciences, Tongji University, Shanghai, 200092 China · School of Mathematical Sciences, Tongji University, Shanghai, 200092 China, Key Laboratory of Intelligent Computing and Applications (Ministry of Education), Tongji University, Shanghai, 200092 China

Dynamical Systems Simulation-Based Inference Bayesian Filtering Optimal Transport

Open paper AskPDF arXiv

A likelihood-free transport filtering method is proposed based on the couplings between state and observation variables. By exploiting a block-triangular structure in the transport map, the analysis step of filtering is reformulated as the minimization of the maximum mean discrepancy (MMD) between the true joint measure and its transport-based approximation. To circumvent the non-convexity in the MMD optimization, we introduce a training-free transport filter method via gradient flows, which leads to an analytic computation for the transport map that implies the steepest descent direction of the MMD. The proposed approach accurately approximates non-Gaussian filtering posteriors and avoids particle collapse. We provide a convergence analysis for the expectation of the MMD between the approximated posterior and the truth posterior. Finally, we extend the method to high-dimensional problems through domain localization. Numerical examples demonstrate the superior performance of our approach over conventional filtering methods in nonlinear, non-Gaussian scenarios.

Abstract

Data Assimilation

Explore similar work

Coupling-Informed Transport Maps for Bayesian Filtering in Nonlinear Dynamical Systems

Flow Matching for Efficient and Scalable Data Assimilation

DAWIS: Data Assimilation with Windowed Inverse Sampling via Multitask Interpolants