Technical Summary: Local MAP Sampling for Diffusion Models

Problem Statement

Diffusion Posterior Sampling (DPS) has emerged as a powerful framework for solving inverse problems by sampling from the posterior distribution $p(x_0 | y)$, offering principled uncertainty quantification and multi-modality. However, many practical applications in imaging and science prioritize accurate point estimation over full distributional sampling. The Maximum A Posteriori (MAP) estimator, defined as $\arg \max p(x_0 | y)$, remains a standard reconstruction objective.

Existing optimization-based methods (e.g., Resample, DiffPIR, DCDP) attempt to improve reconstruction performance by alternating between denoising and optimization steps. However, the theoretical connection between these iterative procedures, the global MAP solution, and posterior sampling remains unclear. It is not established whether these methods converge to the global MAP or if they are consistent with DPS. Furthermore, the specific nature of the optimization steps in these methods—whether they solve a global or local objective—requires clarification.

Methodology: Local MAP Sampling (LMAPS)

The authors introduce Local MAP Sampling (LMAPS), a new inference framework that iteratively solves local maximum-a-posteriori subproblems along the diffusion trajectory.

1. Theoretical Framework

LMAPS operates by solving a sequence of local MAP problems rather than a single global one. At each diffusion step $t$, given a noisy state $x_t$ and measurement $y$, the method seeks:
$$x^*_0(t, x_t, y) := \arg \max_{x_0} p(x_0 | x_t, y)$$
This is followed by a standard DDIM transition step:
$$x_{t-\Delta t} = g(x^*_0, x_t, \epsilon)$$

Distinctions from Existing Methods:

• vs. Global MAP: Global MAP seeks the mode of the full posterior $p(x_0 | y)$. LMAPS seeks the mode of the conditional posterior $p(x_0 | x_t, y)$ at each step. The authors argue that the sequence of local maximizers does not necessarily converge to the global maximizer of $p(x_0 | y)$.
• vs. DPS: DPS uses the conditional mean $E[x_0 | x_t, y]$ in the update. LMAPS replaces this mean with the conditional mode. These are equivalent only if the conditional posterior is Gaussian; otherwise, LMAPS introduces a mode-seeking bias, avoiding low-density regions between modes.
• vs. DMAP: While DMAP also utilizes a local MAP perspective, it optimizes in the latent diffusion variable space ($x^*_{t-\Delta t} = \arg \max p(x_{t-\Delta t} | x_t, y)$). LMAPS poses the problem directly in the clean signal space ($x^*_0$), providing a complementary view that explicitly connects to the measurement likelihood $p(y|x_0)$.

2. Practical Algorithm and Approximations

To make the local MAP problem tractable for general non-linear inverse problems, the authors propose specific approximations and reformulations:

• Covariance Approximation: The authors approximate the conditional covariance $\Sigma_{0|t}$ of the prior $p(x_0 | x_t)$ using an isotropic form motivated by a Gaussian prior assumption. As the noise level $\sigma_t \to 0$, the posterior covariance becomes asymptotically isotropic. They approximate $\Sigma_{0|t} \approx \frac{1}{SNR} I$, where $SNR = \alpha_t^2 / \sigma_t^2$.

• Objective Reformulation: The local MAP objective is formulated as minimizing a weighted sum of the prior distance and data fidelity:
  $$x^*_0 = \arg \min_{x_0} \left\{ \frac{SNR}{k} \|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2} \|y - H(x_0)\|^2 \right\}$$
  To improve numerical stability and interpretability, this is reformulated into a scale-invariant convex combination:
  $$x^*_0 = \arg \min_{x_0} \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 = \sigma_t^2 / (\sigma_t^2 + k_1^2)$. This formulation allows for automatic annealing: as $t$ decreases ($\sigma_t \to 0$), the objective shifts from being measurement-driven to prior-driven.

• Linear vs. Non-Linear: For linear inverse problems, the local MAP solution can be derived in closed form, recovering Tweedie Moment Projected Diffusion (TMPD) as a special case. For non-linear problems, the solution is obtained via gradient descent.

Key Contributions

1. Theoretical Unification: The paper formulates LMAPS to unify optimization-based inverse problem methods and Tweedie Moment Projected Diffusion (TMPD) under a single probabilistic framework. It clarifies that optimization-based methods inherently solve local MAP problems rather than global MAP or full posterior sampling.
2. Methodological Advances: The authors introduce a Gaussian-motivated covariance approximation and a reformulated objective function that enhances numerical stability and interpretability through automatic annealing.
3. Empirical Validation: The framework is validated across 10 image restoration tasks (including linear, non-linear, and non-differentiable operators) and 3 scientific inverse problems.

Experimental Results

The authors evaluated LMAPS on FFHQ and ImageNet datasets for image restoration and on benchmarks for scientific inverse problems (Linear Inverse Scattering, Compressed Sensing MRI, and Black Hole Imaging).

• Image Restoration: LMAPS achieved state-of-the-art performance in 43 out of 60 cases across PSNR, SSIM, and LPIPS metrics on FFHQ and ImageNet. It consistently outperformed the closely related DMAP baseline, showing improvements of over 2 dB PSNR in tasks like motion deblurring, JPEG restoration, and quantization.
• Scientific Inverse Problems: LMAPS consistently attained the highest PSNR across all scientific tasks. Notably, it achieved gains of greater than 1.5 dB PSNR on three Linear Inverse Scattering (LIS) tasks compared to previous methods.
• Efficiency: LMAPS demonstrated superior efficiency compared to DAPS, achieving better results with fewer computational resources (e.g., fewer optimization steps per diffusion step).

Significance and Claims

The paper claims that LMAPS provides a principled probabilistic interpretation for optimization-based methods, bridging the gap between posterior sampling (DPS) and point estimation (MAP). By explicitly solving local MAP subproblems, the method offers a unified view that explains the behavior of existing optimization-based solvers.

The authors emphasize that while LMAPS does not sample from the full posterior distribution (and thus does not capture multi-modality in the same way as DPS), it is specifically designed for scenarios where accurate point estimation is the primary goal. The framework offers a robust alternative that balances reconstruction fidelity with the structural priors of diffusion models, achieving state-of-the-art results in both standard image restoration and challenging scientific reconstruction tasks. The work suggests that future research could further explore the utility of global MAP estimation with diffusion priors, which remains an open challenge.