AugMax: Adversarial Composition of Random Augmentations for Robust Training

1 Introduction . Out-of-distribution ( OOD ) samples present a challenge when deploying AI models in the real world . Examples include natural corruptions ( e.g. , due to camera blurs or noise , snow , rain , or fog in image data ) , sensory perturbations ( e.g. , sensor transient error , electromagnetic interference ) and domain shifts ( e.g. , summer→ winter ) . However , deep networks are often trained on limited amounts of data which may not cover sufficient scenarios . As a result , they are vulnerable to unforeseen distributional changes despite achieving high performance on standard benchmarks [ 1–3 ] . This jeopardizes their trustworthiness as well as safe deployment in real-world environments . Thus it is critical to develop techniques that improve robustness even when training with relatively clean datasets . Several techniques have been proposed to consolidate the robustness against unforeseen corruptions , including robust data augmentation [ 5 , 3 , 8 , 6 ] , Lipschitz continuity [ 9–11 ] , stability training [ 12 ] , pre-training [ 13–16 ] , and robust network structures [ 17–19 ] , to name a few . Among these techniques , data augmentation is of particular interest due to its empirical effectiveness , ease of implementation , low computational overhead , and plug-and-play nature . ∗Work partially done during an internship at NVIDIA . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . There are mainly two categories of data augmentation approaches : The first category aims to increase diversity of the training data by composing multiple random transformations [ 20–22 , 5 , 3 ] . While standard data augmentation methods ( e.g. , random flipping and translation ) lead to poor robustness [ 5 ] , more aggressive combinations of multiple augmentations have shown promise . One such successful example is AugMix [ 5 ] , which stochastically samples from diverse augmentation operations and randomly mixing them to produce highly diverse augmented images . AugMix can increase the sample diversity coverage compared to standard augmentations , as shown in Figure 1 ( a ) and ( b ) . The second category aims to boost the hardness of the training data by sampling from the worst-case augmentations that tries to fool the model into misclassifying the samples . A common technique to achieve this goal is adversarial perturbation [ 6 , 23 , 8 ] . Training over such worst-case samples allows a model to actively fix its generalization weaknesses [ 24 ] and empirically improves its robustness against corruptions [ 25 , 26 , 7 ] . Due to the extra complexity to generate adversarial perturbations , the improved robustness usually comes at the cost of largely increased training time compared with non-adversarial methods [ 6 , 27 ] . An example falling into this category is the PGD attack [ 6 ] , as is shown in Figure 1 ( c ) . Previous work has focused on leveraging one of these category to improve robustness . In this paper , we unify both approaches in a single framework . In particular , we show that diversity and hardness are complementary and that a unification between the two is necessary to achieve robustness . We propose a new strategy to achieve this unification and successfully increase robustness . Summary of contributions : • We propose AugMax , a novel augmentation framework which achieves robustness through a unification between diversity and hardness , by searching for the worst-case mixing strategy . • Being a stronger form of data augmentation , AugMax leads to a significantly augmented and more heterogeneous input distribution , which also makes model training more challenging . To solve this problem , we design a new normalization strategy , termed DuBIN , to disentangle the instance-wise feature heterogeneity of AugMax samples . • We show that combination of AugMax and DuBIN ( AugMax-DuBIN ) achieves state-of-the-art robustness against corruptions and improves robustness against other common distribution shifts . AugMax achieves a good unification between sample diversity and hard corner-case generation during data augmentation . AugMax is built on top of the AugMix framework [ 5 ] which mixes multiple data augmentation operators in a multi-branch and layered pipeline . However , different from AugMix where augmentation operators and mixing weights are both randomly sampled , operators are first randomly sampled , followed by adversarially trained mixture of the selected operators in Augmax . This simple change from AugMix to AugMax results in considerable difference in their feature distributions . From the visualizations ( in Figure 1 ) , AugMax generates more adversarially “ hard samples '' , while still keeping a good amount of diversity compared to AugMix and PGD Attack . Searching for adversarial mixing strategies in AugMax is slightly more expensive since it involves adversarial mixing . To make this efficient , we adopt an efficient adversarial training strategy [ 28 ] . As a result , AugMax adds only a reasonable amount of extra complexity compared to non-adversarial training methods , while improving robustness significantly . For example , AugMax training time is only ∼ 1.5 times that of AugMix on ImageNet ( see Table 7 ) . Moreover , Augmax is significantly more efficient than traditional adversarial training which has ∼ 10 times training time overhead compared with AugMix . Being a stronger form of data augmentation with adversarial sample generation , AugMax leads to a significantly augmented input distribution which makes model training more challenging . This naturally motivates us to propose a novel and finer-grained normalization scheme termed Dual-Batchand-Instance Normalization ( DuBIN ) . As illustrated in Figure 3 , DuBIN adds an additional instance normalization ( IN ) in parallel to the traditional Dual Batch Normalization ( DuBN ) [ 7 , 29 ] used in adversarial training , in order to better model and disentangle the instance-wise feature heterogeneity arising from AugMax . We show that adding the instance normalization is an important knob to promote the capability of AugMax in boosting model robustness . Our framework , AugMax-DuBIN , is illustrated in Figure 2 . AugMax-DuBIN trains on clean images and achieves state-of-the-art robustness on natural corruption benchmarks [ 1 ] , and also improves model robustness against other common distribution shifts [ 30 , 2 ] . In particular , our method surpasses state-of-the-art method on CIFAR10-C , CIFAR100-C , Tiny ImageNet-C and ImageNet-C by 3.03 % , 3.49 % , 1.82 % and 0.71 % respectively . 2 Related Work . 2.1 Robustness to Distributional Shifts . Hendrycks et al . [ 1 ] pioneer the study of prediction errors exhibited by machine learning models on unseen natural corruptions and perturbations . They introduce two variants of the original ImageNet validation set to evaluate the model ’ s robustness : the ImageNet-C for input corruption robustness , and the ImageNet-P for input perturbation robustness . The former applies 15 diverse corruptions drawn from four main categories ( noise , blur , weather , and digital ) on ImageNet validation images . The latter contains perturbation sequences generated from each image . Recht et al . [ 30 , 2 ] collect CIFAR10.1 and ImageNet-V2 as a reproduction of the original CIFAR10 [ 31 ] and ImageNet [ 32 ] test sets , and find that a minute natural distribution shift caused by the minutiae in dataset collection process leads to a large drop in accuracy for a broad range of image classifiers . Geirhos et al . [ 33 ] observe that deep models trained on ImageNet are biased towards textures , and the model robustness can be improved by emphasizing more on global shape features [ 33–36 ] . Wang et al . [ 37 ] find that the benchmark test performance does not always translate to the real-world generalizability , and propose an efficient method to troubleshoot trained models using real-world unlabeled images . Hendrycks et al . [ 38 ] collect ImageNet-A with the new concept of natural adversarial examples : unmodified real-world images that falsify common machine learning models . Other representative benchmark efforts include [ 39–42 , 3 ] , all demonstrating the brittleness of deep models under various distribution shifts . Readers of interest are referred to a recent survey [ 43 ] . 2.2 Data Augmentation : Random and Adversarial . Data augmentations have been widely used to increase training set diversity and improve model generalization ability [ 44 , 21 , 33 , 22 ] . Among those methods , AugMix [ 5 ] achieves the outstanding robustness against natural corruptions by randomly mixing multiple diverse augmentations . It establishes high performance bars on many natural corruption benchmarks including CIFAR10-C , CIFAR100-C and ImageNet-C [ 1 ] . DeepAugment [ 3 ] feeds clean images to image-to-image models with randomly perturbed model weights to generate visually diverse augmented images . Combined with AugMix , it achieves state-of-the-art robustness on ImageNet-C [ 1 ] . Lately , a concurrent work MaxUp [ 45 ] generates hard training samples by selecting the worst-case weights in the MixUp [ 20 ] pipeline , to improve both model generalization ability and adversarial robustness . Our work shared a similar mindset to MaxUp , but we build on the more sophisticated AugMix pipeline and handle the resulting higher feature heterogeneity using a new normalization module . Besides , Adversarial training ( AT ) [ 6 , 46–52 ] utilizes adversarial samples as a special data augmentation method , and has also shown promise in improving model robustness against natural corruptions [ 23 ] and domain gaps [ 25 ] , usually at the cost of increasing training time [ 27 ] and degrading the standard accuracy on clean images [ 53 ] . Volpi et al . [ 25 ] and Zhao et al . [ 26 ] augment the training set with samples from a fictitious adversarial domain to improve domain adaptation performance . Xie et al . [ 7 ] show that adversarial training can improve both standard accuracy and robustness against distribution shifts , with the help of auxiliary batch normalization . Adversarial Noise Training ( ANT ) [ 8 ] uses random Gaussian noise with learned adversarial hyperparameters ( mean and variance ) as the augmentation . Gowal et al . [ 54 ] propose to generate new augmented images by adversarially composing the representations of different original images . Wang et al . [ 55 ] improve the robustness of pose estimation models by adversarially mixing different corrupted images as data augmentation . Based on DeepAugment [ 3 ] , Calian et al . [ 56 ] propose AdversarialAugment , which optimizes the parameters of image-to-image models to generate adversarially augmented images , achieving promising results against natural corruptions . Robey et al . [ 57 ] proposed model-based robust learning which generates augmented images with adversarial natural variations to improve model robustness against naturally-occurring conditions ( e.g. , snow , decolorization , shadow , and many others ) . The authors further generalized such idea to the domain generalization problem and achieved impressive results [ 58 ] . 2.3 Normalization . Batch Normalization ( BN ) [ 59 ] has been a cornerstone for training deep networks . Inspired by BN , more task-specific modifications are proposed by exploiting different normalization axes , such as Instance Normalization ( IN ) [ 60 ] , Layer Normalization ( LN ) [ 61 ] and Group Normalization ( GN ) [ 62 ] . A few latest works investigate to use multiple normalization layers instead of one BN . Zajkac et al . [ 63 ] use two separate BNs to handle the domain shift between labeled and unlabeled data in semi-supervised learning . Xie et al . [ 64 , 7 ] observe the difference between standard and adversarial feature statistics during AT , and craft dual batch normalization ( DuBN ) to disentangle the standard and adversarial feature statistic to improve both the standard accuracy and robustness . While most deep models just use one normalization type throughout the network , IBN-Net [ 29 ] unifies IN and BN in one model . It finds that when applied together , IN tends to learn features invariant to appearance ( colors , styles , etc . ) changes , while BN is essential for preserving high-level semantic contents . IBN-Net shows that combining IN and BN in an appropriate manner can improve model generalization . A similar combination was also explored in the style transfer field [ 65 ] , learning to selectively normalize only disturbing styles while preserving useful styles .

Summary: This paper proposes a novel data augmentation method AugMax which adversarially learns the parameters for integrating randomly selected augmentation operators to achieve a unification of diversity and hardness. To solve the issue of feature heterogeneity incurred by AugMax, this paper proposes DuBIN which simultaneously uses DuBN and IN instead of only DuBN. AugMax-DuBIN can improve OOD robustness while almost maintaining generalization.