Ch 10. The transformer block

Two pieces are new. The ”+” circles mean attention and the neuron layer (the MLP) add a change to each token’s vector instead of replacing it, so the original passes straight through on a highway.^[2] Layer norm rescales the numbers to a steady size before each step.^[3] Both make big networks much easier to train.

Two pieces are new:

• The ”+” circles. Attention and the neuron layer (called the MLP) don’t replace a token’s vector; they add a change to it. The original passes straight through on a highway, and each block adds its contribution. This trick, called a residual connection, makes very deep networks much easier to train.^[2][2]
• Layer norm. Before each step, the numbers are rescaled to a steady size, which also helps training stay stable.^[3][3] GPT-2 puts it at the start of each step, plus once more at the very end.^[4][4]

A block, written down

With $x$ a token’s vector on the highway, one GPT-2 block computes

where $\mathrm{Attn}$ is the attention from Ch 9, $\mathrm{MLP}$ is the two layers of neurons from Ch 6, and $x^{\prime \prime }$ goes on to the next block. Layer norm rescales the 768 numbers of a vector by their own mean $\mu$ and variance ${\sigma }^2$:

where $\gamma$ and $\beta$ are 768 learned numbers each, $\odot$ multiplies number by number, and $ϵ$ is a tiny constant (0.00001 in GPT-2) that avoids dividing by zero.^[3][5] The original transformer put the norm after each addition instead, as $\mathrm{LN}(x+\mathrm{Attn}(x))$; GPT-2 moved it to the start of each step.^[1][4]

The counter shows where GPT-2’s 124 million weights live: most are in the token table at the start and in the MLPs. Attention, the star of the last chapter, has fewer than a quarter.

The counter shows where GPT-2’s 124 million weights live. About a third are in the token and position tables at the start. Almost half are in the MLPs, which take up two thirds of every block. Attention, the star of the last chapter, has fewer than a quarter.

Counted exactly, each block has 2,362,368 weights in attention ($4\times {768}^2$ plus $4\times 768$ biases), 4,722,432 in the MLP, and 3,072 in its two layer norms: 7,087,872 in all. Adding the token table (50,257 × 768), the position table (1,024 × 768), the 12 blocks, and the final layer norm (2 × 768) gives the total on the counter:

There’s no separate output table: the last step reuses the token table, so it isn’t counted twice.

That highway is called the residual stream. Every block reads from it and writes back to it, and the final stream becomes the prediction. You can also stop partway and ask for a prediction from the middle, a trick called the logit lens.^[4] Here it is on “The cat sat on the”:

That highway is called the residual stream. Every block reads from it and writes back to it, and the final stream becomes the prediction. So what if you stopped early and asked for a prediction from the middle of the stack? Researchers do exactly that, a trick called the logit lens.^[5][6] Here it is on “The cat sat on the”:

Early on, the top guess is ” same”, as in “sat on the same”: a sensible kind of word, but not about cats. Later, ” floor” rises to the top and furniture appears. The last block ends at ” floor”, ” bed”, and ” couch”, the guesses from Ch 0.

At the start, the stream is just the vectors for ” the” and its position, and read as a guess it’s nonsense. From block 1 to block 7 the top guess is ” same”, as in “sat on the same”: a sensible kind of word, but not about cats. From block 8 on, ” floor” is at or near the top (” ground” edges ahead once), and furniture appears: ” table”, ” sofa”, ” bed”. The last block spreads the bets out again, ending at ” floor”, ” bed”, and ” couch”, the same guesses you saw in Ch 0.

The lens reads block $\ell$’s stream $x_{\ell }$ the way the model reads its last one: apply the final layer norm, then score every token against the token table $W_E$ and take softmax,

where ${\mathrm{LN}}_f$ is the final layer norm and $W_E^{\top }$ is the token table turned on its side, so each token’s score is its vector times the stream.^[6] The lens is only a rough view: the middle blocks weren’t trained to be read this way.