Paper: BayesDiff: Estimating Pixel-wise Uncertainty in Diffusion via Bayesian Inference (ICLR-2024)

Reviewer comment (Reproducibility):
Some details that might be useful for potential readers are glossed over. E.g. how is $\mathrm{diag}(\gamma_\theta^2(x_t, t))$ computed? What approximation of the Hessian is used? etc.

Author response:
Thank you for pointing out that we lack implementation details regarding the Laplace approximation, we have added this to our paper. 

We adopt the most lightweight diagonal factorization and ignore off-diagonal elements for Hessian approximation [1]. Regarding the covariance $\gamma_\theta^2(x, t) \in R^{h d^2 * h d^2}$ of the posterior predictive distribution, as you point out, to get the accurate covariance of outputs is indeed time-consuming (17s when the image size is 64*64). Luckily, in our settings, we do not need to estimate the entire covariance matrix; we only need to estimate the diagonal. So we adopt the Monte Carlo method to approximate the accurate variance of outputs $diag(\gamma_\theta^2(x, t))$ directly by the variance of samples { $ f_{\theta_i}(x,t) $ }, $\theta_i  \sim p(\theta|\mathcal{D}) =\mathcal{N}(\theta ; \theta_{\text{MAP}}$, $\boldsymbol{\Sigma})$. This avoids storing large Jacobian matrix, resulting in faster computation speed, while still maintaining a reasonable level of accuracy. 

The number of samples is chosen as 100 in practice. That is, by performing 99 additional inferences of the last layer, we can accurately estimate the posterior distribution without significant computational or time constraints (the computational time decreases from 17s to 0.1s). Furthermore, when comparing the results obtained from the accurate variance and the approximated variance, the average approximation error per pixel is 10%. It is minor and can be negligible because the magnitude of the variance relative to the mean is significantly small.

Reviewer comment (Experiments):
Sec. 4.2 and 4.3 are a bit handwavy---it would be much better if the authors could make them more quantitative, like Sec. 4.1. Handpicked examples are not useful to instill confidence about the benefits of BayesDiff.

Author response:
For flawed samples exhibiting artifacts and misalignments selected by humans, we can provide possibilities for generating potential refined samples with BayesDiff. However, it is important to note that quantitatively measuring such value on a large scale poses challenges as there are no benchmarks or quantitative metrics to assess the presence of high-granularity artifacts. Therefore, the effectiveness of our method can only be demonstrated through the case study based on human evaluations. To provide a more intuitive demonstration of the significance of our resampling method, we have included additional prompts and corresponding images containing artifacts annotated by humans in Appendix A.7. By observing the resampling results, we find that the success rate of resampling flawed samples into desirable samples is approximately at least 60%. We will involve more human experts in the evaluation process in the final version.

Reviewer comment (Presentation):
Fig. 2 is quite hard to understand. The caption is not descriptive at all and the colors are hard to see in print.

Author response:
Thanks for your useful suggestions! We have made the modifications to the figures based on your suggestions. If you have any further advice, please feel free to provide them.

Reviewer comment (Reproducibility):
The last-layer Laplace approximation can still be very expensive for networks with high output dimensionality, e.g. in text or image generation. For example, if the image is $d \times d$, then the network has an output dim of $d^2$. Assuming the last-layer feature dim of $h$, this means the last-layer weight matrix is $d^2 \times h$ and so the Hessian is $hd^2 \times hd^2$. Then, to get the variance over outputs $\gamma^2(x, t)$, you need to multiply the Hessian with the last-layer Jacobian, which itself is large---$d^2 \times hd^2$. Can the authors elaborate on how BayesDiff overcomes this issue in practice?

Author response:
We adopt the most lightweight diagonal factorization and ignore off-diagonal elements for Hessian approximation [1]. Regarding the covariance $\gamma_\theta^2(x, t) \in R^{h d^2 * h d^2}$ of the posterior predictive distribution, as you point out, to get the accurate covariance of outputs is indeed time-consuming (17s when the image size is 64*64). Luckily, in our settings, we do not need to estimate the entire covariance matrix; we only need to estimate the diagonal. So we adopt the Monte Carlo method to approximate the accurate variance of outputs $diag(\gamma_\theta^2(x, t))$ directly by the variance of samples { $ f_{\theta_i}(x,t) $ }, $\theta_i  \sim p(\theta|\mathcal{D}) =\mathcal{N}(\theta ; \theta_{\text{MAP}}$, $\boldsymbol{\Sigma})$. This avoids storing large Jacobian matrix, resulting in faster computation speed, while still maintaining a reasonable level of accuracy. 

The number of samples is chosen as 100 in practice. That is, by performing 99 additional inferences of the last layer, we can accurately estimate the posterior distribution without significant computational or time constraints (the computational time decreases from 17s to 0.1s). Furthermore, when comparing the results obtained from the accurate variance and the approximated variance, the average approximation error per pixel is 10%. It is minor and can be negligible because the magnitude of the variance relative to the mean is significantly small.

Reviewer comment (Experiments):
There are no baselines. I appreciate that there may not have been any prior work in this direction (although I would imagine that there would be some non-probabilistic filtering techniques. Perhaps from the literature on GANs?), however given that the iterative sampling process involves a Gaussian at every step, if I am understanding things correctly I would think that a deterministic diffusion model would also give us pixelwise variances that could be used as a baseline.

Author response:
W2 & Q3: Could the sampling variances be used to create a baseline with a deterministic diffusion model?

Thank you for your appreciation of our work as the first one in this direction. Regarding the baseline, as far as I know, most filtering techniques do not apply to our task because they focus on detecting deepfake images or filtering out low-quality images from training datasets. [1, 2] There is currently no widely accepted sample-wise metric that can be used to detect low-quality images generated by large-scale generative models, so we adopt the randomly filtered data as a comparison in our quantitative result. On the other hand, **the sampling variances with a deterministic diffusion model cannot be used to create a baseline.** This is because, for an ODE solver-based reverse generative process, the sampling process does not follow a Gaussian distribution but is deterministic, lacking sample variance. As for an SDE solver-based reverse generative process, all samples' sampling variances are equal, making it unable to distinguish between samples.

Reviewer comment (Experiments):
Alternatively, it might have been interesting to experiment with different covariance structures for the Laplace approximation to see if those make a difference.

Author response:
W3: Other approximation methods for Hessian.
Indeed, our method supports the use of other approximation methods for Hessian, although that is not the primary focus of our contribution. The experimental section primarily emphasizes the understanding and utilization of pixel-wise variance. However, your point about exploring alternative approximation methods is intriguing, and we will explore it as part of our future work.

Reviewer comment (Writing):
I did not find the background section on diffusion models (2.1) particularly helpful as it relies on a lot of terminology on SDEs. The opening paragraph is good, perhaps something similar that briefly summarizes things from an algorithmic perspective (what kind of network are we typically training to predict what and w.r.t. what objective, what is being sampled, ...).

Author response:
W4: Suggestions on Section 2.1.

We have incorporated the common architectural selection of the neural network and specific sampling methods in Section 2.1 based on your suggestions. If you have any further advice, please feel free to share.

Reviewer comment (Evaluation):
I would like to see error bars for (some of) the quantitative results.

Author response:
W1 & Q1: Error bar of results in Table 1.

Table 1 showcases a comparison of metrics between randomly selected 50,000 images and the remaining 50,000 images after filtering out those with the highest uncertainty from 60,000 generated images. In response to your feedback, we conducted multiple random selections (10 different random seeds) of 50,000 images and provided error bars for the FID metric of the randomly selected images. **Results added in Table 1 show that the improvement in the overall distribution metrics by filtering out low-quality images based on the sample-wise uncertainty metric is significant compared to the variance of the measurement.**

Moreover,  we would like to reiterate the significance of our method. There are still low-quality generations in large-scale models currently, and it is challenging to identify such individual images using traditional metrics (FID, IS, LIPSIS) that measure the quality of a large set of images distribution-wisely. The core value of our approach lies in introducing a sample-wise metric based on pixel-wise uncertainty. This sample-wise metric enables us to effectively filter out cluttered low-quality images.

Reviewer comment (Theory):
Is there a strict need to use a Laplace approximation for inference? The likelihood is a Gaussian, so if you are only estimating uncertainty over the final layer weights, shouldn't the posterior be Gaussian as well? Or is this not the case due to the iterative sampling process?

Author response:
Q2：Is there a strict need to use a Laplace approximation for inference?

Sorry that we did not explain this perfectly. You are correct that the posterior distribution of the parameters is indeed a Gaussian distribution in our settings. However, conventional methods require a significant amount of data to fit the Gaussian posterior distribution, which can be computationally expensive. **The reason we opt for the Laplace approximation is to directly leverage the parameters of the pre-trained model.** Under this approximation, the parameters of the pre-trained model can serve as the Maximum A-Posteriori (MAP) estimate of the Gaussian posterior distribution, i.e. the mean of the Gaussian distribution is approximated by the pre-trained weights. It is training-free and facilitates practical applications.

Reviewer comment (Experiments):
Could the sampling variances be used to create a baseline with a deterministic diffusion model?

Author response:
W2 & Q3: Could the sampling variances be used to create a baseline with a deterministic diffusion model?

Thank you for your appreciation of our work as the first one in this direction. Regarding the baseline, as far as I know, most filtering techniques do not apply to our task because they focus on detecting deepfake images or filtering out low-quality images from training datasets. [1, 2] There is currently no widely accepted sample-wise metric that can be used to detect low-quality images generated by large-scale generative models, so we adopt the randomly filtered data as a comparison in our quantitative result. On the other hand, **the sampling variances with a deterministic diffusion model cannot be used to create a baseline.** This is because, for an ODE solver-based reverse generative process, the sampling process does not follow a Gaussian distribution but is deterministic, lacking sample variance. As for an SDE solver-based reverse generative process, all samples' sampling variances are equal, making it unable to distinguish between samples.

Reviewer comment (Reproducibility):
In Section 4.2, it is shown that pixel-wise uncertainty can be used to correct bad portions or artefacts in the original images. It is unclear whether the bounding boxes for the artefacts were determined automatically based on pixel uncertainty or identified manually.

Author response:
The flawed samples are identified by humans and the bounding boxes are manually annotated. This is reasonable and feasible due to a typical pattern of chat logs in the Discord server of Stable Diffusion: users provide a prompt, and the channel bot generates multiple images based on the prompt. Subsequently, individuals manually select one image for further processing [1]. We have included additional prompts and corresponding images containing artifacts annotated by humans in Appendix A.7. By observing the resampling results, we find that the success rate of resampling flawed samples into desirable samples is approximately at least 60%. We will involve more human experts in the evaluation process in the final version.

Reviewer comment (Evaluation):
In the same Section 4.2 experiment, it is not specified how often the refined samples, obtained using rejection sampling on the pixel uncertainty score, are better than the original images.

Author response:
The flawed samples are identified by humans and the bounding boxes are manually annotated. This is reasonable and feasible due to a typical pattern of chat logs in the Discord server of Stable Diffusion: users provide a prompt, and the channel bot generates multiple images based on the prompt. Subsequently, individuals manually select one image for further processing [1]. We have included additional prompts and corresponding images containing artifacts annotated by humans in Appendix A.7. By observing the resampling results, we find that the success rate of resampling flawed samples into desirable samples is approximately at least 60%. We will involve more human experts in the evaluation process in the final version.

Reviewer comment (Evaluation):
It is unclear whether the examples shown in Figure 8 are cherry-picked or representative of the pixel-wise uncertainty rejection sampling method.

Author response:
The flawed samples are identified by humans and the bounding boxes are manually annotated. This is reasonable and feasible due to a typical pattern of chat logs in the Discord server of Stable Diffusion: users provide a prompt, and the channel bot generates multiple images based on the prompt. Subsequently, individuals manually select one image for further processing [1]. We have included additional prompts and corresponding images containing artifacts annotated by humans in Appendix A.7. By observing the resampling results, we find that the success rate of resampling flawed samples into desirable samples is approximately at least 60%. We will involve more human experts in the evaluation process in the final version.

Reviewer comment (Evaluation):
In Figure 2, although it tries to demonstrate that the "pixel-variance sum" statistic, even when skipping pixels, can separate high uncertainty images from low uncertainty ones, at skipping=3 and skipping=4 the two clusters appear to be mixed considerably.

Author response:
For samples with uncertainty around the mean, it is indeed challenging to differentiate them steadily. However, as mentioned in the article, the top and bottom 10% of samples based on uncertainty can be correctly classified. We compare the top and bottom 10% images based on the highest and lowest uncertainty in the new version of Figure 2, which fulfills the requirement of accurately excluding samples with the highest uncertainty in practical applications. Furthermore, we have directly visualized the sorting of samples with different skipping sizes in Appendix A.6. These visualizations demonstrate that regardless of the skipping size, there is a consistent trend from cluttered to clean images. This finding confirms that skip algorithms can be employed in practical applications.

Reviewer comment (Evaluation):
It is not clear why the mean of the scores decreases for skipping=5 and skipping=6 when it was increasing monotonically from skipping=1 to skipping=4.

Author response:
The uncertainty dynamics in diffusion models are complex, as the relationship between Cov$(x_t, \epsilon_t)$ and Var$(\epsilon_t)$ with respect to $t$ is not necessarily a simple monotonic relationship. Therefore, we cannot consider the length of the skipping interval as a simple monotonic relationship with the sum of pixel-wise variance. The exact relationship between them is not easily determined through theoretical analysis. However, this does not hinder our application of skipping algorithms.

Reviewer comment (Experiments):
Why not use the model likelihood to rule out low-fidelity images? The likelihood of diffusion models is tractable as done in [Song et al.](https://openreview.net/pdf/ef0eadbe07115b0853e964f17aa09d811cd490f1.pdf)

Author response:
Thank you for providing this insight, but as far as we know, a model can have poor log-likelihood and produce great samples, or have great log-likelihood and produce poor samples. **Log-likelihood and visual appearance of samples are therefore largely independent.** You can refer to this article [[1](https://arxiv.org/abs/1511.01844)] for details.

Based on [the official code of the paper you provided](https://github.com/yang-song/score_sde_pytorch), we computed the per-sample negative log-likelihood (NLL) w.r.t the ODE reverse process, i.e., $- log (p_\theta^{ODE}(x_0))$ for two datasets. Detailed experimental results can be found in Appendix A.5. **Experimental results demonstrate that model likelihood is indeed unsuitable for filtering out low-quality images. In comparison, the sample-wise metric based on uncertainty enables us to effectively eliminate low-quality images with cluttered backgrounds.** Due to time constraints, we are currently in the process of measuring the quantitative metric for comparison, which will be added to the article as soon as we complete it.