Vol. V · The Transformer Plate 3 of 10pre-norm, post-norm, LayerNorm placement, training stability, warmup, learning rateAL

Pre-Norm and Post-Norm.

Move one norm from after the addition to before the sub-layer, and a deep model trains at learning rates that used to destroy it.

There are two places to put the norm in a residual block. Post-norm, the 2017 original, normalises after the addition: x ← Norm(x + f(x)). Every layer then rescales the whole stream, so the gradient travelling back must pass through a norm at every layer and the identity path of plate III.7 is broken. Pre-norm normalises only the sub-layer’s input: x ← x + f(Norm x). The stream itself is never rescaled, and the gradient has a clean path from the loss to the first layer.

Post-norm Transformers trained well but needed a careful learning-rate warmup and were fragile at depth and at high learning rates. Pre-norm, which became standard with GPT-2 in 2019, removes most of that fragility. The workshop shows recorded runs of an eight-layer model on the course corpus — pre-norm and post-norm, each at a moderate and an aggressive learning rate, with a very short warmup and no gradient clipping — and one post-norm run with a long warmup.

Below: training loss (light) and held-out loss (dark) for an 8-layer, 2.6-million-parameter model with the norm before or after each sub-layer, at peak learning rates of 0.001 and 0.006 with a 30-step warmup and no gradient clipping — and one more post-norm run at 0.001 with a 1 000-step warmup. All five were trained on a GPU and recorded.

What you will learn

Write pre-norm and post-norm blocks.
Explain why pre-norm keeps the residual path clean.
Compare their training stability from recorded runs.
Explain why post-norm needed warmup.

Workshop

Recorded runs of an 8-layer model: pre-norm (blue) and post-norm (red) at learning rates 0.001 and 0.006 with a 30-step warmup, and post-norm at 0.001 with a 1 000-step warmup (green); light lines are training loss, dark lines held-out loss. The faint line is the loss of predicting each token by its overall frequency.

Where the norm sits

With a 30-step warmup, pre-norm trains normally at both learning rates, reaching 4.575 nats held-out at 0.001 and 4.437 at 0.006. Post-norm never gets started: at both rates its loss sits at 7.58 for all 3 000 steps, which is the 7.60 nats of predicting every token by its overall frequency — the model has learned which tokens are common and nothing about context. Give post-norm a 1 000-step warmup at 0.001 and it trains properly, to 4.531, as well as pre-norm at the same rate. That is the dependence on warmup the literature describes: post-norm can work, but only if the first steps are taken gently.

Pre-norm is not free of problems: since the stream is never normalised, it grows with depth (plate III.6), so later layers’ updates are relatively small and some capacity is wasted; a final norm before the head is required. Variants keep trying to have both: “sandwich” norms on both sides of a sub-layer, extra norms on queries and keys (plate VII.7), and residual branches scaled down at initialisation (plate III.5). The practical consequence of pre-norm was large: together with warmup and clipping (plates XI.1 and XI.4), it made the deep, high-learning-rate runs of modern pretraining routine.

Ruibin Xiong and colleagues showed in 2020 that the gradients of post-norm Transformers at initialisation are largest near the output, which is why they need a learning-rate warmup and why pre-norm, whose gradients are well behaved, can do without most of it.

Terms

Post-norm: Normalising after each residual addition: x ← Norm(x + f(x)).
Pre-norm: Normalising each sub-layer’s input: x ← x + f(Norm x).
Warmup: Raising the learning rate gradually from near zero at the start of training.
Clean residual path: An identity route from the loss to the first layer that no layer rescales.

Pre-norm keeps an unnormalised residual path from output to input, and trains stably where post-norm, which normalises that path at every layer, breaks.