[Paper Review] AugMax: Adversarial Composition of Random Augmentations for Robust Training

Haotao Wang, Chaowei Xiao+4 more|arXiv (Cornell University)|Oct 26, 2021

Adversarial Robustness in Machine LearningComputer Science69 references35 citations

TL;DR

AugMax unifies diversity and hardness in data augmentation by adversarially mixing randomly sampled augmentations, aided by a novel DuBIN normalization, to achieve state-of-the-art robustness on multiple corruption benchmarks.

ABSTRACT

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.

Motivation & Objective

• Motivate robust learning under distribution shifts and corruptions by combining diversity (AugMix-style augmentation) with hardness (adversarial mixing).
• Propose AugMax to learn adversarial Mixing weights and parameters that maximize training loss.
• Reduce training complexity via accelerated adversarial strategies while maintaining gains over baseline augmentations.
• Introduce a normalization module DuBIN to disentangle instance-wise feature heterogeneity induced by AugMax.

Proposed method

• Extend AugMix by first randomly sampling augmentation operators, then learning an adversarial mixture of the selected operators (m, w).
• Formulate a minimax objective to maximize the loss with respect to augmentation parameters while minimizing the classifier parameters.
• Reparameterize w as w = softmax(p) to enable differentiable optimization.
• Use an accelerated adversarial attack method to solve the inner maximization efficiently, keeping training overhead modest.
• Introduce DuBIN (Dual Batch-and-Instance Normalization) to disentangle instance-level heterogeneity, with a parallel IN branch and a DuBN branch routed per augmentation type.
• Train AugMax-DuBIN and compare to AugMix and baselines across CIFAR-10/100, Tiny ImageNet, and ImageNet with corruption benchmarks.

Experimental results

Research questions

• RQ1 Can diversity and hardness in data augmentation be effectively unified to improve robustness to distribution shifts and corruptions?
• RQ2 Does learning an adversarial mixing of augmentations (AugMax) yield improved out-of-distribution robustness compared to random augmentations alone?
• RQ3 How does a specialized normalization (DuBIN) affect feature statistics and robustness when using AugMax?

Key findings

• AugMax-DuBIN achieves new state-of-the-art robustness on CIFAR10-C and CIFAR100-C, improving by 3.03% and 3.49% over AugMix respectively.
• On Tiny ImageNet-C and ImageNet-C, AugMax-DuBIN improves robustness by 1.82% and 0.71% over AugMix, respectively.
• AugMax training time is about 1.5x that of AugMix on ImageNet, offering a favorable efficiency vs traditional adversarial training which is ~10x slower.
• Combining AugMax with DuBIN reduces feature variance and improves BN statistics handling, enabling better robustness.
• AugMax-DuBIN outperforms both augmentation-based and adversarial training baselines across multiple datasets and models, including ResNet18, WRN40-2, and ResNeXt29.
• AugMax-DuBIN combined with DeepAugment can surpass DeepAugment + AugMix on ImageNet-C in mCE by about 0.71–0.82 points in reported comparisons.