Local MAP Sampling for Diffusion Models

TL;DR

This work addresses inverse problems under diffusion priors by reframing inference as a sequence of local MAP subproblems along the diffusion trajectory. The proposed Local MAP Sampling (LMAPS) unifies aspects of global MAP, TPS-based approaches, and DPS, and introduces a probabilistic covariance approximation, a principled reformulation of the objective, and a gradient strategy for non-differentiable forward operators. Empirically, LMAPS achieves state-of-the-art results across 10 image restoration tasks and 3 scientific inverse problems, with notable PSNR gains on motion deblurring, JPEG restoration, and quantization, and strong performance on linear LIS benchmarks. The method provides a principled, efficient alternative to posterior sampling for high-fidelity reconstructions and lays groundwork for future MAP-focused diffusion-based inference in complex forward models.

Abstract

Diffusion Posterior Sampling (DPS) provides a principled Bayesian approach to inverse problems by sampling from $p(x_0 \mid y)$. However, in practice, the goal of inverse problem solving is not to cover the posterior but to recover the most accurate reconstruction, where optimization-based diffusion solvers often excel despite lacking a clear probabilistic foundation. We introduce Local MAP Sampling (LMAPS), a new inference framework that iteratively solving local MAP subproblems along the diffusion trajectory. This perspective clarifies their connection to global MAP estimation and DPS, offering a unified probabilistic interpretation for optimization-based methods. Building on this foundation, we develop practical algorithms with a probabilistically interpretable covariance approximation, a reformulated objective for stability and interpretability, and a gradient approximation for non-differentiable operators. Across a broad set of image restoration and scientific tasks, LMAPS achieves state-of-the-art performance, including $\geq 2$ dB gains on motion deblurring, JPEG restoration, and quantization, and $>1.5$ dB improvements on inverse scattering benchmarks.

Figures (14)

• Figure 1: Comparison of LMAPS with other methods. (a). The relationship between different alignment approaches; (b). The generation process of unconditional diffusion model; (c). The generation process of LMAPS.
• Figure 2: DPS
• Figure 3: Comparison of LMAPS, DPS, DAPS and Global MAP on 2D synthetic data. LMAPS is less likely to generate samples in the between-mode regions or low-density regions.
• Figure 4: Ablation study on optimization steps vs. diffusion steps (NFEs) for Gaussian Deblurring.
• Figure 5: Visualization of CS-MRI restoration ($4\times$ raw).