Beyond MMSE: Enhancing PnP Restoration with ProxiMAP

TL;DR

ProxiMAP is a novel iterative MAP approximation method that improves PnP image restoration by maintaining the denoiser's reliability through noise schedule control, outperforming traditional MMSE-based approaches.

Contribution

The paper introduces ProxiMAP, a modular method that enhances PnP restoration by aligning residual noise with the denoiser's training noise, enabling implicit early stopping and better results.

Findings

ProxiMAP consistently improves image quality across various tasks.

The hybrid variant achieves similar or better results at lower computational cost.

ProxiMAP outperforms standard MMSE denoisers in PnP frameworks.

Abstract

Plug-and-Play (PnP) methods have become standard tools for solving imaging inverse problems by replacing the intractable maximum a posteriori (MAP) denoiser with the MMSE one. While this mismatch has been widely treated as unavoidable, recent works have sought to close this gap by targeting the MAP with diffusion-model scores. We show this is problematic in practice: learned scores do not match the true ones, so MAP-targeting iterations converge to cartoon-like images rather than realistic ones, and better results are obtained by stopping short of convergence. We turn this observation into a design principle and introduce ProxiMAP, an iterative MAP approximation whose noise schedule keeps the iterate's residual noise matched to the denoiser's training noise. This keeps the denoiser in-distribution where its score is reliable, and yields implicit early stopping that avoids the failure mode above. ProxiMAP is a modular drop-in replacement for MMSE denoisers in standard PnP algorithms and consistently sharpens reconstructions across deblurring, inpainting, super-resolution, and phase retrieval. Building on the same principle, we propose a hybrid variant that applies ProxiMAP only in the late iterations of PnP, where the denoiser is most reliable -- matching or exceeding the full-replacement variant at a fraction of the cost.