AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

Abstract and Introduction

While Transformers have revolutionized Natural Language Processing (NLP), their application in computer vision has largely been limited to augmenting or replacing components within convolutional networks. In this work, we challenge this reliance on CNNs by demonstrating that a pure Transformer, applied directly to sequences of image patches, can achieve excellent performance on image classification tasks. Our Vision Transformer (ViT) model, when pre-trained on large datasets and then transferred to various mid-sized or small image recognition benchmarks (such as ImageNet, CIFAR-100, and VTAB), consistently matches or surpasses state-of-the-art convolutional networks. Crucially, ViT achieves these results while demanding substantially fewer computational resources for training. This success highlights that large-scale training can overcome the inductive biases inherent in CNNs, paving the way for a new paradigm in computer vision.

Vision Transformer Method

Our Vision Transformer (ViT) architecture closely follows the original Transformer design, enabling us to leverage scalable NLP Transformer implementations. To process 2D images, we first reshape an image $\mathbf{x} \in \mathbb{R}^{H \times W \times C}$ into a sequence of flattened 2D patches $\mathbf{x}_p \in \mathbb{R}^{N \times (P^2 \cdot C)}$, where $(H,W)$ is the image resolution, $C$ is the number of channels, $(P,P)$ is the patch resolution, and $N = HW/P^2$ is the number of patches. These patches are then linearly projected to a constant latent vector size $D$, forming patch embeddings.

Similar to BERT's [class] token, we prepend a learnable classification embedding to this sequence. Position embeddings are added to the patch embeddings to preserve spatial information. The resulting sequence is fed into a standard Transformer encoder, which consists of alternating layers of multiheaded self-attention (MSA) and MLP blocks. Layernorm (LN) is applied before each block, followed by residual connections. The MLP contains two layers with a GELU non-linearity. The final state of the classification token at the Transformer encoder's output serves as the image representation for classification.

Unlike CNNs, ViT has minimal image-specific inductive biases, relying on the self-attention mechanism to learn spatial relationships from data. For fine-tuning, we replace the pre-trained prediction head with a zero-initialized feedforward layer. We often fine-tune at higher resolutions, adapting the pre-trained position embeddings via 2D interpolation to accommodate the increased sequence length. We also explored a hybrid architecture where the input sequence is formed from CNN feature maps instead of raw image patches. Figure 1 provides a visual overview of our model.

Model overview. We split an image into fixed-size patches, linearly embed each of them, add position embeddings, and feed the resulting sequence of vectors to a standard Transformer encoder. In order to perform classification, we use the standard approach of adding an extra learnable 'classification token' to the sequence. The illustration of the Transformer encoder was inspired by Vaswani et al. (2017).

Experimental Setup

To thoroughly evaluate ViT, we conducted extensive experiments across various datasets and model configurations. For pre-training, we utilized ImageNet (1.3M images, 1k classes), its superset ImageNet-21k (14M images, 21k classes), and the large-scale JFT-300M dataset (303M high-resolution images, 18k classes). We carefully de-duplicated these datasets against downstream task test sets. For fine-tuning, we evaluated on popular benchmarks including ImageNet (original and ReaL labels), CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102, and the 19-task VTAB suite, which assesses low-data transfer with 1,000 training examples per task.

Our ViT model variants, summarized in Table 1, are based on BERT configurations: ViT-Base, ViT-Large, and an even larger ViT-Huge. We denote models by their size and patch size, e.g., ViT-L/16 for the 'Large' variant with 16x16 patches. Smaller patch sizes lead to longer sequence lengths and higher computational costs. As baselines, we used 'ResNet (BiT)' models, which are ResNets with Group Normalization and standardized convolutions, known for improved transfer learning. For hybrid models, we fed intermediate CNN feature maps into ViT, effectively using a patch size of one 'pixel' from the feature map.

All models, including ResNets, were pre-trained using Adam with specific hyperparameters ($\beta_1=0.9, \beta_2=0.999$, batch size 4096, weight decay 0.1) and a linear learning rate warmup and decay. Fine-tuning was performed using SGD with momentum and a batch size of 512. For top ImageNet results, we fine-tuned at higher resolutions (512 for ViT-L/16, 518 for ViT-H/14) and applied Polyak averaging. We report both fine-tuning and few-shot accuracies, with the latter obtained via regularized least-squares regression for efficient evaluation.

Experimental Results

Our experiments demonstrate that Vision Transformers achieve state-of-the-art performance, especially when pre-trained on large datasets. As shown in Table 2, our ViT-L/16 model, pre-trained on JFT-300M, consistently outperforms BiT-L across all tasks while requiring significantly less computational resources for pre-training. The larger ViT-H/14 further improves the performance, particularly on challenging datasets like ImageNet, CIFAR-100, and the VTAB suite, all with a lower pre-training cost than prior SOTA models. Figure 2 further breaks down VTAB performance, showing ViT-H/14's superiority on Natural and Structured tasks.

Breakdown of VTAB performance in Natural , Specialized , and Structured task groups.

We investigated the impact of pre-training data size. Figure 3 illustrates that while large ViT models underperform BiT ResNets when pre-trained on smaller datasets like ImageNet, they truly shine with larger datasets such as ImageNet-21k and JFT-300M, where they surpass CNNs. Figure 4 reinforces this, showing that ViT models overfit more than ResNets on smaller subsets of JFT-300M but perform better on larger subsets, indicating that the convolutional inductive bias is beneficial for limited data, but large-scale learning is sufficient and even advantageous for ViT.

Transfer to ImageNet. While large ViT models perform worse than BiT ResNets (shaded area) when pre-trained on small datasets, they shine when pre-trained on larger datasets. Similarly, larger ViT variants overtake smaller ones as the dataset grows.

Linear few-shot evaluation on ImageNet versus pre-training size. ResNets perform better with smaller pre-training datasets but plateau sooner than ViT, which performs better with larger pre-training. ViT-b is ViT-B with all hidden dimensions halved.

Our scaling study, presented in Figure 5, reveals that Vision Transformers generally dominate ResNets in terms of performance versus pre-training compute. ViT models typically require 2-4x less compute to achieve comparable performance. While hybrid models show a slight advantage at smaller computational budgets, this gap diminishes for larger models. Importantly, Vision Transformers do not appear to saturate within the tested range, suggesting further scaling potential.

Performance versus pre-training compute for different architectures: Vision Transformers, ResNets, and hybrids. Vision Transformers generally outperform ResNets with the same computational budget. Hybrids improve upon pure Transformers for smaller model sizes, but the gap vanishes for larger models.

To understand ViT's internal workings, we inspected its representations. The initial linear embedding filters (Figure 7, left) resemble plausible basis functions for image patches. The learned position embeddings (Figure 7, center) effectively encode 2D image topology, with closer patches having more similar embeddings, and even revealing row-column structures. Analyzing attention distances (Figure 7, right) shows that some attention heads integrate information globally even in early layers, while others maintain localized attention, potentially mimicking early convolutional layers. Overall, the model attends to semantically relevant image regions for classification, as exemplified in Figure 6.

Left: Filters of the initial linear embedding of RGB values of ViT-L/32. Center: Similarity of position embeddings of ViT-L/32. Tiles show the cosine similarity between the position embedding of the patch with the indicated row and column and the position embeddings of all other patches. Right: Size of attended area by head and network depth. Each dot shows the mean attention distance across images for one of 16 heads at one layer. See Appendix D.7 for details.

Representative examples of attention from the output token to the input space. See Appendix D.7 for details.

Finally, our preliminary exploration into self-supervised pre-training using masked patch prediction, akin to BERT's masked language modeling, yielded promising results. Our smaller ViT-B/16 model achieved 79.9% accuracy on ImageNet, a 2% improvement over training from scratch, though still behind supervised pre-training.

Conclusion and Outlook

We have successfully demonstrated the direct application of standard Transformers to image recognition. By treating images as sequences of patches and processing them with a Transformer encoder, we've shown that image-specific inductive biases, beyond the initial patch extraction, are not strictly necessary. When coupled with large-scale pre-training, our Vision Transformer (ViT) models match or exceed the state of the art on numerous image classification benchmarks, all while being more computationally efficient to pre-train. These encouraging initial results open several exciting avenues for future research. We aim to apply ViT to other computer vision tasks, such as object detection and segmentation, building on the promise indicated by our findings and related works. Further exploration and improvement of self-supervised pre-training methods are also crucial, as our initial experiments show a significant gap between self-supervised and large-scale supervised pre-training. Lastly, we anticipate that continued scaling of ViT models will lead to even greater performance gains.

Cite our work

@article{dosovitskiy2020image,
    title={An Image Is Worth 16x16 Words: Transformers For Image Recognition At Scale},
    author={Dosovitskiy, Alexey and Beyer, Lucas and Kolesnikov, Alexander and Weissenborn, Dirk and Zhai, Xiaohua and Unterthiner, Thomas and Dehghani, Mostafa and Minderer, Matthias and Heigold, Georg and Gelly, Sylvain and Uszkoreit, Jakob and Houlsby, Neil},
    journal={arXiv preprint arXiv:2010.11929},
    year={2020}
}