Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance

Q3: Why we choose 50 sampling steps?

We implement our techniques based on an open-source diffusion codebase k-diffusion (https://github.com/crowsonkb/k-diffusion), where using 50 sampling steps is a default choice. To further demonstrate the effectiveness of our method, we conducted additional experiments with 100 and 1000 sampling steps using Euler's one order sampler (deterministic) following the conventional choice of DPS and $\Pi$ GDM, shown in the following tables.

Inpainting:

Gaussian debluring:

Motion debluring:

Super resolution:

Q4 & Q5: Grammar issues

We thank the reviewer for the suggestions. We have fixed these issues in our revised version.

Reference

[1] Song, Y., et al. (2020). "Score-based generative modeling through stochastic differential equations." arXiv preprint arXiv:2011.13456.

[2] Karras, T., et al. (2022). "Elucidating the design space of diffusion-based generative models." arXiv preprint arXiv:2206.00364.

[3] Lu, Cheng, et al. "Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps." Advances in Neural Information Processing Systems 35 (2022): 5775-5787.

[4] Bao, F., et al. (2022). "Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models." arXiv preprint arXiv:2201.06503.

[5] Bao, Fan, et al. "Estimating the optimal covariance with imperfect mean in diffusion probabilistic models." arXiv preprint arXiv:2206.07309 (2022).

W6: How bad are the samples when using optimal variances for all steps

The experiment results for Type I guidance using Analytic-type posterior variances for all sampling steps (referred to as Analytic (all steps)) are shown in the following tables and have been included in Appendix E of our revised version. We do not conduct experiments for Convert-type posterior variances since we cannot obtain accurate posterior variances prediction using Equation 22 at high noise level region (see 5.1 Validation of theoretical results, posterior variances prediction).

Inpainting:

Gaussian debluring:

Motion debluring:

Super resolution:

As can be seen, while using optimal variances in all steps shows good performance in tasks such as inpainting and super-resolution, there is a significant performance decrease for the deblurring tasks. Example samples are given in the file “debluring_with_analytic_variances_for_all_steps.png” of Supplementary Material. Our solution, which considers sampling with optimal variance only in regions with low noise levels, demonstrates robust performance across all tasks, including the new colorization tasks brought up by Reviewer rzGL.

Q1: Computational cost of equation 24

Similar to [4], a small number of Monte Carlo samples is typically adequate for obtaining an accurate $\hat{r}_t$. With a pre-trained unconditional model and training dataset, we compute $\hat{r}_t$ offline and deploy it alongside the pre-trained unconditional model. As a result, the online inference cost of our method is the same as previous methods. In our paper, we initially calculate $\hat{r}_t$ using a randomly selected $0.5 \%$ of the training set and reuse it for all our experiments. This process takes a couple of hours on a single 1080Ti GPU.

Q2: Number of sampling steps with optimal posterior variances

As we sample more densely in regions with low noise levels suggested by [2], it corresponds to 12 out of 50 steps. We have added this information in footnote 7 of our revised version.

W1: Limited contribution

We thank the reviewer for comments on our paper. We would like to emphasize that the main contribution of this paper is to provide a novel framework that can unify existing diffusion-based zero-shot methods for inverse problems. This framework can address all the cases, including training-from-scratch, pre-trained models with and without reverse covariances. As recognized by the reviewer, it is useful to convert reverse covariances for many existing models and this novel idea has been long neglected. It is non-trivial to reveal the concept of using the optimality theorem of DDPM [8,9] to transform reverse variance predictions into posterior variance predictions, despite that the construction of Theorem 1 mostly follows [4]. We have also claimed this point as our contribution in the paper.

W2 & W3: Add Citation & Highlight difference between type I and II

We thank the reviewer for the constructive suggestion. In our revised version, the citation and discussion have been included in the last paragraph of Section 3.1, and the difference has been highlighted in the first paragraph of Section 3.2.

W4: Confusion of validation of posterior variances prediction

The computation of the ground-truth square error involves the following steps:

1. We sample a clean image $x_0$ from the training set.
2. We inject Gaussian noise to $x_0$ to generate a noisy image $x_t=x_0 + \sigma_t \epsilon, \epsilon\sim \mathcal{N}(0, I)$.
3. We employ the pretrained diffusion model to denoise $x_t$, resulting in $\hat{x}_0 = D_t(x_t) \approx \mathbb{E}[x_0|x_t]$.
4. We calculate the pixel-wise square errors as $e = (x_0 - \hat{x}_0)^2$.

We call $e$ the ground-truth square errors as it is the target (ground truth) that we want the posterior variances prediction $\hat{r}_t^2(x_t)$ to approximate. We compare $e$ with $\hat{r}_t^2(x_t)$ because posterior variances $r_t^{*2}(x_t)$ is the MMSE estimator of $e$ given $x_t$ (assuming $D_t(x_t) = \mathbb{E}[x_0|x_t]$), as suggested by equation 17, implying that a good posterior variances prediction $\hat{r}_t^2(x_t)$ should be a reliable predictor of $e$. Therefore, we may use the comparison between $e$ and $\hat{r}_t^2(x_t)$ as a sanity check to gauge the effectiveness of equation 22 in practice. As shown in Figure 2, $e$ is indeed accurately predicted by $\hat{r}_t^2(x_t)$ computed using equation 22 in regions with low noise levels. However, we encounter numerical issues using equation 22 in regions with high noise levels, and we attribute this to the fact that equation 22 approaches a 0/0 limit in these regions.

W5: Sampling using SDE for Type II

As noted in footnote 1, replacing the posterior mean with the conditional posterior mean (e.g., from equation 2 to equation 4 for ODE) to achieve conditional sampling is not limited to diffusion models formalized by ODEs. The extension to DDPM, DDIM, or SDE-formalized diffusion models can be similarly proved by leveraging conditional Tweedie’s formula (Lemma 2). Therefore, there is no discrepancy between theory and practice. We chose ODEs (from equations 2 to equation 4) due to the comparatively simpler mathematical form than SDEs and discrete diffusion models like DDPM or DDIM.

However, due to discretization in practice, the characteristics of samples based on ODEs and SDEs differ. When discretizing SDEs, a large step size (small number of steps) often causes non-convergence [3]. Hence, ODEs are more favored for designing fast samplers than SDEs [2,3]. However, SDEs excel at correcting errors made at each step [2], which ODEs lack. This may explain why Type II performs better using a stochastic sampler than a deterministic one, as it can correct errors caused not only by discretization but also by the approximation from the conditional posterior mean $\mathbb{E}_q[x_0|x_t,y]$ to the true one $\mathbb{E}[x_0|x_t,y]$.

To see example samples of Type II using a deterministic sampler, please refer to "comparison_between_deterministic_and_stochastic_sampler_for_type_II.png" in the Supplementary Materials, which is generated by running DiffPIR with $\lambda=10$ using deterministic and stochastic_samplers. To see more results of Type II using a deterministic sampler, please run

bash quick_start/rebuttal/eval_guidance_II_ode.sh

using our source code provided in the Supplementary Material.

Q1: Analysis on inference speed

The speed of inference of our method depends on whether Analytic or Convert is chosen.

Analytic: Our method does not introduce any additional computational cost for online inference compared to previous methods, since the only computational cost arises from the offline pre-calculation of $\hat{r}_t^2$ using equation (24). For the offline pre-calculation, similar to [3], a small number of Monte Carlo samples is adequate for obtaining an accurate $\hat{r}_t$ from a pre-trained unconditional model and training dataset. In the paper, we pre-compute $\hat{r}_t$ using a randomly selected $0.5\%$ of the training set and reuse it for all the experiments. This process takes a couple of hours on a single 1080Ti GPU. When pre-computed, $\hat{r}_t$ is deployed along with the pre-trained unconditional model. As a result, our method is the same as previous methods in computational cost for online inference.

Convert: The conjugate gradient (CG) method could introduce additional computational cost. However, since CG converges rapidly and is only employed during the last few sampling steps, its additional computational cost is negligible compared to the overall inference time. We have included the discussion in Appendix E.3 in our revised version.

The results of inference speed on 1080Ti GPU are shown in the following tables.

Inference time (sec) for Type I using 50 steps Heun's 2nd sampler

Inference time (sec) for Type II using 50 steps Heun's 2nd sampler

Q2: Concern to only using optimal variances at the last few sampling steps

The Gaussian distribution effectively approximates the posterior only in the infinitesimal limit of the noise level [1,2]. As the noise level increases, the posterior distribution can evolve into a non-Gaussian multimodal distribution [2]. We also observed this phenomenon in our initial experiments (refer to 4.2 Quantitative results, Paragraph 2). To address this, we heuristically employed our techniques only in the last few sampling steps, where the noise level is sufficiently low for the Gaussian approximation to be effective. To decide when to switch to the optimal variance, we empirically observe that switching when $\sigma_t < 0.2$ is a very robust choice. We use this heuristic to switch to the optimal variance across all inverse problems and sampler setups and observe consistent performance improvements over baselines.

Q3: Performance on colorization and denoising

We have performed additional experiments on colorization and demonstrate that our method has significant performance improvement over baselines like DPS, $\Pi$ GDM, and DiffPIR. Denoising is a special case of noisy inpainting when the elements of the mask are all one and we have already demonstrated the superiority of our method in inpainting. We believe that current experimental results on image inpainting, deblurring, super-resolution, and colorization are sufficient to validate the effectiveness of the proposed method. Below, we elaborate on the evaluations for colorization.

The colorization experiments are performed on the FFHQ dataset. We provide quantitative results in terms of PSNR, SSIM, LPIPS, and FID to compare the proposed method with Type I guidance (i.e., DPS and $\Pi$ GDM) and Type II guidance (e.g., DiffPIR). Here, PSNR refers to how close the reconstructed image is to the ground truth in the sense of Euclidean distance, and SSIM, LPIPS, and FID are subjective metrics that reflect the perceptual quality. As shown in the following tables, our method yields evident performance gains, especially in terms of subjective metrics, over baselines on the colorization task. The proposed method significantly outperforms Type I guidance (i.e., DPS and $\Pi$ GDM) in all the metrics. Regarding Type II guidance like DiffPIR, we yield evident gains in LPIPS, FID, and SSIM and comparable PSNR. Note that, considering that multiple solutions can be consistent with the measurement, PSNR could not be a desirable metric for the task of colorization. To further validate our method, we provide the qualitative results by visualizing the colorized images in "colorization_type_I.png" and "colorization_type_II_lambda=1.png" in the Supplementary Material. Figure "colorization_type_I.png" shows that our methods significantly outperform baselines in Type I. Figure "colorization_type_II_lambda=1.png" demonstrates that although DiffPIR is slightly better than ours in PSNR, its image contains more artifacts than ours.

Quantitative results (PSNR, SSIM, LPIPS, FID) on Colorization for Type I guidance:

Quantitative results (SSIM) on Colorization for Type II guidance under different $\lambda$:

Quantitative results (LPIPS) on Colorization for Type II guidance under different $\lambda$:

Quantitative results (FID) on Colorization for Type II guidance under different $\lambda$:

Quantitative results (PSNR) on Colorization for Type II guidance under different $\lambda$:

Reference

[1] Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International conference on machine learning. PMLR, 2015.

[2] Xiao, Zhisheng, Karsten Kreis, and Arash Vahdat. "Tackling the generative learning trilemma with denoising diffusion gans." arXiv preprint arXiv:2112.07804 (2021).

[3] Bao, F., et al. (2022). "Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models." arXiv preprint arXiv:2201.06503.

Q5: Moving results for "complete version of prior works" to the main paper

We thank the reviewer for the constructive suggestion. We put these results into Appendix due to page limitation. We will try to highlight these results in our final version. For the performance concerns, please refer to our reply for "Lack of performance improvement".

Q6: Why equation 22 is a 0/0 limit?

When noise levels are high (or when $t$ is large), the posterior variances cannot be zero (LHS of equation 22) due to the considerable uncertainty of the clean image $x_0$ given $x_t$. Therefore, if the numerator $\hat{v}_t^2(x_t) - \tilde{\beta}_t$ in equation 22 approaches zero, it necessitates that the denominator $(\frac{\sqrt{\bar{\alpha}\_{t-1}}\beta_t}{1-\bar{\alpha}_t})^2$ also approaches zero. This can be understood by directly examining $(\frac{\sqrt{\bar{\alpha}\_{t-1}}\beta_t}{1-\bar{\alpha}_t})^2$. In the case of DDPM, where $x_t=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon$, and for large $t$, $\sqrt{\bar{\alpha}_t}\rightarrow 0$, it follows that $(\frac{\sqrt{\bar{\alpha}\_{t-1}}\beta_t}{1-\bar{\alpha}_t})^2 \rightarrow 0$ for large $t$.

Reference

[1] Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International conference on machine learning. PMLR, 2015.

[2] Xiao, Zhisheng, Karsten Kreis, and Arash Vahdat. "Tackling the generative learning trilemma with denoising diffusion gans." arXiv preprint arXiv:2112.07804 (2021).

[3] Boys, Benjamin, et al. "Tweedie Moment Projected Diffusions For Inverse Problems." arXiv preprint arXiv:2310.06721 (2023).

[4] Song, Jiaming, et al. "Pseudoinverse-guided diffusion models for inverse problems." International Conference on Learning Representations. 2022.

[5] Zhu, Y., et al. (2023). "Denoising Diffusion Models for Plug-and-Play Image Restoration." arXiv preprint arXiv:2305.08995.

[6] Chung, Hyungjin, et al. "Diffusion posterior sampling for general noisy inverse problems." arXiv preprint arXiv:2209.14687 (2022).

[7] Wang, Yinhuai, Jiwen Yu, and Jian Zhang. "Zero-shot image restoration using denoising diffusion null-space model." arXiv preprint arXiv:2212.00490 (2022).

[8] Bao, F., et al. (2022). "Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models." arXiv preprint arXiv:2201.06503.

[9] Bao, Fan, et al. "Estimating the optimal covariance with imperfect mean in diffusion probabilistic models." arXiv preprint arXiv:2206.07309 (2022).

W3 & Q2 & Q3 & Q4 & Q5: Lack of performance improvement and missing metrics

We sincerely apologize for any confusion caused by incorrect experimental results. These incorrectnesses were the consequence of the numerical issue from the implementation of our closed-form solutions for debluring. We have presented the corrected results for debluring tasks in the revised version of our paper. To provide a more thorough evaluation, we have included additional image quality metrics, such as PSNR, FID, and SSIM, in Appendix E of the revised version. To facilitate the reproduction of our results, we have provided our source code and quick start bash scripts in (Supplementary Materials, src). We are grateful for your understanding and patience.

In Figures 3 and 4 in our revised version, we clearly find that our proposed method achieves better LPIPS performance than baselines under almost all hyperparameters. We have also provided the additional results on PSNR, SSIM and FID in Appendix E, which show the same robustness and the performance tendency compared to baselines.

We can also find that DiffPIR is inferior to our method under most values of $\lambda$ due to its handcrafted design. In equation (10) in the revised version, $\sigma_t$ is related to $\lambda$ and affects the performance with varying $\lambda$. Our method optimizes $\sigma_t$ for better performance. In the case of $\lambda=10$, DiffPIR coincidently finds a similar value to our estimation and yields a similar performance (see "diffpir_var_vs_analytic_var.png" in the Supplementary Material for details).

W4: Limited novelty

Our paper proposes a unified framework for diffusion-based zero-shot methods. As claimed in the paper, our novelty is three-fold.

First, we provide an insight of interpreting existing methods from the view of approximating conditional posterior mean for the reverse process and reveal that they are equivalent to making isotropic Gaussian approximation with different handcrafted isotropic posterior covariances. This finding enables us to transform the handcrafted design in existing methods (classified as Type I and Type II guidance) to a general solution to optimization via maximum likelihood estimation. It is especially non-trivial for Type II guidance in which posterior variances could not be optimized due to the implicit assumption of a Gaussian distribution.

Second, we develop three approaches to address all the cases for realizing diffusion-based zero-shot methods, including training-from-scratch, and pre-trained models with and without reverse covariances, respectively. Remarkably, both pre-trained models with and without reverse covariances are useful to existing methods as elaborated below.

Pretrained models with reverse covariance. In Section 4.2, we are the first to reveal the concept of utilizing the optimality theorem of DDPM [8,9] to transform reverse variance predictions into posterior variance predictions. This result allows one to leverage previous works’ pre-trained models that estimate reverse covariance to predict posterior variance for inverse problems and enhance performance without incurring additional costs. Actually, unconditional reverse covariance is ignored in many of the existing methods for inverse problems, despite it can be calculated.

Pretrained models without reverse covariance. In Section 4.3, optimizing $r_t^2$ is not trivial for pre-trained models without reverse covariances. In fact, existing methods [4,5,6,7] all focus on finding optimal variance but they resort to handcrafted design. Contrary to them, we theoretically develop the optimal solution to the problem and empirically achieve MMSE estimation via Monte Carlo samples.

Third, as highlighted in the revised version, we offer an efficient way of calculating guidance in noisy inverse problems. $\Pi GDM$ does not provide an efficient method for implementing guidance in noisy inverse problems, as it requires the computation of high-dimensional matrix inversion for the vector-Jacobian product (equation 91). As a result, $\Pi GDM$ did not evaluate the performance on the noisy deblurring and noisy super-resolution (they conduct the noisy inpainting experiments as the guidance calculation is trivial). In contrast, we propose closed-form solutions to Type I guidance in Appendix B.1 (for inpainting, deblurring, super-resolution, and colorization) that can be efficiently evaluated. In the cases that the closed-form solutions to guidance are complex to derive, we further present numerical solutions in Appendix C based on the conjugate gradient (CG) method for Type I and II guidance, thereby enhancing the practicality of our approach.

W1 & Q1: Justification of Gaussian approximation

As discussed in Section 3 in the paper, recent zero-shot methods, including DPS [4], $\Pi GDM$ [6], DiffPIR [5] and DDNM [7], all make Gaussian approximation with isotropic covariance $q_t(x_0|x_t)$ for the true unconditional posterior $p_t(x_0|x_t)$. They explicitly or implicitly leverage HANDCRAFTED design of isotropic posterior covariances and have achieved empirical success. Song et al. (2023) justify that the score of Gaussian approximation is close to that of the true posterior under the assumption that the denoiser behaves like a pseudo linear filter (see [4], Appendix A.6).

Our approximation $q_t(x_0|x_t)$ is more accurate than previous methods by optimizing posterior variance (rather than using handcrafted design) without incurring additional computational costs and has demonstrated significant empirical improvements in solving various inverse problems. It should be noted that we do not use standard normal distribution to approximate the conditional posterior distribution. Actually, when $p(y|x_0)$ (by equation 1) and $q_t(x_0|x_t)$ are both Gaussians, the approximated conditional posterior $q_t(x_0|x_t,y)\propto p(y|x_0)q_t(x_0|x_t)$ is also a Gaussian (but not standard).

Using complex distributions to model the posterior accurately may improve performance, but it will significantly increase computational costs. Gaussian approximation for the posterior is a reasonable choice for balancing computational cost and the accuracy of conditional posterior mean approximation. It is known that the Gaussian distribution is an effective approximation of the posterior in the infinitesimal limit of the noise level [1,2]. As the noise level increases, the posterior distribution can evolve into a non-Gaussian multimodal distribution [2]. We have also observed this phenomenon in our initial experiments (refer to 4.2 Quantitative results, Paragraph 2). To circumvent this problem, in this paper, we employ our techniques only in the last few sampling steps, where the noise level is sufficiently low for the Gaussian approximation to be effective.

W2: Justification of diagonal covariance constraint

In fact, there are currently no tractable methods to implement a full covariance matrix-based approach for high-dimensional signals like images, despite that using a more expressive posterior covariance could potentially enhance performance. We identify two main challenges: (1) the computational cost of predicting the full posterior covariance and (2) the computational cost of implementing guidance based on the full posterior covariance. The concurrent work [3] (also submitted to ICLR 2024) proposed a method to optimize the posterior covariance by computing the Hessian matrix of $\log p_t(x_t)$ (i.e., the derivative of the score function) to obtain a full covariance prediction from a pre-trained unconditional diffusion model. However, this method is only evaluated on low-dimensional synthesis data but still uses diagonal covariance for high-dimensional signals like images (see Appendix E.1 in [3]). Taking RGB images with a resolution of 256 $\times$ 256 as an example, it is cost-ineffective as it requires computing the Hessian matrix and performing $d$ (dimension of signals) times backpropagation (in our case, 256 $\times$ 256 $\times$ 3 times backpropagation). This process is significantly slower than our approach since we do not introduce additional computational costs to compute the posterior covariance. Since our paper focuses on real-world inverse problems, choosing a diagonal covariance is reasonable for balancing computational costs and performance.

In addition, we focus on zero-shot solutions to inverse problems using available checkpoints for unconditional diffusion models since it is costly to train them from scratch. As mentioned in the response to W1 & Q1, existing works [4,5,6,7] make the Gaussian approximation with diagonal covariance for the true unconditional posterior. Our methods are built upon these milestone works, of which the improvement comes from carefully designed training-free methods for obtaining optimal diagonal covariances.

Thank you for your thoughtful response! Regarding the diagonal covariance, we have to emphasize that our paper aims to propose a unified framework for inverse problems and also considers to accommodate to existing diffusion-based zero-shot methods. To this end, it is hard to directly leverage widely adopted checkpoint of unconditional diffusion models used by existing methods without using diagonal covariance. The use of the diagonal constraint has another significant advantage: it enables our method to be seamlessly integrated with existing methods. This plug-and-play nature of our approach makes it the resonable step after recent works.