🤖 machine learning

Beyond MMSE: Enhancing PnP Restoration with ProxiMAP

This paper introduces ProxiMAP, a modular Plug-and-Play framework that approximates maximum a posteriori (MAP) restoration by aligning residual noise with the denoiser's training distribution to ensure in-distribution reliability and implicit early stopping, thereby avoiding the cartoon-like artifacts of traditional MAP-targeting methods while consistently enhancing image reconstruction quality across various inverse problems.

Original authors: Kenta Vert, Giacomo Meanti, Scott Pesme, Michael Arbel, Julien Mairal

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kenta Vert, Giacomo Meanti, Scott Pesme, Michael Arbel, Julien Mairal

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to restore a blurry, noisy photograph. You have a very smart AI assistant (a "denoiser") whose job is to guess what the original, clean picture looked like.

For years, the standard way to use this assistant was to ask it: "What is the average of all the possible clean pictures that could look like this blurry mess?"

The paper argues that this "average" approach has a major flaw: it tends to smooth out the picture too much. It loses the fine details, like the texture of hair or the sharp edge of a building, making everything look like a soft, plastic toy.

The researchers tried a different approach: instead of asking for the average, they asked the AI: "What is the single most likely specific picture that could look like this?" (This is called finding the "MAP").

However, they discovered a trap. When they pushed the AI to find this "most likely" picture, the AI started hallucinating. Because the AI was trained on millions of photos, it got stuck in a weird loop where it produced images that looked like cartoons—flat colors, exaggerated lines, and fake textures. The AI was technically finding the "most likely" answer according to its own internal logic, but that logic was slightly broken for this specific task.

The Solution: "ProxiMAP" (The Goldilocks Strategy)

The authors realized the problem wasn't the goal (finding the best picture), but the timing. The AI is only reliable when the noise level in the picture matches the noise level it was trained on. If you ask it to work on a picture that is too clean or too noisy compared to its training, it starts guessing wildly and creating those cartoon artifacts.

They created a new method called ProxiMAP. Think of it like a guided tour for the AI:

The Matchmaker: Instead of letting the AI wander freely, ProxiMAP carefully controls the "noise level" at every single step of the restoration process. It ensures the picture the AI is looking at always has the exact amount of "grain" or "blur" that the AI is an expert at handling.
The Early Stopper: The most important trick is knowing when to stop. In the past, people let the AI run until it thought it was "done." ProxiMAP stops the process just before the AI gets confused and starts making cartoons. It stops while the AI is still in its "comfort zone," where it is most reliable.

The "Hybrid" Shortcut (Fast ProxiMAP)

Running this careful, step-by-step guide for the entire restoration process is computationally expensive (it takes a lot of computer power).

The authors realized that the AI only needs this careful guidance at the very end, when the picture is almost clean and the details are being refined. In the beginning, when the picture is still very blurry, a standard, rougher guess is fine.

So, they created a Hybrid version:

Start: Use the standard, fast, "average" method to get the general shape of the image.
End: Switch to the careful ProxiMAP method for the final few steps to sharpen the details and remove the last bit of blur without turning the image into a cartoon.

The Results

By using this method, the researchers found they could:

Sharpen images significantly better than the old "average" method.
Avoid the cartoon effect that plagued the "most likely" method.
Save time and money by only using the expensive, careful method for the last few steps of the process.

They tested this on various tasks like removing motion blur, filling in missing parts of an image (inpainting), and making low-resolution images high-resolution. In all cases, the new method produced crisper, more realistic-looking photos than the previous best tools.

In short: They fixed a broken tool by teaching it exactly when to work and when to stop, ensuring it stays in its "zone of expertise" to produce realistic, sharp images without the weird cartoon side effects.

Technical Summary: Beyond MMSE: Enhancing PnP Restoration with ProxiMAP

Problem Statement

Plug-and-Play (PnP) methods have become standard for solving imaging inverse problems (e.g., deblurring, inpainting) by replacing the intractable Maximum A Posteriori (MAP) denoiser in proximal splitting algorithms with a learned Minimum Mean Square Error (MMSE) denoiser. While effective, this substitution introduces a fundamental mismatch: the MMSE estimator approximates the posterior mean, which tends to produce oversmoothed reconstructions due to "regression-to-the-mean," whereas the variational formulation of PnP theoretically requires the MAP estimator.

Recent attempts to close this gap by targeting the MAP directly using scores from diffusion models have proven problematic in practice. The paper demonstrates that when using learned scores, MAP-targeting iterations do not converge to realistic images. Instead, they drift toward high-density regions of the learned distribution that do not align with the true data distribution, resulting in "cartoon-like" artifacts (flat colors, exaggerated edges). The paper posits that this failure is not intrinsic to MAP estimation itself but is a consequence of score approximation errors in learned models.

Methodology: ProxiMAP

The authors introduce ProxiMAP, an iterative MAP approximation algorithm designed to operate within the reliability regime of learned denoisers. The core design principle is noise matching: at every iteration $k$ , the residual noise level of the current iterate must match the noise level $\sigma_k$ at which the denoiser was trained and queried.

Key Mechanisms

Noise Schedule and Residual Tracking: Unlike standard gradient descent on the proximal objective where the noise schedule is decoupled from the iterate's state, ProxiMAP explicitly tracks the residual noise $\sigma_k$ of the iterate $x_k$ . The algorithm enforces the constraint that the denoiser is always queried with an input whose noise statistics match its training distribution.
Implicit Early Stopping: By maintaining this match, the algorithm prevents the iterate from drifting into regions where the learned score is unreliable (the "failure regime"). This acts as a principled early stopping mechanism. The iteration slows down as it approaches the data manifold, avoiding the mismatched queries that cause the cartoon effect.
Algorithmic Formulation: ProxiMAP modifies the smoothed gradient descent recursion (based on Pesme et al., 2025) to ensure $\sigma_k$ (the noise level of the iterate) equals the conditioning noise of the denoiser. The step size $\gamma_k$ and noise schedule are derived to satisfy this constraint, ensuring the denoiser remains "in-distribution."
Hybrid Variant (Fast ProxiMAP): Recognizing that full replacement of MMSE with ProxiMAP in every PnP iteration is computationally expensive, the authors propose a hybrid approach. This variant uses standard MMSE denoisers in the early iterations (where the iterate is far from the manifold and score accuracy is less critical) and switches to ProxiMAP only in the late iterations (where the iterate is close to the manifold and the denoiser is most reliable). This achieves comparable or superior results with a fraction of the computational cost.

Key Contributions

Diagnosis of MAP Estimation with Learned Priors: Through controlled comparisons using Gaussian Mixture Models (GMMs) with exact scores versus diffusion models with learned scores, the paper establishes that the "cartoon effect" in MAP targeting stems from the implicit biases of diffusion training displacing high-density regions, not from the MAP estimation process itself.
ProxiMAP Algorithm: A modular, drop-in replacement for MMSE denoisers in standard PnP frameworks. It is derived from a single principle (noise matching) that keeps the denoiser in-distribution and yields implicit early stopping.
Systematic Improvements: The method demonstrates consistent improvements in perceptual quality (LPIPS) while maintaining competitive distortion (PSNR) across multiple PnP algorithms (DPIR, DiffPIR, DAPS), tasks (deblurring, inpainting, super-resolution, phase retrieval), and datasets (FFHQ, ImageNet).
Efficient Hybrid Variant: The "Fast ProxiMAP" variant matches or exceeds the performance of full-replacement ProxiMAP at a significantly reduced computational cost (e.g., reducing function evaluations from 160 to 27 in some configurations).

Experimental Results

The authors evaluated ProxiMAP on FFHQ and ImageNet datasets across various inverse problems.

Standalone Denoising: ProxiMAP outperforms conditional MMSE and other deterministic diffusion samplers (DDIM, Flow Matching, InDI) on the perception-distortion trade-off curve, achieving lower LPIPS (sharper images) for comparable PSNR.
PnP Integration: When integrated into DPIR, DiffPIR, and DAPS, ProxiMAP variants consistently shift the Pareto front toward the optimal corner (high PSNR, low LPIPS).
Hybrid Efficiency: The Fast ProxiMAP variant (switching to ProxiMAP only in the final PnP iteration) achieves results comparable to the full-replacement variant but with drastically fewer network evaluations. For instance, in DiffPIR, it reduces the number of function evaluations (NFEs) from 160 to 27 while maintaining or improving reconstruction quality.
Failure Mode Resolution: In cases where full ProxiMAP replacement led to high-frequency artifacts (e.g., in super-resolution on FFHQ), the hybrid variant successfully resolved these issues by relying on the robust MMSE estimator during the early, unstable iterations.

Significance and Claims

The paper claims that ProxiMAP offers a practical solution to the long-standing mismatch between the MMSE denoisers used in PnP and the MAP estimators required by the underlying variational formulation. Rather than attempting to force convergence to a potentially ill-defined MAP of a learned distribution, ProxiMAP leverages the reliability of learned scores by keeping the optimization trajectory within the denoiser's training distribution.

The authors emphasize that their contribution is a modular component that enhances existing PnP families without requiring retraining of the underlying diffusion models or the development of new specialized inverse problem solvers. They note that while ProxiMAP improves perceptual quality, the reconstructions still inherit the biases of the pre-trained diffusion models, a caveat particularly relevant for safety-critical applications where visual plausibility should not be the sole determinant of correctness. The work suggests that managing the relationship between the iterate's noise level and the denoiser's training regime is a more effective design lever than pursuing exact MAP convergence with inexact scores.