🤖 machine learning

Inverse problems with diffusion models: MAP estimation via mode-seeking loss

This paper introduces Variational Mode-Seeking Loss (VML), a theoretically grounded and analytically derivable objective for minimizing KL divergence that enables efficient, high-performance Maximum A Posteriori (MAP) estimation for solving arbitrary inverse problems using pre-trained unconditional diffusion models without task-specific training.

Original authors: Sai Bharath Chandra Gutha, Ricardo Vinuesa, Hossein Azizpour

Published 2026-02-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sai Bharath Chandra Gutha, Ricardo Vinuesa, Hossein Azizpour

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a beautiful, high-resolution photograph, but someone has accidentally spilled coffee on it, torn out a chunk, or blurred it so badly you can barely recognize the scene. In the world of computer science, this is called an inverse problem: you have the damaged result (the coffee-stained photo), and you need to figure out what the original, pristine image looked like.

For a long time, solving this required training a specific AI for every single type of damage. If you wanted to fix a blur, you trained one AI; if you wanted to fill in a missing piece, you trained another.

This paper introduces a new way to use a pre-trained diffusion model (a type of AI that learned how to generate images from scratch) to fix any kind of damage without needing to be retrained. Think of this pre-trained AI as a master painter who has seen millions of photos and knows exactly what a "real" face or landscape should look like.

The Problem with Current Methods

The authors explain that while we can use this master painter to fix damaged photos, the current methods are like trying to navigate a dark room with a flashlight that only shows you a blurry outline.

Posterior Sampling: Some methods try to generate many possible versions of the original photo to cover all bases. It's like asking the painter to draw 100 different versions of the missing piece. While this is thorough, it's computationally heavy and often produces results that are "average" rather than sharp and realistic.
MAP Estimation: Other methods try to find the single most likely original image (the "Maximum A Posteriori" or MAP estimate). This is like asking the painter, "What is the one most probable thing that was here?" This usually gives a sharper, more realistic result, but existing ways to do this often rely on rough guesses (approximations) that can lead to errors or take a very long time to compute.

The New Solution: "Mode-Seeking"

The authors propose a new strategy called VML-MAP (Variational Mode-Seeking Loss).

Here is the core idea using an analogy:
Imagine the "original image" exists in a vast, mountainous landscape. The "peaks" of the mountains represent the most likely, realistic images (the modes). The "valleys" represent impossible or weird images.

The damaged photo gives you a starting point in this landscape, but you don't know exactly which peak you are aiming for.
The current methods often wander around the landscape, sometimes getting stuck in small hills or taking a very long, winding path to find the highest peak.
VML-MAP introduces a new "compass" (the Variational Mode-Seeking Loss). This compass doesn't just tell you which way is "up"; it specifically points toward the highest, most prominent peaks (the modes) of the landscape.

How It Works (The "Mode-Seeking Loss")

The paper introduces a mathematical formula called VML.

The Goal: The AI starts with a noisy, blurry version of the image and gradually cleans it up step-by-step (this is the "reverse diffusion" process).
The Compass: At every single step of this cleaning process, the VML formula calculates a "loss" (a score of how wrong the current guess is).
The Magic: The authors prove that if you minimize this specific score at every step, you are mathematically guaranteed to steer the image toward the most probable, sharpest version of the original.
- For simple, linear problems (like blurring or standard resizing), they found they can write this compass down as a perfect, exact formula. No guessing required!
- They call this "Mode-Seeking" because it actively hunts for the "modes" (the peaks) of the probability landscape, ensuring the final image is the most plausible one.

Why It's Better

The authors tested this on various tasks: filling in missing parts of faces (inpainting), making small images huge (super-resolution), and un-blurring photos.

Speed: Their method is faster. It finds the best answer with fewer steps than previous methods.
Quality: The images look more realistic. Because the method focuses on the "most probable" peak rather than an average of many possibilities, the details are sharper and less "mushy."
No Approximations: For many common tasks, their math is exact. They don't need to make shortcuts that other methods rely on, which reduces errors.

A Special Trick for Tough Problems

Sometimes, the "landscape" is tricky (mathematically called "ill-conditioned"), meaning the path to the peak is steep and confusing. The authors also created a preconditioner (a special tool in their algorithm) that acts like a pair of high-tech hiking boots. These boots help the AI climb steep, slippery slopes much faster and more stably, ensuring it doesn't get stuck or take a wrong turn.

Summary

In short, this paper gives us a new, highly efficient "GPS" for AI image restoration. Instead of wandering around or guessing, this new method (VML-MAP) uses a precise mathematical compass to guide the AI directly to the most realistic, sharp, and probable version of a damaged image, doing so faster and more accurately than previous techniques.

Technical Summary: Inverse Problems with Diffusion Models via Mode-Seeking Loss

Problem Definition

The paper addresses the challenge of solving inverse problems using pre-trained unconditional diffusion models. An inverse problem involves estimating an original data sample $x_0$ given a partially degraded measurement $y$ , governed by the equation $y = A(x_0) + \eta$ , where $A$ is a known degradation operator and $\eta$ is Gaussian noise.

While pre-trained diffusion models can solve these problems in a zero-shot (plug-and-play) manner, existing approaches face significant limitations:

Posterior Sampling Methods: These aim to draw samples proportional to the posterior probability $p(x_0|y)$ . However, they often rely on modeling approximations (e.g., Gaussian approximations of the conditional score) or sequential Monte-Carlo methods, which can be computationally demanding and may not strictly adhere to the true posterior without ground-truth validation.
MAP Estimation Methods: These aim to find the Maximum A Posteriori (MAP) estimate (the mode of the posterior). Existing solvers often use heuristic approximations, rely on variable splitting (e.g., HQS, ADMM), or require consistency models. These approaches frequently involve expensive gradient computations and modeling assumptions that may not hold for all linear operators.

The authors argue that many image restoration tasks prioritize generating a plausible image consistent with measurements rather than covering the entire posterior distribution. This aligns more closely with MAP estimation, motivating the need for a more principled, approximation-free MAP strategy.

Methodology: Variational Mode-Seeking Loss (VML)

The core contribution is the Variational Mode-Seeking Loss (VML), a new objective function designed to steer generated samples toward the MAP estimate during the reverse diffusion process.

Theoretical Foundation

The authors derive VML from a variational perspective by minimizing the Kullback-Leibler (KL) divergence between the diffusion posterior $p(x_0|x_t)$ and the measurement posterior $p(x_0|y)$ at each reverse diffusion time step $t$ :
$\text{VML}(x_t, t) = D_{KL}(p(x_0|x_t) \parallel p(x_0|y))$

The intuition is that as $t \to 0$ , the distribution $p(x_0|x_t)$ becomes increasingly peaked around $x_t$ . By finding an optimal $x_t^*$ that minimizes this KL divergence at each step, the process guides the sample toward the modes of the posterior $p(x_0|y)$ .

Closed-Form Derivation for Linear Operators

A key theoretical achievement is the derivation of a closed-form expression for VML in the case of linear inverse problems ($A(x) = Hx$) without any modeling approximations. The full VML expression includes:

A measurement consistency term: $\|y - HD(x_t, t)\|^2 / (2\sigma_y^2)$ .
Prior terms: Involving the denoiser $D(x_t, t)$ , the noise schedule $\sigma_t$ , and the covariance of the posterior $Cov[x_0|x_t]$ .

Simplified and Reweighted Objectives

The full VML contains higher-order terms involving the covariance matrix $Cov[x_0|x_t]$ , which are computationally expensive to compute and differentiate. The authors prove that these higher-order terms converge uniformly to zero (or a constant) as $t \to 0$ . Consequently, they propose a Simplified-VML ( $VML_S$ ) that drops these terms, retaining only the measurement consistency and prior terms.

To further improve optimization stability, they introduce a reweighted simplified-VML ( $VML_S^\rho$ ):
$VML_S^\rho(x_t, t) = \frac{\|y - HD(x_t, t)\|^2}{2\sigma_y^2} - \rho_t \log p(x_t) - \rho_t \frac{\|D(x_t, t) - x_t\|^2}{2\sigma_t^2} + C$
where $\rho_t$ is a prior weight schedule such that $\rho_t \to 1$ as $t \to 0$ . This formulation allows for efficient gradient computation using standard automatic differentiation.

Algorithms

Based on VML minimization, the authors propose:

VML-MAP: An algorithm using gradient descent to minimize $VML_S^\rho$ at each reverse diffusion step.
VML-MAP $_{pre}$ : A preconditioned variant for ill-conditioned linear operators. It utilizes the Singular Value Decomposition (SVD) of the degradation matrix $H$ to construct a preconditioner $P$ , accelerating convergence for difficult tasks like deblurring.
VML-MAP $^\tau$ : A strategy for noisy image restoration. To avoid sharp artifacts caused by optimization challenges at low noise levels, the algorithm minimizes VML only until a diffusion time threshold $\tau$ , then switches to standard DDIM sampling to generate the final image.
LatentVML-MAP: An extension to Latent Diffusion Models (LDMs), approximating the VML in the latent space.

Key Contributions

Novel Loss Function: Introduction of VML, which theoretically steers samples toward MAP estimates by minimizing the KL divergence between the diffusion and measurement posteriors.
Approximation-Free Derivation: For linear inverse problems, VML is derived in closed form without heuristic modeling approximations, providing a formal theoretical basis for MAP estimation with diffusion models.
Convergence Analysis: Formal proof that the simplified and reweighted VML objectives converge to the negative log-posterior (the MAP objective) as $t \to 0$ under mild assumptions.
Practical Algorithms: Development of VML-MAP and its preconditioned variant, offering an empirically effective solution for diverse image restoration tasks.
Extensive Validation: Demonstration of superior performance over existing baselines (e.g., DDRM, $\Pi$ GDM, MAPGA, DAPS) across multiple datasets (ImageNet, FFHQ, CelebA) and tasks (inpainting, super-resolution, deblurring).

Experimental Results

The authors evaluated their methods on image restoration tasks using pre-trained models on ImageNet256, FFHQ256, and CelebA.

Metrics: The paper emphasizes perceptual metrics (LPIPS and FID) over pixel-wise metrics (PSNR, SSIM), arguing that the latter often fail to capture the perceptual quality of plausible reconstructions in ill-posed inverse problems.
Performance:
- VML-MAP consistently outperformed posterior sampling methods (DDRM, $\Pi$ GDM) and other MAP estimators (MAPGA, DAPS) in terms of LPIPS and FID across inpainting, super-resolution, and deblurring tasks.
- VML-MAP $_{pre}$ showed significant improvements over the non-preconditioned version for ill-conditioned tasks (e.g., deblurring), achieving lower LPIPS and FID scores.
- Computational Efficiency: VML-MAP achieved better perceptual quality with lower runtime (fewer neural function evaluations) compared to methods requiring many reverse diffusion steps.
Noisy Restoration: The VML-MAP $^\tau$ strategy successfully handled noisy measurements, producing high-quality reconstructions where standard VML-MAP struggled with artifacts.
Latent Models: LatentVML-MAP demonstrated effectiveness on CelebA, though the authors note that optimization remains more challenging in the latent space due to decoder non-linearity.

Significance and Claims

The paper claims that VML represents the first approximation-free modeling of MAP estimation for linear inverse problems leveraging only a pre-trained unconditional diffusion model, supported by formal theoretical analysis.

The authors emphasize that their approach differs fundamentally from posterior sampling (which targets distribution coverage) and existing MAP solvers (which often rely on heuristics or consistency models). By minimizing VML, the method effectively targets the "modes" of the posterior, which the authors argue is the desired behavior for generating plausible, high-quality image restorations.

The work is modest regarding its limitations, acknowledging that:

The effectiveness relies on the availability of practically efficient optimizers for the VML objective.
Optimization challenges persist for ill-conditioned or non-linear operators, though preconditioning helps.
Future work is needed to develop higher-order optimizers to further mitigate optimization issues without prohibitive computational costs.

The paper concludes that VML-MAP provides a principled, theoretically grounded, and empirically effective framework for solving inverse problems with diffusion models, bridging the gap between theoretical MAP estimation and practical image restoration.