Measurement-Consistent Langevin Corrector for Stabilizing Latent Diffusion Inverse Problem Solvers

This paper proposes the Measurement-Consistent Langevin Corrector (MCLC), a theoretically grounded plug-and-play module that stabilizes latent diffusion inverse problem solvers by bridging the gap between solver dynamics and stable reverse diffusion processes through measurement-consistent Langevin updates.

The Gist

Imagine you are trying to solve a puzzle, but you only have a few scattered pieces and a blurry photo of the finished picture. This is what scientists call an Inverse Problem: trying to figure out the original, perfect image (the "signal") based on a damaged or incomplete version (the "measurement").

In recent years, AI models called Latent Diffusion Models (LDMs) have become the best "puzzle solvers" we have. They are like super-smart artists who have seen millions of pictures and know exactly what a face, a landscape, or a car should look like. They use this knowledge to fill in the missing parts of your puzzle.

However, there's a big problem: These AI artists get dizzy and unstable.

The Problem: The "Dizzy Artist"

When these AI models try to fix your blurry photo, they sometimes take a step in the wrong direction. Imagine the artist is trying to draw a perfect circle, but every time they add a new line, they accidentally spin the paper a little bit. By the time they finish, the circle is a messy scribble.

In technical terms, the AI's internal "dynamics" (how it moves from a blurry guess to a clear image) get out of sync with the stable, learned rules it was trained on. This causes artifacts—weird blobs, distorted faces, or strange patterns that don't belong.

Previous attempts to fix this were like trying to force the artist to stay on a straight, flat road (a "linear manifold"). But the AI's world is actually a bumpy, twisting mountain range (a "non-linear latent space"). Forcing it to stay on a flat road just doesn't work; the artist keeps falling off.

The Solution: The "Measurement-Consistent Langevin Corrector" (MCLC)

The authors of this paper came up with a clever new tool called MCLC. Let's break down what it does using a simple analogy.

The Analogy: The GPS and the Compass

Imagine the AI artist is a hiker trying to find a hidden treasure (the perfect image).

1. The Compass (The AI Model): The AI has a compass that points toward "beautiful, realistic images." It knows the general direction.
2. The GPS (The Measurement): You give the hiker a specific clue, like "The treasure is 10 meters north of this tree." This is your measurement.

The Problem:
When the hiker tries to move toward the tree (the measurement), they sometimes get pushed off the trail. They wander into a swamp (instability), and the compass spins wildly, making them draw a mess instead of finding the treasure.

The Old Fix:
Previous methods tried to build a fence around the trail, forcing the hiker to stay on a straight line. But since the trail is actually curvy and mountainous, the fence just traps the hiker or breaks.

The MCLC Fix:
MCLC is like a smart guide who steps in every time the hiker gets pushed off course.

1. Check the Clue: The guide looks at the GPS clue (the measurement). "Okay, we need to stay 10 meters north of the tree."
2. The Correction: The guide says, "You moved too far left! Let's take a small step back, but only in a direction that doesn't mess up your distance from the tree."
3. The Magic Move: The guide uses a special rule (Langevin dynamics) to nudge the hiker back onto the "safe path" of the AI's training, but they do it in a way that never changes the distance to the tree.

Why is this special?

• It respects the rules: It ensures the final image still matches your original blurry photo (Measurement Consistency).
• It respects the artist: It gently guides the AI back to its "comfort zone" where it knows how to draw realistic things, without forcing it into a straight line that doesn't exist.
• It's a "Plug-and-Play" tool: You don't need to rebuild the whole AI. You just attach this "smart guide" module to any existing AI solver, and it instantly becomes more stable and produces cleaner, sharper images.

The Results

When the researchers tested this "smart guide" on various tasks—like removing blur from photos, filling in missing parts of an image (inpainting), or making low-res images high-res—the results were amazing:

• Fewer weird blobs: The "blob artifacts" disappeared.
• Sharper details: Faces looked more natural.
• More reliable: The AI didn't crash or produce nonsense as often.

In a Nutshell

Think of MCLC as a stabilizer for AI artists. When an AI tries to fix a broken image, it often gets confused and starts hallucinating weird details. MCLC acts like a gentle hand on the artist's shoulder, whispering, "You're getting too far from the original photo! Take a small step back, but keep looking at the beautiful style you learned."

This allows the AI to use its incredible creativity to solve difficult puzzles without losing its mind, resulting in clearer, more accurate, and more trustworthy images.

Technical Summary

1. Problem Statement

Latent Diffusion Models (LDMs) have become state-of-the-art priors for solving inverse problems (e.g., image deblurring, super-resolution, inpainting) due to their ability to learn complex data distributions from large-scale datasets. However, existing LDM-based inverse solvers suffer from instability.

• The Symptom: This instability manifests as visual artifacts, degraded reconstruction quality, and "off-manifold" behavior where the sampling trajectory drifts away from the data distribution.
• The Limitation of Prior Work: Previous approaches attempted to fix this by relying on the Manifold Hypothesis, specifically assuming that the latent space behaves as a linear manifold. They proposed manifold-preserving solvers to keep the sampling path on this assumed linear surface.
• The Core Issue: The authors argue that the linear manifold assumption fails in latent space. Because LDMs utilize highly nonlinear decoders (e.g., in Stable Diffusion), the latent space does not satisfy linear manifold properties. Consequently, methods enforcing linear constraints often fail to stabilize the solver effectively.

2. Methodology: Measurement-Consistent Langevin Corrector (MCLC)

The paper proposes a new perspective on instability and a plug-and-play solution called MCLC.

A. New Perspective: Discrepancy in Dynamics

Instead of viewing instability as a geometric deviation from a manifold, the authors characterize it as a discrepancy between the solver-induced dynamics and the stable reverse diffusion dynamics learned by the pre-trained model.

• The Metric: They quantify this instability using the Kullback–Leibler (KL) divergence between the time-evolving distribution of the solver ($q_t^\#$) and the target time-marginal distribution of the pre-trained diffusion model ($p_t$).
• Observation: Naive solvers exhibit a significant gap (high KL divergence) from the target distribution, leading to instability.

B. The Solution: MCLC

MCLC is a theoretically grounded, plug-and-play module inserted after the measurement-consistency step in existing solvers. It consists of two main components:

1. Langevin Correction:

   • Based on the theory that Langevin dynamics driven by the score function ($\nabla \log p_t$) converges to the target distribution $p_t$.
   • The method applies a Langevin update to pull the latent variable back toward the stable reverse diffusion distribution, thereby reducing the KL divergence.

2. Measurement-Consistent Projection:

   • The Challenge: A standard Langevin update might improve distributional stability but could violate measurement consistency (i.e., the reconstructed image might no longer match the observed measurements $y$).
   • The Fix: MCLC projects the Langevin update onto the orthogonal complement of the measurement gradient ($\nabla_z r(z)$).
   • Mathematical Formulation:
     $$z_t^c \leftarrow z_t^\# + \eta_t \cdot P_{\perp g} \nabla \log p_t(z_t^\#) + \sqrt{2\eta_t} \cdot P_{\perp g}(\epsilon)$$
     Where $P_{\perp g} = (I - gg^T)$ is the projection matrix onto the subspace orthogonal to the normalized measurement gradient $g$.
   • Result: This ensures that the correction step reduces the KL divergence (stabilizing the dynamics) while preserving measurement consistency up to a controlled first-order bound.

3. Key Contributions

1. Redefining Instability: The paper shifts the paradigm from geometric manifold assumptions to a dynamical discrepancy perspective, identifying instability as a deviation from the learned time-marginal distributions.
2. Theoretical Justification: They provide a rigorous proof that applying Langevin dynamics with an orthogonal projection reduces the KL divergence while maintaining measurement fidelity.
3. Plug-and-Play Module: MCLC is designed to be integrated into any existing LDM-based inverse solver (e.g., LDPS, PSLD, ReSample) without modifying the core solver architecture.
4. Robustness: Unlike prior methods, MCLC does not rely on the invalid linear manifold assumption, making it effective in the highly nonlinear latent spaces of modern LDMs.

4. Experimental Results

The authors evaluated MCLC across various tasks (Gaussian/Motion Deblur, Super Resolution, Inpainting, HDR) on FFHQ and ImageNet datasets, using different base solvers (LDPS, PSLD, ReSample, LatentDAPS) and priors (Stable Diffusion v1.5, v2.1, Realistic Vision).

• Quantitative Improvements:

  • Stability: MCLC consistently improves FID (Fréchet Inception Distance) and LPIPS (Learned Perceptual Image Patch Similarity), indicating fewer artifacts and higher perceptual quality.
  • Fidelity: It maintains or slightly improves PSNR (Peak Signal-to-Noise Ratio), proving that stability is achieved without sacrificing data consistency.
  • Comparison: MCLC outperforms DiffStateGrad (a state-of-the-art manifold-preserving method) and non-pluggable approaches like MPGD and SILO.
  • Robustness: The method generalizes well across different LDM priors and even extends to Flow-based models (FlowChef).

• Qualitative Improvements:

  • Visual results show a significant reduction in "blob artifacts" and structural distortions common in naive LDM solvers.
  • The PSNR histograms shift toward higher values, indicating a reduction in catastrophic failure cases.

• Efficiency:

  • MCLC introduces minimal computational overhead (approx. 3% runtime increase for most solvers) because it reuses gradients from the measurement-consistency step and avoids expensive backward passes through the decoder.

5. Significance and Impact

• Theoretical Advancement: The paper provides a principled framework for understanding and stabilizing inverse solvers in latent space, moving away from restrictive geometric assumptions toward dynamical consistency.
• Practical Utility: As a plug-and-play module, MCLC allows researchers and practitioners to immediately stabilize existing LDM-based solvers, leading to more reliable zero-shot inverse problem solving.
• Generalizability: The approach is not limited to diffusion models; the authors demonstrate its applicability to flow-based generative models, suggesting a universal mechanism for stabilizing generative inverse solvers.

In summary, MCLC solves the critical instability problem in latent diffusion inverse solvers by mathematically aligning the solver's trajectory with the pre-trained model's stable dynamics while strictly adhering to measurement constraints, offering a superior alternative to previous manifold-based approaches.