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

Computer vision still relies mainly on convolutional architectures, limiting the demonstrated scalability of pure Transformers for large-scale image recognition. This paper applies a standard Transformer to sequences of image patches and finds that large-scale pretraining enables ViT to match or exceed state-of-the-art performance across multiple benchmarks while using less training compute.

  • Problem

    Large-scale image recognition remains dominated by CNNs because pure attention-based vision models have not scaled effectively on modern hardware.

  • Method

    Vision Transformer splits images into patches, embeds them as a sequence, and processes them with a standard Transformer encoder for supervised image classification.

  • Results

    ViT approaches or beats state of the art across multiple image-recognition benchmarks, reaching 88.55% on ImageNet and using approximately 2−4× less compute than ResNets for equal performance.

  • Takeaways & Limitations

    Large-scale pretraining can outweigh CNN-specific inductive bias, making a simple pure Transformer a competitive and relatively inexpensive approach to image classification.

  • Takeaways & Limitations

    ViT remains to be evaluated more broadly on tasks such as detection and segmentation, and self-supervised pretraining still trails large-scale supervised pretraining.

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

Transformers dominate NLP, but CNNs remain dominant in vision because attention-based alternatives have not scaled effectively. ViT applies a standard Transformer to image-patch sequences and performs strongly when pretrained at large scale and transferred to recognition tasks.

  • Background: Transformer-based NLP systems commonly use large-corpus pretraining followed by fine-tuning on smaller task-specific datasets.Their computational efficiency and scalability have enabled models exceeding 100B parameters.
  • Motivation: CNNs remain dominant in computer vision, while prior attention-based or convolution-replacing models have not scaled effectively on modern hardware accelerators.This motivates testing a minimally modified, pure Transformer architecture for vision.
  • Approach: ViT applies a standard Transformer directly to images by treating linearly embedded image patches as token sequences for supervised classification.The approach makes few architectural modifications and transfers the token-based formulation from NLP to vision.
  • Limitations: Without strong regularization, ViT trained on mid-sized datasets achieves accuracies a few percentage points below comparably sized ResNets.The shortfall is attributed to Transformers lacking CNN inductive biases such as translation equivariance and locality, limiting generalization with insufficient data.
  • Results: When pretrained on 14M–300M images, ViT approaches or beats state of the art across multiple image-recognition benchmarks.The paper reports this result for pretraining on public ImageNet-21k or in-house JFT-300M, concluding that large-scale training can outweigh inductive bias.

2 RELATED WORK

Transformers originated in machine translation and became central to NLP, while vision research explored approximations, CNN-attention hybrids, and pixel-based models to address self-attention’s quadratic image cost. ViT most closely relates to patch-based full-attention models but emphasizes large-scale pre-training to make vanilla Transformers competitive with state-of-the-art CNNs.

  • Large Transformer models typically use large-corpus pre-training followed by task-specific fine-tuning, with BERT using denoising self-supervision and GPT using language modeling.
  • Applying global self-attention directly to image pixels scales quadratically with pixel count, motivating local-attention and other approximations.
  • Cordonnier et al. (2020) apply full self-attention to 2 × 2 image patches, whereas ViT extends this direction by using large-scale pre-training to make vanilla Transformers competitive with or better than state-of-the-art CNNs.The 2 × 2 patch design limits Cordonnier et al. (2020) to small-resolution images.
  • CNN-attention hybrids augment convolutional feature maps or process CNN outputs for tasks including classification, detection, video processing, and object discovery.
  • iGPT trains an unsupervised generative Transformer on resolution- and color-reduced pixels, reaching 72% maximal ImageNet accuracy after fine-tuning or linear probing.
  • Additional data sources have enabled state-of-the-art results on standard benchmarks, alongside empirical studies of CNN scaling and transfer learning from larger datasets.

3 METHOD

Vision Transformer reshapes images into patch-token sequences and processes them with a largely standard Transformer encoder, using minimal image-specific inductive bias. It supports both raw-patch and CNN-feature-map inputs, with pre-training followed by downstream fine-tuning.

  • Transformer encoder: The encoder follows the original Transformer design, alternating multiheaded self-attention and MLP blocks with pre-block layer normalization and post-block residual connections (Vaswani et al., 2017).The MLP uses two layers and a GELU non-linearity.
  • Patch-based input: ViT reshapes an image into a sequence of flattened 2D patches, whose count N = HW/P^2 determines the Transformer’s effective sequence length.The patch sequence is formed by projecting flattened patches into the Transformer dimension.
  • Representation and classification: A learnable [class] token provides the image representation, while learnable 1D position embeddings retain patch order before classification.The [class] token’s final encoder state feeds an MLP head during pre-training and a linear head during fine-tuning.
  • Inductive bias: Unlike CNNs, ViT has much less image-specific inductive bias: self-attention is global, while locality and translation equivariance are not built into every layer.Two-dimensional neighborhood structure is used primarily when forming patches and near the model’s end.
  • Training and variants: ViT is typically pre-trained on large datasets, then fine-tuned on smaller tasks with a new zero-initialized D × K classification layer and potentially higher-resolution inputs.Keeping the patch size fixed at higher resolution increases the effective sequence length; a hybrid alternative derives input patches from CNN feature maps, including 1x1 spatial patches.

4 EXPERIMENTS

The experiments evaluate ViT, ResNet, and hybrid models across varying pre-training scales and many transfer benchmarks. ViT achieves strong recognition performance, often with substantially lower pre-training compute, while scaling and attention analyses characterize its behavior.

  • 4 EXPERIMENTS: The evaluation uses ImageNet, ImageNet-21k, and JFT, then transfers models to multiple downstream recognition benchmarks with deduplicated pre-training data.The datasets contain approximately 1.3M, 14M, and 303M images respectively, with 1k, 21k, and 18k classes.
  • 4.1 Main results: ViT-L/16 pre-trained on JFT-300M outperforms BiT-L on all reported tasks with substantially less training compute, while ViT-H/14 improves further on challenging benchmarks.Table 2 reports mean and standard deviation over three fine-tuning runs; ViT-H/14 particularly improves results on ImageNet, CIFAR-100, and VTAB.
  • 4.1 Main results: ViT-H/14 outperforms BiT-R152x4 and other prior methods on VTAB’s Natural and Structured groups, while the top two models perform similarly on Specialized tasks.VTAB contains 19 low-data classification tasks with 1,000 training examples per task, spanning Natural, Specialized, and Structured groups.
  • 4.2 Data requirements: BiT CNNs outperform ViT on ImageNet at smaller pre-training scales, but ViT overtakes them when trained on larger datasets.The experiments vary pre-training dataset size to assess how ViT’s lower inductive bias affects data requirements.
  • 4.4 Pre-training compute: ViT uses approximately 2–4× less pre-training compute than ResNets to attain the same performance on average across five datasets.Hybrids slightly outperform ViT at small computational budgets, while the broader performance/compute trade-off favors Vision Transformers.
  • 4.5 Analysis: Attention analyses show that some ViT heads integrate information across most of the image in the lowest layers, while learned position embeddings encode spatial distance and row-column structure.Attention distance is measured analogously to CNN receptive-field size, and closer patches tend to have more similar position embeddings.

5 CONCLUSION · APPENDIX · A MULTIHEAD SELF-ATTENTION

The paper shows that a standard Transformer can process image-patch sequences without image-specific inductive biases beyond patch extraction, particularly when pretrained on large datasets. It also identifies detection, segmentation, and improved self-supervised pretraining as remaining challenges, while defining multihead self-attention as parallel attention operations whose outputs are concatenated and projected.

  • 5 CONCLUSION: A standard Transformer encoder can recognize images by processing patch sequences, without image-specific inductive biases beyond initial patch extraction.The approach works well when coupled with pretraining on large datasets.
  • 5 CONCLUSION: Applying the approach to detection and segmentation remains an important open challenge.The authors cite their results and Carion et al. (2020) as indicating promise for this direction.
  • 5 CONCLUSION: Self-supervised pretraining improves initial results, but still trails large-scale supervised pretraining.The authors identify closing this gap as a remaining challenge.
  • A MULTIHEAD SELF-ATTENTION: Standard self-attention computes weighted sums over sequence values using attention weights based on pairwise similarities between queries and keys.For each input element, the operation uses its query and the corresponding keys to determine weighting.
  • APPENDIX: Multihead self-attention runs k self-attention operations in parallel and projects their concatenated outputs.The heads are separate attention operations whose outputs are combined before the final projection.
  • A MULTIHEAD SELF-ATTENTION: When the number of heads changes, setting D_h to D/k typically keeps computation and parameter count constant.This scaling adjusts the per-head dimensionality as the number of parallel heads changes.

B EXPERIMENT DETAILS … B.1.2 SELF-SUPERVISION

The experiments emphasize strong regularization, standardized fine-tuning, and masked patch prediction, with self-supervision showing similar ImageNet gains using 100k steps or ImageNet pretraining rather than requiring JFT-scale training. Fine-tuning replaces the pretrained head and uses controlled validation-based hyperparameter selection across transfer tasks.

  • B.1 TRAINING: Models trained from scratch on ImageNet required strong regularization, with dropout inserted after dense layers and after adding positional to patch embeddings.Dropout was excluded from qkv-projections, and hybrid models used the exact setup of their ViT counterparts; training resolution was 224.
  • B.1.1 FINE-TUNING: ViT fine-tuning used SGD with momentum 0.9, selected learning rates through development-set sweeps, and retrained final models on the full training sets before test evaluation.Development subsets were 10% for Pets and Flowers, 2% for CIFAR, and 1% for ImageNet.
  • B.1.1 FINE-TUNING: Fine-tuning generally used cosine decay, batch size 512, no weight decay, global-norm gradient clipping at 1, and resolution 384 rather than the 224 training resolution.The same setup was used for ResNets and hybrid models except for the stated ResNet-specific additional run.
  • B.1.1 FINE-TUNING: When transferring ViT models, the two-layer head was replaced by a zero-initialized linear layer sized for the target classes, which was more robust than reinitializing only the final layer.This transfer-specific replacement removes the entire original head.
  • B.1.1 FINE-TUNING: VTAB used one shared hyperparameter setting across tasks: learning rate 0.01 for 2500 steps, selected by a small sweep using 200-example validation sets.The protocol followed Kolesnikov et al. (2020), including its preprocessing except for task-specific input resolutions.
  • B.1.2 SELF-SUPERVISION: The masked patch objective corrupted 50% of patch embeddings and predicted each corrupted patch’s 3-bit mean color from its representation.Corruption used a learnable [mask] embedding 80% of the time, a random patch 10%, and no replacement 10%; the model was trained for 1M steps on JFT with Adam.
  • B.1.2 SELF-SUPERVISION: Self-supervised masked patch prediction showed similar ImageNet performance gains after roughly 100k pretraining steps or when pretrained on ImageNet, with diminishing returns beyond 100k steps.The authors therefore state that this approach does not require 1M steps or a large dataset such as JFT to achieve similar gains.

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

Additional results detail how pre-training scale and model size affect ViT transfer, while analyses justify Adam pre-training, identify depth as the strongest scaling dimension, and clarify classifier-head behavior. These findings span C ADDITIONAL RESULTS and D.1–D.3 analyses.

  • C ADDITIONAL RESULTS: Tables 5 and 6 provide detailed ViT transfer results across pre-training datasets, model scales, evaluation datasets, and estimated pre-training compute.Table 5 covers ImageNet, ImageNet-21k, and JFT-300M pre-training; Table 6 reports transfer accuracy and compute for scaling experiments.
  • D.1 SGD VS. ADAM FOR RESNETS: Adam pre-training outperforms SGD pre-training on most datasets and on average for the evaluated ResNet models, supporting Adam for JFT pre-training.The comparison uses ResNet-50x1 and ResNet-152x2 and only seven pre-training epochs, so absolute results are below those reported for 30 epochs.
  • D ADDITIONAL ANALYSES: The additional analyses examine optimizer choice, Transformer shape, and classifier-head design as factors affecting ViT and ResNet transfer performance.These analyses include ResNet optimizer comparisons, Transformer dimension-scaling ablations, and class-token versus global-average-pooling evaluations.
  • D.2 TRANSFORMER SHAPE: Scaling Transformer depth produces the largest ImageNet 5-shot improvements, with gains clearly visible through 64 layers.The ablation fixes an eight-layer, D = 1024, DMLP = 2048, patch-size-32 configuration at the intersection of all comparisons.
  • D.3 HEAD TYPE AND CLASS TOKEN: ViT uses a text-Transformer-style class token whose output is passed through a one-hidden-layer tanh MLP for classification.The class token is used as the image representation throughout the main paper.
  • D.3 HEAD TYPE AND CLASS TOKEN: Class-token and global-average-pooling classifiers perform similarly, but they require different learning rates.A preliminary GAP-plus-linear-classifier setup performed poorly, but the performance difference was attributed to learning-rate requirements rather than the token or pooling operation itself.

D.4 POSITIONAL EMBEDDING

The positional-embedding ablation found that adding positional information substantially improves ViT-B/16 performance, while the specific encoding scheme has little effect. The study compared absolute, relative, one-dimensional, two-dimensional, and layerwise positional-embedding strategies.

  • Encoding schemes: The ablation compared no positional information, one-dimensional and two-dimensional absolute embeddings, and relative embeddings based on pairwise patch distances.Two-dimensional embeddings concatenate learned X- and Y-axis embeddings, each of size D/2; relative embeddings associate offsets p_q − p_k with embeddings.
  • Embedding incorporation: The default one-dimensional embedding treats patches as a raster-order sequence and is added after the model stem before Transformer encoding.The study also tested adding embeddings at each layer, either separately or with a learned embedding shared across layers.
  • Results: Adding positional embeddings creates a large performance gain over treating patches as a bag, while encoding choices produce little to no difference in ViT-B/16.These results are summarized in Table 8; the authors speculate that patch-level rather than pixel-level inputs make spatial-encoding differences less important.

D.5 EMPIRICAL COMPUTATIONAL COSTS

On TPUv3 hardware, ViT models achieve inference speed comparable to similar ResNets while offering clearly better memory efficiency. The expected bi-quadratic scaling with image size appears only for the largest models at the largest resolutions.

  • D.5 EMPIRICAL COMPUTATIONAL COSTS: Inference speed was measured for the main models on a TPUv3 accelerator, with backpropagation speed differing by a constant model-independent factor.The measurements target real-world hardware behavior, which theoretical FLOPs may not predict because of lane widths and cache sizes.
  • D.5 EMPIRICAL COMPUTATIONAL COSTS: The theoretical bi-quadratic scaling of ViT with image size only barely appears for the largest models at the largest resolutions.Throughput measurements use peak performance across a wide range of batch sizes and input sizes.
  • D.5 EMPIRICAL COMPUTATIONAL COSTS: ViT models have speed comparable to similar ResNets and are clearly more memory-efficient in largest per-core batch-size capacity.Figure 12 reports real wall-clock timings and the largest batch-size fitting on one device core across input sizes.

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

The section evaluates axial attention as a 2D alternative for ViT, analyzes how attention distance changes across depth, and describes attention-rollout maps for tracing information flow.

  • D.6 AXIAL ATTENTION: Axial attention applies separate self-attention operations along individual tensor axes rather than flattening the input into one-dimensional sequences.Each operation mixes information along one axis while preserving the others.
  • D.6 AXIAL ATTENTION: The modified ViT processes inputs in 2D using Axial Transformer blocks that apply row self-attention plus an MLP followed by column self-attention plus an MLP.
  • D.6 AXIAL ATTENTION: The naive AxialResNet implementation is extremely slow on TPUs, limiting extensive large-scale experiments despite reproducing reported accuracy scores.The authors note that a carefully optimized implementation may unlock larger experiments.
  • D.6 AXIAL ATTENTION: Axial-ViT-B/32 and Axial-ViT-B/16 outperform their ViT-B counterpart on ImageNet 5-shot linear, but require more compute.The comparison is reported against pretraining compute in FLOPs and inference time; Figure 13 presents the corresponding accuracy–compute trade-offs.
  • D.6 AXIAL ATTENTION: The axial models’ additional compute arises because each global-attention Transformer block is replaced by separate row and column blocks, each containing an MLP.Although axial self-attention uses shorter sequences, the extra MLP contributes to the higher cost.
  • D.7 ATTENTION DISTANCE: Attention distance varies substantially across lower-layer heads, with some attending broadly across the image and others focusing near the query location.Attention distance is interpreted analogously to receptive-field size in CNNs.
  • D.7 ATTENTION DISTANCE: As depth increases, attention distance increases across heads, indicating progressively broader spatial integration in ViT.
  • D.8 ATTENTION MAPS: Attention maps are computed with Attention Rollout by averaging ViT-L/16 attention across heads and recursively multiplying layer weight matrices.This procedure accounts for attention mixing across tokens through all layers and produces maps from the output token to the input space.

D.9 OBJECTNET RESULTS · D.10 VTAB BREAKDOWN

ViT-H/14 achieves 82.1% top-5 and 61.7% top-1 accuracy on ObjectNet under the evaluation setup of Kolesnikov et al. (2020). The VTAB-1k results are broken down by task in Table 9, with additional example attention maps shown in Figure 14.

  • D.9 OBJECTNET RESULTS: 82.1% top-5 and 61.7% top-1 accuracy are achieved by ViT-H/14 on ObjectNet.The evaluation follows the setup in Kolesnikov et al. (2020).
  • D.9 OBJECTNET RESULTS: The ObjectNet evaluation uses the flagship ViT-H/14 model.The reported benchmark results are specifically for ViT-H/14.
  • D.9 OBJECTNET RESULTS: ObjectNet performance is reported using both top-5 and top-1 accuracy.The two reported values are 82.1% top-5 accuracy and 61.7% top-1 accuracy.
  • D.9 OBJECTNET RESULTS: The ObjectNet evaluation follows the setup in Kolesnikov et al. (2020).This benchmark result is presented under the cited prior evaluation protocol.
  • D.10 VTAB BREAKDOWN: Table 9 reports scores for each VTAB-1k task.The supplied passage identifies the table as the source of the per-task VTAB-1k breakdown.
  • D.10 VTAB BREAKDOWN: Figure 14 presents further randomly selected example attention maps as in Figure 6.The figure is described as an extension of the attention-map examples shown in Figure 6.
  • D.10 VTAB BREAKDOWN: Table 9 provides a breakdown of VTAB-1k performance across tasks.No individual task scores are supplied in the passage.
Loading 2010.11929v2…