AugMax: Adversarial Composition of Random Augmentations for Robust Training (2110.13771v3)

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.

AugMax: Adversarial Composition of Random Augmentations for Robust Training

The paper introduces AugMax, a novel data augmentation strategy designed to enhance the robustness of deep neural networks (DNNs) against out-of-distribution (OOD) data and unforeseen distributional changes. Data augmentation has been a well-established tool for improving model robustness by increasing the diversity of training data and introducing challenging, or 'hard', examples. AugMax ingeniously combines these two distinct aspects—diversity and hardness—into a unified framework.

Key Contributions

1. Novel Data Augmentation Framework: AugMax builds on the foundation of AugMix, which utilizes stochastic mixing of diverse augmentation operations to introduce diversity into the training data. AugMax adds to this a layer of adversarial training, wherein augmentation operators are first randomly sampled and subsequently combined in a manner that maximizes the model's classification loss. This adversarial approach introduces 'hard' examples, providing a comprehensive strategy that tackles both diversity and hardness.
2. DuBIN Normalization: The introduction of AugMax naturally leads to a broader and more heterogeneous input distribution. To address the challenges associated with this complexity, the authors propose a new normalization module, Dual-Batch-and-Instance Normalization (DuBIN). DuBIN disentangles instance-wise feature heterogeneity, thereby facilitating effective training under AugMax.
3. Empirical Evaluation and Results: Comprehensive experiments on several corrupted datasets—CIFAR10-C, CIFAR100-C, Tiny ImageNet-C, and ImageNet-C—demonstrate that models trained with AugMax and DuBIN (AugMax-DuBIN) exhibit superior performance in terms of robustness. These models outperform state-of-the-art baselines in RA (robustness accuracy) significantly, with the most notable improvements reported as 3.03% on CIFAR10-C and 3.49% on CIFAR100-C.

Theoretical and Practical Implications

From a theoretical perspective, AugMax showcases a successful integration of diversity and hardness, two traditionally separate streams of data augmentation for robustness. This integration could serve as a template for future developments in data augmentation, potentially leading to more sophisticated models of robustness that take into account a wider array of distributional challenges.

Practically, AugMax provides an accessible augmentation framework useful in various real-world applications where models must remain reliable under diverse and unpredictable conditions. The efficiency of AugMax, especially with the slight computational overhead added relative to standard non-adversarial training methods, makes it a feasible choice for robust model training.

Future Directions

Future research could explore alternative strategies for parameterizing the adversarial mixing weights and operators, potentially improving both computational efficiency and robustness further. Additionally, expanding the framework's applications beyond image classification to include other domains, such as natural language processing and time-series analysis, would be valuable. More extensive validation across different architectures and data domains could confirm the generalizability of the AugMax approach.

In summary, AugMax represents a significant step forward in the development of robust DNNs, providing a comprehensive augmentation strategy that effectively addresses both diversity and hardness, thereby enhancing model performance against a plethora of distributional shifts.