Learn · 10

Vision Transformer

A convolution looks at neighbours. A transformer lets every patch look at every other patch and decide what matters. Cutting the image into patches is one convolution with the stride set to the patch size — after that there is no image left, only a sequence.

This is also the lesson where borch.ts is narrower than torch in four places. They are named where you meet them rather than in a footnote, because each one is a line copied from a torch tutorial that will not run.

Patches are tokens

Conv2d(1, D, patch, patch, 0) gives one D-length vector per patch, laid out as [N, D, H', W']. A sequence wants [N, tokens, D], so the map is flattened and the last two axes swapped.

Divergence 1 — transpose is two-dimensional here. torch's x.transpose(1, 2) on a 3-D tensor has no equivalent; permute([0, 2, 1]) says the same thing and does work.

Attention, and a linear layer that is not batched

nn.MultiheadAttention(D, heads) takes query, key and value. Feeding it the same tensor three times is self-attention: each token asks every token what it has.

Divergence 2 — the layer returns a bare tensor. torch returns (output, weights), so const [out] = att(...) destructures a tensor here and gives you something unusable. The weights are not gone, though — nn.multiHeadAttentionForward is torch's multi_head_attention_forward and hands back both, averaged over heads or one set per head. Two things to know before you reach for it: it returns an object, { output, weights }, not a pair, and it takes length-first (L, N, E) where the layer is batch-first. Feed it batch-first and it mixes the wrong axes without complaining.

Divergence 3 — nn.Linear is two-dimensional. Handed a [N, tokens, D] tensor it refuses, and says why: mm is 2-D by 2-D. Batching is not here yet. Folding the batch and token axes together and unfolding afterwards is exact, not an approximation — a linear layer treats every row independently anyway.

The class token

Something has to turn a sequence into one prediction. Averaging the tokens is the obvious move and here it is the wrong one: in the task below exactly one patch is bright, and its contribution to the average is the same wherever it sits — the answer is averaged away. A transformer instead prepends one learned token that belongs to no patch, lets attention fill it, and reads the class off that.

Divergence 4 — expand and repeat take loose arguments, not an array. repeat([n, 1, 1]) is read as a one-element shape and fails claiming a dimension would shrink; repeat(n, 1, 1) is the same call torch would take. The message points at the wrong thing, which is worth knowing before it costs you ten minutes.

Put together

Eight-by-eight images, one of the four quadrants bright, and the label is which quadrant. Attention has something real to do: find the patch that is on, and report where it was. Position is not in the patch contents at all — it is in the learned positional embedding added to the sequence.

This is the slowest block on the site — around ten seconds. The progress lines are printed as it goes so you can see it moving rather than wonder whether it hung.

Make it narrower and watch it stop

The run above passes a loss near 0.347 on its way down and leaves. At D = 16 with two heads it does not leave. Same code, same task — only the width and the head count are different, and three seeds are run so the result is not one draw's luck.

The number it stops at is not arbitrary. ln(2) / 2 is what a model earns when it splits evenly between two answers on half its examples and is right on the rest — which is what the accuracy above says it is doing: it separates one pair of quadrants and gives up on the other, so argmax takes the same one of the pair every time and three of the four classes come out right.

Before blaming the library it was worth checking, and it was checked: cat, narrow and permute all carry the right gradients, and the head count does change the answer. The plateau is the model's, not the runtime's.

This block exists because the sentence it replaces asserted those numbers instead of producing them. It said the model settles at 0.3466 on every seed, measured on three — and a page whose argument is that its examples run should not be making a claim its examples do not make.

That is the last lesson. The tutorials take one problem to the end instead, and the playground is the same runtime with nothing written in it yet.