Stochastic Generative Plug-and-Play Priors

Chicago Y. Park; Edward P. Chandler; Yuyang Hu; Michael T. McCann; Cristina Garcia-Cardona; Brendt Wohlberg; Ulugbek S. Kamilov

cs.CVApr 4, 2026

Stochastic Generative Plug-and-Play Priors

Authors: Chicago Y. Park, Edward P. Chandler, Yuyang Hu, Michael T. McCann, Cristina Garcia-Cardona, Brendt Wohlberg, Ulugbek S. Kamilov

Organizations: University of Wisconsin–Madison · Washington University in St. Louis · Los Alamos National Laboratory

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 the PnP framework without relying on reverse diffusion sampling. In this paper, we establish a score-based interpretation of PnP that justifies using pretrained SBDMs 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, thereby 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.

Figures & tables

Figure 1: Comparison of deterministic PnP [ 43 ] , stochastic PnP [ 27 ] , and SGPnP on three inpainting tasks: random masks, box masks with small missing regions, and box masks with large missing regions. Deterministic PnP succeeds only on random inpainting, while stochastic PnP extends this to small box masks but fails for large missing regions. In contrast, the proposed SGPnP approach performs reliably across all three settings.

Figure 2

Testing data		Input	DPIR	SNORE	DPS	DiffPIR	SGPnP

Inpainting	PSNR $\uparrow$	$18.17$	$18.42$	$18.50$	$23.69$	$25.09$	$\mathbf{25.21}$
	SSIM $\uparrow$	$0.766$	$0.797$	$0.799$	$0.810$	$0.860$	$\mathbf{0.874}$
	LPIPS $\downarrow$	$0.289$	$0.264$	$0.250$	$0.176$	$0.115$	$\mathbf{0.108}$

Super-resolution	PSNR $\uparrow$	$24.61$	$26.54$	$26.40$	$28.66$	$29.69$	$\mathbf{29.90}$

Table 1: Quantitative comparison between several PnP baselines and diffusion-based solvers on FFHQ (inpainting,

4\times

super-resolution, deblurring) and fastMRI (

4\times

CS-MRI). Best values are highlighted for each metric and inverse problem. SGPnP consistently improves upon conventional PnP methods and is comparable to diffusion-based solvers.

Figure 2: Qualitative visual comparison corresponding to the quantitative results in Table 1 . Note that for challenging box inpainting, deterministic PnP method (DPIR) and stochastic DRUNet-based PnP (SNORE) produce incomplete reconstructions, whereas SGPnP produces more plausible image completions.

Panel of Table 2 — Table 2: Quantitative evaluation of image deblurring comparing PnP methods using DRUNet with score-based variants obtained by replacing the denoiser with a diffusion model. Best values are color-coded for each metric.

Testing data		Input	SDPnP-DPIR	SGPnP-DPIR	SDPnP-ADMM	SGPnP-ADMM	SDPnP-PGM	SGPnP-PGM

Inpainting	PSNR $\uparrow$	$18.17$	$18.29$	$\mathbf{25.08}$	$20.51$	$\mathbf{24.12}$	$22.35$	$\mathbf{25.21}$
	SSIM $\uparrow$	$0.766$	$0.789$	$\mathbf{0.869}$	$0.811$	$\mathbf{0.852}$	$0.847$	$\mathbf{0.874}$
	LPIPS $\downarrow$	$0.289$	$0.259$	$\mathbf{0.109}$	$0.227$	$\mathbf{0.140}$	$0.141$	$\mathbf{0.108}$

Super-resolution	PSNR $\uparrow$	$24.61$	$27.22$	$\mathbf{29.76}$	$27.24$	$\mathbf{29.60}$	$28.75$	$\mathbf{29.90}$

Table 3: Quantitative comparison between deterministic and stochastic PnP using the same score-based prior on FFHQ (inpainting,

4\times

super-resolution, deblurring) and fastMRI (

4\times

CS-MRI). Best values are highlighted for each metric and inverse problem. Note how SGPnP leads to better performance than traditional PnP even using the same prior.

Figure 3: Qualitative visual comparison corresponding to the quantitative results in Table 3 . Using the same score-based prior, stochastic noise injection produces more realistic completions in challenging box inpainting and improved reconstructions for accelerated MRI.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Problem	Method	$\gamma$	$\tau$	$\sigma_{0}^{\text{cond}}$	$\sigma_{0}^{\text{inject}}$	Iter
Inpainting	SGPnP-ADMM	1.7	–	15	15	200
	SGPnP-DPIR	1.5	–	15	15	200
	SGPnP-PGM	0.22	0.4	20	20	200
	SDPnP-ADMM	1.32	–	5	0	50
	SDPnP-DPIR	2.0	–	0.3	0	10
	SDPnP-PGM	2.0	0.4	20	0	200

Appendix

Table 4: Hyperparameters for PnP-based reconstruction methods. SGPnP denotes the proposed stochastic generative plug-and-play framework with noise injection (

\sigma^{\text{inject}}>0

). SDPnP denotes the proposed score-based deterministic plug-and-play framework without noise injection (

\sigma^{\text{inject}}=0

). DPIR denotes a deterministic PnP method based on DRUNet denoisers.

Problem	Method	$\lambda$	$\gamma$	Iter
Inpainting	DiffPIR	2.0	1.0	200
	DPS	0.4	–	1000
Deblur	DiffPIR	1.0	1.0	200
	DPS	10.0	–	1000
SR	DiffPIR	1.0	1.0	200
	DPS	3.0	–	1000

Appendix

Table 5: Hyperparameters for diffusion-based inverse problem solvers.

		SGPnP-DPIR		SGPnP-ADMM		SGPnP-PGM
	Input	Match	Decouple	Match	Decouple	Match	Decouple

PSNR $\uparrow$	$22.43$	$30.99$	$\mathbf{32.91}$	$31.44$	$\mathbf{33.32}$	$30.20$	$\mathbf{32.54}$
SSIM $\uparrow$	$0.622$	$0.856$	$\mathbf{0.880}$	$0.854$	$\mathbf{0.895}$	$0.823$	$\mathbf{0.881}$
LPIPS $\downarrow$	$0.333$	$0.122$	$\mathbf{0.151}$	$0.137$	$\mathbf{0.111}$	$0.169$	$\mathbf{0.119}$

Appendix

Table 6: Ablation on noise-level coupling for fastMRI (

4\times

CS-MRI) using score-based PnP solvers. Results compare the matched setting (

\sigma^{\text{cond}}=\sigma^{\text{inject}}

) and the decoupled setting (

\sigma^{\text{cond}}\neq\sigma^{\text{inject}}

). Best values are highlighted for each metric.

	Input	DPIR ( $\sigma_{\text{low}}$ )	SNORE ( $\sigma_{\text{low}}$ )	SGPnP-PGM ( $\sigma_{\text{low}}$ )	SGPnP-PGM ( $\sigma_{\text{wide}}$ )

PSNR $\uparrow$	$18.17$	$18.42$	$18.50$	$18.27$	$\mathbf{25.21}$
SSIM $\uparrow$	$0.766$	$0.797$	$0.799$	$0.797$	$\mathbf{0.874}$
LPIPS $\downarrow$	$0.289$	$0.264$	$0.250$	$0.257$	$\mathbf{0.108}$

Appendix

Table 7: Ablation study on noise-level coverage for box inpainting. We compare deterministic PnP (DPIR), a stochastic PnP baseline (SNORE), and SGPnP-PGM restricted to the same low-noise regime used by classical denoisers, and SGPnP-PGM leveraging the full noise range available during score-based model training. Best values are highlighted for each metric.

Figure 4: Box inpainting results obtained with SGPnP-PGM from repeated runs on the same measurements. The method produces consistent reconstructions across runs.

Figure 5: Additional box inpainting results on more measurements. For this challenging task, DPIR and SNORE often produce incomplete reconstructions. SDPnP-PGM improves the result but still remains incomplete in some cases, whereas SGPnP-PGM produces more plausible image completions.

	SNORE	SGPnP (Ours)
Denoiser prior	Classical denoiser (e.g., DRUNet), trained for low-noise restoration.	Score-based diffusion prior (SBDM), trained across a wide range of noise levels.
Noise injection	Injects Gaussian noise and denoises at the same noise level.	Injects Gaussian noise while decoupling injected and conditioning noise levels to account for corruption introduced by data-consistency updates.
Theoretical role of noise	Enables convergence toward critical points.	Promotes escape from strict saddle points while enabling convergence toward critical points.
Severely ill-posed tasks	Typically yields modest improvements over deterministic PnP.	Enables more robust generative reconstruction via expressive SBDM priors and decoupled noise control.

Appendix

Table 8: Comparison of SNORE [ 27 ] and our PnP Score framework. Our approach provides an explicit score-based interpretation of PnP with diffusion priors and motivates noise injection as a mechanism for escaping undesirable stationary points.

Explore similar work

Plug-and-Play Diffusion Meets ADMM: Dual-Variable Coupling for Robust Medical Image Reconstruction
Feb 26, 2026Chenhe Du, Xuanyu Tian, Qing Wu +4Image Reconstruction Diffusion-Based Inverse Problems
ShanghaiTech University · Shanghai Jiao Tong University
$\mathrm{S}^3$ -PnP: Plug-and-Play Denoising for Scientific Imaging with Self-Supervised Score Priors
Oct 13, 2026Yukun Cui, Jiangshe Zhang, Shuang Xu +6Image Restoration Image Denoising
Xi’an Jiaotong University · Northwestern Polytechnical University · Nanyang Technological University +3
A Posterior-Dynamics Framework for Imaging Inverse Problems with Pretrained Diffusion Priors
Aug 15, 2026Zhaoqiang Liu, Tongyao Pang, Ruibing Wang +1Diffusion Models for Image Restoration Bayesian Inverse Problems
School of Computer Science and Engineering, University of Electronic Science and Technology of China, Chengdu 611731, China · Yau Mathematical Sciences Center, Shuangqing Complex, Tsinghua University, Beijing 100084, China

Feb 26, 2026J/K move · Enter open · S save

Plug-and-Play Diffusion Meets ADMM: Dual-Variable Coupling for Robust Medical Image Reconstruction

Chenhe Du, Xuanyu Tian, Qing Wu, Muyu Liu+3

ShanghaiTech University · Shanghai Jiao Tong University

Image Reconstruction Diffusion-Based Inverse Problems Diffusion MRI Alternating Direction Method of Multipliers

Open paper AskPDF arXiv

Plug-and-Play diffusion prior (PnPDP) frameworks have emerged as a powerful paradigm for solving imaging inverse problems by treating pretrained generative models as modular priors. However, we identify a critical flaw in prevailing PnP solvers (e.g., based on HQS or Proximal Gradient): they function as memoryless operators, updating estimates solely based on instantaneous gradients. This lack of historical tracking inevitably leads to non-vanishing steady-state bias, where the reconstruction fails to strictly satisfy physical measurements under heavy corruption. To resolve this, we propose Dual-Coupled PnP Diffusion (DC-PnPDP), which restores the classical dual variable to provide integral feedback, progressively enforce agreement between the data-consistency and prior. However, this rigorous geometric coupling introduces a secondary challenge: the accumulated dual residuals exhibit spectrally colored, structured artifacts that violate the Additive White Gaussian Noise (AWGN) assumption of diffusion priors, causing severe hallucinations. To bridge this gap, we introduce Spectral Homogenization (SH), a frequency-domain adaptation mechanism that modulates these structured residuals into statistically compliant pseudo-AWGN inputs. This effectively aligns the solver's rigorous optimization trajectory with the denoiser's valid statistical manifold. Extensive experiments on CT and MRI reconstruction demonstrate that our approach resolves the bias-hallucination trade-off, achieving state-of-the-art fidelity with significantly accelerated convergence. The code is available at https://github.com/duchenhe/DC-PnPDP

Abstract

Explore similar work

Plug-and-Play Diffusion Meets ADMM: Dual-Variable Coupling for Robust Medical Image Reconstruction

S3\mathrm{S}^3S3-PnP: Plug-and-Play Denoising for Scientific Imaging with Self-Supervised Score Priors

A Posterior-Dynamics Framework for Imaging Inverse Problems with Pretrained Diffusion Priors

$\mathrm{S}^3$ -PnP: Plug-and-Play Denoising for Scientific Imaging with Self-Supervised Score Priors