DiffAug: A Diffuse-and-Denoise Augmentation for Training Robust Classifiers

We thank you for your insightful reviews and affirmative evaluation of our work. We thank you for your appreciation of our presentation and method’s simplicity/computational efficiency. We agree with you that it is indeed interesting to learn that classifiers can be improved without sacrificing test accuracy through the use of degraded images from the diffusion process! We answer your questions in the following:

[Q1] Thank you for this question. The reviewer is correct that we meant no prior work on training with diffused-and-denoised data, and in particular by “denoised” we meant “partially denoised”; we apologize for this confusion, and we will make sure to clarify this. In [38], the generated synthetic images are of high-quality since they use iterative denoising: in fact, to ensure better quality, they first fine-tune the Imagen model on ImageNet data and then use FID to select the model checkpoints and various hyperparameters of the sampling algorithm.

---

[Q2] Additional optimization objective refers to extra loss term. When extending previous methods with DiffAug, we modify their official code by simply adding $\mathcal{L}$ in Eq. 5 to the original loss-terms. For example, AugMix involves 2 loss-terms: a cross-entropy loss term and a Jensen-Shannon divergence loss term. When extending AugMix with DiffAug, we introduce Eq.5 as the third loss-term.

---

[Q3] Similar to many standard augmentation techniques, we do not rely upon labels for DiffAug broadening its applicability for different tasks and scenarios. For example, DiffAug can be applied at both train-time and test-time (when class labels are unknown). Furthermore, it enables direct use of unlabeled data or a mix of labelled and unlabeled data to train the diffusion model. In theory, it should be possible to achieve performance enhancements with DiffAug by using conditional diffusion models, especially with innovative choices of prompt. However, this introduces additional hyperparameters such as the guidance-strength and choice of prompt. Additionally, it requires two forward passes through the diffusion model (one unconditional and one conditional) per training step. The augmentation diversity may certainly be affected especially if the guidance-strength is very high: in such cases, it is also possible that the classifier could cheat by exploiting imperceptible image statistics (e.g., due to adaptive layer-norms).

---

[Q4] We use the official pretrained checkpoints for these models and all of them are trained using a ResNet-50 backbone. We will clarify this.

---

[Q5] As compared to DE, DDA is a more complex technique that utilises multiple-diffusion steps to transform a test-sample from an unknown distribution into the source distribution. We find that the multiple-step denoising can be beneficial when dealing with severe corruptions in ImageNet-C. For example, we find negligible difference between DDA and DE over uncorrupted ImageNet test examples when considering the DiffAug-trained models (76.60 & 76.67 respectively). Yes, an optimal multiple-step denoising approach that effectively transforms the image into the source distribution may further enhance classifier robustness on ImageNet-C. For other distribution shifts such as ImageNet-R and ImageNet-S, we observe that DE offers better test-time adaptation on average as compared to DDA.

---

[Q6,7] We thank you for sharing these works on novel methods for diffusion-based adversarial example generation. We agree with your analysis that DiffAug’s certified adversarial accuracy results may generalise to Diff-PGD. From our understanding of DiffAttack, it produces transferable untargeted adversarial attacks at a high success rate and hence, we will need to conduct an empirical analysis to study DiffAug’s effectiveness against these examples. While we are interested in understanding the performance of DiffAug and DiffAug-Ensemble against DiffAttack, we were unable to perform this experiment during the rebuttal period due to time and resource constraints. We will include a discussion on these methods and aim to provide a preliminary empirical analysis in the final submission for completeness.

In the following, we describe additional results on ImageNet-D, another stable-diffusion generated dataset designed to evaluate robustness (similar to DiffAttack). From Table 2 (global PDF), we find that extra synthetic data (from stable-diffusion, as described in global response) offers the most improvements against ImageNet-D --– similar to the suggestion in the DiffAttack paper. Interestingly, DiffAug training and DiffAug-Ensemble (DE) inference can offer further improvements: for example, RN50+Synth achieves 17.52% accuracy on ImageNet-D while RN50+Synth+DiffAug achieves 19.18% accuracy and can be further improved to 21.41%. Since these results are encouraging, we are curious to evaluate the performance of DiffAug against DiffAttack.

---

[Q8] This is a good question! DE utilises a set of different diffusion time-steps ($\\mathcal{S}$ in Eq. 6) while DDS uses a single diffusion time-step. Crucially, DDS is based on the randomized smoothing theory in [1] whereas DE is inspired from test-time augmentation techniques [2]. Additionally, DE uses average voting for prediction whereas DDS uses majority voting. Further, the number of samples generated for each input is typically higher in DDS (e.g., 100 samples for ImageNet) as compared to DE (9 new samples when using $\\mathcal{S}=\\{0,50,100 … 450\\}$). In Figure 11, we compare the performance between DE and DiffAug: since DiffAug is equivalent to DDS using 1 sample, we expect DDS results to be similar to Fig. 11b.

[1] Certified Adversarial Robustness via Randomized Smoothing

[2] Understanding test-time augmentation.

---

We hope this resolves all concerns and look forward to resolving any outstanding concerns during the discussion period.

We thank you for your detailed reviews. We thank you for appreciating the strengths of our method’s simplicity and computational efficiency and its ability to be combined with other augmentations to further improve robustness. In the following responses, we aim to address the weaknesses and answer your questions:

[Q1] Absence of ablation study of sampling methods and sampling processes. Did the use of better sampling methods and more steps achieve better results? Why adopt a single diffusion denoising step when the current sampling method, such as DPM-Solver, can generate good results within 10 steps? In theory, better sample generation can achieve better results.

This is a good question! To answer this question, we performed additional experiments during the rebuttal week using high-quality synthetic data from stable-diffusion as described in our global response (see [E1]). Specifically, we use 1.3M synthetic imagenet images in addition to the original ImageNet training data to finetune the torchvision resnet-50 model for 10 epochs. We call this RN50+Synth. To understand the utility of DiffAug in this case, we finetune another instance of the torchvision resnet-50 model for 10 epochs using DiffAug as well as additional synthetic data. We call this RN50+Synth+DiffAug.

From our experimental evaluation across several datasets, we find that DiffAug training and DiffAug-Ensemble (DE) inference offers complementary benefits to extra synthetic training data. We agree with you that, theoretically, DiffAug augmentations generated using a DDPM model trained on ImageNet should not provide further improvements when high-quality synthetic data from SD — trained on LAION-5B, that also covers ImageNet — is already available. Yet, we surprisingly observe that DiffAug improves over and beyond additional high-quality synthetic data! To explain this, we first note that DiffAug is qualitatively different from previous diffusion-based augmentation techniques. Depending on the diffusion time used to generate the DiffAug augmentation, the resulting image can vary greatly in quality (as shown in Fig. 1). As a result, classifying some of these augmented images is more challenging as compared to the original examples producing a regularizing effect that leads to empirical robustness improvements. Lastly, yet importantly, the complexity introduced by DiffAug training does not sacrifice test accuracy despite training on poor quality examples. This is an unusual and noteworthy property that can be understood by interpreting denoised examples to lie on the image manifold following recent theoretical studies [E.g., refs. 9 & 34 in the main paper].

In summary, DiffAug can enhance robustness even when efficient sampling techniques are available to synthesize high-quality images since its performance improvement can be attributed to the regularizing effect from learning to classify partially synthesized train examples (i.e., diffused-and-denoised examples). We will include these results in the final version of the paper and supplementary materials.

---

[Q2] Absence of ablation study of diffusion model. Can other diffusion model such as DiT also be used to apply DiffAug.

While we consider the DDPM model (a variance-preserving (VP) SDE) for all our ImageNet experiments, we use a variance-exploding (VE) SDE diffusion in our classifier-guided generation experiments for the CIFAR10 dataset. We find that DiffAug introduces identical improvements for both DDPM and VE-SDE in terms of classifier generalization, perceptually aligned gradients, and improved image-generation performance. Yes, in theory, it should be possible to apply DiffAug with a diffusion-transformer model instead of the UNet architecture.

---

We hope this resolves all concerns and look forward to resolving any outstanding concerns during the discussion period.

Augmentation with Diffusion-Transformers (DiTs): In theory, DiffAug can be applied using DiTs. Since DiTs are latent-diffusion models, we note that the augmentations can be generated by directly applying the VAE decoder over the one-step denoised latent representation. Previous works (e.g., [1,2]) follow this technique to use external guidance functions — that operate on images/piano-rolls as input — to guide the latent-diffusion sampling. Anecdotally, our informal experiments have suggested that DiffAug works with a variety of models. While we are greatly interested to train/evaluate classifiers with DiffAug using DiT, we are unable to report results of this experiment during the rebuttal discussion week due to time and resource constraints. We have already started working on this and we will certainly include results with DiT in the final revision.

[1] Universal Guidance for Diffusion Models. (ICLR 2024)

[2] Symbolic Music Generation with Non-Differentiable Rule Guided Diffusion (ICML 2024)

---

We hope these responses help address your concerns. If there are any other concerns, including other ablation analyses that you feel are important for inclusion in the final version, please let us know and we will gladly address your concerns/recommendations.

---

We thank the reviewer for promptly reading our rebuttal, stating their outstanding concerns, and giving us another opportunity to address their concerns. We truly appreciate this level of engagement.

Better Sampling Methods and More Steps.

We originally interpreted your question as follows: better images should yield better results, so wouldn’t we get “even better” results by using the “even better” generated images? To answer this, we considered a strong baseline: a large-scale high-quality synthetic dataset generated with a stronger diffusion model. In particular, we used high-quality synthetic data generated using Stable-Diffusion with 50 reverse-diffusion steps of the PNDM sampler. Then, we studied the role of DiffAug augmentations — generated with a single reverse-diffusion step (of DDIM, or equivalently DPM-Solver-1) — when extra synthetic images are already available. Our experimental results showed that, surprisingly, the single-step DiffAug augmentations introduce further improvements beyond extra synthetic images.

We believe we now understand your question better and it seems that you were interested in applying DiffAug with more denoising iterations (on the diffused training image), or using a different solver (e.g. DPM-solver-2). We address both of these questions directly, followed by a discussion:

(1) More denoising steps: In fact, we originally had the same initial intuition as the reviewer: that taking more denoising steps might yield better performance! In our preliminary analyses, we therefore did indeed explore a multi-step extension to DiffAug (e.g., DDIM). Interestingly, we found that multiple de-noising steps did not help as much as samples from a single reverse-diffusion step, which is why we did not pursue it further. This initially surprised us and prompted us to think about it more carefully — we provide a detailed discussion about this below in addition to our discussion in lines 162-167 of our original paper submission, where we briefly explain why we do not consider multi-step denoising techniques despite their potential to improve sample quality. However, we would be very glad to expand upon this explanation in the final version!

(2) Different sampling method: Additionally, in the last few days, based on this reviewer’s suggestion, we tried using a single-step of reverse-diffusion of DPM-solver-2, both for training the classifier, and also as the sampling method at test-time for the DiffAug ensemble-technique (DE) that we describe in the paper. In particular, we first diffused the train example to a random diffusion-time $t$ and used a single reverse-diffusion step of the order-2 DPM solver to integrate the diffusion ODE backwards from diffusion time $t$ to $t=\epsilon$ with $\epsilon=0.001$ instead of $0$ for numerical stability. The results are shown in the Table below: in both cases, while the resulting augmented images are of high visual quality, this did not work as well as the sampling strategy that we are already using (again, discussion is below). This was very interesting for us to try out, and we feel that having tried this is clearly strengthening our paper and results, as it is consistent with our understanding so far of what DiffAug is doing.

(We use $\downarrow$ to denote lower performance due to the use of DPM-Solver-2 instead of the default DPM-Solver-1 used in our paper.)

We now describe the benefits of augmentation with a single reverse-diffusion step in classifier-training (apart from computational efficiency) in the following:

i. A single reverse-diffusion step generates examples on the image manifold that cannot be generated by multi-step samplers. The generated sample lies in regions between high-quality samples and can be understood mathematically from Eq. 4, where we observe $\hat{\bf x}_t$ is a convex-sum over examples from the data distribution. Also, see Fig. 5 in the Appendix for an illustration on a toy dataset.
In fact, surprisingly, these "low-visual quality" samples are valuable in a way that "high-visual-quality" samples are not! It turns out that we get the best results by including the full range of visual quality: in Appendix B.5.2, we include an ablation study to identify the comparative advantages of “stronger” DiffAug augmentations ($t \in [500,1000]$) and “weaker” DiffAug augmentations ($t \in [0,500]$). While stronger augmentations offer greater performance improvements on ImageNet-C and OOD detection, we find using the entire diffusion time-range tends to yield better performance across all evaluated tasks.

We thank you for the detailed review with suggestions to strengthen the contribution. We thank you for recognizing the strengths of our method’s novelty, simplicity and computational efficiency. We are happy to see that you found our paper clearly written with nice visualisations. In the following responses, we aim to address the weaknesses and answer your questions.

Evaluation on other severity

Following this reviewer’s suggestion, we evaluate the models on all severity levels and include the results in the global PDF. We notice similar overall trends as those that we observed for severity=5. On average, we find that the DiffAug training and DiffAug-Ensemble inference leads to improved performances. We will include this in the final version of the paper.

---

Evaluation on Imagenet-A [1], ImageNet-D [2], Imagenet-S [3] and ImageNet-R [4].

As per this reviewer’s suggestion, we carried out evaluations on ImageNet-A and ImageNet-D. For this evaluation, we also considered the additional classifiers we trained during the rebuttal week: in particular, we trained classifiers using synthetic datasets generated with stable-diffusion (please refer to the global response for more details) and denote this as RN50+Synth. We use RN50+Synth+DiffAug when trained using DiffAug in addition to the synthetic data. We note that ImageNet-R and ImageNet-S results are already included in the Appendix (Table 6) – we will label the paragraph describing these results on page 5 (line 220) to ensure that these results are easily accessible. We summarize our findings as follows:

• DiffAug training introduces slight improvements in the default evaluation mode (i.e., directly evaluating on test examples) for these datasets while DiffAug-Ensemble (DE) inference introduces notable improvements for Imagenet-R, ImageNet-D and ImageNet-S.
• On ImageNet-S, the average performance improves from 15.45 (using default evaluation) to 17.99 (using DE). For DiffAug-trained models, the average performance improves from 15.51 (default) to 18.26 (DE). In particular, DiffAug training and DE inference improves the RN50 performance from 7.12 to 12.52.
• On ImageNet-R, we find similar observations as above. For example, DiffAug training and DE inference improves the accuracy of RN50 from 36.16 to 41.61. On average, DE enhances the performance of DiffAug-trained models from 44.63 to 46.35. The best performance on ImageNet-R is obtained when using additional synthetic data: RN50+Synth has an accuracy of 49.28. Using both DiffAug training and DE inference, we can improve the performance to 54.71.
• On ImageNet-D, we observe that extra synthetic data from stable-diffusion helps enhance the robustness most as you had predicted. Interestingly, DiffAug helps to further improve robustness in this case. For example, RN50+Synth obtains an accuracy of 17.52 on ImageNet-D. Using DiffAug training in addition to extra synthetic data improves the performance to 19.18. When using DE-evaluation, we observe further improvements up to 21.41.
• On ImageNet-A, we generally find that all models except ViT struggle on this task. In this case, we find that DE inference leads to a reduction in performance on this task (on average). DiffAug training neither improves nor negatively affects the model performance in this case although we find slight improvements in many cases.

In summary, we find that DiffAug training improves classifier robustness in general and this can be further enhanced using DiffAug-Ensemble (DE) inference.

---

[W2] Not clear how useful is test-time augmentation, can add the ablation study.

[A] This is a good question, and we acknowledge that while we included this ablation in appendix B.6, the reviewers are not expected to look at appendix materials. We will gladly add a comment in the main body of the paper that highlights and refers the reader to these results.

To summarize these ablation results here: we find in multiple experiments that the DiffAug-Ensemble (DE) inference technique does indeed improve robustness. DE has two hyperparameters: maximum diffusion time-step and step-size controlling the number of augmentations. Overall, we also find that DE is robust to hyperparameter choices.

---

[W3] Unclear why the proposed method can handle covariate shifts. Can the author explain intuitively?

[A] We include a detailed answer in the global response providing an intuitive explanation of DiffAug's performance (in Q2). In summary, DiffAug generates augmentations of varying sample-qualities, each presenting a different level of challenge for the classifier. Intuitively, this produces a regularisation effect that helps improve various aspects of classifier robustness including covariate-shift. Notably, the complexity introduced by DiffAug does not lower the test accuracy despite the explicit use of low sample-quality images for classifier training.

---

[Q1] How do the hyper parameters [...] tradeoff if there is any?

[A] This is a very good question, and the reviewer is correct that smaller diffusion times produce DiffAug examples with localised modifications whereas larger diffusion times produce global modifications. In all our experiments, we use the entire diffusion time range to generate the DiffAug examples. To understand the influence of diffusion time-range, we conducted an ablation study (see Appendix B.5.2) comparing between weaker DiffAug augmentations ($t \in [0,500]$) and stronger DiffAug augmentations ($t \in [500,1000]$). While stronger and weaker augmentations help enhance different aspects of classifier robustness — for example, stronger augmentations are more useful to enhance performance on ImageNet-C and OOD detection — using the entire diffusion time scale tends to work better across all tasks.

---

We hope we have addressed all your concerns and look forward to discussing any outstanding concerns in the discussion period.

We thank you for the insightful review with suggestions to strengthen the contribution. We thank you for your appreciation of our method’s novelty, simplicity and computational efficiency, our presentation/figures and extensive ablation experiments. In the following responses, we aim to address the weaknesses and answer your questions.

[W1-A] Comparison [...] is lacking.

We address your suggestion for a comparison in the experiment [E1] described in the global response and include the results in Table 1 of the global PDF. In summary, we find that DiffAug training and DiffAug-Ensemble (DE) inference offer performance improvements over and beyond additional high-quality synthetic datasets.

---

[W1-B] Missing Related Work.

We thank you for bringing these papers to our attention, and we will certainly include discussion in the final version. We identify important distinctions and similarities between our work and the suggested papers. While the mentioned works focus exclusively on adversarial robustness, we focus on regular non-adversarial classifier training with an aim to improve different aspects of robustness (e.g., covariate shifts, OOD detection and certified adversarial accuracy). Nevertheless, both example papers demonstrate use of synthetic data generated with diffusion models trained with no extra data to improve the adversarial training technique and this is in fact very close to one of our research objectives: exploring the augmentation potential of diffusion models trained with no extra data. On the contrary, previous generative data augmentation approaches for regular non-adversarial classifier training often rely upon synthetic datasets from large diffusion models such as Stable-Diffusion or Imagen.

---

[W1-C] Specifically, [...] theory?

Thank you for this interesting question about performance and sample quality vs compute budget! We provide a detailed answers to this in the global response (please see answers to Q1 and Q2). In the context of adversarial training, we imagine that corresponding enhancements in performance are likely possible using DiffAug. Additional exploration on adapting DiffAug to adversarial training and empirical validation is out of scope, but is an interesting avenue for future work.

---

[W2] We appreciate the reviewer’s dedication to provide a well-informed review of this submission! We use official checkpoints and confirm our reproduction is error-free; we will release both the source-code and model checkpoints for reproducibility.

The DeepAugment+AugMix result in the original paper was obtained by considering all 5 severities of ImageNet-C, while Table 1 is for severity=5 (please see Fig. 1 in the global PDF for results over entire ImageNet-C). Our reproduction of the ImageNet-test and ImageNet-R accuracies of these models do also match the official results. Since the ImageNet-C dataset contains 750k test samples for each severity level, considering all 5 severities is computationally expensive for some methods (e.g., DDA) and in those cases, the standard practice is to evaluate on severity=5 since this represents the most challenging examples.

As we have described in lines 259-267 and as correctly identified by this reviewer, DDS already achieves state-of-the-art certified robustness results by using a pretrained 305M-parameter BeIT-L network. It is intuitive that training a classifier with DiffAug enhances the DDS certified accuracy and we empirically confirm this in Table 8 for a 86.6M parameter ViT-B model since this is computationally feasible.

---

[Q1] Additional optimization objective refers to extra loss term. When extending previous methods with DiffAug, we modify their official code by simply adding $\mathcal{L}$ in Eq. 5 to the original loss-terms. For example, AugMix involves 2 loss-terms: a cross-entropy loss term and a Jensen-Shannon divergence loss term. When combining it with DiffAug, we simply introduce Eq.5 as the third loss-term.

Stacking refers to the sequential application of distinct augmentations on the same image. While stacking augmentations may intuitively offer more robustness enhancements, our preliminary analysis showed no advantage of stacking augmentations over simply including the DiffAug loss in addition to the original training loss. Therefore, we choose to implement DiffAug as an additional loss term since this requires minimal change to the original code and fewer design choices (e.g., order of augmentations).

---

[Q2] Similar to many standard augmentation techniques, we do not rely upon labels for DiffAug broadening its applicability -- e.g., DiffAug can also be applied at test-time when labels are unknown. In theory, DiffAug with conditional diffusion models should not cause any problem. However, this introduces additional hyperparameters such as the guidance-strength and choice of prompt. Additionally, it requires two forward passes through the diffusion model (one unconditional and one conditional) per training step. If the guidance-strength is very high, the classifier could cheat by exploiting imperceptible image statistics (e.g., due to adaptive layer-norms). This is an interesting extension for future exploration.

---

[Q3] In summary, DiffAug generates augmentations of varying sample-qualities, each presenting a different level of challenge for the classifier. Intuitively, this produces a regularisation effect that helps improve various aspects of classifier robustness including covariate-shift. Since both DiffAug and DDS utilise a single reverse-diffusion step, we can observe improvements in certified adversarial accuracy. Yes, we can also improve empirical adversarial robustness against $l_2$ perturbations by using a majority vote following the Algorithm 2 in [1].

[1] (CERTIFIED!!) ADVERSARIAL ROBUSTNESS FOR FREE!

---

We hope we have addressed all your concerns and look forward to discussing any outstanding concerns in the discussion period.