Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors

Thank you for reviewing our paper. Below we provide point-by-point responses to your comments.

Weakness 1.1

We provide the confidence intervals of all the results in Tables 1 and 2 in the attached PDF. Based on these results, we claim that our method achieves comparable or superior outcomes on most tasks, particularly for complex nonlinear inverse problems such as Fourier phase retrieval. Compared to the baseline methods, our standard deviation is slightly lower, indicating comparable and slightly better robustness of our approach.

To better show the significance of the improvement of our method over baselines, we further conducted a sample-wise PSNR comparison between PnP-DM (EDM) and the most competitive baseline DPnP. The table below shows the sample-wise improvement rate PnP-DM over DPnP and $p$-value of the one-sided t-test for each inverse problem. Note that our method improves on all samples in the test set and the $p$-values indicate high statistical significance.

Response Table 2: Sample-wise comparison between our method and DPnP [2].

Weakness 1.2

As mentioned in Page 7 line 244-246, we used the same model checkpoints for all diffusion model-based methods, which includes DDRM, DPS, PnP-SGS, DPnP, and our PnP-DM. This suggests that these sampling methods should share the same target posterior distribution. For our method, we fine-tuned the parameters in the annealing schedule $\rho_0$, $\alpha$, and \rho_\min (notations from Appendix C.3) by performing a grid search over 20 FFHQ images outside of the test set we used for final reporting. Please also see the response to Question 2 by Reviewer u5tS.

Weakness 1.3

Please refer to the response to all reviewers above for a comparison on computational efficiency. To clarify, as we mentioned in Appendix C.2 "Pseudocode" paragraph, the total number of time steps for the entire diffusion process is 100, but we only start running from some intermediate noise level. Therefore, the number of NFEs for each iteration is less than 100. In fact, according to the table, the average number of NFEs is only around 30 per iteration given the current annealing schedule.

Weakness 2.1

Although we agree that there is a gap between the stationary process and the true posterior, we would like to argue that the convergence guarantees with small $\rho$ are not weaker and that the $\frac{1}{\rho}$ factor is in fact expected. For SGS-base methods, $\rho$ can be interpreted as the "step size." The step size has a similar presence in the bound of prior non-asymptotic analysis on the convergence of Langevin-based MCMC algorithm [8], [9]. For these algorithms, it is common to have an $O(1/T)$ convergence rate where $T$ is the total diffusion time and scales with the step size, so there will be a $\frac{1}{\text{step size}}$ factor on the right hand side. This also justifies the annealing strategy, which starts with large $\rho$ (hence large step size) at the beginning when the distribution is far away from the target and gradually decreases.

Weakness 2.2

Thank you for your suggestion. In Figure 2 of the attached PDF file, we show some visual examples of intermediate $\boldsymbol{x}$ and $\boldsymbol{z}$ iterates (left) and convergence plots of PSNR, SSIM, and LPIPS for $\boldsymbol{x}$ iterates (right) on the super-resolution problem. As $\rho_k$ decreases, the $\boldsymbol{x}$ iterate becomes closer to the ground truth and the $\boldsymbol{z}$ iterate gets less noisy. Both the visual quality and metric curves stabilize after the minimum coupling strength $\rho_{\min}$ is achieved. Despite being run for 100 iterations in total, our method generates good images in around 40 iterations, which is around 30 seconds and 1600 NFEs. We can include these results in our supplemental material of the revised manuscript.

Question 1

We will change the writing on this part in the final version of the paper.

Question 2

Thank you for bringing this paper to our attention. We notice that one reviewer's comment of this work also touches on the naming of the method. According to the authors' response, this happens to be a clash of terminology. However, we will include this paper in the related work. As for the "PnP-DM" name, we understand that "PnP" is so commonly used that it may cause some ambiguity. We used it mainly to highlight its similarity to deterministic PnP methods that incorporate an image prior via denoising.

Thank you for your evaluation of our work. Below we provide point-by-point responses to your comments.

Computational cost: I believe that the proposed method is much slower than existing alternatives. This is not that much of an issue for me (but of course this depends on the reader and practitioner). In my opinion the compute time is not properly addressed; although the authors discuss it in the bottom of page 21, I would like to see a proper runtime comparison with existing methods. I would like to emphasize that this is only a minor weakness.

We provide a comparison on computational efficiency for all methods in Response Table 1 above. Our method PnP-DM achieves superior performance while remaining comparable in runtime with DPS, with a runtime no more than 1.5 times that of DPS. In fact, for the linear inverse problem experiments, PnP-DM was faster than DPS because we were able to take multiple samples along one Markov chain generated by PnP-DM but DPS requires running multiple diffusion processes.

Originality: I believe that the proposed method is related to SDEdit [7]; it can be thought of as a variant of SDEdit. Indeed, to see why this is the case, when the present algorithm is applied to inpainting, the first iteration, the likelihood step, will result in a noisy image with the observation already present In it. This image is then denoised using the backward process during the prior step. This is in fact very similar to sdedit which starts with some initial image, say the pseudo-inverse of the observation, noises it to some given noise level and denoises it using the backward Diffusion process. See Algorithm 3, [7]. I believe that the main difference with the algorithm presented in your paper is the use of the decreasing noise schedule, which makes sense theoretically. For a fair treatment I suggest you mention this in your paper.

Thank you for bringing this work to our attention. This work indeed shares some similarity to our own in that it employs a process of noising followed by denoising. Nevertheless, our method is more theoretically grounded and meant for rigorously estimating the posterior of an inverse problem, whereas SDEdit is more empircally driven and designed for solving image synthesis/editing problems. We will cite and discuss [7] in the related work section in the final version.

Thank you for your feedback on our work. Below we provide point-by-point responses to your comments.

While the method demonstrates superior performance, the computational cost and efficiency are not thoroughly discussed. A comparison of computational resources required compared to other methods would be beneficial. Can the authors compare the computational resources and time required for PnP-DM versus existing DM-based methods?

We provide a comparison on computational efficiency for all methods in Response Table 1 above. Our method PnP-DM achieves superior performance while remaining comparable in runtime with DPS, with a runtime no more than 1.5 times that of DPS. In fact, for the linear inverse problem experiments, PnP-DM was faster than DPS because we were able to take multiple samples along one Markov chain generated by PnP-DM but DPS requires running multiple diffusion processes.

The impact of various parameters, such as the annealing schedule for the coupling parameter ρ, on the method's performance is not deeply investigated. A sensitivity analysis could provide insights into the method's robustness. How sensitive is the performance of PnP-DM to the choice of the annealing schedule for the coupling parameter ρ? Is there an optimal range or strategy for selecting these parameters?

Following your suggestion, we have included a sensitivity analysis on the annealing schedule. In Figure 1 of the attached PDF file, we show the PSNR curves of different exponential decay rates $\alpha$ (left) and minimum coupling levels $\rho_{\min}$ (rate) for one linear (super-resolution) and one nonlinear (coded diffraction patterns) problem.

We have the following conclusions based on the results. First, different decay rates lead to different rates of convergence, which corroborates with our theoretical insights that $\rho$ plays the same role as the step size. The final level of PSNR is not sensitive to different decay rates, as all curves converge to the same level. Second, as $\rho_{\min}$ decreases, the final PSNR becomes higher as the stationary distribution converges to the true target posterior.

Our strategy for choosing the annealing parameters is as follows. The starting coupling strength $\rho_0$ should be large to overcome the ill-posedness of the problem (usually around 5 to 10 is sufficient). The minimum coupling strength $\rho_{\min}$ should be small to ensure that the stationary distribution is close to the target posterior; empirically around 0.1 to 0.3 works the best. The range of decay rate $\alpha$ is generally flexible; a value around 0.9 usually leads to good results.

Thank you for your efforts on reviewing our paper. Below we provide point-by-point responses to your comments:

Weakness 1

We believe that our work has several significant contributions over the existing works [1] and [2]. Moreover, we highlight that according to NeurIPS 2024 policy, our work should be considered as a concurrent work with [2] and "not be rejected on the basis of the comparison to contemporaneous work" as [2] was first published on arXiv on March 25th, 2024, which is within two months of the submission deadline of NeurIPS 2024.

• Compared to [1], the significance of our work lies in three aspects:

  • Nonlinear inverse problems are beyond the scope of [1], while we have investigated three different nonlinear inverse problems.
  • We have adopted a more rigorous formulation than [1]. As also pointed out by [2], the denoiser design based on diffusion models in [1] is heuristic.
  • Unlike PnP-SGS [1], we consider an annealing strategy, which is important for ill-posed inverse problems.

• Compared to [2], our work have the following three contributions:

  • Our work proposes a more general formulation based on EDM. To the best of our knowledge, using EDM as a rigorous Bayesian denoiser for solving inverse problems has not been explored in prior literature.
  • Our proposed method has a significant improvement in performance over DPnP [2]. Response Table 2 in the response to reviewer 43sJ below further shows the statistical significance of the improvement. One reason for the performance gain is that we inherited the optimized design choices from EDM [3], which provides better image quality with less diffusion steps. Another reason is that the flexibility of the framework provides a larger design space with more flexible parameter choices.

Weakness 2

We provide a comparison on computational efficiency for all methods in Response Table 1 above. Our method PnP-DM achieves superior performance while remaining comparable in runtime with DPS, with a runtime no more than 1.5 times that of DPS. In fact, for the linear inverse problem experiments, PnP-DM was faster than DPS because we were able to take multiple samples along one Markov chain generated by PnP-DM but DPS requires running multiple diffusion processes.

Weakness 3

• Fixed $\rho$. Our current analysis considers a fixed $\rho$ across different iterations, but we find that a variable schedule for $\rho$ is more practical. Theoretically, our analysis could be extended to the setting with a variable $\rho$ by viewing the iterations with larger $\rho$ values as warm-up iterations that produce favorable initial conditions for the algorithm with smaller $\rho$ values.

• Approximate score bound. We can extend our convergence bound to the case where the learned score has error, which we can include in our revised manuscript. More precisely, assume we implement the EDM reverse diffusion step (as given by equation (10) in the manuscript) as: $\mathrm{d} \boldsymbol{x}\_t=\left[u(t) \boldsymbol{x}\_t- v(t)^2 s\_t\left(\boldsymbol{x}\_t\right)\right] \mathrm{d} t + v(t)\mathrm{d}\bar{\boldsymbol{w}}\_t,$

  where $s\_t$ is the approximate score. Then, by differentiating the KL divergence along the two dynamics or using the Girsanov theorem (in the spirit of the proof in [4]), we can prove the following: for $\tau \in [k(1+t^\ast)+1, (k+1)(1+t^\ast)]$, which corresponds to the prior step in the Split Gibbs Sampler, we have: $\frac{1}{4}\int\_{k(1+t^\ast)+1}^{(k+1)(1+t^\ast)} \lambda(\tau)\text{FI}(\pi\_\tau\Vert \nu\_\tau)\mathrm{d}\tau \leq \epsilon\_{\text{score}}(k) + \text{KL}(\pi\_{(k+1)(1+t^\ast)}\Vert \nu\_{(k+1)(1+t^\ast)}) - \text{KL}(\pi\_{k(1+t^\ast)+1}\Vert \nu\_{k(1+t^\ast)+1})$

  where $\epsilon\_{\text{score}}(k) := \int\_{0}^{t^\ast} v(t)^2\mathbb{E}||s\_t(\boldsymbol{x}\_t^{(k)}) - \nabla \log p\_t(\boldsymbol{x}\_t^{(k)})||^2 \mathrm{d} t$ and $\boldsymbol{x}\_t^{(k)}, 0\leq t \leq t^\ast$ is the exact EDM reverse process, which satisfies equation (10) in the manuscript, starting at $\boldsymbol{x}\_0^{(k)} \sim \nu\_{k}^Z$ and ending up at $\boldsymbol{x}\_{t^\ast}^{(k)} \sim \nu\_{k+1}^X$. Therefore, up to score approximation errors, a convergence bound analogous to Theorem 3.1 will hold.

• Convergence to the true posterior. Indeed, the asymptotic convergence to the true posterior as $\rho\to 0$ is difficult to obtain (see Section 5 of [5]). We agree that this is out of the scope of the paper, so we leave it to future work.

Question 1

We trained a denoiser for the Gaussian image prior. More details are presented in Appendix C.2 "Model checkpoint" paragraph.

Question 2

Please refer to Appendix D for how we chose the hyperparameters for all the methods. For our method, we fine-tuned the parameters in the annealing schedule $\rho_0$, $\alpha$, and \rho_\min (notations from Appendix C.3) by performing a grid search over 20 FFHQ images outside of the test set we used for final reporting (referred to as the validation set hereafter). For DDRM and DPS on linear inverse problems, we used the default parameters in their official repositories. For all other cases, we performed a grid search of the main parameter(s) of each method on the validation set. Specifically, for DPIR, we used the same annealing function as that in the official repository of [6] but fine-tuned the starting/ending noise level and number of iterations.

Question 3

We did not employ any explicit regularization parameters. Within the Bayesian framework, the data-fidelity (likelihood) and regularization (prior) should be automatically balanced based on the noise distribution.

Question 4

The discretization time steps are the same for all values $\rho_k$. The number of iterations that are actually run depends on the specific value of $\rho_k$. See Appendix C.2 “Pseudocode” paragraph and Algorithm 3 for more details.

Thank you for your response and comments. Below we answer them point by point.

W1.1

Indeed, we have tried to apply PnP-SGS [1] to nonlinear inverse problems by combining our likelihood step with its prior step design, which was how we obtained the PnP-SGS results for Table 2. While it could still handle the coded diffraction pattern problem, it failed on the more challenging Fourier phase retrieval problem. We also tried to include the annealing strategy for PnP-SGS but found that the method diverged with large $\rho$, probably due to its heuristic design of the prior step, and thus did not benefit from annealing. Overall, our method significantly outperforms PnP-SGS by at least 1dB in PSNR for all linear problems and coded diffraction patterns, and 15dB for Fourier phase retrieval. Our experimental results indicate that PnP-SGS struggles with challenging nonlinear inverse problems, such as Fourier phase retrieval.

W1.2

There are mainly two aspects. Here we explain using the notations in the original DDPM paper. First, to rigorously implement the prior step of SGS, one need to assume the input as an observation of the unscaled image $\boldsymbol{x}_0$ with additive white Gaussian noise. However, in DDPM, the mean of state $\boldsymbol{x}\_t$ is not $\boldsymbol{x}\_0$, but $\sqrt{\bar{\alpha}\_t}\boldsymbol{x}\_0$, which is a down-scaled version of $\boldsymbol{x}\_0$. So, it is inaccurate to directly use DDPM as a denoiser. This mismatch is particularly significant when starting from a large $t^\ast$ as $\sqrt{\bar{\alpha}\_{t^\ast}}$ is close to 0. Second, according to the SGS formulation, the noise estimation module is unnecessary, and one should always denoise $\boldsymbol{z}^{(k)}$ (the $\boldsymbol{z}$ iterate of SGS at iteration $k$) assuming a noise level $\rho$. However, PnP-SGS does not take into account the hyperparameter $\rho$ in the denoising problem. It is unclear how the generated $\boldsymbol{x}^{(k+1)}$ relates to the target conditional distribution $\pi^{X|Z=\boldsymbol{z}^{(k)}}(\boldsymbol{x})$ for each prior step of SGS. Therefore, the posterior distribution sampled from the denoiser in [1] is not the desired posterior distribution even with a perfect score function.

W1.3

We never claimed that the idea of annealing is our contribution. Nevertheless, we believe that it is our contribution to propose a general framework that accommodates annealing in an easy yet rigorous way for SGS and leads to significant performance improvement. Our focus is more on the algorithm design and experimental validation. To the best of our knowledge, no existing theory on SGS provides a non-asymptotic convergence bound with annealing, so this is a a challenging open question. Although our current theory assumes fixed $\rho$, it still provides some theoretical insights on the algorithm behavior, such as the interpretation of $\rho$ as step size, and potentially opens up new theoretical directions for SGS based on the Fisher information. We hope to extend the analysis to the annealing case in the future.

W2

We respectfully disagree with your last statement. DPS requires backpropagating through the entire denoiser network for each diffusion step after the forward pass of the denoiser, introducing a significant computational overhead. On the other hand, our method does not require doing so. This is the main reason why the clock time does not scale linearly with NFE for these two methods. Moreover, unlike DPS that applies a likelihood update for every function evaluation. Our method does so only every multiple function evaluations.

W3.1

As we showed in Appendix C.3 "Annealing schedule for $\rho$" paragraph, our annealing schedule decreases to a fixed minimum level at $\rho_{\min}$, so we indeed fixed $\rho$ at small values after certain numbers of iterations. Furthermore, we did run our algorithms for the synthetic prior experiments with a fixed schedule of $\rho$ throughout the process, as shown in Table 4. Experimental results show that our method can accurately sample the posterior, which corroborates our theory. For experiments on linear inverse problems and coded diffraction patterns with FFHQ images, the results with fixed $\rho$ are on par with those with annealing, probably because the problems are not highly non-convex. For more challenging nonlinear problems, such as the Fourier phase retrieval, we empirically find that an annealing schedule is essential to overcome the non-convexity of the problem and provide accurate reconstruction.

W3.2

By combining the above argument with those in the manuscript, we will get the following bound with score function error: $\frac{1}{T} \int_0^{T} \mathsf{FI}\left(\pi_\tau || \nu_\tau\right) \mathrm{d} \tau \leq \frac{4\mathsf{KL}(\pi^X||\nu_0^X)}{K(1+t^\ast) \min(\rho, \delta)^2} + \frac{1}{K(1+t^\ast)\delta^2}\sum_{k=1}^K\epsilon_{\text{score}}(k).$

Thank you for your evaluation of our paper. We respond to your comment below.

There are no apparent weaknesses. However, I am curious about the reconstruction speed comparison (seconds/image) between DPS and the proposed method. It would be practically very attractive for the community to use this algorithm if a rigorous comparison of reconstruction speed is made.

We provide a comparison on computational efficiency for all methods in Response Table 1 above. Our method PnP-DM achieves superior performance while remaining comparable in runtime with DPS, with a runtime no more than 1.5 times that of DPS. In fact, for the linear inverse problem experiments, PnP-DM was faster than DPS because we were able to take multiple samples along one Markov chain generated by PnP-DM but DPS requires running multiple diffusion processes.

Thank you for sharing the sampling speed results. The shared results make sense to me. I strongly feel that this paper provides novel enough contributions to be considered for acceptance. I have raised my confidence score to 4. Hope this helps AC decide about the paper.