The complete transformer

The full pass, equation by equation

For a sequence of $T$ tokens with numbers $t_1,\dots,t_T$, an embedding table $E$ (one row of $d$ numbers per vocabulary word), and $L$ blocks, in GPT-2’s “pre-normalization” variant, which is the one our toy uses ($d=8$, $L=2$):

and, in each block, with $\mathrm{LN}$ the layer normalization:

Those two sums are the residual connections: each part adds its contribution to $x$. At the exit of the last block, the last token’s list gives the scores (logits) and softmax the probabilities:

“Comparing with every vocabulary word” is exactly this: a dot product with each row of $E$, the same table as at the input. Reusing it saves $V\cdot d$ parameters; the 2017 paper already did it (section 3.4), and the technique is studied in Press and Wolf (2017). In training all $T$ positions are used at once: each one predicts the next token.

Masked attention

With $Q$, $K$, and $V$ as in lesson 21 and the mask $M_{ij}=0$ if $j\le i$ and $-\infty$ if $j>i$:

Since $e^{-\infty}=0$, the future gets zero weight. That is how the original paper describes the decoder: every softmax input that corresponds to an “illegal” connection is set to $-\infty$ to preserve the auto-regressive property (section 3.2.3). Our toy has two heads of $d_h=4$ numbers per block; the table in the scene is one head of the second block.

The experiment. We trained two identical models (same seed, same data, 4,000 steps) and removed the mask from one. On new sentences, the masked model has a loss of 0.527 nats per word, almost the grammar’s entropy (0.526, the lowest possible). The unmasked model drives its training loss down to 0.059 by copying, but used the way it is used when writing (masked, with no future) it scores 3.088, worse than spreading evenly over twelve words: $\ln 12 = 2.485$.

Each block’s network

$\mathrm{MLP}(x)=W_2\,\mathrm{GELU}(W_1x+b_1)+b_2$, with $W_1$ of $d\times 4d$ and $W_2$ of $4d\times d$, the same for every position in the block (“applied to each position separately and identically,” section 3.3 of the paper). The original paper used ReLU; GPT-1 and GPT-2 use GELU (Hendrycks and Gimpel, 2016), a smooth ReLU: $\mathrm{GELU}(x)=x\,\Phi(x)$, with $\Phi$ the normal cumulative distribution.

Layer normalization

Over the $d$ numbers of a list (Ba, Kiros, and Hinton, 2016):

$\gamma$ and $\beta$ are the “two adjustments that are also learned” (gain and bias, $2d$ parameters per normalization). The 2017 paper normalized after adding, $\mathrm{LN}(x+\mathrm{Sublayer}(x))$; GPT-2 moved it to the input of each part and added one more after the last block. Xiong and colleagues (2020) show that with pre-normalization the gradients are well behaved from the start, so the learning-rate warm-up stage can be dropped. In the scene, the list for “eats” entering the second block has an average of 1.10 and a standard deviation of 0.78.

Positional encoding

The 2017 paper’s (section 3.5), with $pos$ the place and $i$ the index of the pair of numbers:

The wavelengths run from $2\pi$ to $10000\cdot2\pi$. The authors also tried learned positions and got nearly identical results; GPT-1 learns its own. “Attention alone does not know order” is exact for unmasked attention: if the tokens are permuted, their outputs are permuted the same way. With a causal mask it is subtler: Haviv and colleagues (2022) show that a causal language model without positional encoding learns an implicit notion of position, thanks to the mask itself. Even so, almost every model uses one. In the scene, the “no position” calculation uses a real head of our model on the embeddings alone, without the mask.

Counting parameters

With $V$ words, $d$ numbers per list, $L$ blocks, $ctx$ learned positions, and a $d\to4d\to d$ network:

Part	Parameters	Toy ($V{=}12$, $d{=}8$, $L{=}2$)
Embeddings (and learned positions)	$V d + ctx\,d$	96 (wave positions: 0)
Attention	$L\,(4d^2+4d)$	576
Network	$L\,(8d^2+5d)$	1,104
Normalizations	$(4L+2)\,d$	80
Total	$\approx 12\,L\,d^2$ when $d$ is large	1,856

The approximation $N\approx 12\,L\,d^2$ is Kaplan and colleagues’ (2020, equation 2.1), not counting embeddings. With GPT-2’s vocabulary (50,257) and $ctx=2048$, the table gives 125,226,240 for 12 blocks with $d=768$ and 174,604,259,328 for 96 blocks with $d=12{,}288$; the GPT-3 paper (table 2.1) lists 125 million for “GPT-3 Small” and 175.0 billion for GPT-3. The split by parts in the scene is our calculation with this table. For the smallest GPT-2 (context 1,024) it gives 124,439,808, the same count as the Hugging Face Transformers library (checked in train_tiny.py); the GPT-2 paper said 117 million.

In code

import numpy as np

def softmax(z):
    e = np.exp(z - z.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)

def causal_attention(Q, K, V):
    T, d = Q.shape
    scores = Q @ K.T / np.sqrt(d)
    scores[np.triu_indices(T, k=1)] = -np.inf      # causal mask: no peeking at the future
    w = softmax(scores)
    return w, w @ V

def layer_norm(x, g, b, eps=1e-5):
    return g * (x - x.mean(-1, keepdims=True)) / np.sqrt(x.var(-1, keepdims=True) + eps) + b

def n_params(d, L, V, ctx, ff=None):
    ff = ff or 4 * d
    return V*d + ctx*d + L*(4*d*d + 4*d) + L*(2*d*ff + ff + d) + (4*L + 2)*d

rng = np.random.default_rng(0)
Q = K = V = rng.normal(size=(3, 4))
w, out = causal_attention(Q, K, V)
print(np.round(w, 2))
x = rng.normal(2.0, 3.0, size=8)
y = layer_norm(x, np.ones(8), np.zeros(8))
print(round(y.mean(), 3), round(y.std(), 3))
print(n_params(8, 2, 12, 0), n_params(12288, 96, 50257, 2048))
# [[1.   0.   0.  ]
#  [0.25 0.75 0.  ]
#  [0.17 0.13 0.7 ]]
# 0.0 1.0
# 1856 174604259328

The lesson’s model is complete in model.ts and is trained with train_tiny.py (PyTorch, on CPU, under a minute): both live next to this lesson, and verify-forward.mts checks that the TypeScript pass matches PyTorch (maximum difference 7·10⁻⁶).

What comes next

A model with this architecture starts with random parameters. Lesson 23 shows what it learns when trained to predict the next word on internet-scale text, and lesson 24 what that costs. How it is served (the KV cache, which avoids recomputing the earlier tokens) arrives in lesson 26.