Skip to content

Output

The path from text to the next token in seven steps: tokens, embeddings, query/key/value, mask and heads, the block, the output and generation. The current step is Output. ATTENTION TRANSFORMERS Tokenstext → ids Embeddingsids → X Q · K · Vone head Masktwo heads BlockLN · MLP Outputlogits → p ▲ you are here Generationpick text → → next token

Step 6 of 7. In: \(y\) from the block, \(5 \times 4\). Out: for every position, a probability for each of the 12 tokens of the vocabulary; the last row says what comes after "the cat climbed the". ← The block · next step: Generation →

A stack of blocks

The toy has one block. A real model stacks \(L\) of them, the output of each one being the input of the next, each taking \(n \times d\) and giving back \(n \times d\): GPT-2 small has \(L = 12\) with \(d = 768\), Llama 3 8B has 32 with \(d = 4096\). Nothing on this page changes when \(L\) grows — it only reads the output of the last block, which in the toy is the \(y\) of the previous page.

The final norm

Before the output layer, one more LayerNorm, the same as inside the block (GPT-2 calls it ln_f): \(f = \text{LN}(y)\), row by row. For the last token, ␣the, \(y_5 = [-0.436, 1.551, -0.461, 1.698]\): mean \(\mu = 2.352/4 = 0.588\), standard deviation \(\sigma = 1.038\), so

\[ f_5 = \frac{y_5 - \mu}{\sigma} = [-0.987,\ 0.928,\ -1.011,\ 1.069] \]

Unembedding: one score per vocabulary token

The output layer is a matrix \(W_U\) of \(d \times V\) numbers, \(4 \times 12\) in the toy: column \(j\) is a direction in the \(d\)-dimensional space that favours token \(j\). Each score — a logit — is a dot product:

\[ z = f\,W_U \qquad\qquad z_j = f \cdot W_U[:, j] \]
\(W_U\) the ␣the ␣cat ␣dog ␣climb ed ␣roof ␣wall ␣sofa ␣floor . <|endoftext|>
row 1 0 1 −0.5 −0.5 1 −0.5 −0.5 0 0 0 1 −0.5
row 2 0 1 −1 −0.5 −0.5 1 1 1 0.5 0.5 0 0
row 3 0 −1 0 0 −0.5 1 0 0 0 −0.5 0 1
row 4 0 −0.5 1 0.5 0 0 1 0.5 0.5 0 −1 0

For ␣roof, column \([-0.5, 1, 0, 1]\):

\[ z_{\texttt{roof}} = (-0.987)(-0.5) + 0.928 \cdot 1 + (-1.011) \cdot 0 + 1.069 \cdot 1 = 0.494 + 0.928 + 1.069 = 2.491 \]

GPT-2's output layer is \(768 \times 50{,}257\), and it is not a separate matrix: it is the embedding table \(E\) transposed (weight tying), the same 38.6 million numbers read in both directions. The toy has its own \(W_U\) for a reason you can check: with \(d = 4\) and one block, the last vector still looks so much like the token it holds that the tied version, \(z = fE^\top\), would put the highest probability on ␣the itself. A deeper, trained model has moved far enough from its input for tying to work.

Softmax: from scores to probabilities

\[ p_j = \frac{e^{z_j}}{\sum_k e^{z_k}} \]
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 1 0.033
<|endoftext|> −0.518 0.596 0.020
␣climb −0.945 0.389 0.013
. −2.056 0.128 0.004
sum 30.521 1

After "the cat climbed the", ␣roof leads with 0.395, and the four places a cat climbs onto share 0.71. The model has not written anything yet: it produced a distribution, and choosing from it is the next step.

Every position predicts

The two lines above run on every row, and each row predicts the token after its own position — the causal mask guarantees that it did not see that token. Training compares every row with the token that really came next and averages \(-\ln p\) of that token, the cross-entropy loss:

position after target \(p\)(target) \(-\ln p\) most likely
1 the ␣cat 0.500 0.693 ␣cat
2 ␣cat ␣climb 0.444 0.811 ␣climb
3 ␣climb ed 0.274 1.295 ed
4 ed ␣the 0.311 1.169 ␣the
5 ␣the ␣roof 0.395 0.928 ␣roof

Mean: 0.979 nats, against \(\ln 12 = 2.485\) for a uniform guess. Training is gradient descent on this number, for every row of every sequence at once; generation reads only the last row.

What attention bought

Switch the attention branch off (\(h = x\) in the block) and the last row knows only that it holds ␣the at position 5: ␣cat 0.455, ␣dog 0.143, ␣roof 0.143 — what follows any "the". The words "cat climbed" are what moved ␣roof from 0.143 to 0.395, and attention is the only way they reached the last row. The lab does the same to GPT-2 and gets the same verdict.

Check yourself — does the model output a word?

No: it outputs 12 probabilities (50,257 in GPT-2) that add up to 1, one per token of its vocabulary. A word appears only after a token is chosen from them and its id is turned back into text — the next page.

Check yourself — add 10 to every logit of the last row. What changes?

Nothing: \(e^{z_j + 10} = e^{10} e^{z_j}\), and the \(e^{10}\) cancels between the numerator and the sum. Only differences between logits matter, which is why a logit has no meaning on its own.

← The block · next step: Generation →