Source-linked AI summary
AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty
Dan Hendrycks, Norman Mu, Ekin D. Cubuk, Barret Zoph, Justin Gilmer, Balaji Lakshminarayanan
TL;DR
Data shifts make image classifiers fragile and uncertainty estimates unreliable, especially under corruptions unseen during training. AugMix combines stochastic mixed augmentations with Jensen–Shannon consistency, achieving state-of-the-art robustness and calibration across corruption and perturbation benchmarks.
Problem
Train–test mismatches are common, but classifiers remain insufficiently robust and can produce overconfident predictions under shifted data.
Method
AugMix mixes diverse stochastic augmentation chains and enforces consistent classifier outputs with a Jensen–Shannon divergence loss.
Results
AugMix achieves state-of-the-art robustness and uncertainty estimation across CIFAR-10/100-C, ImageNet-C, CIFAR-10/100-P, and ImageNet-P.
Takeaways & Limitations
AugMix is simple to implement and can maintain calibration as the data distribution shifts.
Abstract
from arXiv · showhide
Modern deep neural networks can achieve high accuracy when the training distribution and test distribution are identically distributed, but this assumption is frequently violated in practice. When the train and test distributions are mismatched, accuracy can plummet. Currently there are few techniques that improve robustness to unforeseen data shifts encountered during deployment. In this work, we propose a technique to improve the robustness and uncertainty estimates of image classifiers. We propose AugMix, a data processing technique that is simple to implement, adds limited computational overhead, and helps models withstand unforeseen corruptions. AugMix significantly improves robustness and uncertainty measures on challenging image classification benchmarks, closing the gap between previous methods and the best possible performance in some cases by more than half.
1 INTRODUCTION
Train–test distribution mismatches are common, yet current classifiers remain fragile and often overconfident under shifted data. AugMix is proposed to improve robustness and uncertainty estimation while preserving standard accuracy.
- Motivation: 22% to 64%: modern-model classification error rises from ImageNet to ImageNet-C under common corruptions.ImageNet-C applies corruptions to test images, exposing sensitivity to unforeseen shifts.
- Motivation: Training against specific corruptions can encourage memorization and fail to generalize to unseen corruptions.This limits the effectiveness of corruption-focused training as a general robustness strategy.
- Contribution: AUGMIX targets both classifier robustness and uncertainty estimates under data shift.The method is designed for mismatched deployment conditions rather than only identically distributed evaluation.
- Results: 28.4% to 12.4% and 54.3% to 37.8%: corruption robustness error is reduced on CIFAR-10 and CIFAR-100, respectively.These reductions are reported relative to standard training procedures.
- Results: 57.2% to 37.4%: AUGMIX decreases ImageNet perturbation instability while achieving state-of-the-art corruption robustness.The reported ImageNet result covers robustness and prediction instability under shift.
2 RELATED WORK
Prior work studies robustness, calibration, and augmentation, but unseen corruptions remain difficult because models can memorize training distortions. Related methods span corruption benchmarks, calibration techniques, mixing, learned augmentation, and adversarial training.
- Robustness under Data Shift: Training on distortions can fail to generalize to unseen distortions or different parameter settings.This motivates evaluating robustness against corruptions not used during training.
- Robustness under Data Shift: ImageNet-C benchmarks generalization to unseen corruptions, while adversarial variants extend this robustness evaluation.These benchmarks frame data-shift robustness as a reliability problem for deployed systems.
- Calibration under Data Shift: Model calibration substantially deteriorates under data shift, motivating uncertainty evaluation beyond standard accuracy.Prior work also reports calibration improvements from ensembling and pre-training.
- Data Augmentation: Data augmentation methods include flips, crops, occlusion, CutMix, Mixup, learned policies, Patch Gaussian, and adversarial training.These approaches modify images or training procedures to improve generalization or robustness.
- Evaluation Context: ImageNet-C corruptions are reserved for testing rather than training, making them a measure of generalization under distribution shift.The benchmark contains corruption types applied at test time and absent from training.
3 AUGMIX
AugMix combines stochastic, layered augmentations with convex mixing and a Jensen–Shannon consistency loss. The design seeks diverse yet semantically related inputs while avoiding degradation from long augmentation chains.
- Method: AUGMIX samples simple operations stochastically and layers them to create diverse augmented images.The method is intended to fit easily into existing training pipelines.
- Augmentation Mixing: Multiple augmentation chains are mixed through convex combinations to preserve diversity without drifting too far from the original image.Directly composing long chains can degrade images and move them off the data manifold.
- Rationale: Mixing augmentation chains addresses the failure mode in which models memorize fixed augmentations rather than generalizing to varied corruptions.The method links augmentation diversity to robustness while controlling image degradation.
- Algorithm: The algorithm samples Dirichlet mixing weights, composes operations at varying depths, and applies a second Beta-sampled interpolation with the original image.These choices implement the augmentation and skip-connection mixing scheme.
- Consistency Loss: The Jensen–Shannon consistency loss minimizes divergence among classifier posteriors for the original image and augmented variants.It encourages consistent classifier embeddings across diverse transformations while using a bounded divergence.
4 EXPERIMENTS
The experiments evaluate AUGMIX across corruption robustness, perturbation stability, uncertainty calibration, and ImageNet-scale benchmarks. AUGMIX improves these measures across datasets and architectures, with component ablations identifying diversity, Jensen-Shannon consistency, and mixing as contributors.
- Evaluation settings: CIFAR-10-C and CIFAR-100-C evaluate corruption robustness across 15 corruption types and five severity levels without training on those corruptions.CIFAR-C datasets are constructed by corrupting the original test sets, while the 15 corruptions are excluded from training.
- Evaluation settings: CIFAR-10-P and CIFAR-100-P measure prediction stability through flip probabilities between adjacent frames under progressively changing perturbations.Their videos contain smaller perturbations than CIFAR-C and quantify prediction-stream jaggedness.
- CIFAR-10 and CIFAR-100: 16.6% lower absolute corruption error is achieved by AUGMIX than the Standard ResNeXt baseline on CIFAR-10-C, while gains transfer across architectures and CIFAR-100-C.The method also improves corruption robustness across individual corruption types and severity levels.
- ImageNet: 68.4% mCE versus 80.6% for the baseline is achieved by AUGMIX on ImageNet-C, and stacking it with SIN lowers mCE to 64.1%.The ImageNet-C table averages corruption error across all 15 corruption types; AUGMIX also improves clean accuracy.
- ImageNet: 37.4% mFR versus 57.2% for the baseline is achieved by AUGMIX on ImageNet-P, alongside more calibrated predictions as corruption severity increases.Calibration is assessed with Brier Score and RMS Calibration Error.
- Ablations: 11.2% error is obtained by full AUGMIX on CIFAR-10-C, compared with 26.9% for the standard setup, 17.0% with random chains, and 14.7% after adding JSD.The ablation attributes robustness to augmentation variety, Jensen-Shannon divergence consistency loss, and mixing; adding AUGMIX on top of Mixup raises error to 13.3%.
5 CONCLUSION
AUGMIX mixes randomly generated augmentations and uses a Jensen-Shannon loss to enforce consistency, achieving state-of-the-art robustness and calibration across multiple corruption and perturbation benchmarks.
- AUGMIX mixes randomly generated augmentations and uses a Jensen-Shannon loss to enforce consistency.
- AUGMIX obtains state-of-the-art performance on CIFAR-10/100-C, ImageNet-C, CIFAR-10/100-P, and ImageNet-P.
- AUGMIX models achieve state-of-the-art calibration and can maintain calibration as the distribution shifts.
- The authors position AUGMIX as a way to support more reliable models in safety-critical environments.
A HYPERPARAMETER ABLATIONS
The ablation study evaluates AUGMIX hyperparameters and reports stable ImageNet-C classification performance across tested settings, indicating limited sensitivity to tuning.
- The baseline configuration trains for 90 epochs with α = 0.5, three Jensen-Shannon examples, chain depth 1–3, and k = 3 augmentation chains.
- AUGMIX’s ImageNet-C classification performance is stable across changes to its hyperparameters.
B FOURIER ANALYSIS
Fourier sensitivity analysis measures error after adding individual frequency-specific perturbations to CIFAR-10 test images. AUGMIX is notably more robust than baseline and Cutout models at mid and high frequencies.
- The analysis adds 32 × 32 Fourier basis vectors individually to CIFAR-10 test images and records the resulting error rate.
- Low-frequency vectors appear at the heatmap center, while high-frequency vectors lie farther from the center.
- Baseline error rates exceed 80% for high-frequency perturbations, and Cutout shows a similar lack of robustness.
- AUGMIX maintains low-frequency robustness and is conspicuously more robust than the baseline and Cutout models at mid and high frequencies.
C AUGMENTATION OPERATIONS
The paper illustrates AUGMIX augmentation operations and warns that augmentation choices require care because some transformations can overlap with evaluation corruptions or alter an image’s class.
- Figure 10 illustrates augmentation operations applied to the same image, with some severities increased for detail.
- Contrast, color, brightness, sharpness, and Cutout are excluded because they may overlap with ImageNet-C test corruptions.
- Histogram color swapping can change a bird’s class, producing manifold intrusion.
D ADDITIONAL RESULTS
Additional results evaluate AugMix on CIFAR-10, CIFAR-10-C, and CIFAR-10-P, covering corruption accuracy, calibration, clean error, and mean flip probability. AugMix improves corruption robustness across CIFAR-10-C corruptions, while avoiding the clean-accuracy penalty associated with adversarial training.
- Figure 12 reports accuracy for each CIFAR-10-C corruption, spanning noise, blur, weather, and digital categories.
- AugMix improves corruption robustness across all CIFAR-10-C noise, blur, weather, and digital corruptions without training on those corruptions.
- Table 5 reports RMS Calibration Error for models and augmentation methods across CIFAR-10 and CIFAR-10-C.
- Table 6 reports CIFAR-10 clean error and CIFAR-10-P mean flip probability, both as percentages.
- Adversarial training performs well on CIFAR-10-P but substantially increases clean CIFAR-10 error, whereas AugMix does not.
E CALIBRATION METRICS
The RMS Calibration Error is estimated from finite test sets by sorting predictions by confidence and partitioning them into contiguous bins. The metric is distinct from classification error, and adding a refinement term yields the Brier Score.
- RMS Calibration Error is estimated by partitioning n test examples into b contiguous bins ordered by prediction confidence.
- This work uses adaptive bins containing 100 predictions, partitioning confidence scores over [0, 1].
- RMS Calibration Error differs from classification error because a random classifier with an approximately uniform posterior can be approximately calibrated.
- Adding the refinement term to squared RMS Calibration Error produces the Brier Score.