Local MAP Sampling for Diffusion Models

Basic Information

• Paper ID: 2510.07343
• Title: Local MAP Sampling for Diffusion Models
• Authors: Shaorong Zhang (UC Riverside), Rob Brekelmans (Vector Institute), Greg Ver Steeg (UC Riverside)
• Classification: cs.GR cs.AI eess.IV
• Publication Date/Venue: Preprint (Under review)
• Paper Link: https://arxiv.org/abs/2510.07343

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 objective of inverse problem solving is not to cover the posterior distribution, but to recover the most accurate reconstruction results. Optimization-based diffusion solvers typically excel at this task, despite lacking clear probabilistic foundations. This paper introduces Local MAP Sampling (LMAPS), a novel inference framework that iteratively solves local MAP subproblems along the diffusion trajectory. This perspective clarifies the connections between these methods and global MAP estimation and DPS, providing a unified probabilistic interpretation for optimization-based approaches. Building on this foundation, we develop practical algorithms featuring probabilistically interpretable covariance approximations, reformulated objectives with improved stability and interpretability, and gradient approximations for non-differentiable operators.

Research Background and Motivation

Problem Background

Diffusion models face two major challenges in inverse problem solving:

1. Objective Mismatch: DPS aims to sample from the posterior distribution $p(x_0|y)$, but the actual objective of inverse problem solving is to obtain the most accurate reconstruction, not diverse samples
2. Missing Theoretical Foundation: Optimization-based diffusion solvers (such as Resample, DiffPIR, DCDP, etc.) demonstrate excellent performance but lack clear probabilistic theoretical foundations

Research Motivation

• Practical Orientation: Inverse problem evaluation protocols typically compare against a single ground truth reference and do not reward coverage or diversity
• Theoretical Unification: There is a need to provide probabilistic interpretations for optimization methods and clarify their relationships with MAP estimation and DPS
• Performance Improvement: Achieve better reconstruction performance while maintaining theoretical foundations

Core Contributions

1. Theoretical Contribution: Proposes the Local MAP Sampling (LMAPS) framework that iteratively solves local MAP subproblems along the diffusion trajectory, analyzes its relationship with global MAP and DPS, and unifies TMPD and optimization-based inverse problem methods
2. Methodological Contribution:
   • Provides probabilistically interpretable covariance approximations to replace heuristic choices in existing solvers
   • Introduces objective reformulation to achieve interpretable parameters and improved stability
   • Develops gradient approximation strategies for non-differentiable operators
3. Experimental Contribution: Validates the method on 10 image restoration tasks and 3 scientific inverse problems, achieving best results in 46/60 FFHQ/ImageNet cases, with ≥2dB PSNR improvements on motion deblurring, JPEG restoration, and quantization tasks

Method Details

Task Definition

The inverse problem aims to recover an unknown image or signal $x_0 \in \mathbb{R}^n$ from a prior distribution $\pi(x_0)$ and noisy measurements $y \in \mathbb{R}^m$: $y = H(x_0) + z$ where $H(\cdot): \mathbb{R}^n \to \mathbb{R}^m$ is the forward operator and $z \sim \mathcal{N}(0, \sigma_y^2 I)$ is measurement noise.

Core Theoretical Framework

Global MAP vs. Local MAP

Global MAP directly optimizes the complete posterior: $x_0^{\text{MAP}} := \arg\max_{x_0} p(x_0|y)$

Local MAP solves a conditional optimization problem at each time step $t$: $x_0^*(t, x_t, y) := \arg\max p(x_0|x_t, y)$$x_{t-\Delta t} = g(x_0^*, x_t, \epsilon), \quad \epsilon \sim \mathcal{N}(0,I)$

Relationship with DPS and DAPS

• DPS: Uses conditional mean $E[x_0|x_t, y]$
• DAPS: Samples from $p(x_0|x_t, y)$
• LMAPS: Uses conditional mode $\arg\max p(x_0|x_t, y)$

Key Insight: DPS and LMAPS are equivalent only when $p(x_0|x_t, y)$ is Gaussian.

Practical Algorithm Design

Covariance Approximation

Employs isotropic approximation: $\Sigma_{0|t} \approx \frac{k}{\text{SNR}} I, \quad \text{SNR} := \frac{\alpha_t^2}{\sigma_t^2}$

Objective Function Reformulation

Reformulates the original objective: $x_0^* = \arg\min \left\{\frac{\text{SNR}}{k}\|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2}\|y - H(x_0)\|^2\right\}$

into: $x_0^* = \arg\min \left\{(1-\mu_t)\frac{1}{2}\|x_0 - m_{0|t}\|^2 + \mu_t k_2\|y - H(x_0)\|^2\right\}$

where $\mu_t = \frac{\sigma_t^2}{\sigma_t^2 + k_1^2} \in (0,1)$, achieving:

• Convex Combination Interpretation: Weights of $(1-\mu_t)$ and $\mu_t$
• Automatic Annealing: Transitions from measurement-driven to prior-driven as $\sigma_t^2$ decreases
• Numerical Stability: Avoids extreme SNR scaling

Non-differentiable Operator Handling

For non-differentiable tasks such as JPEG restoration and quantization, uses surrogate gradients: $\nabla_{x_0}\|y - H(x_0)\|^2 \approx 2J_{H'}(x_0)^T(H(x_0) - y)$

For quantization, adopts $H'(x_0) = x_0$, simplifying to: $\nabla_{x_0}\|y - H(x_0)\|^2 \approx 2(H(x_0) - y)$

Experimental Setup

Datasets

• Image Restoration: FFHQ 256×256 and ImageNet 256×256, using 100 test images each
• Scientific Inverse Problems: InverseBench dataset including fluorescence microscopy images (linear inverse scattering), GRMHD data (black hole imaging), fastMRI knee data (compressed sensing MRI)

Evaluation Metrics

• Image Restoration: PSNR, SSIM, LPIPS
• Scientific Inverse Problems: Primarily PSNR, supplemented with task-specific metrics

Comparison Methods

Includes DDNM, DDRM, ΠGDM, DPS, LGD, PnP-DM, FPS, MCG-diff, RedDiff, DAPS, DiffPIR, DCDP, DMPlug and 12 other methods

Implementation Details

• Diffusion steps: typically 200
• Gradient update steps: 20-200 (task-dependent)
• Learning rate: 0.01-1.0 (task-related)
• Parameters $k_1$: 0-10, $k_2$: 0.01-30000

Experimental Results

Main Results

Image Restoration Tasks

In Table 1, LMAPS achieves best performance in 49 out of 60 results:

• Motion Deblurring: 32.62 dB on FFHQ vs. 29.66 dB for DAPS (+2.96 dB)
• JPEG Restoration: 27.25 dB on FFHQ vs. 25.04 dB for ΠGDM (+2.21 dB)
• Quantization: 29.51 dB on FFHQ vs. 25.82 dB for ΠGDM (+3.69 dB)

Scientific Inverse Problems

In Table 2, LMAPS achieves best PSNR on all tasks:

• Linear Inverse Scattering (NR=360): 38.07 dB vs. 36.56 dB for RED-diff (+1.51 dB)
• Linear Inverse Scattering (NR=180): 37.19 dB vs. 35.41 dB for RED-diff (+1.78 dB)
• Linear Inverse Scattering (NR=60): 30.75 dB vs. 27.07 dB for RED-diff (+3.68 dB)

Ablation Studies

Figure 4 shows the trade-off between optimization steps and diffusion steps:

• Best performance typically observed at NFE=200-500
• Increasing optimization steps per diffusion step significantly improves performance
• LMAPS achieves similar performance to SITCOM (600 NFEs) with fewer computational resources

Computational Efficiency

Table 3 shows LMAPS sampling time on deblurring tasks:

• LMAPS (200 diffusion steps, 100 optimization steps): 61 seconds/image, 30.88 dB
• DAPS (200 diffusion steps, 100 optimization steps): 110 seconds/image, 29.19 dB
• SITCOM (600 steps): 73 seconds/image, 29.93 dB

Related Work

Diffusion Posterior Sampling

DPS and its variants solve inverse problems by directly sampling from the posterior distribution $p(x_0|y)$, including methods such as TMPD, DDNM, ΠGDM, etc.

Optimization-based Methods

Methods such as Resample, DiffPIR, DCDP, DMPlug solve inverse problems through alternating denoising, optimization, and resampling, demonstrating excellent performance but lacking theoretical foundations.

MAP Estimation Methods

Recent work has begun focusing on MAP estimation under diffusion priors, but is primarily limited to linear inverse problems.

Conclusions and Discussion

Main Conclusions

1. Theoretical Unification: LMAPS provides a unified probabilistic interpretation for optimization-based diffusion methods
2. Performance Improvement: Achieves significant PSNR improvements on multiple tasks, particularly on challenging non-linear and non-differentiable tasks
3. Computational Efficiency: Demonstrates better computational efficiency compared to existing methods

Limitations

1. Convergence: Local MAP sequences do not necessarily converge to global MAP
2. Diversity: LMAPS may produce less output diversity compared to DPS
3. Hyperparameter Sensitivity: Requires task-specific tuning of parameters $k_1$ and $k_2$

Future Directions

The paper notes that the critical role of global MAP in Bayesian inference has been greatly overlooked, and efficiently solving global MAP under diffusion priors remains an open challenge. While MAP may reduce output diversity by concentrating on a single mode, it provides higher determinism and better alignment with observed data.

In-depth Evaluation

Strengths

1. Significant Theoretical Contribution: First to provide clear probabilistic interpretation for optimization-based diffusion methods
2. Comprehensive Experiments: Covers 10 image restoration tasks and 3 scientific inverse problems with complete experimental setup
3. Obvious Performance Improvements: Achieves significant improvements of over 2dB on multiple challenging tasks
4. Practical Method: Provides effective strategies for handling non-differentiable operators
5. Clear Writing: Both theoretical analysis and method descriptions are well-presented

Weaknesses

1. Limited Theoretical Analysis Depth: While providing probabilistic interpretation, analysis of convergence and theoretical guarantees is relatively limited
2. Hyperparameter Complexity: Requires task-specific tuning of multiple hyperparameters, potentially affecting method generalization
3. Computational Overhead: While more efficient than some methods, still requires multiple gradient updates at each time step
4. Evaluation Limitations: Primarily focuses on reconstruction quality with insufficient assessment of uncertainty quantification

Impact

1. Academic Value: Provides new theoretical perspective for diffusion model applications in inverse problems
2. Practical Value: Has direct application value in image restoration and scientific computing domains
3. Inspirational: May inspire further research on MAP estimation for diffusion models

Applicable Scenarios

• Inverse problems requiring high-quality reconstruction rather than diversity
• Tasks involving non-differentiable forward operators (e.g., JPEG restoration, quantization)
• Inverse problem solving in scientific computing
• Real-time applications with certain computational efficiency requirements

References

The paper cites important works in diffusion models, inverse problem solving, and Bayesian inference, including original papers of key methods such as DPS, DAPS, and TMPD, providing a solid literature foundation for related research.