DiffAug: A Diffuse-and-Denoise Augmentation for Training Robust Classifiers
This work addresses the challenge of training robust image classifiers for computer vision applications, offering an incremental improvement by complementing existing augmentation methods without requiring additional data.
The authors tackled the problem of improving classifier robustness by introducing DiffAug, a diffusion-based augmentation technique that applies one forward-diffusion step and one reverse-diffusion step to training examples, resulting in enhanced robustness to covariate shifts, certified adversarial accuracy, and out-of-distribution detection, with further improvements when combined with other augmentations like AugMix and DeepAugment.
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.