How the bill becomes a nudge on every one of four hundred billion numbers.
The forward pass is built: ids in, distribution out, bill paid. Every matrix along the way, W_E, W_Q, W_K, W_V, W_O, W_in, W_out, W_U, the norm gains, is full of numbers we have been calling "learned" without saying how. This chapter says how. It is the one place in the book where calculus is unavoidable, and the calculus needed is one idea (a slope) applied through one rule (the chain rule). Everything in Parts 2 and 3, from a trillion-token pretraining run to reinforcement learning on preferences, is this chapter's loop with different data and different bills.
The question this chapter answers: given the loss on one prediction, how does the model know which of its numbers to change, in which direction, and by how much?
A factory with a hundred thousand dials, all set at random, and a single quality gauge at the end of the line. Your job is to tune the dials so the gauge reads low. You cannot see inside the machine. You can only read the gauge.
The naive method: nudge one dial, read the gauge, note whether it improved, put the dial back, move to the next dial. A hundred thousand readings per round of adjustment. For a model with four hundred billion dials, a single round would take longer than the universe has existed.
The method that works: blame. When the gauge reads high, ask the last station on the line "how much of this is your fault, and would a little more or a little less of what you did have helped?". That station knows its own recipe, so it can answer, and it can pass the question upstream: "I was too high because my input was too high; input, how much of that was your fault?". The question travels backwards through every station, each one multiplying in its own local sensitivity, until every dial has a number: "turning me up by a hair changes the gauge by this much". Then turn every dial a hair in the direction that lowers the gauge. All hundred thousand at once, from one reading.
That backward pass of blame is backpropagation. The number each dial receives is its gradient. The hair's breadth is the learning rate. And the reason it costs about as much as one forward reading, rather than a hundred thousand, is the chain rule.
| In the picture | In the machine | The word we will use |
|---|---|---|
| A dial | One number in θ | parameter (weight) |
| The gauge reading | −ln p(target), summed over the tokens in a batch | loss L(θ) |
| "A hair's turn changes the gauge by this much" | The derivative of the loss with respect to that parameter | gradient component ∂L/∂w |
| All the dials' answers together | A vector with one entry per parameter, same shape as θ | the gradient ∇L |
| Passing the blame question upstream | Multiplying local derivatives from the loss back to each parameter | backpropagation (the chain rule, organised) |
| A station's own sensitivity | The derivative of one operation's output with respect to its input | local derivative |
| Turning every dial a hair downhill | θ ← θ − η ∇L | update; η is the learning rate |
| Doing the whole thing once | Forward, loss, backward, update | a training step |
Freeze the data for a moment: one batch of text, with its true next tokens. Then the loss is a function of θ alone. Feed the batch through the model with these parameters, get this loss; change a parameter, get a different loss. L(θ) is a landscape with as many dimensions as there are parameters, and training is walking downhill on it.
You cannot picture four hundred billion dimensions, so picture one. Fix every parameter except a single weight w and plot the loss as w varies. You get a curve. The current value of w is a point on it. The question "should w go up or down?" is the question "which way is downhill from here?", and that is answered by the slope of the curve at that point.
For a function L(w), the derivative at a point, written dL/dw, is the slope of the curve there: how much L changes per unit change of w, for tiny changes.
dL/dw ≈ ( L(w + ε) − L(w) ) / ε for a very small ε
Positive slope: increasing w increases the loss, so go down. Negative slope: go up. Zero: you are at a flat spot, possibly the bottom. The size of the slope says how steep the hill is.
Rules you need, and only these: the derivative of a·w with respect to w is a; of w² is 2w; of ln w is 1/w; of e^w is e^w. The derivative of a sum is the sum of derivatives. And the chain rule, which gets its own box below.
With many parameters, take the slope along each one separately, holding the others fixed (a partial derivative, ∂L/∂wᵢ), and collect them in a vector ∇L, the gradient. It has the same shape as θ and it points in the direction of steepest ascent. Minus the gradient is the direction of steepest descent.
Given the gradient, the update is one line:
θ ← θ − η · ∇L(θ)
Every parameter moves against its own slope, scaled by the learning rate η. A steep coordinate moves more; a flat one barely moves. Do this once and the loss on that batch goes down a little (if η is small enough). Do it a million times on a million batches and you have a trained model.
Play with the learning rate, because it is the single most important knob in all of training and the widget shows why. Too small: progress is glacial. Too large: each step overshoots the bottom, lands on the far slope, and the next step overshoots back; past a threshold the steps grow and the loss diverges. Chapter 13 spends a whole section on how η is scheduled over a frontier run. The intuition is entirely in this one-dimensional picture.
The update needs ∂L/∂w for every w. The loss is computed from the logits, which are computed from the last block's output, which is computed from the block before, and so on, through dozens of operations, back to w. The chain rule is how a slope passes through a chain of operations.
If L depends on z, and z depends on w, then
dL/dw = dL/dz · dz/dw
"How much the loss changes per unit of w" equals "how much the loss changes per unit of z" times "how much z changes per unit of w". Read it as a relay: the loss tells z how much it matters; z tells w how much it matters, by multiplying in its own local sensitivity. For a longer chain, multiply all the way along.
When a quantity feeds several downstream places, its total derivative is the sum of the contributions along each path. This is the only other rule, and it is why residual connections behave as they do (§5.4).
Backpropagation is the chain rule organised so that nothing is computed twice. Do the forward pass and keep every intermediate value. Then start at the loss with ∂L/∂L = 1 and walk backwards. At each operation, multiply the incoming derivative by that operation's local derivative (which you can compute because you kept its inputs) and pass the result to its inputs. Parameters are inputs too; when the walk reaches one, the number that arrives is its gradient. One forward pass, one backward pass of about the same cost, and every parameter has its slope.
The smallest network that has the whole structure: an input, a weight, a ReLU, another weight, a loss. Squared error stands in for −ln p; the mechanics are identical and the algebra is shorter.
x = 1.5 target y = 1.0 w₁ = 0.8 w₂ = −0.5 forward, keep everything a = w₁·x = 0.8 × 1.5 = 1.20 h = ReLU(a) = max(0, 1.20) = 1.20 z = w₂·h = −0.5 × 1.20 = −0.60 L = (z − y)² = (−0.60 − 1.0)² = 2.56 backward, from the loss to each weight ∂L/∂z = 2(z − y) = 2 × (−1.60) = −3.20 "z is too low; raising it helps" ∂L/∂w₂ = ∂L/∂z · h = −3.20 × 1.20 = −3.84 (z = w₂·h, so ∂z/∂w₂ = h) ∂L/∂h = ∂L/∂z · w₂ = −3.20 × (−0.5) = 1.60 (∂z/∂h = w₂; note the sign flip) ∂L/∂a = ∂L/∂h · [a > 0] = 1.60 × 1 = 1.60 ReLU was open, gradient passes ∂L/∂w₁ = ∂L/∂a · x = 1.60 × 1.5 = 2.40 (a = w₁·x, so ∂a/∂w₁ = x) update with η = 0.1 w₂ ← −0.5 − 0.1 × (−3.84) = −0.116 w₁ ← 0.8 − 0.1 × ( 2.40) = 0.560 forward again a = 0.84 h = 0.84 z = −0.097 L = 1.204 lower than 2.56 ✓
$ python code/ch05/tiny_backprop.py
forward : a=1.200 h=1.200 z=-0.600 loss=2.560 backward: dL/dz=-3.200 dL/dw2=-3.840 dL/dh=1.600 dL/da=1.600 dL/dw1=2.400 update : w1=0.560 w2=-0.116 forward : a=0.840 h=0.840 z=-0.097 loss=1.204 (lower)
Read the sign flip at ∂L/∂h. The loss wants z higher. But z = w₂·h with w₂ negative, so making h higher makes z lower. The chain rule tracks this automatically: multiplying by w₂ = −0.5 flips the sign. Nobody reasoned about it; the arithmetic did. That is the whole appeal.
Read also what ReLU did. Because a was positive, the gate was open and the gradient passed through unchanged. Had a been negative, ∂L/∂a would be zero, and w₁ would get no update from this example: a closed gate takes no blame. This is what "the MLP's detectors learn only from the inputs they fire on" means, mechanically.
The chain rule needs a local derivative for every operation in the forward pass. There are only a handful of operation types in a transformer, and each has a local derivative worth seeing once, because each explains a design choice made earlier in the book.
p − onehotThe loss is −ln p(target) and p = softmax(z). Work the chain rule through both (the algebra is in the further reading; the result is what matters) and the gradient on the logits is
∂L/∂z = p − onehot(target)
A vector of V numbers: each wrong token's logit gets a push down equal to its probability; the true token's logit gets a push up equal to 1 − p(true). If the model already put 0.99 on the right answer the gradient is nearly zero everywhere: nothing to learn. If it put 0.01 on the right answer and 0.9 on a wrong one, the wrong one gets a push of −0.9 and the right one +0.99. The size of the push is the size of the mistake. This is the mechanical content of Chapter 1's claim that loss is dominated by the tokens the model got most wrong.
$ python code/ch05/softmax_grad.py
p [0.644 0.237 0.087 0.032] onehot [0. 0. 1. 0.] dL/dz [ 0.644 0.237 -0.913 0.032] (sums to -0.0 ) numeric -0.913 vs analytic -0.913
The last line is the check every gradient implementation is tested with: nudge the input by a tiny ε, measure the change in loss, divide. The analytic formula and the numeric nudge agree to three decimals. Note also why the loss had to be −ln p and not accuracy: accuracy is flat almost everywhere, so its gradient is zero almost everywhere, and there would be nothing to backpropagate.
For y = x·W with x of shape [n], W of shape [n, m], and incoming gradient ∂L/∂y of shape [m]:
∂L/∂W = xᵀ · (∂L/∂y) shape [n, m]: entry (i, j) = xᵢ · ∂L/∂yⱼ ∂L/∂x = (∂L/∂y) · Wᵀ shape [n]: passes the blame back to the input in words: the gradient on weight Wᵢⱼ is (the input it multiplied) × (the blame on the output it fed)
Two things follow. A weight that multiplied a zero input gets zero gradient: it did nothing, so it takes no blame. And with a batch of T tokens stacked in X of shape [T, n], the weight gradient is Xᵀ · ∂L/∂Y, a sum over all T tokens' contributions, computed as one matmul of the same cost as the forward one. Every matrix in the model, from W_Q to W_out, gets its gradient this way. That is why the backward pass costs about twice the forward pass: one matmul for ∂L/∂W, one for ∂L/∂x, per forward matmul. Forward plus backward is therefore about 6N operations per token, against 2N for forward alone. Chapter 8 builds the cost of a training run from that 6N.
ReLU's local derivative is 1 where the input was positive and 0 where it was not. The gradient passes through open gates untouched and is blocked at closed ones. A detector in the MLP is updated only by the examples it fired on, which is what lets thousands of detectors specialise on different inputs instead of all learning the same thing. GELU and SiLU are the same with soft edges: a little gradient leaks through nearly-closed gates, which helps units that are stuck closed find their way back.
For y = x + f(x), the local derivative is 1 + f′(x). Two paths from y back to x: straight through the plus sign, with derivative exactly 1, and through the sub-layer, with derivative f′. By the sum rule they add. So whatever gradient arrives at the top of the residual stream reaches the bottom at least intact, no matter how many blocks are in between, because along the highway every local derivative is 1 and the product of ones is one. Without the highway, the gradient at layer 1 is the product of forty sub-layer derivatives, and a product of forty numbers that are each a bit less than 1 is nothing, while a product of forty numbers a bit more than 1 is everything.
This is the deferred reason from Chapter 4. Residual connections are not primarily about preserving information forward, though they do. They are about delivering gradient backward. Pre-norm keeps the highway free of norms for the same reason: a norm on the highway would insert a non-unit derivative on the one path that was supposed to be clean.
The attention output is Σⱼ wⱼ vⱼ. Its gradient with respect to each value vⱼ is the incoming gradient times wⱼ: values that were attended to strongly get most of the blame, values that were ignored get almost none. The gradient with respect to the weights then flows back through the softmax (the same p − something structure as above) to the scores, and from the scores to q and k and hence to W_Q and W_K. Read it as: if attending to "cat" from "it" lowered the loss, the update makes q_it and k_cat more aligned; if it raised the loss, less aligned. That is how the routing rules of Chapter 3 are learned, one nudge per example, and why a head that helped on many pronoun examples ends up as a "referent" head.
The gradient of the loss on one token is noisy: it says how to do better on that token. The gradient on a batch of B sequences of T tokens is the average of B·T such signals, and averaging cancels the noise while keeping what the tokens agree on. Bigger batches give cleaner gradients at proportionally higher cost per step; smaller batches give more steps for the same compute, each noisier. Frontier runs use batches of millions of tokens public (Llama 3 405B: up to 16M tokens per step), because at that scale the gradient is clean enough that each step is worth its cost, and because millions of tokens is what keeps thousands of GPUs busy at once (Chapter 11).
Plain gradient descent uses one η for every parameter. But different parameters live at wildly different scales of gradient: an embedding row for a rare token gets a gradient a thousand times smaller than a norm gain that every token touches. One η is too big for some and too small for others.
m ← β₁·m + (1 − β₁)·g running average of the gradient (momentum; β₁ ≈ 0.9) v ← β₂·v + (1 − β₂)·g² running average of the squared gradient (scale; β₂ ≈ 0.95–0.999) θ ← θ − η · m / (√v + ε) step: momentum direction, divided by typical size, per parameter weight decay: θ ← θ − η·λ·θ a gentle pull toward zero (AdamW)
Every parameter gets its own effective learning rate η/√v: parameters with small gradients take bigger steps, parameters with large ones take smaller steps, and the step size is roughly η in units of "typical gradient". m smooths the direction across steps. The cost: two extra numbers stored per parameter, which at frontier scale is measured in terabytes (§5.7). Every published frontier model is trained with AdamW or a close relative public.
Even with Adam, η is scheduled: small at first while the randomly initialised model is in a chaotic region of the landscape (warm-up), then at its peak for most of the run, then decayed toward zero so the final parameters settle into a minimum rather than bouncing around it. The 1-D widget above shows both failure modes: too small is slow, too big bounces. A schedule is the compromise over time. Chapter 13 shows real schedules from published runs.
Run the loop and watch what happens to each component, because this is where the abstract "learned" of Chapters 2 to 4 becomes concrete.
" cat" is followed by " sat", the gradient on W_E[" cat"] pushes it in whatever direction made " sat" more probable. " dog" gets pushed the same way on its own examples. After enough of both, the two rows point similarly: the geometry of Chapter 2 is the residue of a billion small pushes.W_Q, W_K happen to align "it" with recent nouns slightly better than chance gets a slightly larger gradient reinforcing that alignment on every pronoun example. Reinforcement compounds. Induction heads in particular are known to appear suddenly, partway through training, and to cause a visible drop in the loss curve when they do public.W_in drifts toward the average of the inputs it fires on, and its row of W_out toward whatever those inputs needed added. It becomes a detector for a pattern and a writer of that pattern's consequence.W_U's column for the true token toward the final residual vector and the columns for probable wrong tokens away from it, by p − onehot.None of these were designed. The only things designed were the shapes (Chapter 4) and the loss (Chapter 1). Everything the model can do is what a million steps of "push every dial a hair against its slope" produced from those two choices plus the data. This is why Part 2 spends a chapter on the data and a chapter on the run: given the architecture, they are the whole of the input.
| Quantity | Llama 3.1 405B | Evidence |
|---|---|---|
| Training tokens | 15.6 T | public |
| Total training compute | 3.8 × 10²⁵ FLOPs | public |
| Optimiser | AdamW, peak η = 8 × 10⁻⁵, cosine decay, 8,000 warm-up steps | public |
| Batch size | 4M → 8M → 16M tokens | public |
| Closed frontier models | Token counts, compute, and schedules not disclosed; the same loop is the consensus assumption | unknown / inferred |
Does the 6N rule reproduce the published compute?
operations per token, forward + backward ≈ 6 × N = 6 × 405 × 10⁹ ≈ 2.4 × 10¹² × training tokens × 15.6 × 10¹² ≈ 3.8 × 10²⁵ ✓ matches the report
How many steps? 15.6T tokens at 16M per step (ignoring the smaller early batches) is about a million steps. A million times: forward, bill, backward, nudge every one of 405 billion dials.
What has to be in memory to take one step? Per parameter: the weight itself, its gradient, and Adam's m and v. In the usual mixed-precision setup (Chapter 11), that is 2 bytes for the working weight, 2 for the gradient, and 12 for the 32-bit master weight and the two Adam moments.
bytes per parameter for training 2 + 2 + 4 + 4 + 4 = 16 bytes × 405 × 10⁹ parameters ≈ 6.5 TB + activations kept for the backward pass (batch-dependent) several TB more largest single GPU (2025) ≈ 0.19 TB
Serving needs the weights alone: 0.8 TB. Training needs eight times that before a single activation is stored. This is why a frontier training run occupies thousands of GPUs, not for speed alone but because the optimiser state does not fit on fewer. Chapter 11 is about how it is split.
Skip the chain rule; nudge each parameter and re-measure. Each gradient component costs one forward pass. For 405 billion parameters, one step costs 405 billion forward passes, and the run needs a million steps. Backpropagation gets all 405 billion components from one forward and one backward pass, at about three times the cost of the forward pass alone. The chain rule is a factor of 10¹¹ in efficiency, which is the difference between training being possible and not.
Use accuracy as the loss. Zero gradient everywhere except at the boundary where the argmax flips, where it is undefined. Backpropagation delivers zeros to every parameter. Nothing moves. Chapter 1's insistence on a smooth, proper scoring rule was for this moment.
Set the learning rate 10× too high. Steps overshoot every valley; the loss climbs instead of falling, then the logits blow up, then a softmax overflows and you get NaN. Frontier runs monitor for this ("loss spikes", Chapter 13) and roll back to a checkpoint when it happens.
Use one sequence per step instead of millions of tokens. The gradient is the slope for that one sequence's loss, which points somewhere different from the average slope. The walk becomes a random stagger; with Adam's normalisation the stagger has unit size in every direction regardless of how meaningless the direction is. Progress per step collapses, and no amount of extra steps recovers it, because each step still costs a full pass through the model.
Initialise every weight to zero. Every hidden unit computes the same thing, receives the same gradient, and updates identically, forever. Symmetry never breaks, and the model has effectively one unit per layer. Random initialisation is what gives the units different starting points to specialise from.
Say it back. Fix a batch of text; the loss is now a function of the parameters alone, a landscape. The gradient is the vector of slopes, one per parameter, and stepping against it lowers the loss. The slopes are computed by the chain rule: do the forward pass, keep every intermediate value, then walk backwards from the loss multiplying in each operation's local derivative, so that one backward pass delivers every parameter's gradient at about twice the forward cost. The local derivatives explain the architecture: softmax with −ln p gives p − onehot, a push proportional to the mistake; a matmul gives an outer product of input and blame, so each weight learns from what it multiplied; ReLU gates the gradient so detectors specialise; the residual connection has derivative exactly 1 along the highway, so deep stacks still deliver gradient to their first layers; attention passes blame to the values in proportion to the weights and thereby tunes who attends to whom. Batches of millions of tokens average out the noise; Adam gives each parameter its own step size; the learning rate is warmed up, held, and decayed. A million such steps, and the embeddings, heads, detectors, and unembedding are what they are. Beacon's run was 15.6 trillion tokens, 3.8 × 10²⁵ operations, and needs 16 bytes per parameter of optimiser state, 6.5 TB, which is why it ran on thousands of GPUs.
w₂ = +0.5 instead of −0.5. Compute the forward values, all five backward values, the update at η = 0.1, and the new loss. Then set w₁ = −0.8 (so ReLU is closed) and explain, in one sentence, why w₁ receives no update no matter what the target is.softmax_grad.py into a full check of the matmul rule: create a random x of shape [3], W of shape [3, 4], a target index, compute loss = −ln softmax(xW)[target], form ∂L/∂W = xᵀ(p − onehot), and verify two entries of it numerically by nudging W. Then take ten steps of gradient descent on W and print the loss each step.p − onehot follows from ∂/∂zⱼ(−ln pₜ) = pⱼ − [j = t], three lines of algebra with the quotient rule.