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

TL;DR

DiffAug is a diffusion-based augmentation method that enhances image classifier robustness against covariate shifts, adversarial attacks, and out-of-distribution data, using a simple diffusion process.

Contribution

The paper introduces DiffAug, a novel diffusion-based augmentation technique that improves classifier robustness without extra data and complements existing methods.

Findings

Single-step reverse diffusion improves robustness.

Combining DiffAug with other augmentations enhances performance.

Improves classifier-guided diffusion in generalization and image quality.

Abstract

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.