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

Reading guide

Problem and motivation

Convolutional architectures (e.g., ResNet) have long been the foundation of modern computer vision due to hardcoded inductive biases like translation equivariance and two-dimensional locality. While Transformers became the standard in natural language processing (NLP) due to computational efficiency and massive scalability, prior attempts to bring self-attention to vision typically coupled attention with CNN backbones or used specialized local/axial attention patterns that were difficult to scale efficiently on hardware accelerators. This paper investigates whether a pure Transformer architecture, applied directly to sequences of image patches with minimal modifications, can rival or outperform state-of-the-art CNNs.

Key ideas

• Images as Sequences of Words: An image $\mathbf{x}\in {\mathbb{R}}^{H\times W\times C}$ is split into non-overlapping patches of size $P\times P$. Flattened patches are treated identically to NLP tokens.
• Minimal Inductive Bias: Unlike CNNs where spatial locality and translation equivariance are present in every layer, ViT uses global self-attention across patches. The 2D structure is only referenced when extracting patches and when interpolating position embeddings for higher-resolution fine-tuning.
• Scale Trumps Inductive Bias: Without strong inductive bias, ViT underperforms ResNets on small-to-medium datasets (like standard ImageNet without heavy regularization). However, when pre-trained on large datasets (ImageNet-21k with 14M images or JFT-300M with 303M images), ViT matches or exceeds CNN performance while consuming substantially fewer training FLOPs.

Method

1. Patch Embedding and Input Construction: Given patch size $(P,P)$, an image yields $N=HW/P^2$ patches. Each flattened patch ${\mathbf{x}}_p^i\in {\mathbb{R}}^{P^2\cdot C}$ is linearly projected by matrix $\mathbf{E}\in {\mathbb{R}}^{(P^2\cdot C)\times D}$ to dimension $D$. A learnable classification token ${\mathbf{x}}_{\text{class}}$ is prepended, and 1D learnable position embeddings ${\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}$ are added: $${\mathbf{z}}_0=[{\mathbf{x}}_{\text{class}};{\mathbf{x}}_p^1\mathbf{E};{\mathbf{x}}_p^2\mathbf{E};\dots ;{\mathbf{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{pos}$$
2. Transformer Encoder Layers: The sequence passes through $L$ identical blocks alternating Multihead Self-Attention (MSA) and MLP blocks, using pre-LayerNorm (LN) and residual connections: $${\mathbf{z}}_{\ell }^{\prime }=\mathrm{MSA}(\mathrm{LN}({\mathbf{z}}_{\ell -1}))+{\mathbf{z}}_{\ell -1},\ell =1\dots L$$ $${\mathbf{z}}_{\ell }=\mathrm{MLP}(\mathrm{LN}({\mathbf{z}}_{\ell }^{\prime }))+{\mathbf{z}}_{\ell }^{\prime },\ell =1\dots L$$
3. Representation and Head: The representation is derived from the first output token: $$\mathbf{y}=\mathrm{LN}({\mathbf{z}}_L^0)$$ A small MLP (with one hidden layer) serves as the classification head during pre-training, which is replaced by a single linear projection ($D\times K$) at fine-tuning.
4. Higher Resolution Fine-Tuning: When fine-tuning at higher image resolutions (e.g., $384\times 384$ or $512\times 512$), patch size $P$ remains fixed, increasing sequence length $N$. The pre-trained 1D position embeddings are adjusted via 2D bicubic interpolation according to their original spatial grid locations.

Results

• Comparison to SOTA:
  • ViT-H/14 pre-trained on JFT-300M achieved 88.55% on ImageNet, 90.72% on ImageNet ReaL, 94.55% on CIFAR-100, and 77.63% on the VTAB 19-task suite, taking 2.5k TPUv3-core-days.
  • ViT-L/16 pre-trained on JFT-300M achieved 87.76% on ImageNet with only 0.68k TPUv3-core-days, outperforming BiT-L (ResNet152x4, 87.54%, 9.9k TPUv3-core-days) and approaching Noisy Student EfficientNet-L2 (88.4/88.5%, 12.3k TPUv3-core-days).
  • Public dataset pre-training: ViT-L/16 pre-trained on ImageNet-21k reached 85.30% on ImageNet.
• Compute vs. Performance Trade-off: Across scaling experiments on JFT-300M, pure ViT dominated ResNets on the Pareto frontier, requiring $2\times$ to $4\times$ less compute to attain equivalent accuracy.
• Internal Representations: Lower layers contain attention heads that attend across the entire image alongside heads with small, localized attention distances. Attention distance generally increases monotonically with depth.

Limitations and open questions

• Data Hunger: ViT generalizes poorly on small datasets without either extensive pre-training data (ImageNet-21k or JFT-300M) or heavy regularization.
• Self-Supervised Gap: Initial experiments with masked patch prediction reached 79.9% on ImageNet (ViT-B/16), leaving a 4% performance gap compared to supervised pre-training.
• Application to Dense Prediction: The initial work focuses exclusively on classification; adapting pure patch transformers to dense spatial tasks such as object detection and semantic segmentation remains an ongoing research challenge.