Taming Score-Based Denoisers in ADMM: A Convergent Plug-and-Play Framework

This paper introduces ADMM-PnP with a novel AC-DC denoiser that resolves manifold mismatch and establishes convergence guarantees for score-based generative models within the ADMM framework, thereby improving solution quality across various inverse problems.

The Gist

Imagine you are trying to solve a giant, complex jigsaw puzzle, but someone has taken the picture, shredded it, mixed the pieces with sand, and then handed you a blurry, incomplete version of the original image. Your goal is to reconstruct the perfect picture from this mess. This is what scientists call an "inverse problem."

In the world of AI, we have a very smart helper: a Score-Based Model. Think of this model as a master art restorer who has seen millions of perfect paintings. It knows exactly what a "real" face, a "real" landscape, or a "real" car looks like. If you show it a blurry, noisy smudge, it can guess, "Ah, that smudge is probably part of a nose," and fill in the details.

However, there's a catch. The restorer was trained on specific types of "messy" pictures (like photos with random static). But when we try to use this restorer inside a mathematical algorithm called ADMM (a step-by-step process to solve the puzzle), the "mess" we give it looks different. It's like handing the restorer a painting covered in mud, when they only know how to clean off water stains. If you just ask them to clean it, they get confused, and the result looks weird or broken.

This paper introduces a new method called AC-DC Denoising to fix this mismatch and make the whole process work smoothly. Here is how it works, using simple analogies:

1. The Problem: The "Wrong Kind of Mess"

The algorithm (ADMM) is like a strict project manager. It takes a guess at the solution, checks the math, and then asks the AI restorer to "clean up" the guess.

• The Issue: The AI restorer expects a specific kind of noise (like static on an old TV). But the project manager's guess has a weird, complex kind of noise because of the math steps involved.
• The Result: If you ask the restorer to clean this "wrong" noise directly, they might hallucinate (imagine things that aren't there) or fail to fix the image.

2. The Solution: The AC-DC Denoiser

The authors created a three-step "pre-cleaning" routine before the AI restorer even sees the image. Think of it as a three-stage car wash before the car gets waxed.

• Stage 1: Auto-Correction (AC) - "The Rain Shower"

  • What it does: The algorithm intentionally adds a little bit of standard noise (like rain) to the image.
  • The Analogy: Imagine the project manager's guess is a muddy car. The restorer only knows how to wash cars that are covered in rain. So, we spray the muddy car with a hose to turn the mud into a "rainy mess." Now, the car looks like the kind of mess the restorer is trained to handle.
  • Why it helps: It tricks the restorer into thinking, "Oh, this is a standard rainy mess! I know how to fix this!"

• Stage 2: Directional Correction (DC) - "The GPS Guide"

  • What it does: The image is now "rainy," but it's still a bit far from where the restorer wants it to be. The algorithm uses a technique called "Langevin dynamics" (a fancy way of saying "guided walking") to nudge the image closer to the perfect "rainy" zone.
  • The Analogy: Imagine you are in a foggy forest (the rainy mess) and you need to find a specific campfire (the perfect data). You have a GPS that says, "Walk 10 steps North, then 5 steps East." The DC step is like following those GPS directions to get you exactly to the spot where the restorer can work their magic. It ensures you don't wander off into the woods.

• Stage 3: Score-Based Denoising - "The Master Restorer"

  • What it does: Now that the image has been "rained on" (AC) and "guided" to the right spot (DC), the AI restorer finally gets to work.
  • The Analogy: The restorer sees a car covered in the exact kind of rain they are experts at cleaning. They quickly and accurately wipe it down, revealing the beautiful car underneath. Because the car was prepped correctly, the result is perfect.

3. Why This Matters: The "Proof"

The authors didn't just build a cool tool; they proved mathematically that it works.

• The Guarantee: They showed that if you follow these steps, the algorithm will never go crazy or get stuck in a loop. It will eventually settle on a solution that is very close to the best possible answer.
• The Flexibility: This method works even when the math gets really hard (like trying to reconstruct an image from just the shadows it casts, known as "Phase Retrieval").

The Big Picture

Before this paper, trying to use these powerful AI restorers inside complex math algorithms was like trying to fit a square peg into a round hole. It was possible, but the results were often shaky or blurry.

This paper built a adapter (the AC-DC Denoiser) that makes the square peg fit perfectly into the round hole.

• Without it: The AI gets confused, and the image looks like a bad Photoshop job.
• With it: The AI works perfectly, producing crystal-clear images from blurry, broken data.

In short: The authors figured out how to "translate" the messy output of a math algorithm into a language that the AI restorer understands, ensuring that the final result is not just a guess, but a high-quality, mathematically guaranteed solution.

Technical Summary

Here is a detailed technical summary of the paper "Taming Score-Based Denoisers in ADMM: A Convergent Plug-and-Play Framework".

1. Problem Statement

The paper addresses the challenge of integrating score-based generative models (diffusion models) into Alternating Direction Method of Multipliers (ADMM) for solving inverse problems (e.g., image inpainting, deblurring, super-resolution).

While score-based models have shown great success as priors in inverse problems, their direct integration into ADMM faces two critical hurdles:

1. Manifold Mismatch: Score functions are trained on noisy data manifolds created by Gaussian perturbations ($x_t = x_0 + \sigma(t)\epsilon$). However, ADMM iterates (specifically the input to the denoiser, $\tilde{z}^{(k)} = x^{(k+1)} + u^{(k)}$) are influenced by dual variables ($u^{(k)}$) and do not naturally lie on these training manifolds. Directly applying a pre-trained score function to these iterates leads to geometry mismatches, resulting in artifacts and degraded performance.
2. Lack of Convergence Theory: Existing Plug-and-Play (PnP) methods using score-based denoisers often lack rigorous convergence guarantees, especially within primal-dual frameworks like ADMM where dual variables complicate the geometry of the iterates.

2. Methodology: The AC-DC Denoiser

The authors propose a novel AC-DC Denoiser embedded within the ADMM framework to bridge the manifold gap and ensure convergence. The denoiser operates in three stages during each ADMM iteration $k$:

1. Auto-Correction (AC):
   • Mechanism: Additive Gaussian noise is injected into the ADMM iterate $\tilde{z}^{(k)}$.
   • Purpose: This step "gaussianizes" the iterate, pulling it closer to the noisy data manifolds $\mathcal{M}_{\sigma(t)}$ on which the score function was trained. It addresses the distributional shift caused by the dual variable.
2. Directional Correction (DC):
   • Mechanism: A few steps of Conditional Langevin Dynamics are performed starting from the AC output. The dynamics target the conditional distribution $p(z_{\sigma(k)} | \tilde{z}^{(k)}_{ac})$.
   • Purpose: This step refines the alignment with the specific noise level manifold without losing the signal information contained in the original iterate. It effectively corrects the residual error introduced by the AC step and the dual variable.
3. Score-Based Denoising:
   • Mechanism: The output of the DC step is passed through the pre-trained score function.
   • Implementation: The authors offer two variants:
     • Ours-tweedie: Uses Tweedie's lemma to compute the denoised estimate directly: $z = z_{dc} + \sigma^2 s_\theta(z_{dc}, \sigma)$.
     • Ours-ode: Solves an Ordinary Differential Equation (ODE) backward from the DC output to recover the clean signal.

3. Key Contributions

A. Theoretical Convergence Analysis

The paper establishes rigorous convergence guarantees for ADMM-PnP using the AC-DC denoiser under two regimes:

• Weakly Non-Expansive Residuals (Fixed Step Size):
  • The authors prove that under proper parameter scheduling, the AC-DC denoiser acts as a weakly non-expansive operator (specifically, a $\delta$-weakly $\theta$-averaged operator).
  • This ensures that the ADMM iterates converge to a fixed-point ball (a neighborhood of the solution) with high probability, even when the loss function is strongly convex.
• Bounded Denoiser (Adaptive Step Size):
  • For non-convex loss functions (common in many inverse problems), the authors relax the strong convexity assumption.
  • They prove that the AC-DC denoiser is bounded with high probability.
  • By employing an adaptive step-size schedule (increasing the penalty parameter $\rho$), they guarantee convergence to a fixed point with high probability, extending prior PnP convergence theory to score-based settings.

B. Empirical Performance

The method was validated on a wide range of inverse problems including:

• Super-resolution (4x)
• Gaussian and Motion Deblurring
• Inpainting (Random and Box masks)
• Phase Retrieval
• High Dynamic Range (HDR) reconstruction
• Nonlinear Deblurring

Results:

• Quantitative: The proposed method (both Tweedie and ODE variants) consistently achieved State-of-the-Art (SOTA) performance in PSNR, SSIM, and LPIPS metrics compared to baselines like DPS, DDRM, DiffPIR, and RED-diff across FFHQ and ImageNet datasets.
• Qualitative: Visual results showed significantly fewer artifacts and better preservation of fine details compared to baselines, which often suffered from noise or measurement inconsistency.
• Ablation: Removing the DC step resulted in severe artifacts, confirming the necessity of the directional correction for manifold alignment.

4. Significance and Impact

• Bridging Theory and Practice: This work provides one of the first rigorous convergence analyses for score-based PnP methods within the ADMM framework, a setting previously considered difficult due to dual variable interference.
• Solving the Manifold Mismatch: The AC-DC framework offers a generic, principled solution to the "manifold mismatch" problem, which has plagued previous attempts to combine diffusion models with optimization algorithms. It moves beyond simple noise injection by actively correcting the iterate's position relative to the score manifold.
• Flexibility: The method demonstrates that ADMM can effectively handle complex, data-driven priors (diffusion models) alongside multiple regularizers and constraints, expanding the applicability of PnP methods to more complex inverse problems.
• Robustness: The theoretical guarantees hold even for non-convex objectives when using adaptive step sizes, making the framework robust for real-world applications where convexity assumptions often fail.

In summary, the paper successfully "tames" score-based denoisers for ADMM by introducing a three-stage correction mechanism that aligns optimization iterates with training manifolds, backed by strong theoretical convergence proofs and superior empirical results.