All notes · Computer vision
Vision Transformer from scratch on MNIST
Patches, heads, and a class token. A small ViT, written out, trained on digits.
Why MNIST
I wanted the transformer pieces — patch embedding, multi-head attention, a class token — without a GPU bill. MNIST is small enough to see whether the code is wrong.
The model
Images become patches, patches become tokens, a class token is prepended, and the encoder is a stack of attention blocks I wrote rather than imported.
What I learned
The bug was almost never the attention math. It was the patch size not dividing the image, and a class token that never got a gradient because I detached it. Once those were fixed, digits were a sanity check, not the point.
The repo is Vision Transformers from scratch.