An Image is Worth 16x16 Words: Mastering the Vision Transformer (ViT)
An Image is Worth 16x16 Words: Mastering the Vision Transformer (ViT)
For decades, Convolutional Neural Networks (CNNs) were the undisputed kings of computer vision. Their inherent "inductive biases"—specifically locality (pixels near each other are related) and translation equivariance (a cat is a cat regardless of where it is in the image)—made them incredibly efficient for image processing.
But what happens if we strip those biases away? What if we treat an image not as a grid of pixels, but as a sequence of words?
Enter the Vision Transformer (ViT). In this post, we will dive deep into the architecture that proved Transformers could scale to image recognition, breaking down the theory, the math, and a full PyTorch implementation.
The Core Intuition: Images as Sequences
The fundamental challenge of applying a Transformer to an image is that Transformers were designed for 1D sequences (like text), while images are 2D grids.
The ViT solves this with a clever trick: Patch Partitioning. Instead of processing individual pixels, the model breaks an image into fixed-size square patches. If you have a $224 \times 224$ image and a patch size of $16 \times 16$, you end up with $14 \times 14 = 196$ patches. Each patch is then flattened and treated as a "visual token," effectively turning the image into a sentence where each "word" is a piece of the picture.
The High-Level Architecture
Technical Deep Dive
1. From Pixels to Embeddings
To convert the raw image into a format the Transformer understands, ViT follows these steps:
- Patching: The image $\mathbf{x} \in \mathbb{R}^{H \times W \times C}$ is split into $N$ patches.
- Linear Projection: Each patch is flattened into a vector $\mathbf{x}_p \in \mathbb{R}^{N \times (P^2 \cdot C)}$ and projected into a latent dimension $D$.
- The [CLS] Token: Borrowing from BERT, a learnable classification token $\mathbf{x}_{class}$ is prepended to the sequence. The model uses this token to aggregate information from all other patches.
- Positional Encoding: Since Transformers have no inherent sense of order, learnable 1D position embeddings $\mathbf{E}$ are added to the tokens.
The resulting input sequence $\mathbf{z}_0$ is defined as: $$\mathbf{z}0 = [\mathbf{x}{class}; \mathbf{x}_p \mathbf{E}]$$
2. The Transformer Encoder
The sequence passes through $L$ layers of Transformer blocks. Each block consists of:
- Multi-Head Self-Attention (MSA): Allows every patch to "attend" to every other patch, capturing global dependencies regardless of distance.
- Multi-Layer Perceptron (MLP): A feed-forward network that processes each token independently.
- Layer Normalization & Residuals: Ensures training stability and prevents vanishing gradients.
3. The Classification Head
After $L$ layers, we ignore all patch tokens and extract only the final state of the classification token $\mathbf{z}_L^0$. This vector is passed through a final MLP head to predict the class: $$\text{Output} = \text{MLP}(\mathbf{z}_L^0)$$
Implementation in PyTorch
Below is a production-ready implementation of the Vision Transformer. For demonstration purposes, this code is configured to run on the CIFAR-10 dataset.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
class PatchEmbedding(nn.Module):
"""
Splits image into patches and projects them into a D-dimensional embedding space.
Implemented via Conv2d for efficiency.
"""
def __init__(self, img_size: int, patch_size: int, in_chans: int, embed_dim: int):
super().__init__()
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# [Batch, C, H, W] -> [Batch, D, H/P, W/P]
x = self.proj(x)
# [Batch, D, N] -> [Batch, N, D]
x = x.flatten(2).transpose(1, 2)
return x
class TransformerBlock(nn.Module):
"""A single Transformer Encoder layer."""
def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=num_heads, batch_first=True)
self.norm2 = nn.LayerNorm(dim)
self.mlp = nn.Sequential(
nn.Linear(dim, int(dim * mlp_ratio)),
nn.GELU(),
nn.Linear(int(dim * mlp_ratio), dim),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Residual 1: Norm -> MSA -> Add
x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
# Residual 2: Norm -> MLP -> Add
x = x + self.mlp(self.norm2(x))
return x
class VisionTransformer(nn.Module):
def __init__(self, img_size=32, patch_size=4, in_chans=3, num_classes=10,
embed_dim=128, depth=6, num_heads=8, mlp_ratio=4.0):
super().__init__()
self.patch_embed = PatchEmbedding(img_size, patch_size, in_chans, embed_dim)
num_patches = (img_size // patch_size) ** 2
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
self.blocks = nn.Sequential(
*[TransformerBlock(embed_dim, num_heads, mlp_ratio) for _ in range(depth)]
)
self.norm = nn.LayerNorm(embed_dim)
self.head = nn.Linear(embed_dim, num_classes)
def forward(self, x: torch.Tensor) -> torch.Tensor:
b = x.shape[0]
x = self.patch_embed(x)
cls_tokens = self.cls_token.expand(b, -1, -1)
x = torch.cat((cls_tokens, x), dim=1)
x = x + self.pos_embed
x = self.blocks(x)
x = self.norm(x)
return self.head(x[:, 0]) # Classify based on [CLS] tokenKey Takeaways & Performance
Why does ViT work?
The magic of ViT isn't that it's "better" than CNNs at small scales—in fact, ViT typically performs worse than CNNs when trained on small datasets because it lacks the inductive biases of convolutions.
However, when trained on massive datasets (like JFT-300M), ViT outperforms CNNs. This suggests that the Transformer architecture has a higher "capacity" to learn complex patterns if provided with enough data to learn the spatial relationships from scratch.
Summary Table: CNN vs. ViT
| Feature | Convolutional Neural Network (CNN) | Vision Transformer (ViT) |
|---|---|---|
| Basic Unit | Convolutional Filter | Self-Attention Mechanism |
| Inductive Bias | Locality & Translation Equivariance | Minimal (Learns from data) |
| Receptive Field | Local (increases with depth) | Global (from the first layer) |
| Data Requirement | Moderate | Very High (needs pre-training) |
| Complexity | Linear with image resolution | Quadratic with number of patches |
Final Thoughts
The Vision Transformer represents a paradigm shift in computer vision, unifying the architectures of NLP and CV. By treating images as sequences, we open the door to multimodal models (like CLIP or GPT-4V) that can "read" and "see" using the same underlying mathematical engine.