Paper: "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" (Dosovitskiy et al., 2020) ArXiv: https://arxiv.org/abs/2010.11929 Key Insight: Replace convolutions with patch embeddings + transformers. When trained on large datasets (JFT-300M), ViTs outperform CNNs on ImageNet despite having no inductive biases.
The Problem: Why Convolution?
Convolutions work great because they have inductive biases:
- Locality: Pixels close together are related
- Translation equivariance: Same features everywhere
- Parameter sharing: Fewer parameters
But these biases limit the model's expressiveness. What if we remove them entirely?
Vision Transformer Architecture
Image (224×224)
↓
Split into patches (16×16): 196 patches
↓
Linear embedding: 196 × 768
↓
Add position embedding: (196, 768)
↓
Add [CLS] token: (197, 768)
↓
Transformer blocks (12 layers, 12 heads)
↓
[CLS] token output
↓
Linear classification head
↓
Logits (1000 classes)
Implementation
import torch
import torch.nn as nn
import math
class PatchEmbedding(nn.Module):
def __init__(self, image_size=224, patch_size=16, in_channels=3, embed_dim=768):
super().__init__()
self.image_size = image_size
self.patch_size = patch_size
self.num_patches = (image_size // patch_size) ** 2
# Convert patches to embedding
self.proj = nn.Conv2d(
in_channels,
embed_dim,
kernel_size=patch_size,
stride=patch_size
)
def forward(self, x):
# x: (batch, 3, 224, 224)
# Conv with kernel=16, stride=16 splits into patches
x = self.proj(x) # (batch, 768, 14, 14)
# Flatten: (batch, 768, 196)
x = x.flatten(2) # (batch, 768, 196)
# Transpose: (batch, 196, 768)
x = x.transpose(1, 2)
return x
class VisionTransformer(nn.Module):
def __init__(self, image_size=224, patch_size=16, num_classes=1000, embed_dim=768, depth=12, heads=12):
super().__init__()
# Patch embedding
self.patch_embed = PatchEmbedding(image_size, patch_size, 3, embed_dim)
num_patches = self.patch_embed.num_patches
# Class token
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
# Position embedding
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
# Transformer
self.transformer = nn.TransformerEncoder(
nn.TransformerEncoderLayer(d_model=embed_dim, nhead=heads, dim_feedforward=3072),
num_layers=depth
)
# Classification head
self.classify = nn.Linear(embed_dim, num_classes)
def forward(self, x):
# Patch embedding
x = self.patch_embed(x) # (batch, 196, 768)
batch_size = x.size(0)
# Add class token
cls = self.cls_token.expand(batch_size, -1, -1) # (batch, 1, 768)
x = torch.cat([cls, x], dim=1) # (batch, 197, 768)
# Add position embedding
x = x + self.pos_embed
# Transformer
x = self.transformer(x) # (batch, 197, 768)
# Classification: use [CLS] token
cls_output = x[:, 0, :] # (batch, 768)
logits = self.classify(cls_output) # (batch, 1000)
return logits
Key Results
ImageNet-1K Classification Accuracy:
ResNet-50 (CNN, trained on ImageNet):
- Accuracy: 76.5%
ViT-Base (trained on ImageNet only):
- Accuracy: 77.9% (better!)
ViT-Base (trained on JFT-300M, then fine-tuned):
- Accuracy: 88.2% (much better!)
Key insight: With large-scale pretraining, ViTs win!
Why ViTs Work
CNNs:
- Efficient with small data (inductive biases)
- Limited representational power
- Poor scaling with data
ViTs:
- Flexible, no hand-crafted biases
- Scales with data
- Better long-range interactions (full attention)
Practical Implications
ViTs require:
- Large training data: ImageNet alone is insufficient (22M→300M helps)
- Longer training: More epochs needed to converge
- Strong regularization: Dropout, data augmentation
But once trained, ViTs transfer better and enable new architectures (like CLIP, LLaVA).
Our Analysis
Vision Transformers are a paradigm shift: they show that inductive biases aren't necessary if you have data. This opened up entire new research directions (multimodal models, masked image modeling). The lesson: with enough data, generic architectures (transformers) beat specialized ones (CNNs).
References
- Paper: An Image is Worth 16x16 Words (Dosovitskiy et al., 2020)
- Code: https://github.com/google-research/vision_transformer
Conclusion
Vision Transformers demonstrate that attention is general: it works for vision, language, and beyond. Understanding patch embeddings and positional encoding for images enables building powerful multimodal systems. This is the foundation of modern vision-language models. Next: understanding modality alignment in multimodal learning.