Backprop Primer

Here is the whole network this primer runs on: one input x = 2, one hidden ReLU neuron, one output, a squared error against y = 5, and four parameters starting at w₁ = 1.5, b₁ = 0.5, w₂ = 0.8, b₂ = 0.1. The prediction is 2.90, so the loss is 4.41. Move any one parameter and that number traces a curve — pick one:

Four curves, all through the same point. The w₁ curve has a corner at −0.25, where the hidden unit's ReLU switches off and the loss stops responding to w₁ at all; b₂'s is a clean parabola. What backprop returns is one number per parameter, and each one is a property of that parameter's own curve.

The property is the steepness where we are standing. Take w₂'s curve, lay a tangent on it, and drag the point along to watch the slope turn:

The tangent is the gradient — there is nothing else to work out. At the opening weight it tilts down at −14.70, which is ∂L/∂w₂ exactly. Slide right to w₂ = 1.40 and it goes flat: zero slope, because the loss is at its minimum. A flat tangent is what "no gradient" looks like.

Where does −14.70 come from? From the definition of a derivative: the loss at two nearby points, over the distance between them. Shrink the gap and watch the secant swing onto the tangent:

Notice the sign at the opening gap of 1.60: the secant reads +4.90, the opposite of the truth, because a gap that wide steps clean over the minimum. Close it to ε = 0.05 and it reads −14.09. That is what torch.autograd.gradcheck does, one parameter at a time — a test, never a training loop, because it costs a whole forward pass per parameter.

The slope exists to say which way to move. Because it points uphill we step against it, by −η · ∂L/∂w₂. The learning rate starts at zero, so the step has no length yet; raise it and watch the weight slide down the curve:

Watch what happens past η = 0.041: the step overshoots the minimum and starts climbing the far wall. At η = 0.082 it lands on exactly the loss it started from, and past that every step makes things worse. That threshold is 2/L″, which for this parabola is 1/a² = 0.082 — the gradient tells you the direction, never the distance.

One thing more before we chain slopes together. Backprop needs the slope of every node, not only of the loss, and for a linear node that slope is free. Move the output weight and watch the node's line pivot around its input:

Because ŷ = w₂·a + b₂ is a straight line, its slope is the same at every input, and it is w₂ — a number the forward passalready had in its hand. Every operation a framework ships has a local derivative of that kind: a short formula in values the forward pass computed anyway. Backprop is cheap because those formulas are cheap.

A network is a stack of small functions, each feeding the next. If a nudge at the input moves the middle by two, and a nudge at the middle moves the loss by eight tenths, then a nudge at the input moves the loss by 0.8 × 2 = 1.6. Nothing deeper is going on.

The picture makes the multiplication literal. The same nudge is drawn three times — once at the input, once after the first slope, once after the second — and each rail is as long as the one above it times that link's slope:

Notice what happens when either slider passes through zero: the bottom rail vanishes entirely, however large the other link is. One dead link kills the whole path, and that is §05's problem in miniature. Push both above one and the bottom rail is longer than the top — a gradient can grow on the way back as easily as it can shrink.

Our network is four such links: x → z → a → ŷ → L. Before anything can flow back, the forward pass has to run and leave a value on every wire. Press play, or scrub the track, and watch the four values appear in order:

Every one of those numbers is needed again. The activation a = 3.50 becomes the local derivative of ŷ with respect to w₂; the input x = 2 becomes the local derivative of z with respect to w₁. That is why the forward pass keeps its intermediates instead of throwing them away — a decision §03 puts a price on.

Now the same graph, in reverse. Each backward step multiplies what came before it by one local derivative and hands the result left. Scrub the track and read the gradient appear under each wire in turn:

Watch the multipliers on the arrows: −4.20, then × 0.80, then × 1, then × 1.50. The gradient at a is −3.36 and it stays −3.36 across the ReLU, because z = 3.5 is positive and ReLU's slope there is exactly 1. Backprop never re-derives anything; it walks a list, multiplying.

The reuse is the reason it is affordable. Every parameter hangs off one of those wires, and the wire's gradient is computed once for all of them. Step through the four parameters and watch the shared part of the path stay put while only one multiply is added:

Computed independently, the four gradients would cost 2 + 2 + 4 + 4 = 12 multiplications. Computed together they cost 7: three wire gradients plus one per parameter. Scale that to a real model and the ratio is the whole story — the deepest parameter in a 96-layer network would re-walk 96 layers on its own, and instead pays for a single multiply on a wire the layer above already filled in.

A product can be evaluated left to right or right to left. Start at the input and you propagate one input's influence forward through everything; start at the loss and you propagate the loss's sensitivity backward through everything. Both are correct. Only one is cheap, and which one depends on the shape.

Our shape is millions of parameters in and one scalar out. Raise the parameter count and watch the forward-mode sweeps stack up while the reverse-mode sweep stays a single pass:

Forward mode needs one sweep per input; reverse mode needs one per output. With twelve parameters that is twelve sweeps against one, and a 175-billion-parameter model would need 175 billion. This is the entire reason training is possible: the loss is a single number, so the cheap end to start from is the far one.

Reverse mode is not free, though — it has to look at every wire after the forward pass has left it, so the forward pass cannot throw its intermediates away. Drag the depth down and back up to see the pile of kept tensors for GPT-3's shape:

Ninety-six layers of GPT-3 at one sequence of 2,048 tokens hold 275 GB of activations — on hardware whose cards carry 80 GB. That figure is Korthikanti et al.'s, and their formula splits it: 82.1 GB is linear in the sequence length and 193 GB is the attention matrices, which are quadratic in it.

That quadratic term is the one that hurts. Both axes below are logarithmic, so a straight line is a power law and its steepness is the power. Step the context length and watch the total bend away from the linear part:

At 512 tokens the two curves are within a factor of two; at 16,384 the total is 13,027 GB against 657 GB of linear term. Long context is a memory problem before it is a compute problem, and that is what FlashAttention exists to fix — it never materialises the attention matrix that this term is measuring.

The other way out is to keep almost nothing and pay to rebuild it. Checkpoint the layer inputs only, and the backward pass re-runs each segment's forward before it needs the activations. Switch between storing and recomputing:

Because the backward pass is roughly twice the forward pass's arithmetic, one extra forward is about a third more work per step — and it buys a drop from 275 GB to 7.7 GB. That trade is why gradient checkpointing is a one-line flag in every training framework, and why turning it on is the first thing anyone does when a model will not fit.

A linear layer contributes its weight, which never changes with the input. An activation is different: its slope depends on where the forward pass happened to land, so the same network hands back different multipliers on different examples.

Four activations, one axis. Pick a function and drag the tangent along the curve; the readout is the number backprop will multiply by:

Notice how flat sigmoid is everywhere except the middle. Its steepest point is the origin, and even there the slope is only 0.250. tanh peaks at 1.000, ReLU is exactly 1 on its live half and exactly 0 on the other, and GELU tops out a little above one, at 1.129. Those four numbers decide what a hundred layers do.

Sigmoid's ceiling is worth drawing on its own. Below the curve is its own slope, plotted against the same pre-activation axis, with the ceiling ruled across it. Move the cursor and read the same number in two places:

The slope curve is σ(z)·(1 − σ(z)), a product of two numbers in [0, 1] that sum to one — so it is largest when they are equal, at z = 0, where it is exactly one quarter. There is no input anywhere on the axis that gets a sigmoid to hand back more than 0.25. That is a bound, not a tendency, and §05 multiplies it by itself.

Push the same unit into its tail and the bound stops mattering, because the real number is far below it. Slide the point out and watch the slope triangle flatten to nothing:

At z = 6 the slope is 0.0025 — a hundredth of the maximum — and at z = 8 it is 0.0003. The unit is not broken: it still outputs a confident 0.9997. It has simply stopped being able to learn, because everything downstream of it gets multiplied by three ten-thousandths. This is saturation, and it is a silent failure — the loss just stops moving.

ReLU has no tail to saturate in, but it has a worse trick. Its slope on the dead side is exactly zero, so a unit whose entire batch lands left of the kink receives nothing at all. Drag the bias down and watch the batch cross over:

Six of the eight inputs start alive. Drag past bias = −2.50 and none are: every example gets slope 0, so ∂L/∂w is exactly zero, the weight does not move, the bias does not move, and the unit is dead forever. Nothing throws. Hold the batch there and change what the unit computes:

Under ReLU the eight slopes sum to exactly 0.000. Leaky ReLU's constant 0.01 on the negative side makes it 0.080, and GELU — whose slope there is small, and briefly negative — makes it 0.152. None of those are large. All of them are non-zero, which is the only property that matters: a unit that still receives something can climb back out.

That is the whole difficulty of depth, stated once. Fifty layers means fifty factors multiplied together, and a product of many numbers does not behave like a sum: it does not drift, it compounds. The two failure modes are the two things a product can do.

The bars below are the gradient's size at each layer, on a logarithmic axis — each line up is a factor of ten, so a straight run of bars is exponential growth or decay. Every link starts at exactly 1.00, so nothing shrinks or grows; nudge the factor either way:

Watch how little it takes. A factor of 0.90 — a ten percent loss per layer, which sounds harmless — reaches 5.2e−3 after fifty layers. At 0.40 the gradient arriving at layer 50 is 1.3e−20, and at 1.60 it is 1.6e+10. Exactly one value of the factor leaves the product alone, and no network sits on it by accident.

For sigmoid the factor is not a choice. §04 showed its slope can never exceed a quarter, so a stack of sigmoids has a guaranteed ceiling on how much gradient can survive. Add layers and watch the run fall through the two lines where half precision gives up:

Both of those lines are exact powers of two, which is why they land on whole layer counts. 0.25⁷ = 2⁻¹⁴ is fp16's smallest normal number, so seven sigmoid layers is where a half-precision gradient starts losing bits. 0.25¹² = 2⁻²⁴ is the last number fp16 has, so at twelve layers the gradient has run out of format, and at thirteen it is exactly zero. That is the best case, with every unit sitting at its steepest point.

There is a standard dodge for exactly this, and it is one multiplication. Scale the loss up before calling backward and every gradient in the graph comes back scaled by the same amount. Lift the run off the floor:

At a loss scale of 16,384 the twelve-layer run lands at 9.8e−4, well inside fp16's normal range, and the optimizer divides by 16,384 again before the update — so the step is unchanged and only the representation moved. Note the second rule: the top bar has to stay under fp16's own maximum, which is whyGradScaler hunts for the scale by doubling until something overflows and then backing off.

Upwards the wall is harder and closer than people expect. Here the chain is 128 layers deep and every link multiplies by more than one, with float32's largest value ruled across the top. Raise the factor until the bars hit it:

At a factor of 2.00 — one doubling per layer — the product passes 3.4e+38 at layer 128 and every bar past it reads inf. Oneinf in the gradient makes one inf in the weight update, and inf − inf is nan: two iterations later every parameter in the model is nan and the loss stops printing a number. Explosion is loud, which is the one mercy in it.

One honest correction before we fix any of this. A layer's factor is not a scalar — it is a Jacobian, and it stretches different directions differently. Turn the direction the gradient arrives from and watch what comes out ride the ellipse:

The circle is every direction a gradient could arrive from; the ellipse is where the layer sends them. The long semi-axis is σ_max and the short one is σ_min, here 0.63 of it. So a factor of 1.00 means the gradient is multiplied by something between 0.63 and 1.00, depending on where it arrived from — and it is σ_max that has to stay near one. Spectral normalisation constrains exactly that.

The obvious first guess is to set every weight to zero. It is symmetric, it is unbiased, it is one line of code, and it kills the network before the first step — for a reason that is a property of backprop itself rather than of any particular loss.

Four hidden neurons, drawn as their incoming weight vectors, with the update each one receives hanging off its tip. They all start on top of each other. Open the spread and watch four appear where there was one:

Identical weights see the same input, so they compute the same output, so they receive the same gradient, so after the step they are still identical. Backprop preserves the symmetry it is given; it has no mechanism for breaking one. A hundred-neuron layer initialised this way is one neuron, a hundred times, forever — and it does not error, it just plateaus at the accuracy a one-neuron layer can reach.

So the weights start random. But the scale of that randomness is not free either, because §05's product applies to the forward activations as much as to the backward gradients. Here is the variance layer by layer for a twelve-layer ReLU stack, with the initialisation scaled off the He recipe. Move the gain:

At a gain of exactly 1.00 the variance is flat: every layer hands on what it was given. At 0.70 it reaches 1.9e−4 by layer 12, and at 1.40 it reaches 3.2e+3. A thirty percent error in one hyperparameter, compounded twelve times. This is why initialisation has named recipes instead of a default.

The recipes differ by a factor of two, and the two is exactly the half ReLU throws away. Below, sixty-four pre-activations arriving at one layer, with the half that ReLU zeroes on the left. Step the layer index, then switch recipes:

Because ReLU deletes half the distribution, the variance leaving a layer is half the variance arriving at it. Xavier's √(1/n) therefore halves the signal every layer — twelve layers in, the spread is 0.022 against the 1.000 it started from. He's √(2/n) puts the two back and the spread never moves. At GPT-2's width of 768 that is a standard deviation of 0.051 against Xavier's 0.036 — the whole difference between a trainable deep ReLU network and one that fades out.

A residual block computes x + f(x) instead of f(x), so its local derivative is 1 + f′(x) rather than f′(x). The 1 is a path the gradient can take that multiplies it by nothing at all.

Twenty-four layers, drawn twice: a plain chain and a residual one, with the identity path ruled across at exactly one. Move the branch's own slope:

At the opening slope of 0.05 — a weak branch, which is what a well-initialised block looks like — the plain chain reaches 6.0e−32 and the residual one reaches 3.2e+0. Notice which way the residual chain errs, though: raise the branch to 0.50 and it reaches 1.7e+4. Residuals do not fix the product, they change its failure from vanishing to growing — and growing is the one normalisation can hold.

That is the job LayerNorm does between blocks: it rescales each block's output back to a fixed variance, so the residual stream cannot inflate without bound however many blocks are stacked on it. The same instinct scales the output projection's initialisation by 1/√(2N) for N layers, which is a smaller f′ at the start rather than a correction after it.

And when all of that fails on one unlucky batch, there is a blunt instrument. A ball of radius one, and the raw gradient that has landed outside it. Drag the tip anywhere and watch what the optimizer actually receives:

Clipping rescales, it does not truncate: drag the tip around the outside of the ball and the used gradient keeps the direction exactly and only loses its length. Bring the tip inside and nothing happens at all. That is the trade — on a normal step it is a no-op, and on the step that would have printed nan it costs you a biased update instead of a dead run. ‖g‖ ≤ 1.0 is the value GPT-3 trained with.

Which leaves the loop itself. Six lines, one of which is the line everybody forgets:

for x, y in loader:
    opt.zero_grad()          # or grads add up
    loss = mse(model(x), y)  # keeps every
    loss.backward()          #   activation
    clip_grad_norm_(p, 1.0)  # the ball above
    opt.step()               # w -= lr * grad

backward() accumulates into .grad rather than overwriting it — deliberately, so one batch can be split across several backward passes. Drop opt.zero_grad() and here is what each step actually applies:

Step two applies g₁ + g₂; step twelve applies twelve steps' worth. The direction is roughly right, so nothing throws and the loss still falls — it is the scale that is wrong, climbing with the iteration count until the run diverges an hour in. Switch it back on and every bar lands on the line at one.

The invariant to hold on to is §05's: the gradient at any layer is the gradient one layer later times one local derivative. The cost follows, because the multiplication is shared; both failure modes follow, because a product compounds; and every recipe here follows, because the only safe product has factors near one.