Loss and gradients without the formulas: backprop is a graph walk

Part 7 ended on four changes to a block, each one baked into the weights. Every one of those weights was learned. This part reads how: a single number that says how wrong the model is, and, for every weight, which way to move it to make that number smaller. That per-weight direction is the gradient.

I read micrograd’s engine.py, Andrej Karpathy’s teaching version of automatic differentiation, which does the whole job for single numbers in 94 lines, and checked it against torch. Then I trained a small model five ways, four of them wrong. None of the wrong ones raised an error.

The loss is average surprise

A language model outputs a probability for every token in the vocabulary. The loss looks only at the token that actually came next and takes minus the log of its probability. That is cross-entropy. A token given probability 1 costs 0; a token given 1% costs 4.6. Call it surprise, measured in nats because the log is base e, about 2.718.

I ran GPT-2 small over 1,024 tokens of an earlier post. The average surprise was 3.595, the 3.60 Part 5 reported. Raise e to that power and you get 36: on average the model was as unsure as a fair pick among 36 tokens. A uniform guess over 50,257 tokens scores 10.8.

The average hides a skewed spread. If you’ve watched a p99 latency, you know the shape. The median token cost 2.85, and the worst tenth of the tokens carried 28% of the loss. The most surprising were names: ” Claude” after “renaming the assistants alone drops” cost 16.0. Among the least surprising were the pieces of a URL, “org” after “arxiv.” among them, at a ten-thousandth of a nat.

Training averages that surprise over every position of every sequence in a batch into one number. Everything below is about getting from that number back to the weights.

The forward builds a graph

micrograd wraps a number in a Value that remembers where it came from. Here is multiplication, abridged from engine.py:24-33:

def __mul__(self, other):
    out = Value(self.data * other.data, (self, other), '*')
    def _backward():
        self.grad += other.data * out.grad    # d(out)/d(self) is other
        other.grad += self.data * out.grad    # d(out)/d(other) is self
    out._backward = _backward
    return out

Every operation does two things. It computes its result, and it leaves behind a node that knows its inputs and a closure that knows its own local derivative. The forward pass computes an answer and builds a graph as a side effect. If you’ve written a tracer that records each call with a link to the calls that fed it, you’ve built half of this.

A small expression, relu(a·b + c)·a, builds 7 nodes. A 41-parameter network scored on four examples builds 394. torch does the same with tensors. Each node keeps whatever its backward will need, and for GPT-2 small one 1,024-token sequence leaves about 820 MB of saved tensors, more than the model’s 498 MB of weights. A batch of eight leaves 6.5 GB.

Backward walks it in reverse

Value.backward, abridged from engine.py:54-70, is the whole algorithm:

topo, visited = [], set()
def build_topo(v):                     # depth-first: children before parents
    if v not in visited:
        visited.add(v)
        for child in v._prev:
            build_topo(child)
        topo.append(v)
build_topo(self)
self.grad = 1
for v in reversed(topo):               # the output first, the inputs last
    v._backward()

A topological sort, then the list in reverse. The order guarantees a node has collected every contribution before it passes anything on. Each node’s closure takes the gradient that reached it, multiplies by its local derivative, and adds the result into its inputs. That’s the chain rule, run as a loop. There is nothing else. If you’ve written a build system that walks a dependency graph, this is the same walk backwards: nothing runs until everything that depends on it has reported in.

Step through the backward pass one node at a time. Then switch to grad = and watch a lose one of its two contributions.

The += is load-bearing. In the expression above, a is used twice, so two nodes send it a gradient. With += it gets −2, which torch agrees with. With =, the second write replaces the first and a gets −6. Nothing raises. The number is just wrong. On the 41-parameter network, micrograd’s gradients matched torch’s to within 0.000007.

torch accumulates the same way, into each weight’s .grad, until something zeroes it. nanoGPT, the minimal GPT-2 this series has been reading, uses that on purpose; train.py:292-305 runs 40 small batches on one GPU, divides each loss by 40, and calls backward 40 times into the same buffers before one update. optimizer.zero_grad at train.py:314 clears them afterwards.

Long chains vanish or explode

Along any path from the loss to a weight, the chain rule multiplies one local derivative per node. Deep networks have long paths, and a long product of numbers a bit below 1 or a bit above 1 is tiny or huge. That’s vanishing and exploding gradients, stated as arithmetic. You do not need calculus to see it.

I built stacks of tanh layers, 256 wide, with no residual connection. tanh squashes every value into −1 to 1. The one knob was how much each layer scales its input. At a scale of 0.5, the gradient reaching the first of 30 layers was 7 × 10⁻⁹ of the one reaching the last, and at 60 layers about 10⁻¹⁷. At a scale of 2, the first of 60 layers got about 7,700 times more.

Part 5’s residual stream, the bus every block reads from and adds back to, fixes the vanishing side. Make each layer residual, x + layer(x), and every layer gets a direct path back to the loss. In the same stacks, the first layer never got less gradient than the last, though at 60 layers and a scale of 2 it still got 992 times more. GPT-2 small is built this way, and on one sequence no block’s gradient was starved or runaway.

Pick 60 layers, set the scale to 0.5, and watch the plain stack’s gradient vanish. Then switch on the residual: at 0.5 nothing vanishes, and at 2 the first layer still gets about a thousand times more.

That’s what clipping is for. clip_grad_norm_ measures the length of the whole gradient, every weight at once, and if it is over a limit, scales all of it down. nanoGPT’s limit is 1.0, train.py:63. GPT-2 small’s gradient on my sequence had length 9.76. Nearly all of it sat in the token table, which GPT-2 reuses as its output layer; the input side got 0.14. Clipped, its length was 1.0. Clipping shortens the gradient, never turns it.

Precision adds its own vanishing. fp16, the 16-bit float format, can’t hold a positive number below about 6 × 10⁻⁸, and 9.3% of GPT-2’s gradient values rounded to zero in it. bf16, a 16-bit format that keeps fp32’s range, lost none.

The fp16 fix is to scale up before backward. nanoGPT switches on a GradScaler for fp16, train.py:196, which starts at 65,536. On my gradients that factor brought back all but 1,743 of those 11.6 million, and pushed 11 of the largest past fp16’s ceiling of 65,504. So the scaler divides the factor back out before the update, and when something overflows, it skips that update and halves the factor.

Bugs without a stack trace

In most software a bug stops the program at a line. A wrong gradient doesn’t stop anything; it trains a slightly different model. I trained a four-block character model on this blog’s text: a small GPT that reads one character per token. It ran 1,000 steps once correctly and once with each of four bugs, and I scored every run on text it never saw during training. Guessing letters by frequency alone scores 3.25; the correct run reached 1.74.

Two bugs showed on the curve. Without zero_grad, gradients piled up across steps: the loss fell to 3.37 by step 50, then climbed to 7.37, worse than a uniform guess over 106 characters. With the targets not shifted, the model was asked for the character it had just been given; training loss went to 0.001, and real next-character prediction scored 10.7.

Two didn’t. A loss multiplied by 40, the bug you get by forgetting nanoGPT’s division, reached 1.735, as good as correct. nanoGPT’s optimizer, AdamW, rescales every weight’s step, so a constant factor cancels out; Part 9 reads why. A position table detached from the graph, so no gradient ever reached it, never learned and reached 1.95, 12% worse, on a curve that looks fine.

Be fair about what a curve can tell you. Two of the four were visible to anyone watching the training loss; one of them because the loss looked too good. The other two left a curve you’d accept. If you only log the loss, you catch half.

What the ML engineer watches

Four habits, and none of them needs calculus.

Log the gradient norm per block. It is the closest thing training has to a stack trace. A spike, a block stuck near zero, or a total that creeps up over thousands of steps is visible there. The detached position table would have logged no gradient from step one. Log it before clipping (clip_grad_norm_ returns the unclipped total); after, the total never reads above the limit and every block’s spike is scaled down with it.

Evaluate on a task the training loss can’t game. The unshifted run had the best training loss of the five. Keep a score on held-out text, text the model never trains on, computed by code that doesn’t share the training loop’s data path, so a bug in one can’t hide in both.

Know which bugs your optimizer hides. Under AdamW a loss scaled by 40 changed nothing. Under plain SGD the same bug is a learning rate 40 times too large. Test the loss scaling directly: one batch, one known gradient.

Budget activations, not only weights. GPT-2 small keeps about 820 MB per 1,024-token sequence for its backward pass. Activation checkpointing, recomputing each block during backward instead of keeping its tensors, cut a batch of eight from 6.5 GB to 2.0 GB and cost 17% more time on my laptop’s CPU. Most of what’s left is 1.65 GB of log-probabilities, one for each of the 50,257 tokens at every position.

None of this is modelling work either. Backprop is a sort and a loop over closures, and every bug above is a bookkeeping error in that loop or around it. What it can’t do is fail loudly, so the logging has to.