AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty

1 INTRODUCTION . Current machine learning models depend on the ability of training data to faithfully represent the data encountered during deployment . In practice , data distributions evolve ( Lipton et al. , 2018 ) , models encounter new scenarios ( Hendrycks & Gimpel , 2017 ) , and data curation procedures may capture only a narrow slice of the underlying data distribution ( Torralba & Efros , 2011 ) . Mismatches between the train and test data are commonplace , yet the study of this problem is not . As it stands , models do not robustly generalize across shifts in the data distribution . If models could identify when they are likely to be mistaken , or estimate uncertainty accurately , then the impact of such fragility might be ameliorated . Unfortunately , modern models already produce overconfident predictions when the training examples are independent and identically distributed to the test distribution . This overconfidence and miscalibration is greatly exacerbated by mismatched training and testing distributions . Small corruptions to the data distribution are enough to subvert existing classifiers , and techniques to improve corruption robustness remain few in number . Hendrycks & Dietterich ( 2019 ) show that classification error of modern models rises from 22 % on the usual ImageNet test set to 64 % on ImageNet-C , a test set consisting of various corruptions applied to ImageNet test images . Even methods which aim to explicitly quantify uncertainty , such as probabilistic and Bayesian neural networks , struggle under data shift , as recently demonstrated by Ovadia et al . ( 2019 ) . Improving performance in this setting has been difficult . One reason is that training against corruptions only encourages networks to memorize the specific corruptions seen during training and leaves models unable to generalize to new corruptions ( Vasiljevic et al. , 2016 ; Geirhos et al. , 2018 ) . Further , networks trained on translation augmentations remain highly sensitive to images shifted by a single pixel ( Gu et al. , 2019 ; Hendrycks & Dietterich , 2019 ) . Others have proposed aggressive data augmentation schemes ( Cubuk et al. , 2018 ) , though at the cost of a computational increase . Chun et al . ( 2019 ) demonstrates that many techniques may improve clean accuracy at the cost of robustness while many techniques which improve robustness harm uncertainty , and contrariwise . In all , existing techniques have considerable trade-offs . ∗Equal Contribution . †Corresponding author . CutOut MixUp CutMix AugMix In this work , we propose a technique to improve both the robustness and uncertainty estimates of classifiers under data shift . We propose AUGMIX , a method which simultaneously achieves new state-of-the-art results for robustness and uncertainty estimation while maintaining or improving accuracy on standard benchmark datasets . AUGMIX utilizes stochasticity and diverse augmentations , a Jensen-Shannon Divergence consistency loss , and a formulation to mix multiple augmented images to achieve state-of-the-art performance . On CIFAR-10 and CIFAR-100 , our method roughly halves the corruption robustness error of standard training procedures from 28.4 % to 12.4 % and 54.3 % to 37.8 % error , respectively . On ImageNet , AUGMIX also achieves state-of-the-art corruption robustness and decreases perturbation instability from 57.2 % to 37.4 % . Code is available at https : //github.com/google-research/augmix . 2 RELATED WORK Robustness under Data Shift . Geirhos et al . ( 2018 ) show that training against distortions can often fail to generalize to unseen distortions , as networks have a tendency to memorize properties of the specific training distortion . Vasiljevic et al . ( 2016 ) show training with various blur augmentations can fail to generalize to unseen blurs or blurs with different parameter settings . Hendrycks & Dietterich ( 2019 ) propose measuring generalization to unseen corruptions and provide benchmarks for doing so . Kang et al . ( 2019 ) construct an adversarial version of the aforementioned benchmark . Gilmer et al . ( 2018 ) ; Gilmer & Hendrycks ( 2019 ) argue that robustness to data shift is a pressing problem which greatly affects the reliability of real-world machine learning systems . Calibration under Data Shift . Guo et al . ( 2017 ) ; Nguyen & O ’ Connor ( 2015 ) propose metrics for determining the calibration of machine learning models . Lakshminarayanan et al . ( 2017 ) find that simply ensembling classifier predictions improves prediction calibration . Hendrycks et al . ( 2019a ) show that pre-training can also improve calibration . Ovadia et al . ( 2019 ) demonstrate that model calibration substantially deteriorates under data shift . Data Augmentation . Data augmentation can greatly improve generalization performance . For image data , random left-right flipping and cropping are commonly used He et al . ( 2015 ) . Random occlusion techniques such as Cutout can also improve accuracy on clean data ( Devries & Taylor , 2017 ; Zhong et al. , 2017 ) . Rather than occluding a portion of an image , CutMix replaces a portion of an image with a portion of a different image ( Yun et al. , 2019 ) . Mixup also uses information from two images . Rather than implanting one portion of an image inside another , Mixup produces an elementwise convex combination of two images ( Zhang et al. , 2017 ; Tokozume et al. , 2018 ) . Guo et al . ( 2019 ) show that Mixup can be improved with an adaptive mixing policy , so as to prevent manifold intrusion . Separate from these approaches are learned augmentation methods such as AutoAugment ( Cubuk et al. , 2018 ) , where a group of augmentations is tuned to optimize performance on a downstream task . Patch Gaussian augments data with Gaussian noise applied to a randomly chosen portion of an image ( Lopes et al. , 2019 ) . A popular way to make networks robust to ` p adversarial examples is with adversarial training ( Madry et al. , 2018 ) , which we use in this paper . However , this tends to increase training time by an order of magnitude and substantially degrades accuracy on non-adversarial images ( Raghunathan et al. , 2019 ) . 3 AUGMIX . AUGMIX is a data augmentation technique which improves model robustness and uncertainty estimates , and slots in easily to existing training pipelines . At a high level , AugMix is characterized by its utilization of simple augmentation operations in concert with a consistency loss . These augmentation operations are sampled stochastically and layered to produce a high diversity of augmented images . We then enforce a consistent embedding by the classifier across diverse augmentations of the same input image through the use of Jensen-Shannon divergence as a consistency loss . Mixing augmentations allows us to generate diverse transformations , which are important for inducing robustness , as a common failure mode of deep models in the arena of corruption robustness is the memorization of fixed augmentations ( Vasiljevic et al. , 2016 ; Geirhos et al. , 2018 ) . Previous methods have attempted to increase diversity by directly composing augmentation primitives in a chain , but this can cause the image to quickly degrade and drift off the data manifold , as depicted in Figure 3 . Such image degradation can be mitigated and the augmentation diversity can be maintained by mixing together the results of several augmentation chains in convex combinations . A concrete account of the algorithm is given in the pseudocode below . Algorithm AUGMIX Pseudocode 1 : Input : Model p̂ , Classification Loss L , Image xorig , Operations O = { rotate , . . . , posterize } 2 : function AugmentAndMix ( xorig , k = 3 , α = 1 ) 3 : Fill xaug with zeros 4 : Sample mixing weights ( w1 , w2 , . . . , wk ) ∼ Dirichlet ( α , α , . . . , α ) 5 : for i = 1 , . . . , k do 6 : Sample operations op1 , op2 , op3 ∼ O 7 : Compose operations with varying depth op12 = op2 ◦ op1 and op123 = op3 ◦ op2 ◦ op1 8 : Sample uniformly from one of these operations chain ∼ { op1 , op12 , op123 } 9 : xaug += wi · chain ( xorig ) . Addition is elementwise 10 : end for 11 : Sample weight m ∼ Beta ( α , α ) 12 : Interpolate with rule xaugmix = mxorig + ( 1−m ) xaug 13 : returnxaugmix 14 : end function 15 : xaugmix1 = AugmentAndMix ( xorig ) . xaugmix1 is stochastically generated 16 : xaugmix2 = AugmentAndMix ( xorig ) . xaugmix1 6= xaugmix2 17 : Loss Output : L ( p̂ ( y | xorig ) , y ) + λ Jensen-Shannon ( p̂ ( y | xorig ) ; p̂ ( y|xaugmix1 ) ; p̂ ( y|xaugmix2 ) ) Augmentations . Our method consists of mixing the results from augmentation chains or compositions of augmentation operations . We use operations from AutoAugment . Each operation is visualized in Appendix C. Crucially , we exclude operations which overlap with ImageNet-C corruptions . In particular , we remove the contrast , color , brightness , sharpness , and Cutout operations so that our set of operations and the ImageNet-C corruptions are disjoint . In turn , we do not use any image noising nor image blurring operations so that ImageNet-C corruptions are encountered only at test time . Operations such as rotate can be realized with varying severities , like 2◦ or −15◦ . For operations with varying severities , we uniformly sample the severity upon each application . Next , we randomly sample k augmentation chains , where k = 3 by default . Each augmentation chain is constructed by composing from one to three randomly selected augmentation operations . Mixing . The resulting images from these augmentation chains are combined by mixing . While we considered mixing by alpha compositing , we chose to use elementwise convex combinations for simplicity . The k-dimensional vector of convex coefficients is randomly sampled from a Dirichlet ( α , . . . , α ) distribution . Once these images are mixed , we use a “ skip connection ” to combine the result of the augmentation chain and the original image through a second random convex combination sampled from a Beta ( α , α ) distribution . The final image incorporates several sources of randomness from the choice of operations , the severity of these operations , the lengths of the augmentation chains , and the mixing weights . Jensen-Shannon Divergence Consistency Loss . We couple with this augmentation scheme a loss that enforces smoother neural network responses . Since the semantic content of an image is approximately preserved with AUGMIX , we should like the model to embed xorig , xaugmix1 , xaugmix2 similarly . Toward this end , we minimize the Jensen-Shannon divergence among the posterior distributions of the original sample xorig and its augmented variants . That is , for porig = p̂ ( y | xorig ) , paugmix1 = p̂ ( y | xaugmix1 ) , paugmix2 = p̂ ( y|xaugmix2 ) , we replace the original loss L with the loss L ( porig , y ) + λ JS ( porig ; paugmix1 ; paugmix2 ) . ( 1 ) To interpret this loss , imagine a sample from one of the three distributions porig , paugmix1 , paugmix2 . The Jensen-Shannon divergence can be understood to measure the average information that the sample reveals about the identity of the distribution from which it was sampled . This loss can be computed by first obtaining M = ( porig + paugmix1 + paugmix2 ) /3 and then computing JS ( porig ; paugmix1 ; paugmix2 ) = 1 3 ( KL [ porig‖M ] + KL [ paugmix1‖M ] + KL [ paugmix2‖M ] ) . ( 2 ) Unlike an arbitrary KL Divergence between porig and paugmix , the Jensen-Shannon divergence is upper bounded , in this case by the logarithm of the number of classes . Note that we could instead compute JS ( porig ; paugmix1 ) , though this does not perform as well . The gain of training with JS ( porig ; paugmix1 ; paugmix2 ; paugmix3 ) is marginal . The Jensen-Shannon Consistency Loss impels to model to be stable , consistent , and insensitive across to a diverse range of inputs ( Zheng et al. , 2016 ; Kannan et al. , 2018 ; Xie et al. , 2019 ) . Ablations are in Section 4.3 and Appendix A .

Summary: This paper proposes a method called AugMix, which is intended to improve model robustness to data distribution shift. AugMix appears fairly simple to implement. Several new images are created by augmenting an original image through chains of sequentially applied transformations (the "Aug" part of AugMix), then the augmented images are combined together, along with the original image, via a weighted sum (the "Mix" part of AugMix). Additionally, a Jensen-Shannon Divergence consistency loss is applied during training to encourage the model to make similar predictions for all augmented variations of a single image. This technique is shown to achieve state-of-the-art performance on standard robustness benchmarks without loss of clean test accuracy, and is also shown to improve calibration of model confidence estimates.