Discrete Representations Strengthen Vision Transformer Robustness

1 INTRODUCTION . Despite their high performance on in-distribution test sets , deep neural networks fail to generalize under real-world distribution shifts ( Barbu et al. , 2019 ) . This gap between training and inference poses many challenges for deploying deep learning models in real-world applications where closedworld assumptions are violated . This lack of robustness can be ascribed to learned representations that are overly sensitive to minor variations in local texture and insufficiently adept at representing more robust scene and object characteristics , such as shape and texture . Vision Transformer ( ViT ) ( Dosovitskiy et al. , 2020 ) has started to rival Convolutional Neural Networks ( CNNs ) in many computer vision tasks . Recent works found that ViTs are more robust than CNNs ( Paul & Chen , 2021 ; Mao et al. , 2021b ; Bhojanapalli et al. , 2021 ) and generalize favorably on a variety of visual robustness benchmarks ( Hendrycks et al. , 2021b ; Hendrycks & Dietterich , 2019 ) . These work suggested that ViTs ’ robustness comes from the self-attention architecture that captures a globally-contextualized inductive bias than CNNs . However , though self-attention is well capable of modeling the global context , recent studies found that ImageNet-trained ViT may not take this advantage . For example , Chen et al . ( 2021b ) compared training ViT on ImageNet with and without using position embedding , and showed that spatial structure captured in position embedding contributes less than ∼3 % of ViT ’ s performance . Without using position embedding , ViT treats the image as an orderless bag-of-patches , but achieves similar performance as ViT with position embedding . This suggests that ViT unintentionally relies on local detail ( e.g. , texture ) than the shape or structure of the object . We hypothesize that this deficiency in robustness comes from the high-dimensional , individually informative , linear tokenization which allows ViT to minimize empirical risk without learning much spatial structure . In this paper , we propose a simple yet novel input layer for vision transformers , where image patches are represented by discrete tokens . To be specific , we discretize an image and represent an image patch as a discrete token or “ visual word ” in a codebook . Our key insight is that discrete tokens capture important features in a low-dimensional space ( Oord et al. , 2017 ) preserving shape and structure of the object ( see Figure 2 ) . Our approach capitalizes on this discrete representation to promote the robustness of ViT . Using discrete tokens drives ViT towards better modeling of spatial interactions between tokens , given that individual tokens no longer carry enough information to depend on . We also concatenate a low dimensional pixel token to the discrete token to compensate for the potentially missed local details encoded by discrete tokens , especially for small objects . Our approach only changes the image patch tokenizer to improve generalization and robustness , which is orthogonal to all existing approaches for robustness , and can be integrated into architectures that extend vision transformer . We call the ViT model using our Discrete representation the Dr. ViT . Our experiments and visualizations show that incorporating discrete tokens in ViT significantly improves generalization accuracy for all seven out-of-distribution ImageNet benchmarks : ImageNetRendition by up to 10 % , Stylized-ImageNet by up to 12 % , ImageNet-Sketch by up to 10 % , and ImageNet-C by up to 10 % . Our method establishes the new state-of-the-art by on four benchmarks without any dataset-specific data augmentation . Our work is the first to connect discrete representations to robustness and demonstrate consistent robustness improvements . We will release the code and models to support comparisons and future research . 2 RELATED WORK . Vision Transformer ( ViT ) ( Dosovitskiy et al. , 2020 ) , inspired by the Transformer ( Vaswani et al. , 2017 ) in NLP , is the first CNN-free architecture that achieves state-of-the-art image classification accuracy . Since its inception , numerous works have proposed improvements to the ViT architecture ( Wang et al. , 2021 ; Chu et al. , 2021 ; Liu et al. , 2021 ; d ’ Ascoli et al. , 2021 ) , objective ( Chen et al. , 2021a ) , training strategy ( Touvron et al. , 2021 ) , etc .. Given the difficulty to study all existing ViT models , this paper focuses on the classical ViT model ( Dosovitskiy et al. , 2020 ) and its recent published versions . Specifically , Steiner et al . ( 2021 ) proposed ViT-AugReg that applies stronger data augmentation and regularization to the ViT model . Tolstikhin et al . ( 2021 ) introduced MLPMixer to replace self-attention in ViT with multi-layer perceptions ( MLP ) . We select the above ViT model family ( Dosovitskiy et al. , 2020 ; Steiner et al. , 2021 ; Tolstikhin et al. , 2021 ) in our robustness study for three reasons . First , they represent both the very first and one-of-the-best vision transformers in the literature . Second , these models demonstrated competitive performance when pre-trained on sufficiently large datasets such as ImageNet-21K and JFT-300M . Finally , unlike other Transformer models , they provide architectures consisting of solely Transformer layers as well as a hybrid of CNN and Transformer layers . These properties improve our understanding of robustness for different types of network layers and datasets . Robustness . Recent works established multiple content robustness datasets to evaluate the out-ofdistribution generalization of deep models ( Barbu et al. , 2019 ; Hendrycks et al. , 2021b ; a ; Wang et al. , 2019 ; Geirhos et al. , 2019 ; Hendrycks & Dietterich , 2019 ; Recht et al. , 2019 ) . In this paper , we consider 7 ImageNet robustness benchmarks of real-world test images ( or proxies ) where deep models trained on ImageNet are shown to suffer from notable performance drop . Existing works on robustness are targeted at closing the gap in a subset of these ImageNet robustness benchmarks and were extensively verified with the CNNs . Among them , carefully-designed data augmentations ( Hendrycks et al. , 2021a ; Cubuk et al. , 2018 ; Steiner et al. , 2021 ; Mao et al. , 2021b ; a ) , model regularization ( Wang et al. , 2019 ; Huang et al. , 2020b ; Hendrycks et al. , 2019 ) , and multitask learning ( Zamir et al. , 2020 ) are effective to address the issue . More recently , a few studies ( Paul & Chen , 2021 ; Bhojanapalli et al. , 2021 ; Naseer et al. , 2021 ; Shao et al. , 2021 ) suggest that ViTs are more robust than CNNs . Existing works mainly focused on analyzing the cause of superior generalizability in the ViT model and exploring new data augmentations . As our work focuses on discrete token input , we train the models using the same data augmentation as the ViT-AugReg baseline ( Steiner et al. , 2021 ) . While tailoring data augmentation ( Mao et al. , 2021b ) may further improve our results , we leave it out of the scope of this paper . Discrete Representation was used as a visual representation prior to the deep learning revolution , such as in bag-of-visual-words model ( Sivic & Zisserman , 2003 ; Csurka et al. , 2004 ) and VLAD model ( Arandjelovic & Zisserman , 2013 ) . Recently , ( Oord et al. , 2017 ; Vahdat et al. , 2018 ) proposed neural discrete representation to encode an image as integer tokens . Recent works used discrete representation mainly for image synthesis ( Ramesh et al. , 2021 ; Esser et al. , 2021 ) . To the best of our knowledge , our work is the first to demonstrate discrete representations strengthening robustness . The closest work to ours is BEiT ( Bao et al. , 2021 ) that pretrains the ViTs to predict the masked tokens . However , the tokens are discarded after pretraining , where the ViT model can still overfit the non-robust nuisances in the pixel tokens at later finetuning stage , undermining its robustness . 3 1 2 5 4 8 0 1 0 3 5 5 3 2 9 5 Linear Projection Fine-tuned CodeBook 0 [ 0.01 , 0.2 , 0.53 , … , 0.04 ] 1 [ 0.6 , 0.22 , 0.83 , … , 0.01 ] 1024 [ 0.26 , 0.75 , 0.27 , … , 0.98 ] Position Embedding Classification Prediction Transformer Encoder Pixel Token Pixel Embedding Vector Quantized Model Discrete Token Discrete Embedding Input Image Reconstructed Image * Class Token Decode Figure 1 : Overview of the proposed ViT using discrete representations . In addition to the pixel embeddings ( orange ) , we introduce discrete tokens and embeddings ( pink ) as the input to the standard Transformer Encoder of the ViT model ( Dosovitskiy et al. , 2020 ) . 3 METHOD . 3.1 PRELIMINARY ON VISION TRANSFORMER . Vision Transformer ( Dosovitskiy et al. , 2020 ) is a pure transformer architecture that operates on a sequence of image patches . The 2D image x ∈ RH×W×C is flattened into a sequence of image patches , following the raster scan , denoted by xp ∈ RL× ( P 2·C ) , where L = H×WP 2 is the effective sequence length and P 2×C is the dimension of image patch . A learnable classification token xclass is prepended to the patch sequence , then the position embedding Epos is added to formulate the final input embedding h0 . h0 = [ xclass ; x 1 pE ; x 2 pE ; · · · ; xLpE ] +Epos , E ∈ R ( P 2·C ) ×D , Epos ∈ R ( L+1 ) ×D ( 1 ) h′ ` = MSA ( LN ( h ` −1 ) ) + h ` −1 , ` = 1 , . . . , Lf ( 2 ) h ` = MLP ( LN ( h′ ` ) ) + h ′ ` , ` = 1 , . . . , Lf ( 3 ) y = LN ( h0L ) , ( 4 ) The architecture of ViT follows that of the Transformer ( Vaswani et al. , 2017 ) , which alternates layers of multi-headed self-attention ( MSA ) and multi-layer perceptron ( MLP ) with LayerNorm ( LN ) and residual connections being applied to every block . We denote the number of blocks as Lf . This paper considers the ViT model family consisting of 4 ViT backbones : the vanilla ViT discussed above , ViT-AugReg ( Steiner et al. , 2021 ) which shares the same ViT architecture but applies stronger data augmentation and regularization , MLP-Mixer ( Tolstikhin et al. , 2021 ) which replaces self-attention in ViT with MLP , a variant called Hybrid-ViT which replaces the raw image patches in Equation 1 with the CNN features extracted by a ResNet-50 ( He et al. , 2016 ) . 3.2 ARCHITECTURE . Existing ViTs represent an image patch as a sequence of pixel tokens , which are linear projections of flattened image pixels . We propose a novel architecture modification to the input layer of the vision transformer , where an image patch xp is represented by a combination of two embeddings . As illustrated in Fig . 1 , in addition to the original pixel-wise linear projection , we discretize an image patch into an discrete token in a codebook V ∈ RK×dc , where K is the codebook size and dc is the dimension of the embedding . The discretization is achieved by a vector quantized ( VQ ) encoder pθ that produces an integer z for an image patch x as : pθ ( z = k|x ) = 1 ( k = argmin j=1 : K ‖ze ( x ) − Vj‖2 ) , ( 5 ) where ze ( x ) denotes the output of the encoder network and 1 ( · ) is the indicator function . The encoder is applied to the patch sequence xp ∈ RL× ( P 2·C ) to obtain an integer sequence zd ∈ { 1 , 2 , ... , K } L. Afterward , we use the embeddings of both discrete and pixel tokens to construct the input embedding to the ViT model . Specifically , the input embedding in Equation 1 is replaced by : h0 = [ xclass ; f ( Vz1d , x 1 pE ) ; f ( Vz2d , x 2 pE ) ; · · · ; f ( VzLd , x L pE ) ] +Epos , ( 6 ) where f is the function , embodied as a neural network layer , to combine the two embeddings . We empirically compared four network designs for f and found that the simplest concatenation works best . Note that our model only modifies the input layer of ViT ( Equation 1 ) and leaves intact the remaining layers depicted in Equation 2-4 . 4 1 2 8 9 7 5 9 2 decode t abby catt abby cat vul t ur e f l ami ngo f l ami ngo cout i nous t okens ( i mage pat ch ) di scr et e t okens Figure 2 : Comparison of pixel tokens ( top ) and the reconstructed image decoded from the discrete tokens ( bottom ) . Discrete tokens capture important shapes and structures but may lose local texture . Comparison of pixel and discrete embeddings : Pixel and discrete embeddings represent different aspects of the input image . Discrete embeddings capture important features in a low-dimension space ( Oord et al. , 2017 ) that preserves the global structure of an object but lose local details . Fig . 2 compares the original image ( top ) and the reconstructed images decoded from the discrete embeddings of our model ( bottom ) . As shown , the decoded images from discrete embeddings reasonably depict the object shape and global context . Due to the quantization , the decoder hallucinates the local textures , e.g. , in the cat ’ s eye , or the text in the “ vulture ” and “ flamingo ” images . It is worth noting that the VQ encoder/decoder is only trained on ImageNet 2012 but they can generalize to out-of-distribution images . Please see more examples in Appendix A . On the flip side , pixel embeddings capture rich details through the linear projection from raw pixels . However , given the expressive power of transformers , ViTs can spend capacity on local textures or nuance patterns that are often circumferential to robust recognition . Since humans recognize images primarily relying on the shape and semantic structure , this discrepancy to human perception undermines ViT ’ s generalization on out-of-distribution data . Our proposed model leverages the power of both embeddings to promote the interaction between modeling global and local features .

Summary: The authors have an observation that ViTs trained on ImageNet heavily depend on local features, but fail to use the global features (shape or structure). To address this issue and improve the robustness of ViTs, the authors proposed to replace the linear embedding layer by the a vector-quantized encoder. The authors claims that it can push the ViTs to learn the global information by this replacement. Experiments are conducted on ImageNet and other ImageNet variant datasets. Results show that the proposed method can improve ViTs' robustness on various benchmarks.