Paper: Training Diffusion Classifiers with Denoising Assistance (ICLR-2024)

Reviewer comment (Presentation):
The presentation of the paper could be improved for better readability. Specifically, the formatting of equations (5) and (6) spanning across three rows appears to be cluttered and may benefit from a more concise representation.

Author response:
1) Presentation: Thank you for the feedback. We will fix the formatting in Eqs. (5,6), and we gladly welcome any other suggestions.

Reviewer comment (Novelty):
The core contribution of employing the denoised sample as input to a classifier may come across as straightforward. The paper could benefit from a clearer articulation of the motivation behind this choice and the specific problems it aims to address. While the one-step denoised sample is utilized, the rationale behind not exploring [x, one-step denoised sample, two-step denoised sample] as inputs could use further clarification. Although the authors provide some explanations, a more robust justification could enhance the perceived significance of the work.

Author response:
2) Motivation for using single-step denoised examples. We use denoised example $\hat{\bf x_t} = {\bf x_t}+\sigma^{2}(t) s_\theta({\bf x_t},t)$
as an additional input to the classifier network because (a) learning to classify denoised image is potentially easier than learning to classify noisy examples and (b) we have access to a pretrained score network $s_\theta({\bf x_t},t)$ that can help us estimate the additive gaussian noise contained in ${\bf x}_t$. Beyond just simplifying the learning task, we observe that this simple change also helps address the following deficiencies of vanilla classifier guidance: 1) Generalisation in fully-supervised and semi-supervised settings, 2) Perceptual-alignment of Classifier-gradients, and 3) Image Generation Performance. We support these findings both theoretically and empirically. Additionally, we also show significant advantages of using denoised-images in semi-supervised training as opposed to simple gaussian-noise augmented images. We will clarify this motivation in the revised paper.

Reviewer comment (Theory):
The discussion on the semi-supervised learning framework introduces an idea of selecting pseudo-label data based on confidence thresholds during the diffusion process. However, the motivation behind this choice could be better elucidated. It may be worth exploring or explaining why traditional uncertainty measures in the raw data space were deemed insufficient, and how the proposed method addresses any identified limitations.

Author response:
3) Uncertainty measures for pseudo-label estimation. To train the diffusion classifier on partially-labelled data, we need to estimate the pseudo-labels of the unlabeled examples in each batch. As described in the paper, we follow FixMatch and use confidence-thresholding, a traditional measure of uncertainty, to estimate the pseudo-labels. In order to estimate the pseudo-labels, we use two types of data samples from forward diffusion to estimate the pseudo-labels — samples diffused to 1) $\tau=0.01$ (i.e., with negligible gaussian noise) and 2) uniformly sampled time $s$. Since the samples diffused to time $\tau=0.01$ have negligible gaussian noise, that is similar to estimating pseudo-labels from raw data using traditional uncertainty measures. In addition, we also use pseudo-labels from confident predictions on samples diffused to random uniform time $s$ as they could be more informative than raw data in the following cases:
   * Additive gaussian noise in samples diffused to time $s$ could potentially help the noisy-classifier/DA-classifier better estimate the uncertainty in labels as compared to samples containing negligible gaussian noise. 
   * Denoised images obtained from unlabeled examples with additive gaussian noise could potentially help classification, even when the DA-classifier is not able to confidently classify the original raw image. 

   At each training step, we estimate pseudo-labels from confident predictions on samples diffused to both $\tau$ and $s$ and use them for computing the cross-entropy loss. We do not resolve disagreements between these two pseudo-labels because we expect that the confident predictions on diffused samples should become consistent over training. We thank the reviewer for this comment and we will clarify this discussion in the revised text.

Reviewer comment (Novelty):
What is the motivation to develop DA classifier? What is the advantages of DA classifier compared to existing classifiers including conventional deep models (VGG, ResNet, etc.) or large models (CLIP, etc.)?

Author response:
> What is the motivation to develop DA classifier?

We systematically expose several deficiencies  of vanilla classifier guidance,
1) Generalization in fully-supervised and semi-supervised settings, 
2) Perceptual-alignment of Classifier-gradients,
3) Image Generation Performance 

and we demonstrate how simply using denoised images as additional input can address all these deficiencies. We call these denoising-assisted (DA) classifiers.


> What is the advantages of DA classifier compared to existed classifiers including conventional deep models (VGG, ResNet, etc.) or large models (CLIP, etc.)?

We believe that this question is asking for some overall context, so here is an overview: For classifier-guided diffusion, we train classifiers on images sampled from the forward diffusion, which generally contain additive gaussian noise. As described in the paper, a DA-Classifier refers to a classifier that takes both noisy and denoised images as simultaneous inputs, in contrast to a noisy-classifier that only receives noisy images as inputs. Noisy Classifier and DA-Classifier can be constructed with any backbone neural architecture: as described in the paper, we use WideResNet-28-2 for all CIFAR10, MNIST and SVHN experiments (following Song et al., 2021) while we use the downsampling-half of the UNET for ImageNet experiments (following Dhariwal and Nichol (2021b) ). The vanilla diffusion classifier and the corresponding denoising-assisted diffusion classifier have the same network architecture after the first layer and differ only in terms of the first convolution module: for example, if the first input-convolution of vanilla diffusion classifier takes an input volume of $(N, C_i, H_i, W_i)$ and gives an output volume of $(N, C_o, H_o, W_o)$, the DA-Classifier takes an input volume of $(N, 2C_i, H_i, W_i)$ and gives the same size output $(N, C_o, H_o, W_o)$.

Reviewer comment (Experiments):
What the authors do is to employ the denoised images as the additional input. To this end, the authors should provide more detail about how to get the denoised images and analysing the influence of different types of denoised images.

Author response:
> “the authors should provide more detail about how to get the denoised images and analysing the influence of different types of denoised images”

As described below Eq. 6, we obtain denoised images $\hat{\bf x}$ using the pretrained score-network $s_\theta$ as $\hat{\bf x} = {\bf x}+\sigma^{2}(t) s_\theta({\bf x},t)$. We would be glad to provide additional details about this in an appendix, but we do not understand exactly what additional details would be helpful for this reader. We are also not sure what the reviewer means by “different types of denoised images”; any clarification on this would also be welcome. As explained above, we are denoising images that have been corrupted with additive gaussian noise.

Reviewer comment (Theory):
Why Theorem 1 explain the improvements of DA classifier?

Author response:
> “Why Theorem 1 explain the improvements of DA classifier?”

As written in Section 3, Theorem-1 “explains the improved perceptual alignment in Fig (3) and Fig (2), since multiplying a vector by $\text{Cov}[{\bf \bar{\bf x}}_t|{\bf x}]$  *stretches* the vector along the principal directions of the conditional distribution $p({\bf \bar{\bf x}}_t|{\bf x})$.” Intuitively, since the conditional distribution $p({\bf \bar{\bf x}}_t|{\bf x})$ corresponds to the distribution of candidate denoised images, the principal directions of variation are perceptually aligned and hence stretching the gradient along these directions will yield perceptually aligned gradients. We attribute the quantitative and qualitative improvements on the image generation task to 1) perceptually aligned classifier gradients and 2) improved generalisation to unseen examples.

Reviewer comment (Writing):
"Cifar10" should be changed to "CIFAR-10".

Author response:
We thank the reviewer for their comments. We will incorporate the minor suggestions & typo corrections in our revision.

Reviewer comment (Writing):
"Imagenet" should be changed to "ImageNet".

Author response:
We thank the reviewer for their comments. We will incorporate the minor suggestions & typo corrections in our revision.

Reviewer comment (Writing):
On Page 3, Sec. 2.2: $y\in [1,C]$ should be changed to $y\in =\{1, ..., C\}$.

Author response:
We thank the reviewer for their comments. We will incorporate the minor suggestions & typo corrections in our revision.

Reviewer comment (Writing):
On Page 3, Sec. 3: "propose to use as input both ..." should be changed to "propose to use both ... as input".

Author response:
We thank the reviewer for their comments. We will incorporate the minor suggestions & typo corrections in our revision.

Reviewer comment (Evaluation):
The numerical results are not good. In Table 2, there are no significant differences between the two methods, especially P and R on the two datasets and all metrics on ImageNet. The visualization in Figures 2, 3, and 4 is impressive.

Author response:
1) Numerical results and consistency between metrics in Table 2 and Table 4: Thank you for asking about this. The Precision, Recall, Density and Coverage metrics compare between the manifold of generated distribution and manifold of the source distribution in terms of nearest-neighbours and can be computed conditionally (i.e., classwise) or unconditionally [1]. Table 2 shows the Precision and Recall metrics computed unconditionally while Table 4 shows the Precision, Recall, Density and Coverage metrics computed classwise. For CIFAR10, we do indeed observe significant improvements in terms of classwise Precision, Recall, Density and Coverage metrics (see Table 4). We also observe that using DA-Classifier instead of vanilla classifier improves FID/IS on both ImageNet and CIFAR-10 at similar precision and recall values – this suggests that both the classifiers have similar manifolds in the nearest-neighbours sense but the generated images with DA-Classifier are more realistic than vanilla-classifiers. In Fig. 5, we qualitatively observe that DA-Classifiers improve over Vanilla-classifiers in terms of overall coherence and since these improvements are generally subtle, we do observe slight but meaningful improvements of FID/IS over the baseline. We will incorporate this discussion into the revised text. Overall, we believe that the numbers are indeed consistent with strong performance, but if there is any remaining result or comparison that the reviewer sees as not being good, then if we failed to provide a clear explanation or interpretation of it, we would be grateful to have the opportunity to correct that. (Thank you!)

   Also, thank you for appreciating the strength of our visualizations! While visualizations can be hard to quantify, we are very excited by these results and agree that they demonstrate a significant contribution of this work!

Reviewer comment (Evaluation):
Why the results of Tables 2 and 4 are not consistent.

Author response:
1) Numerical results and consistency between metrics in Table 2 and Table 4: Thank you for asking about this. The Precision, Recall, Density and Coverage metrics compare between the manifold of generated distribution and manifold of the source distribution in terms of nearest-neighbours and can be computed conditionally (i.e., classwise) or unconditionally [1]. Table 2 shows the Precision and Recall metrics computed unconditionally while Table 4 shows the Precision, Recall, Density and Coverage metrics computed classwise. For CIFAR10, we do indeed observe significant improvements in terms of classwise Precision, Recall, Density and Coverage metrics (see Table 4). We also observe that using DA-Classifier instead of vanilla classifier improves FID/IS on both ImageNet and CIFAR-10 at similar precision and recall values – this suggests that both the classifiers have similar manifolds in the nearest-neighbours sense but the generated images with DA-Classifier are more realistic than vanilla-classifiers. In Fig. 5, we qualitatively observe that DA-Classifiers improve over Vanilla-classifiers in terms of overall coherence and since these improvements are generally subtle, we do observe slight but meaningful improvements of FID/IS over the baseline. We will incorporate this discussion into the revised text. Overall, we believe that the numbers are indeed consistent with strong performance, but if there is any remaining result or comparison that the reviewer sees as not being good, then if we failed to provide a clear explanation or interpretation of it, we would be grateful to have the opportunity to correct that. (Thank you!)

   Also, thank you for appreciating the strength of our visualizations! While visualizations can be hard to quantify, we are very excited by these results and agree that they demonstrate a significant contribution of this work!

Reviewer comment (Experiments):
The paper lacks sufficient baseline comparisons. To strengthen their argument, the authors should compare their method with other relevant classifier guidance methods such as DLSM, ED, and Robust-Guidance. While they mention that these methods are complementary, it would be beneficial to present their performance results, similar to how standard semi-supervised discriminative models are presented in Table 3. In addition, the authors should extend the application of DA classifiers to these methods and report the results to demonstrate the orthogonal nature of their work.

Author response:
Even though DLSM-loss and DA-Classifier represent orthogonal improvements over vanilla classifier, as recognized by the reviewer, we directly compare between the two methods on CIFAR10 (trained on VE-SDE) based on the reviewer’s request:
||||||||
|---|:---:|:---:|:---:|:---:|:---:|:---:|
|  | FID  | IS  | CW Precision  | CW Recall  | CW Density  | CW Coverage  |
| **DLSM** | 2.25 | 9.90 | 0.56 | 0.61 | 0.76 | 0.71 |
| **Ours** | 2.33 | 9.88 | 0.63 | 0.64 | 0.92 | 0.77 |

While the FID and IS scores are comparable (especially in relation to the typical FID/IS scores observed in the literature), we note that our class-wise Precision, Recall, Density and Coverage metrics are either comparable or demonstrate a significant improvement. Although we are very interested to analyse the fusion of each methodology, this is out of scope of this current paper: furthermore, a careful analysis of fusion in terms of generalisation, classifier-gradients and image-generation performance also requires specific experiments for finding optimal hyperparameters for each fusion -- for example, the hyper parameters controlling the adversarial attack in Robust-Guidance, the softmax temperature in Entropy-driven Training/Sampling, and so on. We are excited to present these kinds of studies in a thorough and systematic future work. But in the present paper, our focus is, rather, to comprehensively analyse the benefits of using denoised examples as an additional input to the classifier both empirically and theoretically while only using the cross-entropy loss objective. The suggested baseline methods such as DLSM, Robust-Guidance and Entropy-Driven Training/Sampling propose novel loss-objectives additional to the usual cross-entropy loss and are hence orthogonal to our method in which we propose using denoised images as additional inputs to the classifier; furthermore, while the baselines are focused on improving classifier-guidance in fully-supervised settings, we also consider semi-supervised settings. In our initial exploration, we tested to see if the DLSM loss objective can help in semi-supervised settings but did not observe any improvements over just using the cross-entropy loss. We agree that these comparisons strengthen the paper and we will add these results in our revision.

Reviewer comment (Evaluation):
There is no FID and IS results for generated images in semi-supervised settings. They only said that "FID and IS metrics are similar to the results described in Table 2". Since the main task of this paper could be seen as conditional image generation, these experimental results are crucial evidence to assess the effectiveness of their method.

Author response:
We thank the reviewer for mentioning this and giving us the opportunity to provide additional numerical evaluations. As mentioned in the paper, the FID and IS scores of images generated with semi-supervised classifiers are very similar to the FID/IS scores of images generated with fully-supervised classifiers: specifically, the FID/IS scores of semi-supervised vanilla classifier and semi-supervised DA-classifier are 2.82/9.61 and 2.35/9.86 respectively. We will include these in the revised version of the paper. The class-wise P/R/D/C values are, however, different for semi-supervised classifiers as compared to fully-supervised classifiers and we do include these in Table 5 (in Appendix B) – where, we can see that the DA-Classifier outperforms noisy-classifier in semi-supervised settings as well.

Reviewer comment (Experiments):
The authors mentioned that semi-supervised training of a classifier-free model is not easy to implement, citing the original paper's recommendation for the supervised settings. It is important to empirically verify this claim, as there is a lack of evidence on the feasibility and effectiveness of classifier-free models in partial label settings.

Author response:
We provide a two-part rebuttal to this comment: first, we argue that comparing between classifier-guidance and classifier-free guidance in fully-supervised settings is not straightforward to begin with. Then, we elaborate upon the challenges in implementation of classifier-free guidance in semi-supervised settings that further complicates comparison between the two methods in semi-supervised settings.

(a) classifier-free guidance and classifier-guidance represent two main categories of guidance methods for diffusion models; classifier-guidance is modular, theoretically motivated and allows for flexible conditioning without requiring to retrain the unconditional score model while classifier-free guidance is a heuristic method that bakes the conditioning into the score-model and requires careful hyperparameter tuning [1]. Comparing between classifier-guidance and classifier-free guidance properly is challenging: for example, previous studies such as GLIDE [2] do this very well, but they rely upon extensive (and expensive) human evaluation of generated samples to show that classifier-guidance often qualitatively underperforms classifier-free guidance; this result is in contrast to the quantitative metrics they computed such as CLIP-scores that usually favoured classifier-guided samples over classifier-free samples. We therefore argue that establishing fair comparison criteria between the two methods would require dedicated research efforts that would be beyond the scope of this paper. Here, we focus on addressing the following deficiencies of vanilla classifier guidance by using denoised images as additional input: 1) Generalization in fully-supervised and semi-supervised settings, 2) Perceptual-alignment of Classifier-gradients, 3) Image Generation Performance. 

(b) Training a classifier-free diffusion model with partially labelled dataset is harder because the training relies heavily on labels, as described in our text: “…Ho and Salimans (2022) recommend training with class-label conditioning for 80-90% of each batch and null-token conditioning (i.e., no class label) for the remaining examples whereas we assume that the class-labels are available for less than 10% of the complete training dataset.” If we directly follow the suggestion in that paper, it is easy to see that we will overfit to the labeled examples as we recycle examples from the labeled pool more often than the examples in the unlabeled pool: 
       
>In the example of CIFAR10, we assume that 4K examples are labelled out of a total 50k samples. If we consider a batch-size of 512 and select the percentage of labeled examples in a batch to be 80%, we will need 410 labeled examples and 112 unlabeled examples for each gradient step. Overall, this implies that 1 epoch over the unlabeled examples corresponds to 41 epochs over the labeled examples (since, 10 gradient steps correspond to 1 epoch over the labeled data). 

Thus, avoiding overfitting to the labelled pool would require separate research efforts–again outside the scope of this paper–-to identify reliable pseudo-labelling strategies that enable a feasible way to follow the recommendation in Ho and Salismans (2022) — in order to get good generative performance.

[1] Bahjat Kawar, Roy Ganz, Michael Elad. Enhancing Diffusion-Based Image Synthesis with Robust Classifier Guidance. TMLR, 2023.

[2] Nichol, A., Dhariwal, P., Ramesh, A., Shyam, P., Mishkin, P., McGrew, B., Sutskever, I. and Chen, M. GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models. ICML, 2022.

Reviewer comment (Reproducibility):
Since the primary objective of this paper is image generation, it would be beneficial to provide a detailed procedure or algorithm for sample generation using the proposed method. Actually, the utilization of DA classifiers in sample generation is not straightforward, given the difference in gradient computation between $\nabla_x \log p(y|x,t)$ and $\frac{d \log p(y|x,\hat{x},t)}{dx}$. A discussion of the implications of this difference would help clarify the approach.

Author response:
Thank you for raising this point: this gives us an opportunity to highlight the fact that the implementation of sampling procedure with DA-Classifiers is indeed as straightforward as with vanilla-classifiers: the gradient of the log-probability estimated by the classifier is used as additional guidance in both cases. In our implementation, we simply concatenate the noisy and denoised example and use this as input to the classifier. In terms of the model architecture, we simply double the ```in_channels``` of the input convolution for DA-Classifiers and the remaining architecture is left unchanged: in fact, for the Imagenet experiments in the paper, we expanded the in-channels of the first convolution layer and continued training the pretrained model for 50k steps. We will incorporate this into the paper.

Reviewer comment (Experiments):
It would be beneficial to include an ablation study to investigate the key factors contributing to the improvement of the DA classifier. Two specific experiments might be considered: (a) Changing $\hat{x}$ to $s_\theta (x,t)$ in the DA classifier inputs: this would help to identify the essential elements, such as denoised examples or information of data scores; (b) Use only denoised examples, i.e. $p_\phi (y|\hat{x}, t)$: this would evaluate the importance of using both perturbed and denoised examples.

Author response:
Thank you for suggesting these ablations. As described in the paper, we evaluated the relative importance of noisy and denoised examples in DA-classifiers by zeroing out one of the input images and measuring its effect upon classification accuracies; we found that while removing either one of the inputs causes the accuracy to drop below the noisy classifier, the drop is higher when the denoised image is masked out. This demonstrates that the DA-classifier makes use of both noisy and denoised examples in the classification. In fact, the intention behind using both noisy and denoised images as input is to allow the model to use parts of both images as needed. Upon acceptance, we will also include an ablation study considering the recommended ablations in the appendix.

We look forward to addressing any outstanding concerns in the discussion period.

[1] Bahjat Kawar, Roy Ganz, Michael Elad. Enhancing Diffusion-Based Image Synthesis with Robust Classifier Guidance. TMLR, 2023.
[2] Nichol, A., Dhariwal, P., Ramesh, A., Shyam, P., Mishkin, P., McGrew, B., Sutskever, I. and Chen, M. GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models. ICML, 2022.