A UG M IX : A S IMPLE M ETHOD TO I MPROVE R OBUSTNESS AND U NCERTAINTY UNDER D ATA S HIFT
Anonymous authors
Paper under double-blind review
A BSTRACT
Modern deep neural networks can achieve high accuracy when the training distribution and test distribution are identically distributed, but this assumption is frequently violated in practice.
When the train and test distributions are mismatched, accuracy can plummet.
Currently there are few techniques that improve robustness to data shift.
In this work, we propose a technique to improve the robustness and uncertainty estimates of image classiﬁers.
We propose A UG M IX , a data processing technique that is simple to implement, adds limited computational overhead, and helps models withstand data shift.
A UG M IX signiﬁcantly improves robustness and uncertainty measures on challenging image classiﬁcation benchmarks, closing the gap between previous methods and the best possible performance in some cases by more than half.
I NTRODUCTION
Current machine learning models depend on the ability of training data to faithfully represent the data encountered during deployment.
In practice, data distributions evolve, models encounter new scenarios, and data curation procedures may capture only a narrow slice of the underlying data distribution.
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 overconﬁdent predictions when the training examples are independent and identically distributed to the test distribution.
This overconﬁdence and miscalibration is greatly exacerbated by mismatched training and testing distributions.
Small shifts in the data distribution are enough to subvert existing classiﬁers, and techniques to improve performance under data shift remain few in number.
For instance, Recht et al.
(2019) demonstrate that small changes in the dataset collection process for ImageNet can lead to signiﬁcant declines in accuracy.
Hendrycks & Dietterich (2019) show that classiﬁcation 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.
Improving performance in this setting has been difﬁcult.
One reason is that training against corruptions only encourages networks to memorize the speciﬁc 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 substantial increase in required computation.
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.
In this work, we propose a technique to improve both the robustness and uncertainty estimates of classiﬁers under data shift.
We propose A UG M IX , 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.
A UG M IX 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
[Figure]
the corruption robustness error of standard training procedures from 28.4% to 12.4% and 54.3% to 37.8% error, respectively.
On ImageNet, A UG M IX also achieves state-of-the-art corruption robustness and decreases perturbation instability from 57.2% to 37.4%.
R ELATED W ORK
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 speciﬁc 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 perturbations 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.
[Figure]
Calibration under Data Shift.
Guo et al.
; Nguyen & O’Connor (2015) propose metrics for determining the calibration of machine learning models.
Lakshminarayanan et al.
(2017) ﬁnd that simply ensembling classiﬁer 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 ﬂipping and cropping are commonly used He et al.
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.
However, determining this group of augmentations can be remarkably computationally expensive.
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 (cid:96) 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).
A UG M IX
A UG M IX 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 classiﬁer 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 ﬁxed 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 A UG M IX Pseudocode
1: Input: Model ˆ p , Classiﬁcation Loss L , Image x orig , Operations O = { rotate , . .
, posterize } 2: function AugmentAndMix( x orig , k = 3 , α = 1 , β = 1 ) 3:
Fill x aug with zeros 4:
Sample mixing weights ( w 1 , w 2 , .
. . , w k ) ∼ Dirichlet ( α, α, . . . , α ) 5: for i = 1 , . .
, k do 6:
Sample operations op 1 , op 2 , op 3 ∼ O 7:
Compose operations with varying depth op 12 = op 2 ◦ op 1 and op 123 = op 3 ◦ op 2 ◦ op8:
Sample uniformly from one of these operations chain ∼ { op 1 , op 12 , op 123 } 9: x aug += w i · chain ( x orig )
(cid:46) Addition is elementwise 10: end for 11:
Sample weight m ∼ Beta ( β, β ) 12:
Interpolate with rule x augmix = mx orig + (1 − m ) x aug 13: return x augmix 14: end function 15: x augmix1 =
AugmentAndMix ( x orig ) (cid:46) x augmix1 is stochastically generated 16: x augmix2 =
AugmentAndMix ( x orig ) (cid:46) x
augmix1 (cid:54) =
x augmix 17: Loss Output: L (ˆ p ( y | x orig ) , y )
+ λ Jensen-Shannon (ˆ p ( y | x orig ); ˆ p ( y | x augmix1 ); ˆ p ( y | x augmix2 ))
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 B. 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 coefﬁcients 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 ﬁnal 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 A UG M IX , we should like the model to embed x orig , x augmix1 , x augmix similarly.
Toward this end, we minimize the Jensen-Shannon divergence among the posterior distributions of the original sample x orig and its augmented variants.
That is, for p orig =
ˆ p ( y | x orig ) , p augmix
 = ˆ p
[Equation] y [Equation] augmix1 ) , p augmix2 =
ˆ p ( y | x augmix2 ) , we replace the original loss L with the loss
[Equation]
To interpret this loss, imagine a sample from one of the three distributions p orig , p augmix1 , p augmix2 .
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 ﬁrst obtaining M
=
( p orig + p augmix
+ p augmix2 ) / 3 and then computing
[Equation] p
[Equation]
Unlike an arbitrary KL Divergence between p orig and p augmix , 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 ( p orig ; p augmix1 ) , though this does not perform as well.
The gain of training with
JS ( p orig ; p augmix1 ; p augmix2 ; p augmix3 ) 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.
E XPERIMENTS
Datasets.
The two CIFAR (Krizhevsky & Hinton, 2009) datasets contain small 32 × 32 × 3 color natural images, both with 50,000 training images and 10,000 testing images.
CIFAR-10 has
In order to measure a model’s resilience to data shift, we evaluate on the CIFAR-10-C , CIFAR-100C , and ImageNet-C datasets (Hendrycks & Dietterich, 2019).
These datasets are constructed by corrupting the original CIFAR and ImageNet test sets.
For each dataset, there are a total of 15 noise, blur, weather, and digital corruption types, each appearing at 5 severity levels or intensities.
Since these datasets are used to measure network behavior under data shift, we take care not to introduce these 15 corruptions into the training procedure.
The CIFAR-10-P , CIFAR-100-P , and ImageNet-P datasets also modify the original CIFAR and
ImageNet datasets.
These datasets contain smaller perturbations than CIFAR-C and are used to measure the classiﬁer’s prediction stability.
Each example in these datasets is a video.
For instance, a video with the brightness perturbation shows an image getting progressively brighter over time.
We should like the network not to give inconsistent or volatile predictions between frames of the video as the brightness increases.
Thus these datasets enable the measurement of the “jaggedness” (Azulay &
Weiss, 2018) of a network’s prediction stream.
Metrics.
The Clean Error is the usual classiﬁcation error on the clean or uncorrupted test data.
In our experiments, corrupted test data appears at ﬁve different intensities or severity levels 1 ≤ s ≤ 5 .
For a given corruption c , the error rate at corruption severity s is E c,s .
We can compute the average error across these severities to create the unnormalized corruption error uCE c
= (cid:80) 5 s =1 E c,s .
On CIFAR-10C and CIFAR-100-C we average these values over all 15 corruptions.
Meanwhile, on ImageNet we follow the convention of normalizing the corruption error by the corruption error of AlexNet (Krizhevsky et al., 2012).
We compute CE c = (cid:80) 5 s =1 E c,s / (cid:80) 5 s =1 E AlexNet c,s .
The average of the 15 corruption errors CE Gaussian Noise , CE Shot Noise , . .
CE Pixelate , CE JPEG gives us the Mean Corruption Error (mCE) .
Perturbation robustness is not measured by accuracy but whether video frame predictions match.
Consequently we compute what is called the ﬂip probability .
Concretely, for videos such as those with steadily increasing brightness, we determine the probability that two adjacent frames, or two frames with slightly different brightness levels, have “ﬂipped” or mismatched predictions.
There are 10 different perturbation types, and the mean across these is the mean Flip Probability (mFP) .
As with
ImageNet-C, we can normalize by AlexNet’s ﬂip probabilities and obtain the mean Flip Rate (mFR) .
In order to assess a model’s uncertainty estimates, we measure its miscalibration.
Classiﬁers capable of reliably forecasting their accuracy are considered “calibrated.”
For instance, a calibrated classiﬁer should be correct 70% of the time on examples to which it assigns 70% conﬁdence.
Let the classiﬁer’s conﬁdence that its prediction ˆ
Y is correct be written C .
Then the idealized RMS Calibration Error is (cid:113)
E C
[( P ( Y =
ˆ Y | C = c ) − c ) 2 ] , which is the squared difference between the accuracy at a given conﬁdence level and actual the conﬁdence level.
In Appendix C, we show how to empirically estimate this quantity and calculate the Brier Score.
4. CIFAR-10 AND CIFAR-
Training Setup.
In the following experiments we show that A UG M IX endows robustness to various architectures including an All Convolutional Network, a DenseNet-BC ( k = 12 , d = 100 ), a 40-2 Wide ResNet, and a ResNeXt-29 ( 32 × 4 ) (Salimans & Kingma, 2016; Huang et al., 2017;
Springenberg et al., 2014; Zagoruyko & Komodakis, 2016; Xie et al., 2016).
All networks use an initial learning rate of 0 . 1 which decays following a cosine learning rate (Loshchilov & Hutter, 2016).
All input images are pre-processed with standard random left-right ﬂipping and cropping prior to any augmentations.
We do not change A UG M IX parameters across CIFAR-10 and CIFARexperiments for consistency.
The All Convolutional Network and Wide ResNet train for 100 epochs, and the DenseNet and ResNeXt require 200 epochs for convergence.
We optimize with stochastic gradient descent using Nesterov momentum and with a weight decay of 0 . 0005 .
Results. Simply mixing random augmentations and using the Jensen-Shannon loss substantially improves robustness and uncertainty estimates. Compared to the “Standard” data augmentation baseline ResNeXt on CIFAR-10-C, A UG M IX achieves 16.6% lower absolute corruption error as shown in Figure 5. In addition to surpassing numerous other data augmentation techniques, Tabledemonstrates that these gains directly transfer across architectures and on CIFAR-100-C with zero
[Figure] [Figure]
Standard
[Table]
CIFAR-
[Table]
AllConvNet
[Table]
[Table]
DenseNet
[Table]
[Table]
[Table]
[Table]
WideResNet
[Table]
[Table]
[Table] [Table] [Table] ResNeXt
[Table]
[Table]
[Table]
[Table]
[Table]
CIFAR-100-
[Table]
AllConvNet
[Table]
DenseNet
[Table]
WideResNet
[Table]
corruption robustness that exceeds the previous state of the art.
additional tuning.
Crucially, the robustness gains do not only exist when measured in aggregate.
Figure 10 shows that A UG M IX improves corruption robustness across every individual corruption and severity level.
Our method additionally achieves the lowest mFP on CIFAR-10-P across three different models all while maintaining accuracy on clean CIFAR-10 Table 6.
Finally, we demonstrate that A UG M IX improves the RMS calibration error on CIFAR-10 and CIFAR-10-C, as shown in
Table 5.
Expanded CIFAR-10-P and calibration results are in Appendix D.
4. I MAGE N ET
Baselines.
To demonstrate the utility of A UG M IX on ImageNet, we compare to many techniques designed for large-scale images.
While techniques such as Cutout (Devries & Taylor, 2017) have not been demonstrated to help on the ImageNet scale, and while few have had success training adversarially robust models on ImageNet (Engstrom et al., 2018), other techniques such as Stylized
ImageNet have been demonstrated to help on ImageNet-C. Patch Uniform (Lopes et al., 2019) is similar to Cutout except that randomly chosen regions of the image are injected with uniform noise; the original paper uses Gaussian noise, but that appears in the ImageNet-C test set so we use uniform noise.
Next, AutoAugment (Cubuk et al., 2018) searches over data augmentation policies to ﬁnd a high-performing data augmentation policy.
We denote AutoAugment results with AutoAugment* since we remove augmentation operations that overlap with ImageNet-C corruptions, as with A UG M IX .
We also test with Random AutoAugment, an augmentation scheme where each image has a randomly sampled augmentation policy using AutoAugment operations.
In contrast to AutoAugment, Random AutoAugment and A UG M IX require far less computation and provide more augmentation variety, which can offset their lack of optimization.
MaxBlur Pooling (Zhang, 2019) is a recently proposed architectural modiﬁcation which smooths the results of pooling.
Now, Stylized ImageNet (SIN) is a technique where models are trained with the original ImageNet images and also
ImageNet images with style transfer applied.
Whereas the original Stylized ImageNet technique pretrains on ImageNet-C and performs style transfer with a content loss coefﬁcient of 0 and a style loss coefﬁcient of 1 , we ﬁnd that using 0 .
5 content and style loss coefﬁcients decreases the mCE
[Figure]
M e a n
F li p P r o b a b ili t y ( % )
Robustness and CIFAR-10-P Data
[Figure]
R M S
C a li b r a t
i
o n E rr o
[Figure]
Calibration on
[Figure] by 0.6%.
Later, we show that SIN and A UG M IX can be combined.
All models were trained from scratch, except MaxBlur Pooling models which had numerous pre-trained models available.
Training Setup.
Methods are trained with ResNet-50 and we follow the standard training scheme of Goyal et al.
(2017), in which we linearly scale the learning rate with the batch size, and use a learning rate warm-up for the ﬁrst 5 epochs, and AutoAugment and A UG M IX train for 180 epochs.
All input images are pre-processed with standard random left-right ﬂipping and cropping prior to any augmentations. Noise Blur
Weather
[Figure]
[Table]
Results.
We demonstrate substantial improvements across the board in terms of corruption robustness, perturbation robustness, calibration, and clean accuracy due to A UG M IX .
On ImageNet-C our method achieves 68.4% mCE as shown in Table 2, down from the baseline 80.6% mCE.
Additionally, we note that A UG M IX allows straightforward stacking with other methods such as SIN to achieve an even lower corruption error of 64.1% mCE.
Across increasing severities of corruptions, our method also produces much more calibrated predictions measured by both the Brier Score and RMS Calibration
Error as shown in Figure 7.
As shown in Table 3, A UG M IX also achieves a state-of-the art result on
ImageNet-P at with an mFR of 37.4%, down from 57.2%.
We demonstrate that scaling up A UG M IX from CIFAR to ImageNet also leads to state-of-the-art results in robustness and uncertainty estimation.
4. A BLATIONS
We locate the utility of A UG M IX in three factors: training set diversity, our Jensen-Shannon divergence consistency loss, and mixing.
Improving training set diversity via increased variety of augmentations can greatly improve robustness.
For instance, augmenting each example with a
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Table]
[Figure]
randomly sampled augmentation chain decreases the error rate of Wide ResNet on CIFAR-10-C from 26.9% to 17.0% Table 4.
Adding in the Jensen-Shannon divergence consistency loss drops error rate further to 14.7%.
Mixing random augmentations without the Jenson-Shannon divergence loss gives us an error rate of 13 .
Finally, re-introducing the Jensen-Shannon divergence gives us A UG M IX with an error rate of 11 . 2% .
Note that adding even more mixing is not necessarily beneﬁcial.
For instance, applying A UG M IX on top of Mixup increases the error rate to 13 .
3% , possibly due to an increased chance of manifold intrusion (Guo et al., 2019).
Hence A UG M IX ’s careful combination of variety, consistency loss, and mixing explain its performance.
[Table]
C ONCLUSION
A UG M IX is a data processing technique which mixes randomly generated augmentations and uses a Jensen-Shannon loss to enforce consistency.
Our technique is simple to implement, adds little computational overhead, and substantially improves robustness and uncertainty.
Speciﬁcally, we found that A UG M IX obtains state-of-the-art performance on CIFAR-10/100-C, ImageNet-C, CIFAR10/100-P, ImageNet-P, and it also achieves state-of-the-art calibration under data shift.
We hope this will enable more reliable models, a necessity for models deployed in safety-critical environments.
A H YPERPARAMETER A BLATIONS
In this section we demonstrate that A UG M IX ’s hyperparameters are not highly sensitive, so
that
A UG M IX performs reliably without careful tuning.
For this set of experiments, the baseline A UG M IX model trains for 90 epochs, has a mixing coefﬁcient of 0 .
5 where α and β are equal, has 3 examples per Jensen-Shannon Divergence (1 clean image, 2 augmented images), has a chain depth stochastically varying from 1 to 3, and has k = 3 augmentation chains.
Figure 8 shows that the performance of various A UG M IX models with different hyperparameters.
Under these hyperparameter changes, the mCE does not change substantially.
[Figure]
B A UGMENTATION O PERATIONS
The augmentation operations we use for A UG M IX are shown in Figure 9.
C C ALIBRATION M ETRICS
Due to the ﬁnite size of empirical test sets, the RMS Calibration Error must be estimated by partitioning all n test set examples into b contiguous bins { B 1 , B 2 , . . .
, B b } ordered by prediction conﬁdence.
In this work we use bins which contain 100 predictions, so that we adaptively partition conﬁdence scores on the interval [0 , 1] (Nguyen & O’Connor, 2015; Hendrycks et al., 2019b).
Other works partition the interval [0 , 1] with 15 bins of uniform length (Guo et al., 2017).
With these b bins, we estimate the RMS Calibration Error empirically with the formula
[Equation]
[Equation]
This is separate from classiﬁcation error because a random classiﬁer with an approximately uniform posterior distribution is approximately calibrated.
Also note that adding the “reﬁnement” E C [( P ( Y =
ˆ Y | C = c )(1 − ( P
( Y = ˆ Y | C = c ))]
to the square of the RMS Calibration Error gives us the Brier
Score (Nguyen & O’Connor, 2015).
D A DDITIONAL RESULTS
In the tables and ﬁgures below, we include various additional results.
[Figure]
I
[Figure]
t
i
[Figure]
r
[Figure]
these corruptions during training
[Table]
[Table]
[Table]
While adversarial training performs well on CIFAR-10-P, it induces a substantial drop in accuracy on clean CIFAR-10 where A UG M IX does not.