Stochastic Generative Plug-and-Play Priors

Stochastic Generative Plug-and-Play Priors: Leveraging Diffusion Models for Robust Inverse Problem Solving

Introduction and Motivation

This paper establishes a theoretical and algorithmic foundation for integrating score-based diffusion models (SBDMs) as priors within the optimization-centric plug-and-play (PnP) framework for imaging inverse problems. The authors present a new stochastic generative PnP (SGPnP) framework that directly enables the use of pretrained SBDMs within iterative PnP solvers, with explicit noise injection mechanisms to mitigate statistical mismatches between intermediate iterates and the denoiser's training regime.

PnP methods have become a foundational tool in computational imaging, offering a mechanism to modularly combine arbitrary denoisers with data-consistency enforcing optimization. Classical CNN-based denoisers such as DRUNet are typically trained for low to moderate noise levels, constraining their effectiveness in severely ill-posed problems where the intermediate solutions are far outside this domain.

In contrast, SBDMs—such as variance-exploding (VE) and variance-preserving (VP) diffusion models—are trained across a full spectrum of noise levels by learning to estimate the score function (gradient of the log probability density) of noisy images. While SBDMs have been harnessed for generative sampling and posterior sampling in inverse problems, their systematic integration as direct priors in PnP optimization algorithms, detached from the burdens of reverse diffusion sampling, has remained theoretically ambiguous.

Key Contributions

The manuscript's main technical contributions are:

• Score-Based Interpretation of PnP: A rigorous mathematical formalism connecting classical iterative PnP updates to score-matching denoising. This provides a principled justification to deploy pretrained SBDM denoisers directly within PnP iterations, agnostic to the classical CNN prior's limitations.
• Stochastic Generative PnP (SGPnP) Framework: A novel stochastic PnP update in which noise is explicitly injected into the denoiser's input on each iteration. This adjustment (i) aligns the input statistics to the SBDM's training distribution, and (ii) injects stochasticity beneficial for exploration and saddle-point escape in the nonconvex PnP optimization landscape.
• Theoretical Guarantees: The first PnP theoretical result proving that stochastic noise injection ensures escape from strict saddle points and that an annealed noise schedule drives the iteration to stationary points of the true objective as the injected noise vanishes. All assumptions and intermediate constructs (variance preservation, denoiser type) are made explicit.

The SGPnP Update and Implementation

In contrast to reverse diffusion sampling or methods that embed generative steps within PnP, the SGPnP framework uses the SBDM denoiser as a "score-adapted" operator within the classical decoupled two-step scheme:

1. Data-consistency update: Enforces fidelity to the measurements, typically via a proximal or gradient step.
2. Stochastic denoising update: The current estimate (or associated variable) is stochastically perturbed by additive noise and then denoised at a user-controlled conditioning noise level by the SBDM.

This update can be generically formulated for a range of proximal-based and splitting-based optimization algorithms (e.g., ADMM, PGM, gradient descent), with flexibility in parameterizing the denoiser (VE or VP) and decoupling the "injected" and "conditioning" noise levels for optimal adaptation.

Theoretical Analysis: Escape from Saddle Points and Consistency

A key weakness of deterministic PnP methods with expressive priors is susceptibility to local suboptimality—being trapped at strict saddles or poor local minima. The paper's stochastic analysis demonstrates that, under standard regularity conditions and assuming the denoiser is an MMSE operator (as in Tweedie's formula for SBDMs), the random noise injection at the denoiser input not only smooths the objective but, crucially, injects energy along directions of negative curvature:

• Saddle Point Escape: Noise ensures, with high probability, that iterates leave neighborhoods of strict saddle points, leveraging variance lower-boundedness along these directions.
• Convergence under Annealing: Under suitable conditions, as the injected noise is reduced (annealed), the set of limiting points aligns with stationary points of the original non-smoothed objective, yielding theoretical consistency.

These are the first formal guarantees of saddle-point avoidance and robust convergence for a PnP algorithm operating with highly nonconvex, generative diffusion priors.

Empirical Validation

The SGPnP framework is empirically validated on high-fidelity inverse tasks: large-hole inpainting, deblurring, super-resolution on FFHQ, and accelerated multi-coil MRI (fastMRI). Comparative baselines include deterministic PnP (with DRUNet denoisers), stochastic PnP (SNORE), and diffusion posterior sampling methods.

Key findings and empirical observations:

• Superior Inpainting and MRI Reconstruction: SGPnP demonstrates state-of-the-art scores across PSNR, SSIM, and LPIPS, most notably on severe box inpainting where all competing PnP methods fail to produce plausible solutions.

  Figure 1: Deterministic and stochastic PnP methods yield incomplete or unrealistic inpainting in the presence of large masks, while the SGPnP approach successfully hallucinates plausible semantic completions.

• Qualitative Consistency and Diversity: Stochasticity of SGPnP does not result in spurious outputs; reconstructions remain highly plausible across repeated runs, confirming the algorithm's stability and efficacy in leveraging generative priors.

  Figure 2: Box inpainting results from repeated SGPnP runs show solution consistency despite stochasticity.

• Ablation on Score Prior Usage: When only score-based denoisers are used in deterministic PnP (SDPnP), performance increases relative to DRUNet, but the stochastic injection in SGPnP is necessary to close the gap with diffusion sampling-based solvers, particularly in ill-posed settings.

  Figure 3: Qualitative comparisons illustrate deterministic and stochastic PnP with score-based priors, highlighting the essential role of noise injection for robust recovery.

• Generalization Across Problem Types: SGPnP results transfer across natural images and MRI without task-specific adaptation, evidencing the flexibility of the score-based prior and optimization-centric integration.

Practical and Theoretical Implications

• Unified Score-Based PnP Framework: This work bridges variational, proximal optimization and the stochastic generative modeling perspectives, making SBDMs accessible to optimization-based inverse solvers without reverse diffusion.
• Stochasticity as Regularization and Exploration: The framework illustrates a rigorous benefit of noise injection—specifically for nonconvex, high-dimensional tasks where deterministic optimization frequently fails.
• Noise Decoupling: Allowing user control over both injected and conditioning noise allows practical adjustment for the artifacts induced by data-consistency steps, going beyond prior PnP noise or diffusion methods.
• Foundation for Future Directions: This approach prepares the ground for scalable, principled posterior sampling in inverse problems, hybrid stochastic-deterministic inference, and advanced use of diffusion generative models within broadly defined optimization algorithms.

Conclusion

The SGPnP framework makes expressive generative diffusion models directly available to PnP-based imaging inverse problem solvers, resolves longstanding distribution mismatch issues via explicit noise injection, and provides both theoretical and empirical evidence that such stochastic PnP methods outperform previous deterministic and stochastic PnP approaches as well as generative sampling-based solutions in highly ill-posed regimes. Critically, this work synthesizes the fields of score-based generative modeling and iterative optimization, underscoring noise injection as a central mechanism for robust and reliable signal recovery across practical computational imaging applications.

Figure 4: Compared to DPIR and SNORE, SGPnP-PGM achieves substantially more plausible image completions in challenging box inpainting cases, even where other methods yield incomplete or physically implausible results.