D-Flow SGLD: Source-Space Posterior Sampling for Scientific Inverse Problems with Flow Matching

Meet Hemant Parikh; Yaqin Chen; Jian-Xun Wang

cs.LGFeb 25, 2026

D-Flow SGLD: Source-Space Posterior Sampling for Scientific Inverse Problems with Flow Matching

Authors: Meet Hemant Parikh, Yaqin Chen, Jian-Xun Wang

Organizations: Sibley School of Mechanical and Aerospace Engineering, Cornell University, Ithaca, NY, USA · Department of Aerospace and Mechanical Engineering, University of Notre Dame, Notre Dame, IN, USA

Abstract

Data assimilation and scientific inverse problems require reconstructing high-dimensional physical states from sparse and noisy observations, ideally with uncertainty-aware posterior samples that remain faithful to learned priors and governing physics. While training-free conditional generation is well developed for diffusion models, corresponding conditioning and posterior sampling strategies for Flow Matching (FM) priors remain comparatively under-explored, especially on scientific benchmarks where fidelity must be assessed beyond measurement misfit. In this work, we study training-free conditional generation for scientific inverse problems under FM priors and organize existing inference-time strategies by where measurement information is injected: (i) guided transport dynamics that perturb sampling trajectories using likelihood information, and (ii) source-distribution inference that performs posterior inference over the source variable while keeping the learned transport fixed. Building on the latter, we propose D-Flow SGLD, a source-space posterior sampling method that augments differentiable source inference with preconditioned stochastic gradient Langevin dynamics, enabling scalable exploration of the source posterior induced by new measurement operators without retraining the prior or modifying the learned FM dynamics. We benchmark representative methods from both families on a hierarchy of problems: 2D toy posteriors, chaotic Kuramoto-Sivashinsky trajectories, and wall-bounded turbulence reconstruction. Across these settings, we quantify trade-offs among measurement assimilation, posterior diversity, and physics/statistics fidelity, and establish D-Flow SGLD as a practical FM-compatible posterior sampler for scientific inverse problems.

Figures & tables

Figure 1: S-curve toy inverse problem. Conditional posterior samples ( $N=1000$ , orange) overlaid on unconditional OT-CFM samples (gray). For velocity-field guidance, the correction strength is $b=3$ .

Figure 2: Moon-curve toy inverse problem. Conditional posterior samples ( $N=1000$ , orange) overlaid on unconditional OT-CFM samples (gray). For velocity-field guidance, the correction strength is $b=3$

	S-Curve		Two-moon
Method	$\mathcal{F}_{1}(\cdot)$	$\mathcal{F}_{2}(\cdot)$	$\mathcal{F}_{1}(\cdot)$	$\mathcal{F}_{2}(\cdot)$
D-Flow	0.189	0.449	0.056	0.251
D-Flow SGLD (Ours)	0.234	0.171	0.073	0.160
Grad (b=3)	0.357	0.625	0.129	0.603
Grad-Free (b=3)	0.442	0.596	0.132	0.559

Table 1: 1-Wasserstein distance

\mathcal{W}_{1}

between the empirical reference posterior (importance resampling) and the empirical conditional distribution for the toy inverse problems (lower is better). Best and second-best results are highlighted in bold and underlined, respectively.

Figure 3: KS temporal forecasting with noise-free measurements. (top row) The ground truth measurement is shown in the first column, followed by a conditionally generated sample from each method. The region below the white line ( $t>r$ ) is the forecasting portion. (middle row) Pointwise data assimilation error (MAE). (bottom row) The PDE residual, $R_{j}^{n}$ , computed over the entire generated field to assess physical plausibility.

Figure 4: KS temporal forecasting with 10% noise corruption. (top row) The ground truth measurement is shown in the first column, followed by a conditionally generated sample from each method. The region below the white line ( $t>r$ ) is the forecasting portion. (middle row) Pointwise data assimilation error (MAE). (bottom row) The PDE residual, $R_{j}^{n}$ , computed over the entire generated field to assess physical plausibility.

Figure 5: Mean absolute PDE residual, $\overline{R}$ , for the different guidance methods applied in KS system. The results are shown for two inverse problems: temporal forecasting (a-b) and reconstruction from sparse sensors (c-d). Each problem is tested with both noise-free measurements (a, c) and measurements corrupted by $10\%$ noise (b, d). Bar heights represent the mean residual, and error bars denote the standard deviation across generated samples.

Figure 6: Turbulence inpainting ( $u_{1}$ ) with noise-free measurements. Top: the observed field (first column) and one representative conditional sample from each method; the hatched region ( $x>2\pi$ ) indicates the missing/inpainted portion. Bottom: pointwise MAE on the observed region.

Figure 7: Turbulence inpainting ( $u_{1}$ ) with noisy measurements ( $10\%$ ). Top: the observed field (first column) and one representative conditional sample from each method; the hatched region ( $x>2\pi$ ) indicates the missing/inpainted portion. Bottom: pointwise MAE on the observed region.

Panel of Figure 8 — Figure 8: Comparison of the streamwise (right) and spanwise (left) one-dimensional energy spectrum for the streamwise velocity component $u_{1}$ at wall-normal location of $y^{+}=40$ .

Figure 9: $1000$ source distribution samples obtained after performing optimization using the D-Flow algorithm for S-curve toy problem for the measurement functions $\mathcal{F}_{1}(\cdot)$ (left) and $\mathcal{F}_{2}(\cdot)$ (right) respectively. The gray dots represent the samples for the unconditional Gaussian distribution, while the orange samples represents the samples generated after optimizing the source terms using the D-Flow Algorithm.

Figure 10: $1000$ source distribution samples obtained after performing optimization using the D-Flow SGLD algorithm for S-curve toy problem for the measurement functions $\mathcal{F}_{1}(\cdot)$ (left) and $\mathcal{F}_{2}(\cdot)$ (right) respectively. The gray dots represent the samples for the unconditional Gaussian distribution. The orange dots represents the samples generated after optimization using the D-Flow Algorithm. The cyan dots represents the samples generated with source term regularization using the D-Flow SGLD Algorithm.

Case	D-Flow	D-Flow SGLD	Grad	Grad-Free
Toy ( $N=1000$ )	65	8	8	5
KS ( $N=100$ )	$12000^{*}$	578	1114	363
Turb ( $N=1000$ )	$200500^{\dagger}$	928	1841	840

Table 2: Wall-clock inference time (seconds) for training-free conditional generation methods on a single NVIDIA L40 GPU for the toy case and a single NVIDIA A100 80GB GPU for the rest of the cases.

N

denotes the number of posterior samples generated.

\*

and

\dagger

denotes extrapolation of the cost by generating 20 and 1 sample(s) respectively.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Case	D-Flow	D-Flow SGLD	Grad/Grad-Free
Toy	$N_{optim}={10\ \text{(S-curve)},\ 20\ \text{(Two-moon)}}$ ODE solver = midpoint, 6 steps optimizer = LBFGS $\alpha=0$ , $\beta=0$	$N_{parallel}=10$ $N_{steps}=500$ , burn= $100$ $\eta(i)=5e^{-2}$ $s(i)=1e^{-2}$ , $\lambda=0.1$ (S-curve), $0.05$ (Two-moon) $\omega=0.99$ $\delta(i)=1e^{-3}$ ODE solver = midpoint, 6 steps $\alpha=0$ , $\beta=0$	$\lambda(\tau)=1$ $b=\{1,3,10\}$ ODE solver = euler, 300 steps
KS	$N_{optim}=10$ ODE solver = midpoint, 6 steps optimizer = LBFGS $\alpha=0$ , $\beta=0$	$N_{parallel}=2$ $N_{steps}=300$ , burn= $100$ $\eta(i)=1e^{-2}$ $s(i)=1e^{-3}$ , $\lambda=1e^{-3}$ $\omega=0.99$ $\delta(i)=1e^{-3}$ ODE solver = midpoint, 6 steps $\alpha=0$ , $\beta=0$	$\lambda(\tau)=1$ $b=1$ ODE solver = euler, 300 steps
Turb	$N_{optim}=100$ ODE solver = midpoint, 6 steps optimizer = LBFGS $\alpha=0$ , $\beta=0$	$N_{parallel}=20$ $N_{steps}=600$ , burn= $100$ $\eta(i)=5e^{-2}$ $s(i)=1e^{-3}$ , $\lambda=5e^{-6}$ $\omega=0.99$ $\delta(i)=1e^{-3}$ ODE solver = midpoint, 6 steps $\alpha=0$ , $\beta=0$	$\lambda(\tau)=1$ $b=1$ ODE solver = euler, 300 steps

Appendix

Table 3: Hyperparameters for Training-Free Guidance Methods Across Different Cases.

Figure 11: Effect of the guidance scale $b$ on conditional samples for the S-curve toy problem with measurement operator $\mathcal{F}_{1}(\cdot)$ measurement function

Figure 12: KS sparse sensor reconstruction with noise-free measurements. (top row) The ground truth measurement is shown in the first column, followed by a conditionally generated sample from each method. The region below the white line ( $t>r$ ) is the forecasting portion. (middle row) Pointwise data assimilation error (MAE). (bottom row) The PDE residual, $R_{j}^{n}$ , computed over the entire generated field to assess physical plausibility.

Figure 13: KS sparse sensor reconstruction with noisy measurements ( $10\%$ ). (top row) The ground truth measurement is shown in the first column, followed by a conditionally generated sample from each method. The region below the white line ( $t>r$ ) is the forecasting portion. (middle row) Pointwise data assimilation error (MAE). (bottom row) The PDE residual, $R_{j}^{n}$ , computed over the entire generated field to assess physical plausibility.

Figure 14: Turbulence inpainting ( $u_{2}$ ) with noise-free measurements. Top: the observed field (first column) and one representative conditional sample from each method; the hatched region ( $x>2\pi$ ) indicates the missing/inpainted portion. Bottom: pointwise MAE on the observed region.

Figure 15: Turbulence inpainting ( $u_{3}$ ) with noise-free measurements. Top: the observed field (first column) and one representative conditional sample from each method; the hatched region ( $x>2\pi$ ) indicates the missing/inpainted portion. Bottom: pointwise MAE on the observed region.

Figure 16: Turbulence inpainting ( $u_{2}$ ) with noisy measurements ( $10\%$ ). Top: the observed field (first column) and one representative conditional sample from each method; the hatched region ( $x>2\pi$ ) indicates the missing/inpainted portion. Bottom: pointwise MAE on the observed region.

Figure 17: Turbulence inpainting ( $u_{3}$ ) with noisy measurements ( $10\%$ ). Top: the observed field (first column) and one representative conditional sample from each method; the hatched region ( $x>2\pi$ ) indicates the missing/inpainted portion. Bottom: pointwise MAE on the observed region.

Panel of Figure 18 — Figure 18: Comparison of the streamwise (right) and spanwise (left) one-dimensional energy spectrum for the wall-normal velocity component $u_{2}$ at wall-normal location of $y^{+}=40$ .

Panel of Figure 19 — Figure 19: Comparison of the streamwise (right) and spanwise (left) one-dimensional energy spectrum for the spanwise velocity component $u_{3}$ at wall-normal location of $y^{+}=40$ .

Posterior Sampling

Posterior sampling generates draws from Bayesian posteriors to quantify uncertainty and support inference, prediction, and decision-making.

See all 69 papers

Sep 17, 2026

FlowSGS: Improving Flow Priors for Inverse Imaging with Stochastic Interpolants

Flow matching has emerged as the state-of-the-art generative model and has been used for plug-and-play (PnP) priors to solve inverse problems in computational imaging. However, existing flow-based inverse solvers make simplifying approximations in posterior sampling and/or assume linear forward models. In contrast, principled diffusion-based methods exist but are computationally expensive. To circumvent these problems, we introduce FlowSGS, a flow-based posterior sampling method using Split Gibbs Sampling (SGS) with annealed coupling, and make flow-specific interventions to the prior sampling and annealing steps, theoretically grounded in the Stochastic Interpolants (SI) framework. Using a novel timestep correction and accelerated warm-up with one-step flow models, our method is faster and more accurate than the state of the art across a range of imaging problems. We demonstrate the first use of flow-based samplers for a nonlinear inverse problem (Fourier phase retrieval) and apply one-step flow models to principled Bayesian inverse imaging.

Tianao Li, Xinhui Qian, Emma Alexander

Department of Computer Science, Northwestern University · NSF-Simons AI Institute for the Sky (SkAI) · Department of Statistics and Data Science, Northwestern University

Flow Matching Markov Chain Monte Carlo Posterior Sampling Flow-Based Generative Modeling Image Inverse Problems

Ask PDF

Oct 1, 2026

Flow Matching for Fast Posterior Sampling in Bayesian Inverse Problems

Sampling from the posterior is the central task of computational Bayesian inverse problems. The standard workhorse in Bayesian inference - Markov chain Monte Carlo (MCMC) - is sequential, yields correlated samples, and must be rerun for each observation. Conditional flow matching offers an amortized alternative: a transport map, trained once on joint samples of parameter and data, that yields independent approximate posterior samples for any observation at negligible online cost, without new likelihood evaluations. We give a careful, MCMC-literate assessment of flow matching for PDE-based inverse problems with function-valued parameters. Exploiting the flow's tractable density, we derive computable accuracy estimates of the underlying approximate posterior in total-variation distance and Kullback-Leibler divergence and, moreover, propose a hybrid sampler that is asymptotically exact by Metropolization. We validate the accuracy estimates and demonstrate the amortization in several numerical examples, including electrical impedance tomography and a likelihood-free Lotka-Volterra model.

Jan Blechschmidt, Oliver G. Ernst, Moritz Poguntke +1

Posterior Sampling Bayesian Inverse Problems Amortized Inference Conditional Flow Matching PDE Inverse Problems

Ask PDF

Oct 9, 2026

Spanning Diffusion and Flow-based Posterior Sampling via Exact Gaussian Conditioning

Flow-matching and diffusion models are increasingly used as priors for linear inverse problems, however, most methods are bespoke to one class of trajectory prior. Plug-and-play methods have been proposed, however these implicitly decorrelate the guided trajectory, resulting in broad priors that can struggle when data is sparse. This work proposes

β

-Flow: a general plug-and-play method indexed by a correlation parameter

β

that generalises over flow, diffusion, and decorrelated priors, and a confidence parameter

κ

that controls how strongly the path distribution is enforced. Each step requires only one forward evaluation of a denoising network, followed by a Gaussian conditioning step. This step is an exact Kalman update that is not dependent on network gradients. We show that many existing guided samplers can be interpreted as instances or approximations of our method, each implicitly using a different trajectory prior and operator in place of the Kalman gain. From this viewpoint, we provide a qualitative comparison of methods in terms of their implicit biases and assumptions. We demonstrate consistently state-of-the-art level performance across 7 inverse problems over 4 datasets in terms of PSNR and SSIM. Moreover, we demonstrate the utility of this prior control, by showing that confident high correlations (strict priors) work best when observations are less informative, while low correlation is preferred when observations are more informative, reflecting the complementary roles of the prior and likelihood in Bayesian inference.

Teo Deveney, Lea Bogensperger, Hok Shing Wong

University of Bath · University of Zurich · LMU Munich

Posterior Sampling Diffusion Posterior Sampling Flow-Based Generative Modeling Bayesian Inverse Problems

Ask PDF

Abstract

Posterior Sampling

Explore similar work

FlowSGS: Improving Flow Priors for Inverse Imaging with Stochastic Interpolants

Flow Matching for Fast Posterior Sampling in Bayesian Inverse Problems

Spanning Diffusion and Flow-based Posterior Sampling via Exact Gaussian Conditioning