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

We thank reviewer uJfQ for the detailed review and insightful comments!

RE Weakness 1:

We appreciate this insightful comment! Indeed, the two papers have similar ideas to ours from a high level, and we will definitely add and discuss them in our revision. However, we still want to highlight several crucial differences between their methods and our method.

• [1] computes the gradient using the adjoint method and the gradient checkpointing trick to mitigate the memory issue, but the two techniques will further slow down the optimization process. In contrast, our method, with only a few steps, can achieve SOTA results.
• [1] only considers linear IPs, but our paper contains three nonlinear IPs and achieves much better results than the current SOTA methods (typically 3-6dB in terms of PSNR).
• Our paper observes the early-learning-then-overfitting and spectral bias phenomena when solving IPs. With the two phenomena, our method can achieve robustness to unknown types and levels of noise.
• [2] also uses the gradient checkpointing trick to mitigate the memory issue, slowing down the whole process.
• Although IPs can be considered as conditional generation problems, the condition is much stronger than the classifier/CLIP guidance. Please refer to GR1 for more information about IPs and conditional generation with the classifier/CLIP guidance.

[1] D-Flow: Differentiating through Flows for Controlled Generation

[2] End-to-End Diffusion Latent Optimization Improves Classifier Guidance

[3] Scalable gradients for stochastic differential equations

[4] https://pytorch.org/docs/stable/checkpoint.html

RE Weakness 2:

We mainly follow the experiment setting of the recent work—Resample [5] (ICLR’24 Spotlight), which also uses 100 images for each task and dataset, but we totally understand this concern. Hence, we try to add more experiment results in the rebuttal period as you and other reviewers suggested (please check General Response). Besides the two datasets, we also try a more complex dataset—LSUN-bedroom, following [5]. Also, we add more qualitative results for ImageNet as you suggested in GR5.

[5] Solving inverse problems with latent diffusion models via hard data consistency

RE Question 1:

This is indeed a valid question! Please refer to GR2.

RE Question 2:

We appreciate this insightful suggestion! Please refer to GR1.

RE Question 3:

Sure, please refer to GR5.

RE Question 4:

We appreciate this very interesting comment! Please refer to GR1.

We thank reviewer uJfQ for the detailed review and insightful comments!

During the rebuttal and discussion period, we incorporated six additional strong baseline methods in GR3, GR7, and GR9, as suggested by you and reviewer yL1C. These methods include FPS, DiffPIR, DDNM, PSLD, RED-diff, and ΠGDM. We hope these results can somehow address your concerns on this point.

We thank reviewer uJfQ for the detailed review and insightful comments!

For the box inpainting task, please refer to GR8 for the blurriness analysis, and both qualitative and quantitative comparisons between our method and several strong baseline methods.

We thank reviewer y52g for the detailed review and insightful comments!

RE Weakness 1:

Please refer to GR3.

RE Weakness 2:

We totally understand this concern that using different DMs may lead to some unfairness concerns, but it is not clear which model the unfairness pertains to. Actually, even the Resample paper [1] compares their LDM-based methods with DM-based methods in most of their experiments. We are happy that our ablation study in Table 5 partially mitigated this concern, so we further add more LDM experiments this time. Please refer to GR4 for using our method with LDMs on inpainting and nonlinear deblurring tasks.

[1] Solving inverse problems with latent diffusion models via hard data consistency

RE Question 1:

This is indeed a valid question! Please refer to GR2.

We thank reviewer wUnV for the detailed review and insightful comments!

RE Weakness 1 & Question 1:

Please refer to GR2.

RE Weakness 2 & Question 4:

Please refer to GR4.

RE Weakness 3 & Question 2:

This is a very good suggestion! We use the standard DDIM/DDPM and LDM models. The pre-trained DDIM models of CelebA, FFHQ are from this link; the pre-trained DDIM model of LSUN-bedroom is from this link; the pre-trained LDM model of CelebA is from this link. We recommend the original papers for the training details. After optimizing the initial noise, we apply the same number of NFEs in R(.) as the training stage, i.e., 3 in most of the experiments of this paper. We will add an appendix to include these details in the revision.

RE Weakness 4 & Question 3:

Thanks for this great suggestion!

Please refer to GR1 for our discussion of the seemingly fundamental difference between inverse problems and conditional image generation.

In our paper, we have tried three challenging nonlinear IPs, including nonlinear deblur, BID and BID with turbulence, but we still want to try the tasks you suggested. For CT reconstruction, we follow the settings in [1] since they provide pre-trained DM for CT images. However, we are still waiting for the data access approval from TCIA on this link. In phase retrieval (PR), the goal is to recover a 2d or 3d image from the oversampled diffraction pattern (squared magnitude) of its Fourier transform. Although we can definitely test our method on PR, we have decided not to do it here, as the setting in the DPS paper [2] and follow-up work may be physically wrong, as Fourier PR never involves color images (Fourier involves only 2d or 3d single-channel images) [3,4,5,6]—in fact, their setting leads to a more difficult, yet unrealistic, problem; we leave this study as future work.

[1] Improving diffusion models for inverse problems using manifold constraints

[2] Diffusion posterior sampling for general noisy inverse problems

[3] The numerics of phase retrieval

[4] SiSPRNet: end-to-end learning for single-shot phase retrieval

[5] Three-dimensional imaging of strain in a single ZnO nanorod

[6] What is Wrong with End-to-End Learning for Phase Retrieval?

The deep image prior phenomena appears in the link is also interesting. One suggestion from me would be presenting the one with the best LPIPS score since high PSNR may lead to blurriness. However, I still feel like the 3-step DDIM needs a better explanation or empirical support from more experiments on challenging tasks especially when the box-inpainting showing the reconstruction slightly blurry and different from original measurement. This casts doubt on whether the solution can keep data-fidelity to the measurement. Nevertheless, current results on super-resolution and deblurring are impressive. I believe the contribution is sufficient as an empirical method paper. So I keep my score as 6: weak accept.

We thank reviewer wUnV for the further discussion and insightful comments!

For the box inpainting task, please refer to GR8 for the blurriness analysis, and both qualitative and quantitative comparisons between our method and several strong baseline methods.

We thank reviewer yL1C for the detailed review and insightful comments!

RE Weakness 1:

We agree that allowing posterior sampling for uncertainty quantification and other purposes would be great. This is, unfortunately, not what the MAP framework can offer. We follow most current DM-based methods for inverse problems, which do not provide posterior sampling or uncertainty quantification. We will acknowledge this as a limitation of the current work, and leave it for future research.

RE Weakness 2:

We will definitely revise the manuscript thoroughly after the rebuttal to improve the presentation!

Optimizing in the space of the clean image directly reduces things to the typical MAP formulations without using the pretrained DM priors. Our consideration is not about whether the score functions are estimated accurately or not; ours is more geometric: we view the whole reverse process as a learned function that characterizes the image manifold. Then, the question here is really about how many steps we need in the pretrained DM model, so that approximate DM priors are incorporated. About this choice, we have tried to address it in Sec 3.1 as well as our GR1 above.

RE Weakness 3 & Question 3:

We totally agree with the intuition, and our choice of 3 steps is a hyperparameter that works reasonably well in typical settings, but in no sense optimal.

We conduct the ablation study as you suggested. We test 40 cases from the CelebA dataset for random inpainting, and try different numbers of reverse steps with the increase of the mask ratio. The table below shows that our method with 7 steps performs slightly better than it with 3 and 11 steps, as you expected. We will add this ablation study in our future version and recommend users ablate the number of steps when facing challenging IPs.

RE Weakness 4:

Please refer to GR2.

RE Weakness 5:

Please refer to GR3.

RE Question 1:

This is a very good idea! We noticed a recent paper [1] (after our submission) that successfully combines optimizing the input of diffusion models and consistency models. We will probably also incorporate this in our future version.

[1] Inverse Problems with Diffusion Models: A MAP Estimation Perspective

Question 2:

This is an interesting idea and definitely worth exploring! Using our 3-step method for super-resolution on 100 cases from CelebA, we experiment with different latent spaces to optimize when t=T, t=⅔ T, t=⅓ T, respectively. The results below show that optimizing the latent space when t<T can lead to unsatisfactory performance, which shows that it is essential to optimize the latent space when t=T to fully take advantage of the pre-trained prior information.

We thank reviewer yL1C for the further discussion and insightful comments!

For the box inpainting task, please refer to GR8 for the blurriness analysis, and both qualitative and quantitative comparisons between our method and several strong baseline methods.

We sincerely thank all the dedicated ACs and reviewers for your time and great efforts. We would like to take this opportunity to re-emphasize our contributions.

• (Novelty) Recently, interleaving the diffusion process with projection or gradient updates has become the dominant framework for solving inverse problems using pretrained diffusion models, with hundreds of papers published in the past few years. In this context, we introduce a completely novel approach, DMPlug, which directly integrates pretrained diffusion models into the Maximum A Posteriori (MAP) framework. We appreciate the recognition of our novelty by Reviewer yL1C, Reviewer wUnV, and Reviewer y52g.

• (Significant improvements in experiments) Our method outperforms most SOTA methods across a range of linear and nonlinear inverse problems, including super-resolution, inpainting, nonlinear deblurring, blind image deblurring (BID), and BID with turbulence. The performance improvements are particularly significant for complex nonlinear inverse problems, with gains typically ranging from 3 to 6 dB in PSNR. We appreciate the recognition of the improvements by all four reviewers.

• (Robustness to unknown noise). To the best of our knowledge, this is the first paper that can address unknown measurement noise with pretrained DMs for solving IPs. Most current popular methods assume that measurement noise is a certain type, like Gaussian, and the noise level is known. We appreciate the recognition of this point by Reviewer wUnV, Reviewer y52g, and Reviewer uJfQ.

• (Interesting phenomena). We demonstrate that our method achieves superior performance with just 3 DDIM reverse steps, whereas conditional generation using a classifier fails to produce reasonable objects in only 3 steps. We further compare and discuss this during the rebuttal period and this may suggest significant differences between conditional generation and IPs. We are delighted that most reviewers have expressed interest in this. Also, we observe early-learning-then-overfitting and spectral bias phenomena when using our method for solving IPs.

In general, we believe that this work has the potential to open up several new research directions for our community. Thanks again for all the ACs and reviewers!

GR8: Box inpainting (Continuing GR5)

Following several reviewers’ suggestions made during the rebuttal, we include further results and analysis on the box inpainting task below.

Why the reconstructions in GR5 appear slightly blurry inside the box mask?

This seems to be the result of the interaction between the spectral bias of our learning process and the PSNR metric we use to select the reported result, as explained below.

We provide the intermediate reconstruction of our method at https://anonymous.4open.science/r/2024_NIPS_rebuttal-AAC3/box_process.png. After the PSNR peak, we can see an obvious visual improvement in image details. So, if we visually pick the best result, we can get much sharper ones.

The spectral bias of learning here means our method picks up the low-frequency components first, and then gradually high-frequency ones. Visually, it picks up smooth blobs first, and then gradually more details, as can be seen from https://anonymous.4open.science/r/2024_NIPS_rebuttal-AAC3/box_process.png. However, the PSNR metric that we use to select the reported results is known to be not sensitive to slight smoothing/blurring—which typically causes sharp changes in perceived visual quality. For example, as [3] points out, “...a well-known example is that blurring causes large perceptual but small L2 change… ". So the reported results selected based on the PSNR metric may not be the one with the sharpest detail inside the box mask; in fact, due to the spectral bias, the sharpest one tends to be after the PSNR peak.

Would performing the selection based on other visual metrics work? Inspired by Reviewer wUnV’s suggestion, we experimented with using the Learned Perceptual Image Patch Similarity (LPIPS) metric instead of PSNR. Unfortunately, the reconstruction with the lowest LPIPS value still did not align with the visually optimal result. So this question indeed turns into which visual metric aligns well with human visual perception, which is an active area of research itself.

Why is only this box-mask inpainting task special? What about other tasks? In the manuscript, we have reported results for most other tasks based on the PSNR metrics, and their visual qualities are typically satisfactory. We believe the gap lies in how much the measurement consistency can constrain the problem: for other tasks, if we blur the target image locally anywhere, the measurement consistency will be substantially affected; for box-mask inpainting, we can blur the target image inside the box mask arbitrarily without affecting the measurement consistency. So the strong emphasis of our method on measurement consistency has avoided the trouble discussed above for other tasks we experiment with, but not box-mask inpainting.

[3] The unreasonable effectiveness of deep features as a perceptual metric.

Thorough comparison of our method against several strong baseline methods.

• The quantitative results are included in the table below, complemented by a qualitative comparison in https://anonymous.4open.science/r/2024_NIPS_rebuttal-AAC3/box_compare.png. The quantitative comparison shows that our method outperforms PSLD, ReSample, and DDNM, while being on par with FPS and DPS for this task. Following Reviewer yL1C's suggestion, we have also increased the number of reverse steps for this challenging task, but the results remained comparable with our method using fewer steps.

• The qualitative results at https://anonymous.4open.science/r/2024_NIPS_rebuttal-AAC3/box_compare.png shows that PSLD and ReSample occasionally generate artifacts.

Overall, considering both quantitative and qualitative results, our method matches DPS in performance and surpasses other competing methods. Although our reconstructions can sometimes appear slightly blurry, they maintain the highest pixel-level fidelity with the highest PSNR, and they appear photorealistic without artifacts.

(For box inpainting, we maintain a box size of 128 × 128 pixels, consistent with the DPS and ReSample papers. In alignment with other experiments in our manuscript and the ReSample paper, we also introduce additional Gaussian noise with 𝜎 = 0.01.)

We thank reviewer yL1C for the valuable suggestion to include more baseline methods. In addition to PSLD and GR7, we have tried our best to obtain the results of RED-diff and ΠGDM for super-resolution before the discussion deadline. The results indicate that our method consistently demonstrates comparable or superior performance to the tested competing methods.