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

“An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale” was written by Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit and Neil Houlsby of Google Research’s Brain Team, in Berlin, Zürich and Amsterdam. It appeared on arXiv on 22 October 2020 and was published at the International Conference on Learning Representations (ICLR) in 2021. The paper applied a standard transformer encoder, with as few changes as possible, to a sequence of image patches, and asked how much pre-training data such a model needs to compete with convolutional networks. Its answer was that large-scale pre-training outweighs the missing inductive bias. The model itself, its later variants and its use as a backbone are described in the Vision Transformer article.

Versions. The arXiv v1 is the preprint submitted for review. The v2 of 3 June 2021 is, per its arXiv comment, the ICLR camera-ready version with two changes: an added appendix comparing the class token with global average pooling, and corrected exaFLOPs in Figure 5 and Table 6, which the authors say leave the models’ relative performance essentially unchanged. We infer that the published paper is v2 without those two changes, and “published version” below means v2 in that sense. Results, appendix letters and figure numbers follow v2; the added appendix shifts later appendix sections and figures by one relative to v1 and, by the same inference, to the published version.

Problem

By 2020, transformers pre-trained on large text corpora dominated natural language processing and had been scaled past 100 billion parameters without saturating. In vision, the authors saw two lines of work: CNNs combined with self-attention, as in non-local networks and DETR, and models that replaced convolutions with local, sparse or axial attention, which they judged theoretically efficient but not yet scaled on accelerators because of the specialized attention patterns. For large-scale image recognition they considered ResNet-like networks still the state of the art; their baselines were BiT (Kolesnikov et al., 2020), built on ResNet, and Noisy Student’s EfficientNet (Xie et al., 2020).

Two earlier models were close. Cordonnier et al. (2020) had applied full self-attention to $2 \times 2$ patches, which the published version names as the most related work, noting that so small a patch limits the model to small images; image GPT (iGPT) had run a transformer on reduced-resolution pixels as a generative model. The v1 preprint instead said the authors knew of no prior transformer with global self-attention on full-sized images and named iGPT as the closest model.

Contribution

• A plain transformer on patches. Cutting a $224 \times 224$ image into $16 \times 16$ patches gives a sequence of 196 tokens, short enough for global attention and for unchanged NLP transformer code.
• A data-scale study. The same models were pre-trained on ImageNet (1.3 million images), ImageNet-21k (14 million) and Google’s internal JFT-300M (303 million), plus random JFT subsets, to measure when the weaker inductive bias stops hurting.
• A controlled compute study of 18 ViTs, ResNets and hybrids trained on JFT-300M.
• State-of-the-art transfer to ImageNet, CIFAR, Oxford Pets and Flowers and the VTAB suite at a fraction of the pre-training cost of the CNN baselines.
• Analyses of the trained model, a pilot of self-supervised pre-training, and released code and checkpoints (google-research/vision_transformer).

Method

The architecture follows the original transformer (Vaswani et al., 2017) as closely as possible, with BERT’s class token and model sizes; the Vision Transformer article gives its equations and sizes. A name such as ViT-L/16 gives the size and the patch side in pixels.

Baselines. The CNNs are “ResNet (BiT)”: ResNets with group normalization and weight-standardized convolutions in place of batch normalization, modifications that Kolesnikov et al. (2020) found to improve transfer.

Hybrids. ResNet-50 features enter the transformer as $1 \times 1$ patches, taken from stage 4 or from a stage 3 extended with stage 4’s layers, which gives four times as many tokens. The number after the slash is then the backbone’s total downsampling, so R50+ViT-B/32 uses stage 4 and R50+ViT-B/16 the extended stage 3.

Training. Every model, ResNets included, was pre-trained with Adam (settings in the subject article), for 7 or 14 epochs on JFT-300M, 30 or 90 on ImageNet-21k and 300 on ImageNet. In Appendix D.1, ResNets pre-trained on JFT with Adam slightly beat SGD after fine-tuning on five datasets (mean 89.33% against 88.79% for ResNet-50, 94.01% against 93.72% for ResNet-152x2), which the authors took to justify this unusual choice. JFT-300M and ImageNet-21k were de-duplicated against the downstream test sets.

Transfer. Fine-tuning uses SGD at $384 \times 384$ by default, with a learning rate picked from a small grid on held-out parts of each training set. VTAB uses 1,000 training examples per task and one setting for all 19 tasks. Where full fine-tuning was too costly, the authors used linear few-shot accuracy: a regularized least-squares map from frozen features of a few training images per class to $\{-1, 1\}^K$ targets, solved in closed form.

Results

Against the state of the art (Table 2). The headline comparison used two models fine-tuned at higher resolution ($512 \times 512$ for ViT-L/16, $518 \times 518$ for ViT-H/14) with Polyak averaging. Entries are top-1 accuracy in percent, averaged over three fine-tuning runs:

Dataset	ViT-H/14 (JFT)	ViT-L/16 (JFT)	ViT-L/16 (ImageNet-21k)	BiT-L (ResNet-152x4)	Noisy Student (EfficientNet-L2)
ImageNet	88.55	87.76	85.30	87.54	88.4 / 88.5
ImageNet ReaL	90.72	90.54	88.62	90.54	90.55
CIFAR-100	94.55	93.90	93.25	93.51	–
Oxford-IIIT Pets	97.56	97.32	94.67	96.62	–
VTAB (19 tasks)	77.63	76.28	72.72	76.29	–
TPUv3-core-days	2.5k	0.68k	0.23k	9.9k	12.3k

The authors state that ViT-L/16 on JFT-300M outperformed BiT-L, pre-trained on the same data, on all tasks, at about a fourteenth of its pre-training cost; in the table it is ahead on most, but level on ImageNet ReaL (90.54%) and VTAB (76.28% against 76.29%). ViT-H/14 gained most on ImageNet, CIFAR-100 and VTAB. The authors cautioned that pre-training cost also depends on schedule, optimizer and weight decay. By VTAB group (Figure 2), ViT-H/14 led BiT-L and two other prior methods on the Natural and Structured tasks and roughly tied BiT-L on the Specialized (medical and satellite) tasks. On ObjectNet, in an appendix absent from v1 (D.9), ViT-H/14 reached 61.7% top-1 and 82.1% top-5 accuracy.

Data requirements (§4.3). Fine-tuned to ImageNet, the Large models trailed the Base models when pre-trained on ImageNet, roughly tied them on ImageNet-21k and pulled ahead only on JFT-300M, while with ImageNet pre-training the BiT ResNets beat ViT (Figure 3). The subject article gives the numbers, and the transformer article the JFT-subset experiment.

Compute (§4.4). On JFT-300M, ViT needed about 2–4 times less pre-training compute than ResNets for the same accuracy, and hybrids helped only at small budgets (the subject article summarizes this). Selected rows of Table 6 (v2 compute values; ImageNet top-1 after fine-tuning at $384 \times 384$):

Model	Epochs	ImageNet (%)	Pre-training exaFLOPs
ResNet-50x1	7	77.54	50
ViT-B/32	7	80.73	55
R50+ViT-B/32	7	84.90	106
ViT-B/16	7	84.15	224
ResNet-152x2	7	84.97	563
ViT-L/16	14	87.12	1,567
R50+ViT-L/16	14	87.12	1,668
ResNet-200x3	14	87.22	3,306
ViT-H/14	14	88.08	4,262

In these rows, ViT-L/16 nearly matches ResNet-200x3 with under half the compute, and the 14-epoch hybrid and pure ViT-L/16 tie, while among small models R50+ViT-B/32 beats ViT-B/16 with half its compute. In v1 each model’s exaFLOPs were 2.9 to 3.3 times the v2 values, so the comparisons change little.

Model shape (Appendix D.2). Starting from an 8-layer, $D = 1024$, patch-32 model and evaluating ImageNet 5-shot accuracy, depth helped most, clearly up to 64 layers though with diminishing returns after 16, width helped least, and smaller patches gave robust gains without new parameters. The authors concluded that compute predicts performance better than parameter count. In Appendix D.6, axial-attention versions of ViT-B were more accurate than ViT-B but cost more compute, and an AxialResNet baseline was very slow on TPUs.

Inspection (§4.5). The top principal components of ViT-L/32’s patch-embedding filters resemble basis functions for the fine structure within a patch. The position embeddings’ similarity pattern, which the authors took as the reason 2D-aware embeddings brought no gain, depends on training hyperparameters (Figure 10). Mean attention distance, for each head and layer, averages over 128 images the distance between the query pixel and all other pixels, weighted by attention. In ViT-L/16 some heads span most of the 224-pixel image already in the lowest layers while others stay local, less so in an R50+ViT-L/16 hybrid; in the second half of the network most heads attend widely. Class-token attention maps computed with attention rollout (Abnar & Zuidema, 2020), which averages heads and multiplies the attention matrices of all layers, highlight class-relevant regions.

Self-supervision pilot (§4.6, Appendix B.1.2). Following BERT, 50% of patch embeddings were corrupted: 80% of those replaced by a learned mask embedding, 10% by another patch’s embedding, 10% left unchanged. The model predicted the 3-bit mean color of each corrupted patch, one of 512 colors. A ViT-B/16 pre-trained this way for 1 million steps on JFT reached 79.9% on ImageNet, 2 points above training from scratch and 4 below supervised pre-training. Predicting a $4 \times 4$ downsized patch or regressing all pixels with an L2 loss also worked, the latter slightly worse; the mean color gave the best few-shot accuracy, and a 15% corruption rate was slightly worse. Gains flattened after 100,000 steps and were similar when pre-training on ImageNet.

The analysis quantities in code. The NumPy script below computes, on a synthetic $224 \times 224$ image, the 512-class mean-color targets of the self-supervision pilot, and the mean attention distance of five artificial heads on the $14 \times 14$ grid. The paper does not give its quantization rule, so the script assumes eight equal levels per channel (our choice); it also measures distance between patch centers, our patch-level approximation to the paper’s pixel-level definition.

import numpy as np

rng = np.random.default_rng(0)

# A 224x224 RGB synthetic image: gray gradient, a red square and a blue disc; 16x16 patches -> 14x14 grid.
S, P = 224, 16
G = S // P
yy, xx = np.mgrid[0:S, 0:S]
img = np.repeat((0.3 + 0.4 * xx / (S - 1))[..., None], 3, axis=2)
img[32:96, 128:192] = [0.9, 0.1, 0.1]                                  # red square
img[(yy - 160) ** 2 + (xx - 64) ** 2 < 40 ** 2] = [0.1, 0.2, 0.9]      # blue disc

# Self-supervision pilot target: the 3-bit mean colour of each patch (8 levels per channel, 512 classes).
patches = img.reshape(G, P, G, P, 3).transpose(0, 2, 1, 3, 4).reshape(G * G, P * P, 3)
levels = np.clip((patches.mean(axis=1) * 8).astype(int), 0, 7)       # (196, 3), each in 0..7
target = levels[:, 0] * 64 + levels[:, 1] * 8 + levels[:, 2]          # one of 512 classes per patch
print(f"{G * G} patches, {len(np.unique(target))} distinct mean-colour classes out of 512")
for name, (r, c) in {"red square": (3, 9), "blue disc": (10, 4), "background": (0, 0)}.items():
    i = r * G + c
    print(f"  {name:10s} patch ({r:2d},{c:2d}): RGB levels {tuple(int(v) for v in levels[i])} -> class {target[i]}")

# Mean attention distance: attention-weighted pixel distance between query and key patch centres,
# averaged over queries (the paper averages over 128 images and plots one dot per head and layer).
ctr = (np.stack(np.mgrid[0:G, 0:G], -1).reshape(-1, 2) + 0.5) * P   # patch centres in pixels
dist = np.linalg.norm(ctr[:, None] - ctr[None], axis=-1)              # (196, 196)

def mean_attention_distance(A):
    return (A * dist).sum(axis=1).mean()

def softmax(z):
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)

heads = {
    "self only (identity)": np.eye(G * G),
    "local, Gaussian sigma = 1 patch": softmax(-(dist / P) ** 2 / 2),
    "local, Gaussian sigma = 3 patches": softmax(-(dist / (3 * P)) ** 2 / 2),
    "uniform over all patches": np.full((G * G, G * G), 1 / (G * G)),
    "random logits, std 3": softmax(rng.normal(0, 3, (G * G, G * G))),
}
print("\nmean attention distance on a 224 x 224 image (pixels)")
for name, A in heads.items():
    print(f"  {name:34s} {mean_attention_distance(A):6.1f}")

196 patches, 12 distinct mean-colour classes out of 512
  red square patch ( 3, 9): RGB levels (7, 0, 0) -> class 448
  blue disc  patch (10, 4): RGB levels (0, 1, 7) -> class 15
  background patch ( 0, 0): RGB levels (2, 2, 2) -> class 146

mean attention distance on a 224 x 224 image (pixels)
  self only (identity)                  0.0
  local, Gaussian sigma = 1 patch      18.4
  local, Gaussian sigma = 3 patches    52.9
  uniform over all patches            116.5
  random logits, std 3                116.2

The mean-color target is coarse: the whole image yields only 12 of the 512 classes, and a patch’s class says nothing about its texture. A uniform head scores about 116 pixels, near the top of the 0–120-pixel axis of the paper’s attention-distance plots, so in our reading heads plotted there attend almost uniformly or to distant patches. The random head shows a limitation of the measure: it attends sharply, to a few patches, yet scores like the uniform head, because attention distance measures how far attention reaches, not how concentrated it is.

Impact

The paper is widely cited, and its model became a standard vision backbone, as the Vision Transformer article describes. Its central claim, in the authors’ words, was that “large scale training trumps inductive bias”, and the ImageNet-21k model showed that a public dataset could give strong results. Its measurements became reference points: DeiT compares with its ImageNet-only figures, MAE quotes its ViT-L ImageNet result as the from-scratch baseline, and Zhai et al. (2022) list the ViT-H/14 results as a baseline.

Limitations

• The headline needs private data. The best results rest on JFT-300M, which is not public; the ImageNet-21k model, at 85.30% on ImageNet, sits below BiT-L’s 87.54%.
• Moderately regularized small-data runs. For pre-training on ImageNet and ImageNet-21k the authors tuned three basic regularization parameters: weight decay, dropout and label smoothing. The regularization described as “heavy” in v1 is “moderate” in the published version, whose introduction adds that ViT trails ResNets on ImageNet “without strong regularization”. In our reading, the small-data conclusion is tied to this recipe.
• Classification only. The conclusion lists detection and segmentation as open challenges; the paper reports no dense-prediction results.
• A large gap for self-supervision. The masked-patch pilot lagged supervised pre-training by 4 points, and contrastive pre-training was left to future work.
• Small inconsistencies. Tables 5 and 6 give different results for ViT-H/14 pre-trained on JFT-300M (88.04% and 88.08% on ImageNet, 97.56% and 97.11% on Pets), although the other JFT-300M rows of Table 5 match Table 6; and v1’s exaFLOPs were about three times too large.

What Came After

DeiT (Touvron et al., 2021), first posted two months after the ViT preprint, quoted its conclusion that transformers generalize poorly with insufficient data and trained the unchanged ViT-B on ImageNet alone with strong augmentation and regularization, reaching 81.8% at $224 \times 224$ and 83.1% at $384 \times 384$ against the 77.9% reported by Dosovitskiy et al.; a distillation token raised this further. Steiner et al. (2022), including four ViT authors, found that with more compute plus augmentation and regularization, ImageNet-21k models could match or beat the paper’s JFT-300M ones.

Masked image modeling later went much further than the pilot. Masked autoencoders (He et al., 2022) masked 75% of patches and reconstructed pixels; they quoted the ViT paper’s 76.5% for ViT-L trained from scratch, reached 82.5% from scratch with their own strongly regularized recipe, and 84.9% after MAE pre-training, all on ImageNet-1k (see masked autoencoder).

The paper’s observation that ViT did not saturate was tested by Zhai et al. (2022), all four of them ViT authors, who scaled ViT to two billion parameters and three billion training images and reached 90.45% top-1 on ImageNet with ViT-G/14. Finding class-token and pooling heads equally accurate, in line with the appendix added in v2, they replaced the class token with attention pooling to save memory and dropped the head’s nonlinear projection. The paper’s other open challenge, detection and segmentation, was taken up by hierarchical and plain-backbone ViTs covered in the Vision Transformer article.