Local MAP Sampling for Diffusion Models

2025-10-07

AI Summary

Optimization-based diffusion solvers lack a rigorous probabilistic foundation, exhibit ambiguous connections to maximum a posteriori (MAP) estimation and diffusion posterior sampling (DPS), and suffer from instability and poor interpretability. This paper proposes a Local MAP Sampling framework that iteratively solves local MAP subproblems along the diffusion trajectory—thereby systematically embedding optimization-based solvers within the Bayesian inference paradigm for the first time. We introduce an interpretable covariance approximation, a stability-enhancing objective reconstruction mechanism, and a gradient approximation technique applicable to non-differentiable forward operators. Evaluated on image deblurring, JPEG artifact removal, quantization recovery, and inverse scattering, our method achieves state-of-the-art performance, with average PSNR gains exceeding 2 dB (≥1.5 dB in several tasks).

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Learning & Optimization for CV

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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.

Problem

• Develops probabilistic framework for inverse problem reconstruction
• Unifies optimization-based methods with diffusion posterior sampling
• Achieves state-of-the-art performance in image restoration tasks

Innovation

• Iteratively solves local MAP subproblems
• Uses probabilistically interpretable covariance approximation
• Reformulates objective for stability and interpretability