DMPlug: A Plug-in Method for Solving Inverse Problems with Diffusion Models

🌈 Abstract

The paper focuses on solving inverse problems (IPs) using pretrained diffusion models (DMs). The authors propose a novel plug-in method called DMPlug that addresses the issues of manifold feasibility and measurement feasibility in a principled manner, and also shows potential for robustness to unknown types and levels of noise. Through extensive experiments, the authors demonstrate that DMPlug consistently outperforms state-of-the-art methods, often by large margins, especially for nonlinear IPs.

🙋 Q&A

[01] Introduction

1. What are the key issues with existing methods for solving IPs using pretrained DMs?

• The existing methods mostly interleave iterative steps in the reverse diffusion process and iterative steps to bring the iterates closer to satisfying the measurement constraint. However, such interleaving methods struggle to produce final results that look like natural objects of interest (i.e., manifold feasibility) and fit the measurement (i.e., measurement feasibility), especially for nonlinear IPs.
• The capabilities of existing methods to deal with noisy IPs with unknown types and levels of measurement noise are unknown.

2. What are the key contributions of the paper?

• The authors propose a novel plug-in method, DMPlug, that views the reverse diffusion process as a function and reparameterizes the object to be recovered as the seed of this function.
• DMPlug addresses the issues of manifold feasibility and measurement feasibility in a principled manner, and also shows great potential for being robust to unknown types and levels of noise.
• The authors demonstrate through extensive experiments that DMPlug consistently outperforms state-of-the-art methods, often by large margins, especially for nonlinear IPs.
• The authors observe an "early-learning-then-overfitting" (ELTO) property in DMPlug, and by integrating an early-stopping method (ES-WMV), they achieve robustness to unknown noise types and levels.

[02] Background and Related Work

1. What are diffusion models (DMs) and how are they used for solving IPs?

• DMs, such as the denoising diffusion probabilistic model (DDPM) and the denoising diffusion implicit model (DDIM), gradually transform an input into total noise and then learn to gradually recover from the noise.
• Ideas for solving IPs with DMs can be classified into two categories: supervised and zero-shot. The zero-shot category makes use of pretrained DMs as data-driven priors.
• Most zero-shot methods follow a common algorithmic template that interleaves iterative reverse diffusion steps and iterative steps to move closer to the feasible set.

2. What are the issues with the prevailing interleaving methods?

• It is unclear whether the interleaving iterative sequence will converge to either the data manifold defined by the pretrained DM (manifold feasibility) or the feasible set (measurement feasibility).
• Most methods assume known noise types (often Gaussian) and known, often very low, noise levels, which is unrealistic in practice.

[03] Method

1. What is the key idea behind DMPlug?

• DMPlug views the whole reverse diffusion process as a function that maps from the seed space to the object space (or the object manifold).
• This allows DMPlug to reparameterize the object of interest as the seed variable and plug this reparameterization into the traditional regularized data-fitting framework.

2. How does DMPlug address the issues of manifold feasibility and measurement feasibility?

• By viewing the reverse diffusion process as a function and optimizing the unified formulation with respect to the seed, DMPlug is expected to produce an object on the object manifold (addressing manifold feasibility) and promote the object to satisfy the measurement constraint (addressing measurement feasibility).

3. How does DMPlug achieve robustness to unknown noise types and levels?

• DMPlug exhibits an "early-learning-then-overfitting" (ELTO) property, where the quality of the estimated object climbs first to a peak and then degrades once the noise is picked up.
• By integrating an early-stopping method (ES-WMV) that stops the estimate sequence near the peak, DMPlug can solve IPs without the exact noise information.

[04] Experiments

1. What are the key findings from the experimental evaluation?

• DMPlug consistently outperforms state-of-the-art methods on both linear and nonlinear IPs, often by large margins, especially for nonlinear IPs.
• For linear IPs, DMPlug can lead the best SOTA methods by about 2 dB in PSNR and 0.02 in SSIM on average.
• For nonlinear IPs, DMPlug can lead the best SOTA methods by about 4 dB in PSNR and 0.05 in SSIM on average.
• DMPlug demonstrates flexibility in employing different priors and optimizers.
• DMPlug is the first to achieve robustness to unknown noise types and levels, leading SOTA methods by about 2 dB and 4 dB in PSNR for linear and nonlinear IPs, respectively.