ViT: An image is 16×16 words: pure-transformer image classification
Click any block in the architecture to see what it does, the math behind it, and the exact source lines. Then walk one sample end-to-end and stress-test yourself with the debug lab.
FrameworkJAX / Flax (NHWC)
Patch → tokenConv kernel=stride=P
Base modelViT-B/16 · D=768 · 12 layers
Tokens (224/16)196 patches + 1 CLS = 197
Click a stageMath cardSample stepperDebug trace
Upstream: https://arxiv.org/abs/2010.11929 · code snapshot google-research/vision_transformer @ 64801f1 · Apache-2.0
Mental model
The ideas that make this code readable
Layer-by-layer math
Every formula links to the source that implements it
Follow one sample
Tensor flow
Shapes must agree across every stage
| Stage | Input | Output | Op | Note |
|---|
Debug lab
Pick the symptom, work one boundary at a time
Code excerpts
Only the lines these notes point at, from the upstream repository at the pinned commit. Gaps are marked.
Select a file
Knowledge check
Click a question to reveal the answer