Chapter 17A Transformer You Can Check By Hand

Why do this arithmetic

Chapter 16 took the transformer apart and explained each piece. This chapter puts one back together and runs it, with numbers small enough to check on paper.

The model here has four dimensions, two heads, one block and a vocabulary of eight. It has 172 parameters. GPT-3 has 175 billion. Every equation is the same.

That last sentence is the reason for the chapter. Scale changes what a model can do. It does not change what a model is. If you can follow four numbers through sixteen steps, you have followed the architecture, and the only thing left to learn about the large ones is what emerges when the same arithmetic is run at a size no one can trace.

The prompt.

Four tokens in, one prediction out.

The vocabulary is a, chuck, how, much, wood, woodchuck, would, and the full stop. Eight tokens, indexed alphabetically from zero.

A warning about the weights.

They are chosen, not trained. Every matrix in this chapter was written by hand so that its job is visible, and the chapter says what each one is for. A trained model reaches its weights by the process of Chapter 8, and they are never this tidy. What is genuinely the same is the arithmetic they go through, which is the subject here.

From words to vectors

Step 1: the prompt becomes integers

Nothing in a model consumes text. The tokenizer of Chapter 5 converts the string to integers, and here every word already sits in the vocabulary, so the mapping is a lookup.

Step 2: each integer indexes one row

The embedding table $E$ is $8 \times 4$: one row per token, four numbers per row. The lookup is not a multiplication. The row is the answer.

To make the arithmetic readable, the four dimensions have been given meanings: thing, action, quantifier, wood. A trained table has no such labelled columns, which is exactly the difficulty of Chapter 9.

Woodchuck is a thing that is half about wood. Chuck is an action that is half about wood. That overlap is what makes the example worth running.

Step 3: position, because attention cannot see order

Self-attention takes a set of vectors. Permute the input and the output permutes with it, unchanged. So a woodchuck and woodchuck a would be identical, which is not a subtlety, it is a fatal flaw.

The fix is to make position part of the vector. Chapter 16 gave the sinusoids:

$$\begin{equation} PE_{(pos,\,2i)} = \sin\!\left(\frac{pos}{10000^{2i/d}}\right), \qquad PE_{(pos,\,2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d}}\right). \label{eq:tt-pe} \quad\text{(17.1)} \end{equation}$$

With $d = 4$ there are two frequencies: one that turns once per token, and one a hundred times slower.

The fast pair separates neighbours. The slow pair barely moves across four tokens, and would be doing the work if the context were a thousand long.

Step 4: the two are added

$$\begin{equation} x_p = E[\mathrm{id}_p] + PE_p. \label{eq:tt-input} \quad\text{(17.2)} \end{equation}$$

Added, not concatenated, which is worth a pause. The model is not given four dimensions for meaning and four for position. It is given four in total, and has to share them.

This is the block’s input. Everything that follows is a function of these sixteen numbers.

The attention sublayer

Step 5: normalise first

$$\begin{equation} \mathrm{LN}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad \mu = \frac{1}{d}\sum_j x_j, \qquad \sigma^2 = \frac{1}{d}\sum_j (x_j - \mu)^2. \label{eq:tt-ln} \quad\text{(17.3)} \end{equation}$$

Each row is normalised on its own, across its four dimensions, with no reference to the other tokens or to the batch. Row a has mean $0.75$ and variance $0.1875$, so its first entry becomes $(0 - 0.75)/\sqrt{0.1875 + \epsilon} = -1.7320$ with $\epsilon = 10^{-5}$. That $\epsilon$ is not decoration. It is what stops a row of identical values, whose variance is zero, from dividing by nothing.

Normalising before the sublayer rather than after it is the pre-norm arrangement. It is the reason a stack of ninety-six blocks can be trained at all: the residual path from the input to the loss passes through no normalisation, so gradients reach the early layers undiminished. This is the gradient highway of Chapter 13, built again in a different architecture.

Steps 6 and 7: three projections, then one comparison

Each head sees the same normalised rows through three different matrices.

$$\begin{equation} Q = X_n W^Q, \qquad K = X_n W^K, \qquad V = X_n W^V. \label{eq:tt-qkv} \quad\text{(17.4)} \end{equation}$$

Head 1 was written to ask what thing am I acting on: its query reads the action dimension, its key reads the thing dimension. Head 2 asks what else here is about wood.

The comparison is a scaled dot product.

$$\begin{equation} S = \frac{QK^{\top}}{\sqrt{d_k}}, \qquad d_k = 2, \qquad \sqrt{d_k} = 1.4142. \label{eq:tt-scores} \quad\text{(17.5)} \end{equation}$$

The division by $\sqrt{d_k}$ looks cosmetic at $d_k = 2$. It is not cosmetic at $d_k = 64$, and Chapter 16 measured why: the dot product of two vectors of unit variance has standard deviation $\sqrt{d_k}$, and a softmax fed large numbers saturates onto one key before training has begun.

Step 8: the mask, then the softmax

The upper triangle is set to $-\infty$ before the softmax, so $e^{-\infty} = 0$ and no weight survives.

$$\begin{equation} A = \mathrm{softmax}\!\left(S + M\right), \qquad M_{ij} = \begin{cases} 0 & j \le i \\ -\infty & j > i \end{cases} \label{eq:tt-mask} \quad\text{(17.6)} \end{equation}$$

Three things to read here.

The first row of each head is $1.0000$, and it has no choice. The first token can attend to nothing but itself, so the softmax is over a single value. Whatever the model knows about a at this point, it knew before attention ran.

Every row sums to one. That is the softmax doing its job, and it is why attention is called a weighted average rather than a sum.

The two heads disagree, which is the argument for having more than one. Looking from woodchuck, head 1 sends three quarters of its weight back to a, the determiner that introduced it. Head 2 keeps $0.8618$ on woodchuck itself, because it is asking about wood and woodchuck is the wood-bearing token so far. One mechanism, two questions, in parallel.

The honest reading of the last row is that both heads keep most of their weight on the current token. That is what untrained attention does. Specialisation is something training produces, and the heads that get published as clean pictures of syntax are heads that were trained into that shape.

Steps 9 and 10: combine, project, and add back

Each head returns a weighted average of its values. The heads are concatenated back to width four and mixed by $W^O$.

$$\begin{equation} \mathrm{MHA}(X) = \left[\mathrm{head}_1 ; \mathrm{head}_2\right] W^O, \qquad \mathrm{head}_h = A_h V_h. \label{eq:tt-mha} \quad\text{(17.7)} \end{equation}$$

Then the residual connection, which is one addition and the most important addition in the architecture.

$$\begin{equation} x' = x + \mathrm{MHA}(\mathrm{LN}(x)). \label{eq:tt-residual} \quad\text{(17.8)} \end{equation}$$

Read the last column. It was near $1.5$ for chuck before attention and is $3.25$ after. The wood dimension has been amplified by looking at woodchuck and wood-ish neighbours. Information has moved between positions, which is the one thing attention is for.

The feed-forward sublayer

Step 11: widen, rectify, narrow

$$\begin{equation} \mathrm{FFN}(z) = \max(0,\, z W_1 + b_1)\, W_2 + b_2. \label{eq:tt-ffn} \quad\text{(17.9)} \end{equation}$$

Four dimensions in, eight hidden, four out. Every position goes through this independently, with the same weights and no reference to any other position.

Count the zeros. Of the 32 hidden values, 24 are zero, and they are zero because the ReLU cut them off. This sparsity is not an accident of a toy: the feed-forward layers of real models are sparsely active in exactly this way, and that observation is what makes mixture-of-experts routing possible, a thread Chapter 24 picks up.

Note also that the division of labour is now complete. Attention is the only part of the block that moves information sideways. The feed-forward layer is the only part that transforms a position on its own. Two thirds of the parameters in a real transformer sit in this second job.

Step 12: the second residual

$$\begin{equation} y = x' + \mathrm{FFN}(\mathrm{LN}(x')). \label{eq:tt-block} \quad\text{(17.10)} \end{equation}$$

Equations (17.8) and (17.10) together are the block. A large model is this pair, ninety-six times, with wider matrices.

From vectors back to words

Step 13: only the last position matters

The block produced four vectors. Only the last one is used, because only the last one has seen the whole prompt. The other three were needed to build it.

After a final layer norm, the last position is

$$z = (-0.6805,\; -0.5283,\; -0.5198,\; 1.7285).$$

Step 14: the tied unembedding

To turn four numbers into a distribution over eight tokens, multiply by a $4 \times 8$ matrix. The transformer uses one it already has: $E^{\top}$, the embedding table read the other way.

$$\begin{equation} \ell = z E^{\top}, \qquad p = \mathrm{softmax}(\ell). \label{eq:tt-logits} \quad\text{(17.11)} \end{equation}$$

This is weight tying, and the argument for it is exactly the distributional hypothesis. If a row of $E$ says what a token means on the way in, the same row can score how well the output matches that token on the way out. It also saves $V \times d$ parameters, which is 32 here and 38 million in GPT-2.

The model’s answer is wood, at 31 per cent.

It is worth being clear about why, because it is not because the model knows the tongue twister. The last position accumulated a large value in the wood dimension, by attending to woodchuck and chuck. The row of $E$ with the largest wood component is wood. The dot product in Equation (17.11) is doing nothing more than measuring that agreement.

Note that how and much tie exactly, at $-0.2599$. They have identical embedding rows, so nothing anywhere in the model can ever separate them. Two tokens with the same vector are the same token, whatever the spelling.

Learning, in one step

Step 15: the loss and its gradient

The target is wood. Cross-entropy asks one question: what probability did you give the right answer?

$$\begin{equation} L = -\log p_{t} = -\log 0.3107 = 1.1690 \text{ nats}. \label{eq:tt-loss} \quad\text{(17.12)} \end{equation}$$

Now the derivative, which is the reason softmax and cross-entropy are always paired. For the softmax followed by cross-entropy the gradient with respect to the logits collapses to something a student can write down without any chain rule at all.

$$\begin{equation} \frac{\partial L}{\partial \ell} = p - y, \label{eq:tt-grad} \quad\text{(17.13)} \end{equation}$$

where $y$ is the one-hot target. Predicted probability minus what it should have been. That is the whole derivative.

One gradient is negative and seven are positive. The negative one raises the right answer, the positive ones lower every wrong one, and their magnitudes are exactly the probabilities that were wrongly assigned. A step of size $0.5$ takes the loss from $1.1690$ to $0.9119$, an improvement of $0.2570$.

That is one step, on one example, on the logits alone. Training pushes this same vector backwards through Equation (17.11), then (17.10), (17.9), (17.8), (17.7), (17.6), (17.5), (17.4), (17.3) and into $E$ itself, by the chain rule of Chapter 8, and repeats it for every token of every document in the corpus.

What scale changes

Step 16: counting this model

The same equations, four sizes

Three numbers changed: the width, the number of blocks, and the vocabulary. No equation in this chapter was added to, removed, or altered. A billion-parameter model is Equations (17.2) through (17.11), run wider and more often.

What does change is what the weights end up containing, and that is the subject of the rest of the book. Chapter 18 asks what happens when this block is trained on a corpus instead of written by hand, and what changes when the causal mask is removed.

Try it yourself. code/worked_examples/tiny_transformer.py produces every table in this chapter. Run it with no arguments for all sixteen steps in order, or take one stage at a time: --embed for the lookup and the sinusoids, --attention for both heads with their masks, --ffn for the sparsity of the hidden layer, --logits for the tied unembedding, --loss for the gradient step, and --params for the count.

Change one weight and run it again. The fastest way to believe the architecture is to break it: set the mask aside and watch the model see its own answer, or zero the positional encoding and watch a woodchuck and woodchuck a become the same prompt.

Run it in ColabNotebookSource

Code for this chapter

Each one reproduces the tables above. Open it in Colab, change an input and watch which numbers move.

• tiny_transformer.pyA transformer, small enough to check

  One prompt through every layer of a four-dimensional model, ending in a prediction, a loss and one step of learning.

  pip install numpy Run it in Colab Notebook Source