Backpropagation and the optimiser

Key idea

Chain rule backwards through every layer; residuals keep gradients alive; Adam scales each step; learning rate warms up then decays

Real world

A relay team loses a race. To fix it, the coach works backwards from the finish: how much time did the last runner lose, and how much of that came from a bad hand-off by the runner before? Each runner gets their share of the blame.

In an LLM

After one forward pass over a batch, one backward pass gives a gradient for every one of GPT-2 small’s 124 million weights, at about twice the cost of the forward pass. Adam then moves each weight by a tiny amount, on the order of the learning rate (a typical peak for a model this size is about 0.0006).

What training needs

The training objective gives one number, the loss, for a batch of text. Gradient descent lowers it by moving every weight a little against its gradient: how much the loss would change if that weight changed. A model has millions or billions of weights, so we need millions or billions of gradients, every step.

Wiggling each weight separately and re-running the model would take one forward pass per weight. Backpropagation gets all of them from a single backward pass.

The chain rule at a glance

A network is a chain of simple steps. If $a$ affects $b$, and $b$ affects the loss $L$, then

how much $L$ changes per unit of $a$ is how much it changes per unit of $b$, times how much $b$ changes per unit of $a$. Backprop applies this from the loss backwards: start with $\partial L / \partial L = 1$, and at every step multiply by one local slope, the derivative of that step on its own. Each step only needs its own inputs from the forward pass, which is why training stores the forward pass’s intermediate values (the activations).

A transformer is the same, just bigger: the “steps” are matrix multiplications, softmaxes, norms and activations, and the local slopes are matrices instead of single numbers. Frameworks like PyTorch record every operation during the forward pass and run this backward sweep automatically.

Gradients flowing down the stack

In a transformer, the gradient starts at the loss, goes back through the unembedding, and then down through every block from the top to the bottom. Here is where the residual stream earns its keep.

In a plain stack, $x_{l+1} = f_l(x_l)$, the gradient reaching layer 1 is a product of one local slope per layer. If those slopes are typically a bit below 1, the product shrinks exponentially with depth: the vanishing gradient. A bit above 1 and it explodes. With a residual, $x_{l+1} = x_l + f_l(x_l)$, the local slope is $1 + f'_l$: there is always a path where the gradient is passed straight through, multiplied by exactly 1. Watch it in a simulated network:

With the default settings, the plain 32-layer stack delivers a gradient of a few times $10^{-8}$ to its first layer: those weights would effectively never change. The residual stack delivers about 3. This is why every deep network since ResNet, and every transformer, has residual connections, and why pre-norm (which keeps the residual path free of norms) trains more reliably than post-norm.

Adam: one step size per weight

Plain gradient descent uses one learning rate for every weight. That works badly for transformers, whose gradients differ in size by orders of magnitude between, say, the embedding of a rare token and a busy attention matrix. Adam (Kingma and Ba, 2014) keeps two running averages for every weight:

• $m$, the average gradient (momentum: keep going in a consistent direction),
• $v$, the average squared gradient (how large and noisy this weight’s gradients usually are),

and moves each weight by $\eta \cdot m / (\sqrt{v} + \epsilon)$. Dividing by $\sqrt{v}$ means every weight takes steps of roughly the size $\eta$, whatever the scale of its gradients. AdamW (Loshchilov and Hutter, 2017) adds weight decay, a gentle pull of every weight towards zero, and is the standard optimiser for LLMs. Its cost: two extra numbers stored per weight, which is why training needs far more memory than inference (see low-precision training).

The learning-rate schedule

The learning rate $\eta$ is not fixed. Almost every LLM uses the same shape:

1. Learning-rate warm-up: start near zero and rise linearly over the first few thousand steps. At the start the weights are random and Adam’s averages are empty, so large steps would be based on bad statistics and can make the loss spike.
2. Decay: after the peak, lower the learning rate along a cosine curve to a small floor. Late in training the model is fine-tuning many small details, and small steps stop it from bouncing around the minimum.

GPT-3 warmed up over its first 375 million tokens and decayed to 10% of the peak. Llama 3 405B warmed up over 8,000 steps to a peak of $8 \times 10^{-5}$ and decayed by a cosine to $8 \times 10^{-7}$. Some recent runs use “warm-up, stable, decay” instead: a long flat stretch at the peak and a short decay at the end, which makes it easy to keep training from a checkpoint.

You can see all of this at work, gradients, Adam and a schedule, in train a tiny transformer live.

opt = torch.optim.AdamW(model.parameters(), lr=6e-4, betas=(0.9, 0.95), weight_decay=0.1)
for step, batch in enumerate(data):
    for g in opt.param_groups:
        g["lr"] = lr_at(step)                  # warm-up + cosine
    loss = model(batch).loss
    loss.backward()                            # backprop: fills p.grad for every weight
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); opt.zero_grad()