Stage-wise Distortion-Perception Traversal in Zero-shot Inverse Problems with Diffusion Models

Technical Summary: Stage-wise Distortion–Perception Traversal in Zero-shot Inverse Problems with Diffusion Models

Problem Statement
Bayesian inverse problems, which aim to reconstruct original signals from degraded observations, are fundamentally constrained by the distortion–perception (D–P) tradeoff. Distortion performance measures sample-wise reconstruction fidelity (e.g., via MSE or PSNR), while perception performance reflects the visual quality and distributional similarity to real data. Theoretical work has established that these objectives are inherently in tension: improving one often necessitates degrading the other. While diffusion models have achieved state-of-the-art (SOTA) results in zero-shot inverse problem solving without task-specific retraining, existing methods lack principled and computationally efficient strategies to traverse the D–P curve at inference time. Current approaches, such as adjusting sampling steps or sample averaging, are often heuristic or computationally expensive, and the mechanisms for controlling this tradeoff in diffusion-based zero-shot algorithms remain inadequately characterized.

Methodology
The authors propose a two-stage framework, termed MAP-RPS (Maximum A Posteriori - Re-noised Posterior Sampling), to enable flexible D–P traversal using a single pretrained diffusion model. The method is later extended to latent space as LMAP-RPS to leverage large-scale latent diffusion backbones.

1. Stage 1: Low-Distortion Initialization (MAP Estimation)
   The first stage approximates the Minimum Mean Squared Error (MMSE) solution, which corresponds to the optimal distortion point on the D–P curve. The authors analyze posteriors that are strongly log-concave and establish that the MAP estimator serves as a valid surrogate for the MMSE solution with a bounded approximation error. They derive a tractable gradient for the log-prior using the pretrained diffusion model's score function at a specific timestep $t_1$, allowing for efficient MAP estimation via stochastic gradient-based optimization. This yields a low-distortion starting point ($x_{MAP}$).

2. Stage 2: Perceptual Quality Improvement (Re-noised Posterior Sampling)
   The second stage traverses the D–P curve by progressively improving perceptual quality. The MAP estimate is re-noised up to a prescribed timestep $t_0$, and posterior sampling is performed from this noisy initialization. Theoretically, the authors show that the Wasserstein-2 ($W_2$) distance (a metric for perception error) between the generated distribution and the true data distribution is controlled by the re-noising timestep $t_0$. By varying $t_0$, the method enables a continuous interpolation between the low-distortion MAP solution (where $t_0=0$) and standard posterior sampling (where $t_0=T$).

Key Contributions

• Theoretical Framework: The paper provides theoretical analyses for both stages. It establishes an error bound for approximating the MMSE solution with the MAP estimator under strong log-concavity assumptions and derives an upper bound on the $W_2$ distance for the re-noised posterior sampling stage, proving that perception error can be controlled by adjusting $t_0$.
• Algorithm Design: The introduction of MAP-RPS and its latent-space variant, LMAP-RPS, offers a computationally efficient mechanism for D–P traversal. Unlike methods requiring repeated sampling or complex hyperparameter tuning, this approach requires only one MAP estimation procedure and one posterior sampling process.
• Latent Space Extension: By extending the framework to latent diffusion models (LMAP-RPS), the authors leverage powerful pretrained backbones (e.g., Stable Diffusion) for broader applicability to real-world, high-resolution tasks.

Experimental Results
Extensive experiments were conducted on the FFHQ dataset (for D–P curve analysis) and the MS-COCO dataset (for near-real-world inverse problems).

• D–P Traversal: On FFHQ tasks (denoising, super-resolution, inpainting, deblurring), both MAP-RPS and LMAP-RPS consistently achieved curves closer to the ideal lower-left corner (optimal distortion and perception) compared to existing methods like VSDPS and PSCGAN.
• Real-World Performance: On MS-COCO, LMAP-RPS achieved SOTA performance across nearly all tasks and metrics (PSNR, SSIM, LPIPS, FID) compared to nine recent latent diffusion-based baselines.
• Efficiency: LMAP-RPS demonstrated significant computational advantages. At $t_0=0$, it was significantly faster than all baselines. Even at $t_0=600$, it remained computationally efficient, often outperforming other methods in speed while maintaining superior reconstruction quality.

Significance and Claims
The paper claims that MAP-RPS and LMAP-RPS provide a novel perspective for realizing D–P traversal in zero-shot inverse problems. The authors emphasize that their method offers a principled, theoretically grounded, and computationally efficient alternative to existing heuristic strategies. By decoupling the distortion and perception objectives into two distinct stages, the framework allows practitioners to flexibly select a reconstruction that best fits specific application requirements without retraining models. The authors also highlight the strong practical potential of LMAP-RPS, noting its ability to deliver SOTA results on complex, real-world inverse problems with lower computational complexity than current state-of-the-art algorithms.

Limitations and Future Work
The authors acknowledge that their theoretical guarantees rely on the assumption of strongly log-concave (approximately unimodal) posteriors. While empirical results suggest the method remains effective in more general settings, the behavior under complex multimodal posteriors requires further investigation. Additionally, the current implementation relies on specific posterior sampling algorithms (like DPS) for Stage 2, which may introduce estimation errors; future work could integrate asymptotically consistent posterior sampling methods to further refine the D–P traversal.