Fast Samplers for Inverse Problems in Iterative Refinement Models

Constructing fast samplers for unconditional diffusion and flow-matching models has received much attention recently; however, existing methods for solving inverse problems, such as super-resolution, inpainting, or deblurring, still require hundreds to thousands of iterative steps to obtain high-quality results. We propose a plug-and-play framework for constructing efficient samplers for inverse problems, requiring only pre-trained diffusion or flow-matching models. We present Conditional Conjugate Integrators, which leverage the specific form of the inverse problem to project the respective conditional diffusion/flow dynamics into a more amenable space for sampling. Our method complements popular posterior approximation methods for solving inverse problems using diffusion/flow models. We evaluate the proposed method's performance on various linear image restoration tasks across multiple datasets, employing diffusion and flow-matching models. Notably, on challenging inverse problems like 4x super-resolution on the ImageNet dataset, our method can generate high-quality samples in as few as 5 conditional sampling steps and outperforms competing baselines requiring 20-1000 steps. Our code will be publicly available at https://github.com/mandt-lab/c-pigdm

Paper

References (63)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer X2Z67/10 · confidence 3/52024-07-05

Summary

This paper introduces a inverse problem approach to reconstruct low-resolution blurred or obfuscated images, from pre-trained generative models. The method leverages conjugate integrators to project the diffusion dynamics in an easier space, solve the inverse problem and then map back.

Strengths

- The integration of Conjugate integrators for inverse problems in the paper is rigorously defined and proved. - The proposed method drastically reduces the number of sampling steps compared to other inverse problem approaches. - The presented approach achieve better reconstructed images compared with state-of-the-art inverse problem methods.

Weaknesses

The paper is extremely dense with equations and mathematical formulations. I would suggest moving sections 2.3 and 3.4, as well as the numerical details of the experiments, to the appendix. This would allow space in the main paper for additional qualitative results and a broader discussion on the intuition behind the conjugate integrators for inverse problems. Including a figure illustrating the method would also be appreciated. While the mathematical formulation is the main essence of the paper, a visual representation or descriptive insight would be appreciated for the understanding.

Questions

Minor details: - Line 608, link the approximation to Eq. (5)

Rating

7

Confidence

3

Soundness

2

Presentation

2

Contribution

3

Limitations

I don't have any concern.

Reviewer PpUR5/10 · confidence 3/52024-07-08

Summary

Inverse problems like super-resolution and deblurring remain computationally inefficient with current diffusion/flow models. The paper introduces a plug-and-play framework with Conditional Conjugate Integrators that use pre-trained iterative refinement models to simplify sampling. The method generates high-quality samples in as few as 5 steps for tasks like 4× super-resolution on ImageNet, outperforming existing methods.

Strengths

1. The paper is well-written, and the literature review is thorough. 2. The method is well-supported by theoretical foundations. 3. The experiments on both diffusion and flow-matching methods across various tasks demonstrate the robustness of the proposed speed-up algorithm.

Weaknesses

1. The reason why the paper is restricted to the $\Pi$GDM paradigm is unclear.

Questions

1. Can the method be applied to nonlinear inverse problems? If not, please clarify the reason. 2. Though the Conditional Conjugate Integrators work well and can speed up the sampling process to 5 steps, a more detailed computational comparison is needed as each iteration becomes more complicated. Can you provide more detailed computational results? 3. Can you clarify the reason why you restrict yourself to the posterior approximation in $\Pi$GDM? Though the experimental results are promising, a direct comparison with DPS is preferred.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the paper mentioned its broader impact at the end, which looks good to me.

Authorsrebuttal2024-08-13

Author Response

We thank the reviewer for their response. Please find our detailed response below: **Comparison of perceptual vs Recovery metrics:** We agree that highlighting robustness in both perceptual and recovery metrics is important in the context of inverse problems. Following previous works like DPS and $\Pi$-GDM, we include only perceptual metrics in the main text since recovery metrics like PSNR and SSIM usually favor blurrier samples over perceptual quality. So, a tradeoff exists between distortion (a.k.a recovery) and perceptual quality. However, whether perceptual quality/recovery is preferred depends on the application. Therefore, for completeness, we provide a comparison between DPS, $\Pi$GDM, and C-$\Pi$GDM (our method) in terms of PSNR, SSIM, FID, and LPIPS in Tables 1 and 2 on the ImageNet-256 and FFHQ-256 datasets on the 4x super-resolution task. It is worth noting that the PSNR and SSIM scores for all methods correspond with the best FID/LPIPS scores presented in the main text for these methods. Tables 1 and 2 show that our method achieves competitive PSNR and SSIM scores for better perceptual quality than competing methods, even for very small sampling budgets. For instance, on the FFHQ dataset, our method achieves a PSNR of 28.97 compared to 28.49 for DPS while achieving better perceptual sample quality (LPIPS: 0.095 for ours vs 0.107 for DPS) and requiring around 200 times less sampling budget (NFE=5 for our method vs 1000 for DPS). Therefore, we argue that our perceptual quality to recovery trade-off is much better than DPS, given our method is significantly faster than DPS. **Traversing the Recovery vs Perceptual trade-off in C-$\Pi$GDM**: In addition to the guidance weight $w$, our method also allows tuning an additional hyperparameter $\lambda$, which controls the dynamics of the projection operator (See Sections 2.3 and 3.2 for more intuition). Therefore, tuning $w$ and $\lambda$ can help traverse the trade-off curve between perceptual quality and distortion for a fixed NFE budget. We illustrate this aspect in Table 3 (fixed $\lambda$ with varying $w$) and Table 4 (fixed $w$ with varying $\lambda$) for the SR(x4) task on the ImageNet-256 dataset using the PSNR, LPIPS, and FID metrics. Therefore, our method offers greater flexibility to tune the sampling process towards either good perceptual quality or good recovery for a given application while maintaining the same number of sampling steps. In contrast, other methods like DPS do not offer such flexibility. Moreover, tuning the guidance weight in DPS is very expensive in the first place due to its high sampling budget requirement (around 1000 NFE). We will extend these comparisons for other tasks and add them in the Appendix section of our revised paper. We would also request the reviewer to reconsider their evaluation. | | PSNR | SSIM | FID | LPIPS | |---------------------|-------|-------|-------|-------| | DPS (NFE=1000) | 23.81 | 0.708 | 38.18 | 0.252 | | $\Pi$GDM (NFE=20) | 21.92 | 0.646 | 37.36 | 0.222 | | C-$\Pi$GDM (NFE=5) | 22.32 | 0.641 | 37.31 | 0.220 | | C-$\Pi$GDM (NFE=10) | 23.00 | 0.651 | 34.22 | 0.206 | | C-$\Pi$GDM (NFE=20) | 23.16 | 0.654 | 34.28 | 0.207 | **Table 1**: Comparison between different methods on ImageNet-256 for the SR(x4) task | | PSNR | SSIM | FID | LPIPS | |---------------------|-------|-------|-------|-------| | DPS (NFE=1000) | 28.49 | 0.834 | 30.86 | 0.107 | | $\Pi$GDM (NFE=20) | 28.26 | 0.818 | 26.17 | 0.087 | | C-$\Pi$GDM (NFE=5) | 28.97 | 0.832 | 32.01 | 0.095 | | C-$\Pi$GDM (NFE=10) | 29.03 | 0.821 | 29.07 | 0.086 | | C-$\Pi$GDM (NFE=20) | 28.79 | 0.809 | 26.37 | 0.083 | **Table 2**: Comparison between different methods on FFHQ-256 for the SR(x4) task. | w | PSNR | LPIPS | FID | |----|-------|-------|-------| | 2 | 22.91 | 0.339 | 48.48 | | 4 | 23.37 | 0.306 | 45.03 | | 6 | 23.49 | 0.274 | 42.68 | | 8 | 23.44 | 0.266 | 40.96 | | 10 | 23.28 | 0.254 | 40.27 | | 12 | 22.89 | 0.246 | 40.13 | | 14 | 22.74 | 0.239 | 40.16 | **Table 3**: Illustration of the impact of $w$ for a fixed $\lambda=0.0$ on the sample recovery (PSNR) vs sample perceptual quality (LPIPS, FID) at NFE=5 for our method. The task is SR(x4) on the ImageNet-256 dataset. | $\lambda$ | PSNR | LPIPS | FID | |-----------|-------|-------|-------| | -1.0 | 20.96 | 0.291 | 42.56 | | -0.8 | 21.33 | 0.265 | 40.97 | | -0.6 | 21.69 | 0.240 | 39.38 | | -0.4 | 22.04 | 0.223 | 37.83 | | -0.2 | 22.32 | 0.220 | 37.31 | | 0.2 | 22.73 | 0.257 | 45.27 | | 0.4 | 22.90 | 0.275 | 48.98 | | 0.6 | 23.03 | 0.283 | 47.47 | | 0.8 | 23.11 | 0.285 | 46.2 | | 1.0 | 23.15 | 0.285 | 46.41 | **Table 4**: Illustration of the impact of $\lambda$ for a fixed $w=15$ on the sample recovery (PSNR) vs sample perceptual quality (LPIPS, FID) at NFE=5 for our method. The task is SR(x4) on the ImageNet-256 dataset.

Reviewer hLds7/10 · confidence 4/52024-07-12

Summary

This work proposes a plug-and-play based method that leverages pretrained diffusion and flow models to solve inverse problems. The proposed method called Conditional Conjugate Integrators adapts previously proposed Conjugate Integrators framework for fast sampling of diffusion and flow models to solve linear inverse problems. The key idea is to project the diffusion (flow) dynamics into another latent space where sampling can be more efficient. Upon completion, the dynamics is projected back to the original pixel space. The paper provides the mathematical forms of the projection operator for conditional diffusion (flow) dynamics, then adapts it to get the projection operator for linear inverse problems. The paper also provides tractable forms of the projection matrix and its inverse. This derivation can be seen as a more general form of previously proposed method $\Pi$GDM ($Pi$GFM). The method can be extended to nonlinear and noisy settings. The paper provides promising results on datasets like LSUN Bedrooms, AFHQ Cats, FFHQ, ImageNet etc. on tasks like inpainting, superresolution, gaussian deblurring etc.

Strengths

1. Methodology: The proposed method seems efficient — it can solve linear inverse problems in 5 steps. In comparison, previous methods need 20-100 or even 1000 steps. 2. Experiments: The paper includes ablation studies on the choice of hyper-parameters and provides comparisons against former state-of-the-art methods like $\Pi$GDM, DPS, DDRM,$\Pi$GFM etc. The proposed method out-performs previous methods by a significant margin while using 5-10 steps and is on par with or better than previous method for >=20 steps. 3. Writing: Paper is well written. The core ideas and the methodology have been presented well, and derivations are easy to follow.

Weaknesses

The benefits of the proposed method for the settings of noisy linear inverse problem setting as well as non-linear inverse problems remain unclear. The paper does not include any quantitative results for these two problem settings. It only provides some limited qualitative results for super-resolution for the noisy setting with $\sigma_y$=0.05 and compression inverse problem and JPEG restoration problem for non-linear settings.

Questions

Questions: 1. Why does Eq. 117 hold true for VPSDE? I’m missing the simplification step here. 2. Why does Table 1 not include results of inpainting for diffusion models on ImageNet? Similarly, Table 4 in the appendix skips results for inpainting on FFHQ. Suggestions: 1. All the tables and figures must explicitly state all the relevant settings of the inverse problems. It is not immediately apparent that some of these quantitative results are only for noiseless linear inverse problems. 2. Consider including an algorithmic box that summarizes C-$\Pi$GFM and C-$\Pi$GDM. This would provide a concise overview of the method to the readers.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper discusses relevant limitations.

Reviewer hLds2024-08-11

I thank the authors for their response. I would like to retain my score.

Reviewer PpUR2024-08-12

I appreciate the authors' efforts in preparing the rebuttal and acknowledge the theoretical contributions of the paper, which have been well-received by the other reviewers. However, I still have some concerns. The paper does not report PSNR and SSIM results, making it difficult to fully assess the quality of the proposed methods. Given that the paper utilizes the $\Pi$GDM paradigm, it is crucial to provide these metrics for a comprehensive evaluation. DPS, for example, demonstrates robustness in both perceptual and recovery metrics across various tasks, highlighting the importance of such evaluations. Considering above, I keep the original rating.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC