Plug-and-Play Diffusion Meets ADMM: Dual-Variable Coupling for Robust Medical Image Reconstruction
| Authors | Chenhe Du et al. |
| Year | 2026 |
| Field | Computer Vision |
| arXiv | 2602.23214 |
| Download | |
| Categories | cs.CV, cs.LG |
Abstract
Plug-and-Play diffusion prior (PnPDP) frameworks have emerged as a powerful paradigm for solving imaging inverse problems by treating pretrained generative models as modular priors. However, we identify a critical flaw in prevailing PnP solvers (e.g., based on HQS or Proximal Gradient): they function as memoryless operators, updating estimates solely based on instantaneous gradients. This lack of historical tracking inevitably leads to non-vanishing steady-state bias, where the reconstruction fails to strictly satisfy physical measurements under heavy corruption. To resolve this, we propose Dual-Coupled PnP Diffusion, which restores the classical dual variable to provide integral feedback, theoretically guaranteeing asymptotic convergence to the exact data manifold. However, this rigorous geometric coupling introduces a secondary challenge: the accumulated dual residuals exhibit spectrally colored, structured artifacts that violate the Additive White Gaussian Noise (AWGN) assumption of diffusion priors, causing severe hallucinations. To bridge this gap, we introduce Spectral Homogenization (SH), a frequency-domain adaptation mechanism that modulates these structured residuals into statistically compliant pseudo-AWGN inputs. This effectively aligns the solver's rigorous optimization trajectory with the denoiser's valid statistical manifold. Extensive experiments on CT and MRI reconstruction demonstrate that our approach resolves the bias-hallucination trade-off, achieving state-of-the-art fidelity with significantly accelerated convergence.
Engineering Breakdown
Plain English
This paper addresses a fundamental problem in medical image reconstruction using diffusion models: existing Plug-and-Play Diffusion (PnP) solvers treat the reconstruction process as a memoryless operation that only looks at current gradients, causing them to fail under heavy image corruption and never fully satisfy physical measurement constraints. The authors propose Dual-Coupled PnP Diffusion, which reintroduces dual variables from classical optimization (ADMM) to create a feedback mechanism with memory across iterations, theoretically guaranteeing convergence to the exact data manifold. This approach combines the representational power of modern diffusion priors with the convergence guarantees of classical optimization theory, enabling more robust and reliable medical image reconstruction even in severely corrupted scenarios.
Core Technical Contribution
The key innovation is replacing memoryless gradient-based updates with a dual-variable coupling scheme inspired by Alternating Direction Method of Multipliers (ADMM). Instead of updating the reconstruction estimate based only on the current iteration's information, the method maintains a dual variable that accumulates historical information across all iterations, providing integral feedback that prevents steady-state bias. This theoretical insight—that classical optimization dual variables address a blind spot in modern deep-learning-based solvers—is novel and represents a principled way to merge pretrained generative models with rigorous optimization guarantees. The authors prove that this dual coupling framework asymptotically converges to solutions that strictly satisfy the physics-based measurement constraints, even when corruption is severe.
How It Works
The system starts with a corrupted medical image and a pretrained diffusion model that acts as an implicit prior on natural images. At each iteration, the Dual-Coupled PnP framework performs three coupled updates: (1) a diffusion sampling step that refines the image estimate using the generative model, (2) a data-fidelity step that pulls the estimate back toward consistency with the physical measurements (e.g., from MRI or CT sensors), and (3) critically, a dual-variable update that accumulates the discrepancy between the image estimate and the measurement-consistent solution. The dual variable acts as an integrator or "memory bank" across iterations—if earlier iterations failed to satisfy the measurement constraint, the dual variable grows and forces subsequent iterations to correct the error. This coupling is formalized through ADMM, where the image estimate and measurement consistency are treated as two coupled variables that must converge to the same value. The output is a medical image that both looks natural (via the diffusion prior) and strictly satisfies the physics-based measurement equations.
Production Impact
For medical imaging teams deploying reconstruction pipelines (MRI, CT, PET), this approach directly improves image quality under extreme noise or sparse measurements, which translates to reduced patient radiation exposure or faster scan times. Engineers integrating this would replace their current PnP solver (typically a proximal gradient or HQS-based loop) with the dual-coupled variant, which requires only a small architectural change—storing and updating one additional dual variable per pixel/voxel. The convergence guarantees mean fewer manual hyperparameter tunings and more predictable behavior in clinical settings, reducing the risk of artifacts that could be misinterpreted as pathology. Compute-wise, the dual-coupled approach adds one extra vector update per iteration but eliminates the need for extensive tuning of step sizes or early-stopping heuristics, potentially saving wall-clock time overall. The main trade-off is code complexity and the need to retrain or validate the diffusion prior on domain-specific medical datasets; generic vision diffusion models may not capture the statistics of medical images, undermining the entire framework.
Limitations and When Not to Use This
The paper assumes a pretrained diffusion model already exists and generalizes well to the specific imaging modality and corruption type in deployment; if the diffusion prior is mismatched (e.g., trained on chest X-rays but applied to brain MRI), the theoretical guarantees break down in practice. The convergence proof likely assumes convex or well-behaved measurement models; real inverse problems (e.g., nonlinear sensor responses, phase retrieval) may not satisfy these assumptions, leaving the practical convergence speed unclear. Computational cost scales with the number of iterations needed to reach convergence—in high-noise regimes this could require 100+ diffusion sampling steps, each itself costly, making real-time clinical deployment challenging. The paper does not address how to choose the coupling strength between the dual variable and the primal updates in non-convex settings, or how to detect convergence without ground truth, both critical for practitioners.
Research Context
This work sits at the intersection of two major research streams: (1) Plug-and-Play priors, which treat any pretrained model as a modular "black box" in optimization loops, and (2) diffusion-based generative modeling as a prior for inverse problems. Prior work (e.g., HQS-based PnP solvers) demonstrated that diffusion priors enable high-fidelity reconstructions but lacked convergence guarantees; this paper closes that gap by borrowing the ADMM dual-variable framework from classical convex optimization. The contribution bridges a 20+ year-old optimization technique with cutting-edge generative models, showing that classical insights remain valuable in the deep-learning era. This opens a research direction of systematically auditing which optimization components are "lost" in modern deep solvers and how to restore them for robustness.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::