Building a transformer from scratch
What you only learn about attention when you refuse to import it.
Placeholder draft; replace with the real post.
I wanted to know what actually happens between input_ids and logits, so I
banned myself from torch.nn.Transformer and wrote the whole thing: tokenizer,
embeddings, causal multi-head attention, residual stream, layer norm, MLP
blocks, and the training loop on Apple MPS.
The parts that surprised me
Attention itself is three matrix multiplications and a softmax. Everything that makes it work lives elsewhere: in the residual stream, in where you put the layer norms, in the learning-rate warmup you copied without understanding.
att = (q @ k.transpose(-2, -1)) / math.sqrt(head_dim)
att = att.masked_fill(mask == 0, float("-inf"))
out = att.softmax(dim=-1) @ vWhat I'd tell past me
Start with a tiny model that overfits one batch. If it can't memorise ten sentences, it will never generalise to ten million.