PaperlayerView exact revision on arXiv

Source-linked AI summary

SNaP: One-Step Posterior Sampling for Noisy Inverse Problems

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

arXiv:2609.34071v1cs.CV

TL;DR

Iterative posterior samplers can require tens to thousands of network evaluations per draw, while one-step MeanFlow leaves no intermediate steps for measurement consistency. SNaP adapts its Gaussian source to measurement geometry; it produces diverse, high-quality samples across natural-image restoration and multi-coil MRI with one evaluation per draw, 30 to 2250× faster than iterative samplers.

  • Problem

    Iterative posterior sampling can require tens to thousands of network evaluations per draw, while one-step MeanFlow offers no intermediate states for enforcing measurement consistency.

    Source: §1 · p. 1

  • Method

    SNaP uses a measurement-adapted Gaussian source with a Tikhonov center and anisotropic spread, then applies conditional MeanFlow to transport its draws.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    30 to 2250× faster than iterative posterior samplers, SNaP produces diverse, high-quality samples across natural-image restoration and multi-coil MRI with one network evaluation per draw.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The exact one-step map transports the measurement-adapted source to the posterior, and SNaP's tested CelebA draws have near-unit calibration ratios.

    Source: §1 · p. 2

    The construction is limited to linear operators; nonlinear forward models require further work, and source sampling may be demanding without exploitable operator structure.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Noisy, often non-unique linear measurements make image recovery depend on prior information, while iterative posterior samplers can require many network evaluations. SNaP shapes a one-step transport around measurement geometry to retain sample diversity and reduce sampling cost.

1 Introduction

  • Ill-conditioned or rank-deficient measurement operators do not uniquely determine images, making prior information necessary for reliable recovery.Diffusion and flow-matching models provide learned image distributions that can serve as priors.Source: §1 · p. 1
  • Iterative generative inverse solvers enforce measurement consistency repeatedly, often requiring tens to thousands of network evaluations per sample.DPS uses 1000 NFEs and 27 seconds for one 128 × 128 CelebA Gaussian-deblurring reconstruction.Source: §1 · p. 1
  • A measurement-conditioned MeanFlow with an isotropic source produces nearly identical draws in inverse problems, so averaging provides no improvement.This motivates shaping one-step transport to reflect forward-model geometry.Source: §1 · pp. 1–2
  • SNaP uses a measurement-adapted Gaussian source that anchors well-measured directions while retaining variation in weakly measured and null-space directions.A single network evaluation maps each source draw to a reconstruction.Source: §1 · p. 2
  • 30 to 2250× faster than iterative posterior samplers, SNaP produces diverse, high-quality reconstructions across natural-image restoration and multi-coil MRI with one network evaluation per draw.Its samples achieve near-unit calibration ratios on tested CelebA tasks.Source: §1 · p. 2

2 Background

The background frames posterior sampling as inference from noisy linear measurements and reviews flow-based transports and one-step models. It also contrasts noisy inverse problems with a measurement-dependent source designed for noiseless constraints.

2 Background

  • The inverse problem is to infer an image from linear measurements with additive Gaussian noise, which may not uniquely determine the image.The goal in posterior sampling is to draw from p(x1 | y).Source: §2 · p. 2
  • Flow matching learns a velocity field that transports a source distribution to the data distribution, while MeanFlow models average velocity to reduce sampling to one evaluation.The source's random initial conditions provide variability propagated by the deterministic flow.Source: §2 · pp. 2–3
  • NullFlow [28] uses a measurement-dependent source for noiseless linear inverse problems, keeping its trajectory in the measurement-consistent affine set.With noisy measurements, the posterior is not supported on that set, and the anchor's row-space noise cannot be changed by a null-space-restricted velocity.Source: §2 · p. 3

3 Method

SNaP builds measurement geometry into a Gaussian source, then uses conditional MeanFlow to transport its draws toward the posterior in one step. The source's center and spread are chosen separately, and the exact transport result depends on stated regularity conditions.

3 Method

  • The source's anisotropic variation follows likelihood geometry: little in well-measured directions and more in weakly measured and null-space directions.Likelihood tubes are narrow in well-measured directions, wider in weakly measured directions, and unbounded along the null space.Source: §3 · pp. 3–4, Figure 2 · p. 4
  • For an isotropic working covariance, the displacement-minimizing linear center is a Tikhonov estimate, and its error covariance supplies the source spread.The working model constructs the source; it does not assume that the true image distribution is isotropic or Gaussian.Source: §3 · p. 4
  • The source variance is suppressed in well-measured directions while preserving variation in weakly measured and null-space directions.As noise tends to zero, the source converges to the NullFlow source when τ = 1.Source: §3 · p. 5
  • Under finite-second-moment and terminal-flow-limit assumptions, the exact one-step map transports the measurement-adapted source to the posterior for every τ > 0.The trained network approximates the exact average velocity, so its samples inherit approximation error.Source: §3 · p. 5

3.1 Training and sampling

  • Perturb-and-solve draws the source exactly and requires one solve with AA⊤+ λI per draw.This avoids directly applying the generally impractical matrix square root used to sample the source covariance.Source: §3.1 · p. 5
  • At test time, SNaP returns a reconstruction from one source draw and one network evaluation, with redrawn randomness giving conditionally independent learned-model samples.The conditional MeanFlow identity supplies the training target, and the derivative term uses one Jacobian–vector product.Source: §3.1 · p. 5

4 Experiments

Across natural-image restoration and multi-coil MRI, SNaP produces competitive posterior samples with one network evaluation per draw. The experiments compare quality, calibration, computational cost, and source choices.

4 Experiments

  • The evaluation covers CelebA, AFHQ-Cat, and multi-coil fastMRI, using natural-image restoration operators and MRI at ×4 and ×8 acceleration.Metrics include PSNR, SSIM, LPIPS, FID, and a calibration ratio for which 1 is ideal; the ratio summarizes second moments and does not certify conditional coverage.Source: §4 · pp. 5–6

4.1 Main Results

  • On CelebA, averaging SNaP draws gives the best SSIM on all four operators and best PSNR on three; its calibration is closest to 1 on all three, while fastMRI results are competitive.A single draw has the best LPIPS on Gaussian deblurring, and SNaP has the best FID on two of three CelebA operators.Source: Table 1 · p. 6, Table 2 · p. 6, §4.1 · p. 7

4.2 Computational cost

  • 30 to 2250× faster per sample on CelebA deblurring and 335 to 563× faster on ×8 MRI, with one network evaluation and a source draw per sample.The source solve is diagonal or Fourier diagonal for the image operators and a short conjugate-gradient solve for multi-coil MRI.Source: §4.2 · p. 7, Table 4 · p. 7

4.3 Ablations

  • With only the source changed, isotropic sources yield a 0.00 calibration ratio and near-identical draws, while the measurement-adapted source preserves dispersion.The isotropic-source draws have higher single-draw PSNR but worse LPIPS; the proposed source's advantage is perceptual rather than distortion-based.Source: §4.3 · pp. 7–8
  • Figure 3 compares samples on CelebA super-resolution and AFHQ-Cat inpainting and deblurring; SNaP has the best LPIPS on all three displayed tasks.Figure 4 shows a multi-coil fastMRI reconstruction at ×8 acceleration and 20 dB input SNR, with SNaP using one NFE versus 1000 for baselines.Source: Figure 3 · p. 8, Figure 4 · p. 8

5 Related work

Related work contrasts iterative posterior samplers with one-step generators and measurement-dependent source constructions. SNaP sits among approaches that incorporate measurements into one-step transport, with a constructed source tailored to the likelihood geometry.

Iterative posterior samplers

  • Iterative generative solvers enforce measurement consistency through data-fidelity or pseudoinverse corrections, or alternating generative and explicit consistency steps, at tens to thousands of network evaluations per sample.These methods reach high sample quality but require repeated evaluations.Source: §5 · pp. 8–9

One-step generative models

  • Consistency models, shortcut models, and MeanFlow map a source sample to data in one evaluation, so inverse-problem use must put measurements in conditioning or the source.MFSR uses conditioning without an explicit forward model or posterior sampling.Source: §5 · p. 9

Measurement-dependent sources

  • Measurement-dependent sources range from degraded-observation bridges without per-direction covariance to NullFlow's null-space restriction and VFM's jointly learned adapter.MeanFlow has also been used for PDE-based Bayesian inverse problems with a measurement-independent, prior-aligned base measure.Source: §5 · p. 9

6 Conclusion

SNaP constructs a measurement-adapted source for one-step posterior sampling in linear inverse problems with Gaussian noise. It targets the posterior with one network evaluation per draw and is reported to be 30 to 2250× faster than iterative samplers.

Conclusion

  • 30 to 2250× faster than iterative samplers, SNaP draws a posterior sample on natural images and multi-coil MRI in one network evaluation.On CelebA, averaging draws gives the best SSIM on all four operators and best PSNR on three; FID is best on two of three operators, and calibration is closest to 1 on all three.Source: §6 · p. 9
  • SNaP centers its measurement-adapted source at the linear estimate minimizing expected transport distance and sets its spread to match remaining uncertainty; the exact one-step map sends it to the target posterior.The construction replaces repeated measurement corrections with measurement information built into the transport.Source: §6 · p. 9

Limitations

  • The method is limited to linear operators, requires a solve that may be demanding without exploitable structure, and currently operates in pixel space.Nonlinear forward models and latent-space flow extensions remain open.Source: §6 · p. 9

A Proofs

The appendix establishes properties of the source center and exact conditional transport, then gives a source-sampling construction and additional AFHQ-Cat results. The exact-map statement is conditional on its stated regularity and terminal-limit assumptions.

A.1 Proof of Proposition 1

  • The source center minimizes expected displacement among linear estimators; for an isotropic Gaussian prior, its center and covariance give the exact Gaussian posterior.In the zero-noise limit, the center approaches the pseudoinverse estimate and the covariance approaches the null-space projection scaled by τ^2.Source: §A.1 · p. 13

A.2 Proof of Proposition 2

  • Under finite-second-moment and terminal-limit assumptions, the exact marginal-velocity flow transports the working source to the target posterior.The intermediate flow is defined for t < 1, and the terminal map agrees with the one-step average-velocity map.Source: §A.2 · pp. 13–14

A.3 Perturb-and-solve sampling

  • Perturb-and-solve draws have the same Gaussian source law without applying the covariance square root.Each draw uses applications of A and A⊤ and one solve with AA⊤ + λI.Source: §A.3 · p. 14

B.1 Full results on AFHQ-Cat

  • On AFHQ-Cat, a single SNaP draw is competitive in LPIPS at one network evaluation, while averaging a few draws recovers distortion at lower cost than iterative samplers.The MSE regressor is omitted for box inpainting because its single output is the blurred conditional mean of the masked region and is not comparable with the other methods.Source: Table 6 · p. 14, §B.1 · p. 14

B.2 Additional visual results

The additional visual results show SNaP reconstructions across natural-image restoration and multi-coil MRI, alongside further independent box-inpainting draws.

B.2 Additional visual results

  • Figures 5 and 6 show box-inpainting results on CelebA and AFHQ-Cat, respectively.Source: §B.2 · p. 15
  • Figure 7 shows multi-coil fastMRI brain reconstruction at ×8 acceleration and 30 dB input SNR.Source: §B.2 · p. 15
  • Figures 8 and 9–11 show CelebA Gaussian deblurring and several independent box-inpainting draws.Source: §B.2 · p. 15

B.3 Ablations

These ablations examine sensitivity to noise and source-design choices, while comparisons with VFM and NullFlow show how operator geometry and measurement noise distinguish their performance regimes.

B.3 Ablations

  • With noise overestimated by up to a factor of 5, PSNR rises from 32.90 to 34.29 dB as measurements get cleaner, while LPIPS stays between 0.017 and 0.020.The gain over the input shrinks from 5.09 to 4.33 dB as mismatch grows; underestimation is not tested here.Source: §B.3 · p. 15, Table 7 · p. 16
  • Changing τ from 0.15 to 0.5 leaves distortion indistinguishable and draw spread essentially unchanged at 5.49% and 5.48%, respectively.At τ = 1.0, distortion falls by about 0.4 dB while LPIPS is best at τ = 0.5.Source: §B.3 · p. 16, Table 8 · p. 20
  • Setting pratio = 0.25 is worse at every M, by about 0.3 dB PSNR at a single draw; 0.5 and 0.75 are indistinguishable within reported precision.The study uses pratio = 0.5 throughout.Source: §B.3 · p. 17, Table 9 · p. 20
  • For CelebA Gaussian deblurring, all metrics peak at k = 2 and decline afterward; k = 8 is worse than one step for M ≥16.The small k = 2 gain leads the authors to use k = 1 throughout.Source: §B.3 · p. 18, Table 10 · p. 21
  • For one CelebA deblurring measurement, the posterior source yields visible draw variation, while both isotropic Gaussian sources produce visually indistinguishable draws.Increasing the Gaussian scale from τ = 0.15 to τ = 1.0 does not restore spread; the source shape dictates it more than its size.Source: §B.3 · p. 19, Figure 12 · p. 21

B.4 Comparison with VFM

  • The diagonal adapter’s source approximation has KL divergence 1.576 nats for the rotated operator at τ = 1, versus 0 nats for the aligned operator.Rotation breaks the coincidence between the measurement basis and coordinate basis.Source: §B.4 · p. 20
  • On the rotated operator, SNaP’s support accuracy is 0.97 ± 0.004 versus 0.81 ± 0.03 for the jointly trained adapter.The learned adapter degrades by roughly sevenfold and the frozen variant by more than an order of magnitude; SNaP changes far less as its source rotates with the operator.Source: §B.4 · p. 20

B.5 Comparison with NullFlow

  • At σn = 0.05, SNaP is better than NullFlow on all three metrics, while NullFlow is stronger at σn = 0.NullFlow’s accuracy falls as noise is added; the two methods occupy different regimes.Source: §B.5 · pp. 20–21

B.6 Irreducible regression error: SNaP vs. isotropic sources

SNaP’s measurement-adapted source yields a lower irreducible velocity-regression error than isotropic sources that are at least as spread, but this comparison does not rank trained samplers. The section evaluates the gap and sets limits on the lemma’s scope.

B.6 Irreducible regression error: SNaP vs. isotropic sources

  • At r = t, the MeanFlow target is the conditional velocity x1 − x0 regressed on (zt, y), and Et is the error remaining for the best predictor.The source center does not affect this lemma’s regression error; its covariance is what lowers the error.Source: §B.6 · pp. 21–22
  • On CelebA Gaussian deblurring, the regression-error gap peaks at 56% under a stationary Gaussian prior with the empirical power spectrum.The two-dimensional mixture comparison uses exact quadrature; there the gap comes entirely from the observed direction.Source: §B.6 · p. 23, Figure 15 · p. 24
  • The lemma applies to the r = t objective, compares sources at least as spread as qτ, and does not directly compare trained samplers.A narrower source can have smaller Et; on the stated mixture, N(0, I) has smaller Et than qτ at every t.Source: §B.6 · p. 24
  • For any posterior with finite second moments, SNaP’s source has lower irreducible velocity-regression error than isotropic sources with σ ≥ τ.A smaller error means an easier regression problem at time t.Source: §B.6 · pp. 21, 24

B.7 Verification on a Gaussian-mixture inverse problem

This verification uses a four-component Gaussian-mixture prior whose partial observation leaves a bimodal posterior, distinguishing posterior sampling from estimating its mean. The source construction and trained one-step flow are compared with the closed-form posterior.

B.7 Verification on a Gaussian-mixture inverse problem

  • At y ≈ 2.5, the unobserved coordinate has posterior modes at ±2.5, while the MMSE estimate lies between them and a mode-collapsed sampler reports variance 0.12 instead of 6.36.The prior has four equal-weight isotropic Gaussian components, and only the first coordinate is observed.Source: §B.7 · p. 24
  • The perturb-and-solve source matches N(my, τ 2W) to Monte Carlo accuracy, with mean error 3.8 × 10−3 and covariance Frobenius error 4.8 × 10−3.Its observed-direction variance is 0.089, while the null direction retains the working-prior variance τ 2 = 9.Source: §B.7 · p. 25
  • The one-step flow reaches null-direction variance 6.66 against the closed-form posterior’s 6.36 and transports samples to both modes while pinning x[0] ≈ y.Figure 16 also shows some within-mode distortion relative to the exact posterior.Source: §B.7 · p. 25, Table 13 · p. 25
  • In the mixture calibration, flow assigns almost all null-direction freedom (6.66, posterior 6.36, source 9.0) and contracts the well-measured direction from source 0.089 toward posterior 0.053, reaching 0.079.Source: §B.7 · p. 25
  • For a matched single-Gaussian prior, the source is the posterior and the optimal map is the identity; retraining gives mean errors ≤ 0.05 and MMD2 ≤ 3 × 10−3.The table caption reports posterior variances overestimated by 37–39%.Source: §B.7 · p. 25, Table 14 · p. 25

C.1 Datasets and evaluation protocol

The evaluation covers CelebA, AFHQ-Cat, and complex multi-coil fastMRI brain data, with dataset-specific splits and preprocessing. Source construction uses operator-specific solves, including conjugate gradients for multi-coil MRI.

C.1 Datasets and evaluation protocol

  • The evaluation uses the first 100 CelebA test images, AFHQ-Cat’s stated train/validation/test split, and patient-level fastMRI brain splits with 4,305/539/539 slices.The MRI data use full multi-coil SENSE on 320 × 320 complex images, retained as real and imaginary channels.Source: §C.1 · p. 25

C.2 Solutions of the source construction

  • For multi-coil MRI, the source solve has no closed form and uses image-space conjugate gradients; the coupled coils make the measurement-space Gram matrix non-diagonal.The image-space formulation avoids the rank-deficient measurement-space solve, which gives a worse anchor at the same iteration budget.Source: §C.1 · p. 26, §C.2 · p. 26
  • Circular Gaussian deblurring uses two FFTs, while inpainting and box-decimated super-resolution use elementwise solves; unobserved inpainting coordinates retain variance τ 2.For inpainting, observed coordinates are shrunk toward the prior mean and retain variance τ 2λ/(1+λ).Source: §C.2 · p. 26
  • For general linear operators, equation 13 uses conjugate gradients, while diagonal or Fourier-diagonal operators can use a diagonal inverse or a single FFT pair.The solve is needed once per draw and does not depend on t, r, or the network.Source: §C.2 · pp. 26–27
  • Figure 16 compares source, closed-form posterior, and one-step samples for y = 2.5 and y = −2.5, including marginals along the unobserved coordinate.The caption says the one-step model captures both posterior modes, with some within-mode distortion.Source: Figure 16 · p. 27

C.3 Training and implementation details

Training uses a MeanFlow network conditioned on the measurement and a measurement-adapted source, with the source scale estimated from training data. The implementation details describe the solve costs and define evaluation metrics for reconstruction quality and sample calibration.

C.3 Training and implementation details

  • The backbone is a DDPM-style residual U-Net with separate embeddings for t, r, and σn, summed before entering the network body.The implementations use 9.1M parameters for CelebA, 110.5M for AFHQ-Cat, and 40.7M for fastMRI brain.Source: §C.3 · pp. 27–28
  • Training interpolates between the source x0 and clean image x1, regressing a MeanFlow target with r = t for a fraction of sampled pairs and an added near-endpoint sample schedule.The stated r = t probability is 0.5; 10% of the other samples are redrawn near r = 0 and t = 1.Source: §C.3 · p. 28
  • The source scale τ 2 is estimated offline from residuals of a linear-MMSE estimate on 256 training images and solved by fixed-point iteration.Because the linear estimate is suboptimal, its residual upper-bounds posterior variance and τ errs toward over-dispersion; reconstruction quality changes little across the reported range.Source: §C.3 · pp. 28–29
  • SNaP’s inference cost is M network evaluations and M source draws for M samples; source draws are O(n) for super-resolution and inpainting, O(n log n) for Gaussian deblurring, and use conjugate gradients for multi-coil MRI.The MRI source solve typically reaches its tolerance in 10–20 iterations, with a cap of 200.Source: §C.3 · p. 29

C.4 Metric details

  • Natural-image PSNR, SSIM, and LPIPS use specified image ranges and clamping, while MRI PSNR and SSIM use per-slice magnitude images and LPIPS is omitted.FID uses InceptionV3 pool3 features from 2,048 single-draw reconstructions and 2,048 ground-truth images.Source: §C.4 · pp. 29–30
  • The calibration ratio tests whether sample spread matches the sampler’s mean error; values below 1 indicate too little spread and values above 1 indicate over-dispersion.The evaluation uses M = 100 draws to estimate the quantities, with a correction for variance remaining in their mean.Source: §C.4 · p. 30

Related papers

Loading 2609.34071v1…