Part III · How LLMs work · Chapter 18

The transformer block

A transformer is a stack of identical blocks. Each block does two things to every token's vector: attention, where the token gathers information from other tokens, and a small feed-forward network, where it processes what it gathered. Residual connections and layer normalisation hold the stack together. This chapter computes each piece with real numbers.

The shape of the model

1. Tokens become vectors of width d, with positions added (chapter 17).
2. N blocks each update every vector: x ← x + Attention(LayerNorm(x)), then x ← x + MLP(LayerNorm(x)).
3. A final layer norm, then one matrix turns the last token's vector into a logit for every vocabulary token, and the softmax gives the next-token probabilities (chapter 16).

The "x +" in each line is a residual connection (chapter 13). Putting the layer norm before each part, as here, is the common modern arrangement ("pre-norm"); the 2017 original normalised after.

Attention with real numbers

Attention lets each token look at the others and pull in what is relevant. Every token's vector x is multiplied by three learned matrices to get three new vectors:

• a query q = x W_Q: what this token is looking for;
• a key k = x W_K: what this token offers to be found by;
• a value v = x W_V: what this token hands over if chosen.

Worked with three tokens, "fund pays it", and d = 2. The embeddings are x₁ = (1, 0), x₂ = (0, 1), x₃ = (1, 1). Take W_Q = [[1, 0], [1, 1]], W_K = the identity, and W_V = [[1, 0], [0, 2]]. Each vector is a row, multiplied on the right.

1. Project. Queries: (1, 0), (1, 1), (2, 1). Keys: (1, 0), (0, 1), (1, 1). Values: (1, 0), (0, 2), (1, 2).
2. Score token 3 ("it") against every key: q₃ · k₁ = 2 × 1 + 1 × 0 = 2, q₃ · k₂ = 1, q₃ · k₃ = 3.
3. Scale by √2 = 1.414: 1.414, 0.707, 2.121.
4. Softmax: e^1.414 = 4.113, e^0.707 = 2.028, e^2.121 = 8.342; the sum is 14.484; the weights are 0.284, 0.140, 0.576.
5. Mix the values: 0.284 × (1, 0) + 0.140 × (0, 2) + 0.576 × (1, 2) = (0.860, 1.432).

Token 3's new vector is mostly its own value and token 1's, because its query matched those keys best. Training adjusts the three matrices so that, for example, a pronoun's query matches the key of the noun it refers to.

Why divide by √d?

A dot product of two random vectors of length d has a spread that grows like √d. With d_k = 128, raw scores of 30 or 40 are common, and a softmax of such numbers puts almost all weight on one token, where the gradient is nearly zero (chapter 9). Dividing by √d_k keeps the scores around size 1, so the softmax stays soft enough to learn.

The causal mask

A model that writes text left to right must not peek at the tokens it is about to predict. The mask sets scores for later positions to −∞, whose exponential is 0. In the example, token 1 sees only itself (weight 1), token 2 sees tokens 1 and 2 (scores 1 and 1, so 0.5 each, output (0.5, 1.0)), and only token 3 sees all three. Models that only read, such as the embedding models used for search, drop the mask and let every token see the whole input.

Multi-head attention

One head computes one pattern of "who looks at whom". A token may need several at once: the subject of its verb, the previous token, the matching opening bracket. Multi-head attention runs h heads side by side, each with its own smaller W_Q, W_K and W_V of width d/h. Their outputs are joined end to end back to width d and multiplied by one more matrix, W_O.

Worked for d = 768 and 12 heads: each head works in 64 dimensions. The four matrices W_Q, W_K, W_V and W_O are each 768 × 768 in total across the heads, so attention costs 4 × 768² = 2,359,296 weights per block, plus 3,072 biases if the model uses them. Splitting into heads adds no weights; it changes what the same weights can express.

The feed-forward network

After attention, each token's vector goes through a two-layer MLP on its own: widen to 4d, apply a nonlinearity, narrow back to d. Attention moves information between tokens; the MLP transforms it within each token. Much of what a model "knows" is thought to be stored in these weights.

Worked with d = 2 and a hidden width of 4. Take the attention output (0.860, 1.432) and a first matrix whose four columns are (1, 0), (0, 1), (1, −1) and (−1, 1): the hidden pre-activations are 0.860, 1.432, −0.572 and 0.572. ReLU zeroes the negative one: (0.860, 1.432, 0, 0.572). A second 4 × 2 matrix brings it back to two numbers, which are added to the token's vector through the residual connection. Real models use smoother activations such as GELU, or gated variants like SwiGLU, but the shape is the same.

Weights per block for d = 768: two matrices of 768 × 3,072, so 2 × 2,359,296 = 4,718,592, plus 3,840 biases. That is twice the attention weights: in most transformers about two thirds of the block's parameters are in the MLP.

Layer norm in the block

Layer normalisation rescales each token's vector to mean 0 and standard deviation 1, then applies a learned scale and shift (worked in chapter 13). In the block it sits before attention and before the MLP, so each part receives well-scaled inputs however large the residual stream has grown. Many recent models use RMSNorm, which skips the mean and only divides by the root mean square.

Counting a whole model

Put it together for a model with d = 768, 12 blocks, an MLP of width 3,072, a vocabulary of 50,257 tokens and 1,024 learned positions:

About 124 million, the size of the smallest model in the GPT-2 family. A useful shortcut: each block has about 12d² weights, so a model has roughly 12 N d² plus its embedding table.

What attention costs

For a sequence of n tokens, the score matrix QK^T has n × n entries in every head of every block. Going from 1,000 to 2,000 tokens takes it from 1 million to 4 million entries: double the input, four times the attention work. Most other work in the block grows only in proportion to n. For short inputs the MLPs dominate the cost; for long ones attention takes over. Much engineering, such as memory-saving attention kernels and windowed or sparse attention, exists to tame this square.

The KV cache

When generating, each new token needs to attend to all earlier tokens, whose keys and values do not change. Servers keep them in memory, the KV cache, instead of recomputing them. Its size is 2 (keys and values) × blocks × tokens × d × bytes per number. For 32 blocks, d = 4,096, 8,192 tokens and 2-byte numbers: 2 × 32 × 8,192 × 4,096 × 2 ≈ 4.3 GB, for one conversation. That is why long contexts are expensive to serve, and why many models share keys and values across groups of heads to shrink the cache.

Reading a real model

A small open embedding model, BAAI/bge-small-en-v1.5, is a real transformer that reads like this chapter. Its public config.json gives 12 blocks, width 384, 12 attention heads (so 32 dimensions each), an MLP of width 1,536 (four times 384) with GELU, and 512 learned positions. Counting as in the table above gives about 33 million parameters. It is a reading-only model, so it uses no causal mask: every token in a question sees every other token.

The language models you call through an API are far larger and run on the provider's servers. Because attention cost grows with the square of the input, careful applications cap how much of a prompt the user can fill: a chat box might refuse a question longer than 2,000 characters, which keeps the user's share of every prompt, and of its attention cost, bounded.