Rethinking Forward Processes for Score-Based Nonlinear Data Assimilation in High Dimensions

Eunbi Yoon; Won Chang; Donghan Kim; Dae Wook Kim

stat.MLApr 3, 2026

Rethinking Forward Processes for Score-Based Nonlinear Data Assimilation in High Dimensions

Authors: Eunbi Yoon, Won Chang, Donghan Kim, Dae Wook Kim

Organizations: KAIST · Seoul National University

Abstract

Data assimilation is the process of estimating the state of a dynamical system over time by combining model predictions with measurements. This task becomes challenging when the system is nonlinear and high-dimensional. To address this, score-based Bayesian filters have recently emerged. However, these methods still show unsatisfactory performance in certain cases, particularly under spatially sparse measurements. Such degradation stems from heuristic approximations of the likelihood score, whose errors can accumulate over time. This limitation arises because the methods simply adopt a classical forward process for generative modeling that transforms a data distribution toward a Gaussian distribution, which is independent of the measurement equation. Here, we propose a forward process tailored for filtering that transforms the system state toward the measurement space, enabling a theoretically sound formulation of the likelihood score. Based on this, we develop the Measurement-Aware Score-based Filter (MASF). We evaluate MASF on Kolmogorov flow, a high-dimensional fluid benchmark with up to $\mathcal{O}(10^5)$ dimensions, under diverse measurement operators, including nonlinear cases with a dimensional mismatch between the state and the measurements. MASF shows improved performance over existing score-based filters and ensemble-type Kalman filters. Notably, MASF achieves up to a $28.2\times$ wall-clock speedup compared with the baselines when using amortized pretraining. Our implementation is available at https://github.com/tcnllab-oss/masf.

Figures & tables

Figure 1: Overview of MASF. The measurement-aware forward process interpolates states toward the measurement space, while guided reverse-time sampling performs the measurement update.

Measurement	Dim.	Property	$H(\mathbf{x})$
Grid mask	$2(H/s)(W/s)$	Spatially sparse mask	$A_{\mathrm{grid}}\odot\mathbf{x}$
Center mask	$2(HW-h^{2})$	Central hole mask	$A_{\mathrm{center}}\odot\mathbf{x}$
Sigmoid	$2HW$	Element-wise nonlinear	$\operatorname{sigmoid}(\mathbf{x})$
Speed	$HW$	Channel-coupled nonlinear	$[H(\mathbf{x})]_{ij}=\operatorname{sigmoid}\!\left(\sum_{c=1}^{2}x_{cij}^{2}\right)$

Table 1: Measurement settings. Dim. denotes the effective dimension of all measurements.

Type	Measurement	Method	Ensemble	RMSE $(\downarrow)$	CSI $(\uparrow)$	Wall-clock (s)	Speedup
Linear	Grid mask	EnKF	250	$0.18\pm 0.04$	$0.76\pm 0.05$	$180.3\pm 0.4$	$1.5\times$
		LETKF	40	$0.18\pm 0.01$	$0.74\pm 0.02$	$389.6\pm 5.4$	$3.3\times$
		SF	100	$0.81\pm 0.04$	$0.00\pm 0.00$	$670.6\pm 0.9$	$5.6\times$
		SSLS	100	$0.47\pm 0.07$	$0.34\pm 0.10$	$611.5\pm 1.1$	$5.1\times$
		MASF	10	$0.15\pm 0.02$	$0.78\pm 0.03$	$118.8\pm 2.5$	–
	Center mask	EnKF	250	$0.75\pm 0.17$	$0.32\pm 0.08$	$185.4\pm 2.5$	$1.1\times$

Table 2: Main results on Kolmogorov flow at

128^{2}

resolution. Mean

\pm

standard deviation over five seeds. Wall-clock time is reported in seconds, and speedup is computed relative to MASF within each measurement setting.

Figure 2: Qualitative results. Rows correspond to different measurements, and columns show the ground truth, measurement, and estimates from EnKF, LETKF, SF, SSLS, and MASF.

Resolution	Method	Ensemble	RMSE $(\downarrow)$	CSI $(\uparrow)$	Wall-clock (s)	Speedup
$256^{2}$	EnKF	250	$0.17\pm 0.03$	$0.76\pm 0.03$	$1276.3\pm 5.8$	$3.7\times$
	LETKF	40	$0.19\pm 0.01$	$0.70\pm 0.01$	$1646.7\pm 15.2$	$4.8\times$
	SF	100	$0.88\pm 0.03$	$0.00\pm 0.00$	$1056.2\pm 4.3$	$3.1\times$
	SSLS	100	$0.72\pm 0.06$	$0.03\pm 0.04$	$1016.7\pm 8.8$	$3.0\times$
	MASF	10	$0.16\pm 0.02$	$0.77\pm 0.03$	$343.2\pm 4.0$	-
$512^{2}$	EnKF	250	$0.18\pm 0.07$	$0.78\pm 0.05$	$11162.5\pm 80.3$	$10.6\times$

Table 3: High-dimensional scaling on grid-mask measurements. Mean

\pm

standard deviation over five seeds. Wall-clock time is reported in seconds, and speedup is computed relative to MASF at the same resolution.

Figure 3: Runtime–accuracy and sensitivity analysis. Comparison under grid-mask measurements across runtime budget, temporal length, temporal gap, and spatial sparsity. Shaded regions denote standard deviation over five seeds.

Symbols	Description	Defined in
$\tau$	The physical time variable of the dynamical system.	Eq. ( 1 )
$t\in[0,1]$	The diffusion time variable of the measurement-aware forward process.	Sec. 3.1
$X_{k}$	The state at the measurement time $\tau_{k}$ .	Sec. 2.1
$Z_{k}$	The measurement at the measurement time $\tau_{k}$ .	Eq. ( 2 )
$\mathbf{z}_{1:k}$	The observed measurement sequence $(\mathbf{z}_{1},\ldots,\mathbf{z}_{k})$ .	Sec. 2.1
$X_{t}$	The measurement-aware forward process.	Eq. ( 12 )

Table 4: Summary of notation.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Resource	Specification
Operating system	Linux, kernel 5.14.0
CPU	Intel Xeon 6530P, 2 sockets
CPU capacity	64 physical cores, 128 logical CPUs total
CPU allocation	8 logical CPUs per Slurm task
System memory	250 GiB RAM, 63 GiB swap
GPU	NVIDIA RTX PRO 6000 Blackwell Server Edition

Appendix

Table 5: Compute resources and software environment used for the experiments.

Asset	Use	Version / source	License / terms
JAX-CFD	Kolmogorov flow simulation	0.2.1	Apache-2.0
JAX	Numerical computation	0.4.33	Apache-2.0
PyTorch	Neural network training	2.9.0	BSD-style license
NumPy	Numerical computation	2.1.1	BSD 3-Clause
SciPy	Numerical computation	1.14.1	BSD 3-Clause
Optuna	Hyperparameter tuning	Cited package	MIT license

Appendix

Table 6: Existing assets used in the experiments.

Parameter	Search values
same normalization	{ true , false }
normalization form	{ affine , scale_only }
statistics mode	{ standard , robust , minmax }
statistics update	{ moving , fixed , adaptive }
momentum	{0.2, 0.4}

Appendix

Table 7: Normalization search space for MASF.

Parameter	Kolmogorov- $128$	Kolmogorov- $256$	Kolmogorov- $512$
base channels	{16, 32}	{16, 32}	{16, 32}
residual blocks	{1, 2, 3}	{1, 2, 3}	{1, 2, 3}
channel multiplier	{ 1-2-4 , 1-2-4-8 }	{ 1-2-4-8 , 1-1-2-3-4 }	{ 1-2-4-8 , 1-1-2-3-4 }
attention resolution	{8, 16}	{8, 16, 32}	{8, 16, 32}

Appendix

Table 8: Model architecture search spaces for MASF on Kolmogorov datasets.

Parameter	Search values
pretraining samples	{1000}
batch size	{16, 32, 64}
epochs	{100, 200, 300, 400}
learning rate	{ $0.0001$ , $0.0003$ , $0.0005$ }

Appendix

Table 9: Pretraining search space for MASF.

Parameter	Search values
NFE	{150, 200, 250}
prior scale min. $\lambda_{s,\min}$	{0.9, 0.95, 1.0}
prior scale max. $\lambda_{s,\max}$	{1.4, 1.5, 1.6}
prior scale power $p_{s}$	{1.0, 1.25, 1.5}
likelihood scale min. $\lambda_{g,\min}$	{1.0, 1.1, 1.2}
likelihood scale max. $\lambda_{g,\max}$	{1.4, 1.5, 1.6}

Appendix

Table 10: Sampling and guidance search space for MASF.

Parameter	Search values
batch size	{10}
learning rate	{ $3\times 10^{-4}$ , $1\times 10^{-4}$ }
full epochs	{80, 90, 100, 110, 120}

Appendix

Table 11: Online fine-tuning search space for MASF.

Category	Field	Grid mask $128^{2}$	Grid mask $256^{2}$	Grid mask $512^{2}$
measurement	type	grid_mask	grid_mask	grid_mask
	same normalization	true	false	false
	normalization form	affine	affine	scale_only
	statistics mode	standard	standard	standard
	statistics update	moving	moving	adaptive
	main parameter	stride $=10$	stride $=15$	stride $=20$

Appendix

Table 12: Selected MASF configurations for grid-mask measurements across resolutions.

Category	Field	Center mask	Sigmoid	Speed
measurement	type	center_mask	sigmoid	speed
	same normalization	false	false	false
	normalization form	scale_only	affine	scale_only
	statistics mode	standard	standard	robust
	statistics update	fixed	fixed	adaptive
	main parameter	hole ratio $=0.5$	$-$	$-$

Appendix

Table 13: Selected MASF configurations for the remaining

128^{2}

measurement settings.

Parameter	Search values
ensemble size	{100, 200, 250, 300, 400}
update mode	{ full }
inflation	{1.0, 1.05, 1.1, 1.15, 1.2, 1.25, 1.3}
epsilon	{ $10^{-5}$ , $10^{-6}$ }

Appendix

Table 14: EnKF hyperparameter search space.

Parameter	Search values
ensemble size	{10, 20, 30, 40, 60, 100}
inflation	{1.0, 1.05, 1.1, 1.15, 1.2, 1.25, 1.3}
localization radius	{3, 5, 7}
epsilon	{ $10^{-5}$ , $10^{-6}$ }

Appendix

Table 15: LETKF hyperparameter search space.

Method	Ensemble size	RMSE ( $\downarrow$ )	Wall-clock (s)
EnKF	$40$	$1.011\pm 0.069$	$42.9\pm 2.2$
	$100$	$0.540\pm 0.192$	$70.6\pm 1.6$
	$200$	$0.295\pm 0.130$	$139.1\pm 0.6$
	$\mathbf{250}$	$\mathbf{0.176\pm 0.039}$	$\mathbf{179.5\pm 2.5}$
	$400$	$0.141\pm 0.030$	$263.5\pm 9.8$
LETKF	$10$	$0.222\pm 0.023$	$273.1\pm 4.0$

Appendix

Table 16: Ensemble-size sensitivity for EnKF and LETKF under grid-mask measurements. The selected default ensemble sizes are highlighted in bold.

Method	Parameter	Search values / selected value
SF, SSLS	NFE	500
SSLS	perturbation std.	{0.001, 0.05, 0.1, 0.2}
SSLS	stepsize	{0.0001, 0.001, 0.002, 0.003, 0.004}
SF, SSLS	tolerance	{10, 50, 100, 200}

Appendix

Table 17: Score-based baseline sampling settings.

Parameter	From	To	RMSE ( $\downarrow$ )
Normalization
current_config	–	–	$\mathbf{0.151\pm 0.019}$
same_normalization	true	false	$0.229\pm 0.017$
normalization_form	affine	scale_only	$0.180\pm 0.013$
stats_mode	standard	robust	$0.157\pm 0.020$
stats_update_mode	moving	fixed	$\mathbf{0.151\pm 0.018}$

Appendix

Table 18: Component ablations. Each block reports the current configuration for that phase and one-parameter changes from it.

Measurement	Method	Strategy	Ensemble	RMSE ( $\downarrow$ )	CSI ( $\uparrow$ )	Wall-clock (s)
Grid mask	SF	Scratch	100	$0.81\pm 0.05$	$0.00\pm 0.00$	$670.6\pm 0.9$
		Scratch	10	$0.87\pm 0.03$	$0.00\pm 0.00$	$100.9\pm 0.7$
		Pretraining	10	$0.87\pm 0.05$	$0.00\pm 0.00$	$81.9\pm 0.6$
	SSLS	Scratch	100	$0.47\pm 0.07$	$0.34\pm 0.10$	$611.5\pm 1.1$
		Scratch	10	$0.58\pm 0.10$	$0.26\pm 0.09$	$89.2\pm 0.6$
		Pretraining	10	$0.39\pm 0.04$	$0.44\pm 0.01$	$81.3\pm 12.0$

Appendix

Table 19: Effect of pretraining. We compare scratch training and pretrained initialization for score-based filtering methods under grid-mask and speed measurements. Wall-clock time excludes offline pretraining and measures only the online data assimilation procedure.

Figure 4: Grid-mask measurements. Qualitative comparison of measurements and posterior samples under grid-mask measurements at selected time steps.

Figure 5: Center-mask measurements. Qualitative comparison of measurements and posterior samples under center-mask measurements at selected time steps.

Figure 6: Sigmoid measurements. Qualitative comparison of measurements and posterior samples under sigmoid measurements at selected time steps.

Figure 7: Speed measurements. Qualitative comparison of measurements and posterior samples under speed measurements at selected time steps.

Figure 8: Ensemble-size sensitivity. Comparison of RMSE, wall-clock time, and runtime–accuracy trade-off as functions of ensemble size for the considered methods.

Figure 9: Extended sensitivity analysis. Sensitivity to temporal length and temporal gap under grid-mask measurements.

Data Assimilation

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

See all 58 papers

AECSF: Adaptive Ensemble Conditional Score Filtering for High-Dimensional Nonlinear Data Assimilation
Sep 26, 2026Yangwen Zhang, Shiwei Ni, Xiaoping Zhang +2Data Assimilation Bayesian Filtering
School of Mathematical Sciences, Tongji University, Shanghai 200092, China · Department of Information and Computational Sciences, Wuhan University, Wuhan 430072, China · Department of Mathematics, University of South Carolina, Columbia, South Carolina 29208, USA
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
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

Sep 26, 2026J/K move · Enter open · S save

AECSF: Adaptive Ensemble Conditional Score Filtering for High-Dimensional Nonlinear Data Assimilation

Yangwen Zhang, Shiwei Ni, Xiaoping Zhang, Xiaofei Guan+1

School of Mathematical Sciences, Tongji University, Shanghai 200092, China · Department of Information and Computational Sciences, Wuhan University, Wuhan 430072, China · Department of Mathematics, University of South Carolina, Columbia, South Carolina 29208, USA

Data Assimilation Bayesian Filtering Score-Based Generative Modeling Adaptive Sampling Conditional Distribution Estimation

Open paper AskPDF arXiv

Bayesian state estimation for high-dimensional nonlinear dynamical systems entails a fundamental tension between statistical fidelity and computational tractability, as particle weights can collapse, while Gaussian ensemble updates can miss non-Gaussian posterior structure. Score-based diffusion filters offer a sampling-based alternative, but existing training-free score filters often rely on heuristic likelihood corrections, which can compromise posterior accuracy by neglecting uncertainty about the system state associated with each noisy reverse particle. To address these issues, we propose AECSF, a training-free adaptive ensemble conditional score filter. AECSF constructs an analytically tractable score estimator from the conditional Tweedie identity, which recasts noisy posterior score estimation as estimating the conditional mean of the system state given a noisy reverse particle and the observation. To estimate these conditional means efficiently, AECSF employs a shared adaptive weighted proposal ensemble, while particle-specific conditional weights yield an estimate for each noisy reverse particle without separate proposal sampling. The proposal ensemble is updated using reverse-particle information within the same reverse-diffusion run to improve conditional-mean estimation. Theoretically, we characterize when a fixed weighted proposal measure yields the exact noisy posterior score. Under stated assumptions, we establish a bound relating conditional-mean estimation errors to reverse-sampling endpoint error. Numerical experiments demonstrate that AECSF improves the accuracy of posterior sampling and nonlinear filtering in high-dimensional problems with limited forecast ensembles.

Abstract

Data Assimilation

Explore similar work

AECSF: Adaptive Ensemble Conditional Score Filtering for High-Dimensional Nonlinear Data Assimilation

Flow Matching for Efficient and Scalable Data Assimilation

Amortized Filtering and Smoothing with Conditional Normalizing Flows