PaperlayerView exact revision on arXiv

Source-linked AI summary

Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors

Zihui Wu, Yu Sun, Yifan Chen, Bingliang Zhang, Yisong Yue, Katherine L. Bouman

arXiv:2405.18782v2eess.IVcs.CVstat.ML

TL;DR

Existing DM-based inverse solvers often rely on approximations or guidance updates that do not target the Bayesian posterior, limiting reliable posterior sampling. PnP-DM uses a split Gibbs sampler and an EDM connection to sample through a denoising problem; it outperforms existing DM-based methods across linear and nonlinear inverse problems. Its current whole-image steps may pose computational challenges for large-scale problems or expensive likelihood evaluations.

  • Problem

    Existing DM-based inverse methods approximate the generative process or use empirical guidance, so they do not generally sample the target Bayesian posterior.

    Source: §1 · p. 2

  • Method

    PnP-DM uses split Gibbs sampling and connects its prior update to a Gaussian denoising problem under the EDM formulation for rigorous posterior sampling with diverse pretrained DMs.

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

  • Results

    PnP-DM outperforms existing DM-based methods across linear and nonlinear inverse problems, including a nonlinear, severely ill-posed black-hole imaging problem.

    Source: §5 · pp. 9–10

  • Takeaways & Limitations

    The unified EDM interface supports using a range of state-of-the-art diffusion models as priors for Bayesian inference.

    Source: §1 · p. 2, §5 · p. 9

    Whole-image likelihood and prior steps may challenge large-scale or expensive-likelihood problems, and the analysis omits finite-iteration likelihood-step approximation error for general nonlinear problems.

    Source: §5 · p. 10

Abstract

from arXiv · show

Diffusion models (DMs) have recently shown outstanding capabilities in modeling complex image distributions, making them expressive image priors for solving Bayesian inverse problems. However, most existing DM-based methods rely on approximations in the generative process to be generic to different inverse problems, leading to inaccurate sample distributions that deviate from the target posterior defined within the Bayesian framework. To harness the generative power of DMs while avoiding such approximations, we propose a Markov chain Monte Carlo algorithm that performs posterior sampling for general inverse problems by reducing it to sampling the posterior of a Gaussian denoising problem. Crucially, we leverage a general DM formulation as a unified interface that allows for rigorously solving the denoising problem with a range of state-of-the-art DMs. We demonstrate the effectiveness of the proposed method on six inverse problems (three linear and three nonlinear), including a real-world black hole imaging problem. Experimental results indicate that our proposed method offers more accurate reconstructions and posterior estimation compared to existing DM-based imaging inverse methods.

1 Introduction

Sparse, noisy measurements leave substantial uncertainty about an image, motivating posterior sampling with expressive diffusion-model priors. PnP-DM is proposed to avoid approximations that can keep existing DM-based methods from sampling the target posterior.

1 Introduction

  • Sparse and noisy measurements can leave substantial uncertainty about the image, so posterior sampling represents possible solutions using a prior rather than selecting one deterministic reconstruction.Diffusion models offer a way to model complex image distributions, but the paper identifies reliable use of them as inverse solvers as an open challenge.Source: §1 · p. 1
  • Existing DM-based inverse methods either approximate an intractable likelihood score, which is exact only for Gaussian priors in the cited best approximation, or use empirical guidance updates that no longer target the Bayesian posterior.The paper characterizes these approaches as guidance methods; earlier sequential Monte Carlo work considered asymptotic consistency but was limited to linear imaging problems.Source: §1 · p. 2
  • PnP-DM uses a Markov chain Monte Carlo algorithm to perform posterior sampling with diffusion-model priors, avoiding the approximation required by a common approach to generic DM-based solvers.The method connects Bayesian denoising to unconditional image generation under the EDM formulation, enabling rigorous denoising-based posterior sampling with a range of pretrained DMs.Source: §1 · p. 2
  • PnP-DM is evaluated on three linear and three nonlinear noisy inverse problems, including real M87 black-hole data, and is reported to outperform baseline methods in posterior-estimation accuracy.The M87 diffusion prior was trained on GRMHD black-hole simulations; the reported posterior samples capture the bright-spot location and ring diameter.Source: §1 · p. 2, Figure 1 · p. 2

2 Preliminaries

The Split Gibbs Sampler uses variable splitting to sample an alternative distribution through separate likelihood and prior updates. This modular structure underpins PnP-DM, whose prior update is connected to diffusion-model denoising.

2 Preliminaries

  • SGS is an MCMC approach for Bayesian inference and serves as the backbone for the GPnP and DPnP frameworks in computational imaging.The paper relates it to the Proximal Sampler and describes it as a sampling analogue of Half-Quadratic Splitting.Source: §2 · p. 3
  • SGS uses variable splitting to sample an alternative distribution whose x-marginal converges in total variation to the target posterior as the coupling parameter ρ approaches zero.The augmented variable z and hyperparameter ρ control the coupling between x and z.Source: §2 · p. 3
  • Split Gibbs sampling decouples the likelihood and prior into separate conditional sampling steps, allowing each step to be designed modularly.The paper describes SGS as a principled posterior-sampling approach when both sampling steps are rigorously implemented.Source: §2 · p. 3
  • Earlier SGS-based methods differ mainly in their prior updates; PnP-DM connects that step to the EDM formulation to integrate state-of-the-art DMs through a unified interface.The paper also distinguishes its Bayesian posterior-sampling perspective from optimization-based variable splitting and notes sparse empirical evidence on real-world nonlinear problems in prior work.Source: §2 · p. 3

3 Method

PnP-DM uses Split Gibbs sampling to alternate likelihood and diffusion-prior steps, with an annealed coupling parameter. Its prior update is recast as Gaussian denoising and solved through the EDM framework; the section also presents theoretical convergence analysis.

3 Method

  • PnP-DM builds on Split Gibbs sampling, alternating data-consistency and diffusion-prior steps under an annealing schedule for the coupling parameter ρ.The EDM connection allows a family of state-of-the-art diffusion models to serve as priors without additional training.Source: §3 · p. 3, Figure 2 · p. 4

3.1 Likelihood step: enforcing data consistency

  • For linear Gaussian measurements, the likelihood update samples a Gaussian conditional; for differentiable nonlinear likelihoods, it uses Langevin dynamics initialized at z0 = x.The nonlinear update targets πZ|X=x.Source: §3.1 · p. 4

3.2 Prior step: denoising via the EDM framework

  • The prior update is a Gaussian denoising posterior with input z(k) and noise ρ; EDM matches it by setting σ(t*) = ρ and reversing to t = 0.This provides a unified interface for converting diffusion models into solvers for the denoising posterior.Source: §3.2 · pp. 4–5

3.3 Putting it all together

  • Against PnP-SGS, PnP-DM uses EDM for a rigorous prior step and anneals ρ, unlike its heuristic prior and constant ρ.The annealing schedule is described as important for highly ill-posed inverse problems.Source: §3.3 · p. 5

3.4 Theoretical insights

  • With exact likelihood steps and fixed ρ, the average Fisher-information gap falls as O(1/K), up to score error; the analysis requires L2-accurate scores under πτ.The convergence statement concerns the non-stationary process relative to the stationary process over K iterations.Source: §3.4 · p. 6
  • On a simulated compressed-sensing problem with a known posterior, PnP-DM samples it more accurately than DPS [17].The posterior is available in closed form for this comparison.Source: Figure 4 · p. 6

4 Experiments

The experiments assess posterior accuracy, benchmark reconstruction and uncertainty quality, and performance on nonlinear and real-world black hole imaging. Results include recovery of both Fourier phase-retrieval orientations and high-fidelity samples for simulated and real black hole data.

4.1 Validation with ground truth posterior

  • In a Gaussian-prior compressed-sensing test with a closed-form posterior, PnP-DM’s estimated standard deviation matches the true posterior in magnitude and spatial distribution, unlike DPS.Both methods estimate the posterior mean accurately.Source: §4.1 · p. 7

4.2 Benchmark experiments

  • 97.46% of ground-truth pixels fall within PnP-DM’s 3-sigma interval, versus 96.20% and 88.77% for baselines, without expecting 99% for a non-Gaussian posterior.The motion-deblur uncertainty analysis also reports the lowest outlier percentage among baselines without overestimating per-pixel standard deviations.Source: §4.2 · p. 8, Figure 6 · p. 8
  • On Fourier phase retrieval, PnP-DM significantly outperforms baselines and recovers both upright and rotated modes; 90 of 100 runs produce either orientation above 28 dB PSNR.The two orientations fit measurements equally because the forward model is invariant to 180° rotation.Source: §4.2 · p. 8, Figure 7 · p. 9

4.3 Experiments on black hole imaging

  • For simulated black hole imaging, PnP-DM captures two high-fidelity modes, while two of DPS’s three modes lack meaningful black hole structure and have larger measurement mismatch.The data are simulated for a black hole at Sagittarius A*’s location.Source: §4.3 · p. 9
  • On real M87 data, PnP-DM produces visually plausible samples consistent with ring diameters in the official EHT reconstruction.The data are from April 6, 2017.Source: §4.3 · p. 9

5 Conclusion

PnP-DM uses split Gibbs sampling to incorporate state-of-the-art diffusion models through a unified EDM interface, and the paper reports results across linear, nonlinear, and black hole imaging problems. The appendix describes an interpolation that alternates likelihood and prior steps, while the authors identify computational and theoretical limitations.

5 Conclusion

  • PnP-DM’s split Gibbs sampler links the prior step to the EDM formulation, integrating diverse state-of-the-art diffusion models through a unified interface.The conclusion reports performance across linear and nonlinear problems, including severely ill-posed black hole interferometric imaging.Source: §5 · pp. 9–10
  • Evaluating likelihood and prior steps for an entire image may challenge large-scale problems such as 3D imaging and problems with expensive likelihoods such as PDE inverse problems.The theoretical analysis also omits finite-iteration likelihood-step approximation error for general nonlinear problems.Source: §5 · p. 10
  • The authors expect expressive priors and uncertainty quantification to support better imaging systems, while noting that diffusion-model bias is inherited in the imaging process.Source: §5 · p. 10

A.1 Interpolation of PnP-DM

  • The interpolation assumes a unique t* satisfying σ(t*) = ρ when the coupling strength ρ is constant.Source: §A.1 · p. 16
  • The interpolation alternates a likelihood-step SDE with an EDM reverse-diffusion SDE for the prior step.The likelihood step runs from t = 0 to t = 1; the prior step runs from t = t* to t = 0.Source: §A.1 · p. 16

A.2 Proof of Theorem 3.1

The proof uses a KL-divergence derivative bound in terms of Fisher information, then applies it across the sampler’s likelihood and prior steps and combines the resulting bounds.

A.2 Proof of Theorem 3.1

  • Lemma A.4 quantifies the time derivative of KL divergence using Fisher information for a pair of general diffusion processes.Source: §A.2 · p. 17
  • The proof applies the lemma to the non-stationary and stationary likelihood steps, then integrates the bound over each step’s time interval.Source: §A.2 · p. 18
  • For prior steps, the proof applies the lemma using the exact and approximate drifts, combines the stepwise bounds across iterations, and divides by total process time.Source: §A.2 · p. 18

A.3 Discussion

Proposition A.5 bounds a weighted average of Fisher information across the sampler’s likelihood and prior processes. The discussion explains when this implies convergence of the unweighted average and how to handle diffusion processes with a zero coefficient at t=0.

A.3 Discussion

  • Proposition A.5 bounds weighted average Fisher information, while Theorem 3.1’s unweighted bound follows by lower-bounding the weight function by its infimum.The proposition applies under the same settings as Theorem 3.1.Source: §A.3 · pp. 18–19
  • With strictly positive δ, the weight remains positive and the unweighted average Fisher information converges to zero.This applies to VP- and VE-SDEs.Source: §A.3 · p. 19
  • When δ = 0, Fisher information may grow as the weight approaches zero, as can occur near t = 0 for iDDPM and EDM.For these processes, v(0) = 0 in the reverse diffusion.Source: §A.3 · p. 19
  • Replacing v(t) with v(t) + ε for ε > 0 makes the coefficient positive and guarantees convergence of the unweighted average Fisher information.The adjusted reverse SDE has the same law at each t as the original process.Source: §A.3 · p. 19

B Inverse problem setup

The experiments use several image datasets and linear and nonlinear forward models, including black hole imaging. The likelihood step uses exact Gaussian sampling for suitable linear models and Langevin Monte Carlo for general differentiable nonlinear models.

B Inverse problem setup

  • The benchmark uses the first 100 FFHQ images, while black hole experiments use simulated data and publicly available 2017 EHT data.Synthetic-prior experiments use grayscale CelebA images to construct a Gaussian prior.Source: B Inverse problem setup · p. 19
  • Gaussian deblurring uses a fixed 61×61 Gaussian kernel, motion deblurring uses per-image random kernels, and super-resolution downsamples by a factor of four.The blur kernel is held constant across compared methods for each test image.Source: B Inverse problem setup · p. 20
  • Nonlinear tests include phase-mask coded diffraction and Fourier phase retrieval; black hole likelihoods use closure measurements plus a total-flux constraint.The black hole measurements are derived from telescope visibilities and closure quantities.Source: B Inverse problem setup · pp. 20–21
  • Black hole data mismatch is the sum of reduced χ² values for closure phase and log closure amplitude; values around 2–3 indicate a good measurement fit.Each reduced χ² should ideally be around 1 for high-SNR data.Source: B Inverse problem setup · p. 21

C.1 Likelihood step

  • For linear models with Gaussian noise, the likelihood step samples exactly from a Gaussian with mean m(x) and covariance Λ^-1.An SVD can make the computation efficient, and a Fourier-transform derivation applies to circulant convolution matrices.Source: §C.1 · p. 21
  • For general differentiable nonlinear forward models, the likelihood step uses Langevin Monte Carlo with the likelihood-potential gradient and Gaussian noise.Algorithm 2 takes a step size and iteration count as hyperparameters; Table 3 lists likelihood-step hyperparameters.Source: §C.1 · p. 21, Table 3 · p. 22

C.2 Prior step

The prior step expresses reverse diffusion through the EDM formulation, enabling pretrained diffusion models to act as denoisers in the split Gibbs sampler. The appendix also records model preparation, solver choices, annealing, sample collection, and computational comparisons.

C.2 Prior step

  • The EDM formulation expresses reverse diffusion through scaling and noise levels, with a denoiser approximating the score via Tweedie’s formula.The denoiser is trained to minimize an ℓ2 denoising objective.Source: §C.2 · p. 22
  • The prior-step algorithm starts reverse diffusion at the first noise level σ(t_i) ≤ ρ, initializes with s(t_i)z, and integrates toward t = 0.It supports an SDE or probability-flow ODE solver.Source: §C.2 · p. 23
  • The black hole prior was trained on 3,068 simulated images, augmented by random flips and resizing to vary flux spin rotation and ring diameter.The paper also describes its model architecture and other training datasets.Source: §C.2 · p. 23
  • VP preconditioning converts the DDPM- or VP-SDE-trained checkpoints into EDM-form denoisers and permits other DM formulations with suitable preconditioning parameters.The conversion also supports comparison using the same pretrained models.Source: §C.2 · p. 23
  • DDS-DDPM is a special case of the prior step using VP preconditioning, VE reverse diffusion, and a particular integration scheme; PnP-DM also allows ODE and SDE solvers.Source: §C.2 · p. 24

C.3 Others

  • The coupling strength ρ uses an exponential annealing schedule, while investigation of other decay schedules is left for future work.The schedule is specified by a starting level, decay rate, and minimum value.Source: §C.3 · p. 24
  • Initialization varies by problem; negative initial pixel values caused problems in black hole experiments, and iteration counts were chosen for convergence.The method was reported as insensitive to reasonable initializations outside that black hole condition.Source: §C.3 · pp. 24–25
  • Across linear super-resolution and nonlinear coded diffraction tests, PnP-DM significantly outperforms DM baselines while achieving DPS-comparable runtimes (≈1.5×) despite ≈3× NFEs; linear experiments generated 20 samples along one batch-size-1 chain.For linear inverse problem experiments, the 20-sample PnP-DM chain took around 50 seconds, versus around 330 seconds for DPS. Runtime is reduced by adapting the starting noise level to ρk at each iteration and annealing ρk.Source: §C.3 · p. 25

D Implementation details of baseline methods

The appendix specifies implementation choices for baseline methods and reviews related plug-and-play, MCMC, and diffusion-model approaches to inverse problems.

D Implementation details of baseline methods

  • Baseline settings include a 500-iteration PnP-ADMM run, 40 DPIR iterations, and DDRM’s default parameters with a 20-step DDIM sampler.Source: D Implementation details of baseline… · p. 25
  • DPS uses a 1,000-step DDPM backbone, while PnP-SGS, DPnP, and HIO use method-specific burn-in, sampling, or initialization procedures.HIO selects the best measurement-fit reconstruction among 100 random initializations before 10,000 further iterations.Source: D Implementation details of baseline… · p. 26

E Additional related works

  • Optimization-based plug-and-play methods substitute learned denoisers for proximal updates, with cited convergence guarantees and empirical performance.Source: E Additional related works · p. 26
  • Learning-prior MCMC methods can generate diverse solutions and quantify uncertainty, while DM inverse solvers often use approximate data-consistency updates or task-specific retraining.Particle filtering and sequential Monte Carlo are also described as approaches to asymptotic consistency.Source: E Additional related works · p. 26

F.1 Synthetic prior experiment

Across these supplementary evaluations, PnP-DM estimates Gaussian-deblurring posterior uncertainty more accurately than DPS and produces fine-grained CDP reconstructions with greater robustness.

F.1 Synthetic prior experiment

  • PnP-DM estimates per-pixel posterior standard deviation more accurately than DPS on Gaussian deblurring, while their posterior mean estimates are roughly equally accurate.The comparison uses a Gaussian prior and a 7×7 Gaussian blur kernel with standard deviation 3.0.Source: Figure 10 · p. 27, §F.1 · p. 27

F.2 Linear inverse problems

  • Visual comparisons show one generated sample per algorithm for Gaussian deblurring and super-resolution with i.i.d. Gaussian noise σy = 0.05, alongside additional examples for deblurring, motion deblurring, and super-resolution.The additional examples are identified in Figures 12–14.Source: Figure 11 · p. 27, §F.2 · p. 27

F.3 Nonlinear inverse problems

  • On CDP reconstruction, PnP-DM recovers fine-grained hair-thread features missing from baseline reconstructions.The comparison visualizes one sample from each method; additional CDP examples are also shown.Source: §F.3 · p. 27
  • DPS failed on around half of the CDP test images, while PnP-DM produced high-fidelity reconstructions on almost all, indicating greater robustness.Each method's best reconstruction out of four runs per test image was selected by PSNR.Source: §F.3 · p. 27

F.4 Black hole imaging

In simulated black-hole imaging, PnP-DM samples both posterior modes and fits the measurements, whereas DPS captures only one mode; real-data samples are diverse and measurement-consistent.

F.4 Black hole imaging

  • On simulated black-hole data, PnP-DM samples both posterior modes and fits the measurements, while DPS captures only one mode and its Mode 2 and Mode 3 samples deviate from measurements and lack expected structure.Additional PnP-DM samples on real M87 data are diverse, fit the measurements with mismatch values around 2, and show a ring diameter consistent with the official EHT reconstruction and a bright spot in the ring's lower half.Source: §F.4 · pp. 28–29

F.5 Further analysis

Further analysis finds that annealing settings affect convergence behavior and final PSNR, while super-resolution images and metrics stabilize after minimum coupling is reached; the paper also lists asset licenses.

F.5 Further analysis

  • Different annealing decay rates yield different convergence rates but the same final PSNR level, while lower minimum coupling strength yields higher final PSNR.The sensitivity analysis covers super-resolution and coded diffraction patterns; as ρ decreases, the stationary distribution πX should converge to the target posterior p(x|y).Source: §F.5 · p. 29, Figure 20 · p. 32
  • On super-resolution, xk approaches the ground truth and zk becomes less noisy as ρk decreases; image quality and metrics stabilize after minimum coupling, with good images in around 40 of 100 iterations.The reported runtime is around 30 seconds and 1600 NFEs.Source: §F.5 · pp. 29–30, Figure 21 · p. 32

G Licenses

  • The listed FFHQ license is Creative Commons BY-NC-SA 4.0, the simulated and real black-hole data license is unknown, and the FFHQ pretrained model license is MIT.The paper also lists the licenses of the code repositories after their links.Source: G Licenses · pp. 30–31

Related papers

Loading 2405.18782v2…