Inside the transformer

Tokens, embeddings, positions

• Tokenisation. Text is split into tokens: sub-word pieces from a fixed vocabulary of size \(V\) (30k–100k in the GPT-2/GPT-3 generation; newer vocabularies are larger, e.g. Llama 3's 128k and GPT-4o's ~200k). "tokenisation" might become token | isation; a rare word becomes several pieces. Each token is an integer id.
• Embedding. A learned matrix \(W_E \in \mathbb{R}^{V \times d_{\text{model}}}\) maps each id to a vector in \(\mathbb{R}^{d_{\text{model}}}\) (one row of \(W_E\)). A length-\(T\) sequence becomes an activation matrix \(X \in \mathbb{R}^{T \times d_{\text{model}}}\), one row per position.
• Position. Self-attention is permutation-equivariant: shuffle the input rows and the output rows shuffle the same way; it has no built-in sense of order (because, as the next section shows, attention is a content-weighted sum with no term that depends on a position's index). So positional information is injected (added positional embeddings, or rotary/relative schemes) before the first block.

From here on, each transformer block maps a \(T \times d_{\text{model}}\) matrix to another of the same shape (internally, as we'll see, it routes through narrower matrices and back). That preserved width is the residual stream.

Step 1 — Project to queries, keys, values

From the input \(X \in \mathbb{R}^{T \times d_{\text{model}}}\), three learned matrices produce three new matrices:

with \(W_Q, W_K, W_V \in \mathbb{R}^{d_{\text{model}} \times d_{\text{head}}}\), so \(Q, K, V \in \mathbb{R}^{T \times d_{\text{head}}}\). Think of each position as emitting a query ("what am I looking for?"), advertising a key ("what do I offer?"), and holding a value ("what I'll pass on if attended to").

Step 2 — Score every pair

The relevance of position \(j\) to position \(i\) is the dot product of query \(i\) with key \(j\). All pairs at once:

Row \(i\) of \(S\) says how much position \(i\) wants to read from each position.

Step 3 — Scale by \(\sqrt{d_{\text{head}}}\)

Divide by \(\sqrt{d_{\text{head}}}\) before the softmax:

Why: if the components of \(q\) and \(k\) have variance ~1, their dot product over \(d_{\text{head}}\) terms has variance ~\(d_{\text{head}}\). Without rescaling, large \(d_{\text{head}}\) pushes the softmax into saturation (one weight ≈ 1, the rest ≈ 0), where gradients vanish. Dividing by \(\sqrt{d_{\text{head}}}\) keeps the scores at a sensible scale.

Step 4 — Causal mask + softmax → attention weights

For a language model predicting the next token, position \(i\) must not see the future. We set \(\tilde{S}_{ij} = -\infty\) for \(j > i\), so those weights will become 0. Then we apply a row-wise softmax: the function that turns a row of real scores \(z\) into a probability distribution, \(\operatorname{softmax}(z)_j = e^{z_j} / \sum_k e^{z_k}\), every entry positive and the whole row summing to 1:

Each row of \(A\) is now a probability distribution over the positions at or before \(i\) (the masked future positions, sent to \(-\infty\), get weight 0), the attention pattern. These weights are explicit numbers you can plot, which is why attention is a natural entry point for interpretability.

Step 5 — Read out a weighted sum of values

Position \(i\)'s output is the average of the value vectors, weighted by how much it attended to each. Putting the steps together gives the familiar formula:

Several heads in parallel, then recombine

A layer runs \(n_h\) heads independently and in parallel, each with its own \(W_Q, W_K, W_V\), so each can learn a different routing pattern (one head tracks subject–verb agreement, another copies rare tokens, and so on). Their outputs are concatenated and projected back to the model width by an output matrix \(W_O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}\):

Because \(d_{\text{head}} = d_{\text{model}}/n_h\), the \(n_h\) head outputs (each \(T \times d_{\text{head}}\)) concatenate to exactly \(T \times d_{\text{model}}\), and the total compute is similar to one big head, but the model gets several independent "read channels". GPT-2 small, the model you'll open in the lab, has \(d_{\text{model}} = 768\) and \(n_h = 12\) heads of dimension 64, across 12 layers.

Per-position computation and storage

After attention moves information between positions, a position-wise MLP does computation within each position, the same two-layer network applied to every row independently:

It up-projects to a wider hidden dimension (commonly \(4 d_{\text{model}}\)), applies a nonlinearity \(\sigma\) (GELU in GPT-2), then down-projects back. The MLP holds much of a transformer's stored "knowledge" and accounts for roughly two-thirds of a block's non-embedding parameters. It is also, so far, harder to interpret than attention.

Mixture of experts

That last figure, the MLP holding most of a block's parameters, is what the following design exploits, and it is now standard in the largest open-weight models.

If you want a model to know more, you need more parameters. If every parameter is used on every token, more parameters means proportionally more computation, for training and for every query afterwards. A mixture-of-experts layer breaks that link. In place of one MLP it holds \(E\) copies of it, the experts, plus a small network called the router. For each token the router scores the experts and sends that token through only the best \(k\) of them, commonly \(k = 2\). Parameters grow with \(E\); computation per token grows with \(k\).

Mixtral 8x7B is the clearest open example. Each layer has 8 experts, the router picks 2 per token, and the model holds 47 billion parameters while using about 13 billion on any given token. The arithmetic in the name does not work (8 × 7 is not 47), because only the feed-forward blocks are replicated. Attention, embeddings and the norms are shared across experts, so the sparsity sits in the part of the block that stores knowledge, and not in the part that moves information between positions.

Two consequences to carry forward. The router is learned along with everything else, and left alone it tends to collapse onto a handful of popular experts, so these models are trained with an extra term that pushes load to spread out. And the question "how big is this model" now has two different correct answers, one for capacity and one for cost. 2.3 shows what that does to the compute arithmetic.

One transformer block

A block wraps attention and MLP with LayerNorm (which rescales each position's vector to zero mean and unit variance, stabilising training) and adds each sub-layer's output back to its input. Modern transformers use the pre-norm arrangement:

Stack \(L\) such blocks (12 in GPT-2 small, ~100 in frontier models). Each block reads the current state, computes a correction, and adds it back.

Unembed, logits, loss

After the final block, a last LayerNorm is applied and the residual vector at each position is mapped to vocabulary scores by the unembedding matrix \(W_U \in \mathbb{R}^{d_{\text{model}} \times V}\):

Row \(i\) is the model's predicted distribution for the token after position \(i\). Training minimises the cross-entropy of \(p\) against the true next token at every position, the loss from Session 2.1. That single objective, next-token prediction, is the only thing the network is ever directly optimised for.