Stochastic Generative Plug-and-Play Priors

TL;DR

Score-based PnP: We show how pretrained score-based diffusion models can be used directly as priors inside classical Plug-and-Play reconstruction algorithms.
Stochastic generative update: We introduce a stochastic generative PnP framework that adds noise before denoising to better match the score model training regime.
Theory: The injected noise smooths the objective and helps the optimization avoid strict saddle points, while annealing enables convergence to a stationary point.
Flexibility: The framework works with multiple PnP algorithms, including ADMM, HQS, and PGM-style updates, and allows the injected noise level and denoiser noise level to differ.
Strong performance: The method improves reconstruction on challenging inverse problems such as large-mask inpainting and multi-coil MRI, outperforming standard PnP baselines and remaining competitive with diffusion-based solvers.

Abstract

Plug-and-play (PnP) methods are widely used for solving imaging inverse problems by incorporating a denoiser into optimization algorithms. Score-based diffusion models (SBDMs) have recently demonstrated strong generative performance through a denoiser trained across a wide range of noise levels. Despite their shared reliance on denoisers, it remains unclear how to systematically use SBDMs as priors within classical PnP frameworks without relying on the reverse diffusion sampling. In this paper, we establish a score-based interpretation of PnP that justifies using pre-trained SBDMs to be used directly within PnP algorithms. Building on this connection, we introduce a stochastic generative PnP (SGPnP) framework that injects noise to better leverage the expressive generative SBDM priors, improving robustness in severely ill-posed inverse problems. We provide a new theory showing that this noise injection induces optimization on a Gaussian-smoothed objective and promotes escape from strict saddle points. Experiments on challenging inverse tasks, such as multi-coil MRI reconstruction and large-mask natural image inpainting, demonstrate consistent improvement over conventional PnP methods and achieve performance competitive with diffusion-based solvers.

Reconstruction Performance

Comparison of our method (SGPnP) with several PnP baselines (DPIR and SNORE) and diffusion-based solvers (DPS and DiffPIR) on FFHQ (inpainting, super-resolution, and deblurring) and fastMRI (compressed sensing MRI). SGPnP consistently outperforms conventional PnP methods and is comparable to diffusion-based solvers.

Implementation Available

SGPnP variants

We share implementations of SGPnP-ADMM, SGPnP-DPIR, and SGPnP-PGM for inverse problems including inpainting, super-resolution, deblurring, and compressed sensing MRI.

Open in GitHub

Citation

@article{park2026SGPnP,
  title={Stochastic Generative Plug-and-Play Priors},
  author={Park, Chicago Y. and Chandler, Edward P. and Hu, Yuyang and McCann, Michael T. and Garcia-Cardona, Cristina and Wohlberg, Brendt and Kamilov, Ulugbek S.},
  journal={arXiv:2604.03603},
  year={2026}
}

TL;DR

Abstract

Reconstruction Performance

Effect of Noise Injection

Implementation Available

Citation