Small-scale proxies for large-scale Transformer training instabilities

Attention logits blow up first, and QK-norm stops them. Without a fix, the learning rate at which runs diverge falls with size, and the maximum attention logit climbs by orders of magnitude before the loss blows up. The growth is in the norms of queries and keys, while their cosine similarity does not rise; the authors' explanation is that attention logits are the only quantity quadratic in the parameter scale, $\langle XW_1,XW_2\rangle$, so they are the first to grow large as the weights grow. Forcing a 10M model's logits to a fixed scale of $10^4$ makes it worse than a bigram model, so the saturated softmax is the cause. Every run whose maximum logit at step 2000 exceeded $10^4$ diverged, and a fit of that maximum against log size predicted the 4.8B divergence at LR 0.01 before that model was trained. Normalizing queries and keys per head before the dot product removes the failure, and a 1.2B model then trains at LR 0.3. The paper uses LayerNorm with learned gains; the speedrun uses a parameter-free RMS form, under which every query and key has norm $\sqrt{d_h}$ and Cauchy-Schwarz bounds each logit by the scale times $d_h$.

Output logits drift unless something holds $\log Z$ near zero. The cross-entropy is invariant to adding a constant to every output logit, so the mean logit lies along a flat direction and nothing pulls it back. Without weight decay, at high learning rates and every size tested, the logits drift strongly negative late in training and the loss degrades. PaLM's z-loss adds a penalty that pulls $\log Z$ back toward zero, $$\mathcal L=\mathrm{CE}(y,\text{target})+10^{-4}\,(\log Z)^2,\qquad Z=\textstyle\sum_je^{y_j},$$ and removes the instability; weight decay also removes it for the larger models. A logit softcap blocks the same drift by bounding the logits (my connection; the paper does not test caps).

Gradients shrink until AdamW's $\epsilon$ dominates the denominator. The gradient RMS falls with model size and learning rate, and at 4.8B and LR 0.3 it reaches about $10^{-8}$, AdamW's default $\epsilon$. In $m/(\sqrt v+\epsilon)$, once $\sqrt v\approx\epsilon$ the update shrinks well below its intended scale and the layer stops learning, and the measured update RMS drops when the gradient RMS crosses $\epsilon$. Lowering $\epsilon$ to $10^{-15}$ improves the 4.8B loss; raising it to $10^{-6}$ makes the run diverge. The proposed mechanism is output-norm growth: the residual stream's RMS grows with depth and learning rate, and the final layernorm divides the gradient it passes back by that RMS.

How standard hyperparameter choices change LR sensitivity. With qk-layernorm and z-loss in place, other interventions barely move the range of stable learning rates but change the loss inside it. Longer warmup lowers LR sensitivity, most for the larger models; weight decay applied independently of the learning rate lowers it relative to the PyTorch default, where a learning-rate sweep also sweeps the decay; and depth raises it faster than width. $\mu$Param keeps the best learning rate fixed across width but neither lowers sensitivity nor improves loss, and it does not remove the need for qk-layernorm. More steps or a larger batch leave sensitivity unchanged.

Which of these fixes the speedrun adopted. QK-norm came in at record 5, alongside other changes. Because it fixes the query and key norms, the logit scale (the softmax temperature) becomes a hyperparameter, and record 20 set it to 0.12 instead of $1/\sqrt{d_h}$ (about 3 s; with $d_h=128$ that caps logits near 15 instead of 11). For the output logits the speedrun uses a cap: $30\tanh(z/30)$ at record 9, cap 15 at record 18 (about 10 s, 1490 to 1390 steps), now an asymmetric $23\,\sigma((z+5)/7.5)$; these gains are steps saved, so the cap does more than prevent divergence, which is outside this paper. Adam's $\epsilon$ went from $10^{-8}$ to $10^{-10}$ at record 20 (about 10 steps), by the same mechanism for a different cause: with a zero-initialized head, the gradients reaching the Adam-trained embeddings start near zero, below $\epsilon$. Muon normalizes whole matrices and has no per-coordinate $\epsilon$ in this role.

Limits. The paper studies only slow divergences (fast loss spikes are out of scope), in one codebase with one data set and optimizer; LR sensitivity does not reflect the overall loss level and meaningless for interventions that change what the learning rate means; and the mechanisms (quadratic dependence, output-norm growth) are hypotheses supported by measurements.