Chapter 17A Transformer You Can Check By Hand

Why do this arithmetic

Chapter 16 took the transformer apart and explained each piece, and 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, which comes to 172 parameters against GPT-3’s 175 billion, and every equation is the same. That last sentence is the reason for the chapter, because scale changes what a model can do without changing what a model is. If you can follow four numbers through sixteen steps then 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 is four tokens in and one prediction out.

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

One warning about the weights before we start. They are chosen rather than trained, because 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, so the tokenizer of Chapter 5 converts the string to integers, and here every word already sits in the vocabulary, which makes the mapping a lookup.

Step 2: each integer indexes one row

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

To make the arithmetic readable the four dimensions have been given meanings, namely thing, action, quantifier and 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 while chuck is an action that is half about wood, and that overlap is what makes the example worth running.

Step 3: position, because attention cannot see order

Self-attention takes a set of vectors, so permuting the input permutes the output with it and changes nothing else. A woodchuck and woodchuck a would therefore be identical, which is a fatal flaw.

The fix is to make position part of the vector, and 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, while 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 rather than concatenated, which is worth a pause. The model gets four dimensions in total, for meaning and position together, and it has to share them.

This is the block’s input, and 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$ 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, and it is the reason a stack of ninety-six blocks can be trained at all, since the residual path from the input to the loss passes through no normalisation and gradients therefore 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, one per projection.

$$\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, so its query reads the action dimension while its key reads the thing dimension, and head 2 asks what else here is about wood. The comparison is then a scaled dot product, and the scores below are head 1’s.

$$\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$ and is nothing of the kind at $d_k = 64$, for the reason Chapter 16 measured: 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 that $e^{-\infty} = 0$ and no weight survives there.

$$\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 are worth reading here. The first row of each head is $1.0000$ and it has no choice, because the first token can attend to nothing but itself, so the softmax is over a single value and whatever the model knows about a at this point it knew before attention ran.

Every row sums to one, which is the softmax doing its job and the reason 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, which 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, and the heads are then 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 comes 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, which was near $1.5$ for chuck before attention and is $3.25$ after, because the wood dimension has been amplified by looking at woodchuck and its 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 and four out, with every position going through this independently, using the same weights and with no reference to any other position.

Count the zeros. Of the 32 hidden values, 22 are zero, which is $69$ per cent, and they are zero because the ReLU cut them off. Note that the pattern is uneven: row a keeps four of its eight units alive while the other three rows keep only two each. This sparsity is not an accident of a toy, because 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, while the feed-forward layer is the only part that transforms a position on its own, and 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, and a large model is this pair ninety-six times over with wider matrices.

From vectors back to words

Step 13: only the last position matters

The block produced four vectors and only the last one is used, because only the last one has seen the whole prompt, while 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, and the transformer uses one it already has in $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, because if a row of $E$ says what a token means on the way in then 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, and it is worth being clear about why, since 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, and 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$, because they have identical embedding rows and 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, and cross-entropy asks one question, which is what probability you gave 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, because for a softmax followed by cross-entropy the gradient with respect to the logits collapses to something you 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, so it is the predicted probability minus what it should have been, and that is the whole derivative.

One gradient is negative and seven are positive, so the negative one raises the right answer while 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

Only three numbers changed the parameter count: the width, the number of blocks and the vocabulary. The head count moved too, and it is free, for the reason Chapter 16 gave. No equation in this chapter was added to, removed or altered, so 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.

Then build the same block as torch.nn.TransformerEncoderLayer(d_model=4, nhead=2, dim_feedforward=8) and sum its parameters. It comes to $172$, the same total as Step 16, and not one of its parts matches. The library spends nothing on an embedding and spends the $32$ that buys it here on attention biases and two layer norms instead. Itemise both columns until you can say where every parameter went, because a total that agrees for the wrong reasons is the easiest kind of number to trust by mistake.

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

From Natural Language Processing: From Corpus Statistics to Grounded LLMs, by Ramaseshan Ramachandran.