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

Teo Deveney; Lea Bogensperger; Hok Shing Wong

cs.LGOct 9, 2026

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

Authors: Teo Deveney, Lea Bogensperger, Hok Shing Wong

Organizations: University of Bath · University of Zurich · LMU Munich

Abstract

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.

Figures & tables

Figure 1: Three guided trajectories on a toy denoising example for varying $\beta_{t}$ . Each line is an independent trajectory, and the black marker is the observation. At $\beta_{t}=-1$ the trajectory is anti-correlated which results in the closest convergence to the observation, at $\beta_{t}=0$ every step re-noises independently, and as $\beta_{t}$ increases consecutive states become correlated, until at $\beta_{t}=1$ they follow a deterministic flow.

$\beta_{t_{i+1}}$	$\tau^{2}$	Coincides with
$0$	$\sigma_{t_{i}}^{2}$	Independently noised $x_{0}$
$\beta^{\star}_{t_{i+1}}$	$\sigma_{t_{i}}^{2}(1-\beta^{\star 2}_{t_{i+1}})$	Reverse diffusion SDE
$1$	$0$	Flow-matching ODE

Table 1: Common cases of the kernel ( 5 ). The diffusion schedule (

\beta^{\star}_{t_{i+1}}:=\sigma_{t_{i}}\alpha_{t_{i+1}}/(\sigma_{t_{i+1}}\alpha_{t_{i}})

) gives a Markovian forward process. Every schedule with

\beta_{t}\in[-1,1]

defines a valid kernel with the correct marginals.

Panel of Figure 2 — Figure 2: Graphical model for one reverse step. The state $\tilde{x}_{t_{i+1}}$ held by the sampler depends on $x_{0}$ and is correlated, with coefficient $\beta_{t_{i+1}}$ , with the reference point $x_{t_{i+1}}$ on the interpolant path (grey, dashed). It is conditionally independent of $x_{t_{i+1}}$ given $x_{0}$ when $\beta_{t_{i+1}}=0$ and coincides with it when $\beta_{t_{i+1}}=1$ . Guided sampling proceeds from $x_{1}$ to $x_{0}$ by iteratively modelling the distribution of $x_{t_{i}}$ given $\tilde{x}_{t_{i+1}}$ and $y$ .

Method	Transport $\mathcal{T}$	Point estimate	Gain $\mathcal{G}$ on the residual	Gradients through $f_{\theta}$
This paper	kernel, $\beta_{t}\in[-1,1]$	one step	$G\big(\kappa(\sigma_{t_{i+1}}/\alpha_{t_{i+1}})^{2}\big)$	none
FLOWER	kernel, $\beta_{t}=0$	one step	$G\big(\sigma_{t_{i+1}}^{2}/(\alpha_{t_{i+1}}^{2}+\sigma_{t_{i+1}}^{2})\big)$	none
PnP-Flow	kernel, $\beta_{t}=0$	one step	$\rho\,H^{\top}$	none
DPS	kernel, $\beta^{\star}_{t}$	one step	$\lambda_{i}\,JH^{\top}$	every step
$\Pi$ GDM	kernel, $\beta^{\star}_{t}$	one step	$\propto J^{\top}H^{\top}(r_{t}^{2}HH^{\top}+\sigma_{y}^{2}I)^{-1}$	every step
D-Flow	flow, $F^{-1}$	integrated, $F(x_{1}^{(k)})$	$2\eta\,\bar{\mathcal{J}}\mathcal{J}^{\top}H^{\top}$	whole trajectory

Table 2: Every method’s update written in the form of ( 13 ), with

r=y-H\hat{x}_{0}

and

G(S)=SH^{\top}(HSH^{\top}+\sigma_{y}^{2}I)^{-1}

. For D-Flow,

\bar{\mathcal{J}}

is the average of

\mathcal{J}

along the step. Appendix A , A.2, derives the DPS and D-Flow entries.

	Denoising			Deblurring			Super-resolution			Random inpainting			Box inpainting
Method	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS
Degraded	20.00	0.314	0.509	23.94	0.517	0.444	11.70	0.208	0.873	13.36	0.223	1.081	21.80	0.740	0.198
PnP-GS	32.58	0.894	0.072	27.91	0.753	0.349	24.15	0.632	0.362	29.42	0.836	0.126	–	–	–
DiffPIR	30.58	0.835	0.189	27.56	0.728	0.342	23.65	0.624	0.402	31.70	0.881	0.062	–	–	–
OT-ODE	30.03	0.815	0.076	27.06	0.713	0.123	25.91	0.716	0.108	29.40	0.839	0.090	24.62	0.875	0.085
D-Flow	26.13	0.574	0.175	27.82	0.721	0.164	24.64	0.601	0.190	32.20	0.894	0.040	26.26	0.842	0.077

Table 3: AFHQ-Cat performance. Most results were using 100 steps, with the exception of some pre-existing PnP-Flow/Flower baselines which used

500^{*}

and

200^{\dagger}

steps in certain tasks.

	Denoising			Deblurring			Super-resolution			Random inpainting			Box inpainting
Method	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS
Degraded	20.00	0.348	0.372	27.83	0.740	0.126	10.26	0.183	0.827	11.95	0.196	1.041	22.27	0.742	0.214
PnP-GS	32.64	0.910	0.035	34.03	0.924	0.041	31.31	0.892	0.064	29.22	0.875	0.070	–	–	–
DiffPIR	31.20	0.885	0.060	32.77	0.912	0.060	31.52	0.895	0.033	31.74	0.917	0.025	–	–	–
OT-ODE	30.54	0.859	0.032	33.01	0.921	0.029	31.46	0.907	0.025	28.68	0.871	0.051	29.40	0.920	0.038
D-Flow	26.04	0.607	0.092	31.25	0.854	0.038	30.47	0.843	0.026	33.67	0.943	0.015	30.70	0.899	0.026

Table 4: Results on 100 test images of the dataset CelebA.

Figure 3: Qualitative reconstructions across datasets and problems. All methods used 100 steps (100 updates in the case of D-Flow), and hyperparameters tuned for each dataset and task. Columns 2-4 represent reconstructions from each class of model. In Column 2 the likelihood score methods OT-ODE/ $\Pi$ GDM are applied to the upper/lower two rows respectively. These excel at producing realistic looking fine detail that does not necessarily match the ground truth. In Column 3 is the ODE rollout method D-Flow, which also produces some fine details, but appears sensitive to the degradation process as can be seen in the noise artifacts in row 2 and vertical smearing in row 3. Column 4 shows the independent-renoising method Flower, which achieves excellent accuracy when measurements are informative, but struggles at 100 steps with the sparser 4x super-resolution problem. In Column 5 $\beta$ -Flow achieves excellent accuracy across all problems by adapting the prior to the task.

	Undersampled MRI			Sparse-view CT
Method	PSNR	SSIM	LPIPS	PSNR	SSIM	LPIPS
DPS	36.19	0.926	0.037	28.79	0.740	0.107
$\Pi$ GDM	36.12	0.924	0.039	35.40	0.838	0.154
D-Flow	34.19	0.871	0.158	27.42	0.740	0.157
PnP-Flow	36.30	0.907	0.094	25.70	0.695	0.186
Flower	37.70	0.942	0.091	33.41	0.905	0.076

Table 5: FastMRI and AAPM CT datasets results.

	$\times 2$		$\times 3$		$\times 4$		$\times 6$
$\beta_{t}$	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
$-0.6$	$\mathbf{33.08}$	$\mathbf{.909}$	$29.87$	$.838$	$27.14$	$.762$	$22.00$	$.609$
$-0.3$	$33.07$	$.908$	$29.93$	$.839$	$27.36$	$.768$	$22.40$	$.620$
$0.0$	$33.04$	$.907$	$29.99$	$\mathbf{.840}$	$27.60$	$.772$	$22.84$	$.632$
$0.3$	$32.96$	$.904$	$\mathbf{30.01}$	$.839$	$27.76$	$.776$	$23.40$	$.645$
$0.6$	$32.76$	$.896$	$29.96$	$.836$	$\mathbf{27.85}$	$\mathbf{.777}$	$\mathbf{23.93}$	$\mathbf{.657}$

Table 6: PSNR (dB) and SSIM on AFHQ-Cat super-resolution across the correlation

\beta_{t}

, with

\kappa=1

fixed. The optimal

\beta_{t}

increases with the downsampling factor, showing reconstructions favour more strongly correlated trajectories at higher scale factors.

	$\times 2$		$\times 3$		$\times 4$		$\times 6$
$\kappa$	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
$0.1$	$32.09$	$.882$	$30.15$	$.845$	$\mathbf{28.40}$	$\mathbf{.798}$	$\mathbf{25.57}$	$\mathbf{.712}$
$0.5$	$33.00$	$.907$	$\mathbf{30.19}$	$\mathbf{.846}$	$28.22$	$.790$	$25.12$	$.696$
$1.0$	$\mathbf{33.08}$	$\mathbf{.909}$	$30.01$	$.839$	$27.85$	$.777$	$23.93$	$.657$
$3.0$	$32.99$	$.904$	$29.49$	$.823$	$26.47$	$.735$	$20.66$	$.553$

Table 7: PSNR (dB) and SSIM on AFHQ-Cat super-resolution for varying confidence scale

\kappa

. For each column

\beta_{t}

is chosen to be the PSNR-optimal value from Table 6 (see Figure 7 ). The optimal

\kappa

decreases as the downsampling factor increases, favouring a more confident prior when observations are less informative.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Trajectories for a two-moons denoising problem under for trajectory priors in $\beta\in\{1,\beta_{t}^{*},0\}$ and endpoint-uncertainty scale $\kappa\in\{0,0.1,1,5\}$ . For $\kappa=0$ the Kalman gain vanishes, so the particles follow the unconditional prior unaffected by the observations. For $\kappa>0$ the prior weakens and lets the Kalman gain pulls the particles towards observation with the magnitude of the pull determined by $\kappa$ . Trajectories in the top row are deterministic given the initial noise and observations, as expected of an ODE update, whereas those in the lower two rows take random transitions at each step.

Figure 5: The conditioning step at six times along a trajectory (columns) for three choices of $\beta_{t}$ (rows): the decorrelated kernel ( $\beta_{t}=0$ ), the diffusion schedule ( $\beta_{t}=\beta^{\star}_{t}$ ) and the flow-matching kernel ( $\beta_{t}=1$ ). Each panel shows the prior over $x_{t_{i}}$ before conditioning on $y$ (dashed orange) and the posterior ( 9 ) (solid green), together with the true signal (red diamond) mapped into current time point (i.e. the red diamonds are showing $\alpha_{t}y$ for the time associated to that column, as this is what the time conditional posterior is targeting). Larger $\beta_{t}$ leaves less fresh noise $\tau^{2}$ in the kernel, so both prior and posterior contract, and at $\beta_{t}=1$ the conditioning step makes only a small correction to a nearly deterministic update.

	CelebA		AFHQ-Cat
Task	$\beta_{t}$	$\kappa$	$\beta_{t}$	$\kappa$
Denoising	$0.1$	$1.0$	$0.4$	$1.0$
Deblurring	$-0.99$	$1.0$	$-0.8$	$5.0$
Super-resolution	$-0.9$	$1.0$	$0.5$	$0.1$
Random inpainting	$-0.9$	$10.0$	$-0.9$	$10.0$
Box inpainting	$0.9$	$0.1$	$0.8$	$0.1$

Appendix

Table 8:

\beta

-Flow hyperparameters (

\beta_{t}

\kappa

) per task, selected on the validation split to maximise PSNR.

Method	MRI	CT
DPS	8.92	7.43
$\Pi$ GDM	9.02	21.75
D-Flow	312.52	283.87
PnP-Flow	15.64	15.95
Flower	23.48	48.76
$\beta$ -Flow (ours)	6.67	27.10

Appendix

Table 9: Wall-clock time per image, in seconds, for the MRI and CT results of Table 5 , under the same fixed step budget of 100 used there (except D-Flow, which uses its own optimisation budget).

Figure 6: Qualitative effect of $\beta_{t}$ on AFHQ-Cat super-resolution with fixed $\kappa=1$ . Leftmost column is the observation, and Columns 2-6 show reconstructions with different $\beta_{t}=\beta$ parameters. As the downsampling factor increases the observation constrains $x_{0}$ less, and reconstructions benefit more from increased $\beta_{t}$ provides. At $\times 4$ and $\times 6$ , $\beta_{t}=-0.6$ have visible grid artefacts and washed out tone that improve when $\beta_{t}$ increases, demonstrating the trend in Table 6 .

Figure 7: Qualitative effect of $\kappa$ on AFHQ-Cat super-resolution using the same test image as Figure 6 . For each scale factor $\beta_{t}$ is chosen to be its own PSNR-optimal $\beta_{t}$ ( $\times 2$ : $-0.6$ , $\times 3$ : $0.3$ , $\times 4,\times 6$ : $0.6$ ) from Figure 6 . A larger $\kappa$ widens the prior and lets the observation dominate the update; by $\kappa=3$ this produces blurring and a strong grid artefact, most pronounced at $\times 4$ and $\times 6$ where the observation is least informative, matching the trend in Table 7 .

Explore similar work

Jun 23, 2026cs.CV

What Do Flow-Based Inverse Solvers Approximate? A Posterior-Transport View

A growing family of training-free solvers -- FlowDPS, FLOWER, PnP-Flow and their diffusion ancestors (DPS, DAPS) -- repurpose a pretrained flow-matching prior to solve imaging inverse problems by adding a measurement-guidance term to the deterministic probability-flow ODE. Despite strong empirical results, what these per-step corrections actually approximate -- and how far the resulting samples are from the true posterior

p(x\mid y)

-- has not been characterized. We give a posterior-transport account of flow-based inverse problem solving. Our starting point is a simple but consequential fact: for a \emph{deterministic} flow prior, Bayesian conditioning is realized entirely by a \emph{reweighting of the source distribution}, not by a drift correction; pushing the reweighted source through the \emph{unmodified} velocity field yields exact posterior samples. From this we show that trajectory-guidance solvers can be read as the minimum-kinetic-energy \emph{correction} field needed to morph the unconditional source into the posterior, and that FlowDPS / FLOWER / PnP-Flow correspond to distinct zeroth-order / Gaussian / proximal approximations of this single object; we bound the resulting posterior bias in Wasserstein distance. A controlled

2

D study with a closed-form posterior confirms the theory decisively: source reweighting matches the true posterior to the Monte-Carlo floor on every metric, whereas trajectory guidance incurs

200

800\times

larger error and collapses posterior modes, \emph{regardless of guidance strength}. Guided by the analysis we propose a cheap, principled velocity-correction solver that is competitive across two in-domain priors (AFHQ, CelebA) and two out-of-distribution settings while, unlike point-estimate source-space optimizers, producing diverse posterior samples with uncertainty that correlates with reconstruction error.

Jian Xu, Delu Zeng, John Paisley +1

1RIKEN iTHEMS · 2RIKEN AIP · 3South China University of Technology +1

Flow Matching Bayesian Inverse Problems Flow-Based Generative Modeling Image Inverse Problems

Ask PDF

Oct 1, 2026math.NA

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

Sep 17, 2026cs.CV

FlowSGS: Improving Flow Matching 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 assume linear forward models and/or make simplifying approximations in posterior sampling. To circumvent these problems, we introduce FlowSGS, a flow-based posterior sampling method using Split Gibbs Sampling (SGS) to decompose the posterior into a likelihood step and a prior step. Specifically, we sample from the likelihood step using Langevin dynamics and leverage the Stochastic Interpolants (SI) framework to integrate a pretrained flow model into the prior step. We provide a form for the prior step that uses SI's reverse-time SDE, and show connections to previous PnP methods. Moreover, with the aid of the flow prior's straight probability paths and a novel timestep correction technique for the reverse-time SDE, FlowSGS requires fewer network evaluations in its prior step than plug-and-play diffusion samplers. Our experiments show state-of-the-art performance on a range of inverse problems. For the first time, we provide an experiment on a nonlinear inverse problem (Fourier phase retrieval) for flow-based inverse solvers.

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