class PatchEmbedding(nn.Module):
"""Image-to-sequence stem of the vision transformer."""
def __init__(self, img_size=96, patch_size=16, num_hiddens=512):
super().__init__()
def _make_tuple(x):
if not isinstance(x, (list, tuple)):
return (x, x)
return x
img_size, patch_size = _make_tuple(img_size), _make_tuple(patch_size)
# A partial trailing patch would silently shrink the token grid
assert img_size[0] % patch_size[0] == 0 and \
img_size[1] % patch_size[1] == 0, \
'image size must be divisible by the patch size'
self.num_patches = (img_size[0] // patch_size[0]) * (
img_size[1] // patch_size[1])
self.conv = nn.LazyConv2d(num_hiddens, kernel_size=patch_size,
stride=patch_size)
def forward(self, X):
# Output shape: (batch size, no. of patches, no. of channels)
return self.conv(X).flatten(2).transpose(1, 2)Vision Transformer
Dive into Deep Learning · §11.5
Vision transformer
patches as tokens · the pre-norm block at work · what small scale teaches
Images were CNN country
Nothing in the block asks what its tokens are
CNNs owned vision on the strength of two built-in commitments: locality and translation equivariance.
Early hybrids kept them: specialized local self-attention (Ramachandran et al., 2019) was hard to run fast; Cordonnier et al. (2020) showed self-attention can learn to act like convolution — with 2×2 patches.
ViT (Dosovitskiy et al., 2021) dropped the caution: 16×16 patches, one embedding per patch, an unmodified transformer encoder. On 300M images it beat the best CNNs.
Architecture
Patchify = one strided convolution
“Split into patches, then linearly project” is a single convolution with kernel_size = stride = patch_size — each output position sees exactly one patch:
Shape check
A 96×96 image with 16×16 patches becomes a sequence of (96/16)^2 = 36 tokens of width num_hiddens:
This is the entire interface between vision and the transformer.
The ViT block
The pre-norm block of the chapter, with two vision-era details: GELU in the MLP and dropout after each of its linear layers — at this scale, regularization earns its keep. No mask: every patch attends to every patch.
class ViTBlock(nn.Module):
"""Pre-norm transformer block with a GELU MLP."""
def __init__(self, num_hiddens, mlp_num_hiddens, num_heads, dropout,
use_bias=False):
super().__init__()
self.ln1 = nn.LayerNorm(num_hiddens)
self.attention = d2l.MultiHeadAttention(num_hiddens, num_heads,
dropout, use_bias)
self.ln2 = nn.LayerNorm(num_hiddens)
self.mlp = nn.Sequential(
nn.LazyLinear(mlp_num_hiddens), nn.GELU(), nn.Dropout(dropout),
nn.LazyLinear(num_hiddens), nn.Dropout(dropout))
def forward(self, X, valid_lens=None):
X = X + self.attention(*([self.ln1(X)] * 3), valid_lens)
return X + self.mlp(self.ln2(X))The full model
Patch embed → prepend learnable <cls> token → add learned position embeddings (initialized to small noise — no 2-D structure given!) → dropout → blocks → classify from the <cls> representation:
class ViT(d2l.Classifier):
"""Vision transformer."""
def __init__(self, img_size, patch_size, num_hiddens, mlp_num_hiddens,
num_heads, num_blks, emb_dropout, blk_dropout, lr=0.1,
use_bias=False, num_classes=10):
super().__init__()
self.save_hyperparameters()
self.patch_embedding = PatchEmbedding(
img_size, patch_size, num_hiddens)
self.cls_token = nn.Parameter(d2l.zeros(1, 1, num_hiddens))
num_steps = self.patch_embedding.num_patches + 1 # Add the cls token
# Positional embeddings are learnable, initialized to small noise
self.pos_embedding = nn.Parameter(
torch.randn(1, num_steps, num_hiddens) * 0.02)
self.dropout = nn.Dropout(emb_dropout)
self.blks = nn.Sequential()
for i in range(num_blks):
self.blks.add_module(f"{i}", ViTBlock(
num_hiddens, mlp_num_hiddens, num_heads, blk_dropout,
use_bias))
self.head = nn.Sequential(nn.LayerNorm(num_hiddens),
nn.Linear(num_hiddens, num_classes))
def forward(self, X):
X = self.patch_embedding(X)
X = d2l.concat((self.cls_token.expand(X.shape[0], -1, -1), X), 1)
X = self.dropout(X + self.pos_embedding)
for blk in self.blks:
X = blk(X)
return self.head(X[:, 0])Training on Fashion-MNIST
2 blocks, width 512, 8 heads — about 6.5M parameters; ten epochs, plain SGD, the recipe the CNNs of the Basics part used:
ViT: 6.5M parameters, validation accuracy 0.86Do the positions discover the grid?
Cosine similarity of each patch position’s embedding with all 36, reshaped into the 6×6 grid — position (i,j) shows its map at (i,j):
mean cosine similarity: grid neighbors 0.05, distant positions -0.03The grid, barely
- Fully trained ViTs show crisp row/column bands: 2-D structure inferred from data alone.
- After ten epochs: neighbors a few hundredths positive, far positions slightly negative — a faint halo, a weak column band here and there.
The geometry the model was never told is beginning to show — and at this scale it has not gotten far.
A CNN at the same budget
Same parameter count (ResNet-style, 7×7 stem + seven residual blocks), same data, same ten epochs:
CNN: 6.5M parameters, validation accuracy 0.91
ViT: 6.5M parameters, validation accuracy 0.86The loss is the lesson
- Everything the CNN knows by construction, the ViT must learn from data — and 60,000 images do not cover it.
- 300M images flip the verdict: learned structure overtakes built-in structure (Dosovitskiy et al., 2021).
- DeiT: augmentation + distillation manufacture the missing invariances on ImageNet-1k alone.
- Swin: reinstate locality and hierarchy, escape quadratic attention at high resolution.
Recap
- ViT = strided-conv patchify → standard pre-norm encoder → classify from
<cls>; the block is untouched, only the tokenization is new. - Learned position embeddings start as noise; the 2-D grid must be discovered, and at small scale it barely is.
- At matched parameters and data, the CNN’s priors win; at large scale the transformer’s capacity wins.
- Today the plain ViT is the standard vision backbone — CLIP-style encoders, detection, segmentation, multimodal models.