PnP-CM: Consistency Models as Plug-and-Play Priors for Inverse Problems

This paper introduces PnP-CM, a plug-and-play framework that reinterprets consistency models as proximal operators to solve diverse linear and nonlinear inverse problems with high-quality reconstructions in as few as two to four neural function evaluations, eliminating the need for task-specific training while outperforming existing methods.

The Gist

Imagine you are trying to solve a giant, complex jigsaw puzzle, but someone has thrown away most of the pieces, smeared the remaining ones with mud, and then crumpled the whole box. This is what scientists call an inverse problem: trying to figure out what the original, perfect picture looked like based on a broken, noisy, and incomplete version.

For a long time, computers have used "Diffusion Models" to solve these puzzles. Think of a Diffusion Model as a very talented, but incredibly slow, artist. This artist knows how to paint a perfect picture from scratch. To fix your broken puzzle, the artist starts with a blank canvas (pure noise) and slowly, step-by-step, peels away layers of "mud" to reveal the image underneath.

The Problem: This artist is too slow. To get a high-quality result, they might need to take 1,000 tiny steps. If you need to fix a medical scan (like an MRI) in a hospital, waiting for 1,000 steps is like waiting for a snail to cross a highway. It's too slow for real life.

The New Idea: The "Consistency" Shortcut
Recently, a new type of AI called a Consistency Model (CM) was invented. Imagine this is a different artist who has memorized the entire painting process. Instead of taking 1,000 steps from start to finish, this artist can look at a muddy piece of the puzzle and instantly say, "Ah, if I clean this specific spot, the final picture will look like this." They can jump straight to the answer in just 2 or 4 steps.

The Catch:
While this "Consistency Artist" is fast, they are a bit stubborn. They are great at painting from scratch, but they aren't very good at following strict instructions like, "Make sure this part of the picture matches the blurry photo we have." If you just let them paint, they might create a beautiful picture that doesn't actually match your broken puzzle.

The Solution: PnP-CM (The Plug-and-Play Team)
This paper introduces PnP-CM, which is like a brilliant project manager who puts the "Consistency Artist" and a "Strict Inspector" in the same room to work together.

Here is how they work, using a simple analogy:

1. The Inspector (The Data Fidelity): This person holds the broken, muddy photo. Their only job is to say, "No, that part doesn't match the photo. Fix it." They ensure the final result actually looks like the input data.
2. The Artist (The Consistency Model): This person is the fast genius who knows how to make things look realistic and sharp.
3. The Process (Plug-and-Play):
   • The Inspector looks at the current guess and says, "This part is too blurry compared to the photo."
   • The Artist takes that suggestion, uses their super-speed to instantly "clean up" the image to make it look realistic, and hands it back.
   • They repeat this back-and-forth conversation very few times (only 2 to 4 times!).

The Secret Sauce: Momentum and Noise
The authors added two special tricks to make this team work even better when they only have time for a few steps:

• Momentum: Imagine you are pushing a heavy shopping cart. If you just push and stop, it stops immediately. But if you give it a little "push" based on where it was going before (momentum), it keeps moving smoothly. The paper uses this to help the computer "remember" its previous steps, so it doesn't waste time correcting small mistakes over and over.
• Controlled Noise: Sometimes, when you are stuck in a maze, you need to shake the table a little to see if there's a hidden path. The team adds a tiny, controlled amount of "shake" (noise) to the image. This helps the Artist escape bad guesses and find the best solution faster.

Why This Matters
The paper tested this method on everything from fixing blurry photos and removing JPEG artifacts to, most importantly, reconstructing MRI scans.

• Speed: Old methods took hundreds of steps. This new method does it in 4 steps (and sometimes even 2!).
• Quality: The results are sharper and more accurate than previous fast methods.
• Real World: They successfully trained this system on real MRI data (knee scans) for the first time. This means doctors could potentially get clear, high-quality MRI images much faster, helping patients get diagnosed sooner.

In Summary
PnP-CM is like hiring a speed-reading genius (the Consistency Model) and pairing them with a strict editor (the data constraints). By adding a little bit of "momentum" and "shaking," they can solve complex image puzzles in the blink of an eye, producing results that are both fast and incredibly accurate. It turns a slow, tedious process into a quick, reliable tool for real-world problems.

Technical Summary

1. Problem Statement

Inverse problems in imaging (e.g., deblurring, super-resolution, inpainting, MRI reconstruction) aim to recover a clean signal $x$ from degraded measurements $y$, modeled as $y = A(x) + n$. These problems are often ill-posed, requiring regularization via a prior distribution $p(x)$.

While Diffusion Models (DMs) have become state-of-the-art priors for solving these problems, they suffer from high computational costs, typically requiring hundreds of Neural Function Evaluations (NFEs) to converge. Consistency Models (CMs) offer a solution by learning to map any point on the diffusion trajectory directly to the clean origin, enabling high-quality sampling in just 1–4 NFEs.

However, existing CM-based solvers for inverse problems face significant limitations:

• Lack of Generality: Methods like CoSIGN require task-specific training (e.g., training a ControlNet for every specific degradation).
• Instability in Ill-Conditioned Systems: Methods like CM4IR rely on pseudo-inverse guidance, which becomes unstable for highly ill-conditioned forward operators (common in MRI) or nonlinear problems.
• Convergence Issues: Existing approaches often lack theoretical convergence guarantees when integrating CMs into optimization frameworks.

2. Methodology: PnP-CM

The authors propose PnP-CM, a framework that reinterprets Consistency Models as proximal operators within a Plug-and-Play (PnP) optimization scheme, specifically using the Alternating Direction Method of Multipliers (ADMM).

Core Framework

The inverse problem is formulated as a Maximum A Posteriori (MAP) estimation:
$$\hat{x} = \arg\min_x f(x) + \lambda g(x)$$
Where $f(x)$ is the data fidelity term and $g(x)$ is the regularizer (prior). PnP-CM solves this via ADMM:

1. Data Fidelity Update ($z$-step): Enforces consistency with measurements $y$. For linear problems, this admits a closed-form solution; for nonlinear or large-scale problems (like MRI), it uses iterative solvers (e.g., Conjugate Gradient).
2. Prior Update ($x$-step): Instead of a traditional proximal operator, the algorithm uses the Consistency Model as a denoiser:
   $$x^{(k+1)} = f_\theta(\nu, t)$$
   Here, the CM maps a noisy input directly to a clean estimate.

Key Innovations for Low-NFE Regime

To ensure high performance with very few iterations (2–4 NFEs), the authors introduce two critical enhancements:

1. Controlled Noise Injection:

   • Standard ADMM assumes a fixed denoiser. PnP-CM injects controlled noise $\eta_k$ into the input of the CM: $x^{(k+1)} = D_\sigma(z + u + \eta_k)$.
   • Theoretical Guarantee: The authors prove (Theorem 1) that if the noise sequence is diminishing and satisfies an energy bound ($\sum \|\eta_k\|^2 < \infty$), the algorithm retains the convergence properties of standard PnP-ADMM.
   • Purpose: This allows the CM to operate at higher noise levels during early iterations, improving its ability to explore the solution space and escape local minima, which is crucial for low-iteration regimes.

2. Momentum-Based Updates:

   • The algorithm incorporates momentum terms for both the primal ($x$) and dual ($u$) variables.
   • Purpose: This accelerates convergence, significantly reducing the number of outer iterations required to reach high-quality reconstructions.

Algorithm Flow

The algorithm iterates in reverse order (from $N$ down to 0) to align with the noise scheduling of Consistency Models. In each step, it performs a data fidelity update, injects noise, applies the CM as a proximal operator, and updates dual variables with momentum.

3. Key Contributions

• Unified Framework: PnP-CM is the first framework to successfully integrate Consistency Models as proximal operators within a PnP-ADMM setting, providing a unified solver for both linear and nonlinear inverse problems without task-specific retraining.
• Theoretical Guarantees: The paper provides a convergence proof for PnP-ADMM with noise injection, ensuring that the acceleration techniques do not compromise stability.
• Efficiency: The method achieves state-of-the-art results in as few as 4 NFEs and produces meaningful outputs in 2 steps, drastically outperforming diffusion-based methods that require hundreds of steps.
• Medical Imaging Application: The authors trained and applied CMs to MRI data for the first time (using the fastMRI dataset), demonstrating the method's viability for large-scale, real-world medical reconstruction.

4. Experimental Results

The authors evaluated PnP-CM on natural images (LSUN Bedroom, CelebA-HQ) and medical images (fastMRI).

• Natural Images:

  • Tasks: Gaussian deblurring, Super-resolution (4x), Inpainting (70% mask), JPEG artifact removal, Nonlinear deblurring, Phase retrieval.
  • Performance: PnP-CM achieved State-of-the-Art (SOTA) PSNR and LPIPS scores with only 4 NFEs.
  • Comparison: It outperformed diffusion-based methods (DPS, DiffPIR, DPnP) which required 100–2000 NFEs, and existing CM-based methods (CoSIGN, CM4IR) which struggled with nonlinear tasks or required retraining.
  • Qualitative: PnP-CM produced sharper details and avoided the over-smoothing seen in DPS and the grid-like artifacts of other methods.

• MRI Reconstruction:

  • Tasks: Multi-coil MRI reconstruction with acceleration factors $R=4$ and $R=8$.
  • Performance: With only 4 NFEs, PnP-CM significantly outperformed DPS (1000 NFEs), DDS (100 NFEs), and CM4IR (4 NFEs) in both PSNR and SSIM.
  • Artifact Reduction: Unlike CM4IR, which suffered from residual aliasing due to unstable pseudo-inverses in multi-coil settings, PnP-CM effectively reduced blurring and structured artifacts.

5. Significance

• Bridging the Efficiency Gap: PnP-CM demonstrates that Consistency Models can be effectively utilized as powerful priors in optimization frameworks, solving the "speed vs. quality" trade-off that has limited the practical adoption of diffusion models in time-sensitive applications.
• Generalizability: By avoiding task-specific training (unlike CoSIGN) and unstable pseudo-inverses (unlike CM4IR), PnP-CM offers a robust, general-purpose solver for a wide range of linear and nonlinear inverse problems.
• Clinical Impact: The successful application to MRI reconstruction suggests that PnP-CM can be deployed in clinical settings where rapid, high-fidelity image reconstruction is critical, potentially reducing scan times or improving diagnostic quality.

In conclusion, PnP-CM represents a significant advancement in inverse problem solving by combining the sampling efficiency of Consistency Models with the robust convergence guarantees of Plug-and-Play ADMM, setting a new benchmark for speed and quality in image reconstruction.