Stochastic Generative Plug-and-Play Priors

1. Problem Statement

Image reconstruction from incomplete and noisy measurements (inverse problems) is fundamental to fields like MRI, super-resolution, and inpainting. Plug-and-Play (PnP) methods have become a standard approach by alternating between data consistency steps and a learned denoiser that acts as a prior.

However, PnP methods face significant limitations in severely ill-posed inverse problems (e.g., large-mask image inpainting):

• Distribution Mismatch: Intermediate PnP iterates often contain structured artifacts rather than the additive Gaussian noise for which standard denoisers are trained.
• Limited Noise Regime: Traditional PnP denoisers (e.g., DRUNet) are typically trained only for low-noise restoration, failing when the input requires "hallucinating" significant missing information.
• Saddle Points: Deterministic optimization in non-convex landscapes often gets trapped in strict saddle points, leading to suboptimal reconstructions.

While Score-Based Diffusion Models (SBDMs) offer powerful generative priors trained across a wide range of noise levels, integrating them directly into PnP without resorting to full reverse-diffusion sampling (which is computationally expensive and distinct from optimization) has remained an open challenge.

2. Methodology

The authors propose a framework that bridges the gap between classical PnP optimization and modern score-based generative models.

A. Score-Based Interpretation of PnP

The paper establishes a mathematical link between PnP iterations and score-based denoising using Tweedie's formula.

• Adaptation: They derive formulas to map pre-trained SBDMs (both Variance-Exploding and Variance-Preserving types) into the PnP denoising step.
• Parameter Matching: Since SBDMs are conditioned on discrete time steps ($t$) while PnP uses continuous noise levels ($\sigma$), the authors introduce a linear interpolation strategy to map the PnP noise level to the specific time step required by the diffusion model.
• SDPnP (Score-based Deterministic PnP): This baseline replaces the standard denoiser in PnP algorithms (ADMM, HQS, PGM) with the adapted SBDM, allowing the use of high-noise priors without reverse sampling.

B. Stochastic Generative PnP (SGPnP)

To address the distribution mismatch and optimization stagnation, the authors introduce noise injection into the denoiser input.

• Mechanism: Before applying the SBDM denoiser, Gaussian noise ($\sigma_{inject} n$) is added to the intermediate iterate ($z_k$).
• Decoupled Noise Levels: A key innovation is allowing the injected noise level ($\sigma_{inject}$) to differ from the conditioning noise level ($\sigma_{cond}$) of the denoiser. This accounts for the fact that intermediate iterates contain corruption from both the injected noise and residual artifacts from data-consistency steps.
• Algorithm: The framework is implemented in standard PnP variants (SGPnP-ADMM, SGPnP-PGM, SGPnP-RED), transforming the deterministic denoising step into a controlled stochastic transition.

3. Key Contributions

1. Theoretical Link: The paper provides the first principled derivation showing how pre-trained SBDMs can be used directly as denoisers within PnP optimization loops, bypassing the need for reverse diffusion sampling at every step.
2. SGPnP Framework: Introduction of a stochastic framework that injects noise to align intermediate iterates with the SBDM's training distribution and to facilitate optimization dynamics.
3. Theoretical Guarantees:
   • Saddle Point Escape: The authors prove that under specific conditions, the injected noise induces a stochastic gradient that satisfies the "Correlated Negative Curvature" (CNC) condition, ensuring the algorithm escapes strict saddle points of the Gaussian-smoothed objective.
   • Convergence: They prove that under an annealed noise schedule (where noise vanishes over iterations), the iterates converge to a stationary point of the original, un-smoothed Maximum A Posteriori (MAP) objective.
4. Decoupled Noise Strategy: Demonstrating that separating the injection noise from the conditioning noise significantly improves stability and performance in ill-posed tasks.

4. Experimental Results

The method was evaluated on FFHQ (natural faces) and fastMRI (multi-coil brain scans) across tasks like box inpainting, super-resolution, deblurring, and compressed sensing MRI.

• Performance vs. Baselines: SGPnP consistently outperforms:
  • Deterministic PnP (DPIR, PnP-ADMM): Significant gains in PSNR, SSIM, and LPIPS.
  • Stochastic PnP (SNORE): While SNORE adds noise, it relies on low-noise denoisers. SGPnP's use of SBDMs yields superior results, especially in large-mask inpainting where SNORE fails to generate plausible content.
  • Diffusion Solvers (DPS, DiffPIR): SGPnP achieves competitive or superior performance to dedicated diffusion-based solvers while maintaining the efficiency of an optimization-based framework.
• Ablation Studies:
  • Score Adaptation: Simply replacing the CNN denoiser with an SBDM (SDPnP) improves results, but adding noise injection (SGPnP) provides further substantial gains.
  • Noise Decoupling: Using different values for $\sigma_{inject}$ and $\sigma_{cond}$ yields better stability than matching them (as done in SNORE).
  • Noise Range: Restricting the SBDM to low-noise regimes negates its benefits; leveraging the full noise spectrum is critical for success.

5. Significance

This work represents a significant step forward in computational imaging by unifying optimization-based PnP with generative diffusion priors.

• Robustness: It solves the "severely ill-posed" problem where traditional methods fail to hallucinate missing structures, enabling high-quality reconstruction in tasks like large-mask inpainting.
• Efficiency: Unlike methods that require full reverse-diffusion sampling (which is slow), SGPnP retains the speed of iterative optimization while leveraging the expressive power of diffusion models.
• Theoretical Foundation: By providing rigorous proofs for saddle-point escape and convergence, the paper moves PnP from a heuristic engineering approach to a theoretically grounded framework capable of navigating complex, non-convex energy landscapes.

In summary, SGPnP demonstrates that injecting noise into PnP iterations, when combined with score-based priors, creates a powerful, robust, and theoretically sound method for solving the most challenging image reconstruction problems.