Tweedie Moment Projected Diffusions for Inverse Problems

Q3. What exactly justifies the increased compute cost of the proposed sampler beyond the theoretical justification? In other words, in what setting would one choose the proposed sampler over others, such as DPS?

As demonstrated in the new PDF, the cost increase is 1.4-1.6x, while this also comes with significant performance increase. We do believe that our method can be further engineered in practical settings to have even higher quality - which means that it could be a practical choice when the sample quality is the ultimate goal.

Another important aspect is that we do not have hyperparameters to search for, therefore, our sampler is much more generic. We also have shown in the synthetic experiments that it produces more diverse samples, which is also important in statistical settings, e.g., the use of our method in scientific applications (as opposed to image generation).

Conclusion

We sincerely thank the reviewer for insightful and thoughtful comments and expressing their concerns. We followed their advice to improve the weak points of our paper and we believe that our updated manuscript has a much better standing as a result. We now think that we have experimental evidence that shows our method has strong qualities (increased performance as well as removal of hyperparameter tuning which is a significant cost) and obtained better theoretical understanding. We kindly request the reviewer to reevaluate the paper in the light of these changes and increase the score if the reviewer agrees that our paper now meets the standards of ICLR. Thank you.

Weakness 2. Qualitative comparison seem to show that reconstruction quality is poor. There are many inconsistencies and artifacts on generated samples both for inpainting and SR (Figures 3 and 4). State-of-the-art diffusion samplers should be able to provide higher quality reconstructions visually, for instance take a look at samples generated by DPS in the original paper. I believe DPS is not set up properly for the experiments as the reconstructions are clearly not consistent with the observations. I recommend tuning the step size for DPS on a small validation set to optimize LPIPS/PSNR.

As we noted in our "general response to reviewers", this caused by an error on our side. We have now comprehensively updated these results, please see our new experimental results.

We have now set up DPS properly and have, for each inverse problem, tuned the step size for DPS on a 128 size validation set to optimize LPIPS, PSNR, MSE and SSIM, please see Fig. 4 for an example.

We illustrate the higher quality reconstructions visually in Figure 3 and Figure 5, and are happy to provide visuals of the full set of samples used to calculate Tables 2, 5 and 6 upon request (as the file size is quite large). In our newly shown experiments, (see Tables 2 and 5), we indeed demonstrate that our method does increase quantitative metrics significantly w.r.t. state-of-the-art.

Weakness 3. Claims such as "having on-par restoration results with PiGDM" and "we found our approach more robust across various problems" are not supported by the paper in its current form.

Thank you for this comment. We agree in the first version of our paper, this claim was not supported well but after fixing our code (see "general response to reviewers") we have been able to make an explicit conclusion that our method does not fail for the problems where the additive Gaussian observation noise is large, and demonstrate on-par restoration results in terms of FID, KID, LPIPS, MSE, PSNR and SSIM in Table 2 and 5 for the low (\sigma_y = 0.01) observation noise problems and better restoration results in Table 2 and 5 for noisy (\sigma = 0.05, 0.1) problems.

Q1. How does the proposed technique compare quantitatively (both perception and distortion metrics) to prior work on the image datasets?

Thank you for this suggestion. We have now compared quantitatively on both perception (FID, KID, LPIPS, SSIM) and distortion (MSE, PSNR) metrics. Table 2 shows that our method is on-par with Pi-GDM and DPS for low noise (\sigma_y = 0.01) problems and significantly better in terms of model consistency (FID and KID) for noisy (\sigma_y = 0.05, \sigma_y = 0.1) problems, and Table 5 shows that our method is on-par or better in terms of observed data consistency (LPIPS, MSE, PSNR, SSIM).

Q2. What has been done to ensure fair comparison with competing techniques in terms of hyperparameter tuning?

Before our update to the paper, the paper uses the hyperparameters suggested by the authors, but since there was a bug adding much more observation noise than the hyperparameters were set to, the experimental setup was not correct. This negatively affected the performance of TMPD as well as the benchmarks DPS and PiGDM, since TMPD and PiGDM use the noise standard deviation in the algorithms, and the DPS hyperparameter was not optimal. Please see our "General response to reviewers" for a full explanation.

After correcting the experimental setup, the noise standard deviation is the correct level for the observations. We have made sure to select optimal DPS hyperparameters by cross-validating over LPIPS, MSE, PSNR and SSIM distance metrics (see, for example, Fig. 4 for a demonstration). We use we the hyperparameters and algorithm suggested in the original paper for PiGDM and can motivate this choice, since the ablation studies of Table 6, Table 10, and Table 11 in Song et al., 2021a show that the use of $\eta=1.0$ in combination with a large number of noise-levels (we use 1000) gives optimal performance, and stronger performance than DPS in for low observation noise problems. The PiGDM methodology is similar to ours in the derivation of step-size but makes an impractical assumption that the image distribution is standard normal distribution, and so there is no need to select a step-size hyperparameter for the PiGDM method.

We thank the reviewer for their insightful comments. We are indeed happy to see that the reviewer observed the main strengths of our method, namely removing hyperparameter tuning requirement, as well as having excellent performance. We are also happy to see that the reviewer found our paper well-written, thank you for these comments.

We address your concerns below point-by-point.

Weakness 1. The experiments are insufficient to verify the practicality of the method. Beyond the synthetic examples, no quantitative comparison is provided for actual image reconstruction problems (inpainting, SR). I recommend performing experiments in the same setting as DPS or in the original papers and compare the following metrics: PSNR, SSIM, LPIPS, FID. The proposed method has greatly increased compute cost, and therefore it would only make practical sense if sample quality is demonstrated to be significantly better than in competing techniques.

We have now a comprehensively updated experimental results section with the improved code and clarity (see our "General response to reviewers"). In particular, we now provide quantitative comparison for the inpainting and super-resolution experiments on CIFAR10 (which is the same setting as in Song et. al "Score-Based Generative Modeling through Stochastic Differential Equations" https://arxiv.org/abs/2011.13456), comparing to DPS and PiGDM benchmarks using perception and distance metrics in Tables 2 and 5, which shows increased performance compared to the alternative methods for problems with large Gaussian observation noise problems, and comparitively good performance in low noise problems. To expand on this point, Table 6 presents a comparison to DPS and PiGDM baselines on CIFAR-10, which demonstrates better consistency with the whole dataset through FID and KID scores, and better consistency with the data through LPIPS, MSE, PSNR and SSIM scores.

For small (\sigma_y=0.01) noise on 2 times 'nearest neighbour' super-resolution: FID 34.5 TMPD-D is better compared to 43.4 DPS-D and 38.7 PiGDM-D. LPIPS 0.140 TMPD-D is comparible compared to 0.136 DPS-D and 0.126 PiGDM-D.

For medium (\sigma_y=0.05) noise on 2 times 'nearest neighbour' super-resolution: FID 37.0 TMPD-D is better compared to 101 DPS-D and 109 PiGDM-D. LPIPS 0.223 TMPD-D is better compared to 0.286 DPS-D and 0.313 PiGDM-D.

For large (\sigma_y=0.1) noise on 2 times 'nearest neighbour' super-resolution: FID 37.2 TMPD-D is better compared to 140 DPS-D and 161 PiGDM-D. LPIPS 0.296 TMPD-D is better compared to 0.420 DPS-D and 0.451 PiGDM-D.

The numerical results for all of the other inverse problems (4 times 'bicubic' super-resolution, ‘half’ mask inpainting and ‘half’ mask inpainting) are similar, showing TMPD has strong performance for large noise compared to DPS and PiGDM that are less (numerically and visually) consistent with the model or data for large noise. We have made sure to select optimal DPS hyperparameters by cross-validating over LPIPS, MSE, PSNR and SSIM distance metrics (see, for example, Fig. 4 for a demonstration) and use we the hyperparameters and algorithm suggested in the original paper for PiGDM. We also would like to clarify that TMPD does not require hyperparameter tuning (as also pointed out by other reviewers) which is a big plus in terms of computational cost.

Furthermore, we would like to note that we also have a very practical advantage as we remove the requirement of hyperparameter tuning (as noted by other reviewers as "strengths"). This means that our method removes a very costly step that is mostly hidden in papers - but very crucial in practice. Also, our inpainting and super-resolution sampler is only 1.4-1.6 times slower when compared to DPS and PiGDM samplers across all of our image experiments.

We thank the reviewer for their careful and informative review. We followed most of their suggestions and addressed questions, which significantly improved the quality of our work. We are happy to see that the reviewer found our work well written and observed that our theoretical results could indeed help analysing similar algorithms for inverse problems. We address weaknesses and questions below.

I have identified two significant limitations:

Weakness 1. Practical Usability and Compute Cost: One major limitation I've identified is the insufficient discussion of practical usability, particularly in terms of the associated compute cost. Although the authors briefly mentioned the drawback that "TMPD is slower, as each iteration costs more memory and compute..." in the concluding section, I believe that this point requires more emphasis to address its impact on real-world applications.

We found that, for real world applications, sampling is 1.4-1.6 times slower than DPS and PiGDM for inpainting and superresolution. However, we would like to point out, with our updated results comparing to DPS and PiGDM benchmarks using perception and distance metrics in Tables 2, 5 and 6, this comes with an increased performance. We also would like to clarify the removal of hyperparameter tuning (as also pointed out by other reviewers) is a big plus in terms of computational cost.

To make these points clear in our updated manuscript, we have added a Section 3.4 that discusses the computational cost of the additional calculation, and suggests a computationally cheaper approximation of the Gaussian projection.

Considering now the practical limitations of our method, for high dimensional problems where the row sum of the jacobian is not approximately equal to the diagonal of the jacobian, which occurs for non-linear problems or problems with lots of mixing components in the obervation matrix H, then computational cost of calculating and inverting the full jacobian is much higher. To treat this case, one could explore more advanced low rank or ensemble approximations to the Jacobian, which we leave to future work. For low dimensional problems this is not an issue, and our method can be used as an efficient monte-carlo sampler, as demonstrated in the synthetic examples, where our method clearly outperforms DPS and PiGDM.

To improve the clarify of writing, I recommend moving a summary of the discussion on computation cost comparison and cheaper approximations from Appendix E.1 into the main text. This change would provide readers with a clearer understanding of the computational trade-offs and make the paper more accessible to a broader audience.

Thank you for this comment. We have summarised Appendix E.1 in the main text, which is now Section 3.4 in the updated manuscript, and improved the clarity in the appendix.

On the practical side, I doubt if this second-order approximation can be practical without resorting to low-rank approximation or other proxy approaches (while the authors assert the computation is not expensive for inpainting but that cannot be generalized to other tasks).

We agree with the reviewer a full second-order approximation comes with a significant cost - but we would like to make a few arguments why this cost might be justified even in the setting of diagonal/low-rank approximations. In low dimensions, of course our method has a significant advantage of having a full-preconditioner in effect. In approximate settings (where this second-order approximation is replaced by a diagonal or low-rank matrices), we still argue that this brings a dimension-wise scaling that adapts the diffusion sampler and improves the condition number during sampling. In fact, low rank or diagonal preconditioning approaches are very popular and successful in their own right (e.g. Adam and RMSprop in stochastic optimisation), therefore we think that developing a principled way to obtain these matrices are invaluable.

Weakness 2. Limited Empirical Evaluation: Another significant limitation is the limited empirical evaluation presented in the paper. While the paper showcases samples generated by the proposed method and other competing methods for inpainting and super-resolution tasks, it is not immediately evident how the performance of the proposed method differs from the alternatives. The paper lacks a comprehensive summary of the qualitative differences observed in these samples. To address this limitation, I suggest incorporating quantitative metrics, such as FID or IS.

We have now a comprehensively updated experimental results section with the improved code and clarity (see our "General response to reviewers"). In particular, we have added results comparing to DPS and PiGDM benchmarks using perception and distance metrics in Tables 2 and 5, which shows increased performance compared to the alternative methods for problems with large Gaussian observation noise problems, and comparitively good performance in low noise problems. To expand on this point, Table 2 presents a comparison to DPS and PiGDM baselines on CIFAR-10, which demonstrates better consistency with the whole dataset through FID scores, and better consistency with the data through LPIPS.

For small (\sigma_y=0.01) noise on 2 times 'nearest neighbour' super-resolution: FID 34.5 TMPD-D is better compared to 43.4 DPS-D and 38.7 PiGDM-D. LPIPS 0.140 TMPD-D is comparible compared to 0.136 DPS-D and 0.126 PiGDM-D.

For medium (\sigma_y=0.05) noise on 2 times 'nearest neighbour' super-resolution: FID 37.0 TMPD-D is better compared to 101 DPS-D and 109 PiGDM-D. LPIPS 0.223 TMPD-D is better compared to 0.286 DPS-D and 0.313 PiGDM-D.

For large (\sigma_y=0.1) noise on 2 times 'nearest neighbour' super-resolution: FID 37.2 TMPD-D is better compared to 140 DPS-D and 161 PiGDM-D. LPIPS 0.296 TMPD-D is better compared to 0.420 DPS-D and 0.451 PiGDM-D.

The results for all of the other inverse problems (4 times 'bicubic' super-resolution, ‘half’ mask inpainting and ‘box’ mask inpainting) are similar, showing TMPD has strong performance for large noise compared to DPS and PiGDM that are less (numerically and visually) consistent with the model or data for large noise. We have made sure to select optimal DPS hyperparameters by cross-validating over LPIPS, MSE, PSNR and SSIM distance metrics (see, for example, Fig. 4 for a demonstration) and we then use the hyperparameters and algorithm suggested in the original paper for PiGDM. We also would like to clarify that TMPD does not require hyperparameter tuning (as also pointed out by other reviewers) which is a big plus in terms of computational cost.

Q1. Page 3, bottom part, "the central to our method" is this a typo? this phrase is not clear to me.

Thank you for pointing this out, we have fixed this. This sentence now reads:

Our method relies on the approximation $p_{0|t}(\mathbf{x}\_{0}|\mathbf{x}\_{t}) \approx \mathcal{N}\left(\mathbf{x}\_0; \mathbf{x}\_{0|t}(\mathbf{x}\_t), \mathbf{C}\_{0 | t}(\mathbf{x}\_t)\right)$ to make the sampling process tractable.

Q2. Theorem 1: can yo unpack how "we can bound the approximation error of our SDE solely due to approximations made for the likelihood."? The bound in the theorem depends on H, and the property of Phi (which characterizes the behavior of how true data density deviates from a Gaussian). Where does the approximation to the likelihood enter in to this upper bound?

Thank you, the passage was indeed confusing and we removed it. $H$ and the noise-level do enter the bounds, but we have not made that dependence explicit. We did so, since we were more interested in the dependence of the approximation error due to our distribution being non-Gaussian (as measured in $M$ and $L$), instead of the dependence on the problem at hand. We leave adding the explicit dependence on $\|H\|$ and $\sigma_y$ for future work.

Q3. On page 8 "our chosen hyperparameter η = 1.0 (Song et al., 2021a)". Have you defined?

Thanks for this comment. We have now made it clear to refer to Algorithm 1 in Song et al., 2021a for the definition of this hyperparameter. We chose not to elaborate further since we only consider DDPM and so the hyperparameter is not of use after setting it to $\eta=1.0$: whilst Song et al., 2021a use a DDIM sampler, they show that by setting $\eta=1.0$, this is equivalent to a DDPM-type sampler. We can motivate this choice, since the ablation studies of Table 6, Table 10, and Table 11 in Song et al., 2021a show that the use of $\eta=1.0$ in combination with a large number of noise-levels (we use 1000) gives optimal performance, and stronger performance than DPS for low observation noise inverse problems.

Q4. On page 9: "explore more efficient sampling techniques, such as ensemble methods or low rank approximations" can you expand on how ensemble methods can be useful here?

Ensemble methods can be used to make an approximation to the covariance with a low rank equal to the number of particles in the ensemble. The covariance we calculate via the jacobian is the covariance of the ensemble cloud of particles at the current time. We think this could be useful but will leave this direction to future work.

Q5. On page 9 "On the positive side, we found our approach more robust across various problems due to the stabilising effect of better approximations.", similar to my comment in Weakness 2, I don't see how this point is supported other than the synthetic results.

We agree with the reviewer that the original submission did not support this comment, but after fixing our code (see "general response to reviewers") we have been able to make an explicit conclusion that our method does not fail for the problems where the additive Gaussian observation noise is large.

Q6. On page 18 the compute cost: what is N? The number of data points to generate?

We apologise for the confusion, $N$ is the number of noise-scales (equivalently the number of steps) the diffusion model is run for. We have now clarified this in Section E.1.

Conclusion

We thank the reviewer for their thoughtful comments once again which helped to increase the quality of our manuscript significantly. We sincerely believe that our paper is now much higher quality. We therefore kindly remind the reviewer to reevaluate our paper, if they agree with us, and update their score if they are happy with our rebuttal. Thank you.

Thank you for your review and discussion. We have updated the synthetic experiments to include the diagonal approximation the full Jacobian with TMPD for DDPM (denoted DTMPD-D in the updated summary Table 2) and we are working on more empirical results on established benchmark tasks, which will be uploaded before the end of the discussion period. All reviewers should note that this revision adds a Table 3 that in turn renames the tables that were part of the discussion, we have already made the edits which correct for this, and all of our comments should have the correct table numbers.

While I agree that low-rank approximation is a common practice in other fields, I don't see in the setting of this paper how to directly derive a low-rank approximation (or develop a good preconditioner for the matrix solve). Can the authors shed some light on this point?

The ensemble Kalman filter provides methodology to do a low rank approximation using J particles, where J is the rank of the approximation, by computing an empirical covariance matrix with J particles. The algorithm will only operate on the J particles, so should have only O(J d_y) memory complexity. We leave this extension to future work. We are happy to provide further references for this methodology upon request.

Thank you again to the reviewer for these questions on the numerical implementation of our method, and we will be updating the synthetic experiments section tomorrow as a result of the suggestions given by the reviewer.

While the diagonal approximation to the Jacobian matrix (termed as DTMPD) seems reasonable in high-dimensional setting, are the experiment results presented in the main paper (mainly Table1, Figure 2, 3, Table 2) from the non-diagonal, i.e. full-rank Jacobian method (TMPD)? If so, can you list the practical memory and time comparison of TMPD(-D) to DPS(-D) and ΠGDM(-D)? I think that would be a more comprehensive comparison for practical reference. In addition, do you have results (generated samples and test statistics e.g. FID) for the DTMPD variant?

The experimental results presented in the synthetic experiments, Table 1, and Figure 2, are using the full-rank jacobian matrix. On the other hand, all of the image experiments, Figure 3 and Table 2, use the row sum approximation. We have listed the theoretical memory and time complexity of TMPD-D compared to DPS-D and PiGDM-D, and, in practice, TMPD-D with the row sum approximation is 1.4-1.6x slower than DPS-D and PiGDM-D and has the same memory costs as DPS-D and PiGDM-D. For the diagonal approximation and the full Jacobian, TMPD-D can be 10x slower than DPS-D and Pi-GDM-D, but the memory cost of computing the full Jacobian is the limiting factor, since it means that only a small batch size can be used for large images.

Thanks for your suggestion to do experiments for the DTMPD variant. In the time period of the review discussion we can only do synthetic experiments for this, and we are now running the synthetic experiments that we will add to Table 1 and we will upload tomorrow. It would indeed be interesting to see how this diagonal approximation performs on the synthetic experiments, and would bring Table 1 in line with Table 4, that does already compare the DTMPD variant, but for the score-based reverse-SDE.

As mentioned above, performing the diagonal approximation variant on the image experiments is computationally infeasible for the timeframe of the review discussion because the best method that is currently possible to compute the diagonal of the jacobian using JAX's transformations involves infact computing the full Jacobian, and then taking the diagonal. We expand on this point by providing commented pseudo-code, below. Our rowsum approximation was used in all of the image experiments, as we noted in section 3.4 of the paper (and also Appendix E.1 in the original submission). The diagonal approximation serves as a stepping stone to give both inspiration and justification for the rowsum approximation.

> def get_grad_estimate_x_0_vmap(observation_map):
>     """Method as suggested in https://stackoverflow.com/questions/70956578/jacobian-diagonal-computation-in-jax """
> 
>   def estimate_x_0_single(val, i, x, t, timestep):
>         x_shape = x.shape
>         x = x.flatten()
>         x = x.at[i].set(val)
>         x = x.reshape(x_shape)
>         x = jnp.expand_dims(x, axis=0)
>         t = jnp.expand_dims(t, axis=0)
>         v = discrete_sigmas[timestep]**2
>         s = score(x, t)
>         x_0 = x + v * s
>         h_x_0 = observation_map(x_0)
>         return h_x_0[i]
>     return vmap(value_and_grad(estimate_x_0_single), in_axes=(0, 0, None, None, None))
> 
> grad_estimate_x_0_vmap = get_grad_estimate_x_0_vmap(observation_map)
> 
> def analysis(y, x, t, timestep, ratio):
>     h_x_0, vjp_h_x_0, (_, x_0) = vjp(
>     lambda x: estimate_h_x_0_vmap(x, t, timestep), x, has_aux=True)
> 
>     # There is no natural way to get the diagonal of a Jacobian with JAX's transforms:
>     # you cannot map the input, because in general each diagonal entry of the jacobian depends on all inputs.
> 
>     # Calculating the full Jacobian seems like the best method, but the batch size cannot be large due to memory constraints, and so numerical evaluation is slow
>     vec_vjp_h_x_0 = vmap(vjp_h_x_0)
>     diag = jnp.diag(batch_observation_map(vec_vjp_h_x_0(jnp.eye(y.shape[0]))[0]))
> 
>     # # This method gives OutOfMemory error
>     # idx = jnp.arange(len(y))
>     # h_x_0, diag = grad_estimate_x_0_vmap(x.flatten(), idx, x, t, timestep)
>     # diag = self.observation_map(diag)
> 
>     # # This method cannot be XLA compiled and is too slow for numerical evaluation experiments
>     # diag = jnp.empty(y.shape[0])
>     # for i in range(y.shape[0]):
>     #   eye = jnp.zeros(y.shape[0])
>     #   diag_i = jnp.dot(self.observation_map(vjp_h_x_0(eye)[0]), eye)
>     #   diag = diag.at[i].set(diag_i)
> 
>     C_yy = diag + self.noise_std**2 / ratio
>     ls = vjp_h_x_0((y - h_x_0) / C_yy)[0]
>     return x_0.squeeze(axis=0) + ls

Q1. Can you explain in which aspects, the added posterior covariance matrix is important for the conditional diffusion models?

The covariance approximation in the Tweedie moment projection leads to the formula in Eq. (8) where we obtain a full covariance in our approximate likelihood, rather than a tunable scalar one as used in PiGDM. Algorithmically, this acts similar to a preconditioner, scaling each dimension appropriately with respect to the uncertainty, improving the conditioning of the drift. Even though we end up approximating this matrix with a diagonal approximation, the effect of this adaptive scaling is still apparent and results in improved performance.

Q2. Can you add some numerical results on the image restoration tasks as I said in the previous section? Thanks very much.

Thanks for this comment. We have now added numerical results on the image restoration tasks, and refer the reviewer to our comment for "weakness 1".

Conclusion

We thank the reviewer for their comments which, in our opinion, significantly improved the quality of our paper and made it more comprehensive. We believe that with the added experimental results as well as clarifications and improvements on theory, our paper is much higher quality, therefore also deserves another evaluation. We would be very happy if the reviewer can reevaluate our paper and increase the score if the updated paper meets the standards of ICLR in their view. Thank you.

We thank the reviewer for their insightful and helpful comments. We are happy to see that the reviewer found our writing quality high and our experimental setting solid. We aim to address their points below point-by-point, especially their concerns about quantitative metrics and theoretical results.

Weakness 1. It would be better to show us some numerical results under some image metrics like PSNR or FID of the generated images, with the results of baseline methods added, so that the audience will know whether the image quality really increases numerically.

We have now comprehensively updated experimental results section with the improved code and clarity (see our "General response to reviewers"). In particular, we have added results comparing to DPS and PiGDM benchmarks using perception and distance metrics in Tables 2, 5 and 6, which shows increased performance compared to the alternative methods for problems with large Gaussian observation noise problems, and comparitively good performance in low noise problems. To expand on this point, Table 6 presents a comparison to DPS and PiGDM baselines on CIFAR-10, which demonstrates better consistency with the whole dataset through FID and KID scores, and better consistency with the observation data through LPIPS, MSE, PSNR and SSIM scores.

For small (\sigma_y=0.01) noise on 2 times 'nearest neighbour' super-resolution: FID 34.5 TMPD-D is better compared to 43.4 DPS-D and 38.7 PiGDM-D. LPIPS 0.140 TMPD-D is comparible compared to 0.136 DPS-D and 0.126 PiGDM-D.

For medium (\sigma_y=0.05) noise on 2 times 'nearest neighbour' super-resolution: FID 37.0 TMPD-D is better compared to 101 DPS-D and 109 PiGDM-D. LPIPS 0.223 TMPD-D is better compared to 0.286 DPS-D and 0.313 PiGDM-D.

For large (\sigma_y=0.1) noise on 2 times 'nearest neighbour' super-resolution: FID 37.2 TMPD-D is better compared to 140 DPS-D and 161 PiGDM-D. LPIPS 0.296 TMPD-D is better compared to 0.420 DPS-D and 0.451 PiGDM-D.

The numerical results for all of the other inverse problems (4 times 'bicubic' super-resolution, ‘half’ mask inpainting and ‘half’ mask inpainting) are similar, showing TMPD has strong performance for large noise compared to DPS and PiGDM that are less (numerically and visually) consistent with the model or data for large noise. We have made sure to select optimal DPS hyperparameters by cross-validating over LPIPS, MSE, PSNR and SSIM distance metrics (see, for example, Fig. 4 for a demonstration) and use we the hyperparameters and algorithm suggested in the original paper for PiGDM. We also would like to clarify that TMPD does not require hyperparameter tuning (as also pointed out by other reviewers) which is a big plus in terms of computational cost.

Weakness 2. In Proposition 3, the authors assumed that the initial distribution $p_{data}$ is Gaussian, which is understandable since $p(x_0 | x_t)$ is approximated by using a posterior Gaussian. However, I would say the assumption is very strong and impractical. It would be better if the authors can figure out a way to propose a better theoretical result with weaker assumptions.

Thank you for this comment. We would like to note Proposition 3 only outlined a specific case to show our approximation's validity in the Gaussian data case. We extended this result in Theorem 1 to general non-Gaussian distributions. To clarify this, we have now added (Gaussian data distribution) note to Proposition 3 and (General data distribution) note to Theorem 1.

In Theorem 1, we assume that if the distribution is non-Gaussian, we still get an upper bound on the distance to the true posterior. How far the distance to the true posterior is, depends on the non-Gaussianity of the distribution. This is however an inherent property of our method, since we made a Gaussian approximation. Note, that none of the methods that we compare to can give similar guarantees to neither Proposition 3 or Theorem 1, since their approximations are not exact in the Gaussian setting.

Weakness 3. Compared with DPS, the authors added a posterior covariance matrix, which makes the contribution quite incremental. Unless the authors can show that the added covariance can greatly improve the performance compared to DPS, or I afraid the contribution is not enough for an "accept" to a more and more competitive ML conference.

In our newly shown experiments, (see Tables 2 and 5), we indeed demonstrate that our method does increase quantitative metrics significantly w.r.t. state-of-the-art. Furthermore, we would like to note that we also have a very practical advantage as we remove the requirement of hyperparameter tuning (as noted by other reviewers as "strengths"). This means that our method removes a very costly step that is mostly hidden in papers which is crucial in practice. Also, our inpainting and super-resolution sampler is only 1.4-1.6 times slower when compared to DPS and PiGDM samplers across all of our image experiments.

Q4. What is the intuition behind the specific family of data distribution considered in Theorem 1? Is it the setting where the authors could prove something? Or does it capture any meaningful data distributions?

The goal of the Theorem is to quantify the error to the posterior. Our method is exact in the Gaussian case (Proposition 3). Due to the derivation, we cannot prove (or expect) exact posteriors in other cases. What we do in Theorem 1 is to quantify the distance of our distribution to the Gaussians in some way, and then give quantitative guarantees on the distance of our generative model to the true posterior.

In general, any distribution that can be written as $\exp(f(x))$ with a $f \in C^1$ satisfies the assumptions of the Theorem, since we can just use

$\exp(f(x)) = \exp(f(x) - \frac{1}{2} \|x - \mu_0\|^2_{\Sigma_0})\exp(\frac{1}{2} \|x - \mu_0\|^2_{\Sigma_0}) = \exp(f(x) - \frac{1}{2} \|x - \mu_0\|^2_{\Sigma_0})\mathcal{N}(x;\mu_0, \Sigma_0).$

Then, the theorem gives theoretical guarantees on the performance of our method, depending on the properties of $f(x) - \frac{1}{2}\|x - \mu_0\|^2$, or equivalently, how similar $\exp(f(x))$ is to a Gaussian.

The argument could be extended to more general distributions by first convolving them with a Gaussian. One then gets a representation as above. The authors in [2] and [C] used similar arguments.

Summarizing, we extend Proposition 3, where we say, the error of our model on Gaussian distributions is 0, to a statement quantifying that if our distribution has distance $\text{dist}$ (measured via $M$ and $L$), then our approximation error can be bounded by an explicit formula in terms of $\text{dist}$. We hope this demonstrates that the choice of distributions is not arbitrary, but natural due to the derivation of our method. Note, that none of the baseline methods that we compare have similar guarantees, also not for strictly Gaussian measures.

[C] Chen, Lee, Lu, "Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions"

Q5. Section 6.2: The largest standard deviation of the added noise is 0.01 and the baselines still perform poorly as per the experiments reported in this paper. However, baselines such as DPS perform much better even with higher noise level (sigma = 0.05).

As we noted in our "general response to reviewers", this caused by an error on our side. We have now comprehensively updated these results, please see our new experimental results.

Q6. Super-resolution results of the proposed method TMPD in Fig 4 are equally bad with the unwanted color artifacts. Maybe take more measurements and compare.

We have repeated the experiment, this time correcting the noise standard deviation parameter in the code (please see our "general response to reviewers") and we have compared FID, KID and LPIPS metrics for 2 times super-resolution and 4 times super-resolution over a 1k CIFAR-10 validation set. We can report significantly better results when compared to DPS and PiGDM, especially when the Gaussian noise is large. TPDM is sucessful where DPS and PiGDM methods fail when the noise is large. We are happy to provide the samples used to generate the Tables upon request (as the files are quite large), and our code has been updated with this experiment that can be run to reproduce all of the tables.

Q7. Typo in the first equation on Page 18. Missing gradient

Thank you - we have fixed this. This is now below Eq. (33) in the new manuscript.

Q8. Typo on pp 14 in KL divergence. Missing gradient

Thank you! This is fixed, this is now on page 15 in the new PDF.

Q9. What is the scale of M in Theorem 1? Is the dependence on M optimal?

The scale of M is determined by the proximity of $p_{\text{data}}$ to a Gaussian. An interesting sidenote might be that the assumption implies that

$\|p_\text{data} - \mathcal{N}(\mu_0, \Sigma_0)\|_{\text{TV}} \le M -1$

Therefore, it can be interpreted as bound on the total variation distance to the Gaussians. We do not know if the $M$ dependence is optimal but happy to report that in the updated proof, we have significantly improved dependence on $M$. We think that this would be an exciting direction for future research.

Conclusion

We sincerely believe that, with the added experimental results, as well as our updated theoretical section in response to the reviewer's comment, our paper has now a much better standing. We also kindly ask reviewer to increase their score if they are satisfied by the response to acknowledge that the points are addressed and that the paper is now in good standing from a scientific perspective. Thank you.

Weakness 4. Page 9: “The ability to handle measurement noise also gives TMPD the potential to address certain real-world applications, such as CT and MRI imaging with Gaussian noise” I am not sure how TMPD handles the noise. The results are not as good as the results reported in the baseline papers, piGDM and DPS in a similar setting. Also, it is slower than piGDM.

As we have noted above, and in our general comment, the comparisons in the original PDF were not correct due to a bug. As many reviewers rightly pointed out, we investigated this issue and now we are confident that our results reflect the improved quality of our samples.

Table 2 presents a comparison to DPS and PiGDM baselines on CIFAR-10, which demonstrates better consistency with the whole dataset through FID scores, and better consistency with the data through LPIPS.

For small ($\sigma_y=0.01$) noise on 2 times 'nearest neighbour' super-resolution: FID 34.5 TMPD-D is better - compared to 43.4 DPS-D and 38.7 PiGDM-D. LPIPS 0.140 TMPD-D is competitive - compared to 0.136 DPS-D and 0.126 PiGDM-D.

For medium ($\sigma_y=0.05$) noise on 2 times 'nearest neighbour' super-resolution: FID 37.0 TMPD-D is better compared to 101 DPS-D and 109 PiGDM-D. LPIPS 0.223 TMPD-D is better compared to 0.286 DPS-D and 0.313 PiGDM-D.

For large ($\sigma_y=0.1$) noise on 2 times 'nearest neighbour' super-resolution: FID 37.2 TMPD-D is better compared to 140 DPS-D and 161 PiGDM-D. LPIPS 0.296 TMPD-D is better compared to 0.420 DPS-D and 0.451 PiGDM-D.

The results for all of the other inverse problems (4 times 'bicubic' super-resolution, ‘half’ mask inpainting and ‘box’ mask inpainting) are similar, showing TMPD has strong performance for large noise compared to DPS and PiGDM that are less (numerically and visually) consistent with the model or data for large noise. We have made sure to select optimal DPS hyperparameters by cross-validating over LPIPS, MSE, PSNR and SSIM distance metrics (see Fig. 4 for a demonstration of our hyperparameters search for DPS) and use we the hyperparameters and algorithm suggested in the original paper for PiGDM.

As for the computational cost, we agree that our method is 1.4-1.6x slower for a single run, this is the cost (in this case) for increased quality, coming from more accurate approximations. But we would like to point out that we have no hyperparameters in our setting (as the reviewer rightly observed in their 2nd point in "Strengths"), this means we have to factor in the cost of hyperparameter search. We claim that our method could be the method of choice due to removing this requirement (which is a necessity every time there is a new inverse problem to solve), despite the increase of the sampling cost. Furthermore, we present a mature and efficient codebase that can be used to reproduce all of our results and will be made public after the review period.

Q1. Missing citation: The score is trained using a denoising score-matching objective [1].

Thank you. We have added this citation in Section 2, technical background "This can be circumvented by approximating the drift using score matching techniques...".

Q2. What is C_{0|n} in the last paragraph on Page 4?

Thanks for pointing out. It's meant to read $C_{0|t}$. This is the same quantity given in Prop. 1, which is the posterior covariance of $x_0|x_t$.

Q3. In equation 9, it seems like the first gradient in the right hand side is over the entire function, not just the mean, which is not correct. I ask the authors to put the brackets appropriately to make it clear.

While putting brackets in the newly updated version posed a challenge of equation being too long, we have now added a sentence right after this equation that reads "where $\nabla_{{\mathbf{x}\_t}}$ only operates on $\mathbf{m}\_{0|t}$ in the above equation."

We thank the reviewer for their careful review. The reviewer's comments resulted in significant increase in the quality of our results and we are grateful for this. We are also happy that the reviewers observed the strengths of our paper, and we reply below to both "weaknesses" and "questions" that come after.

Weakness 1. The theoretical analysis is a simple extension of Chen et al. [2].

We bound the TV-Distance to the target via a triangle inequality and then bound one of the terms using Pinskers-Inequality, to get a KL-Term. This is an attractive argument in this setting, introduced in the compelling work [2]. However, the major part of the proofs (from the conceptual difficulties as well as length-wise) consists in bounding the resulting terms, and in their case showing that one can apply Girsanov's theorem, for them as well as for us. However, these steps are inherently different:

In our case, the approximation in our analytical extra drift comes from the fact that our distribution is not Gaussian (in which case the approximation would be exact). Therefore, the core of the proof consists in bounding these approximations under our assumptions on $p_\text{data}$. We do not treat time discretization. On the other hand, in [2], the errors stem from the drift being approximated by a neural network and the time-discretization of the SDE. The inherent different setting makes the contents of the proofs very different.

We already cite [2], but we agree that another citation, directly at the beginning of our proof where we use the argument, is beneficial to the reader. We added the citation to the uploaded PDF. Furthermore, we hope that the preceding discussion highlighted that the proofs are distinct otherwise.

Weakness 2. The authors use the results from Girsanov’s theorem. But there is no proof whether the condition to apply Girsanov holds in this case.

Thank you for reading the proof so carefully. You are right, that we should justify our use of Girsanov. We now do so by using Novikov's condition. Again, we can make use of the Gaussianity, and basically just apply the bounds on the differences in the drift that we have proven before: Since we have bounded $|\nabla \log p_t + f_t - p_t|$ by a quantity depending only on $M$, $L$ and $\|x\|$, Novikovs condition reduces to showing that

$\mathbb{E}[\exp(\frac{1}{2} c \int_0^T \|x_t\|^2 \mathrm{d}t)] < \infty.$

However, the expectation is taken with respect to the Gaussian reference process. We want to apply Novikov's condition iteratively and therefore prove that

$\mathbb{E}[\exp(\frac{1}{2} c \int_{t_i}^{t_i+1} \|x_t\|^2 \mathrm{d}t)] < \infty$

for $\Delta_i = t_{i+1} - t_i$ small enough. Since our process is well behaved (Gaussian process with bounded mean and lower/upper bounded covariance) we can show that we can lower bound the size of $\Delta_i$ and iterate this procedure. This argument is taken from Corollary 3.5.16 in [B].

[B] Karatzas and Shreve, "Brownian motion and stochastic calculus"

Weakness 3. The experimental results do not seem correct. In fact, I suspect the baseline results, e.g. DPS results are wrong. I have tested many of the baselines myself and the results are not as bad as what the authors have shown in this paper, zoom in Fig 3, Fig 7 for instance to see the artifacts.

The experimental results for the image experiments were indeed not correct due to a typo/bug in the code (please see our "general response to reviewers").

The colour artifacts in Fig 3 and Fig 7 were due to the noise_std being very small compared to the observation noise (due to a typo in the code), as a result the DPS hyperparameter was not set correctly. We have now corrected the bug and have also optimised over LPIPS, SSIM, MSE and PSNR to choose the DPS scale hyperparameter in all image experiments. This is the only part of the code that has changed in our updated experiments. We believe the DPS implementation is correct as it closely follows the original code by Chung et al. Although the code has not been released for \PiGDM, we believe our code is correct as it follows their algorithm exactly.

Dear reviewers,

We thank you for your initial reviews, and thank Reviewer uiGE for their timely response and discussion to our rebuttal.

We kindly ask all other reviewers (NB1r, ixot, 3wW3) to read our rebuttal (general and specific responses) and please reevaluate our paper. We genuinely believe that we have addressed all main concerns, with a significant update to empirical results (including all quantitative metrics requested) as well as mathematical clarification of theoretical results. We would really appreciate if you could take a look and reevaluate your scores as we think the current form of paper deserves a reevaluation.

Thank you for your service!