Vol. VII · Modern Architectures Plate 7 of 10training stability, loss spikes, QK-norm, z-loss, logit soft-capping, learning rate, divergenceLE

Small Changes That Stabilise.

When a run blows up after weeks, the culprit is usually a logit that grew too large — so cap it, normalise it, or penalise it.

A pretraining run that costs millions can be lost to one bad hour: the loss spikes, recovers or does not, and engineers roll back to a checkpoint and skip the batch. Reports from PaLM, OPT and others describe exactly this. Many instabilities trace back to logits growing too large — inside attention, where q · k scores saturate the softmax into one-hot patterns (plate IV.3), or at the output, where the log-normaliser drifts.

Three small fixes are now common. QK-norm normalises queries and keys before their dot product, so attention scores cannot grow just because the vectors do (used in Gemma 3, OLMo 2 and others). z-loss adds a penalty on the square of the output log-normaliser log Z, keeping it near zero (PaLM). Logit soft-capping passes logits through c · tanh(x / c), so none can exceed c (Gemma 2). Wortsman and colleagues (2023) showed that small models reproduce the instabilities of large ones if the learning rate is pushed high enough — which is what the recorded runs do.

Below: left, the final held-out loss of an 8-layer model trained with no gradient clipping at three learning rates, plain and with each fix. Right, the largest gradient norm seen in each run.

What you will learn
Workshop
Recorded runs: 8 layers, d = 128, 3 000 steps, no gradient clipping, at learning rates 0.01, 0.03 and 0.1. Left: final held-out loss (a run that diverged is plotted at 7.6 nats, the loss of predicting token frequencies alone). Right: largest gradient norm among the steps logged in each run (one in thirty), log scale.

One fix for one failure

At a learning rate of 0.01 all four variants train normally and end within a few hundredths of a nat of each other — the plain model at 4.381. Triple the rate to 0.03 and the plain model fails to learn properly: it ends at 6.104 nats, and its gradient norm, under 1 at the lower rate, reaches 25.1. With QK-norm the same rate is harmless: 4.377 nats, essentially as good as at 0.01, with gradient norms that stay below 1. At 0.1 the plain model ends at 6.663 and QK-norm at 4.610.

z-loss and soft-capping, at 6.826 and 6.748 at the highest rate, do not rescue the run, and that is informative rather than disappointing: both act on the output logits, while QK-norm acts inside attention, so the failure here is attention logits growing until the softmax saturates — the instability Wortsman and colleagues found to dominate in small models at high learning rates, and which QK-norm fixed in their runs too. Output-logit divergence, which z-loss targets, is a separate failure that these runs did not show. The general lesson carries: diagnose which logits are growing, then bound those.

Gradient clipping, warmup and a lower learning rate are the older answers to instability. They work, but each costs something — slower progress or a smaller usable learning rate — which is why fixes that remove the cause are preferred.

Terms

Loss spike
A sudden jump in training loss, sometimes recovered from, sometimes not.
QK-norm
Normalising queries and keys before their dot product, bounding attention scores.
z-loss
A small penalty on the squared log-normaliser of the output softmax.
Logit soft-capping
Squashing logits through c tanh(x / c) so none exceeds c.
Instabilities usually come from logits that grow too large; normalising, penalising or capping them widens the range of learning rates that train well.

Sources

  1. Wortsman et al., “Small-scale proxies for large-scale Transformer training instabilities”, arXiv:2309.14322, 2023.
  2. Chowdhery et al., “PaLM: Scaling Language Modeling with Pathways”, JMLR 2023, section 5 (z-loss, loss spikes).
  3. Gemma Team, “Gemma 2: Improving Open Language Models at a Practical Size”, arXiv:2408.00118, 2024 (logit soft-capping).
  4. Dehghani et al., “Scaling Vision Transformers to 22 Billion Parameters”, ICML 2023 (QK normalisation).