Paper: Robust Classification via a Single Diffusion Model (ICLR-2024)

Reviewer comment (Evaluation):
The manuscript mentions multiple times the expensive inference in both time and memory - however, a quantitative analysis is missing. I would like to see a table with inference time and memory for this method in comparison to other generative classifiers and regular ones.

Author response:
**Question 1: Qualitative analysis of time complexity in terms of real-time cost.**

Thank you for your advice. We test the memory and real-time cost on a single 3090 GPU. The results are shown below:


| Method     | NFEs | Real Time | Memory Cost (MB) | 
|:-----------|:----:|:---------:|:----------------:|
| WRN-70-2   |  1   |   0.01    |       973        |
| AT-DDPM    |  1   |   0.01    |       973        |
| AT-EDM     |  1   |   0.01    |       973        |
| JEM        |  1   |   0.01    |       1655       |
| HybViT     |  1   |   0.01    |       2502       |
| SBGC       | 30TK |   15.78   |       3557       |
| DiffPure   | 101  |   0.60    |       3089       |
| LM (ours)  |  T   |   0.10    |       3401       | 
| DC (ours)  |  TK  |   9.76    |       2664       | 
| RDC (ours) | N+TK |   9.86    |       3661       | 
| RDC (ours) | N+T  |   1.43    |       3661       | 


We have also supplemented this analysis in the Table 3 in Appendix. As shown, our multi-head diffusion greatly improves the efficiency while maintains the accuracy and robustness.

Reviewer comment (Experiments):
The method appears to work well on "unseen" threats. However, all threats are limited to adversarial attacks. Is there any evidence for an increased robustness to other robustness aspects such as common corruptions (e.g., CIFAR10-C)?

Author response:
**Question 2: Experiments on CIFAR10-C.**

Thank you for your advice. The results on CIFAR10-C are shown below.


|                   | WRN70-2  | AT-DDPM | AT-EDM | DiffPure| DC   |
| :---------------- | :-------: | :-------------------: | :----------------: | :------: | :-------: |
| Glass blur        | 41\.2     | 52\.2                 | **55\.4**          | 54\.0    | 36\.4     |
| Gaussian noise    | 45\.8     | 69\.6                 | **73\.0**          | 72\.6    | 54\.0     |
| Shot noise        | 65\.4     | 87\.0                 | **90\.0**          | 89\.0    | 69\.0     |
| Speckle noise     | 69\.0     | 87\.2                 | **90\.8**          | 87\.8    | 67\.0     |
| Impulse noise     | 50\.4     | 75\.0                 | **77\.4**          | 81\.4    | 33\.4     |
| Defocus blur      | 59\.0     | 55\.6                 | **59\.6**          | 56\.0    | 58\.8     |
| Gaussian blur     | 71\.6     | 83\.4                 | **89\.0**          | 84\.4    | 88\.2     |
| Motion blur       | 84\.8     | 83\.2                 | **88\.0**          | 83\.4    | 85\.8     |
| Zoom blur         | 81\.8     | 81\.8                 | **90\.2**          | 84\.2    | 88\.6     |
| Snow              | 91\.0     | 85\.4                 | **91\.2**          | 87\.8    | 86\.2     |
| Fog               | **61\.6**     | 47\.8                 | 51\.6              | 47\.0    | 58\.4 |
| Brightness        | **62\.6**     | 57\.2                 | 60\.2              | 58\.4    | 60\.6 |
| Contrast          | 57\.9     | 37\.2                 | 42\.0              | 42\.8    | **58\.0** |
| Elastic transform | 58\.0     | 54\.4                 | **59\.2**          | 54\.6    | 56\.2     |
| Pixelate          | 82\.2     | 89\.4                 | **92\.2**          | 87\.4    | 83\.0     |
| JPEG compression  | 83\.8     | 88\.8                 | **92\.6**          | 88\.0    | 87\.2     |
| Spatter           | 90\.8     | 86\.8                 | **91\.6**          | 86\.6    | 86\.8     |
| Saturate          | **95\.0** | 86\.8                 | 91\.0              | 88\.2    | 91\.4     |
| Frost             | 57\.8     | 52\.8                 | 57\.8              | 56\.6    | **58\.2** |
| Average           | 69\.0     | 71\.2                 | 75\.9              | 73\.1    | 69\.4     |



The performance of our method does not exceed the baselines. We observe that our diffusion classifier assigns extremely low likelihood to out-of-distribution (OOD) samples. This may be because our diffusion classifier accurately captures the distribution of the training set but does not generalize well to data that is outside of this distribution. We intend to investigate this issue further and aim to enhance the OOD accuracy in future work.

Reviewer comment (Experiments):
The experimental evaluation is performed on a clearly separated and limited number of classes. Are there any results or theoretical insights on how this method would scale to more and potentially more fine-grained classes?

Author response:
**Question 3: Scalability to dataset with more fine-grained classes.**


Thank you for the suggestion. To verify the scalability of our method to fine-grained dataset, We test the robustness of different methods against $\ell_\infty$-norm threat model with $\epsilon_\infty=8/255$ on the CIFAR-100 dataset, following the same experimental settings as CIFAR-10. Due to the time limit, we only random sample 128 images for evaluation. The results are shown below.


| Method | Clean Acc | Robust Acc|
|---------|:-------:|:--------:|
| WRN-40-2 |   78.13     | 0.00 |   
| AT-DDPM |    63.56     |   34.64   |
| AT-EDM|  75.22        |   42.67 |
| DiffPure  |         39.06         |     7.81     |
| DC  |    79.69    |  39.06|
|RDC |      80.47       |  53.12|




We find that RDC still achieves superior result compared with the state-of-the-art adversarially trained models and DiffPure.
More surprisingly, we discover that DiffPure does not work well on CIFAR-100. We guess this is because CIFAR-100 has more fine-grained classes, and thus a small amount of noise will make the image lose its semantic information of a specific class. Hence, DiffPure is not suitable for datasets with more fine-grained classes and small resolution. This experiment indicates that our methods could be scaled to datasets with more fine-grained classes.

Reviewer comment (Reproducibility):
“... [the mthod] ... leverages an off-the-shelf diffusion model” - how is this possible when the last layer is modified? Which diffusion model is this exactly? How does it compare to the diffusion models previous works have used?

Author response:
**Question 5: Clarification of usage of the off-the-shelf diffusion model.**


We apologize for the ambiguity. In this paper, we initially explore the robustness of a generative classifier using a single off-the-shelf model. Subsequently, we propose the multi-head diffusion, which requires training and is not an off-the-shelf model. We have revised this in the revision.

Reviewer comment (Reproducibility):
Some details regarding the diffusion model need clarification. For example, the sampling strategy is not very clear to me. It seems the authors deploy a VP sampling with uniform sliding timesteps similar to the one used in Nichol and Dhariwal, (2021). In the appendix and the code the authors also seem to leverage some implementation from Karras et al. (2022). As we note the sampler in Karras et al. (2022) involves additional correction steps, the NFE of the diffusion backbone would be more than T.

Author response:
We would like to clarify that the time complexity of our diffusion classifier is not related to the sampling complexity. Our diffusion classifier initially calculates the Evidence Lower Bound (ELBO) to approximate the log likelihood, and then computes the classification probability using $p(y|\textbf{x}) \propto p(\textbf{x}|y)$. Therefore, the time complexity of our method is determined by the time required to calculate the ELBO. For each class, we need $T$ NFEs to calculate its ELBO, hence we need $K\times T$ NFEs to calculate $p(y|\textbf{x})$.

Reviewer comment (Theory):
The theoretical results are based on the hypothesis that the evidence lower bound is tight. It would be intriguing to explore the implications of the gap between the likelihood and this lower bound, especially when viewed through the lens of the Bayesian framework for uncertainty quantification. Exploring how this gap influences robust classification performance could enrich the paper's depth and utility.

Author response:
We agree that exploring the gap between ELBO and log likelihood on the robustness of diffusion classifier is quite interesting. However, calculating the exact likelihood of  the diffusion model is extremely hard. Song et al. (2021) suggest using the instantaneous change of variable theorem (Chen et al., 2018) to compute the likelihood that a diffusion ODE generates certain images. However, this likelihood is not equivalent to the one in the evidence lower bound (ELBO), which actually represents the likelihood of the diffusion SDE generating such an image. These two likelihoods are only equal when the score function is optimal. In practice, we train a UNet to approximate the score of the real data distribution, which is certainly not optimal. Consequently, these two scores are not identical, and we have not found any method to calculate the exact likelihood of the diffusion SDE. We will strive to explore new ways to calculate the exact likelihood and the gap between ELBO and log likelihood in  future work.

Reviewer comment (Experiments):
While the current experiments are limited to relatively small-scale datasets, there is inherent value in examining the method's scalability. It would be beneficial if the authors could present results or potential methodologies to apply their approach to larger datasets, drawing inspiration from other generative classifiers in the Bayesian paradigm, such as those highlighted by Heek and Nal (2019) and Han et al. (2022).

Author response:
Thank you for the suggestion. In Appendix B.3, we have provided an experiment on Restricted ImageNet, which shows the effectiveness of our approach to larger datasets. During the rebuttal phase, as suggested by Reviewer Xkzq, we supplement an experiment on CIFAR-100. The results of different methods are shown below.

| Method | Clean Acc | Robust Acc|
|---------|:-------:|:--------:|
| WRN-40-2 |78.13| 0.00 |   
| AT-DDPM|63.56| 34.64|
| AT-EDM  | 75.22 |42.67|
| DiffPure | 39.06|7.81|
| DC  |79.69|39.06|
|RDC |80.47|53.12|

Our proposed RDC achieves superior result compared with state-of-the-art adversarially trained models and DiffPure.
These two experiments imply the scalability of our method to larger datasets with more fine-grained classes.


We apologize for the confusion regarding the application of the methods proposed by Heek and Nal (2019) and Han et al. (2022) to enhance the scalability of our approach. Our method, being a generative classifier, utilizes the ELBO to approximate the log likelihood. We argue that the primary challenge in adapting our method to larger datasets is finding a more efficient way to compute the ELBO. The ATMC sampler proposed by Heek and Nal (2019) is designed to accelerate the training of Bayesian networks, while Han et al. (2022) focus on constructing a diffusion bridge for converting logits from discriminative classifiers into more accurate and calibrated probabilities. We reckon that these methods are quite different from our paradigm since they are discriminative approaches and hence their method could not be directly used to improve the scalability of our diffusion classifier.

Reviewer comment (Novelty):
This paper also has close relation to other generative classifiers not specified for robust classification. Although the authors have discussed some concurrent works, some prior works may need to include and discuss potential connections.

Author response:
Thank you for your suggestion. We have incorporated a discussion of more popular generative classifiers (Mackowiak et al., 2021; Hoogeboom et al., 2021; Han et al., 2022) into Section 2 of the revision.

Reviewer comment (Novelty):
The current form of the paper draws parallels to several generative classifiers, though not specifically designed for robust classification. While some works have been discussed as concurrent works, it might be helpful in enhancing the completeness to integrate and discuss other prior works, emphasizing their relevance and potential connections to the proposed methodology.

Author response:
Thank you for your suggestion. We have incorporated a discussion of more popular generative classifiers (Mackowiak et al., 2021; Hoogeboom et al., 2021; Han et al., 2022) into Section 2 of the revision.

Reviewer comment (Theory):
Even with the proposed multi-head U-Net and other complexity reduction measures, the computational complexity still seems to be high. While the likelihood maximization step only requires $N=5$ forward and backward passes, approximating $p(x|y)$ requires $T$ U-Net queries. That being said, this drawback can be left for future work.

Author response:
***Question 1: Computational complexity.***

We agree that the high computational complexity is a drawback of diffusion classifiers, and more generally generative classifiers. In our paper, we have made significant efforts to reduce the computational complexity by proposing multi-head UNet and various accelerated sampling strategies, as detailed in Section 3.5. Recently, we also made some progress for further reducing the time complexity of diffusion classifier. We found that a significant drop in robust accuracy, primarily caused by reducing $T'$ (i.e., number of sampling timesteps), is due to the large variance in diffusion loss. Owing to this substantial variance, more instances of diffusion loss must be calculated to achieve stable predictions. However, if we manage to reduce the variance of the diffusion loss (for example, by using the same $\textbf{x}_t$ for different classes), it becomes feasible to decrease $T'$ and thus achieves faster inference speeds. By implementing this approach, **we successfully reduce the time complexity of the diffusion classifier by 200 times while maintaining the accuracy of clean data.** In the future, we will continue our work with the diffusion classifier, striving to further reduce its time complexity without compromising either robust or clean accuracy.

Reviewer comment (Theory):
Theorem 3.2 discusses "optimal diffusion models". In what sense is the diffusion model optimal? The proof to this theorem clarifies that such a model minimizes the noise estimation error, but this should be in the main text.

Author response:
***Question 2: Unclear definition of optimal diffusion model.***

Thank you for pointing this out. Yes, the optimal diffusion model  has minimal noise estimation error. We have revised our paper to make this clearer.

Reviewer comment (Theory):
There is a softmax operation in Theorem 3.2, but the quantity on which it operates is a scalar (a norm square divided by some variance). Clarification is needed on what softmax exactly means here. The same applies for Corollary 3.3.

Author response:
***Question 3: Unclear definition of softmax operator.***

Sorry for the ambiguity. The softmax function is operated over all data samples. We have revised our paper to make this clearer.

---

R3: ***Question 1: Softmax in Corollary 3.3.***

Yes, the softmax operator in Corollary 3.3 is defined as the softmax across all classes. As demonstrated in Appendix A.1, $p(y|\textbf{x})$ is the $y$-th output of the softmax over the log likelihoods of different classes. Our diffusion classifier $f(\textbf{x})_y$ utilizes the evidence lower bound to approximate the log likelihood of class $y$, therefore, the softmax here refers to the softmax applied across all classes.

Reviewer comment (Experiments):
The sentence "We find that the optimal diffusion classifier achieves 100% robust accuracy in both cases, validating our hypothesize that accurate density estimation of diffusion models facilitates robust classification" raises questions. How was this finding obtained? Also, it should be "hypothesis", not "hypothesize". Furthermore, what is the main gap between an optimal diffusion classifier and an empirical diffusion classifier? Is it the limited amount of training data? Or is it how well the U-Net is optimized?

Author response:
**Question 4: The gap between optimal diffusion classifier and empirical diffusion classifier.**


The primary difference between an optimal diffusion classifier and an empirical diffusion classifier is that the empirical diffusion classifier does not generalize well to the test data. In essence, the diffusion model fails to provide accurate score estimations for the noised versions of the test data. Conversely, an optimal diffusion model can yield accurate score estimations for test data. We empirically test the optimal diffusion classifier against AutoAttack, which achieves 100\% robustness. This result indicates that if a diffusion model can generalize  well to the test set, it would attain 100\% adversarial robustness.
The gap between the optimal diffusion model and the empirical one could be reduced by using more training data or designing advanced training techniques, which would further improve the performance of diffusion classifier.
In this paper, we focus on an alternative strategy (i.e., likelihood maximization) that moves the input $\mathbf{x}$ to $\hat{\mathbf{x}}$ that diffusion model could give a more accurate score estimation.

Reviewer comment (Reproducibility):
Section 3.5 states, "instead of calculating the diffusion loss using all timesteps like Eq. (9), we only sample a single timestep". Clarification is needed on how this timestep is sampled (e.g., uniformly randomly?).

Author response:
**Question 5: Lack of sampling details in likelihood maximization.**

Thank you for pointing this out. As you mentioned, we uniformly sample a timestep and then calculate the diffusion loss at that specific timestep. This process is used to approximate the expected diffusion loss across all timesteps. We have made this clearer in the revision.

Reviewer comment (Reproducibility):
Section 3.5 also states, "(BPDA) approximates the gradient with an identity mapping". Clarification is needed on how exactly the identity mapping is applied. A pseudo-code or Python code explanation would be appreciated.

Author response:
**Question 6: Specific details of BPDA**

Thank you for pointing this out. Here is the python code of BPDA:

```
class LikelihoodMaximization(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: Tensor):
        x = likelihood_maximization(x.detach().clone())
        return x

    @staticmethod
    def backward(ctx, upstream_grad):
        return upstream_grad
```

In the code, we do the forward pass normally, but in the backward pass, we directly return the upstream gradient. In other words, we view the Jacobian of likelihood maximization as an identity matrix.

Reviewer comment (Experiments):
It would be beneficial to have experimental results from the CIFAR-100 dataset to introduce more diversity in the evaluation.

Author response:
**Question 7: Result on CIFAR-100.**

Thank you for the suggestion. We further conduct experiments on CIFAR-100 and test the robustness of different methods against $\ell_\infty$-norm threat model with $\epsilon_\infty=8/255$, following the same experimental settings as CIFAR-10. Due to the time limit, we only random sample 128 images for evaluation. The results are shown below.


| Method | Clean Acc | Robust Acc|
|---------|:-------:|:--------:|
| WRN-40-2 |   78.13     | 0.00 |   
| AT-DDPM |    63.56     |   34.64   |
| AT-EDM  |  75.22        |   42.67 |
| DiffPure   |         39.06         |    7.81    |
| DC  |    79.69    |  39.06|
|RDC |      80.47       |  53.12|





We find that RDC still achieves superior result compared with the state-of-the-art adversarially trained models and DiffPure.
More surprisingly, we discover that DiffPure does not work well on CIFAR-100. We guess this is because CIFAR-100 has more fine-grained classes, and thus a small amount of noise will make the image lose its semantic information of a specific class. Hence, DiffPure is not suitable for datasets with more fine-grained classes and small resolution.

Reviewer comment (Evaluation):
Since AutoAttack with BPDA is used for evaluation, clarification is needed on whether the Square Attack component of AutoAttack is also included in the evaluation, and if so, whether Square Attack finds additional examples on top of the BPDA gradient-based attacks.

Author response:
**Question 8: Square attacks can not find more adversarial examples.**

Yes, we used Square Attack in AutoAttack for robustness evaluation. We found that Square Attack can not find additional examples on top of the BPDA adaptive attack.

Reviewer comment (Experiments):
In Figure 2a, the robust accuracy is barely over 30%. What is the value of $T$ and $T'$ for the main results (Table 1)? Does the result improve if $T'$ is larger than 1000? How large is $T$ during the training of the diffusion model? If a diffusion model is trained with fewer time steps (i.e., discretizing the trajectory into less than 1000 steps during training), should we expect the resulting classifier obtained via the proposed method to work well with a smaller $T'$?

Author response:
**Question 9: Experiments about $T'$.**

In Figure 2a, we show the robustness of Diffusion Classifier (DC), but not the Robust Diffusion Classifier (RDC). DC can only achieve 35.94\% robustness against the $\ell_\infty$ threat model with $\epsilon=8/255$, as shown in Table 1.

During training, $T$ is set as 1000 in DDPM (Ho et al., 2020). EDM (Karras et al., 2022) extends the discrete diffusion model into continuous diffusion model, where $t$ represents the standard deviation of isotropic Gaussian noises added to images. Consequently, in the context of EDM, there is no equivalent concept of $T$. Instead, during training, a $t$ value is randomly sampled from the range $[\sigma_{min}, \sigma_{max}]$.
$T'$ is set as 1000 in Table 1.


As argued in Karras et al. (2022), the training and inference of diffusion model can be decoupled. During training, a continuous $t$ could provide more accurate score estimation for different $t$, or even play a role of regularization, making the diffusion model generalize better. Therefore, $T'$ during inference is not related with $T$ during training. 
Therefore, we assert that reducing $T$ during training does not contribute positively to the performance of the diffusion classifier. We believe that a significant drop in robust accuracy, primarily caused by reducing $T'$, is due to the large variance in diffusion loss (See Question 1 for a solution). 




We agree that adopting a large $T'$ may further improve the robustness of the diffusion classifier. We conduct another experiment where we set $T'=2000$. As shown below, the robustness is further improved. 


| $T'$ |   Robustness  |
|---------|:-------:|
|  1000   |    35.94    |
| 2000    |    39.06       |

Reviewer comment (Experiments):
With multi-head diffusion, is it true that only the last layer receives the class condition signal? How does this affect performance compared to injecting class conditioning into various locations in the U-Net? It would also be useful to see some generations from this diffusion model.

Author response:
**Question 10: Details about multi-head diffusion.**

We did not add any class condition signal to the multi-head diffusion. The last layer of multi-head diffusion directly maps the feature into $3\times K \times H \times D$ dimension to output predictions for $K$ classes. If we inject the class condition signal into some intermediate locations in UNet, then we have to calculate $K$ times for the rest part of the UNet, because the activations of each class are not same after the injection.


Thank you for your advice on including the generation of multi-head diffusion. We were surprised to find that our multi-head diffusion achieves a 1.96 FID on the CIFAR-10 dataset, demonstrating its generative ability, which is comparable to the state-of-the-art generative models.