Chapter 8 of 16

Stability

A loss spike reproduced and cured: clipping, z-loss, QK-norm.

Slides for this chapter: Transformer Training 07, Stability

Large training runs fail in two ways that small ones rarely show: the loss suddenly jumps (a spike) and either recovers or never comes back, or it quietly stops improving because part of the model has saturated. Both can be reproduced on the Decoder's model. Below, the model trains at ten times the default learning rate, and at a few saved steps the site runs its own forward pass on the saved weights: watch one attention head's pattern collapse onto single characters as its scores grow by orders of magnitude.

Loading the animation…

Concept

Attention-logit growth. 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 turns into a hard maximum: each position attends to exactly one other, the attention entropy falls towards zero, and the gradient through the softmax vanishes, so attention stops learning (Decoder, chapter 3 shows the healthy version). Dehghani et al. (2023) saw exactly this diverge an 8-billion-parameter vision transformer; Wortsman et al. (2023) showed the same instability appears in small models at high learning rates. Here, at a peak learning rate of 0.1 (warmup, cosine, clip 1.0, three seeds), the largest attention score reaches 16,105 on average and the mean attention entropy falls to 0.05 nats (at the default 0.01: 42 and 0.62). The run does not blow up; it just learns less: 2.571 ± 0.102 against 2.277 ± 0.090 at 0.01.

The cures. QK-norm applies 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). With it, the run at 0.1 keeps its largest score at 24 and reaches 2.102 ± 0.039: the best of every run in this chapter, because the higher learning rate is now usable. z-loss (Chowdhery et al. 2022: PaLM added $10^{-4}\log^2 Z$ to the loss; here the coefficient is $10^{-3}$ ) pulls the output softmax's normaliser $\log Z$ towards 0; it targets a different instability, the output logits drifting, and does nothing for attention: with it the largest score still reaches 19,649 and the loss 2.589 ± 0.102. The plainest cure is a lower learning rate, which costs speed.

A loss spike: one bad batch. The second experiment trains with momentum (no adaptive scaling) at a learning rate of 0.5, and at step 1,000 feeds one batch that is nothing like the book: every window the character "@" repeated. Its gradient is huge. Without protection the step throws the weights far off: the validation loss jumps from 2.33 to 17.1 and the run then diverges (seed 1; all three seeds do the same, results.md §8a). Gradient clipping (Pascanu et al. 2013) rescales any gradient whose norm exceeds a threshold (1.0 here, and in the default run): the same bad batch now moves the loss from 2.424 to 2.477, and the run carries on. Skipping a batch whose gradient norm is far above the recent average (here 4×) avoids even that. AdamW at its usual learning rate barely notices the batch even unclipped (2.367 → 2.377): Adam divides each step by the running size of past gradients, so one large gradient cannot make one large step.

momentum, no clipping · momentum, clip at 1.0 · momentum, skip the batch · AdamW (lr 0.01), no clipping.Seed 1; all three seeds are in results.md §8a.

What large runs report. PaLM's largest model spiked about 20 times despite clipping; the fix was to restart from a checkpoint about 100 steps before the spike and skip the 200–500 batches around it. Its authors note the spikes were not caused by bad data as such: the same batches did not spike from another checkpoint (Chowdhery et al. 2022, §5.1). Molybog et al. (2023) argue that such divergence is an artifact of Adam itself: it can enter a state where its update is large yet essentially uncorrelated with the direction that lowers the loss, more often in deep models trained with large batches. So the usual toolkit is several guards at once: warmup, clipping, QK-norm or z-loss, a lower $\beta_2$ (LLaMA used 0.95; Touvron et al. 2023), monitoring the gradient norm, and checkpoints to restart from (chapter 15).

Run it yourself on a GPU. Notebook 08 (open in Colab) repeats the baseline / QK-norm / z-loss comparison on a 6-layer, width-512 model on a GPU (about 45 minutes on an A100). The notebook is one of the planned GPU replications listed under Runs to come; every number on this page comes from the in-browser model.

Maths

Attention entropy of row $i$ : $H_i = -\sum_{j \le i} p_{ij}\log p_{ij}$ , between 0 (one-hot) and $\log(i+1)$ (uniform over the $i+1$ visible positions). The softmax's backward pass $\delta s_{ij} = p_{ij}(\delta p_{ij} - \sum_k \delta p_{ik} p_{ik})$ (chapter 4) is $\approx 0$ when $p_i$ is one-hot: saturated heads stop learning.

z-loss: $L = L_{\text{CE}} + c\,\overline{(\log Z)^2}$ with $\log Z = \operatorname{logsumexp}(z)$ ; $\partial/\partial z = 2c \log Z \cdot \operatorname{softmax}(z)$ per position, added to $p - y$ .

Clipping by global norm: $g \leftarrow g \cdot \min\big(1, \tfrac{c}{\|g\| + 10^{-6}}\big)$ (PyTorch's clip_grad_norm_). Momentum's 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 is $\eta\, \hat m/\sqrt{\hat v}$ , bounded by about $\eta$ per coordinate.

Code

QK-norm and z-loss in the reference model, scripts/training/model.py (gradient-checked against PyTorch autograd):

qk_cache = None
if qk is not None:
    Qh, cq = layernorm_fwd(Qh, qk[0], qk[1])
    Kh, ck = layernorm_fwd(Kh, qk[2], qk[3])
    Qh, Kh = cast.a(Qh), cast.a(Kh)
    qk_cache = (cq, ck)

def z_loss_terms(logits, probs, coeff):
    """PaLM's z-loss, coeff * mean(lse^2), and its gradient 2 coeff lse p / n."""
    m = logits.max(-1, keepdims=True)
    lse = m[..., 0] + np.log(np.exp(logits - m).sum(-1))
    n = lse.size
    return coeff * float((lse * lse).mean()), (2 * coeff / n) * lse[..., None] * probs

Comments

Be the first to leave a comment on this section.