Understanding and Improving Robustness of Vision Transformers through Patch-based Negative Augmentation

1 INTRODUCTION . Building vision models that are robust , i.e. , that are highly accurate even on unexpected and out-ofdistribution images , is increasingly a requirement to trusting vision models and a strong benchmark for progress in the field . Recently , Vision Transformers ( ViTs , Dosovitskiy et al . ( 2021 ) ) sparked great interest in the literature , as a radically new model architecture offering significant accuracy improvements and with hope of new robustness benefits . Over the past decade , there has been extensive work on understanding the robustness of convolution-based neural architectures , as the dominant design for visual tasks ; researchers have explored adversarial robustness ( Szegedy et al. , 2013 ) , domain generalization ( Xiao et al. , 2020 ; Khani & Liang , 2021 ) , feature biases ( Brendel & Bethge , 2019 ; Geirhos et al. , 2018 ; Hermann et al. , 2020 ) . As a result , with the new promise of vision transformers , it is critical to understand their properties and in particular their robustness . Recent early studies ( Naseer et al. , 2021 ; Paul & Chen , 2021 ; Bhojanapalli et al. , 2021 ) have found ViTs be more robust than ConvNets in some scenarios , with the hypothesis that the non-local attention based interactions enabled ViTs to capture more global and semantic features . In contrast , we add to this line of research showing a different side of the challenge : we find ViTs are still vulnerable to relying on non-robust features impeding out-of-distribution performance . In this paper , we first demonstrate ViTs rely on specific non-robust features and then show how to reduce the reliance on these non-robust features , enabling improved out-of-distribution performance . To understand the robustness properties of ViTs , we start with the architectural traits of ViTs – ViTs operate on non-overlapping image patches and allow long range interaction between patches even in lower layers . It is hypothesized in recent studies ( Naseer et al. , 2021 ; Paul & Chen , 2021 ; Bhojanapalli et al. , 2021 ) that the non-local attention based interactions contribute to better robustness of ViTs than ConvNets . To study the ability of ViTs to integrate global semantics structures across patches , we design and apply patch-based image transformations , such as random patch rotation , shuffling , and background-infilling ( Figure 1 ) . Those transformations destroy the spatial relationship between patches and corrupted the global semantics , and the resultant images are often visually unrecognizable . However , we find that ViTs are surprisingly insensitive to these transformations and can make highly accurate predictions on these transformed images . This suggests that ViTs use features that survive such transformations but are generally not indicative of the semantic class to humans . Going one step further , we find that those features are useful but not robust , as ViTs trained on them achieved high in-distribution accuracy , but suffered significantly on robustness benchmarks . With this understanding of ViTs ’ reliance on non-robust features captured by patch-based transformations , we still must answer : ( a ) how can we train ViTs to not rely on such features ? and ( b ) will reducing reliance on such features meaningfully improve out-of-distribution performance and not sacrifice in-distribution accuracy ? A majority of past robust training algorithms encourage the smoothness of model predictions on augmented images with semantic preserving transformations ( Hendrycks et al. , 2020b ; Cubuk et al. , 2019 ) . However , the patch-based transformations deliberately destroy the semantic meaning and only leave non-robust features . Taking inspiration from recent research on generative modeling ( Sinha et al. , 2020 ) , we propose a family of robust training algorithms based on patch-based negative augmentations that regularize the training from relying on non-robust features surviving patch-based transformations . Through extensive evaluation on a wide set of ImageNet-based benchmarks , we find that our methods consistently improve the robustness of the trained ViTs . Furthermore , our patch-based negative augmentation can be combined with the traditional ( positive ) data augmentation to boost the performance further . With this we get a more complete picture : training models both to be insensitive to spurious changes ( as in positive augmentation ) but also to not rely on non-robust features ( as in negative augmentation ) together can meaningfully improve robustness of ViTs . Our key contributions are as follows : • Understanding Non-Robust Features in ViT : We show that ViTs heavily rely on nonrobust features surviving patch-based transformations but are not indicative of the semantic classes to humans . • Modeling : We propose a set of patch-based operations as negatively augmented views , complementary to existing works that focus on semantic-preserving ( “ positive ” ) augmentations , to regularize the training away from using these specific non-robust features ; • Improved Robustness of ViT : We show across a wide set of robustness benchmarks that our proposed negative augmented views can consistently improve ViT ’ s robustness and complementary to “ positive ” augmentation . 2 PRELIMINARIES . Vision Transformers Vision transformers ( Dosovitskiy et al. , 2021 ) are a family of architectures adapted from Transformers in natural language processing ( Vaswani et al. , 2017 ) , which di- rectly process visual tokens constructed from image patch embeddings . To construct visual tokens , ViT ( Dosovitskiy et al. , 2021 ) first splits an image into a grid of patches . Each patch is linearly projected into a hidden representation vector , and combined with a positional embedding . A learnable class token is also added . Transformers ( Vaswani et al. , 2017 ) are then directly applied on this set of visual and class token embeddings as if they are word embeddings in the original Transformer formulation . Finally , a linear projection of class token is used to calculate the class probability . Model Variants We consider ViT models pretrained on either ILSVRC-2012 ImageNet-1k , with ∼1.3 million images or ImageNet-21k , with ∼14 million images ( Russakovsky et al. , 2015 ) . All models are fine-tuned on ImageNet-1k dataset . We adopt the notations used in ( Dosovitskiy et al. , 2021 ) to denote model size and input patch size . For example , ViT-B/16 denotes the “ Base ” model variant with input patch size 16× 16 . Robustness Benchmarks To evaluate models ’ robustness , we mainly focus on three ImageNetbased robustness benchmarks , ImageNet-A ( Hendrycks et al. , 2019 ) , ImageNet-C ( Hendrycks & Dietterich , 2019 ) , ImageNet-R ( Hendrycks et al. , 2020a ) . Specifically , ImageNet-A contains challenging natural images from a distribution unlike ImageNet training distribution , which can easily fool models to make a misclassification . ImageNet-C consists of 19 types of corruptions that are frequently encountered in natural images and each corruption has 5 levels of severity . It is widely used to measure models ’ robustness under distributional shift . ImageNet-R is composed of images obtained by artistic rendition of ImageNet classes , e.g. , cartoons , and is widely used to evaluate model ’ s robustness on out-of-distribution data . 3 UNDERSTANDING ROBUSTNESS OF VISION TRANSFORMERS . Recent works ( Naseer et al. , 2021 ; Paul & Chen , 2021 ; Bhojanapalli et al. , 2021 ) have shown that vision transformers achieve better robustness compared to standard convolutional networks . One explanation for vision transformers ’ stronger robustness is that the attention mechanism can capture better global structures . To investigate if ViT has successfully taken advantage of the long range interactions between patches , we design a series of patch-based transformations which significantly destroys the global structure of images . The patch-based transformations ( see Fig . 1 ) are : • Patch-based Shuffle ( P-Shuffle ) : we randomly shuffle the input image patches to change their positions1 . • Patch-based Rotate ( P-Rotate ) : we randomly select a rotation degree from the set Ω = { 0◦ , 90◦ , 180◦ , 270◦ } and rotate each image patch independently . • Patch-based Infill ( P-Infill ) : we replace the image patches in the center region of an image with the patches on the image boundary2 . 1P-Shuffle is equivalent to shuffling the position embeddings . 2For example , given an image with size 384× 384 , input patch size is 16× 16 and replace rate 0.25 , we in total have 576 patches xi , j , where i and j denotes the row and column index and 1 ≤ i , j ≤ 24 . The patches in the center xm , n , 7 ≤ m , n ≤ 18 are replaced by the remaining patches . Each patch-based transformation is performed to a single image . We make sure the patch size of our patch-based transformation is a multiple of the input image patch of ViT so that the content within each patch is well-maintained . For P-infill , we use “ replace rate ” to denote the ratio of replaced patches in the center over the total number of patches in an image . Examples of transformed images are shown in Fig . 1 ( see Appendix F for more examples ) . In most cases , it is challenging to recognize the semantic classes after those transformations . Do ViTs rely on features not indicative of the semantic classes to humans ? To validate if ViTs behave similarly as humans on these patch-based transformed images , we evaluate ViT models ( Dosovitskiy et al. , 2021 ) on these patch-based transformed image . Specifically , we apply each patchbased transformation to ImageNet-1k validation set and report the test accuracy of each ViT on the transformed images . The test accuracy is computed by using the semantic class of the corresponding original image as the ground-truth . As shown in Figure 2 , the accuracy achieved by ViTs are significantly higher than random guessing ( 0.1 % ) . In addition , as shown in Figure 1 , ViT gives these patch-based transformed images a very high-confident prediction even when the transformation largely destroys the semantics and make the image unrecognizable by humans 3 . This strongly indicates that ViT models heavily rely on the features that survive these transformations to make a prediction . However , these features are not indicative of the semantic class to humans . Do features preserved in patch-based transformations impede robustness ? Taking one step further , we want to know if the features preserved by simple patch-based transformations , which are not indicative of the semantic class to humans , result in robustness issues . To this end , we train a vision transformer , e.g. , ViT-B/16 , on patch-based transformed images with original semantic class assigned as their ground-truth . Note that all the training images are patch-based transformed images . In this way , we force the model to fully exploit the features preserved in patch-based transformations . In addition , we reuse all the training details and hyperparameters in ( Dosovitskiy et al. , 2021 ) to make sure the “ only ” difference between our models and the baseline ViT-B/16 ( Dosovitskiy et al. , 2021 ) is the training images . Then , we test the model on ImageNet-1k validation set and three robustness benchmarks , ImageNet-A , ImageNet-C and ImageNet-R without any transformation . First , we can observe that for the baseline ViT-B/16 , compared to in-distribution accuracy , all the out-of-distribution accuracies have suffered from a significant drop ( the 4 blue bars in Figure 3 ) . This trend has been observed both for ViTs and convolution-based networks ( Paul & Chen , 2021 ; Zhai et al. , 2021 ) . Second , if we compare the accuracy between the baseline model and models trained on patch-based transformations ( i.e. , the difference between the blue bar and one of the red/green/orange bars in Figure 3 ) , we find that ViTs ’ in-distribution accuracy drops only slightly , but the robustness drop is significant when models are trained on these patch-based transformations . Take P-Shuffle as an example , the model trained on patch-based shuffled images can still achieve 79.1 % accuracy on ImageNet-1k , only 5 percentage point ( pp ) drop in in-distribution accuracy . In contrast , the accuracy drop on robustness datasets is much more significant , e.g. , 17pp on ImageNetR . The deterioration rate in robustness is close to 50 % of the baseline ViT-B/16 . This strongly 3Similar patterns are observed when ViT models are pretrained on ImageNet-1k . suggests that the features preserved in patch-based transformations are sufficient for high-accurate in-distribution prediction but are not robust under distributional shifts . Taking the above results together , we conclude that even though ViTs are shown to be more robust than ConvNets in previous studies ( Naseer et al. , 2021 ; Paul & Chen , 2021 ) , they still heavily rely on features that are not indicative of the semantic classes to humans . These features , captured by patch-based transformations , are useful but non-robust , as ViTs trained on them achieve high indistribution accuracy but suffer significantly on robustness benchmarks . With this knowledge we now ask : based on these understandings , can we train ViTs to not rely on such non-robust features ? And if we do , will it improve their robustness ?

Summary: This paper investigates the robustness of vision transformers. They first found that ViTs heavily use features that survived patch-based transformations but are generally not indicative of the semantic class to humans. However, these features are non-robust and the authors propose to use the images transformed with our patch-based operations as negatively augmented views and offer losses to regularize the training away from using non-robust features. Experiments on ImageNet show that patch-based negative augmentation consistently improves the robustness of ViTs.