A tiny GPT, trained in your browser

The problem: continuing what the text started

A language model is mostly shown text it has never seen: a new name, a list, a variable in someone's code, a few examples of a task. A good one carries on in kind. It spells the name the same way each time, keeps the format of the list, uses the variable again. This is called in-context learning, learning from the text in front of it without changing a single weight. It is what made GPT-3 famous in 2020: shown a few examples of a task in its prompt, it could often do the task. Knowing what usually follows what, from training, is not enough for it, because whatever is new, by definition, never followed anything in training.

Why is it hard?

Chapter 7's attention can bring back a token from anywhere in the story, if its key matches what the reader is looking for. But a key is made from one token's row, so it can only say what that token is. Give a story's hero an invented name, “vorna”, cut into the pieces “v”, “##or” and “##na”. At a later mention, after “v”, the right next piece is the one that came after “v” before, so the network has to find a token by what came before it, and one layer of attention cannot express that. Two layers can: the first writes into each position which token came before it, and the second looks that up. But stacks of layers can be hard to train (in chapter 4, gradients faded over six layers of σ neurons, and needed well-scaled starting weights even with others), and the transformer's answer comes in two small parts, residual connections and layer norms. Here is a test story the networks never saw, with its hero renamed. Marked pieces come after the first piece of a later mention, where the story itself shows what comes next.

What a stack does differently

• Every block adds to every row. A token's row passes up through the blocks, and each block adds what it found to the row instead of replacing it. This is the residual connection. So the row a later block reads still holds the token itself, plus everything the blocks below it added.
• Later blocks read what earlier ones wrote. If a head in block 1 has written “the token before me was z” into a row, a key in block 2 can offer it. Heads in different blocks can work together, which is where the GPT's new abilities come from.
• Layer norm keeps the numbers in range. Before each part reads a row, the row is rescaled to an average of 0 and a spread of 1, so each part sees numbers of the same size, however many blocks have added to the row.

Measuring…

How two blocks finish a name, step by step

Before the real network, a toy small enough to see every number. It has two heads, one in each of two blocks, and its numbers are made up: a token's row is a row of cells, one cell for each kind of token, so that a dot product of two of these rows is 1 when they name the same token and 0 when they do not. The arithmetic is real, computed as you watch. What the toy leaves out (the MLPs, layer norms, and rows of numbers that are not so tidy) the rest of the page covers.

That is the whole trick. A head in block 1 reads the token just before, by a high distance score at distance 1, and its result is added to each row, so every row now says what it is and what came before it. A head in block 2 asks, with its query, “whose previous token was the one I am?”, and compares that with keys made from the added part. It finds the token after the earlier copy and brings it back. Researchers who studied small transformers named the pair a previous-token head and an induction head. Without the residual connection, block 1's result would replace the row, and block 2's query would no longer know which token it is reading.

What this is

Now the real network, on a test story whose hero is given the invented name “vorna”, cut into “v”, “##or” and “##na”. The film follows the second mention: which heads work together, and what they bring back.

The question

Here is a short story whose first sentence names its hero. At its end the hero comes back, and the first piece of the name is already written. Type any name, or try the presets: which network finishes it?

The bars are each network's guesses for the piece right after the first piece of the name. Every name here is cut into pieces of chapter 6's vocabulary, so a name like “lily” is a single token and there is nothing to finish. The tools below take the GPT apart.

Tool 1: the residual stream

A token's row is 64 numbers, and it goes up through the 4 blocks. Each block adds to it twice: once with what attention brought back from earlier positions, and once with what the MLP made of the result. Nothing is replaced, so the row works like a shared notepad that every part reads and adds to; it is often called the residual stream. Type a phrase and follow its last token up the stack.

Each strip is drawn on its own scale, so compare colors within a strip, and sizes with the numbers beside it. The size of a row is its length: the square root of the sum of its numbers squared.

Tool 2: layer norm

Since every part adds to the rows, the rows can grow as they go up, and different numbers in a row can drift to very different sizes. Before attention and before the MLP, each block passes the row through a layer norm: subtract the row's average from each of its 64 numbers, divide by their spread (the standard deviation), then multiply each number by a learned gain and add a learned shift. The part then reads numbers of a steady size, whatever happened below it. The row itself, in the residual stream, is not changed; only the copy the part reads is.

Tool 3: a block, and the stack

Each of the 4 blocks has the same parts: a layer norm, attention with 4 heads (each with queries, keys and values of 16 numbers, and chapter 7's score for each distance back), a table W_o that turns the 4 heads' results back into a row of 64 so it can be added, a second layer norm, and an MLP of 256 neurons. After the last block, one more layer norm, and an output layer scores the 2,000 tokens. Type a phrase and see what comes out.

In all, 461,136 numbers to learn, 1.3 times chapter 7's 352,848. The output layer and the table take more than half of them, as in chapters 6 and 7; the 4 blocks take 203,008. Training is chapter 7's: the same loss, the same Adam, every position of whole stories at once, with 16 stories per step instead of 8 and a learning rate of 0.004 instead of 0.008. One thing in the data is new. In TinyStories, about two thirds of the stories have as their hero one of about two dozen names, such as “lily” or “tom”, each a single token. So each time such a story was drawn for a training step, with chance one half, its hero got a new made-up name, such as “nuto” or “gakobe”, built from random syllables, cut into pieces by the tokenizer, and put everywhere the name stands as a word, “tom's” included. The six invented names this page tests with, and every pair of their pieces, were never used. Secret 3 shows why this matters.

From attention to a GPT

Chapter 7 placed attention inside the transformer. This animation goes the other way, from chapter 7's network to this one, and shows how its parts work together: the residual stream that every part reads and adds to, a block, heads in different blocks working together, and the whole GPT.

How big real GPTs are, next to this one:

This page's blocks are GPT-2's, layer norm first, with three differences: positions are marked by chapter 7's score for each distance back, not by a row added for each position; the tables of queries, keys, values and W_o have no biases; and the output layer has a table of its own, where GPT-2 reuses the embedding table.

The words, and how they fit

Residual stream

The row of numbers each position carries up through the blocks. Every part reads it and adds to it, through a residual connection, so nothing written into it is lost unless a later part adds the opposite.

Layer norm

Rescaling a copy of each row to an average of 0 and a spread of 1, then a learned gain and shift for each number. Each block has two, one before attention and one before the MLP, and one more comes after the last block.

MLP

A small network applied to each position alone: 64 numbers in, 256 neurons, 64 numbers out, added to the row. Its neurons use GELU, a smooth relative of the ReLU: close to 0 for large negative inputs (never below −0.17), and about the input itself for large positive ones.

Block

Layer norm, multi-head attention, add; layer norm, MLP, add. Rows in, rows out, so blocks can be stacked.

Previous-token head

A head that reads mostly the token just before. Its result tells each position what came before it.

Induction head

A head in a later block that reads the position right after an earlier copy of the current token, and brings that token back: after “z”, what came after “z” before. It relies on a previous-token head below it, because the keys it matches are made from what the previous-token head added.

In-context learning

Doing better on text because of what came earlier in the same text, with no change to the weights. Finishing an invented name is a small case of it.

Logit lens

Reading a row partway up the stack as if it were the last: pass it through the last layer norm and the output layer, and see what the network would guess at that point. (A logit is a score before softmax.)

GPT

A generative pre-trained transformer: an embedding table, a way of marking positions, a stack of blocks with causal attention, a last layer norm and an output layer, trained to guess the next token. This page's is one, at a tiny size.

Secret 1: two blocks finish an invented name

Chapter 7's network could not finish an invented name: after the first piece of a later mention, it guessed the next piece only for pairs of pieces that appear in the training stories. Here is the same test, given to the GPT. Each of the page's name stories gets the invented name, and the table shows the GPT's first guess for each piece, at each mention.

Measuring…

Which heads do it? At every later mention of the invented names, here is each head's average weight on the token just before, and on the token right after the earlier copy of the one being read (from the second “v”, the earlier “##or”). Heads that put at least 30% of their weight on the token just before are tinted blue; heads in block 2 or above that put at least 20% on the token after the earlier copy, orange. Click a head to see where it looks in a story.

Measuring…

Secret 2: switch the helper off

Reading the right token does not prove a head is needed. To test it, switch a head off (its value table set to 0, so it adds nothing to the rows) and measure again. If the induction head relies on the previous-token head below it, switching off the previous-token head should break finishing names too, though that head reads mostly the token just before.

Measuring…

Secret 3: when it appeared

Induction heads are known for appearing suddenly. In small transformers trained on large amounts of text, there is a point early in training where they form and the network's use of earlier context jumps, often with a visible bump in the loss. Here is this GPT, measured at checkpoints saved during its training, next to the same network trained for as many steps on the stories as they are, without made-up names.

You can watch a jump like this happen in your browser. Below, a fresh network of 2 blocks with 2 heads each, attention only (no MLP neurons), trains on sequences of 8 to 24 different words drawn at random, each followed by the same words again. In the second copy, from its second word on, every word can be found by induction and by nothing else: the words are random and the length changes, so neither knowing words nor counting places helps. Press Train.

The grids show who each position reads in the test sequence below them: a row for each position that reads, a column for each position read, darker for more weight. A previous-token head draws a line just below the diagonal; an induction head draws a line parallel to it, in the second copy, one place to the right of the earlier copy.

Secret 4: what it thinks along the way

The residual stream can be read before the top. Take a token's row as it comes out of block 1, 2 or 3, pass it through the last layer norm and the output layer as if it were the last, and see what the network would guess at that point. Interpretability researchers call this the logit lens. It is not exact, because the output layer was trained to read only the top row, but it often shows a guess taking shape.

Measuring…

Secret 5: why the residual connections matter

The same network was trained three times more, for 3,000 steps, from the same seed, with the same stories and learning rate: without its residual connections (each block's output is only what its parts make, with nothing added to the row), and without its layer norms (each part reads the row as it is); the third run, without residual connections and with ordinary starting sizes for its tables, is told in the text below. Here is how well each guessed the next token on the validation stories during those steps, next to the GPT's own first 3,000.

The algorithm: a GPT

Five parts, from the weights to a training step. This is the JavaScript this page actually runs, or the same program in Python. The tokenizer is chapter 6's.

The playground

Write the start of a story, click any token to see where a head looks from there, and let the GPT write on. Or train your own GPT from scratch, and watch its heads find their jobs.