Fast samplers for Inverse Problems in Iterative Refinement models

We thank the reviewer for their feedback. We provide a detailed response to several concerns below:

I would suggest moving sections 2.3 and 3.4, as well as the numerical details of the experiments, to the appendix.....

We thank the reviewer for their suggestion. We stress that the final version allows for an extra page, which we will use to make the paper more easily accessible. Re. Section 2.3, we would like to highlight that this section provides an intuitive understanding of our method's tractability and working. Therefore, we think this section serves as an intuitive premise for the observations presented in our experimental section. Re. Section 3.4, we will consider moving this section to the appendix since our presented results for non-linear and noisy inverse problems serve as a proof of concept. Additionally, we plan to add an overview figure in the revised manuscript. We also provide a sample overview figure as Figure 1 in the rebuttal pdf (Also see the second point of our shared response).

Line 608, link the approximation to Eq. (5)

Thanks for pointing it out. We will revise our manuscript accordingly

\newcommand\dag\dagger We thank the reviewer for their feedback. We provide a detailed response to several concerns below:

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

We chose to work with the $\Pi$GDM baseline due to the following reasons:

1. Firstly, $\Pi$GDM is a very strong and widely-used baseline for tackling inverse problems, providing a more expressive posterior approximation $p(x_0|x_t)$ compared to other methods such as DPS or MCG. This makes it an excellent starting point for low NFE budget sampling. As demonstrated in our experiments (See Table 1), DPS typically requires 1000 NFEs to achieve optimal results. In contrast, $\Pi$GDM serves as a competitive baseline and a strong foundation for our method. Additionally, $\Pi$GDM offers greater flexibility since it does not require a differentiable image degradation transform, unlike DPS. Thus, PGDM is a natural choice for us to demonstrate the effectiveness of our method.

2. Secondly, in the context of popular diffusions and flows, using the posterior approximation in $\Pi$GDM results in a closed-form derivation of the conjugate operator and its inverse (See Eqns. 11 and 12 in the main text). Note that this is not a limitation of our method since the projection operator can also be derived for DPS, as highlighted in the following theoretical result (for brevity, we omit the full proof and would include it in our revision).

Proposition. For a time-dependent design matrix $B: [0,1] \rightarrow \mathbb{R}^{d\times d}$ and the posterior approximation proposed in DPS [Chung et al.] with guidance step size $\rho$ such that the conditional score $\nabla_{x_t} \log p(y|x_t) = \rho \frac{\partial \hat{x}_0}{\rho x_t}^\top H^\top (y - H\hat{x}_0)$, introducing the transformation $\hat{x}_t=A_tx_t$, where $A_t = \exp{\Big[\int_0^t B_s - \Big(F_s + \frac{\rho}{2\mu_s^2}G_s G_s^\top H^\top H\Big)ds \Big]}$ induces the following projected diffusion dynamics $d\hat{x}_t = A_t B_t A_t^{-1}\hat{x}_t dt + d \Phi_y y + d \Phi_s \epsilon(x_t, t) + d \Phi_j \Big[\partial\epsilon(x_t, t)H^\top (y - H\hat{x}_0)\Big]$ where $\exp(.)$ denotes the matrix exponential, $\hat{x}_0$ represents the Tweedie's moment estimate and $H$ denotes the degradation operator. The proof roughly follows similar ideas from Appendix A.3 in our paper and we will include a complete proof in our revision.

Despite our method applying to DPS, it is worth noting that in contrast to $\Pi$GDM, the projection operator $A_t$ and its inverse for DPS do not exhibit a closed-form solution and need to be approximated using perturbation analysis (if the step size $\rho$ is small) or computed using standard routines in packages like PyTorch. Therefore, though our method is also applicable to DPS, we choose to stick with the $\Pi$GDM framework in this work.

Can the method be applied to nonlinear inverse problems? If not, please clarify the reason.

While our presentation of the proposed method is primarily in the context of linear-noiseless inverse problems, we also present an extension of our method to noisy-linear and non-linear inverse problems in Section 3.4. Furthermore, in Figure 11 of our paper, we illustrate the applicability of our method to non-linear inverse problems in the context of challenging problems like JPEG and lossy neural compression-based restoration. We would also like to point the reviewer to our common response (Point 1) for a more comprehensive justification of the utility of our method for noisy and non-linear inverse problems.

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?

We provide a computational comparison in Table 1 using NFE as the metric, which is standard in the existing literature on inverse problems like DPS [Chung et al.] and PiGDM [Song et al.]. It is important to note that the compared models utilize the same pre-trained backbone diffusion/flow models, ensuring no differences in model inference speed. The only additional computation incurred involves a simple linear transform $A_t$ and some scalar integral calculations, which are computationally inexpensive. Moreover, these integrals can be precomputed offline (i.e., before sampling starts, as only the noise schedule influences the integral results). Therefore, the computational cost can be amortized between samples, reducing potential overhead. We will make this point more explicit in our revised manuscript.

Can you clarify the reason why you restrict yourself to the posterior approximation in $\Pi$GDM

Please see our response to your first comment.

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.

Table 1: Comparison between different methods on ImageNet-256 for the SR(x4) task

Table 2: Comparison between different methods on FFHQ-256 for the SR(x4) task.

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.

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.

\newcommand\dag\dagger We thank the reviewer for their insightful comments and feedback. We provide a detailed response to several concerns below:

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 =0.05 and compression inverse problem and JPEG restoration problem for non-linear settings.

As illustrated in Figures 7 and 11 in the main text, we would like to highlight that our proposed method can be applied to challenging noisy and non-linear inverse problems in as few as 5-10 sampling steps. Therefore, the proposed method can also be beneficial for these settings. However, we acknowledge that our qualitative results for noisy and nonlinear problems serve as a proof of concept and that further evaluation and refinements are necessary. We intend to pursue more detailed investigations and improvements in future work and will highlight the same in the Conclusion section of a subsequent revision. We would also like to point the reviewer to our common response (Point 1) for a more comprehensive justification of the utility of our method for noisy and non-linear inverse problems.

Why does Eq. 117 hold true for VPSDE? I’m missing the simplification step here.

For a given degradation operator $H$, the core idea in this step is the property of the projection matrix $P=H^\top(HH^\top)^{-1}H$, which shows: $PH^\dag = \big[H^\top(HH^\top)^{-1}H\big]\big[H^\top(HH^\top)^{-1}\big] = H^\top(HH^\top)^{-1} = H^\dag$ Now, Eqn. 116 reads as: $\Phi_y = -\int_0^t \frac{w\beta_t\mu_s}{2}\exp(\kappa_1(s))\Big[H^\dag + (\exp(\kappa_2(s)) - 1)PH^\dag\Big] ds$ Replacing $PH^\dag = H^\dag$ in the above equation, we get,

$

\Phi_y &= -\int_0^t \frac{w\beta_t\mu_s}{2}\exp(\kappa_1(s))\Big[H^\dag + (\exp(\kappa_2(s)) - 1)H^\dag\Big] ds\\ &= -\int_0^t \frac{w\beta_t\mu_s}{2}\exp(\kappa_1(s))\exp(\kappa_2(s))H^\dag\Big] ds \\ &= -\big[\int_0^t \frac{w\beta_t\mu_s}{2}\exp(\kappa_1(s) + \kappa_2(s))ds\big]H^\dag

$

Therefore, we arrive at Eqn. 117. We hope this resolves any confusion. We will also add these simplifying steps in our subsequent revision.

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.

We thank the reviewer for pointing this out. We couldn't include these comparisons due to time constraints but will include them in a subsequent revision.

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.

We thank the reviewer for pointing this out. We are planning to update the paper with more detailed captions and agree that we should mention the settings of the problems more explicitly in the figures and tables

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.

We thank the reviewer for this suggestion. We plan to add a pseudocode/algorithm box summarizing the proposed samplers in the main text. We would also like to point the reviewer to our common response (Point 2) for more details regarding improved illustrations of our method.