SNaP: One-Step Posterior Sampling for Noisy Inverse Problems

Shirin Shoushtari; Edward P. Chandler; Xiao Shi; Ulugbek S. Kamilov

cs.CVSep 28, 2026

SNaP: One-Step Posterior Sampling for Noisy Inverse Problems

Authors: Shirin Shoushtari, Edward P. Chandler, Xiao Shi, Ulugbek S. Kamilov

Organizations: WashU · UW-Madison

Abstract

Diffusion and flow-matching models can produce high-quality posterior samples for inverse problems, but typically require tens to thousands of network evaluations per draw. MeanFlow enables one-step generation, yet applying it to inverse problems leaves no intermediate steps at which to enforce measurement consistency. We introduce SNaP, a one-step MeanFlow posterior sampler for linear inverse problems with Gaussian noise. Its central innovation is a measurement-adapted source: a Gaussian distribution whose mean and anisotropic covariance are determined by the measurement operator, observation, and noise level. The source anchors well-measured directions while preserving variation where the measurements are weak or uninformative. We show that the exact conditional flow transports this source to the true posterior. Across natural-image restoration and multi-coil MRI, SNaP produces diverse, high-quality samples with one network evaluation per draw, 30 to 2250 $\times$ faster than iterative samplers.

Figures & tables

Figure 1: Left: LPIPS versus wall-clock time per sample for CelebA Gaussian deblurring. SNaP achieves lowest LPIPS at a fraction of the cost of iterative samplers. Right: $4\times$ super-resolution on AFHQ-Cat. Columns show one SNaP draw, averages of $M=4$ , $16$ , and $100$ draws, and the ground truth; zoomed crops, error maps, and PSNR are shown below. Each draw requires one network evaluation. Averaging reduces pixel error and improves PSNR while smoothing fine detail.

Figure 2: Geometry-aware one-step transport (schematic). Left: Likelihood superlevel sets form tubes around $\mathcal{S}_{{\bm{y}}}$ : narrow in well-measured directions, wider in weakly measured directions, and unbounded along $\nullsp({\bm{A}})$ . Gray points represent posterior samples. Middle: An isotropic source spreads samples across directions regardless of measurement strength, leading to large source-to-target displacements. Right: The SNaP source concentrates samples in well-measured directions while retaining variation in weakly measured and null-space directions, shortening displacements.

Method	NFE	Deblurring			Super-resolution			Random inpainting			Box inpainting
Method	NFE	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$

Degraded	–	27.26	0.838	0.214	11.67	0.182	0.859	11.98	0.198	1.070	22.33	0.754	0.217
MSE regressor	1	35.74	0.953	0.029	34.17	0.945	0.033	34.42	0.955	0.022	–	–	–
PnP-GS	23	33.97	0.924	0.041	31.23	0.890	0.065	29.17	0.874	0.066	–	–	–
DDRM	20	35.02	0.946	0.027	32.24	0.927	0.027	32.40	0.942	0.029	–	–	–

Table 1: CelebA,

100

test images. We report single draw (

M{=}1

) and the average of

M

draws for SNaP . Note that a single SNaP draw is competitive in LPIPS at one network evaluation, while averaging draws attains the best SSIM on every task and the best PSNR on three of the four. Best and second best are shown in blue and red .

Method	NFE	Gaussian deblurring		Super-resolution		Box inpainting
Method	NFE	FID $\downarrow$	$\Delta_{\mathrm{cal}}\,(\to 1)$	FID $\downarrow$	$\Delta_{\mathrm{cal}}\,(\to 1)$	FID $\downarrow$	$\Delta_{\mathrm{cal}}\,(\to 1)$
MSE regressor	1	22.06	–	24.76	–	–	–
DDRM	20	20.26	0.31	18.94	0.19	–	–
Flower1-OT	100	17.86	0.23	23.79	0.24	18.72	0.24
DPS	1000	18.25	0.64	19.94	0.67	15.42	2.15
SNaP (ours)	1	15.80	1.03	16.80	0.99	15.53	1.08

Table 2: Distributional and calibration metrics for CelebA.

\Delta_{\mathrm{cal}}

is the calibration ratio of App. C.4 , for which

1

is ideal. Note that SNaP attains the best or second-best FID among the posterior-sampling baselines in one network evaluation. Best and second best are shown in blue and red .

		$\times 4$				$\times 8$
		$20$ dB		$30$ dB		$20$ dB		$30$ dB
Method	NFE	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Zero-filled	—	25.35	0.791	25.37	0.793	21.98	0.694	21.99	0.694
Wavelet $+\ell_{1}$	—	26.57	0.672	27.75	0.768	23.02	0.551	23.81	0.643
TV	—	26.19	0.798	26.27	0.799	22.62	0.689	22.66	0.691
MSE regressor	1	33.80	0.909	33.85	0.912	30.71	0.875	30.73	0.878

Table 3: fastMRI brain, multi-coil. PSNR and SSIM over

100

test slices, at

\times 4

and

\times 8

acceleration and two input SNRs. For SNaP , we report results for a single draw

M=1

and average of

M

draws. Note that among posterior samplers SNaP is very competitive even using one NFE. Best and second best are shown in blue and red .

Figure 3: Visual comparison of samples for three imaging tasks, CelebA super-resolution and AFHQ-Cat box inpainting and deblurring, with PSNR and LPIPS for the displayed image. Note that SNaP requires one NFE, while DiffPIR requires $100$ and DPS and Flower require $1000$ . SNaP attains the best LPIPS on all three.

Figure 4: fastMRI brain, multi-coil, $\times 8$ acceleration at $20$ dB input SNR. Reconstructions with PSNR and SSIM, and error magnitude and a zoomed in region is displayed. Note that SNaP achieves competitive reconstruction using 1 NFE, while baselines use $1000$ .

CelebA, Gaussian deblurring
	DPS	OT-ODE	Flower	DDRM	SNaP
NFE $\downarrow$	$1000$	$180$	$100$	$20$	$\mathbf{1}$
Time / sample (s) $\downarrow$	27.014	6.575	2.614	0.360	0.012
Slowdown vs. ours $\downarrow$	$2251\times$	$548\times$	$218\times$	$30\times$	–
fastMRI brain, multi-coil $\times 8$
	CSGM	DAPS	PnP-DM	DPS	SNaP

Table 4: Computational cost. NFE and 1-batch wall-clock time per posterior sample. Best in each row is in bold.

Source	PSNR $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	$\Delta_{\mathrm{cal}}\,(\to 1)$
	$M=1$		$M=100$
${\mathcal{N}}({\bm{0}},{\bm{I}})$	35.85	0.030	35.86	0.030	0.00
${\mathcal{N}}({\bm{0}},\tau^{2}{\bm{I}})$	35.83	0.030	35.84	0.030	0.00
${\mathcal{N}}({\bm{m}}_{{\bm{y}}},\tau^{2}{\bm{W}})$ (ours)	32.90	0.018	35.91	0.030	1.03

Table 5: Source distribution. CelebA, Gaussian deblurring at

\sigma_{\mathrm{n}}=0.05

; only the source differs across rows.

M

is the number of averaged draws, and

\Delta_{\mathrm{cal}}

is the calibration ratio (App. C.4 ), for which

1

is ideal. An isotropic source collapses onto the posterior mean, where draws are near-identical and averaging brings no gain.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Method	NFE	Deblurring			Super-resolution			Random inpainting			Box inpainting
Method	NFE	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$

Degraded	–	24.39	0.532	0.543	11.98	0.212	0.900	13.25	0.214	1.091	21.57	0.736	0.214
MSE regressor	1	29.82	0.806	0.289	29.16	0.817	0.201	33.23	0.915	0.066	–	–	–
PnP-GS	23	28.39	0.787	0.387	24.44	0.639	0.411	29.89	0.841	0.123	–	–	–
PnP-Flow5	500–2500	29.01	0.785	0.312	28.01	0.791	0.167	34.47	0.933	0.044	27.14	0.900	0.127

Appendix

Table 6: AFHQ-Cat,

100

test images. We report single draw (

M{=}1

) and the average of

M

draws for SNaP . Note that a single SNaP draw is competitive in LPIPS at one network evaluation, and averaging a few draws recovers distortion at a cost still far below the iterative samplers. On AFHQ-Cat the number of time steps used by PnP-Flow and Flower depends on the operator, so their NFE does too; the per-operator budgets are listed in App. C.3 . The MSE regressor is omitted on box inpainting, where its single output is the blurred conditional mean of the masked region and is not comparable with the other methods. Best and second best are shown in blue and red .

Figure 7: fastMRI brain, multi-coil, $\times 8$ acceleration at $30$ dB input SNR.

Figure 8: Gaussian deblurring on CelebA.

Figure 9: Independent SNaP draws for box inpainting.

Figure 10: Further independent SNaP draws for box inpainting.

Figure 11: Further independent SNaP draws for box inpainting.

Measurement noise $\sigma_{\mathrm{n}}$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	input PSNR	Gain
$0.01$	34.29	0.940	0.020	29.96	+4.33
$0.02$	34.15	0.939	0.019	29.62	+4.53
$0.03$	33.89	0.936	0.018	29.12	+4.77
$0.04$	33.50	0.930	0.017	28.52	+4.98
$0.05^{\dagger}$	32.90	0.919	0.018	27.81	+5.09

Appendix

Table 7: CelebA Gaussian deblurring trained at

\sigma_{\mathrm{n}}=0.05

and evaluated on

y

with different measurement noise level.

{\dagger}

marks the matched case, where the test noise equals the training noise.

	$\tau=0.15$			$\tau=0.5$			$\tau=1.0$
$M$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
$1$	32.941	0.9198	0.0190	32.922	0.9187	0.0186	32.512	0.9060	0.0244
$4$	34.966	0.9463	0.0209	34.948	0.9457	0.0181	34.588	0.9385	0.0197
$16$	35.688	0.9537	0.0275	35.663	0.9532	0.0234	35.323	0.9475	0.0249
$100$	35.913	0.9558	0.0300	35.885	0.9553	0.0256	35.552	0.9500	0.0273

Appendix

Table 8: Working-prior scale

\tau

. CelebA Gaussian deblurring,

100

test measurements, input PSNR

27.24

dB.

M

is the number of averaged draws

	$p_{\mathrm{ratio}}=0.25$			$p_{\mathrm{ratio}}=0.5$			$p_{\mathrm{ratio}}=0.75$
$M$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
$1$	32.567	0.9130	0.0216	32.941	0.9198	0.0190	32.976	0.9210	0.0189
$4$	34.679	0.9429	0.0197	34.966	0.9463	0.0209	34.976	0.9465	0.0206
$16$	35.438	0.9513	0.0266	35.688	0.9537	0.0275	35.682	0.9536	0.0268
$100$	35.676	0.9537	0.0296	35.913	0.9558	0.0300	35.902	0.9557	0.0292

Appendix

Table 9:

(r,t)

sampling ratio. CelebA Gaussian deblurring,

\tau=0.15

100

test measurements, input PSNR

27.24

dB.

M

is the number of averaged draws.

	$k=1$			$k=2$			$k=4$			$k=8$
$M$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
$1$	32.941	0.9198	0.0190	33.084	0.9211	0.0172	33.037	0.9205	0.0172	32.992	0.9200	0.0173
$4$	34.966	0.9463	0.0209	35.044	0.9469	0.0206	35.001	0.9465	0.0206	34.942	0.9458	0.0210
$16$	35.688	0.9537	0.0275	35.729	0.9540	0.0275	35.688	0.9536	0.0276	35.621	0.9529	0.0281
$100$	35.913	0.9558	0.0300	35.943	0.9561	0.0302	35.902	0.9557	0.0303	35.833	0.9550	0.0308

Appendix

Table 10: Trajectory steps

k

. CelebA Gaussian deblurring,

\tau=0.15

100

test measurements, input PSNR

27.24

dB.

M

is the number of averaged draws. Best in each row and metric is in bold.

Figure 12: Source distribution, qualitative. Six independent draws from a fixed CelebA Gaussian-deblurring measurement under the three sources, with the per-pixel standard deviation over draws on the same $\times 20$ scale in every row. The posterior source varies where the measurement is uninformative; both Gaussian sources collapse onto a single reconstruction.

Figure 13: Source distribution, averaging budget. PSNR (left, higher is better) and LPIPS (right, lower is better) against the number of averaged draws $M$ on CelebA Gaussian deblurring. Only the posterior source responds to $M$ ; the Gaussian sources are flat because their draws carry almost no variance to average away.

Figure 14: One-step ( $K{=}1$ ) posterior samples on the $4{\times}4$ checkerboard. Rows are the two forward operators and columns are the learned adapter of VFM, the SNaP source, and the exact rejection-sampled posterior. The green line is the measurement constraint $\bm{a}^{\top}{\bm{x}}={\bm{y}}$ , and the dashed ellipse on each method panel is that method’s $2\sigma$ source covariance, the learned diagonal $q_{\phi}({\bm{z}}\,|\,{\bm{y}})$ for VFM and the closed-form $\tau^{2}{\bm{W}}$ for SNaP . In the top row the measurement direction is axis-aligned, both sources align with the constraint, and both recover the posterior. In the bottom row the constraint tilts. The SNaP source tilts with it, whereas the diagonal adapter cannot and degenerates to a near-isotropic blob, so its samples concentrate on one mode and spill off the checkerboard support.

Method	Aligned (parity)	Rotated (discriminating)
SNaP	$0.0027\pm 0.0003$	$0.0171\pm 0.0012$
VFM (joint)	$0.0109\pm 0.0007$	$0.0737\pm 0.0199$
VFM (frozen ${\bm{\theta}}$ )	$0.0102\pm 0.0003$	$0.1900\pm 0.0015$

Appendix

Table 11: Posterior MMD 2 (

\downarrow

) on the two-dimensional checkerboard benchmark at

K{=}1

, mean

\pm

standard deviation over three seeds, each at its validation-selected setting. The aligned operator is a parity check, where a diagonal adapter is exactly expressive, and the rotated operator differs only in the direction of

\bm{a}

	NullFlow			SNaP
Test $\sigma_{\mathrm{n}}$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
$0.00$	$33.50$	$0.977$	$0.012$	$31.37$	$0.947$	$0.025$
$0.01$	$33.27$	$0.970$	$0.012$	$31.35$	$0.947$	$0.024$
$0.02$	$32.65$	$0.952$	$0.016$	$31.30$	$0.946$	$0.023$
$0.05$	$30.11$	$0.855$	$0.055$	$30.74$	$0.934$	$0.021$

Appendix

Table 12: NullFlow vs. SNaP on CelebA box inpainting (

M{=}1

) as test noise grows.

Figure 15: Irreducible velocity-regression error $\mathcal{E}_{t}$ of equation 21 under the SNaP source $q_{\tau}(\cdot\,|\,{\bm{y}})$ and the isotropic source ${\mathcal{N}}({\bm{0}},\tau^{2}{\bm{I}})$ . Left: two-dimensional mixture of App. B.7 ( $\tau=3$ , $\sigma_{\mathrm{n}}=0.3$ , $y=2.5$ ), exact. Right: CelebA Gaussian deblurring ( $\tau=0.15$ , $\sigma_{\mathrm{n}}=0.05$ ) under a stationary Gaussian prior. Lemma 2 guarantees the ordering; the shaded region is the gap.

$y$	mean (post)	mean ( SNaP )	$\left\lVert\Delta\text{mean}\right\rVert$	$\Var$ (observed, null) post $\to$ SNaP	$\mathrm{MMD}^{2}$
$2.5$	$(2.5,0)$	$(2.474,0.065)$	$0.069$	$(0.052,6.37)\!\to\!(0.079,6.66)$	$4.0{\times}10^{-3}$
$2.0$	$(2.212,0)$	$(2.045,0.051)$	$0.174$	$(0.052,6.37)\!\to\!(0.084,6.64)$	$1.1{\times}10^{-2}$
$-2.5$	$(-2.5,0)$	$(-2.444,0.108)$	$0.121$	$(0.052,6.37)\!\to\!(0.081,6.59)$	$2.6{\times}10^{-3}$

Appendix

Table 13: Posterior match at three in-distribution measurements.

$y$	mean (post)	mean ( SNaP )	$\left\lVert\Delta\text{mean}\right\rVert$	$\Var$ (observed, null) post $\to$ SNaP	$\mathrm{MMD}^{2}$
$0$	$(0,0)$	$(-0.040,0.026)$	$0.048$	$(0.052,0.123)\!\to\!(0.071,0.171)$	$2.7{\times}10^{-3}$
$1$	$(0.576,0)$	$(0.553,0.012)$	$0.026$	$(0.052,0.123)\!\to\!(0.069,0.160)$	$1.7{\times}10^{-3}$

Appendix

Table 14: Single-Gaussian ground-truth check. The map reduces to the Wiener estimate, with the posterior mean recovered to within

0.05

and the posterior variances over-estimated by

37

–

39\%

	${\bm{m}}_{{\bm{y}}}$	$\tau^{2}{\bm{W}}$ (diag)	error
closed form	$[2.4752,\ 0]$	$[0.0891,\ 9.000]$	—
RTO equation 13	$[2.4752,\ 0.0038]$	$[0.0889,\ 8.995]$	$\left\lVert\Delta{\bm{m}}_{{\bm{y}}}\right\rVert=3.8{\times}10^{-3}$ , $\left\lVert\Delta\mathrm{cov}\right\rVert_{F}=4.8{\times}10^{-3}$

Appendix

Table 15: The square-root-free sampler equation 13 reproduces

{\mathcal{N}}({\bm{m}}_{{\bm{y}}},\tau^{2}{\bm{W}})

to Monte-Carlo accuracy at

y=2.5

direction	$\Var$ posterior	$\Var$ one-step SNaP	$\Var$ source
${\bm{e}}_{0}$ (observed, $s=1$ )	$0.0518$	$0.0791$	$0.0891$
${\bm{e}}_{1}$ (null, $s=0$ )	$6.372$	$6.656$	$9.000$

Appendix

Table 16: Per-direction variance at

y=2.5

(equation 10 ). The source over-disperses both directions relative to the posterior; the trained flow moves both toward it, almost exactly in the null direction and partially in the observed one.

Figure 16: Posterior recovery for observations $y=2.5$ (left) and $y=-2.5$ (right). The top row compares samples from the source distribution (green), the closed-form posterior (blue), and the one-step model (red). Dashed vertical lines indicate the observed values. The bottom row shows the corresponding marginals along the unobserved coordinate $x_{1}$ ; the dashed line marks the posterior mean. The one-step model captures both posterior modes, although it exhibits some within-mode distortion relative to the exact posterior.

Problem	${\bm{A}}{\bm{A}}^{\!\top}$ structure	Cost	Notes
Inpainting (random)	${\bm{I}}_{m}$	free	${\bm{W}}$ diagonal
Inpainting (box)	${\bm{I}}_{m}$	free	identical structure
MRI, single-coil	${\bm{I}}_{m}$	free	diagonal in $k$ -space
SR, average-pool	$d^{-2}{\bm{I}}$	free	disjoint rows of norm $1/d$
Deblurring, Gaussian	diag. in Fourier	2 FFTs	${\bm{A}}{\bm{A}}^{\!\top}=\left\lvert\hat{h}\right\rvert^{2}$
Deblurring, uniform	diag. in Fourier	2 FFTs	$\hat{h}$ has exact zeros

Appendix

Table 17: Cost of the source construction equation 13 . Larger

\lambda

(noisier data) improves conditioning and accelerates CG.

Explore similar work

FlowSGS: Improving Flow Matching Priors for Inverse Imaging with Stochastic Interpolants
Sep 17, 2026Tianao Li, Xinhui Qian, Emma AlexanderBayesian Inverse Problems Posterior Sampling
Department of Computer Science, Northwestern University · NSF-Simons AI Institute for the Sky (SkAI) · Department of Statistics and Data Science, Northwestern University
What Do Flow-Based Inverse Solvers Approximate? A Posterior-Transport View
Jun 23, 2026Jian Xu, Delu Zeng, John Paisley +1Bayesian Inverse Problems Posterior Sampling
1RIKEN iTHEMS · 2RIKEN AIP · 3South China University of Technology +1
Provable diffusion-based posterior sampling for linear inverse problems via DDIM
Jul 21, 2026Yuchen Jiao, Na Li, Changxiao Cai +2Bayesian Inverse Problems Posterior Sampling
Department of Statistics and Data Science, Chinese University of Hong Kong, Hong Kong · Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, USA · Department of Statistics and Data Science, the Wharton School, University of Pennsylvania, PA, USA

Sep 17, 2026, cs.CVJ/K move · Enter open · S save

FlowSGS: Improving Flow Matching Priors for Inverse Imaging with Stochastic Interpolants

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

Bayesian Inverse Problems Posterior Sampling Stochastic Interpolants

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

FlowSGS: Improving Flow Matching Priors for Inverse Imaging with Stochastic Interpolants

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

Provable diffusion-based posterior sampling for linear inverse problems via DDIM