SNaP: One-Step Posterior Sampling for Noisy Inverse Problems
Shirin Shoushtari, Edward P. Chandler, Xiao Shi, Ulugbek S. Kamilov
cs.CV
2026-09-28
SNaP maps a measurement-adapted Gaussian source through MeanFlow in one NFE. CelebA deblurring takes 12 ms, 2250× faster than DPS, with calibration ratio 1.03.
Recovering an image from blur, undersampling or occlusion is a linear inverse problem y=Ax+n. A is often rank-deficient or ill-conditioned, so y does not determine a unique x. Posterior sampling draws from p(x|y) to show uncertainty and to average. Diffusion and flow-matching models are strong image priors, but DPS, DDRM, DiffPIR and Flower enforce measurement consistency over a long trajectory. On CelebA Gaussian deblurring, DPS needs 1000 network evaluations and 27 seconds for one 128×128 image.
MeanFlow learns average velocity over an interval and finishes generation in one evaluation. Collapsing the trajectory also removes the states where consistency used to be applied. Condition on y and keep an isotropic Gaussian source, and the draws become nearly copies of each other. The calibration ratio Δcal falls to 0.00, and averaging does not raise PSNR. One-step speed arrives without the directional variance a posterior should have.
SNaP changes the source. The target is still p(x|y). The source is a measurement-adapted Gaussian qτ(x|y)=N(my, τ²W), with mean and covariance closed-form in A, y and noise level σn.
The mean my is the Tikhonov estimate. Inside the family of linear centers plus independent noise, it minimizes expected squared displacement to posterior samples. Well-measured directions sit near the observation. Weak directions are regularized, so the construction does not amplify noise the way the pseudoinverse A†y does.
Covariance is diagonal in the right singular basis of A. Direction i gets weight wi=σn²/(σn²+τ²si²): large singular values, small variance; weak directions open up; the null space carries the working-prior scale τ². Likelihood superlevel sets are tubes around the least-squares set, and the source is narrow where those tubes are narrow. The working model only builds the source. It does not say natural images are isotropic Gaussians.
Training draws the source by perturb-and-solve: perturb, then solve AA⊤+λI, no square root of W. Super-resolution and inpainting cost O(n), Gaussian deblurring two FFTs, multi-coil MRI a conjugate-gradient solve of typically 10 to 20 steps. The network learns conditional average velocity. At test time one evaluation returns x̂=x0+uθ(x0,0,1|y,σn). The exact average velocity transports this source to the true posterior for any τ>0; the trained net approximates that map. As σn goes to 0, the source becomes NullFlow's pseudoinverse plus null-space noise.
CelebA, four operators, 100 test images. A single draw is already competitive on perceptual scores. Averaging raises distortion metrics and flattens fine detail.
| Method | NFE | Deblur LPIPS↓ | Deblur FID↓ | Δcal |
| DPS | 1000 | 0.019 | 18.25 | 0.64 |
| DDRM | 20 | 0.027 | 20.26 | 0.31 |
| SNaP (M=1) | 1 | 0.018 | 15.80 | 1.03 |
Single-draw deblurring LPIPS is 0.018, the best in the table. After M=100, SSIM leads on all four tasks and PSNR on three (35.91 dB on deblurring, against 35.74 for an MSE regressor and 35.65 for Flower5-OT). Super-resolution single-draw PSNR is 31.27 dB (LPIPS 0.023); box inpainting 30.73 dB (0.021). Averaging 100 draws takes those to 34.16 and 33.69 dB. Calibration ratios are 1.03, 0.99 and 1.08 on deblurring, super-resolution and box inpainting. DPS is over-dispersed at 2.15 on box inpainting. On random inpainting, DPS still has the better LPIPS: 0.013 against SNaP's 0.018.
Wall clock on CelebA deblurring is 12 ms per sample: 2251× versus DPS, 548× versus OT-ODE, 218× versus Flower, 30× versus DDRM. Multi-coil fastMRI at ×8 is 0.227 s, 563× versus CSGM and 335× versus DPS. At ×4 and 20 dB input SNR, a single draw reaches 32.06 dB PSNR, next to PnP-DM at 32.16 and below the supervised MSE regressor at 33.80. Averaging 16 draws gets to 32.85 dB, still short of the regressor. On AFHQ-Cat deblurring, OT-ODE's LPIPS 0.126 beats SNaP's 0.159.
Replace the source with N(0,I) or N(0,τ²I) and keep y as a conditioner: Δcal becomes 0.00, and averaging 100 draws barely moves PSNR. Diversity comes from source anisotropy. On a 2D checkerboard, rotating the operator raises SNaP's MMD² from 0.0027 to 0.017, while VFM's learned diagonal adapter goes from 0.011 to 0.074. On a 2D Gaussian-mixture problem with a closed-form posterior, one-step samples have null-space variance 6.66 against 6.36 for the true posterior, not the collapsed component variance 0.12.
If the product wants one image and the score is PSNR, a supervised MSE regressor is often enough. SNaP is for drawing many samples at one forward each: uncertainty maps, averaging to cut distortion, MRI reconstructions where variation is the point. Averaging itself trades LPIPS for PSNR. Deblurring LPIPS goes from 0.018 at M=1 to 0.030 at M=100, as the mean walks toward the posterior mean. Against zero-shot iterative samplers such as DPS, SNaP needs paired data and a retrain when A changes. Against NullFlow, which is exact only for noiseless measurements, box-inpaint SSIM at σn=0.05 is 0.934 versus 0.855.
For anyone already putting MeanFlow on restoration, the failure mode is specific. An isotropic source turns the one-step model into a point estimator. Source shape matters more than scale: τ from 0.15 to 0.5 barely moves the metrics.
Three limits are stated. The construction needs a linear operator, because the working posterior is closed-form only then. Nonlinear forwards are future work. Drawing the source solves AA⊤+λI, which can be expensive without structure. Everything is pixel-space; how to put this source on a latent flow model is open.
Calibration ratio is a second-moment necessary condition, not a coverage certificate. Each table uses 100 test images. FID uses 2048 single draws. SNaP is a supervised conditional model, so the speedup against DPS mixes training assumptions. On MRI the MSE regressor wins PSNR across the board. The argument for posterior sampling is diversity and calibration, not distortion scores. Only overestimates of σn were tested. Splitting the jump into k steps peaks at k=2 and is worse at k=8 than one step, so uθ at intermediate times is not a reliable ODE.