Matt Strong

Matt Strong

CS PhD at Stanford University

← Paper Digest

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

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby

Seminal October 2020

TL;DR Cut an image into 16x16 patches and feed them, like word tokens, to a standard Transformer encoder. Pre-trained on 300M images, ViT-H/14 reaches 88.55% on ImageNet and 77.63% on VTAB with about a quarter of the pre-training compute of the best ResNet.

What they do

They apply a standard NLP Transformer to image classification with as few image-specific changes as possible. An image becomes a sequence of patches, each treated like a word token. The model is pre-trained with supervised classification on large datasets and fine-tuned on smaller benchmarks. The central finding is that large-scale pre-training makes up for the convolutional inductive biases (locality, translation equivariance) that ViT lacks.

How it works

  • Patch embedding. The image is cut into P×P patches (e.g. 16×16), each flattened and linearly projected. A learnable [class] token is prepended and learned 1D position embeddings are added; 2D-aware and relative variants gave no clear gain.
  • Encoder. A standard Transformer with pre-LayerNorm, alternating multi-head self-attention and two-layer GELU MLP blocks. Base/Large/Huge sizes follow BERT (86M, 307M, 632M parameters).
  • Fine-tuning. The head is replaced by a zero-initialized linear layer. At higher resolution the patch size stays fixed, so the sequence gets longer, and position embeddings are interpolated in 2D.
  • Hybrid. Optionally, the patches come from a ResNet feature map instead of raw pixels.

Results

  • ViT-H/14 pre-trained on JFT-300M: 88.55% ImageNet, 94.55% CIFAR-100, 77.63% VTAB, using 2.5k TPUv3-core-days versus 9.9k for BiT-L and 12.3k for Noisy Student.
  • Pre-trained on ImageNet alone, ViT-Large underperforms ViT-Base and BiT ResNets win; with JFT-300M, ViT overtakes. ViT-B/32 does much worse than ResNet50 on a 9M-image subset but better on 90M+.
  • At equal transfer accuracy, ViT uses about 2-4× less pre-training compute than ResNets; hybrids help only at small budgets.
  • Masked patch prediction self-supervision gives ViT-B/16 79.9% on ImageNet, 2% above training from scratch but 4% below supervised pre-training.

Limitations

Trained on mid-sized data like ImageNet without strong regularization, ViT lands a few points below comparable ResNets, so the strong results depend on 14M-300M images of pre-training, the largest set being the in-house JFT-300M. Only classification is evaluated; detection and segmentation are left open, as is the gap between self-supervised and supervised pre-training.

Figures

From Figure 1: The image is split into fixed-size patches, each linearly embedded; position embeddings and a learnabl...
From Figure 1: The image is split into fixed-size patches, each linearly embedded; position embeddings and a learnable classification token are added, and the sequence goes through a standard Transformer encoder.
From Figure 5: Transfer accuracy versus pre-training compute for ViTs, ResNets and hybrids. ViTs beat ResNets at equa...
From Figure 5: Transfer accuracy versus pre-training compute for ViTs, ResNets and hybrids. ViTs beat ResNets at equal compute; hybrids help only for smaller models.
From Figure 6: Attention from the output token back to the input image, for three examples.
From Figure 6: Attention from the output token back to the input image, for three examples.
Abstract

While the Transformer architecture has become the de-facto standard for natural language processing tasks, its applications to computer vision remain limited. In vision, attention is either applied in conjunction with convolutional networks, or used to replace certain components of convolutional networks while keeping their overall structure in place. We show that this reliance on CNNs is not necessary and a pure transformer applied directly to sequences of image patches can perform very well on image classification tasks. When pre-trained on large amounts of data and transferred to multiple mid-sized or small image recognition benchmarks (ImageNet, CIFAR-100, VTAB, etc.), Vision Transformer (ViT) attains excellent results compared to state-of-the-art convolutional networks while requiring substantially fewer computational resources to train.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.