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

TL;DR

This paper advocates viewing the reverse process in DMs as a function and proposes a novel plug-in method for solving IPs using pretrained DMs, dubbed DMPlug, which addresses the issues of manifold feasibility and measurement feasibility in a principled manner and shows great potential for being robust to unknown types and levels of noise.

Abstract

Pretrained diffusion models (DMs) have recently been popularly used in solving inverse problems (IPs). 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. Moreover, their capabilities to deal with noisy IPs with unknown types and levels of measurement noise are unknown. In this paper, we advocate viewing the reverse process in DMs as a function and propose a novel plug-in method for solving IPs using pretrained DMs, dubbed DMPlug. 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. Through extensive experiments across various IP tasks, including two linear and three nonlinear IPs, we demonstrate that DMPlug consistently outperforms state-of-the-art methods, often by large margins especially for nonlinear IPs. The code is available at https://github.com/sun-umn/DMPlug.

Figures (19)

• Figure 1: Visualization of sample results from our DMPlug method (Ours) and main competing methods (DPSchung_diffusion_2023 and Resamplesong_solving_2023 for super-resolution, inpainting, and nonlinear deblurring; BlindDPSchung_parallel_2022 and Stripformertsai_stripformer_2022 for blind image deblurring (BID) and BID with turbulence) on IPs we focus on in this paper. All measurements contain Gaussian noise with $\sigma = 0.01$.
• Figure 2: Evolution of the data-fitting loss $\| \boldsymbol y - \mathcal{A}\pqty{\boldsymbol x} \|_2^2$ of our DMPlug method vs. SOTA methods over percentage progress, for noiseless nonlinear deblurring on the CelebA dataset. Here, the percentage progress is calculated with respect to the total number of iterations taken by each method. The shadow regions indicate the ranges of the loss over $50$ instances.
• Figure 3: Interleaving methods (left) vs. our DMPlug method (right) for solving IPs using pretrained DMs. While interleaving methods cannot ensure the feasibility of the final estimate for either the object manifold $\mathcal{M}$ or the feasible set $\left\{ \boldsymbol x | \boldsymbol y = \mathcal{A}(\boldsymbol x) \right\}$, our DMPlug method ensures the manifold feasibility while promoting $\boldsymbol y \approx \mathcal{A}(\boldsymbol x)$ through global optimization.
• Figure 4: Template for interleaving methods
• Figure 5: PSNR (dB) vs. per-iteration wall-clock time (s) running on an NVIDIA A100, for various reverse steps in $\mathcal{R}$. Experiments on CelebA for $4 \times$ super-resolution; solver: ADAM; maximum iterations: $6,000$