Inside a transformer

Training is one question asked millions of times: given these characters, what comes next? The model starts with random numbers, guesses, and measures how surprised it was by the real answer. Then it nudges every one of its numbers a tiny step in the direction that would have made it less surprised. That nudge is backpropagation, and the step rule is AdamW.

This one took 5,000 steps, each on 32 random 64-character windows of about a million characters of Shakespeare. It ran in about a quarter of an hour on a laptop's CPU, in a training script written for this page. Its surprise is measured in nats, and lower is better. Picking any of the 65 characters at random scores 4.2, guessing by letter frequency scores 3.3, and guessing from the previous letter alone scores 2.5. It learned spelling, the shape of a speech and who speaks when, but not meaning.

GPT-2, from 2019, has exactly this layout: tokens, embeddings, then a stack of blocks that each run attention and an MLP, adding their results into the residual stream. It is about a thousand times bigger, with 12 blocks of 12 heads, vectors of 768 numbers and 124 million parameters. Today's large models are bigger again by a thousand or more and read word pieces rather than letters, but the step from one token to the next is still the one on this page.

• Heads and neurons rarely mean one tidy thing. Some heads here clearly track the previous character or the start of a word. Most mix several jobs, which is what interpretability research works to untangle.
• The logit lens is a rough probe. The middle of the model was never trained to be read out directly, so its early guesses are a hint at what is building up, not a measurement.
• It is tiny. With 108,000 parameters and 64 characters of memory it can sound like Shakespeare for a line, but it cannot follow an argument.

• Vaswani et al. 2017 · A. Vaswani et al., "Attention Is All You Need", NeurIPS (2017): the transformer
• Radford et al. 2019 · A. Radford et al., "Language Models are Unsupervised Multitask Learners" (2019): GPT-2, whose layout this model copies
• Karpathy 2015 · A. Karpathy, char-rnn (2015): character-level language models, and the Tiny Shakespeare corpus used here
• Ba, Kiros & Hinton 2016 · J. L. Ba, J. R. Kiros and G. E. Hinton, "Layer Normalization" (2016)
• Hendrycks & Gimpel 2016 · D. Hendrycks and K. Gimpel, "Gaussian Error Linear Units (GELUs)" (2016): the MLP's activation
• Press & Wolf 2017 · O. Press and L. Wolf, "Using the Output Embedding to Improve Language Models", EACL (2017): tying the input and output embeddings
• Loshchilov & Hutter 2019 · I. Loshchilov and F. Hutter, "Decoupled Weight Decay Regularization", ICLR (2019): AdamW, the optimizer used to train it
• nostalgebraist 2020 · nostalgebraist, "interpreting GPT: the logit lens" (2020): reading predictions out of the middle of the model
• Elhage et al. 2021 · N. Elhage et al., "A Mathematical Framework for Transformer Circuits" (2021): the residual stream as a shared channel, and attention heads as readers and writers of it