Many heads, a residual stream, normalisation, positions, and the full pass from token to logits.
You now hold every component: embeddings and matrices (Chapter 2), the MLP (Chapter 2), one attention head (Chapter 3), and the softmax-and-loss at the end (Chapter 1). This chapter assembles them into the repeating unit of every frontier model, the block, stacks the blocks, and follows one token all the way through to a distribution. When you finish, the box on the map labelled f(context; θ) is fully open. Chapter 5 then asks how its numbers were found.
The question this chapter answers: how are attention and the MLP arranged so that stacking dozens of them produces a model that can be trained, and what does the whole forward pass look like?
A document passes through a chain of specialists. Each specialist reads it, then staples a note to it, and passes it on. Nobody rewrites the document. Nobody removes an earlier note. By the time it reaches the end, the document is the original page plus forty stapled notes, and the last reader can consult any of them, including the original.
Some specialists are gossips: they read other documents in the same batch and staple "document 3 mentions the same client" onto this one. Others are analysts: they read only this document and its notes and staple "this is an invoice, probably overdue". Before each specialist reads, an assistant flattens the pile so it is a readable thickness, because forty notes of wildly different sizes would be unmanageable.
The gossips are attention. The analysts are MLPs. The stapling is the residual connection. The flattening is normalisation. The pile itself, original plus notes, is the residual stream: the token's vector as it travels through the model.
| In the picture | In the machine | The word we will use |
|---|---|---|
| The document plus all notes stapled so far | The token's vector at a given depth, shape [C] | residual stream |
| Stapling a note instead of rewriting | x ← x + f(x): the sub-layer's output is added, not substituted | residual connection |
| A gossip who reads other documents | Multi-head attention: several heads reading other positions at once | attention sub-layer |
| An analyst who reads only this document | The MLP from Chapter 2, applied per token | MLP sub-layer |
| Flattening the pile before reading | Rescale the vector to a standard size before each sub-layer | normalisation (RMSNorm) |
| One gossip plus one analyst | Attention sub-layer then MLP sub-layer, each with its residual | block (also "layer") |
| The specialists need to know page order | Attention is blind to order; positions are injected by rotating q and k | positional encoding (RoPE) |
| The last reader scoring every possible verdict | Final norm, then W_U to V logits, softmax | unembedding |
Chapter 3 built one head of width d = 128 from a token of width C = 4096. One head can implement one kind of lookup: previous token, or subject-of-verb, or "find my earlier occurrence". A token needs several of those at the same time. So the block runs H heads in parallel, each with its own W_Q, W_K, W_V, each producing a [T, d] output. The outputs are concatenated side by side into [T, H·d], and H·d is chosen to equal C: 32 heads × 128 = 4096 for Llama 3 8B public. One more matrix, W_O of shape [C, C], mixes the concatenated heads into the vector that gets added to the stream.
┌── head 1: Q₁K₁V₁ → out₁ [T, d] ──┐
X [T, C] ── norm ───┼── head 2: Q₂K₂V₂ → out₂ [T, d] ──┼── concat [T, H·d = C] ── · W_O [C, C] ──► [T, C]
│ … │
└── head H: Q_H K_H V_H → out_H ────┘
Why not one wide head of width C? Because one head has one softmax, so it produces one weighting over positions. It can attend mostly to the previous token or mostly to the subject, and a blend of the two is a compromise on both. Splitting the width into H independent heads gives H independent weightings for the same cost in matrix parameters: H matrices of [C, d] hold exactly as many numbers as one matrix of [C, H·d]. Multi-head attention is free in parameters and expensive only in the T² score matrices, of which there are now H per layer.
The projection W_O matters more than it looks. Each head writes into its own 128-dimensional slot of the concatenation; W_O lets the block route any head's finding into any direction of the residual stream, and lets heads' findings interact. Without it, head 3 could only ever write into coordinates 256 to 383.
Here is the design decision that makes deep transformers trainable, and it is one plus sign. Each sub-layer does not replace the token's vector with its output. It adds its output to the vector:
x ← x + Attention(norm(x)) then x ← x + MLP(norm(x))
The vector x that flows down the middle, receiving additions from every sub-layer, is the residual stream. Nothing in the stack can erase what an earlier layer wrote; it can only add something that cancels it. So the stream is an accumulating record: the embedding, plus what attention in layer 1 fetched, plus what the MLP in layer 1 concluded, and so on. By layer 30, the vector at position "it" contains the token identity, its position, what it refers to, its syntactic role, and a growing guess about the next token, all as a sum of arrows in the sense of Chapter 2.
Two consequences, one now and one deferred. Now: because every sub-layer reads from and writes to the same stream, the stream is a shared workspace. A head in layer 8 can read what an MLP in layer 3 wrote, provided it wrote it in a direction the head's W_K can detect. Interpretability research treats the stream as exactly this, a bus that components communicate over public. Deferred: without residual connections, networks deeper than a dozen layers were extremely hard to train, and the reason is gradient flow. Chapter 5 makes that precise, and it is the main reason the plus sign is there.
The stream grows by addition, forty or a hundred times. Left alone, its size would drift: some tokens' vectors would become ten times larger than others depending on how many sub-layers had something to say. The matrices downstream would then see inputs of wildly varying scale, and the softmax inside attention, which is sensitive to score magnitude (Chapter 3), would sharpen or flatten for no semantic reason.
The fix is to rescale the vector to a standard size just before each sub-layer reads it. Modern models use RMSNorm: divide the vector by its root-mean-square, then multiply each coordinate by a learned gain.
rms(x) = √( (x₁² + x₂² + … + x_C²) / C ) a single number: the typical size of a coordinate x̂ = x / rms(x) same direction, coordinates now have rms 1 output = x̂ ⊙ g g: a learned vector of C gains, one per coordinate
Only the length changes; the direction, which is where meaning lives (Chapter 2), is untouched. The gains g let training re-amplify particular coordinates if the model finds that useful. Older architectures used LayerNorm, which also subtracts the mean; RMSNorm drops that step because it turned out not to matter and costs time public (Llama 3 uses RMSNorm).
Where the norm sits matters. The original 2017 transformer normalised after the addition ("post-norm"). Every frontier model since about 2020 normalises the copy that goes into the sub-layer and leaves the highway itself un-normalised ("pre-norm") public. The difference: with pre-norm the residual stream is a pure sum from embedding to output, never rescaled in between, and that keeps the backward path of Chapter 5 clean. With post-norm each layer's norm sits on the highway and distorts it.
A fact hiding in Chapter 3: nothing in attention knows where a token is. The score between "cat" and "it" is q_it · k_cat, which depends on the two vectors and not on whether "cat" is 2 or 200 tokens back. Shuffle the keys and every query gets exactly the same weights, reordered. The mask blocks the future, but among the past, attention is a bag. Yet "the dog bit the man" and "the man bit the dog" must get different next-token distributions. Position has to be put in.
Early transformers added a position vector to each token's embedding public. Current frontier models do something cleverer, called RoPE (rotary position embedding): they rotate the query and key vectors by an angle that grows with position, inside each attention head, just before the dot product.
Take two coordinates of a vector, (a, b), as a point in the plane. Rotating it by angle φ gives
(a, b) → (a·cos φ − b·sin φ, a·sin φ + b·cos φ)
This is a matrix, so it is linear, and it preserves length. The key fact: rotate two vectors by the same angle and their dot product does not change (the angle between them is what the dot product measures, and both moved together). Rotate them by different angles φ₁ and φ₂, and the dot product changes as if only one had been rotated by φ₁ − φ₂.
RoPE splits the d coordinates of q and k into d/2 pairs, and rotates pair i of the token at position m by m·θᵢ. Each pair has its own base angle θᵢ, ranging from fast-turning (sensitive to nearby offsets) to very slow (sensitive to long-range offsets).
Now the score between a query at position m and a key at position n depends on the two content vectors and on m − n, the distance between them, and on nothing else. Not on m. Not on n. The head can learn "attend two positions back" as a specific rotation to match, and that rule works identically at position 5 and position 50,000. That is why RoPE is what every current open frontier model uses public and why it generalises to contexts longer than it was trained on better than the alternatives, with some adjustment (Chapter 12).
Assemble it. Embed the tokens, run L blocks, normalise, unembed, softmax. Every frontier model that has been described publicly is this, with the block details varying at the margins public.
ids [T] │ W_E lookup ▼ x [T, C] ─────────────────────────────┐ │ │ │ ┌─ block 1 ─────────────────┐ │ ├─►│ x += Attn(RMSNorm(x)) │ │ the residual stream: ├─►│ x += MLP(RMSNorm(x)) │ │ the same [T, C] tensor │ └───────────────────────────┘ │ all the way down │ ┌─ block 2 ─────────────────┐ │ ├─►│ … │ │ │ ⋮ (L blocks) │ ▼ │ RMSNorm ◄─────────────────────────────┘ │ ▼ · W_U [C, V] logits [T, V] │ softmax per row ▼ p [T, V] T distributions, one per position, each over the whole vocabulary
Read the last line again. The forward pass produces a next-token distribution at every position, all at once, from one pass. Position 3's distribution used only positions 1 to 3, thanks to the mask. During training every one of those T distributions is billed against its true next token; during generation only the last row is used. Chapter 7 is about that asymmetry.
Notice what the stack does not contain. No recurrence: block 5 does not run again after block 6. No loop: a token's vector passes through each block exactly once. The depth L is a fixed number of refinement steps, and the only way to "think longer" about a token is to generate more tokens, each of which gets its own pass. That constraint is the starting point for the reasoning models of Chapter 16.
You can now count every number in θ from the shapes, and the count reproduces published model sizes. Per block:
attention W_Q [C, C] + W_O [C, C] + W_K [C, C·r] + W_V [C, C·r] r = kv heads / q heads
= 2C² + 2C²·r (r = 1 without GQA)
MLP W_in [C, H] + W_out [H, C] (+ W_gate [C, H] if gated) = 2CH or 3CH
norms 2 gain vectors of C (negligible)
whole model L · (attention + MLP) + W_E [V, C] + W_U [C, V]
Llama 3 8B: C = 4096, L = 32, H = 14336, V = 128256, 32 query heads and 8 KV heads so r = 0.25, gated MLP public.
attention per block 2·4096² + 2·4096²·0.25 = 33.6M + 8.4M = 41.9 M MLP per block 3 · 4096 · 14336 = 176.2 M one block = 218.1 M × 32 blocks = 6.98 B embedding + unembedding 2 · 128256 · 4096 = 1.05 B total ≈ 8.03 B
$ python code/ch04/param_count.pyLlama 3 8B total 8.03 B per-layer attn 41.9 M mlp 176.2 M emb 1.05 B Llama 3 70B total 70.55 B per-layer attn 151.0 M mlp 704.6 M emb 2.10 B Llama 3.1 405B total 405.85 B per-layer attn 570.4 M mlp 2617.2 M emb 4.20 B
Three shapes and a formula give 8.03 billion, and the model is called 8B. The same formula gives 70.55B and 405.85B for the larger two. Two things to notice: the MLPs hold about four times the parameters of attention in every size, and the embedding tables, which dominate tiny models, are 1% of the 405B. When someone says "a 400-billion-parameter model", this is what they are counting, and now you can count it yourself.
| Quantity | Llama 3 8B | Llama 3.1 405B | Evidence |
|---|---|---|---|
Blocks L | 32 | 126 | public |
Width C | 4,096 | 16,384 | public |
| Query heads / KV heads | 32 / 8 | 128 / 8 | public |
| Normalisation, positions, activation | pre-norm RMSNorm, RoPE, SwiGLU | public | |
| Parameters | 8.03 B | 405.85 B | computed above; matches release |
| Closed frontier models | Same block structure is the consensus assumption; depths, widths, and any departures are not disclosed | inferred / unknown | |
What does one token cost, and where do the weights live? Chapter 2 found that a matmul costs about two operations per matrix entry per token. Almost all of θ is matrix entries, so a forward pass costs about 2N operations per token, where N is the parameter count. Attention's T² terms come on top and dominate only at long contexts.
Llama 3.1 405B operations per token, forward 2 × 405 × 10⁹ ≈ 8 × 10¹¹ (0.8 TFLOP) weights in memory at 2 bytes 405 × 10⁹ × 2 ≈ 810 GB largest single GPU memory (2025) ≈ 192 GB ⇒ the weights alone need at least 5 GPUs before a single token can be processed ⇒ a 1,000-token reply is ≈ 10¹⁵ operations of matrix arithmetic, about one second of one GPU's peak
The first arrow is why Part 2 exists: a frontier model does not fit on one device, for training or for serving, and everything about how it is built follows from that. The second arrow is why a reply costs what it costs.
Remove the residual connections. Each sub-layer now replaces the vector. After one block the embedding is gone; after two, whatever block 1 found is gone. More importantly, Chapter 5's gradients must pass through every sub-layer's matrices in sequence, and the product of forty matrices is either vanishingly small or explosively large. Networks this deep without residuals did not train at all before 2015, and the residual connection is the reason they do now public.
Use post-norm instead of pre-norm. Trainable, but the norm now sits on the highway, rescaling the accumulated stream at every layer. Deep post-norm transformers need careful warm-up schedules and are prone to diverging early in training; pre-norm removed that fragility, which is why the switch was universal public.
Remove positional encoding. Attention sees a bag of past tokens. "Dog bites man" and "man bites dog" produce identical distributions at the final position. The MLP cannot help: it never sees more than one token. The model can still learn a great deal (which words co-occur) but nothing about order, and word order is most of grammar. (A subtle exception: with a causal mask, the model can in principle infer something about position from how many tokens are attendable, and models trained without explicit positions do work somewhat public. They are worse.)
Use one head of width C per layer. One weighting over positions per layer. The 32 simultaneous lookups become 32 sequential ones spread across layers, so a fact that needed "previous token and subject at the same time" now needs two layers, and the model is effectively shallower. Published ablations show quality dropping steadily as heads are removed at fixed width public.
Make it very deep and very narrow, or very wide and shallow. Same parameter count either way. The narrow model's residual stream has too few directions to hold the accumulating notes. The shallow model has too few refinement steps to compose them. The published sweet spot has depth growing slowly with width: 32 layers at 4k wide, 126 at 16k wide public. Chapter 8 returns to how such ratios are chosen.
Say it back. A block is two sub-layers, each of which reads a normalised copy of the token's vector, computes an update, and adds it back onto the vector. The attention sub-layer runs H heads of width C/H in parallel, each attending with its own softmax to a different aspect of the past, concatenates their outputs, and mixes them with W_O. The MLP sub-layer is Chapter 2's expand-bend-compress, per token. The vector that flows through, growing by addition, is the residual stream, and nothing is ever erased from it. RMSNorm resets its size before each read so that scale cannot drift; pre-norm keeps the highway itself untouched. Attention is blind to order, so RoPE rotates queries and keys by position-dependent angles, making every score a function of the distance between tokens. Stack L blocks, normalise, multiply by W_U to get V logits at every position, softmax. Count the matrices and you get the model's size to the second decimal; multiply by two and you have the operations per token. Beacon is 126 of these blocks at width 16k, 405 billion numbers, 810 GB of weights that do not fit on a GPU.
C = 4, one head of width 4, and this residual stream after the embedding: x = [1, 0, 0, 0]. The attention sub-layer outputs [0, 0.5, 0, 0] and the MLP sub-layer outputs [−0.2, 0, 0.8, 0]. Write the stream after each addition. Then compute RMSNorm of the final vector (ignore gains). What would the final vector be without residuals? Which of the two can still tell you which token it started as?V = 128k, gated MLP with hidden 3.5·C, r = 0.125, and depth-to-width in the Llama 3 range. Give C, L, and the head count. Then compute the weights' size in gigabytes at 2 bytes per parameter and at 1 byte, and the forward operations per token.block(X, params) that implements one pre-norm block with a single head and a ReLU MLP, using the attention code from code/ch03/one_head_by_hand.py and the MLP from code/ch02/mlp_by_hand.py. Run it twice in sequence on a 4 × 4 input and confirm the output shape is unchanged. Then delete the two additions and observe what the second block's input has lost.