Skip to content

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (Vision Transformer)

Authors: Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, et al. (Google Research) Published: October 2020 (ICLR 2021) Paper: arxiv.org/abs/2010.11929

Why This Matters

Vision Transformer (ViT) revolutionized computer vision by proving that the pure Transformer architecture—without any convolutions—could match or exceed state-of-the-art CNNs on image classification. This breakthrough:

  • Unified vision and language: The same architecture now works for both domains
  • Enabled multimodal models: Paved the way for models like CLIP, Flamingo, and GPT-4 Vision
  • Changed computer vision research: Shifted focus from CNN architectures to Transformers
  • Improved scalability: Transformers scale better with data than CNNs

Real-world impact: Powers modern vision models in autonomous vehicles, medical imaging, content moderation, visual search, and all major multimodal AI systems.

The Problem

Before ViT: - Computer vision relied almost exclusively on Convolutional Neural Networks (CNNs) - CNNs have inductive biases (assumptions about locality and translation equivariance) baked into their architecture - While Transformers dominated NLP, applying them to vision seemed impractical due to: - Images have many pixels (high resolution = huge computational cost) - Self-attention scales quadratically with sequence length - No clear way to tokenize images

The question: Could pure Transformers work for vision if we had enough data and compute?

Core Innovation

Treating Images as Sequences of Patches

The key insight: Split an image into fixed-size patches and treat them like words in a sentence.

How it works:

  1. Patch Embedding
  2. Split image (e.g., 224×224) into 16×16 patches
  3. Results in 14×14 = 196 patches
  4. Flatten each patch into a vector

  5. Linearly project to embedding dimension

  6. Position Embeddings

  7. Add learnable position embeddings (like in NLP Transformers)

  8. Tells the model where each patch is located

  9. Class Token

  10. Prepend a special [CLS] token (borrowed from BERT)

  11. Its final representation used for classification

  12. Standard Transformer

  13. Feed patch embeddings through standard Transformer encoder
  14. Multi-head self-attention learns relationships between patches
  15. MLP layers process the representations

Mathematical representation: 

Input image: x ∈ R^(H×W×C)
Patches: x_p ∈ R^(N×(P²·C))
  where N = HW/P² (number of patches)

Patch embeddings: z_0 = [x_class; x_p¹E; x_p²E; ...; x_pᴺE] + E_pos

Transformer encoder: z_ℓ = Transformer(z_(ℓ-1))

Classification: y = LayerNorm(z_L⁰)

Architecture Details

Standard ViT Configuration

ViT-Base (ViT-B/16): - Patch size: 16×16 - Embedding dimension: 768 - Layers: 12 - Attention heads: 12 - MLP hidden size: 3072 - Parameters: ~86M

ViT-Large (ViT-L/16): - Embedding dimension: 1024 - Layers: 24 - Attention heads: 16 - Parameters: ~307M

ViT-Huge (ViT-H/14): - Patch size: 14×14 - Embedding dimension: 1280 - Layers: 32 - Attention heads: 16 - Parameters: ~632M

Why This Works

  1. Self-attention captures global context
  2. Unlike CNNs with limited receptive fields

  3. Each patch can attend to all other patches from layer 1

  4. Fewer inductive biases

  5. Model must learn spatial relationships from data

  6. With enough data, this is more flexible than hard-coded CNN biases

  7. Scalability

  8. Transformers scale better with data than CNNs
  9. Performance continues improving with larger datasets

Training Approach

Pre-training on Large Datasets

Critical finding: ViT requires more data than CNNs to reach comparable performance.

Datasets used: - ImageNet-21k: 14M images, 21k classes - JFT-300M: 300M images, 18k classes (Google internal)

Training details: - Optimizer: Adam with β₁=0.9, β₂=0.999 - Learning rate: Linear warmup + cosine decay - Regularization: Dropout, weight decay - Data augmentation: RandAugment, MixUp

Fine-tuning on Downstream Tasks

After pre-training: 1. Remove pre-training classification head 2. Add new head for target task 3. Fine-tune at higher resolution (e.g., 384×384) 4. Use fewer training steps than pre-training

Key trick: When fine-tuning at higher resolution, interpolate position embeddings to match new patch count.

Results and Impact

Performance Comparison

ImageNet (top-1 accuracy): - ViT-H/14 (JFT-300M pre-training): 88.55% - BiT-L (ResNet-152x4): 87.54% - Noisy Student (EfficientNet): 88.4%

Transfer learning (average over 19 tasks): - ViT-H/14: 77.91% - BiT-L: 76.08%

Computational Efficiency

Pre-training cost: - ViT-H/14: ~2.5k TPUv3-core-days - BiT-L: ~9.9k TPUv3-core-days

ViT is 4× more compute-efficient than ResNet-based models!

What Makes ViT Better?

  1. Scales better with data
  2. Performance gap widens with larger pre-training datasets

  3. CNNs plateau faster

  4. More efficient attention patterns

  5. Lower layers: Local attention (similar to convolution)

  6. Higher layers: Global attention (captures long-range dependencies)

  7. Transfer learning

  8. Pre-trained ViT transfers better across diverse tasks

Attention Visualizations

Key findings from attention analysis:

  1. Spatial awareness emerges naturally
  2. Model learns to attend to spatially nearby patches early

  3. No convolution needed to learn locality

  4. Global context from layer 1

  5. Some attention heads attend globally even in early layers

  6. Captures long-range dependencies CNNs struggle with

  7. Semantic grouping

  8. Attention heads learn to group semantically similar regions
  9. Example: All patches of "dog" attend to each other

Limitations

1. Data Hungry

  • Requires large-scale pre-training (millions of images)
  • On small datasets (e.g., ImageNet-1k only), CNNs perform better
  • Why: Fewer inductive biases means more data needed to learn patterns

2. Computational Cost

  • Quadratic complexity in number of patches
  • High-resolution images expensive (e.g., 1024×1024 medical images)
  • Solution: Hierarchical variants like Swin Transformer

3. Less Interpretable Position Embeddings

  • Position embeddings are learned, not fixed
  • Less clear how model understands spatial structure
  • 2D structure not explicitly encoded

4. Fixed Input Size During Training

  • Must interpolate position embeddings for different resolutions
  • Not as naturally flexible as CNNs with varying input sizes

Practical Applications

Where ViT Excels

  1. Image Classification
  2. ImageNet, iNaturalist, CIFAR benchmarks

  3. Medical image classification

  4. Object Detection

  5. ViT backbone in Detection Transformers (DETR)

  6. Faster R-CNN with ViT features

  7. Semantic Segmentation

  8. Vision Transformer backbones in segmentation models

  9. Medical imaging segmentation

  10. Multimodal Models

  11. CLIP: Vision encoder using ViT
  12. Flamingo: Visual language models

  13. GPT-4 Vision: Image understanding

  14. Video Understanding

  15. Video Vision Transformer (ViViT)
  16. Temporal attention over frame patches

Variants and Improvements

Hierarchical Vision Transformers

Swin Transformer (2021): - Shifted window attention (local then global) - Hierarchical feature maps (like CNNs) - Better for dense prediction tasks

PVT (Pyramid Vision Transformer) (2021): - Multi-scale feature pyramid - Spatial-reduction attention for efficiency

Efficient Vision Transformers

DeiT (Data-efficient ViT) (2020): - Distillation from CNN teachers - Works well with ImageNet-1k only

Mobile-ViT (2021): - Lightweight for mobile devices - Combines convolutions and transformers

Hybrid Approaches

Early Convolutions + Transformers: - Use CNN stem to reduce patch count - Then apply Transformer - Combines CNN inductive biases with Transformer flexibility

Implementation Details

Pseudocode

class VisionTransformer:
    def __init__(self, image_size=224, patch_size=16, num_classes=1000,
                 dim=768, depth=12, heads=12, mlp_dim=3072):
        self.patch_embed = PatchEmbedding(patch_size, dim)
        self.cls_token = Parameter(torch.randn(1, 1, dim))
        self.pos_embed = Parameter(torch.randn(1, num_patches + 1, dim))

        self.transformer = TransformerEncoder(dim, depth, heads, mlp_dim)
        self.classifier = Linear(dim, num_classes)

    def forward(self, img):
        # Split into patches and embed
        x = self.patch_embed(img)  # (B, N, dim)

        # Prepend class token
        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)  # (B, N+1, dim)

        # Add position embeddings
        x = x + self.pos_embed

        # Apply Transformer
        x = self.transformer(x)

        # Classify using class token
        return self.classifier(x[:, 0])

Key Hyperparameters

Pre-training: - Batch size: 4096 - Learning rate: 0.001 (with linear warmup) - Weight decay: 0.1 - Optimizer: Adam - Augmentation: RandAugment(2, 15)

Fine-tuning: - Batch size: 512 - Learning rate: 0.003 (higher resolution: 0.01) - Fine-tuning steps: ~20k (much less than pre-training)

Key Takeaways

  1. Pure Transformers work for vision when trained on sufficient data
  2. Patches as tokens is the key insight for applying Transformers to images
  3. Scales better than CNNs with increased data and model size
  4. Unified architecture enables seamless multimodal learning
  5. Pre-training is critical - ViT needs large-scale pre-training to excel

Influence on Later Work

Direct Descendants

  • CLIP (2021): ViT as image encoder
  • DALL-E (2021): ViT-based image tokenization
  • Swin Transformer (2021): Hierarchical ViT
  • BEiT (2021): BERT-style pre-training for images
  • MAE (2021): Masked autoencoder for images

Conceptual Impact

  • Showed Transformers are universal architectures (not just for sequences)
  • Inspired transformer applications in video, 3D, audio, biology, reinforcement learning
  • Established that data > inductive biases at scale

Further Reading

Original Paper

  • An Image is Worth 16x16 Words: https://arxiv.org/abs/2010.11929

Follow-up Papers

  • DeiT: https://arxiv.org/abs/2012.12877 (Data-efficient training)
  • Swin Transformer: https://arxiv.org/abs/2103.14030 (Hierarchical ViT)
  • BEiT: https://arxiv.org/abs/2106.08254 (Self-supervised pre-training)
  • MAE: https://arxiv.org/abs/2111.06377 (Masked image modeling)

Code Implementations

  • Official JAX implementation: https://github.com/google-research/vision_transformer
  • PyTorch (timm): https://github.com/huggingface/pytorch-image-models
  • Hugging Face Transformers: https://huggingface.co/docs/transformers/model_doc/vit

Tutorials

  • The Illustrated Vision Transformer: https://jalammar.github.io/illustrated-vit/
  • ViT Paper Walkthrough (Yannic Kilcher): YouTube search
  • Fine-tuning ViT on custom datasets: Hugging Face documentation

Published: October 2020 Impact Factor: 10,000+ citations (as of 2025) Legacy: Made Transformers the dominant architecture in computer vision, enabling the multimodal AI revolution.