We introduce DiffAug, a simple and efficient diffusion-based augmentation technique to train image classifiers for the crucial yet challenging goal of improved classifier robustness. Applying DiffAug to a given example consists of one forward-diffusion step followed by one reverse-diffusion step. Using both ResNet-50 and Vision Transformer architectures, we comprehensively evaluate classifiers trained with DiffAug and demonstrate the surprising effectiveness of single-step reverse diffusion in improving robustness to covariate shifts, certified adversarial accuracy and out of distribution detection. When we combine DiffAug with other augmentations such as AugMix and DeepAugment we demonstrate further improved robustness. Finally, building on this approach, we also improve classifier-guided diffusion wherein we observe improvements in: (i) classifier-generalization, (ii) gradient quality (i.e., improved perceptual alignment) and (iii) image generation performance. We thus introduce a computationally efficient technique for training with improved robustness that does not require any additional data, and effectively complements existing augmentation approaches.
Paper
Similar papers
Peer review
Summary
The paper proposes DiffAug, a new method for training image classifiers that are more robust. DiffAug is based on diffusion models and effective at improving classifier robustness in several ways, such as resistance to variations in the data. It also can improve the performance of classifier-guided diffusion models. Furthermore, DiffAug is computationally efficient and can be combined with other augmentation techniques.
Strengths
1. The writing and the presentation of the work is good and easy to follow. 2. The finding, that training with degraded images from the diffusion process won’t harm the classifier’s performance but even improve it, is interesting. 3. The experiments are comprehensive and diverse, showing the effectiveness of the proposed method.
Weaknesses
While I currently find the paper favorable and have no major weaknesses to point out, I do have some questions that the authors' responses could help clarify (see the part below). I am open to raising the score based on their explanations.
Questions
1. The authors state on Line 106 that they are unaware of prior studies training classifiers with denoised examples. However, Line 20 mentions previous work using "synthetic images generated using Imagen [38]". Does reference [38] not qualify as a "previous study on training classifiers with denoised examples"? Perhaps the authors meant there's no prior work on training with diffused-and-denoised data? 2. The authors mention "include Eq. 5 as an additional optimization objective" (Line 128-129). Could they elaborate on why it's considered "additional"? How does Eq. 5 differ from the standard classification loss typically used during classifier training? 3. In Section 4, the experiments explore "unconditional improved-DDPM." Have the authors investigated using DiffAug with "conditional" diffusion models? Can such text guidance improve DiffAug's performance, or would the "conditional" process limit augmentation diversity and hinder performance? 4. In Table 1, which model structure is trained on when evaluating “AM, DA, and DAM”? It seems the first column didn’t present enough information on that. 5. Table 1 results (columns 3, 4 & 8, 9) consistently show DE underperforming compared to DDA(SE). What explains this difference? Could it be due to DDA's "self-ensemble" strategy (although DE also uses ensembles)? Does this suggest that using even multiple-step-denoised examples can enhance classifier robustness? 6. While the paper focuses on improving classifier’s robustness based on diffusion models, there's a complementary area of research: diffusion-based attacks (DiffAttack [1], Diff-PGD [2]). How effective is DiffAug against these attacks? Given Diff-PGD's constraints with perturbation norms, the results in the paper on certified adversarial accuracy might generalize. However, DiffAttack is unrestricted. Can DiffAug still be applied in such cases? The authors should supplement these discussions in the paper for better evidence of DiffAug’s robustness. 7. About the related work, the idea of leveraging a diffusion model for augmentation and incorporating denoised examples to improve classifier robustness was also mentioned in the discussion section of DiffAttack [1]. Citing this related work would be appropriate. 8. Both DE and DDS employ a single reverse diffusion step during testing. What are the key distinctions between these methods? There seem no comparisons between them or descriptions about their differences. [1] Chen, Jianqi, et al. "Diffusion models for imperceptible and transferable adversarial attack." arXiv preprint arXiv:2305.08192 (2023). [2] Xue, Haotian, et al. "Diffusion-Based Adversarial Sample Generation for Improved Stealthiness and Controllability." arXiv preprint arXiv:2305.16494 (2023).
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
Yes. Limitations and social impacts have been involved in the paper.
Summary
This paper applies a diffusion-based data augmentation method to enhance the robustness of classifier. First, a gaussian perturbation is applied to train examples and then a single diffusion denoising step is applied to generate the augmentations. Besides, DiffAug can also be combined with other augmentations to further improve robustness. Empirically, DiffAug can achieve improvements in classifier generalization, gradient quality and image generation performance.
Strengths
1. DiffAug can be combined with other augmentations to improve robustness. Besides, DiffAug can be also used to improve many other performance, such as classifier generalization, gradient quality and image generation performance 2. DiffAug is simple, computational and easy to follow to improove robustness.
Weaknesses
1. Absence of ablation study of sampling methods and sampling processes. Did the use of better sampling methods and more steps achieve better results? 2. Absence of ablation study of diffusion model. Can other diffusion model such as DiT also be used to apply DiffAug
Questions
As shown in Weakness 1, 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
Rating
6
Confidence
4
Soundness
2
Presentation
2
Contribution
2
Limitations
yes
Thank you for your response!
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. | Training Method | | |Evaluation | |:---:|:---:|:---:|:---:| | | Default | DiffAug-Ensemble (DPM-Solver-1) | DiffAug-Ensemble (DPM-Solver-2) | | AM | 26.72 | 34.08 | 31.77$\downarrow$ | | AM+DiffAug/DPM-Solver-1 | 29.47 | 38.58 | 35.56$\downarrow$ | | AM+DiffAug/DPM-Solver-2 | 22.96 $\downarrow$ | 29.69$\downarrow$ | 25.16$\downarrow$ | (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.
Thank you for your response! (contd.)
**ii. Improved samplers generate higher quality examples faster but may not offer regularization effects similar to DiffAug.** To demonstrate this, we conducted experiments using additional high-quality synthetic samples generated with Stable-Diffusion. Improved samplers such as DPM can help generate new synthetic data faster. While large-scale synthetic data can improve performance on ImageNet-R, ImageNet-Sketch and ImageNet-D, we show that both DiffAug training and also DiffAug-Ensemble (DE) inference offer additional improvements in each of these cases. We attribute the performance improvements from extra synthetic images to the manually designed prompts (Table 8 of [1]) and large-scale upstream dataset (LAION-5B) used to train stable-diffusion. On the other hand, the benefits of DiffAug training can be attributed to the regularization effect introduced by training over examples lying in regions between clean samples (and on the image manifold). **iii.** Informally, we also found that there is another subtle catch with increasing the number of iterations. Our current method requires essentially no hyper-parameter tuning: it just works. However, we found that **DiffAug with multiple reverse-diffusion steps tended to generate augmented examples that effectively belong to a different class than that of the original training image.** This happens because we use the entire diffusion time range by default — in particular, images diffused farther could get augmented to an image that belongs to a different class. On one hand, we were able to address this, by, e.g. defining an upper limit $\tau$ such that the forward-diffusion step is applied with $t \le \tau$. However, any such solution that we considered ultimately introduced additional hyperparameters that required tuning, but without noticeable gains, and in fact generally missing out on the robustness improvements from stronger augmentations. Interestingly, while DiffAug with one reverse-diffusion step also does not preserve the class-label for larger diffusion time $t$ and injects label-noise while training, the model can still "learn" from these images without underfitting since the label-noise is correlated with visual quality. For more details, please refer to lines 148-167 on page 4 of the main paper; our current discussion in the paper is brief, but we would be glad to expand it. For test time image adaptation, DDA uses multiple reverse-diffusion steps to deal with severe ImageNet-C corruptions and an improved sampler such as DPM can help accelerate sampling in this case. However, we note that DDA contains hyperparameters (in addition to the diffusion range $\tau$) that can strongly influence the test time performance. DiffAug-Ensemble is an alternative approach that is robust to hyperparameter choices (Appendix B.6) and demonstrates comparable or better improvements as compared to DDA using a single reverse-diffusion step. We feel that it may be advantageous to implement and evaluate a multi-step DiffAug-Ensemble technique using DPM-solver. *** In summary, the benefit (in terms of regularization and robustness) of our augmentations appears to be in where the augmented data points sit in relation to the data manifold, and the “effective” augmentations do not necessarily correlate with visual quality (as supported by our experiments including: synthetic data, multi-step denoising, DPM-solver-2 denoising, Appendix B.5.2 ablation). For this reason we find that using a single reverse-diffusion DDIM step is both crucial for improvements introduced by DiffAug training, and is robust in that it does not require any hyper-parameter tuning. We appreciate your insightful questions and recommendations, and we look forward to including appropriate parts of this discussion in the final version of the paper. For the final version, we will also implement DiffAug-Ensemble (for inference) using the DPM-solver and also include an analysis of hyperparameters (e.g., diffusion-range) and runtime. We will also include the results of our preliminary analysis where we generated DiffAug train augmentations using multiple steps (e.g., DDIM, DPM-Solver-2). We hope that these additional experiments can inspire multi-step extensions to DiffAug in future work. More specifically, while many multi-step samplers are impressively optimised for improved efficiency and sample quality (i.e., for humans), multi-step samplers for improved regularizations (i.e., for downstream neural models) likely require alternative designs. Finding these designs, and using them to build on DiffAug, would be a very exciting future research direction. [1] Leaving Reality to Imagination: Robust Classification via Generated Datasets ***
Thank you for your response! (contd.)
**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. ***
Thank you very much for the detailed reply, which has solved my main concern, I will update my score.
Thank you!
We thank you for carefully considering our response and engaging with us during the discussion period. We are confident that the experiments prompted by your reviews not only strengthen our contributions but also help provide a better understanding of our new augmentation technique.
Summary
The paper explores the use of diffusion as a data augmentation technique to train robust classifiers. Specifically, it investigates whether a diffusion model trained without additional data can be leveraged to enhance classifier performance. The study shows utilizing a diffusion model trained without extra data to improve classifier robustness on Imagenet-C by 2 points, while train improved classifiers with just a single step of reverse diffusion. The proposed method is evaluated on challenging benchmarks such as Imagenet-C and Imagine, utilizing Vision Transformer (ViT) and ResNet models.
Strengths
1. Originality: The paper proposes an effective method of one-step diffusion augmentation for training robust image classifiers. While similar augmentation techniques using generative models have been explored previously, the utilization of a single-step diffusion process in this context is interesting. The idea is straightforward yet innovative, offering a novel approach to enhancing classifier robustness with minimal computational overhead. 2. Clarity: the paper is clearly written. 3. Quality: the paper includes ablation study and visualization. Say on diffusion time.
Weaknesses
1. Soundness of results: the robustness is only evaluated on Imagenet-C (severity 5 only), which is only local corruption. Does the proposed method work when viewpoint, background, scale, style, texture change? What about the other severity? Say can test on Imagenet-A [1], ImageNet-D [2], Imagenet-S [3], ImageNet-R [4]. Since it is diffusion-based augmentation, the reviewer wondered if it will help ImageNet-D the most since this one is diffusion-based testing. Evaluation on those will help understand the strengths and weakness of the proposed approach. 2. Not clear how useful is test-time augmentation, can add the ablation study. 3. Unclear why the proposed method can handle covariate shifts. Can the author explain intuitively? The score will be updated based on whether evaluation concerns are addressed. [1] https://openaccess.thecvf.com/content/CVPR2021/html/Hendrycks_Natural_Adversarial_Examples_CVPR_2021_paper.html [2] https://openaccess.thecvf.com/content/CVPR2024/html/Zhang_ImageNet-D_Benchmarking_Neural_Network_Robustness_on_Diffusion_Synthetic_Object_CVPR_2024_paper.html [3] https://arxiv.org/abs/1905.13549 [4] https://openaccess.thecvf.com/content/ICCV2021/html/Hendrycks_The_Many_Faces_of_Robustness_A_Critical_Analysis_of_Out-of-Distribution_ICCV_2021_paper.html
Questions
How do the hyper parameters of the one-step diffusion set? If you add more noise, you may change globally, if only local noise, you can only handle local noise corruption. Can the author show an ablation study on this tradeoff if there is any?
Rating
6
Confidence
4
Soundness
3
Presentation
2
Contribution
2
Limitations
Yes.
Summary
This paper proposes a data augmentation method DiffAug for training robust classifiers. The method is very simple: first, add Gaussian noise to one training image, and then denoise the noisy images with a pre-trained diffusion model. They also propose methods for test-time augmentation using DiffAug. Experiments show that DiffAug is able to improve robust accuracy on a wide range of datasets such as ImageNet-C, R,S, OOD detection, and certification accuracy. This paper also shows that the proposed method can improve the generation quality of classifier-guided diffusion generation.
Strengths
1. To the best of my knowledge, this method is new. Most of the existing methods use diffusion models to generate new training examples instead of augmenting existing ones. The proposed method is simple and easy to follow and implement. 2. This paper provides extensive ablation experiments to demonstrate the effectiveness of the proposed method. The proposed method is able to achieve significant improvements even combined with SoTA augmentation for robustness, say AugMix and DeepAugment, which indicates the proposed method provides new information for the model to learn. DiffAug can also be used to improve the performance of guidance classifiers in diffusion generation, thus lead to better results. 3. Models trained with this method have perceptually aligned gradients and this paper provides nice figures to demonstrate this.
Weaknesses
1. **Missing related work and comparison** Comparison with existing work that uses large synthetic data datasets with diffusion models to improve training is lacking. For example [1] uses diffusion models to generate extra training examples to improve empirical adversarial robustness and [2] for certificated robustness. Both two works show significant improvements for robustness. Specifically, suppose DiffAug train the model with E epochs on a dataset of N images. SoTA diffusion models can generate new images with K steps (K can be small (K<= 10) using DDIM, even 1 step using some flow based methods). If we use the diffusion model to generate EN/K new images, and train the model on the original dataset with the generated images. The extra computational cost is the same as DiffAug, if I understand it correctly. Will DiffAug achieve better results over this method? If yes, what could be the motivation/theory? 2. **Concerns regarding some numbers reported** In the DeepAugment paper, they report results for ResNet50 on ImageNet-C ([3], Figure 5). The accuracy of DeepAugment+AugMix (DAM in this paper) is close to 60%, which is much higher than the numbers in this paper (around 40%). Similarly, the results of AugMix reported in the original paper are also much higher than the numbers in this paper. (I do not work in the corruption robustness field, so there may be errors in citing these numbers) Can the authors explain why the gap is so large? If it is due to different settings. I suggest following the original settings for a fair comparison. Similarly, the results on Table 8 are also significantly higher than numbers in [4], and far from SoTA. [1] Better Diffusion Models Further Improve Adversarial Training. ICML23 [2] A Recipe for Improved Certifiable Robustness. ICLR 2024 [3] The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution Generalization [4] (CERTIFIED!!) ADVERSARIAL ROBUSTNESS FOR FREE!
Questions
1. I do not pretty understand the sentence in line 128: what does "additional optimization objective" refer to? ``` When combining DiffAug with such novel augmentation techniques, we simply include Eq. 5 as an additional optimization objective instead of stacking augmentations ``` and line 131: why stacking augmentations are not good and what is the difference between stacking augmentations the the proposed method? ``` our preliminary analysis on stacking augmentations showed limited gains over simply training the network to classify independently augmented samples likely because training on independent augmentations implicitly generalizes to stacked augmentations. ``` 2. Why use unconditional diffusion models? Will using conditional diffusion models raise to some problem? 3. Can authors provide some insights why the proposed method can improve both corruption robustness and certificated adversarial robustness? Will this method also improve empirical adversarial robustness?
Rating
5
Confidence
3
Soundness
3
Presentation
3
Contribution
2
Limitations
Yes, authors adequately addressed the limitations.
Thank you!
We thank you for your prompt acknowledgement of our new results and insights. We feel that these new experiments, prompted by your review, indeed strengthen the contributions of the paper.
In fact, I believe the authors haven't addressed my concerns. For the first question, I was hoping to see attempts at different sampling methods. The authors proposed using high-quality synthetic data but neglected to explore sampling approaches. Regarding the second question, my main concern is with the network architecture of the diffusion model, not the training method (VPSDE or VESDE). I hope the authors can still address my concerns. Thank you.
Thanks for the response.
After carefully considering the authors' response, I have no further questions. I am pleased to recommend the acceptance of this paper, and I have raised my rating to 7. In the revised manuscript, I would be delighted to see not only the experiments on robustness against diffusion-based models as mentioned in Q6, but also those on the effect of DiffAug using conditional diffusion models, as discussed in Q3. Including these experiments would offer readers a more comprehensive understanding of the method's scope of application. This seems to be a point of curiosity for many, as also highlighted by Reviewer uEtP, and would make a valuable addition to the paper.
Thank you!
We thank you again for your insightful reviews and positive evaluation of our submission. We agree with your recommendation that DiffAug experiments with conditional diffusion models would be valuable. Based on both your and Reviewer uEtP’s interest in conditional diffusion models, and based on Reviewer 3vHc’s interest in Diffusion-Transformer models, we will include in the final version additional DiffAug experiments with Diffusion-Transformer (DiT). Since DiTs are trained for classifier-free guidance (i.e., both class-conditioning as well as null-conditioning), we will include additional settings with class-conditioning using DiTs in the final version.
Additional Information regarding Q2
We thank you again for your thorough review of our submission and insightful questions. Based on both Reviewer QY34’s recent positive recommendation and your question (Q2) about conditional augmentation, we want to inform you that we will include experiments on conditional DiffAug in the final version. In particular, we will explore class-conditioning with Diffusion-Transformer (DiTs) models, as that is also consistent with the DiT experiments we committed to include in the final version to Reviewer 3vHc. We hope we have addressed all your concerns and look forward to discussing any outstanding concerns in the remaining discussion period.
Thank authors for the response, and my concerns are addressed. I will keep my score as my original score is based on that the numbers are reported correctly and the proposed method should be better than using generated data to increase the dataset size.
Decision
Accept (poster)