Large runs fail in two ways small ones rarely show: the loss jumps, or part of the model quietly saturates. Both reproduced on a tiny model: attention-logit growth and QK-norm, z-loss, one bad batch, clipping and skipping, and what large runs report doing.
This deck follows chapter 8, stability, where a run at ten times the default learning rate shows one attention head collapse onto single characters. Measured numbers: results.md §8a–8b.
The loss suddenly jumps, and either recovers or never comes back. Often one step with a huge gradient throws the weights somewhere bad.
The loss stops improving because part of the model has saturated: attention heads that put all their weight on one position pass no gradient back, so they stop learning.
Both can be reproduced on the Decoder’s 5,376-parameter model, three seeds each, and both have standard cures.
At a high learning rate the query and key weights grow, the scores \(q^\top k/\sqrt{d_k}\) grow with their product, and the softmax becomes a hard maximum: each position attends to exactly one other, the attention entropy falls towards zero, and the gradient through the softmax vanishes (deck 03). Dehghani et al. (2023) saw exactly this diverge an 8-billion-parameter vision transformer; Wortsman et al. (2023) showed the same instability in small models at high learning rates.
| Peak lr | Run | Final validation loss | Largest score | Entropy (nats) |
|---|---|---|---|---|
| 0.01 | baseline | 2.277 ± 0.090 | 42.2 | 0.62 |
| 0.01 | QK-norm | 2.272 ± 0.033 | 11.0 | 0.60 |
| 0.01 | z-loss 1e-3 | 2.310 ± 0.113 | 67.5 | 0.38 |
| 0.1 | baseline | 2.571 ± 0.102 | 16,104.7 | 0.05 |
| 0.1 | QK-norm | 2.102 ± 0.039 | 24.4 | 0.34 |
| 0.1 | z-loss 1e-3 | 2.589 ± 0.102 | 19,648.6 | 0.05 |
AdamW, warmup 50 + cosine, 2,000 steps, clip 1.0, three seeds; largest attention score and mean attention entropy on the validation windows at the end (results.md §8b). At 0.1 the baseline does not blow up; it learns less.
Eight attention scores for one query, multiplied by a growing scale (as the \(W_Q, W_K\) weights grow). Watch the weights, the entropy \(H = -\sum_j p_j \log p_j\) and the size of the gradient the softmax passes back.
The upstream gradient \(\delta P\) is fixed in this demo; only the scale changes. Toy numbers, chosen for the picture; the measured run is on the previous slide.
Apply a LayerNorm to each head’s queries and keys before the dot product, so the scores cannot grow just because the weights do (only the norm’s gain can scale them). At learning rate 0.1 the largest score stays at 24.4 and the run reaches 2.102 ± 0.039: the best of every run here, because the higher learning rate is now usable.
PaLM added \(10^{-4}\log^2 Z\) to the loss (Chowdhery et al. 2022; here \(10^{-3}\)), pulling the output softmax’s normaliser \(\log Z\) towards 0. It targets a different instability, output logits drifting, and does nothing for attention: the largest score still reaches 19,649.
The plainest cure is a lower learning rate, which costs speed.
Train with momentum (no adaptive scaling) at learning rate 0.5, and at step 1,000 feed one batch that is nothing like the book: every window the character “@” repeated. Its gradient is huge. 1,400 steps, constant learning rate, three seeds (results.md §8a):
| Run | lr | Validation loss before | Just after | At the end |
|---|---|---|---|---|
| momentum, no clipping | 0.5 | 2.334, 2.424, 2.296 | 17.138, 8.175, 14.494 | diverged, diverged, diverged |
| momentum, clip at 1.0 | 0.5 | 2.424, 2.454, 2.359 | 2.477, 2.598, 2.448 | 2.325, 2.351, 2.318 |
| momentum, skip the batch | 0.5 | 2.334, 2.424, 2.296 | 2.334, 2.424, 2.296 | 2.322, 2.276, 2.306 |
| AdamW (lr 0.01), no clipping | 0.01 | 2.367, 2.329, 2.435 | 2.377, 2.342, 2.441 | 2.229, 2.250, 2.282 |
Without protection the step throws the weights far off and every seed diverges. AdamW at its usual learning rate barely notices the batch even unclipped: Adam divides each step by the running size of past gradients, so one large gradient cannot make one large step.
Rescale any gradient whose global norm exceeds a threshold \(c\) (Pascanu et al. 2013; PyTorch’s clip_grad_norm_; \(c = 1\) here and in the default run). The bad batch now moves the loss only a little, and the run carries on.
Skip a batch whose gradient norm is far above the recent average (here 4×): the step never happens, so the loss does not move at all.
Why momentum is fragile: its step is \(\eta\) times a running sum of gradients, so one gradient \(k\) times too large moves the weights \(k\) times too far. Adam’s step \(\eta\,\hat m/\sqrt{\hat v}\) is bounded by about \(\eta\) per coordinate.
Warmup, clipping, QK-norm or z-loss, a lower \(\beta_2\), monitoring the gradient norm, and checkpoints to restart from (deck 13). The site’s notebook 08 repeats the QK-norm / z-loss comparison on a width-512 model on a GPU: GPU run pending.
Animated: chapter 8, stability · next: deck 08, mixed precision