Loss and Backpropagation

Cross-entropy loss: turning “wrong” into one number

It's tempting to picture training as grading a quiz. The model guesses a word, you mark it right or wrong, tally the score. That picture is wrong in a way that matters. The forward pass covered in the first seven articles of this series never outputs a single guess. It outputs a full probability distribution over the entire vocabulary, one number per possible next token, all summing to 1, the same softmax output from the very first article. Training needs to turn that whole distribution, compared against the one token that was actually correct, into a single scalar it can push down. That scalar is the loss.

The function that does this is cross-entropy loss, and the formula is smaller than people expect. Find the probability the model assigned to the correct token, call it p, and compute:

1import math
2
3def cross_entropy(probabilities, correct_index):
4    p = probabilities[correct_index]
5    return -math.log(p)
6
7# model is confident and right
8cross_entropy(probabilities=[0.01, 0.02, 0.90, 0.07], correct_index=2)
9# model is confident and wrong
10cross_entropy(probabilities=[0.01, 0.02, 0.02, 0.95], correct_index=2)

That is the entire computation. No averaging over classes, no comparison against every wrong answer individually, just the negative log of the probability assigned to the one token that actually came next. Everything else the model output that step, the other tens of thousands of vocabulary entries, gets ignored by the loss entirely. It only cares how much probability mass landed on the right answer.

Think about a weather forecaster instead of a quiz. One forecaster says 90% chance of rain, it doesn't rain. Embarrassing, but not a disaster, they hedged. A second forecaster says 99.99% chance of no rain, and it pours. That forecaster deserves to be punished far harder, not a little harder, because they staked almost all their credibility on being wrong. -log(p) is exactly this scoring rule. When p is near 1, the correct answer, -log(p) is near 0, barely a penalty. As pcreeps toward 0, the penalty doesn't rise gently, it explodes toward infinity. Confidently wrong costs dramatically more than mildly wrong, and that asymmetry is exactly the behavior you want out of a training signal.

In practice a training batch has thousands of these per-token losses, one for every position in every sequence, and they get averaged into one number for the whole batch. That single average is the number the rest of this article exists to shrink. Everything from here on is about answering one question: given that this number is higher than we'd like, which direction do we nudge every weight in the model to bring it down.

Gradient descent: which direction reduces the loss

It's tempting to imagine training as solving for the right weights directly, algebra, set the loss to zero, solve. That works for a line of best fit with two parameters. It does not work here. The loss from the last section is a function of every single weight in the model at once, the same weights from the neural network article, and a real model has billions of them. That's not a curve, it's a landscape with billions of dimensions, and there is no algebra that solves it in one step.

What you can compute, at any specific point in that landscape, is the gradient: the direction that increases the loss fastest, for every weight simultaneously. Since you want the loss to go down, not up, you move each weight a small step in the opposite direction. That's the entire idea. Not solve, just repeatedly ask “which way is downhill from here” and take a step.

Picture hiking down a mountain in thick fog. You can't see the valley, you can't see the whole terrain, but you can feel which way the ground slopes under your feet right now. So you take a step in the steepest downhill direction, feel the new slope, take another step. The size of each stride is the learning rate. Too large a stride and you don't settle into the valley, you step clean over it and land partway up the opposite slope, sometimes farther from the bottom than where you started. Too small a stride and you'll get there eventually, but eventually might mean more compute than anyone is willing to pay for. This is iterative by nature. One step is not training, it is one update out of the hundreds of thousands a real run performs.

Start the weight above away from the valley floor and hit Step a few times with the learning rate left low. Each click nudges the point downhill and the loss number ticks down. Now drag the learning rate slider up past where the stride overshoots the bottom of the curve. Keep clicking. Instead of settling, the point starts bouncing from one side of the valley to the other, and if the rate is high enough each bounce lands farther out than the last, the exact mechanical shape of a training run that diverges instead of converges.

Backpropagation: the chain rule, computed automatically

So the gradient tells you which way to step. But how do you actually compute it for a weight sitting in an early transformer block, buried under dozens of stacked blocks like the ones from the transformer block article, when all you can directly measure is the loss at the very end. The naive approach: nudge that one weight slightly, rerun the entire forward pass, see how much the loss moved, divide. That gives you the gradient for exactly one weight, at the cost of a full forward pass. Repeat that for every weight in a billion-parameter model and you'd need a billion forward passes just to take a single training step. That's not slow, it's not happening.

Backpropagationgets every weight's gradient in roughly the cost of one extra pass through the network, and the trick is a calculus rule you likely met before you ever touched machine learning: the chain rule. The loss depends on the final layer's output, which depends on the second-to-last layer's output, which depends on the layer before that, all the way back to the first embedding. The chain rule says the effect of an early layer on the final loss is just the product of each layer's local effect on the next one, multiplied together, link by link.

Think of it as an assembly line where something comes out defective at the end. You don't re-inspect every station from scratch to find the cause. You start at the last station and ask “how much did you change what came in,” then move one station back and ask the same question, reusing what you already worked out downstream instead of recomputing it. Backpropagation runs exactly this pass, once, backward through the whole network: start at the loss, compute how much the last layer's output affects it, then move one layer back and combine that with how much this layer affects the next one, then the next layer back, reusing every intermediate result instead of starting over.

1# y1 = layer1(x, w1)
2# y2 = layer2(y1, w2)
3# loss = cross_entropy(y2, target)
4
5# forward pass, left to right, values get cached
6y1 = layer1(x, w1)
7y2 = layer2(y1, w2)
8loss = cross_entropy(y2, target)
9
10# backward pass, right to left, reusing cached values
11d_loss_d_y2 = cross_entropy_grad(y2, target)
12d_loss_d_w2 = d_loss_d_y2 * layer2_grad_w(y1, w2)
13d_loss_d_y1 = d_loss_d_y2 * layer2_grad_x(y1, w2)
14d_loss_d_w1 = d_loss_d_y1 * layer1_grad_w(x, w1)
15# d_loss_d_y1 already folds in everything downstream,
16# layer 1's gradient never has to look past layer 2

Notice d_loss_d_y1in that last block. It's computed once and it already accounts for everything that happens after layer 1, the whole rest of the network, folded into a single number through the chain rule multiplication. Layer 1's own gradient just multiplies that one number by its own local derivative. It never needs to know layer 2 exists in any more detail than that. Stack forty transformer blocks instead of two and the pattern is identical, just a longer chain, each link computed once and handed backward.

This reuse is also why backpropagation needs memory, not just compute. Every intermediate value computed during the forward pass has to stay cached until the backward pass reaches it, because the chain rule needs those exact numbers to multiply against. That cache is a big share of why training a model takes so much more memory than just running it afterward.

This is the piece that makes training billions of parameters tractable at all. Not one gradient computation per weight, one backward pass, shared across every weight in the model, each one picking up exactly the piece of the chain rule product that belongs to it.

What one training step actually changes

Put the three previous sections in order and a single training step is just four moves, repeated: run the forward pass and get a probability distribution, score it against the correct token with cross-entropy loss, run backpropagation to get a gradient for every single parameter in the model, then nudge every weight a small amount opposite its gradient. Then do it again, on the next batch of text.

1def training_step(model, batch, optimizer):
2    logits = model.forward(batch.tokens)          # forward pass
3    loss = cross_entropy(logits, batch.targets)     # one scalar
4    gradients = backprop(loss, model.parameters)    # one gradient per weight
5    optimizer.step(model.parameters, gradients)     # nudge every weight
6    return loss

That optimizer.step line is doing plain gradient descent under the hood in the simplest case, subtract learning rate times gradient, but almost no real model trains on plain gradient descent anymore. The near-universal default is Adam, which keeps a running memory of each parameter's recent gradients and adjusts its effective step size per parameter instead of using one fixed learning rate for every weight in the model. The mechanism this article covers, forward, loss, backward, step, doesn't change. Adam just makes each step smarter about its own size, and that's as deep as this series needs to go on it.

Here's the part that surprises people the first time they actually watch it happen: one step barely moves anything. A single weight might shift by a few thousandths, an amount that changes the model's output on that specific example by a hair and does essentially nothing measurable to its behavior on anything else. No single step teaches the model to write code or hold a conversation. It just makes tomorrow's wrong answer on that one example fractionally less wrong.

That's exactly why this loop needs to run so many times. Not because any individual step is doing something dramatic, but because a model trains on hundreds of billions of tokens, one tiny nudge per batch, compounding across hundreds of thousands of steps. There is no shortcut version of this and no hidden insight step that does more than the others. It's the same four moves this article walked through, over and over, at a scale that the next article, Pretraining at Scale, covers directly.