The Backward Pass and Gradient Flow

Training a transformer means computing, for every one of its billions of parameters, how much a tiny nudge to that parameter would change the loss. That quantity is the gradient, and the algorithm that computes all of them in one pass, no matter how many parameters there are, is backpropagation (Rumelhart, Hinton, Williams, 1986). It is not a separate learning method from the forward pass; it is the same computational graph traversed in reverse, applying the chain rule at every step.

The chain rule, mechanically

If y = f(x) and loss = g(y), the chain rule says d(loss)/dx = d(loss)/dy * dy/dx. Backpropagation applies this repeatedly, starting from d(loss)/d(loss) = 1 at the output and working backward through every operation in the forward pass, each time multiplying the incoming gradient by that operation's local derivative. This is exactly why the forward pass's computational graph needs to be tracked (autodiff frameworks record every operation and its inputs) and why the backward pass costs roughly the same order of FLOPs as the forward pass, generally estimated at about twice as much, since every matmul in the forward direction implies roughly two matmuls (with respect to the input and with respect to the weights) in the backward direction (see backpropagation-autodiff).

What flows well and what does not

Two transformer components exist specifically because of how they treat gradient magnitude on the way backward.

Residual connections are additive, and the derivative of a sum with respect to either term is 1. So d(loss)/dx for y = x + Sublayer(x) includes a term that passes straight through unmultiplied, regardless of what Sublayer's own gradient looks like. That unbroken path is why residual networks can be stacked far deeper than networks without them: even if every sublayer's local gradient were tiny, the residual path alone still delivers signal to early layers (see residual-connections-skip).

Normalisation layers rescale the forward activations to a controlled range, and by extension keep the backward gradients from drifting to extreme scales as they pass through many layers. Without normalisation, small per-layer multiplicative effects compound across depth: repeated multiplication by numbers slightly above 1 explodes, repeated multiplication by numbers slightly below 1 vanishes. This is the vanishing/exploding gradient problem, and it is a direct, mechanical consequence of the chain rule applied through many layers, not a mysterious training instability (see layernorm-residual-connections).

Where gradients get multiplied, not added

Attention and the feed-forward non-linearity are where the chain rule genuinely multiplies gradients by non-trivial, input-dependent factors rather than passing them through. Softmax saturation is the clearest example: when softmax output is near 0 or near 1, its local derivative is near zero, so gradient flowing backward through a saturated softmax is crushed regardless of what happens elsewhere (this is exactly the failure mode that motivated scaling attention scores by 1/sqrt(d_k); see softmax-logits). Activation functions have the same property: a ReLU unit that output zero on the forward pass has zero derivative, so no gradient reaches the weights that fed it on that example, an effect known as a dead unit (see activation-functions).

forward:   x -> LN -> Attn -> (+x) -> LN -> FFN -> (+x) -> ...
backward:  ... <- (+1 grad) <- FFN' <- LN' <- (+1 grad) <- Attn' <- LN' <- x

When it falls down

• Gradient checkpointing trades compute for memory. Storing every intermediate activation for the backward pass is expensive, so large-scale training often recomputes some activations during the backward pass instead of storing them, roughly a 20 to 30 percent compute overhead for a substantial memory saving (see training-memory-footprint).
• Gradient clipping is a patch, not a fix. Exploding gradients are usually controlled by clipping the global gradient norm before the optimiser step; this stabilises training but does not address why the gradient exploded in the first place, and a model that needs aggressive clipping every step usually has a deeper initialisation or learning-rate problem (see signal-propagation-and-init).
• Mixed precision introduces its own gradient hazard. Gradients computed in fp16 can underflow to zero for small values, which is why mixed-precision training scales the loss up before the backward pass and back down before the optimiser step, a workaround for the limited dynamic range of fp16, less necessary with bf16.
• The backward graph is not free to build. Every operation recorded during the forward pass, tensors, intermediate values, must be kept in memory (or recomputed) until its corresponding backward operation runs, which is the direct source of the memory-versus-depth tradeoff in large models.

Further reading

• Rumelhart, Hinton, Williams, 1986, Learning representations by back-propagating errors - the original statement of backpropagation.
• Micikevicius et al., 2017, Mixed Precision Training, arXiv:1710.03740 - loss scaling and the fp16 gradient underflow problem.