End to End

The two chapters opened the path one page at a time. This page closes it in one place, on the layout of Transformer Explainer and with the numbers of the toy model. The path is one numbered sequence: it starts before attention, with the embedding table and the matrix built from it, goes through the heads (choose how many) and ends in the next token. Click a step: that part of the diagram opens in place and the box alongside explains it, with its calculations. At the end of the page, the same path is done by hand.

← Generation · next: What changed since 2017 →

Where the weights are

piece	shape	toy	GPT-2 small (12 blocks)
embedding table \(E\)	\(V \times d\)	12 × 4 = 48	38.6 million
position table \(P\)	positions × \(d\)	8 × 4 = 32	786 thousand
\(W_Q\), \(W_K\), \(W_V\)	\(3 \times (d \times d)\)	48	21.2 million
\(W_O\)	\(d \times d\)	16	7.1 million
MLP: \(W_1\), \(W_2\)	\(d \times\) width and width \(\times\, d\)	32 + 32 = 64	56.6 million
output \(W_U\)	\(d \times V\)	48	tied to \(E\)
total		256	≈ 124 million

The GPT-2 counts leave out the biases and the LayerNorms, a small share. In the toy the LayerNorms have no learned weight and the weights were picked by hand, so that every sum can be checked. The full table compares the two models.

To check

Check yourself — if all the keys are identical, what is the output?

The scores come out equal, the softmax turns uniform and the output is the plain average of the values, the same for any query.

Check yourself — what is left if you remove attention from every block?

Each token goes through the stack without knowing its neighbours. After "the", the model would always say the same thing, whether the sentence was about a cat or a bank.

Check yourself — does the temperature change which token is the most probable?

No. Dividing the logits by the same positive \(T\) keeps their order. Only how much mass the other tokens get changes.

Check yourself — add 10 to all the logits of the last row. What changes?

Nothing: the \(e^{10}\) cancels between the numerator and the sum. Only differences between logits matter.

A worked example, step by step

The whole path by hand, from "the cat climbed the" to the next token, with \(h = 2\) heads, the mask on and \(T = 1\). These are the numbers of the simulator with the sentence "the cat climbed the". They appear rounded to 3 places, but each sum carries all the places, so a last digit may differ by one.

Step 1. From text to ids

The tokenizer cuts the text into pieces of a vocabulary fixed beforehand, and each piece becomes the number of its row:

token	`the`	`␣cat`	`␣climb`	`ed`	`␣the`
id	0	2	4	5	1

Step 2. The table \(E\): how it is calculated

\(E\) exists before any sentence and before attention. It has one row per token of the vocabulary and one column per dimension: \(V \times d = 12 \times 4\). Its numbers come from three calculations, in this order.

a. Creation: random numbers. A real model starts each number from a draw. GPT-2 draws from \(\mathcal{N}(0,\ 0.02^2)\), with \(V \times d = 50257 \times 768\). A row drawn this way, with \(d = 4\):

\[[0.003, -0.003, 0.013, 0.002]\]

This row still means nothing: at this point the vector of ␣roof is as good as that of any other token.

b. Training: each number moves against its derivative. With the loss \(L\) of a prediction and a learning rate \(\eta\), the update is

\[E_{i,c} \leftarrow E_{i,c} - \eta\, \frac{\partial L}{\partial E_{i,c}}\]

The toy skips the draw and the many steps: its table was picked by hand, so that every sum can be checked. This is it, as if already trained, with the rows of this sentence in bold:

id	token	\(c_1\)	\(c_2\)	\(c_3\)	\(c_4\)
0	`the`	−0.5	−0.5	0	1
1	`␣the`	−0.5	0	−0.5	1
2	`␣cat`	1	−0.5	0	−0.5
3	`␣dog`	1	−0.5	−0.5	0
4	`␣climb`	−0.5	1	0	−0.5
5	`ed`	0	0.5	−0.5	0
6	`␣roof`	0	−0.5	1	−0.5
7	`␣wall`	1	0	−1	0
8	`␣sofa`	0.5	−0.5	−0.5	0.5
9	`␣floor`	0.5	0	−0.5	0
10	`.`	0	0	1	−1
11	`<eot>`	0.5	−0.5	0.5	−0.5

One training step on it, to see the calculation. The position "the cat climbed the" should predict ␣roof, so the loss is

\[L = -\ln p(\text{␣roof}) = -\ln 0.395 = 0.928\]

The derivative of \(L\) for each row of the sentence is:

id	token	\(\partial L / \partial E_{i}\)
0	`the`	[−0.001, 0.000, 0.001, 0.000]
2	`␣cat`	[0.000, 0.000, 0.026, −0.026]
4	`␣climb`	[−0.001, 0.000, 0.001, 0.000]
5	`ed`	[0.002, 0.002, 0.000, −0.004]
1	`␣the`	[−0.011, 0.011, 0.027, −0.028]

The other 7 rows have derivative exactly 0: a token that is not in the sentence does not move. With \(\eta = 0.1\), the row of ␣the (id 1) becomes:

\[\begin{aligned} -0.5 - 0.1 \cdot (-0.011) &= -0.499 \\ 0.0 - 0.1 \cdot (0.011) &= -0.001 \\ -0.5 - 0.1 \cdot (0.027) &= -0.503 \\ 1.0 - 0.1 \cdot (-0.028) &= 1.003 \end{aligned}\]

After this one step, with every row updated, \(p(\text{␣roof})\) goes from \(0.395\) to \(0.396\) and the loss from \(0.928\) to \(0.927\). Repeating this over many sentences is what fills the table. The rest of this example uses the table above, without the step.

c. Reading: \(E[\text{id}]\). With the table ready, reading a row is a product by a vector of zeros with a 1 at the id. For ␣the, id 1:

\[E[1] = [0,\ 1,\ 0,\ 0,\ 0,\ 0,\ 0,\ 0,\ 0,\ 0,\ 0,\ 0]\ E = [-0.5, 0.0, -0.5, 1.0]\]

The five reads, one per token, make the matrix \(E[\text{ids}]\):

token	id	\(E[\text{id}]\)
`the`	0	[−0.5, −0.5, 0.0, 1.0]
`␣cat`	2	[1.0, −0.5, 0.0, −0.5]
`␣climb`	4	[−0.5, 1.0, 0.0, −0.5]
`ed`	5	[0.0, 0.5, −0.5, 0.0]
`␣the`	1	[−0.5, 0.0, −0.5, 1.0]

Step 3. The position table \(P\) and the matrix \(X\)

\(E\) does not say where the token is. Each position has a row of \(P\), also a table of weights, and the sum of the two rows is the row of \(X\), the only thing the block receives:

token	\(E[\text{id}]\)	\(P_t\)	\(x_t = E[\text{id}] + P_t\)
`the`	[−0.5, −0.5, 0.0, 1.0]	[0.0, 0.0, 0.0, 0.0]	[−0.5, −0.5, 0.0, 1.0]
`␣cat`	[1.0, −0.5, 0.0, −0.5]	[0.5, 0.0, −0.5, 0.0]	[1.5, −0.5, −0.5, −0.5]
`␣climb`	[−0.5, 1.0, 0.0, −0.5]	[0.0, 0.5, 0.0, −0.5]	[−0.5, 1.5, 0.0, −1.0]
`ed`	[0.0, 0.5, −0.5, 0.0]	[−0.5, 0.0, 0.5, 0.0]	[−0.5, 0.5, 0.0, 0.0]
`␣the`	[−0.5, 0.0, −0.5, 1.0]	[0.0, −0.5, 0.0, 0.5]	[−0.5, −0.5, −0.5, 1.5]

Step 4. LayerNorm

From here on we follow the last row, the one that predicts the next token, ␣the. Each row is normalized on its own:

\[\mu = 0.000,\quad \sigma = 0.866,\quad u_5 = \frac{x_5 - \mu}{\sigma} = [-0.577, -0.577, -0.577, 1.732]\]

The five rows of \(u\):

token	\(u_1\)	\(u_2\)	\(u_3\)	\(u_4\)
`the`	−0.816	−0.816	0.000	1.633
`␣cat`	1.732	−0.577	−0.577	−0.577
`␣climb`	−0.535	1.604	0.000	−1.069
`ed`	−1.414	1.414	0.000	0.000
`␣the`	−0.577	−0.577	−0.577	1.732

Step 5. Query, Key, Value

Three weight matrices multiply \(u\). In the toy they only pick columns: \(Q = [u_4, u_2, u_4, u_1]\), \(K = [u_2, u_4, u_1, u_4]\) and \(V = [u_1, u_2, u_3, u_4]\). Head 1 reads columns 1 and 2 and head 2 reads columns 3 and 4.

token	\(Q\)	\(K\)	\(V\)
`the`	[1.633, −0.816, 1.633, −0.816]	[−0.816, 1.633, −0.816, 1.633]	[−0.816, −0.816, 0.000, 1.633]
`␣cat`	[−0.577, −0.577, −0.577, 1.732]	[−0.577, −0.577, 1.732, −0.577]	[1.732, −0.577, −0.577, −0.577]
`␣climb`	[−1.069, 1.604, −1.069, −0.535]	[1.604, −1.069, −0.535, −1.069]	[−0.535, 1.604, 0.000, −1.069]
`ed`	[0.000, 1.414, 0.000, −1.414]	[1.414, 0.000, −1.414, 0.000]	[−1.414, 1.414, 0.000, 0.000]
`␣the`	[1.732, −0.577, 1.732, −0.577]	[−0.577, 1.732, −0.577, 1.732]	[−0.577, −0.577, −0.577, 1.732]

Step 6. The attention of each head

The row of ␣the against each column, \(s_{5,j} = q_5 \cdot k_j / \sqrt{2}\), then \(a_{5,j} = e^{s_{5,j}} / \sum e^{s}\). It is the last row, so the mask erases nothing here.

Head 1, with \(q_5 = [1.732, -0.577]\):

\(j\)	\(k_j\)	\(q_5 \cdot k_j\)	\(s_{5,j}\)	\(e^{s}\)	\(a_{5,j}\)
`the`	[−0.816, 1.633]	−2.357	−1.667	0.189	0.011
`␣cat`	[−0.577, −0.577]	−0.667	−0.471	0.624	0.035
`␣climb`	[1.604, −1.069]	3.395	2.400	11.027	0.622
`ed`	[1.414, 0.000]	2.449	1.732	5.652	0.319
`␣the`	[−0.577, 1.732]	−2.000	−1.414	0.243	0.014

Head 2, with \(q_5 = [1.732, -0.577]\):

\(j\)	\(k_j\)	\(q_5 \cdot k_j\)	\(s_{5,j}\)	\(e^{s}\)	\(a_{5,j}\)
`the`	[−0.816, 1.633]	−2.357	−1.667	0.189	0.016
`␣cat`	[1.732, −0.577]	3.333	2.357	10.559	0.882
`␣climb`	[−0.535, −1.069]	−0.309	−0.218	0.804	0.067
`ed`	[−1.414, 0.000]	−2.449	−1.732	0.177	0.015
`␣the`	[−0.577, 1.732]	−2.000	−1.414	0.243	0.020

Head 1 looks at the verb, which the tokenizer cut in two (␣climb + ed = 0.94). Head 2 looks at the subject, ␣cat with 0.88. In the row of ␣climb, the third, the mask writes \(-\infty\) in positions 4 and 5 and their weights are 0: [0.930, 0.063, 0.007, 0.000, 0.000].

Step 7. Mix the values and join the heads

Each head averages the values with its weights, the two results are placed side by side and \(W_O = I\) leaves them as they are:

\[o_5^{(1)} = \sum_j a_{5,j}\, v_j = [-0.739, 1.411],\quad o_5^{(2)} = [-0.521, -0.520],\quad a_5 = [-0.739, 1.411, -0.521, -0.520]\]

Step 8. First sum, LayerNorm and MLP

\[\begin{aligned} h_5 &= x_5 + a_5 = [-1.239, 0.911, -1.021, 0.980] \\ u'_5 &= \mathrm{LN}(h_5) = [-1.102, 0.964, -0.892, 1.030] \\ u'_5 W_1 &= [u'_5\ |\ -u'_5] = [-1.102, 0.964, -0.892, 1.030, 1.102, -0.964, 0.892, -1.030] \\ \mathrm{GELU} &\to [-0.149, 0.802, -0.166, 0.874, 0.952, -0.162, 0.726, -0.156] \\ m_5 &= \mathrm{GELU}(\cdot)\, W_2 = [0.803, 0.640, 0.560, 0.718] \\ y_5 &= h_5 + m_5 = [-0.436, 1.551, -0.461, 1.698] \end{aligned}\]

Step 9. The next token

One more LayerNorm, then the product by \(W_U\) gives one score per token of the vocabulary. For ␣roof:

\[f_5 = \mathrm{LN}(y_5) = [-0.987, 0.928, -1.011, 1.069],\quad z_{\text{roof}} = f_5 \cdot W_U^{(6)} = -0.987 \cdot (-0.5) + 0.928 \cdot (1.0) + -1.011 \cdot (0.0) + 1.069 \cdot (1.0) = 2.491\]

The softmax turns the 12 scores into probabilities:

token	\(z\)	\(e^{z}\)	\(p\)
`␣roof`	2.491	12.069	0.395
`␣wall`	1.463	4.318	0.141
`␣sofa`	0.999	2.715	0.089
`␣floor`	0.970	2.637	0.086
`␣cat`	0.634	1.886	0.062
`␣dog`	0.564	1.757	0.058
`␣the`	0.418	1.519	0.050
`ed`	0.410	1.507	0.049
`the`	0.000	1.000	0.033
`<eot>`	−0.518	0.596	0.020
`␣climb`	−0.945	0.389	0.013
`.`	−2.056	0.128	0.004

The most probable is ␣roof, with \(0.395\): the sentence continues as "the cat climbed the roof". Its loss, \(-\ln 0.395 = 0.928\), is the \(L\) that step 2 used to adjust the table.

← Generation · next: What changed since 2017 →