Local MAP Sampling for Diffusion Models

Diffusion Posterior Sampling (DPS) provides a principled Bayesian approach to inverse problems by sampling from $p(x_0 \mid y)$. While posterior sampling is valuable for capturing uncertainty and multi-modality, many classical and practical inverse problem settings ultimately prioritize accurate point estimation—most notably the MAP estimator, which has long served as a standard reconstruction objective in imaging and scientific applications. 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 and DPS, offering a unified probabilistic interpretation for optimization-based methods. Building on this foundation, we develop practical algorithms with a covariance approximation motivated by Gaussian prior assumption, a reformulated objective for stability and interpretability. Across a broad set of image restoration and scientific tasks, LMAPS achieves state-of-the-art performance.

AI systems called diffusion models can generate highly realistic images by gradually turning random noise into structured content. Beyond creating images from scratch, they are increasingly used to solve "inverse problems" — tasks such as sharpening blurry photos, filling in missing regions of an image, reconstructing MRI scans from limited measurements, or imaging black holes from sparse radio telescope data. We introduce Local MAP Sampling (LMAPS), a method that at every step of the diffusion process searches for the single most likely clean image consistent with both the current noisy state and the observed measurement. This perspective unifies and clarifies several earlier approaches under one principled framework. Across a wide range of image restoration tasks and scientific problems — including MRI reconstruction and black hole imaging — LMAPS produces sharper and more accurate reconstructions than existing methods, often at lower computational cost.