Cut a 224×224 image into 16×16 patches, project each patch into 768 numbers, prepend a class token and add a position table — and a plain transformer encoder trains on images. No convolution, no pooling, no pyramid. Five chapters of arithmetic turn the picture into a sequence, then the recipe that makes it work: DeiT, MAE, and how to fine-tune a pretrained ViT with timm.
Conv2d(3, 768, kernel_size=16, stride=16) reads 3×224×224 and writes 768×14×14: 14×14 = 196 non-overlapping patches, each projected to a 768-dim vector in the same operation. It costs 3·768·16² + 768 = 590,592 weights — and 196 × 768 = 150,528, exactly the number of values the image had, so patching changes shape, not size.
224 ÷ 16 = 14 → 196 patches · 197 tokens with [CLS] · 38,809 attention pairsSelf-attention is permutation-equivariant: without a position signal the image is a bag of patches and it can never learn that a wheel sits under a car. ViT prepends a learned [CLS] vector (the summary seat) and adds a learned 197×768 position table — 151,296 numbers, 0.17% of the model — before the first block. The table is fixed at 197 rows, so a resolution change needs interpolation.
197 × 768 = 151,296 position numbers · 0.175% of 86,567,656ViT-B/16 scored 77.9% on ImageNet-1k alone and 84.2% pretrained on JFT-300M. DeiT's augmentation and distillation lifted the same 86.6M model to 81.8% on ImageNet-1k alone (83.4% distilled); MAE's 75%-masked pretraining reached 83.6% after fine-tuning and 68.0% as a linear probe — the honest 15.6-point gap between frozen and unfrozen features.
86,567,656 params · 12 blocks · 12 heads · MLP is 65.5% of the weightsA ViT is a text transformer that reads a sentence of patches — the patch embedding is the word, the position table is the word’s place in the sentence, and [CLS] is the seat that collects the summary. Everything else (blocks, heads, pre-LN, fine-tuning) is the standard transformer stack you already know.
By the end you will be able to derive the token count of any ViT from its input size and patch size (224/16 → 196 + 1), count its parameters from the config (86,567,656 for ViT-B/16, 65.5% of them in the MLPs), say why the position table must be interpolated at a new resolution, explain why ViT needed JFT-300M in 2020 and why DeiT and MAE removed that requirement, compare the ViT/Swin/ConvNeXt priors and their costs, and walk the linear-probe → last-block → full-fine-tune ladder with the right learning rates.
A ViT refuses to look at pixels. Its first layer slices the image into non-overlapping 16×16 patches and projects each one to a 768-dim vector — and that single convolution is the entire image-to-tokens step.
The whole architecture is seven steps, and every variant you will meet — DeiT, Swin, ConvNeXt, MAE — changes one or two of them and leaves the rest alone:
Start with the arithmetic that makes the first line work. A 224×224 image with 16×16 patches divides into 224 ÷ 16 = 14 positions per side, so the grid is 14 × 14 = 196 patches. Each patch holds 16 × 16 × 3 = 768 numbers — exactly the width of a ViT-B token. And here is a coincidence worth checking: 196 × 768 = 150,528, and 224 × 224 × 3 = 150,528 too. Patch embedding does not compress the image at all; it changes the shape of the data from a 3-D tensor into a sequence of vectors, so the transformer machinery from text can run on it. (The two 150,528s are a coincidence of 16² × 3 = 768; at 8px patches the token sequence holds 784 × 768 = 602,112 numbers, 4× the image, and at 32px patches it holds 49 × 768 = 37,632.)
The whole transformation is one convolution with a trick: kernel_size = stride = patch_size, no padding. Because the kernel is exactly as wide as the step it takes, the windows never overlap and never share a pixel — each patch is read once. Because it is a convolution, the patch is also linearly projected at the same time: 3 input channels in, dim channels out. The layer is literally nn.Conv2d(3, 768, kernel_size=16, stride=16), and it learns 3 · 768 · 16² + 768 = 590,592 weights — about 0.7% of the model.
The patch size is the one dial that decides how many tokens the transformer sees. Smaller patches mean more tokens, finer detail and a quadratic jump in attention cost; larger patches mean fewer tokens and a coarser grid:
The source’s tiny ViT uses the same code with a smaller dial: a 64×64 image with 16×16 patches gives a 4×4 grid = 16 patches + CLS = 17 tokens. That is the model you will build in chapter 04, and it is small enough to run on a laptop CPU.
import torch
import torch.nn as nn
class PatchEmbedding(nn.Module):
def __init__(self, in_channels=3, patch_size=16, dim=192, image_size=64):
super().__init__()
assert image_size % patch_size == 0
# kernel size = stride = patch size: non-overlapping, and the
# projection to dim happens in the same convolution.
self.proj = nn.Conv2d(in_channels, dim, kernel_size=patch_size, stride=patch_size)
self.num_patches = (image_size // patch_size) ** 2
def forward(self, x):
x = self.proj(x) # (N, dim, H/P, W/P)
return x.flatten(2).transpose(1, 2) # (N, num_patches, dim)
# shape trace, tiny model: 64 / 16 = 4 -> 4×4 = 16 patches
patch = PatchEmbedding(image_size=64, patch_size=16, dim=192)
x = torch.randn(2, 3, 64, 64)
print(patch(x).shape) # torch.Size([2, 16, 192]) -> 16 tokens, 192 dims
# the real thing: 224 / 16 = 14 -> 14×14 = 196 patches, 768 dims
vit_b = PatchEmbedding(image_size=224, patch_size=16, dim=768)
print(vit_b.proj.weight.shape) # torch.Size([768, 3, 16, 16])
print(sum(p.numel() for p in vit_b.parameters())) # 590,592The first layer of a ViT is a single convolution with kernel size = stride = patch size: it slices the image into non-overlapping patches and projects each one to a 768-dimensional token in the same operation. Click a patch to follow it into the token sequence.
A 16px patch on a 14×14 grid is 196 tokens, so the tokens hold 196 × 768 = 150,528 numbers against the image’s 150,528 — a 1× change in size. Only at 16px does 196 × 768 equal the input’s 150,528 exactly (because 16² × 3 = 768): there the conv is a pure reshape, and at every other patch size the projection expands or shrinks the payload.
A 384×384 image goes through a ViT whose patch embedding is Conv2d(3, 768, kernel_size=16, stride=16). How many tokens (including [CLS]) does the transformer see?
Flattening gives 196 patch vectors, but a transformer cannot tell position and has no single vector to classify. Two learned tensors fix both problems: the [CLS] token and the position table.
After the patch conv, flatten(2).transpose(1, 2) turns the 768×14×14 feature map into a sequence of 196 vectors. Two operations then prepare it, both copied straight from BERT:
Why a class token at all? The classifier needs one vector per image, not 196. You could average the patch outputs — and many later models do — but the original design prepends a single learned vector and lets self-attention do the aggregation: because every token can attend to every other, the CLS token’s output after the last block is a weighted summary of the whole image. The classifier is then a single Linear(768, num_classes) reading row 0. The CLS token costs 768 parameters and one extra row in the sequence; it is a convention, not a machine.
Why position embeddings? Self-attention is permutation-equivariant: feed the same 196 tokens in a different order and every output comes back in that same different order, with identical values. Without a position signal the model sees a bag of patches — it could never learn that a wheel sits under a car, because “under” is not expressible. The fix is a learned table with one row per sequence slot, added to the tokens before the first block. It holds 197 × 768 = 151,296 numbers, which is 0.17% of ViT-B’s 86.6M parameters — tiny, and load-bearing.
The ViT paper tried the alternative too: fixed 2-D sinusoidal embeddings, the hand-built patterns that encode row and column, needed no training at all. They performed almost identically; learned embeddings won by a small margin and became the default. What matters is that something breaks the symmetry — the network can learn the rest from data.
class ViT(nn.Module):
def __init__(self, image_size=64, patch_size=16, num_classes=10,
dim=192, depth=6, num_heads=3, mlp_ratio=4):
super().__init__()
self.patch = PatchEmbedding(3, patch_size, dim, image_size)
num_patches = self.patch.num_patches
# one learned vector, broadcast to every batch element
self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
# one learned vector per sequence slot: patches + CLS
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, dim))
nn.init.trunc_normal_(self.cls_token, std=0.02)
nn.init.trunc_normal_(self.pos_embed, std=0.02)
# ... blocks, final LayerNorm, head ...
def forward(self, x):
x = self.patch(x) # (N, 16, 192)
cls = self.cls_token.expand(x.size(0), -1, -1) # (N, 1, 192)
x = torch.cat([cls, x], dim=1) # (N, 17, 192)
x = x + self.pos_embed # (N, 17, 192)
# ... blocks, then head(self.ln(x[:, 0])) ...
# the real model: 197 slots × 768 dims = 151,296 position numbers
# 151,296 / 86,567,656 = 0.175% of ViT-B's parametersSelf-attention is permutation-equivariant: with no position signal it cannot tell patch 3 from patch 137. Switch the mode and watch the 196 slot vectors — then read what the model can still recover.
The tiles are a teaching rendering of learned vectors (or the smooth fixed 2-D sinusoidal pattern); whichever mode you pick, the arithmetic underneath is the same: 197 × 768 = 151,296 numbers, added to the tokens before the first block.
You feed a trained ViT a batch of images, then shuffle the order of the 196 patch tokens inside the sequence — but you keep the same position embeddings in the same slots. What happens to the output?
Each token builds a query, a key and a value. The query is compared with all 197 keys, and the result is a weighted average of the values — a global mixing operation with no notion of locality and a quadratic bill.
Self-attention is the same operation you met in Phase 7; ViT only changes what the tokens are. For each token, three linear projections produce a query, a key and a value. The query of token i is dotted with the key of every token j, the scores are scaled by 1/√d and pushed through a softmax, and the outputs are summed into a new vector for i:
Take token i = the [CLS] row. Its query vector is 64 numbers. Multiplying by the key matrix (64 × 197) costs 197 × 64 = 12,608 multiply-accumulates and produces one row of 197 scores. Scaling by √64 = 8, softmaxing, and multiplying by V (197 × 64) is another 12,608 MACs. Every one of the 197 tokens does the same, so one head costs about 2 × 197 × 197 × 64 = 4.97M MACs for scores and values; twelve heads cost 59.6M, and twelve blocks 715.3M MACs for the attention core — per image, at batch 1.
The score matrix itself is the shape to remember: 197 × 197 = 38,809 entries per head, stored as 152 KiB of float32, and 12 heads make it 1.78 MiB per block per image. The same tensor at Swin’s 4px patches (3,136 tokens) would need 9,834,496 entries × 4 bytes = 37.5 MiB per head, or 450 MiB per block — before a single multiplication, which is why nobody runs full attention over thousands of tokens.
Notice how modest the quadratic term still is at ViT-B’s size. The score computation costs n² · d = 29.8M MACs per block, while the MLP sublayer costs 8 · n · d² = 930M MACs — about 31× more. Most of the attention sublayer’s own cost is not the scores at all but the three linear projections that build Q, K and V, which are linear in n, not quadratic. The n² term becomes the wall when you keep resolution: at 8px patches (785 tokens) scores grow 15.9× to 473M; at 4px patches (3,136 tokens) the score matrix alone is 51% of the MLP and the attention sublayer as a whole matches it. Scaling the token count is what makes attention expensive, and that is exactly what dense prediction tasks demand.
What does attention actually learn? The ViT paper measured the mean attention distance — how far, in pixels, each head reaches — and found a clean pattern: in the earliest layers some heads attend only to the neighbouring few patches, while others already reach across the image; by the last layers attention is global across the board. Early layers do the work convolution does with a 3×3 kernel; late layers mix semantics — sky to sky, tree to tree. The explorer below lets you feel that gradient: slide the layer and watch locality give way to meaning.
Pick a patch and a layer. Early layers of a ViT attend locally — neighbours dominate — while late layers attend by meaning: sun to sun, tree to tree. Lines are scaled to the strongest target; the panel shows the true softmax weights.
Attention is O(n²): every one of the 197 tokens scores every other, so the matrix is 38,809 entries per head. At 224×224 with 16px patches that is cheap; go to 8px patches and n grows to 785, so the same matrix grows to 616,225 entries — 15.9× bigger.
At ViT-B's 197 tokens, which part of one encoder block spends the most multiply-accumulates?
The block is not a new idea: multi-head self-attention and an MLP, each wrapped in LayerNorm and a residual connection. Stack twelve of them and you have ViT-B — 86,567,656 parameters, of which the attention scores are almost none.
Here is the whole block, in two lines of pseudocode and one design decision:
The decision is pre-LayerNorm: normalization happens before each sublayer, not after the residual add. Early transformers used post-LN, x = LN(x + sublayer(x)), and struggled to train past 6–8 layers without a learning-rate warmup — gradients through the stacked normalizations are unstable. Pre-LN leaves the residual stream un-normalized, so information can flow from the input to the output through additions alone, and every sublayer sees a well-scaled input. It trains deep stacks without warmup, which is why ViT, GPT-2 and every modern LLM use it. The 12 blocks of ViT-B would be a research project with post-LN.
Now the parameter budget. This is the table people get wrong: the model is called a transformer, so intuition says attention dominates the weights, but two-thirds of every block is the MLP.
The same code with smaller dials gives the source’s tiny ViT: a 64×64 image, 16 patches, dim 192, 6 blocks, 3 heads. One block is 444,864 parameters; the whole model is 2,822,602 — small enough to train on a CPU in minutes, and structurally identical to the 86.6M model that beat ResNet.
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4, dropout=0.0):
super().__init__()
self.ln1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout,
batch_first=True)
self.ln2 = nn.LayerNorm(dim)
self.mlp = nn.Sequential(
nn.Linear(dim, dim * mlp_ratio),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(dim * mlp_ratio, dim),
nn.Dropout(dropout),
)
def forward(self, x):
normed = self.ln1(x) # pre-LN: normalize first
a, _ = self.attn(normed, normed, normed, need_weights=False)
x = x + a # residual
x = x + self.mlp(self.ln2(x)) # residual
return x
# nn.MultiheadAttention does the split into heads, the scaled
# dot-product, and the output projection.
class ViT(nn.Module):
def __init__(self, image_size=64, patch_size=16, in_channels=3,
num_classes=10, dim=192, depth=6, num_heads=3, mlp_ratio=4):
super().__init__()
self.patch = PatchEmbedding(in_channels, patch_size, dim, image_size)
num_patches = self.patch.num_patches
self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, dim))
self.blocks = nn.ModuleList([
Block(dim, num_heads, mlp_ratio) for _ in range(depth)
])
self.ln = nn.LayerNorm(dim)
self.head = nn.Linear(dim, num_classes)
nn.init.trunc_normal_(self.pos_embed, std=0.02)
nn.init.trunc_normal_(self.cls_token, std=0.02)
def forward(self, x):
x = self.patch(x)
cls = self.cls_token.expand(x.size(0), -1, -1)
x = torch.cat([cls, x], dim=1)
x = x + self.pos_embed
for blk in self.blocks:
x = blk(x)
x = self.ln(x[:, 0]) # the [CLS] row, after LayerNorm
return self.head(x)
vit = ViT(image_size=64, patch_size=16, num_classes=10,
dim=192, depth=6, num_heads=3)
print(f"params: {sum(p.numel() for p in vit.parameters()):,}")
# params: 2,822,602
x = torch.randn(2, 3, 64, 64)
print(vit(x).shape) # torch.Size([2, 10])
print(vit(x).softmax(-1).sum(-1)) # tensor([1.0000, 1.0000])The source’s step 4 is to run a single image and look at the numbers, because shape errors and uninitialized parameters hide until the first real forward pass:
ViT arrived with no locality and no translation equivariance, and it lost to ResNet on ImageNet. Two years later DeiT and MAE had shown that the data hunger was a recipe problem — and Swin and ConvNeXt showed it was never really attention versus convolution.
For a decade, convolution and computer vision were synonyms — CNNs had the priors, and nobody expected to replace them. Then the 2020 ViT paper was blunt about its own weakness. Trained on ImageNet-1k’s 1.28M images, ViT-B/16 reached 77.9% top-1 and lost to the best CNNs of the day. The network had given away the two assumptions that make convolutions efficient — locality (neighbouring pixels matter most) and translation equivariance (the same pattern can appear anywhere) — so it had to learn them from data. The paper’s own summary: large scale training trumps inductive bias. On ImageNet-21k (14M images) the same model reached 84.0%; on JFT-300M (300M images), ViT-B/16 hit 84.2% at 384² and the headline ViT-H/14 reached 88.55%. From 1.28M to 300M images is 234× more data, and the ranking flipped.
That story created a myth worth killing: “ViTs need hundreds of millions of images.” DeiT (2020) trained the same architecture on ImageNet-1k alone to 81.8% — up 6.3 points from the original ViT-B in a comparable setting (at 224²; the 77.9% quoted above is the 384² fine-tune) — with four changes to the recipe, none to the model:
The point of stochastic depth is worth spelling out: dropping entire blocks at random means the network can never rely on any single layer, which regularizes a model that otherwise memorizes 1.28M images. Plus distillation: the ViT learns from a convnet’s soft outputs and inherits some of its priors. Every modern ViT recipe descends from DeiT.
MAE (2021) went further and removed the labels: mask 75% of patches at random, train the encoder on only the visible 25% (49 of 196 patches, plus CLS = 50 tokens instead of 197), and train a small decoder to reconstruct the 147 missing patches from those 50 tokens. Then discard the decoder and fine-tune the encoder. The mask ratio is the interesting part. BERT masks 15% of tokens in text; MAE masks 75% of patches in images, five times more. Why the difference? Language is information-dense — delete 15% of the words and the sentence is hard. Images are spatially redundant: if a model can see a patch’s neighbours, it can interpolate the missing colour almost perfectly without understanding anything. You have to hide three-quarters of the picture before reconstruction requires real visual structure. MAE pretraining on ImageNet-1k alone reaches 83.6% after fine-tuning — the neighbourhood of supervised pretraining on far larger corpora — and its linear probe is 68.0%. Those two numbers are the honest measure of the frozen-vs- fine-tuned gap you will meet in the next chapter, and masked autoencoding remains the default self-supervised recipe for vision: VideoMAE extends it to video, and V-JEPA turns the same idea into a prediction task in representation space.
Meanwhile two other architectures answered a question the ViT paper had left open. Swin puts locality back: 4×4 patches give a 56×56 grid of 3,136 tokens, and full attention there would mean 9,834,496 pairs per layer. Instead each block attends inside 7×7 windows — 2,401 pairs per window — and alternating blocks shift the window by half its size so information crosses window borders over a few layers. The saving is 64×. ConvNeXt asks whether you need attention at all: it takes a ResNet and rebuilds it with the transformer era’s choices — depthwise 7×7 convolutions, LayerNorm, GELU, inverted bottlenecks, a 4-stage pyramid, and the modern training recipe. It reaches 83.8% on ImageNet-1k alone, ahead of the from-scratch ViT, and 86.8% with ImageNet-21k pretraining. The lesson of all three, in one line: the 2020s gap was mostly architecture plus training recipe, not attention versus convolution.
Where does that leave the field? ConvNeXt-V2 and Swin-V2 are both production-grade, and pure convnets still win on edge devices because inference stacks compile them best. Everything else has converged on the ViT block: segmentation (Mask2Former, SegFormer), detection (DETR, RT-DETR), multimodal (CLIP, SigLIP, LLaVA, BLIP-2) and video (VideoMAE, V-JEPA) all use it, and often the image encoder is the only part that changes between them. That is why this is the vision architecture worth knowing end to end.
The same accuracy board read three ways: what each architecture assumes about images, what it costs, and which one to reach for given your data and deployment target.
JFT-300M (300M images), 384² fine-tune · 197 (196 patches + CLS)
ImageNet-22k, 384² fine-tune · 3,136 at 4×4 patches (56×56 grid)
ImageNet-22k, 384² fine-tune · no tokens: 4-stage feature pyramid (7×7 depthwise convs)
Numbers from the papers’ own tables: ViT-B/16 ImageNet-only 77.9 and JFT-300M 84.2 (at 384²), Swin-B 83.5 / 86.4 (ImageNet-22k at 384²), ConvNeXt-B 83.8 / 86.8 (ImageNet-22k at 384²). All three are trained on ImageNet-1k in the left bar; only the right bar sees extra data.
Rules of thumb, not laws — the honest 2026 answer is that the gap between a well-trained convnet and a well-trained ViT is small, and pretraining corpus plus inference stack usually decides. Label the models you compare, not the families you like.
Pretrained ViTs ship through timm in one line. The craft is choosing how much of the network to unfreeze, and how gently — a full fine-tune at the wrong learning rate destroys exactly what you paid for.
timm’s API has not changed in years because it does not need to: it is the production default for pretrained vision models in 2026, and one line gets you a ViT, a DeiT, a Swin or Swin-V2, a ConvNeXt or ConvNeXt-V2, a MaxViT, an MViT or an EfficientFormer with ImageNet weights. For multimodal work, transformers ships CLIP, SigLIP, BLIP-2 and LLaVA — and the image tower inside every one of them is a ViT variant, usually at 14px patches.
import timm
import torch
import torch.nn as nn
# one line: architecture + ImageNet weights + a new 10-class head
model = timm.create_model("vit_base_patch16_224", pretrained=True, num_classes=10)
# what just happened:
# patch conv, cls_token, pos_embed, 12 blocks -> pretrained weights
# head: old Linear(768, 1000) -> new Linear(768, 10)
# the old head's knowledge is gone; the features are not
print(sum(p.numel() for p in model.parameters()))
# 85,806,346 = 86,567,656 - 769,000 (old head) + 7,690 (new head)
# recipe 1: linear probe - freeze everything, train the head
for name, param in model.named_parameters():
param.requires_grad = name.startswith("head")
optimizer = torch.optim.AdamW(
[p for p in model.parameters() if p.requires_grad], lr=1e-3
)
# recipe 3: full fine-tune - everything trainable, 10-100x smaller LR
for param in model.parameters():
param.requires_grad = True
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)There is a ladder of how much to unfreeze, and it is worth walking in order, because each step tells you something the previous step cannot:
Start with the linear probe. Freeze the backbone, train the new head for 10–30 epochs at 1e-3, and read the accuracy: that number is the quality of the frozen features for your task. It is also the cheapest useful experiment you can run — minutes, not hours. MAE’s own ViT-B scores 68.0% this way on ImageNet-1k, and 83.6% after full fine-tuning: the gap between frozen and unfrozen features can be 15.6 points, so a good probe is a floor, never a ceiling.
Then unfreeze the last block and the final LayerNorm: that is where the global, task-specific mixing concentrates, so about 8% of the parameters buys most of the remaining gap. Finally, the full fine-tune: every weight moves, at a learning rate 10–100× lower than from-scratch training (1e-5 to 1e-4 for ViT-B), with cosine decay and the DeiT augmentation stack. Lower is not timidity; the pretrained solution sits in a good basin, and a large step walks out of it.
Same pretrained ViT-B/16, same 10-class dataset — the only difference is how much of the network you let the optimizer touch. Walk the three steps and watch trainable parameters, learning rate and memory change together.
Trainable share of the 86,567,656-parameter model. Bar widths are √share so the 7,690-parameter head stays visible next to 100%.
The MAE paper’s own ViT-B/16 numbers on ImageNet-1k: 68.0% linear probe, 83.6% after full fine-tuning — a 15.6-point gap between frozen features and the same features with every weight allowed to move. (Bars span 60–90%.)
· freeze the backbone and fit the head first: 10 epochs, a couple of minutes each, and it tells you whether the pretrained features already separate your classes.
· if the probe is close to what you need, stop there — with 5k images the fine-tune gap is usually a few points, not the 15.6 points MAE measured on ImageNet.
· if you fine-tune, use the lowest learning rate in the table (1e-5–1e-4) with cosine decay, and watch the validation loss for the first epoch; a pretrained ViT can be damaged by one large step.
The backbone is frozen and stays in eval mode; only the new classifier learns. This measures how good the frozen features are — nothing more. It is the cheapest recipe and the right first run, because if a linear probe works, fine-tuning will work better.
You fine-tune a pretrained ViT-B/16 on a small dataset and the accuracy after epoch 1 is worse than the linear probe's final accuracy. What is the most likely cause?
The patch-count question and the [CLS] question are the two you will be asked to do in your head for the rest of the phase. The DeiT question and the Swin question separate the history from the mechanism; the pre-LN question is the one that shows up in every transformer you build after this.
0 / 5 answered · 0 correct
01How does a ViT turn an image into a sequence of tokens?
02What is the [CLS] token in ViT and why is it needed?
03Why did the original ViT paper need JFT-300M pretraining to beat ResNet, and why does DeiT not?
04Pre-LayerNorm (`x = x + sublayer(LN(x))`) vs post-LayerNorm (`x = LN(x + sublayer(x))`) — which is used in modern transformers and why?
05Swin Transformer introduces windowed attention. What problem does it solve?
Click a card to swap the lazy description for what it actually means.
Three problems with exact numbers — trace every tensor shape through the 2.8M-parameter tiny ViT, fine-tune a pretrained ViT-S/16 against ResNet-18 and explain the parameter gap, and implement MAE pretraining with the 75% mask ratio. Try first; a worked answer is one click away.
Trace the source code line by line with image_size=64, patch_size=16, dim=192, depth=6, heads=3, num_classes=10. The patch conv is Conv2d(3, 192, kernel_size=16, stride=16), so 64/16 = 4 gives a 4×4 grid of patches: (2, 3, 64, 64) → conv → (2, 192, 4, 4) → flatten(2).transpose(1, 2) → (2, 16, 192). cls_token.expand(2, -1, -1) makes (2, 1, 192); torch.cat([cls, x], dim=1) gives (2, 17, 192) — 16 patches + 1 CLS. Adding pos_embed (1, 17, 192) broadcasts to (2, 17, 192) with no shape change. Six blocks keep the shape: self-attention on (2, 17, 192) with 3 heads of 192/3 = 64 dims, then the MLP 192 → 768 → 192. After the final LayerNorm, x[:, 0] selects the CLS row: (2, 192). The classifier is Linear(192, 10), so the output is (2, 10) logits, and softmax over dim=-1 sums to 1. The parameter count by hand: patch conv 3·192·16² + 192 = 147,648; CLS 192; position 17·192 = 3,264; one block = two LayerNorms (2·2·192 = 768) + qkv (192·576 + 576 = 111,168) + output projection (192·192 + 192 = 37,056) + MLP fc1 (192·768 + 768 = 148,224) + fc2 (768·192 + 192 = 147,648) = 444,864, times 6 = 2,669,184; final LayerNorm 384; head 192·10 + 10 = 1,930. Total 147,648 + 192 + 3,264 + 2,669,184 + 384 + 1,930 = 2,822,602, which is the 'about 2.8M' the source quotes. Run it with torch.manual_seed(0) and every print matches.
Build the data once with a fixed split, then run both models through the same loop. The timm model is one line: model = timm.create_model('vit_small_patch16_224', pretrained=True, num_classes=10); the ResNet comparison is torchvision.models.resnet18(weights='IMAGENET1K_V1') with its fc replaced by nn.Linear(512, 10). Fine-tune both with AdamW at lr=1e-4 (the pretrained ViT) and lr=1e-3 (ResNet, or 1e-4 with SGD momentum 0.9), cosine decay over 10–20 epochs, batch 64, and the same augmentation. Report wall-clock per epoch and final top-1 on the held-out split. Parameter counts: ViT-S/16 is dim=384, depth=12, heads=6 — embeddings 295,296 (patch) + 384 (CLS) + 197·384 = 75,648 (position) = 371,328, blocks 12 × 1,774,464 = 21,293,568, final LayerNorm 768, head 385,000, total 22,050,664 ≈ 22.1M; ResNet-18 is 11,689,512 ≈ 11.7M. ViT-S is 1.9× bigger yet typically within a point of ResNet-18 on this small dataset, because its 197-token sequence (16×16 patches on the 224² resize the lesson uses) has to learn locality from pretraining. The training-time difference is the real story: the ViT's attention is n² over 197 tokens, so per-image compute is several times ResNet-18's at the same batch size, even though both converge in similar epochs.
The shape arithmetic first: the tiny ViT has 16 patches; mask 75% means 12 masked and 4 visible. Keep the visible patches' order (shuffle with a fixed permutation, keep the first 12 indices as 'masked'), add the CLS token, and feed the encoder (16 → 4 tokens + CLS = 5 tokens) — the encoder never sees the masked pixels, so the compute saving is real (roughly 4× fewer tokens here, 49 vs 196 in ViT-B). Then a small decoder: project the encoder output to a decoder dim (96 is plenty for the tiny model), append a shared mask token for each of the 12 missing positions, add the decoder's positional embeddings for all 17 slots, and run 2 decoder blocks. The loss is mean squared error on the reconstructed pixels, computed only at masked patches — 12 × 16 × 16 × 3 = 9,216 pixel values per image you actually score. After pretraining, discard the decoder: keep the encoder, attach a Linear(192, 10) head, and compare two runs on the same synthetic split — (a) linear probe with the encoder frozen, (b) full fine-tune. The expected pattern is the MAE paper's at a smaller scale: the probe improves over a randomly initialized encoder but sits well below full fine-tuning, and the gap narrows as the reconstruction loss falls. Report both numbers plus the mask ratio curve (try 50%, 75%, 90%): 75% is not a magic constant, it is the point where the task is hard enough to force semantic features but still solvable from neighbouring patches.
You do not need to master these here. Each one gets a proper treatment in its own lesson; the one-line meaning is enough to keep reading. Orange dotted underlines in the prose point back to this list.
This lesson is a port of an open course. Everything here traces back to it — and the next step is running the code yourself.
Lesson text adapted from AI Engineering from Scratch (Phase 04, Lesson 14) and the Math Foundations Notebook reference build. The five labs — the patch-embedding visualizer, the clickable attention explorer, the position-embedding board, the architecture comparator and the fine-tuning stepper — are original to this page, as is the exact parameter budget (86,567,656 for ViT-B/16 broken into 33.3% attention, 66.6% MLP per block and 2,822,602 for the source's tiny ViT), the patch-size and attention-pair table (2,500 / 38,809 / 616,225 pairs), the honest note that at 197 tokens the n² scores cost 29.8M MACs against the MLP's 930M, the Swin window arithmetic (3,136 tokens, 9.8M full-attention pairs, 64× fewer with 7×7 windows), MAE's 50-token encoder input and 75%-vs-15% mask comparison, the fine-tune ladder with 120 KiB / 108 MiB / 1.29 GiB of optimizer state, and the position-resolution trap with the memory hook. Every number shown is computed live by the labs or verified by hand in the prose.