class PatchEmbedding(nnx.Module):
"""Image-to-sequence stem of the vision transformer."""
def __init__(self, img_size=96, patch_size=16, num_hiddens=512,
num_channels=3, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
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 = nnx.Conv(num_channels, num_hiddens, kernel_size=patch_size,
strides=patch_size, padding='VALID', rngs=rngs)
def __call__(self, X):
# Output shape: (batch size, no. of patches, no. of channels)
X = self.conv(X)
return X.reshape((X.shape[0], -1, X.shape[3]))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(nnx.Module):
"""Pre-norm transformer block with a GELU MLP."""
def __init__(self, num_hiddens, mlp_num_hiddens, num_heads, dropout,
use_bias=False, rngs=None):
rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs
self.ln1 = nnx.LayerNorm(num_hiddens, rngs=rngs)
self.attention = d2l.MultiHeadAttention(
num_hiddens, num_heads, dropout, use_bias, rngs=rngs)
self.ln2 = nnx.LayerNorm(num_hiddens, rngs=rngs)
self.mlp = nnx.Sequential(
nnx.Linear(num_hiddens, mlp_num_hiddens, rngs=rngs), nnx.gelu,
nnx.Dropout(dropout, rngs=rngs),
nnx.Linear(mlp_num_hiddens, num_hiddens, rngs=rngs),
nnx.Dropout(dropout, rngs=rngs))
def __call__(self, X, valid_lens=None):
X_norm = self.ln1(X)
X = X + self.attention(X_norm, X_norm, X_norm, valid_lens)[0]
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, num_channels=1, rngs=None):
super().__init__()
self.save_hyperparameters(ignore=['rngs'])
rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs
self.patch_embedding = PatchEmbedding(
img_size, patch_size, num_hiddens, num_channels, rngs=rngs)
self.cls_token = nnx.Param(jnp.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 = nnx.Param(
rngs.params.normal((1, num_steps, num_hiddens)) * 0.02)
self.embedding_dropout = nnx.Dropout(emb_dropout, rngs=rngs)
self.blks = nnx.List([
ViTBlock(num_hiddens, mlp_num_hiddens, num_heads, blk_dropout,
use_bias, rngs=rngs) for _ in range(num_blks)])
self.head = nnx.Sequential(
nnx.LayerNorm(num_hiddens, rngs=rngs),
nnx.Linear(num_hiddens, num_classes, rngs=rngs))
def forward(self, X):
X = self.patch_embedding(X)
X = d2l.concat((jnp.tile(self.cls_token, (X.shape[0], 1, 1)), X), 1)
X = self.embedding_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.87Do 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.03, distant positions -0.02The 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.87The 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.