Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Rajesh Shrestha, Xiao Fu

TL;DR

This work proposes ADMM plug-and-play (ADMM-PnP) with the AC-DC denoiser, a new framework that embeds a three-stage denoiser into ADMM: (1) auto-correction via additive Gaussian noise, (2) directional correction using conditional Langevin dynamics, and (3) score-based denoising.

Abstract

While score-based generative models have emerged as powerful priors for solving inverse problems, directly integrating them into optimization algorithms such as ADMM remains nontrivial. Two central challenges arise: i) the mismatch between the noisy data manifolds used to train the score functions and the geometry of ADMM iterates, especially due to the influence of dual variables, and ii) the lack of convergence understanding when ADMM is equipped with score-based denoisers. To address the manifold mismatch issue, we propose ADMM plug-and-play (ADMM-PnP) with the AC-DC denoiser, a new framework that embeds a three-stage denoiser into ADMM: (1) auto-correction (AC) via additive Gaussian noise, (2) directional correction (DC) using conditional Langevin dynamics, and (3) score-based denoising. In terms of convergence, we establish two results: first, under proper denoiser parameters, each ADMM iteration is a weakly nonexpansive operator, ensuring high-probability fixed-point $\textit{ball convergence}$ using a constant step size; second, under more relaxed conditions, the AC-DC denoiser is a bounded denoiser, which leads to convergence under an adaptive step size schedule. Experiments on a range of inverse problems demonstrate that our method consistently improves solution quality over a variety of baselines.

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

TL;DR

This work proposes ADMM plug-and-play (ADMM-PnP) with the AC-DC denoiser, a new framework that embeds a three-stage denoiser into ADMM: (1) auto-correction via additive Gaussian noise, (2) directional correction using conditional Langevin dynamics, and (3) score-based denoising.

Abstract

While score-based generative models have emerged as powerful priors for solving inverse problems, directly integrating them into optimization algorithms such as ADMM remains nontrivial. Two central challenges arise: i) the mismatch between the noisy data manifolds used to train the score functions and the geometry of ADMM iterates, especially due to the influence of dual variables, and ii) the lack of convergence understanding when ADMM is equipped with score-based denoisers. To address the manifold mismatch issue, we propose ADMM plug-and-play (ADMM-PnP) with the AC-DC denoiser, a new framework that embeds a three-stage denoiser into ADMM: (1) auto-correction (AC) via additive Gaussian noise, (2) directional correction (DC) using conditional Langevin dynamics, and (3) score-based denoising. In terms of convergence, we establish two results: first, under proper denoiser parameters, each ADMM iteration is a weakly nonexpansive operator, ensuring high-probability fixed-point using a constant step size; second, under more relaxed conditions, the AC-DC denoiser is a bounded denoiser, which leads to convergence under an adaptive step size schedule. Experiments on a range of inverse problems demonstrate that our method consistently improves solution quality over a variety of baselines.
Paper Structure (49 sections, 27 theorems, 142 equations, 18 figures, 5 tables, 1 algorithm)

This paper contains 49 sections, 27 theorems, 142 equations, 18 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Proposed Approach
  4. Convergence Analysis
  5. Convergence of Under Weakly Non-Expansive Residuals
  6. Convergence Under Weakly Non-Expansiveness with AC-DC
  7. Convergence without Convexity of $\ell$
  8. Related Works
  9. Experiments
  10. Conclusion
  11. Preliminaries
  12. Notation
  13. Definitions
  14. $\mu$-strongly convex function boyd2004convex.
  15. Nonexpansive function bowder1965nonexpansive.
  16. ...and 34 more sections

Key Result

Theorem 1

Under Assumption assumpt: delta weakly residual, assume that $\ell$ is $\mu$-strongly convex. Then, there exists ${\bm x}^*$, ${\bm u}^*$ and $K>0$ such that the sequences $\{ {\bm x}^{(k)} \}_{k\in \mathbb{N}^+}$ and $\{ {\bm u}^{(k)} \}_{k \in \mathbb{N}^+}$ generated by ADMM-PnP using a fixed ste

Figures (18)

  • Figure 1: Left: direct denoising of ${\widetilde{\bm z}}^{(k)}$ using score functions could lead to unnatural recovered signals with artifacts. Right: AC-DC denoising brings ${\widetilde{\bm z}}^{(k)}$ closer to $\mathcal{M}_{{\sigma}^{(k)}}$, and then uses the score function to bring ${\widetilde{\bm z}}^{(k)}$ to the data manifold $\mathcal{M}_{\rm data}$.
  • Figure 2: Inpainting under random missings.
  • Figure 3: Inpainting under box missing.
  • Figure 4: Recovery under motion blurring.
  • Figure 5: Influence of DC steps in the denoiser.
  • ...and 13 more figures

Theorems & Definitions (53)

  • Definition 1: Fixed point convergence
  • Definition 2: Sequence convergence to a $\delta$-ball
  • Definition 3: ADMM convergence to a $\delta$-ball
  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Lemma 1: Tweedie's lemma robbins1992anempirical
  • Lemma 2
  • proof
  • Lemma 3: Concatenation of $\delta$-weakly $\theta$-averaged functions
  • ...and 43 more