PRIME: A Few Primitives Can Boost Robustness to Common Corruptions

We consider the CIFAR-10 (C-10), CIFAR-100 (C-100) [25], ImageNet-100 (IN-100) and ImageNet (IN) [11] datasets. IN-100 is a $100$-class subset of IN obtained by selecting every $10$^th class in WordNet ID order. We train a ResNet-18 [19] on C-10, C-100 and IN-100; and a ResNet-50 on IN for $100$ epochs. Following AugMix, and for a complete comparison, we also integrate the Jensen-Shannon divergence (JSD)-based consistency loss in PRIME which compels the network to learn similar representations for differently augmented versions of the same input image. Detailed training setup appears in Appendix 0.C. We evaluate our trained models on the common corrupted versions (C-10-C, C-100-C, IN-100-C, IN-C) of the aforementioned datasets. The common corruptions [21] constitute $15$ image distortions each applied with $5$ different severity levels. These corruptions can be grouped into four categories, viz. noise, blur, weather and digital.

In order to assess the effectiveness of PRIME, we evaluate its performance against C-10, C-100, IN-100 and IN common corruptions. The results are summarized in Tab. 1³³3We provide the per-corruption performance of every method in Appendix 0.H.. Amongst individual methods, PRIME yields superior results compared to those obtained by AugMix and DeepAugment alone and advances the baseline performance on the corrupted counterparts of the four datasets. As listed, PRIME pushes the corruption accuracy by $1.2\%$ and $3.3\%$ on C-10-C and C-100-C respectively over AugMix. On IN-100-C, a more complicated dataset, we observe significant improvements wherein PRIME outperforms AugMix by $10.9\%$. In fact, this increase in performance hints that our primitive transformations are actually able to cover a larger space of image corruptions, compared to the restricted set of AugMix. Interestingly, the random transformations in PRIME also lead to a $3.9\%$ boost in corruptions accuracy over DeepAugment despite the fact that DeepAugment leverages additional knowledge to augment the training data via its use of pre-trained architectures. Moreover, PRIME provides cumulative gains when combined with DeepAugment, reducing the mean corruption error (mCE) of prior art (DA+AugMix) by $2.7\%$ on IN-100-C. Lastly, we also evaluate the performance of PRIME on full IN-C. However, we do not use JSD in order to reduce computational complexity. Yet, even without the JSD loss, PRIME outperforms, in terms of corruption accuracy, both AugMix (with JSD) and DeepAugment by $6.7\%$ and $2.4\%$ respectively, while the mCE is reduced by $7.8\%$ and $2.9\%$. And last, when PRIME is combined with DeepAugment, it also surpasses the performance of DA+AugMix (with JSD), reaching a corruption accuracy of almost $60\%$ and an mCE of $51.3\%$. Note here, that, not only PRIME achieves superior robustness, but it does so efficiently. Compared to standard training on IN-100, AugMix requires 1.20x time and PRIME requires 1.27x. In contrast, DA is tedious and we do not measure its runtime since it also requires the training of two large image-to-image networks for producing augmentations, and can only be applied offline.

Recently, robustness to common corruptions has also been of significant interest in the field of unsupervised domain adaptation [2, 39]. The main difference is that, in domain adaptation, one exploits the limited access to test-time corrupted samples to adjust certain network parameters. Hence, it would be interesting to investigate the utility of PRIME under the setting of domain adaption.

To that end, we combine our method with the adaption trick of [39]. Specifically, we adjust the batch normalization (BN) statistics of our models using a few corrupted samples. Suppose $z_{s}\in\{\mu_{s}$, $\sigma_{s}\}$ are the BN mean and variance estimated from the training data, and $z_{t}\in\{\mu_{t}$, $\sigma_{t}\}$ are the corresponding statistics computed from $n$ unlabelled, corrupted test samples, then we re-estimate the BN statistics as follows.

We consider three adaptation scenarios: single sample ($n=1,N=16$), partial ($n=8,N=16$) and full ($n=400,N=0$) adaptation. Here, we do not perform parameter tuning for $N$. As shown in Tab. 2, simply correcting BN statistics using as little as $8$ corrupted samples pushes the corruption accuracy of PRIME from $71.6\%$ to $75.3\%$. In general, PRIME yields cumulative gains in combination with adaptation and has the best IN-100-C accuracy.

We want to understand how the transformations in each domain of Eq. 1 contribute to the overall robustness. To that end, we conduct an ablation study on IN-100-C by training a ResNet-18 with the max-entropy transformations of PRIME individually or in combination. As shown in Tab. 3, spectral transformations mainly help against blur, weather and digital corruptions. Spatial operations also improve on blurs, but on elastic transforms as well (digital). On the contrary, color transformations excel on noises and certain high frequency digital distortions, e.g., pixelate and JPEG artefacts, and have minor effect on weather changes. Besides, incrementally combining the transformations lead to cumulative gains e.g., spatial+color help on both noises and blurs. Yet, for obtaining the best results, the combination of all transformations is required. This means that each transformation increases the coverage over the space of possible distortions and the increase in robustness comes from their cumulative contribution.

In most data augmentation methods, besides the importance of the transformations themselves, mixing has been claimed as an essential module for increasing diversity in the training process [48, 47, 22, 44]. In our attempt to provide insights on the role of mixing in the context of common corruptions, we found out that it is capable of constructing augmented images that look perceptually similar to their corrupted counterparts. In fact, the improvements on specific corruption types observed in Tab. 3 can be largely attributed to mixing. As exemplified in Fig. 3, careful combinations of spectral transformations with the clean image introduce brightness and contrast-like artefacts that look similar to the corresponding corruptions in IN-C. Also, combining spatial transformations creates blur-like artefacts that look identical to zoom blur in IN-C. Finally, notice how mixing color transformations helps fabricate corruptions of the “noise” category. This means that the max-entropy color model of PRIME enables robustness to different types of noise without explicitly adding any during training.

Note that one of the main goals of data augmentation is to achieve maximum coverage of the space of possible distortions using a limited transformation budget, i.e., within a few training epochs. The principle of max-entropy guarantees this within each primitive, but the effect of mixing on the overall space is harder to quantify. In this regard, we can use the distance in the embedding space, $\phi$, of a SimCLRv2 [7] model as a proxy for visual similarity [49, 30]. We are interested in measuring how mixing the base transformations changes the likelihood that an augmentation scheme generates some sample during training that is visually similar to some of the common corruptions. To that end, we randomly select $N=1000$ training images $\{\bm{x}_{n}\}_{n=1}^{N}$ from IN, along with their $C=75$ ($15$ corruptions of $5$ severity levels) associated common corruptions $\{\hat{\bm{x}}_{n}^{c}\}_{c=1}^{C}$, and generate for each of the clean images another $T=100$ transformed samples $\{\tilde{\bm{x}}^{t}_{n}\}_{t=1}^{T}$ using each augmentation scheme. Moreover, for each corruption $\hat{\bm{x}}_{n}^{c}$ we find its closest neighbor $\tilde{\bm{x}}_{n}^{t}$ from the set of generated samples using the cosine distance in the embedding space. Our overall measure of fitness is

Table 4 shows the values of this measure applied to AugMix and PRIME, with and without mixing. For reference, we also report the values of the clean (no transform) images $\{\bm{x}_{n}\}_{n=1}^{N}$. More percentile scores can be found in Appendix 0.F. Clearly, mixing helps reduce the distance between the common corruptions and the augmented samples from both methods. We also observe that PRIME, even with only $100$ augmentations per image – in the order of the number of training epochs – can generate samples that are twice as close to the common corruptions as AugMix. In fact, the feature similarity between training augmentations and test corruptions was also studied in [29], with an attempt to justify the good performance of AugMix on C-10. Yet, we see that the fundamental transformations of AugMix are not enough to span a broad space guaranteeing high perceptual similarity to IN-C. The significant difference in terms of perceptual similarity in Tab. 4 between AugMix and PRIME may explain the superior performance of PRIME on IN-100-C and IN-C (cf. Tab. 1)⁴⁴4A visualization of the augmented space using PCA can be found in Appendix 0.G..

An important phenomenon observed in the literature of adversarial robustness is the so-called robustness-accuracy trade-off [16, 43, 35], where technically adversarial training [27] with smaller perturbations (typically smaller $\varepsilon$) results in models with higher standard but lower adversarial accuracy, and vice versa. In this sense, we want to understand if the strength of the image transformations introduced through data augmentations in PRIME can also cause such phenomenon in the context of robustness to common corruptions. As described in Sec. 2, each of the transformations of PRIME has a strength parameter $\sigma$, which can be seen as the analogue of $\varepsilon$ in adversarial robustness. Hence, we can easily reduce or increase the strength of the transformations by setting $\hat{\sigma}=\alpha\sigma$, where $\alpha\in\mathbb{R}^{+}$. Then, by training a network for different values of $\alpha$ we can monitor its accuracy on the clean and the corrupted datasets.

We train a ResNet-18 on C-10 and IN-100 using the setup of Sec. 4.1. For reducing complexity, we do not use the JSD loss and train for $30$ epochs. This sub-optimal setting could cause some performance drop compared to the results of Tab. 1, but we expect the overall trends in terms of accuracy and robustness to be preserved. Regarding the scaling of the parameters’ strength, for C-10 we set $\alpha\in[10^{-3},10^{2}]$ and sample $100$ values spaced evenly on a log-scale, while for IN-100 we set $\alpha\in[10^{-2},10^{2}]$ and we sample $20$ values.

The results are presented in Fig. 4. For both C-10 and IN-100, it seems that there is a sweet spot for the scale around $\alpha=0.2$ and $\alpha=1$ respectively, where the accuracy on common corruptions reaches its maximum. For $\alpha$ smaller than these values, we observe a clear trade-off between validation and robust accuracy. While the robustness to common corruptions increases, the validation accuracy decays. However, for $\alpha$ greater than the sweet-spot values, we observe that the trade-off ceases to exist since both the validation and robust accuracy present similar behaviour (slight decay). In fact, these observations indicate that robust and validation accuracies are not always positively correlated and that one might have to slightly sacrifice validation accuracy in order to achieve robustness.

Finally, we investigate the necessity of performing augmentation during training (on-line augmentation), compared to statically augmenting the dataset before training (off-line augmentation). On the one hand, on-line augmentation is useful when the dataset is huge and storing augmented versions requires a lot of memory. Besides, there are cases where offline augmentation is not feasible as it relies on pre-trained or generative models which are unavailable in certain scenarios, e.g., DeepAugment [20] or AdA [5] cannot be applied on C-100. On the other hand, off-line augmentation may be necessary to avoid the computational cost of generating augmentations during training.

To this end, for each of the C-10 and IN-100 training sets, we augment them off-line with $k=1,2,\dots,10$ i.i.d. PRIME transformed versions. Afterwards, for different values of $k$, we train a ResNet-18 on the corresponding augmented dataset and report the accuracy on the validation set and the common corruptions. For the training setup, we follow the settings of Sec. 4.1, but without JSD loss. Also, since we increase the size of the training set by $(k+1)$, we also divide the number of training epochs by the same factor, in order to keep the same overall number of gradient updates.

The performance on common corruptions is presented in Fig. 5. The first thing to notice is that, even for $k=1$, the obtained robustness to common corruptions is already quite good. In fact, for IN-100 the accuracy ($65\%$) is already better than AugMix ($60.7\%$ with JSD loss cf. Tab. 1). Regarding C-10, we observe that for $k=4$ the actual difference with respect to the on-line augmentation is almost negligible ($88.8\%$ vs. $89.3\%$), especially considering the overhead of transforming the data at every epoch. Technically, this means that augmenting C-10 with $4$ PRIME counterparts is enough for achieving good robustness to common corruptions. Finally, we also see in Fig. 5 that the corruption accuracy on IN-100 presents a very slow improvement after $k=4$. Comparing the accuracy at this point ($67.2\%$) to the one obtained with on-line augmentation and without JSD ($68.8\%$ cf. Tab. 3) we observe a gap of $1.6\%$. Hence, given the cost of on-line augmentation on such large scale datasets, simply augmenting the training with $4$ extra PRIME samples presents a good compromise for achieving competitive robustness. Nevertheless, the increase of $1.6\%$ introduced by on-line augmentation is rather significant, hinting that generating transformed samples during training might be necessary for maximizing performance. In this regard, the lower computational complexity of PRIME allows it to easily achieve this $+1.6\%$ gain through on-line augmentation, since it only requires $1.27\times$ additional training time compared to standard training, and only $1.06\times$ compared to AugMix, but with much better performance. This can be a significant advantage with respect to complex methods, like DeepAugment, that cannot be even applied on-line (require heavy pretraining).