ICLR 2021 Landmark Paper

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

Authors: Alexey Dosovitskiy*, Lucas Beyer*, Alexander Kolesnikov*, Dirk Weissenborn*, Xiaohua Zhai*, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby* (*equal contribution)

TL;DR

The Vision Transformer (ViT) shows that the dominant reliance on Convolutional Neural Networks (CNNs) for computer vision is unnecessary. By splitting an image into flat patches and feeding them directly into a standard, unmodified NLP Transformer encoder, ViT achieves state-of-the-art accuracy on major benchmarks when pre-trained on massive datasets (like JFT-300M), while requiring significantly less compute to train.

#vision-transformer #self-attention #image-classification #scaling-laws

Why This Matters

For nearly a decade, Convolutional Neural Networks (CNNs) were the undisputed kings of computer vision. CNNs are hardcoded with powerful inductive biases: locality (pixels close to each other are related) and translation equivariance (an object is the same whether it is in the top-left or bottom-right corner).

Meanwhile, the Natural Language Processing (NLP) community had completely transitioned to standard Transformers, which rely on global self-attention. Self-attention has no inherent spatial bias; it forces the model to learn relationships between all tokens from scratch.

The critical question was: Can we throw away the specialized, hand-designed architectural assumptions of convolutions and let a pure Transformer learn the structure of the visual world entirely from data?

The Big Idea

Instead of processing pixels individually (which is computationally impossible due to the quadratic complexity of self-attention), ViT slices an image into a grid of fixed-size patches (e.g., $16 \times 16$ pixels). Each patch is flattened, projected into a vector space, and treated exactly like a "word" in a sentence.

Mathematical Formulation of Patch Embedding

$$ \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}_{\text{pos}} $$

$\mathbf{x}_p^i\mathbf{E}$: Flattened image patch $i$ multiplied by trainable projection matrix $\mathbf{E}$ to map it to dimension $D$.

$\mathbf{x}_{\text{class}}$: A prepended learnable token whose state at the output represents the aggregate image representation.

$\mathbf{E}_{\text{pos}}$: Learnable 1D position embeddings added to preserve spatial relationships.

Interactive Playground

1. Image Patching & Linear Projection

See how an image is tokenized into "words" of size $P \times P$.

Patch Size ($P$):

Input Image ($226 \times 226$ px canvas)

Resulting Sequence (Tokens)

Sequence Length ($N$): 196 tokens Patch Dimensions: 16x16x3

Note: Each patch is flattened into a 1D vector of size $P^2 \cdot C$ (where $C=3$ for RGB), then multiplied by a learnable linear projection matrix $\mathbf{E}$ to produce $D$-dimensional embeddings.

2. The "Scale Trumps Bias" Phenomenon

Drag the slider to see how ViT overtakes CNNs once pre-training dataset size scales up.

Dataset Scale ImageNet-21K (Medium)

ImageNet-1K (1.3M images) ImageNet-21K (14M images) JFT-300M (300M images)

Accuracy (%) on ImageNet-1K (Transfer)

Key Insight:

At medium scale (ImageNet-21K), ViT performs on par with ResNet. The lack of inductive bias is compensated for by the larger dataset.

Convolutions have strong priors that help them learn extremely fast on small data. However, these priors act as a ceiling. When data is abundant, ViT's global receptive field allows it to learn superior representations.

3. Receptive Field & "Attention Distance"

Click on the layers below to see how the average attention distance behaves as depth increases.

Select Layer Depth:

Layer Type Early Layer

Early layers contain heads that focus both locally and globally.

Attention Distance (Simulated Figure 7 Right)

How It Works

1. Patch Extraction and Flattening

An input image $\mathbf{x} \in \mathbb{R}^{H \times W \times C}$ is reshaped into a sequence of flattened 2D patches $\mathbf{x}_p \in \mathbb{R}^{N \times (P^2 \cdot C)}$, where $(P, P)$ is the patch resolution and $N = HW/P^2$ is the resulting number of patches (effective sequence length).

2. Linear Projection and Class Token

We map the flattened patches to $D$ dimensions using a trainable linear projection $\mathbf{E}$. Similar to BERT's `[class]` token, we prepend a learnable embedding $\mathbf{x}_{\text{class}}$ to the sequence. The output state of this token serves as the aggregate representation for classification.

3. Learnable 1D Position Embeddings

Since Transformers are permutation-invariant, we add learnable 1D position embeddings to the patch representations to retain spatial context. Interestingly, the authors found that 2D-aware position embeddings did not yield significant performance gains over standard 1D ones.

4. Transformer Encoder & MLP Head

The sequence of vectors is passed through a standard Transformer encoder consisting of alternating layers of Multi-Head Self-Attention (MSA) and MLP blocks. Layer Normalization (LN) is applied before every block, and residual connections are added after.

Results & Performance

When pre-trained on Google's proprietary JFT-300M dataset, ViT models match or exceed state-of-the-art ResNet-based baselines while requiring significantly fewer computational resources to train.

ImageNet Top-1

88.55%

ViT-H/14 (JFT-300M)

ImageNet ReaL

90.72%

ViT-H/14 (JFT-300M)

CIFAR-100

94.55%

ViT-H/14 (JFT-300M)

VTAB (19 Tasks)

77.63%

ViT-H/14 (JFT-300M)

Pre-training Compute Efficiency

Not only does ViT perform better, but it is also much cheaper to train than comparable CNNs. For example, ViT-L/16 pre-trained on JFT-300M requires 2-4x less compute (0.68k TPUv3-core-days) to achieve the same or better performance than BiT-L (ResNet152x4), which took 9.9k TPUv3-core-days.

Limitations & Open Questions

Data Hunger: When trained on mid-sized datasets (like ImageNet-1K) without strong regularization, ViT underperforms ResNets by a few percentage points. It requires massive pre-training (ImageNet-21K or JFT-300M) to shine.
High Memory Footprint: Global self-attention scales quadratically ($O(N^2)$) with the number of patches. This makes applying ViT directly to high-resolution images or pixel-level tasks (like segmentation) challenging without specialized modifications.
Self-Supervised Gap: While NLP relies heavily on self-supervised pre-training (like BERT or GPT), ViT's self-supervised variant (masked patch prediction) achieved 79.9% on ImageNet—a solid improvement over training from scratch, but still 4% behind supervised pre-training.

Glossary

Inductive Bias

The set of assumptions that a learning algorithm uses to predict outputs for unseen inputs. For example, CNNs assume translation equivariance and spatial locality.

Multi-Head Self-Attention (MSA)

An attention mechanism that relates different positions of a single sequence to compute a representation of the sequence. "Multi-head" allows the model to jointly attend to information from different representation subspaces.

MLP (Multi-Layer Perceptron)

A class of feedforward artificial neural network. In ViT, the MLP contains two layers with a GELU non-linearity.

Cite This Paper

@inproceedings{
  dosovitskiy2021an,
  title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},
  author={Alexey Dosovitskiy and Lucas Beyer and Alexander Kolesnikov and Dirk Weissenborn and Xiaohua Zhai and Thomas Unterthiner and Mostafa Dehghani and Matthias Minderer and Georg Heigold and Sylvain Gelly and Jakob Uszkoreit and Neil Houlsby},
  booktitle={International Conference on Learning Representations},
  year={2021},
  url={https://openreview.net/forum?id=YicbFdNTTy}
}