The transformer, visually: building a tiny GPT

The network I trained for the digit-recognition exercise in Andrew Ng’s class in 2011 had 400 inputs, one hidden layer of 25 units and 10 outputs. Counting the bias weights, that’s (400 + 1) × 25 + (25 + 1) × 10 = 10,285 parameters. Last night I trained one with a bit under 11 million, about a thousand times as many, and it wrote me fake Shakespeare. Neither number is big by current standards, but the second one is the same kind of model that’s behind ChatGPT, and I understood every line of it, which I couldn’t have said a week ago.

The reason is Andrej Karpathy’s Let’s build GPT: from scratch, in code, spelled out, which went up on Tuesday. It’s almost two hours long. He starts from an empty Python file and a 1MB text file of Shakespeare, writes the dumbest possible model (predict the next character from only the current one), and then adds one piece at a time until he has a small transformer. It goes with nanoGPT, the repo he put out a few weeks ago, which is about 300 lines of model and 300 lines of training loop and can reproduce the smallest GPT-2.

I wrote an explainer on attention when the transformer paper came out in 2017, and I’m not going to redo it here. Queries, keys and values, multiple heads, positional encodings, that’s all in the old post and it all still holds. What that post didn’t cover, because the 2017 paper was about translation, is what you get when you throw away the encoder and keep only the decoder, which is what GPT is. That turns out to be the simpler model, and seeing it built line by line filled in the parts I’d been hand-waving for five years.

The whole thing on one page

Here’s the model end to end. Text goes in at the left as a sequence of tokens, a guess at the next token comes out at the right, and everything in between is the same block repeated.

I’ll go left to right.

Tokens

Karpathy’s video works on characters. Tiny Shakespeare has 65 distinct characters (upper and lower case letters, ten punctuation marks, a stray digit 3, space and newline), so each one gets an integer from 0 to 64, and that’s the tokenizer. The file is about 1.1 million characters, split 90/10 into training and validation.

Real GPTs don’t use characters. GPT-2 uses byte-pair encoding with a vocabulary of 50,257 tokens, where a token is a common chunk of text, often a whole word with its leading space, sometimes a piece of one. That makes sequences about four times shorter for English, which matters a lot when attention cost grows with the square of the length. It’s also why these models are weirdly bad at questions like “how many letters are in this word”, which I’d seen people poke ChatGPT about and hadn’t understood until this week. The model never sees the letters. It sees an id for “ lollipop”, or for two or three pieces of it, and has to have learned the spelling from context, like a fact.

nanoGPT’s README has a section titled “i only have a macbook” that makes the same point from the other side. With the GPT-2 vocabulary, even a tiny model spends most of its parameters on the first and last layers, because those are 50,257 rows long.

Embeddings, and the residual stream

Each token id indexes into a table, and out comes a vector of 384 numbers. That’s the token embedding. It starts random and is learned like everything else. A second table, indexed by position (0 to 255, since the context is 256 characters), gives a position embedding, and the two are added. The 2017 paper used fixed sine waves for positions. GPT learns them, which I’d wondered about in the old post, and the answer seems to be that it doesn’t much matter.

The thing I didn’t have a picture of before was what happens to that vector afterwards. The diagram above draws it as a horizontal line, and that’s how Karpathy describes it: there’s a stream, one vector per position, running straight through the network. Every attention layer and every MLP reads from the stream and adds its output back in. Nothing ever replaces the stream, it only gets things added to it. That’s what the residual connections in the paper’s diagram are, and drawn this way they stop looking like a training trick and start looking like the main design. Each block makes a small edit to a shared scratchpad.

One detail here is different from the 2017 paper. The original applied layer normalization after each sub-layer. GPT-2 moved it to the input of each sub-layer, plus one at the end, and nanoGPT does the same. It’s supposed to make deep stacks train more reliably. I’m taking that on trust.

Masked attention, and the triangle

In translation, the encoder could let every word look at every other word, because the whole English sentence was available. A GPT is generating text left to right, so when it’s predicting character 6 it mustn’t be allowed to see character 7. The fix is a mask.

In code it’s one line. You compute the full grid of query-key scores, the same way as in the 2017 post, then overwrite everything above the diagonal with minus infinity, then take the softmax of each row. e to the minus infinity is zero, so those positions get zero weight and the rest renormalize among themselves. Karpathy builds up to it slowly, starting with a version that just averages all previous positions using a lower-triangular matrix of ones, then a version that uses the mask with uniform weights, and only then lets the weights come from queries and keys. By the time the real thing shows up it’s obvious, which is the best kind of teaching.

The mask also explains something about training that I’d missed completely. Because position 5 can’t see position 6, every position in a sequence is an independent training example. One chunk of 256 characters gives you 256 predictions to score at once: the character after the first, the character after the first two, and so on up to the character after all 256. Nothing about it is sequential, so the GPU does all of them in parallel. That’s the “designed around the hardware” point from 2017 again, and it’s a big part of why these models can be trained on so much text.

The MLP, which is most of the model

After attention, each position goes through a small two-layer network on its own: expand from 384 to 1,536 (four times wider), apply a nonlinearity (GELU, a smoother cousin of the ReLU), and project back to 384. No position talks to any other in this step. Karpathy’s framing is that attention is where positions communicate, and the MLP is where each one thinks about what it’s gathered.

I counted the weights in one block of the Shakespeare model to see where they go:

• attention: the query, key and value projections are 384 × 1,152, plus the output projection at 384 × 384, about 590,000 with biases
• MLP: 384 × 1,536 and then 1,536 × 384, about 1.18 million with biases
• layer norms: 1,536

So two thirds of each block is the MLP. All the attention diagrams, including mine, give the impression that attention is the transformer, and by parameter count it’s the smaller half. Six blocks at about 1.77 million each is 10.6 million, and the embedding tables add about 120,000 more, which is where “a bit under 11 million” comes from. (The output layer reuses the token embedding table instead of having its own, a trick called weight tying that GPT-2 also uses. At 65 characters it barely matters. At 50,257 tokens it saves 38 million parameters in the small GPT-2.)

Out the other end

The final vector at each position goes through one more layer norm and a linear layer that produces 65 numbers, one score per character. Softmax turns those into probabilities. During training, the loss is the negative log of the probability the model gave to the character that actually came next. Before any training, with random weights, every character is about equally likely, so the loss should start near −ln(1/65), which is 4.17. It does, and checking that is the first thing he does after writing the model, which is a habit I wish I’d had in 2011.

Generation is a loop around all of that. Feed in the context, take the probabilities at the last position, sample a character from them, append it, repeat. Sampling instead of always taking the most likely character is why you get different text each time, and dividing the scores by a “temperature” before the softmax makes the distribution sharper or flatter. That’s the temperature setting in every one of these APIs, and it’s two lines of code.

What I actually ran

nanoGPT ships a config for character-level Shakespeare: 6 layers, 6 heads, 384 dimensions, 256 characters of context, 5,000 iterations. On the M1 laptop with the CPU settings from the README it was too slow to be worth waiting for, so I shrank it to run on the laptop for a sanity check, and then rented a GPU instance in our AWS account for an evening and ran the real config. It took well under half an hour. Validation loss got down to about 1.5 and then the model started to overfit, which the config warns you about. The samples look like this, more or less:

KING RICHARD III:
What say you, my lord? who shall be so good
As shall be sent for to your majesty?

That’s not Shakespeare and it doesn’t quite mean anything, but it’s names in capitals, a colon, verse-length lines, and words that are all real. A model that only ever saw characters learned what words are, where line breaks go and that kings talk to lords. With 10 million parameters and an hour of compute, from scratch.

The step from there to ChatGPT is, as far as I can tell from what’s public, mostly scale plus one more stage. GPT-3 is this architecture at 175 billion parameters with BPE tokens and a 2,048-token context. ChatGPT is a model from that family that’s been further trained on human-written example conversations and then on human rankings of its answers, which OpenAI calls reinforcement learning from human feedback. I don’t understand that last stage the way I now understand the block above, and Karpathy is clear at the end of the video that he’s only covered the pretraining part. That’s next on my list. OpenAI’s ChatGPT announcement says it uses the same methods as their InstructGPT work with some differences in the data, so I’ll be learning it from the InstructGPT paper and filling in the rest by guessing.

One thing I’d tell anyone who writes software and has been nodding along to transformer explanations without quite following them: watch this one with the code open, and type it in. I’ve read the 2017 paper four or five times now and I still learned more in two hours of typing along than from any of those reads.*

* I found one bug of my own along the way. I’d applied the mask in the wrong place, zeroing out the future weights after the softmax instead of setting the scores to minus infinity before it, so the weights in each row no longer summed to 1 and the loss stalled well above his. It’s the same species of bug as the transposed matrix that cost me an evening in 2011, a model that runs fine and quietly learns less, and I found it the same way, by printing intermediate shapes and values until something looked wrong.