An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

1 INTRODUCTION . Self-attention-based architectures , in particular Transformers ( Vaswani et al. , 2017 ) , have become the model of choice in natural language processing ( NLP ) . The dominant approach is to pre-train on a large text corpus and then fine-tune on a smaller task-specific dataset ( Devlin et al. , 2019 ) . Thanks to Transformers ’ computational efficiency and scalability , it has become possible to train models of unprecedented size , with over 100B parameters ( Brown et al. , 2020 ; Lepikhin et al. , 2020 ) . With the models and datasets growing , there is still no sign of saturating performance . In computer vision , however , convolutional architectures remain dominant ( LeCun et al. , 1989 ; Krizhevsky et al. , 2012 ; He et al. , 2016 ) . Inspired by NLP successes , multiple works try combining CNN-like architectures with self-attention ( Wang et al. , 2018 ; Carion et al. , 2020 ) , some replacing the convolutions entirely ( Ramachandran et al. , 2019 ; Wang et al. , 2020a ) . The latter models , while theoretically efficient , have not yet been scaled effectively on modern hardware accelerators due to the use of specialized attention patterns . Therefore , in large-scale image recognition , classic ResNetlike architectures are still state of the art ( Mahajan et al. , 2018 ; Xie et al. , 2020 ; Kolesnikov et al. , 2020 ) . Inspired by the Transformer scaling successes in NLP , we experiment with applying a standard Transformer directly to images , with the fewest possible modifications . To do so , we split an image into patches and provide the sequence of linear embeddings of these patches as an input to a Transformer . Image patches are treated the same way as tokens ( words ) in an NLP application . We train the model on image classification in supervised fashion . When trained on mid-sized datasets such as ImageNet without strong regularization , these models yield modest accuracies of a few percentage points below ResNets of comparable size . This seemingly discouraging outcome may be expected : Transformers lack some of the inductive biases 1Fine-tuning code and pre-trained models are available at https : //github.com/ google-research/vision_transformer inherent to CNNs , such as translation equivariance and locality , and therefore do not generalize well when trained on insufficient amounts of data . However , the picture changes if the models are trained on larger datasets ( 14M-300M images ) . We find that large scale training trumps inductive bias . Our Vision Transformer ( ViT ) attains excellent results when pre-trained at sufficient scale and transferred to tasks with fewer datapoints . When pre-trained on the public ImageNet-21k dataset or the in-house JFT-300M dataset , ViT approaches or beats state of the art on multiple image recognition benchmarks . In particular , the best model reaches the accuracy of 88.55 % on ImageNet , 90.72 % on ImageNet-ReaL , 94.55 % on CIFAR-100 , and 77.63 % on the VTAB suite of 19 tasks . 2 RELATED WORK . Transformers were proposed by Vaswani et al . ( 2017 ) for machine translation , and have since become the state of the art method in many NLP tasks . Large Transformer-based models are often pre-trained on large corpora and then fine-tuned for the task at hand : BERT ( Devlin et al. , 2019 ) uses a denoising self-supervised pre-training task , while the GPT line of work uses language modeling as its pre-training task ( Radford et al. , 2018 ; 2019 ; Brown et al. , 2020 ) . Naive application of self-attention to images would require that each pixel attends to every other pixel . With quadratic cost in the number of pixels , this does not scale to realistic input sizes . Thus , to apply Transformers in the context of image processing , several approximations have been tried in the past . Parmar et al . ( 2018 ) applied the self-attention only in local neighborhoods for each query pixel instead of globally . Such local multi-head dot-product self attention blocks can completely replace convolutions ( Hu et al. , 2019 ; Ramachandran et al. , 2019 ; Zhao et al. , 2020 ) . In a different line of work , Sparse Transformers ( Child et al. , 2019 ) employ scalable approximations to global selfattention in order to be applicable to images . An alternative way to scale attention is to apply it in blocks of varying sizes ( Weissenborn et al. , 2019 ) , in the extreme case only along individual axes ( Ho et al. , 2019 ; Wang et al. , 2020a ) . Many of these specialized attention architectures demonstrate promising results on computer vision tasks , but require complex engineering to be implemented efficiently on hardware accelerators . Most related to ours is the model of Cordonnier et al . ( 2020 ) , which extracts patches of size 2 × 2 from the input image and applies full self-attention on top . This model is very similar to ViT , but our work goes further to demonstrate that large scale pre-training makes vanilla transformers competitive with ( or even better than ) state-of-the-art CNNs . Moreover , Cordonnier et al . ( 2020 ) use a small patch size of 2 × 2 pixels , which makes the model applicable only to small-resolution images , while we handle medium-resolution images as well . There has also been a lot of interest in combining convolutional neural networks ( CNNs ) with forms of self-attention , e.g . by augmenting feature maps for image classification ( Bello et al. , 2019 ) or by further processing the output of a CNN using self-attention , e.g . for object detection ( Hu et al. , 2018 ; Carion et al. , 2020 ) , video processing ( Wang et al. , 2018 ; Sun et al. , 2019 ) , image classification ( Wu et al. , 2020 ) , unsupervised object discovery ( Locatello et al. , 2020 ) , or unified text-vision tasks ( Chen et al. , 2020c ; Lu et al. , 2019 ; Li et al. , 2019 ) . Another recent related model is image GPT ( iGPT ) ( Chen et al. , 2020a ) , which applies Transformers to image pixels after reducing image resolution and color space . The model is trained in an unsupervised fashion as a generative model , and the resulting representation can then be fine-tuned or probed linearly for classification performance , achieving a maximal accuracy of 72 % on ImageNet . Our work adds to the increasing collection of papers that explore image recognition at larger scales than the standard ImageNet dataset . The use of additional data sources allows to achieve state-ofthe-art results on standard benchmarks ( Mahajan et al. , 2018 ; Touvron et al. , 2019 ; Xie et al. , 2020 ) . Moreover , Sun et al . ( 2017 ) study how CNN performance scales with dataset size , and Kolesnikov et al . ( 2020 ) ; Djolonga et al . ( 2020 ) perform an empirical exploration of CNN transfer learning from large scale datasets such as ImageNet-21k and JFT-300M . We focus on these two latter datasets as well , but train Transformers instead of ResNet-based models used in prior works . 3 METHOD . In model design we follow the original Transformer ( Vaswani et al. , 2017 ) as closely as possible . An advantage of this intentionally simple setup is that scalable NLP Transformer architectures – and their efficient implementations – can be used almost out of the box . 3.1 VISION TRANSFORMER ( VIT ) . An overview of the model is depicted in Figure 1 . The standard Transformer receives as input a 1D sequence of token embeddings . To handle 2D images , we reshape the image x ∈ RH×W×C into a sequence of flattened 2D patches xp ∈ RN× ( P 2·C ) , where ( H , W ) is the resolution of the original image , C is the number of channels , ( P , P ) is the resolution of each image patch , andN = HW/P 2 is the resulting number of patches , which also serves as the effective input sequence length for the Transformer . The Transformer uses constant latent vector size D through all of its layers , so we flatten the patches and map to D dimensions with a trainable linear projection ( Eq . 1 ) . We refer to the output of this projection as the patch embeddings . Similar to BERT ’ s [ class ] token , we prepend a learnable embedding to the sequence of embedded patches ( z00 = xclass ) , whose state at the output of the Transformer encoder ( z 0 L ) serves as the image representation y ( Eq . 4 ) . Both during pre-training and fine-tuning , a classification head is attached to z0L . The classification head is implemented by a MLP with one hidden layer at pre-training time and by a single linear layer at fine-tuning time . Position embeddings are added to the patch embeddings to retain positional information . We use standard learnable 1D position embeddings , since we have not observed significant performance gains from using more advanced 2D-aware position embeddings ( Appendix D.3 ) . The resulting sequence of embedding vectors serves as input to the encoder . The Transformer encoder ( Vaswani et al. , 2017 ) consists of alternating layers of multiheaded selfattention ( MSA , see Appendix A ) and MLP blocks ( Eq . 2 , 3 ) . Layernorm ( LN ) is applied before every block , and residual connections after every block ( Wang et al. , 2019 ; Baevski & Auli , 2019 ) . The MLP contains two layers with a GELU non-linearity . z0 = [ xclass ; x 1 pE ; x 2 pE ; · · · ; xNp E ] +Epos , E ∈ R ( P 2·C ) ×D , Epos ∈ R ( N+1 ) ×D ( 1 ) z′ ` = MSA ( LN ( z ` −1 ) ) + z ` −1 , ` = 1 . . . L ( 2 ) z ` = MLP ( LN ( z ′ ` ) ) + z ′ ` , ` = 1 . . . L ( 3 ) y = LN ( z0L ) ( 4 ) Inductive bias . We note that Vision Transformer has much less image-specific inductive bias than CNNs . In CNNs , locality , two-dimensional neighborhood structure , and translation equivariance are baked into each layer throughout the whole model . In ViT , only MLP layers are local and translationally equivariant , while the self-attention layers are global . The two-dimensional neighborhood structure is used very sparingly : in the beginning of the model by cutting the image into patches and at fine-tuning time for adjusting the position embeddings for images of different resolution ( as described below ) . Other than that , the position embeddings at initialization time carry no information about the 2D positions of the patches and all spatial relations between the patches have to be learned from scratch . Hybrid Architecture . As an alternative to raw image patches , the input sequence can be formed from feature maps of a CNN ( LeCun et al. , 1989 ) . In this hybrid model , the patch embedding projection E ( Eq . 1 ) is applied to patches extracted from a CNN feature map . As a special case , the patches can have spatial size 1x1 , which means that the input sequence is obtained by simply flattening the spatial dimensions of the feature map and projecting to the Transformer dimension . The classification input embedding and position embeddings are added as described above .

Summary: This paper introduces a Transformer-based image recognition model that is fully built on the Transformer layers (multi-head self-attention + point-wise MLP) without any standard convolution layers. Basically, it splits an image into patches and takes as input the set of linear embeddings of the patches and their positions. For classification, a learnable class token is added to the input and a classification head (MLP) is attached to the class token output of the final Transformer layer. Extensive experiments of transfer learning show that when pretrained on a sufficiently large dataset (100~300M), the proposed Vision Transformer can outperform state-of-the-art convolutional networks with less training cost as well as less number of parameters.