PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2010.11929v2cs.CVcs.AIcs.LG

TL;DR

CNNs dominated large-scale image recognition, with limited evidence for scalable attention-based alternatives. ViT applies a standard Transformer to image-patch sequences and, after large-scale pre-training, approaches or beats state of the art across image recognition benchmarks.

  • Problem

    Large-scale image recognition was dominated by CNNs, with limited evidence that attention-based alternatives could scale effectively on modern hardware accelerators.

    Source: §1 · p. 1

  • Method

    ViT applies a standard Transformer to sequences of linearly embedded image patches and trains it using supervised image classification.

    Source: §1 · p. 1, §3.1 · p. 3

  • Results

    After large-scale pre-training, ViT approaches or beats state of the art across image recognition benchmarks, reaching 88.55% accuracy on ImageNet.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Approximately 2–4× less compute than ResNets for equal transfer performance, averaged over five datasets, positions ViT as a compute-efficient alternative.

    Source: §4.4 · p. 8

    Applying ViT to computer vision tasks beyond classification, such as detection and segmentation, remains an open challenge.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Vision Transformer (ViT) applies a standard Transformer directly to linearly embedded image patches for supervised classification. Although it underperforms comparable ResNets on mid-sized datasets without strong regularization, sufficiently large-scale pre-training enables strong transfer results across image recognition benchmarks.

1 INTRODUCTION

  • Specialized attention patterns hindered effective scaling of earlier convolution-replacing vision models on modern accelerators (Ramachandran et al., 2019).ResNet-like architectures remained state of the art in large-scale image recognition (Mahajan et al., 2018; Kolesnikov et al., 2020).Source: §1 · p. 1
  • ViT splits images into patches and feeds their linear embeddings into a standard Transformer, treating patches like word tokens and training through supervised image classification.The approach makes as few modifications as possible to the standard Transformer.Source: §1 · p. 1
  • A few percentage points lower accuracy than comparably sized ResNets results from training ViT on mid-sized datasets such as ImageNet without strong regularization.Transformers lack CNN inductive biases such as translation equivariance and locality, limiting generalization when training data are insufficient.Source: §1 · pp. 1–2
  • 88.55% on ImageNet, 90.72% on ImageNet-ReaL, 94.55% on CIFAR-100, and 77.63% on VTAB are the best ViT model’s accuracies.Pre-training on ImageNet-21k or JFT-300M enables ViT to approach or beat state of the art across multiple benchmarks. Training datasets contain 14M-300M images, and the VTAB suite comprises 19 tasks.Source: §1 · p. 2

2 RELATED WORK

ViT extends prior patch-based image Transformers through large-scale pre-training, making vanilla Transformers competitive with state-of-the-art CNNs while handling medium-resolution images. Its related work spans NLP pre-training, specialized visual attention, CNN–attention hybrids, generative image modeling, and large-scale transfer learning.

2 RELATED WORK

3 METHOD

ViT closely follows the original Transformer, allowing scalable NLP architectures and efficient implementations to be reused with minimal changes. It processes images as patch embeddings and supports downstream fine-tuning at higher resolutions.

3.1 VISION TRANSFORMER (VIT)

  • ViT linearly projects flattened image patches into D-dimensional embeddings and adds learnable 1D position embeddings before encoding.For image dimensions H and W, P x P patches produce N = HW/P^2 tokens. More advanced 2D-aware position embeddings showed no significant performance gains.Source: §3.1 · p. 3
  • A prepended learnable classification token supplies the final image representation, which feeds a classification head during both pre-training and fine-tuning.The head is an MLP with one hidden layer during pre-training and a single linear layer during fine-tuning.Source: §3.1 · p. 3
  • The Transformer encoder alternates multiheaded self-attention and MLP blocks, with layer normalization before each block and residual connections afterward.Each MLP contains two layers with a GELU nonlinearity. The normalization and residual arrangement follows (Baevski & Auli, 2019).Source: §3.1 · pp. 3–4
  • ViT has less image-specific inductive bias than CNNs: self-attention is global, and spatial relationships between patches must be learned from scratch.Only MLP layers are local and translationally equivariant. Explicit 2D structure enters only through patch extraction and position-embedding adjustment for different image resolutions.Source: §3.1 · p. 4, §3.2 · p. 4
  • The hybrid architecture replaces raw image patches with patches extracted from CNN feature maps, then projects them into Transformer embeddings.With 1x1 feature-map patches, spatial dimensions are simply flattened and projected; classification and position embeddings are added as in the raw-patch model.Source: §3.1 · p. 4

3.2 FINE-TUNING AND HIGHER RESOLUTION

  • Higher-resolution fine-tuning is often beneficial (Kolesnikov et al., 2020), retaining patch size while applying 2D interpolation to pretrained position embeddings.The pretrained prediction head is replaced with a zero-initialized D x K layer for K downstream classes. Higher resolution increases sequence length, which ViT can accommodate subject to memory constraints.Source: §3.2 · p. 4

4 EXPERIMENTS

With large-scale pre-training, ViT outperforms ResNet-based baselines across the reported benchmarks while requiring less pre-training compute; its advantage depends strongly on pre-training dataset size.

4.1 SETUP

  • Experiments compare ViT, modified ResNets, and hybrids across pre-training scales using downstream fine-tuning accuracy and frozen-representation few-shot accuracy.Pre-training spans ImageNet, ImageNet-21k, and JFT, with test-set deduplication; transfer covers ImageNet/ReaL, CIFAR-10/100, Pets, Flowers, and VTAB’s 19 tasks using 1,000 training examples each. ResNet baselines use Group Normalization and standardized convolutions, hybrids feed intermediate CNN features into ViT, and few-shot evaluation uses closed-form regularized least squares.Source: §4.1 · pp. 4–5

4.2 COMPARISON TO STATE OF THE ART

  • JFT-300M-pretrained ViT-L/16 outperforms BiT-L on all reported tasks with substantially less pre-training compute, and ViT-H/14 improves further.Efficiency differences may also reflect training schedules, optimizers, and weight decay rather than architecture alone. ViT-L/16 pre-trained on public ImageNet-21k also performs well, requiring approximately 30 days on eight TPUv3 cores.Source: §4.2 · pp. 5–6, Table 2 · p. 6
  • ViT-H/14 outperforms BiT-R152x4 and other compared methods on VTAB’s Natural and Structured tasks, while the top two models perform similarly on Specialized tasks.The comparison includes VIVI, co-trained on ImageNet and YouTube, and S4L, combining supervised and semi-supervised ImageNet learning.Source: Figure 2 · p. 6, §4.2 · p. 6

4.4 SCALING STUDY

  • 2–4× less pre-training compute gives ViT the same performance as ResNets, averaged over five datasets in the controlled JFT-300M scaling study.Hybrids slightly outperform pure ViT at small computational budgets, but their advantage disappears for larger models. ViT shows no apparent performance saturation within the tested range.Source: Figure 5 · p. 7, §4.4 · p. 8

4.3 PRE-TRAINING DATA REQUIREMENTS

  • Larger pre-training datasets reverse ViT’s disadvantage against BiT and unlock the benefits of larger ViT models.Despite tuning weight decay, dropout, and label smoothing, ViT-Large trails ViT-Base with ImageNet pre-training; they perform similarly with ImageNet-21k, while JFT-300M reveals larger models’ full benefit.Source: §4.3 · pp. 6–7, Figure 3 · p. 7
  • ViT overfits more than comparably costly ResNets on small JFT subsets but overtakes them with larger pre-training datasets.With fixed hyperparameters and early stopping, ViT-B/32 performs much worse than ResNet50 on 9M images but better on 90M+ subsets, using frozen-feature few-shot evaluation. These findings and VTAB results suggest promise for very low-data transfer.Source: §4.3 · p. 7, Figure 4 · p. 7

4.5 INSPECTING VISION TRANSFORMER

  • ViT learns two-dimensional positional structure and uses both local and global attention, with some heads integrating most of the image in the lowest layers.Initial projection filters resemble patch-level basis functions; position embeddings encode distance and row-column structure, explaining why handcrafted two-dimensional variants do not improve results. Attention distance grows with depth, localized early attention is less pronounced in hybrids, and attention targets semantically relevant regions.Source: §4.5 · p. 8, Figure 7 · p. 9

4.6 SELF-SUPERVISION

  • 79.9% ImageNet accuracy follows masked-patch self-supervised pre-training of ViT-B/16, improving by 2% over training from scratch but remaining 4% behind supervised pre-training.This preliminary experiment predicts masked image patches; contrastive pre-training is left for future work.Source: §4.6 · p. 9

5 CONCLUSION · APPENDIX · A MULTIHEAD SELF-ATTENTION

ViT processes image patches with a standard Transformer encoder and, with large-scale pre-training, matches or exceeds state-of-the-art image classification performance while remaining relatively cheap to pre-train. The appendix explains self-attention and its multihead extension.

5 CONCLUSION

  • ViT matches or exceeds the state of the art on many image classification datasets while remaining relatively cheap to pre-train with large datasets.These results depend on coupling the architecture with large-scale pre-training.Source: §5 · p. 9
  • ViT processes images as patch sequences using a standard Transformer encoder, introducing no image-specific architectural inductive biases beyond initial patch extraction.The encoder follows the standard architecture used in NLP.Source: §5 · p. 9
  • Initial experiments show improvements from self-supervised pre-training, but a large gap remains relative to large-scale supervised pre-training.Further exploration of self-supervised pre-training remains an open challenge.Source: §5 · p. 9
  • Extending ViT to other computer vision tasks, including detection and segmentation, remains an open challenge.The authors see promise in this direction based on their results and those of Carion et al. (2020).Source: §5 · p. 9
  • The authors anticipate that further scaling of ViT would improve performance, rather than presenting this as an established result.Source: §5 · p. 9

A MULTIHEAD SELF-ATTENTION

  • Self-attention (Vaswani et al. (2017)) weights sequence values using query–key similarities; multihead attention runs these operations in parallel and projects their concatenated outputs.With k heads, setting Dh = D/k keeps compute and parameter counts constant when k changes.Source: A MULTIHEAD SELF-ATTENTION · p. 13

B EXPERIMENT DETAILS … B.1.2 SELF-SUPERVISION

The experiments specify regularized training, transfer fine-tuning, and preliminary masked-patch self-supervision. Masked-patch pretraining shows diminishing downstream returns after 100k steps and similar ImageNet-classification gains when pretrained on ImageNet.

B.1 TRAINING

  • 224 resolution is used for all training, with strong regularization essential for training from scratch on ImageNet and identical setups for hybrids and corresponding ViTs.When enabled, dropout follows every dense layer except qkv projections and is also applied immediately after positional embeddings are added to patch embeddings.Source: §B.1 · p. 13

B.1.1 FINE-TUNING

  • 0.9 momentum SGD fine-tunes ViT, with learning rates selected on held-out training subsets before full-training-set fitting and test evaluation.Defaults are cosine decay, batch size 512, no weight decay, global-norm clipping at 1, and resolution 384; development subsets are 10% for Pets/Flowers, 2% for CIFAR, and 1% for ImageNet. ResNets and hybrids share this setup but add 0.06 to the ImageNet learning-rate sweep; ResNets additionally select the best result across that sweep and the setup of Kolesnikov et al. (2020).Source: §B.1.1 · pp. 13–14, Table 4 · p. 14
  • Transfer replaces ViT’s entire two-layer head with one zero-initialized linear classifier, which is slightly more robust than reinitializing only the final layer.The replacement classifier outputs the number of classes required by the target dataset.Source: §B.1.1 · p. 14
  • 0.01 learning rate and 2500 training steps are shared across VTAB tasks, with 384 × 384 inputs benefiting ViT most instead of task-specific resolutions.Following Kolesnikov et al. (2020), the setting is selected by sweeping two learning rates and two schedules on 200-example validation sets. Preprocessing follows that work except for the resolution choice.Source: §B.1.1 · p. 14

B.1.2 SELF-SUPERVISION

  • 50% of patch embeddings are corrupted for masked-patch prediction: 80% receive a learnable mask embedding, 10% a random other patch embedding, and 10% remain unchanged.Each corrupted patch’s representation predicts its mean 3-bit color, choosing among 512 colors. The corruption setup resembles the language procedure of Devlin et al. (2019).Source: §B.1.2 · p. 14
  • 1M steps, approximately 14 epochs, of self-supervised JFT pretraining use Adam and batch size 4096.The base learning rate is 2 × 10^-4, with 10k warmup steps and cosine learning-rate decay.Source: §B.1.2 · p. 14
  • 15% corruption slightly worsened few-shot metrics; mean 3-bit color prediction achieved the best few-shot performance among tested targets, while full-patch L2 regression was slightly worse.All three targets worked reasonably well: one mean-color prediction, sixteen color predictions for a 4 × 4 downsized patch, and 256 RGB regressions for the full patch. Final results use only mean-color prediction.Source: §B.1.2 · p. 14
  • 100k pretraining steps were followed by diminishing downstream returns, and pretraining on ImageNet yielded similar ImageNet-classification gains without requiring JFT.The authors conclude that their masked-patch prediction approach does not require extensive pretraining or a dataset as large as JFT to obtain these gains.Source: §B.1.2 · p. 14

C ADDITIONAL RESULTS … D.3 HEAD TYPE AND CLASS TOKEN

Additional results document ViT transfer performance and model-scaling compute costs. Ablations support Adam for ResNet pre-training, emphasize depth over width when scaling ViT, and explain classifier-head differences through learning-rate requirements.

C ADDITIONAL RESULTS

  • Table 5 reports transfer top-1 accuracy for ViT models pre-trained on ImageNet, ImageNet-21k, or JFT-300M.Models are fine-tuned at 384 resolution. ImageNet results exclude the Polyak averaging and 512-resolution images used for Table 2.Source: C ADDITIONAL RESULTS · p. 14, Table 5 · p. 15
  • Table 6 documents transfer accuracy and estimated pre-training compute in exaFLOPs for ViT, ResNet, and hybrid models of varying sizes.These detailed scaling results correspond to Figure 5 in the main paper.Source: Table 6 · p. 15, C ADDITIONAL RESULTS · p. 15

D.1 SGD VS. ADAM FOR RESNETS

  • Adam pre-training outperforms SGD on most transfer datasets and on average for JFT-pre-trained ResNet models.The comparison covers ResNet models 50x1 and 152x2. Absolute results are lower than those of Kolesnikov et al. (2020) because pre-training lasts 7 rather than 30 epochs.Source: Table 7 · p. 16, §D.1 · p. 16

D.2 TRANSFORMER SHAPE

  • 64 layers still show gains from depth scaling, the strongest ImageNet 5-shot scaling direction, although diminishing returns appear after 16 layers.The reference configuration has 8 layers, D = 1024, DMLP = 2048, and patch size 32. Width scaling produces the smallest changes, while scaling all dimensions proportionally yields robust improvements.Source: §D.2 · p. 16, Figure 8 · p. 16
  • Smaller patches robustly improve ImageNet 5-shot performance by increasing sequence length without adding parameters.These findings suggest that compute may predict performance better than parameter count.Source: §D.2 · p. 16

D.3 HEAD TYPE AND CLASS TOKEN

  • Class-token and global-average-pooling classifiers perform similarly well but require different learning rates.The main-paper design uses an additional [class] token followed by an MLP with one tanh hidden layer. Initially poor global-average-pooling-plus-linear performance is fully explained by learning-rate choice, not pooling or token removal.Source: §D.3 · pp. 16–17, Figure 9 · p. 17

D.4 POSITIONAL EMBEDDING

Positional embeddings substantially improve ViT-B/16 performance, but different ways of encoding spatial information yield little to no performance difference.

D.4 POSITIONAL EMBEDDING

  • The default uses one-dimensional raster-order embeddings; two-dimensional embeddings concatenate learned X- and Y-axis vectors of size D/2 each.The two-dimensional construction assigns each patch its concatenated embedding according to its grid coordinates.Source: §D.4 · p. 17
  • Relative positional embeddings encode patch-pair distances through additional attention whose logits bias the main content-based attention before softmax.Each query–key offset has an embedding, used as a key with the original content query in the additional attention.Source: §D.4 · p. 17
  • Absolute positional embeddings were added once before the encoder, separately at each layer, or at each layer with shared embeddings.Adding embeddings after the stem and before the Transformer encoder is the default throughout the other experiments.Source: §D.4 · pp. 17–18
  • Positional embeddings substantially outperform treating inputs as a bag of patches, while different spatial encoding strategies perform similarly in ViT-B/16.The authors speculate that spatial relationships are similarly easy to learn across encoding strategies on smaller patch grids, such as 14 × 14 rather than 224 × 224 pixels.Source: §D.4 · pp. 17–18
  • Learned positional-embedding similarity patterns depend on training hyperparameters despite the similar performance of different spatial encoding strategies.Figure 10 illustrates positional embeddings from models trained with different hyperparameters.Source: §D.4 · p. 18, Figure 10 · p. 18

D.5 EMPIRICAL COMPUTATIONAL COSTS

ViT models achieve inference speeds comparable to similar ResNets on TPUv3, while large ViT models show a clear memory-efficiency advantage.

D.5 EMPIRICAL COMPUTATIONAL COSTS

  • Inference timing on TPUv3 measures hardware-dependent speed that theoretical FLOPs do not always predict because of factors such as lane widths and cache sizes.The difference between inference and backpropagation speed is a constant, model-independent factor.Source: §D.5 · pp. 18–19
  • ViT’s theoretical bi-quadratic scaling with image size only barely begins to appear for the largest models at the largest tested resolutions.Reported throughput is the peak images per second per core measured across a wide range of batch sizes.Source: §D.5 · p. 19
  • ViT models achieve speeds comparable to similar ResNets, while large ViT models show a clear memory-efficiency advantage in maximum per-core batch size.Larger per-core batches are better for scaling to large datasets.Source: Figure 12 · p. 19, §D.5 · p. 19

D.6 AXIAL ATTENTION · D.7 ATTENTION DISTANCE · D.8 ATTENTION MAPS

Axial-ViT variants improve ImageNet 5-shot linear performance at greater compute cost, while attention-distance analysis shows broader attention with depth. Output-token attention maps are computed by averaging heads and propagating attention through layers.

D.6 AXIAL ATTENTION

  • Axial attention (Ho et al., 2019) attends along individual tensor axes rather than flattened inputs, keeping other axes independent.The implemented AxialResNet baseline from Wang et al. (2020b) replaces all ResNet50 3 × 3 convolutions with row and column attention augmented by relative positional encoding.Source: §D.6 · p. 19
  • Axial-ViT processes two-dimensional inputs with row-attention-plus-MLP followed by column-attention-plus-MLP, adding an extra MLP that increases compute despite shorter attention sequences.Each original global self-attention Transformer block is replaced by two axial blocks.Source: §D.6 · pp. 19–20
  • AxialResNet reproduces the accuracy reported by Wang et al. (2020b), but its slow TPU implementation prevents extensive large-scale experiments.Its accuracy/compute trade-off appears reasonable; the authors suggest that careful implementation optimization could enable larger-scale experiments.Source: §D.6 · pp. 19–20
  • Axial-ViT-B/32 and Axial-ViT-B/16 achieve higher ImageNet 5-shot linear top-1 accuracy than their ViT-B counterparts after JFT pretraining, but require more compute.Figure 13 compares accuracy against FLOPs and inference speed.Source: §D.6 · pp. 19–20, Figure 13 · p. 20

D.7 ATTENTION DISTANCE

  • Attention distance varies across heads in early layers, from local regions to much of the image, and increases with depth for all heads.In the network’s second half, most heads attend widely across tokens. Attention distance is analogous to receptive-field size in CNNs.Source: §D.7 · p. 20

D.8 ATTENTION MAPS

  • Attention Rollout (Abnar & Zuidema, 2020) maps ViT-L/16 output-token attention to input space by averaging heads and recursively multiplying attention matrices across layers.This accounts for attention mixing across tokens through all layers.Source: §D.8 · p. 20

D.9 OBJECTNET RESULTS · D.10 VTAB BREAKDOWN

ViT-H/14 achieves 82.1% top-5 accuracy and 61.7% top-1 accuracy on ObjectNet. VTAB-1k results are reported separately for each task.

D.9 OBJECTNET RESULTS

D.10 VTAB BREAKDOWN

Related papers

Loading 2010.11929v2…