Source-linked AI summary
AugMax: Adversarial Composition of Random Augmentations for Robust Training
Haotao Wang, Chaowei Xiao, Jean Kossaifi, Zhiding Yu, Anima Anandkumar, Zhangyang Wang
TL;DR
Deep networks trained on limited data remain vulnerable to unforeseen distribution shifts, motivating augmentation that combines diversity with hardness. AugMax adversarially mixes randomly selected operators, while DuBIN disentangles the resulting feature heterogeneity. AugMax-DuBIN improves robustness across natural-corruption benchmarks and other distribution shifts, with reported gains over prior state of the art.
Problem
Deep networks trained on limited datasets remain vulnerable to unforeseen distribution shifts, motivating a unified treatment of augmentation diversity and sample hardness.
Method
AugMax randomly selects augmentation operators and adversarially learns their mixing strategy, while DuBIN disentangles the resulting instance-wise feature heterogeneity.
Results
AugMax-DuBIN achieves state-of-the-art robustness on natural-corruption benchmarks and improves robustness against other common distribution shifts, surpassing prior methods by 3.03% on CIFAR10-C.
Takeaways & Limitations
The combination consistently improves robustness across several natural-corruption benchmarks while training from clean images.
Takeaways & Limitations
AugMax uses a heuristic unification of diversity and hardness, and better trade-offs may exist.
Abstract
from arXiv · showhide
Data augmentation is a simple yet effective way to improve the robustness of deep neural networks (DNNs). Diversity and hardness are two complementary dimensions of data augmentation to achieve robustness. For example, AugMix explores random compositions of a diverse set of augmentations to enhance broader coverage, while adversarial training generates adversarially hard samples to spot the weakness. Motivated by this, we propose a data augmentation framework, termed AugMax, to unify the two aspects of diversity and hardness. AugMax first randomly samples multiple augmentation operators and then learns an adversarial mixture of the selected operators. Being a stronger form of data augmentation, AugMax leads to a significantly augmented input distribution which makes model training more challenging. To solve this problem, we further design a disentangled normalization module, termed DuBIN (Dual-Batch-and-Instance Normalization), that disentangles the instance-wise feature heterogeneity arising from AugMax. Experiments show that AugMax-DuBIN leads to significantly improved out-of-distribution robustness, outperforming prior arts by 3.03%, 3.49%, 1.82% and 0.71% on CIFAR10-C, CIFAR100-C, Tiny ImageNet-C and ImageNet-C. Codes and pretrained models are available: https://github.com/VITA-Group/AugMax.
1 Introduction
AugMax unifies diverse augmentation coverage with adversarially hard samples, while DuBIN addresses the resulting heterogeneous features; together, AugMax-DuBIN improves robustness to distribution shifts.
- 1 Introduction: Out-of-distribution shifts from corruptions, perturbations, and changing domains expose weaknesses in models trained on limited clean data.This motivates robustness methods that work without requiring extensive coverage of every deployment scenario.
- 1 Introduction: AugMix increases augmentation diversity, whereas adversarial training targets hard worst-case samples but typically incurs higher training cost.The paper motivates combining these complementary properties rather than relying on either category alone.
- 1 Introduction: AugMax unifies diversity and hardness by randomly selecting augmentation operators and adversarially learning their mixing strategy.Unlike AugMix, AugMax learns the mixing parameters rather than sampling them randomly.
- 1 Introduction: DuBIN adds parallel instance normalization to dual batch normalization to disentangle instance-wise heterogeneity from AugMax samples.The stronger augmentation distribution makes training more challenging, motivating this finer-grained normalization scheme.
- 1 Introduction: AugMax-DuBIN surpasses state-of-the-art methods on CIFAR10-C, CIFAR100-C, Tiny ImageNet-C, and ImageNet-C by 3.03%, 3.49%, 1.82%, and 0.71%, respectively.The framework trains on clean images and also improves robustness against other common distribution shifts.
2 Related Work
Prior robustness work uses augmentation to increase diversity or adversarial training to generate hard samples, alongside normalization methods designed for different feature statistics.
- 2 Related Work: Robustness research also includes corruption benchmarks, perturbation benchmarks, natural adversarial examples, and methods for diagnosing failures under distribution shifts.These efforts document brittleness beyond standard benchmark performance.
- 2 Related Work: AugMix randomly mixes multiple diverse augmentations and established strong robustness on natural-corruption benchmarks.It represents the diversity-oriented augmentation category.
- 2 Related Work: Adversarial training uses adversarial samples as augmentation and improves robustness to corruptions and domain gaps, but can increase training time and reduce clean accuracy.These trade-offs distinguish hardness-oriented methods from ordinary augmentation.
- 2 Related Work: Normalization research has extended batch normalization with instance, layer, group, and multiple-normalization variants for task-specific feature handling.This work provides context for introducing a normalization scheme tailored to heterogeneous augmented features.
3 Method
AugMax generates augmented samples by optimizing adversarial mixing parameters over randomly selected augmentation operators, while DuBIN addresses the resulting feature heterogeneity through parallel instance and dual-batch normalization.
- 3.2 AugMax: Augmented Training with Unified Diversity and Hardness: AugMix creates diverse samples through multiple augmentation chains and random combinations, whereas AugMax changes the mixing-parameter selection to an optimization procedure.This comparison identifies the core methodological difference between the two augmentation frameworks.
- 3.2 AugMax: Augmented Training with Unified Diversity and Hardness: AugMax generates an augmented image by learning mixing parameters that maximize the model loss under bounded mixing constraints.The optimization is reparameterized with a softmax so the mixing weights remain normalized.
- 3.2 AugMax: Augmented Training with Unified Diversity and Hardness: AugMax training is formulated as minimax optimization, with the adversarially generated image used in the training objective.The learned parameters determine the augmentation output x∗ for each training image.
- 3.2 AugMax: Augmented Training with Unified Diversity and Hardness: Figure 1 shows clean-image clusters, broader AugMix coverage, and AugMax samples that better approach decision boundaries while retaining diversity.The visualization motivates combining broad coverage with harder samples.
- 3.3 DuBIN: Disentangled Normalization for Heterogeneous Features: DuBIN splits each feature tensor along channels, sends one half through instance normalization and the other through dual batch normalization, then concatenates the outputs.The design targets instance-level heterogeneity from augmentation diversity while DuBN handles group-level statistics of clean and corner-case samples.
- 3.3 DuBIN: Disentangled Normalization for Heterogeneous Features: In AugMax-trained networks, DuBIN produces lower average batch-normalization variances than DuBN, indicating that its instance-normalization branch reduces the feature diversity modeled by batch normalization.The supplied table reports BN statistics across layers of WRN40-2 trained on CIFAR100 with AugMax.
4 Experiments
Experiments evaluate AugMax-DuBIN across natural-corruption benchmarks, distribution shifts, ablations, stability, and training cost. The method achieves state-of-the-art robustness, with gains over AugMix and competing diversity- or hardness-based strategies.
- Robustness against Natural Corruptions: 3.49%: AugMax-DuBIN improves accuracy over AugMix on CIFAR100-C, while also achieving a 3.03% gain on CIFAR10-C.These results hold across different model structures, with the largest robustness gains on ResNeXt29.
- Robustness against Natural Corruptions: 0.52%: AugMax-DuBIN reduces mCE versus AugMix on ImageNet-C, while the DeepAugment combination improves mCE by 0.82% over AugMix + DeepAugment.The combined method achieves new state-of-the-art performance on ImageNet-C.
- Robustness against Natural Corruptions: 2.32%: AugMax-DuBIN improves mCE over AugMix on Tiny ImageNet-C, with an additional 0.45% improvement when combined with DeepAugment.It also outperforms ANT by a considerable margin on the evaluated Tiny ImageNet-C subset.
- Ablation Study: AugMax-DuBIN achieves the best performance among diversity, hardness, and naive combined baselines, supporting complementarity between diversity and hardness.The poorer results of naive combinations indicate that balancing the two properties is nontrivial.
- Robustness against Other Distribution Shifts: AugMax-DuBIN outperforms AugMix on CIFAR10.1 and CIFAR10-STA, showing improved robustness against dataset and spatial-transform distribution shifts.The evaluation uses the same ResNeXt29 models trained on CIFAR10.
5 Conclusion
The paper presents AugMax and DuBIN as a unified approach to robustness, combining augmentation diversity and hardness while addressing heterogeneous features. It reports state-of-the-art robustness on natural corruption benchmarks, while acknowledging that the diversity–hardness trade-off is heuristic.
- AugMax-DuBIN consistently demonstrates state-of-the-art robustness on several natural corruption benchmarks.
- AugMax unifies diversity and hardness, while DuBIN enables efficient training with the resulting heterogeneous features.
- The diversity–hardness trade-off remains heuristic, and better trade-offs may exist.
A AugMax Algorithm
The AugMax algorithm generates adversarially mixed images from clean inputs and uses them for robust classifier training. Its implementation applies accelerated adversarial optimization with early stopping and specified iteration, step-size, and batch-training inputs.
- AugMax images are generated from clean images using an accelerated adversarial attack that early-stops after misclassification.The experiments use k = 1, n = 5, and α = 0.1.
- Algorithm 1 takes an image, label, attack settings, augmentation function, classifier, and loss, then outputs an AugMax image.
- The adversarial updates ascend the loss with respect to the mixture variables m∗ and p∗.
- Algorithm 2 trains a robust classifier from dataset, iteration, batch-size, model, loss, learning-rate, and loss-trade-off inputs.
- The training procedure initializes the classifier parameters randomly before robust learning with AugMax.
B Details on AdvMix and AdvMax
AdvMix and AdvMax attack augmentation hyperparameters using worst-of-k selection. They choose the sampled configuration producing the largest classification loss.
- AdvMix and AdvMax use worst-of-k attacks on augmentation hyperparameters, with k set to 5.Examples include rotation angles and translation pixel numbers.
- For AdvMix, the method samples operation types and mixing parameters as in AugMix, then selects the hyperparameter set with largest classification loss.