Part 1 · Chapter 5

Learning

How the bill becomes a nudge on every one of four hundred billion numbers.

Where we are

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?

Picture this

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.

Map it
In the pictureIn the machineThe word we will use
A dialOne number in θparameter (weight)
The gauge reading−ln p(target), summed over the tokens in a batchloss L(θ)
"A hair's turn changes the gauge by this much"The derivative of the loss with respect to that parametergradient component ∂L/∂w
All the dials' answers togetherA vector with one entry per parameter, same shape as θthe gradient ∇L
Passing the blame question upstreamMultiplying local derivatives from the loss back to each parameterbackpropagation (the chain rule, organised)
A station's own sensitivityThe derivative of one operation's output with respect to its inputlocal derivative
Turning every dial a hair downhillθ ← θ − η ∇Lupdate; η is the learning rate
Doing the whole thing onceForward, loss, backward, updatea training step

5.1The loss is a function of the dials

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.

Math box · the derivative is a slope

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.

5.2One step downhill

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.

Static view of the widget. A wavy bowl-shaped loss curve. From w = −1.6 with η = 0.2, twenty steps reach the bottom near w = 1.4, each step shorter than the last as the slope flattens. At η = 1.0 the ball overshoots the bottom and oscillates; at η = 1.2 it leaves the plot.

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.

5.3The chain rule: blame travels backwards

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.

Math box · the chain rule

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.

By hand

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.

Static view of the widget. The graph x → a → h → z → loss with w₁ feeding a and w₂ feeding z. Forward values in blue: 1.5, 1.2, 1.2, −0.6, 2.56. Backward values in red: ∂L/∂z = −3.2, ∂L/∂w₂ = −3.84, ∂L/∂h = 1.6, ∂L/∂a = 1.6, ∂L/∂w₁ = 2.4. Applying the update and re-running shows the loss at 1.204.
embedW_E block 1W_Q…W_out …blocks 2…L unembedW_U loss−ln p forward: values, kept for later → ← backward: ∂L/∂(everything), each block multiplies in its local derivative every matrix the backward pass touches receives its gradient, same shape as the matrix
Where does the backward pass go? Back along exactly the path the forward pass took, in reverse, through every block to the embedding. Each matrix it passes gets a gradient of its own shape. The two passes cost about the same.

5.4Gradients through the pieces you already know

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.

Through softmax and the bill: p − onehot

The 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.

Static view of the widget. Logits [2, 1, 0, −1] give p = [.644, .237, .087, .032]. With "mat" (index 2) as the target, the gradient is [+.644, +.237, −.913, +.032]: cat's logit is pushed down hardest because it stole the most probability, mat's is pushed up by 0.913, and the four pushes sum to zero.
$ 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.

Through a matrix multiply: an outer product

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.

Through ReLU: a gate

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.

Through the residual connection: an identity

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.

Static view of the widget. Twenty-four layers with a per-layer gain of 0.7. Without residuals, the gradient reaching layer 1 is 0.7²⁴ ≈ 2 × 10⁻⁴; layer 1 effectively does not train. With residuals, the factor per layer is about 1.07 and the gradient at layer 1 is about 5: every layer receives a usable signal.

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.

Through attention: blame flows along the weights

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.

5.5Batches, learning rates, and Adam

Minibatches

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).

Per-parameter step sizes: Adam

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.

Static view of the widget. A valley where one direction is 40× steeper than the other. Plain descent at η = 0.124 zigzags across the steep direction, still bouncing at a ≈ ±0.4 after 60 steps, while crawling along the shallow one. Adam at the same η reaches the minimum, because it scales each coordinate's step by that coordinate's own gradient history.
Math box · Adam, in four lines
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.

The learning rate over time

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.

5.6What a million steps do to the pieces

Run the loop and watch what happens to each component, because this is where the abstract "learned" of Chapters 2 to 4 becomes concrete.

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.

5.7Beacon's numbers

QuantityLlama 3.1 405BEvidence
Training tokens15.6 Tpublic
Total training compute3.8 × 10²⁵ FLOPspublic
OptimiserAdamW, peak η = 8 × 10⁻⁵, cosine decay, 8,000 warm-up stepspublic
Batch size4M → 8M → 16M tokenspublic
Closed frontier modelsToken counts, compute, and schedules not disclosed; the same loop is the consensus assumptionunknown / inferred
Back of the envelope

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.

Break it

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.

Rebuild the model

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.

forwardkeep values loss−ln p, batch mean backwardchain rule, ∇L for all θ Adam updateθ −= η·m/√v × 1,000,000 stepsnew batch each time cost per token ≈ 6N operations · memory per parameter ≈ 16 bytes
What is the whole chapter in one line? Forward, bill, backward, nudge, repeat. Everything a frontier model knows is the accumulated effect of that loop on the shapes of Chapter 4.
Exercises
  1. By hand. Redo the two-weight example with 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.
  2. Calculation. A lab trains a 70B model on 15T tokens. Using the 6N rule, how many operations is the run? If a GPU sustains 4 × 10¹⁴ useful operations per second, how many GPU-seconds, and how many days on 8,000 GPUs? How many bytes of optimiser state at 16 bytes per parameter, and how many 80 GB GPUs does that alone need?
  3. Code. Extend 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.
Further reading