Four gauges worth logging
Loss is a lagging indicator. By the time it spikes, the variance problem is several hundred steps old. Four cheap scalars, logged every 50–100 steps via forward hooks, catch it early:
| Gauge | Definition | Healthy |
|---|---|---|
| Activation RMS | sqrt(mean(x^2)) at each block’s input | smooth, monotone in depth |
| Residual growth | RMS at block L ÷ RMS at block 1 | a few×, no knee |
| Global grad norm | √(Σ ||∇W||^2) pre-clip | flat or slowly falling |
| Update ratio | RMS(ΔW) ÷ RMS(W) per tensor | ≈1e-3 under Adam |
Use RMS, not mean or variance. Mean hides sign-symmetric blow-up entirely, and RMS is in the same units as the activations, so a jump from 3 to 30 reads as what it is. The update ratio is the one people skip and the one that most often explains a stalled run: a tensor moving at 1e-6 of its own scale is frozen, and one moving at 1e-1 is being destroyed.
The residual-stream growth curve
The single most informative plot in transformer training is activation RMS against block index. Signal-propagation theory tells you the shape to expect: the residual stream is additive, so roughly independent branch contributions make variance grow linearly and RMS grow like √L. That is the reference curve. Everything you learn comes from how the measured curve departs from it.
A healthy curve is smooth and gently concave: each block adds a similar quantum of variance to an already-large stream, so the marginal effect shrinks. Three pathologies are recognizable at a glance. A knee, where RMS is flat for twenty blocks then hooks upward in the last three, means one or two blocks dominate the stream and the rest have been squeezed into irrelevance. Exponential growth (a straight line on a log axis) means the recursion is multiplicative somewhere it should not be. A flat curve near the input means early blocks contribute nothing and their gradients have vanished.
Worked example: a residual-stream budget
Take a 24-layer, d = 768 Pre-LN model, so 2L = 48 residual branches (attention and MLP per block). Each branch reads a LayerNorm’d input of unit RMS, so with a variance-preserving init each branch writes a contribution of variance ≈ 1. Variances add:
unscaled: q_L = q_0 + 2L · q_F = 1 + 48 · 1 = 49 → RMS ≈ 7.0
scaled: branch std × 1/√(2L) ⇒ q_F = 1/48
q_L = 1 + 48 · (1/48) = 2 → RMS ≈ 1.41So the 1/√(2L) convention is not cosmetic: it is the difference between a stream that grows 7× across depth and one that grows 1.4×. As a ratio it becomes a test you can run on any checkpoint. Log RMS at block 1 and block 24; at initialization the scaled model should show well under 2×. Training will push that ratio up — real models commonly reach 10× or more — but it should climb slowly and smoothly. A ratio that doubles within a few hundred steps is the alarm.
Normalization placement decides what you must control
Where you put the norm determines which quantity is protected and which one you have to police yourself. Pre-LN normalizes each branch’s input and leaves the residual stream itself unnormalized. That makes every block robust to the stream’s absolute scale, and it is why Pre-LN trains with little warmup — but it also means nothing bounds the stream, so the growth curve above is your responsibility, and a final norm before the output head is mandatory.
Post-LN puts the norm on the residual path itself, so the stream is pinned to unit RMS at every block and the growth curve is trivially flat. You pay for that with a multiplicative gradient recursion through depth, which is exactly the instability that makes warmup non-negotiable. Variants in between — normalizing a branch’s output too, or scaling the residual path by a depth-dependent constant — all try to keep Post-LN’s bounded stream without its gradient behaviour.
Attention logit growth and QK-norm
One variance failure hides from every activation gauge, because it happens inside the softmax. For query and key vectors whose entries have standard deviations σ_q and σ_k, the scaled logit has Var(q·k / √d_h) = d_h · σ_q^2 σ_k^2 / d_h = σ_q^2 σ_k^2. The 1/√d_h cancels the head dimension exactly — but it does not cancel growth in the projections. If training grows W_Q and W_K until the q and k entries they produce reach std 3, logit std jumps from 1 to 9.
Past that point the softmax saturates: attention entropy collapses toward a one-hot pick, the softmax Jacobian goes to zero, and the head stops learning while contributing a large, brittle activation. Log max attention logit and attention entropy. The structural fix is QK-norm — an RMSNorm on q and k before the dot product — which pins σ_qσ_k to a learned gain and caps logit scale by construction. The output head has the same disease; a small z-loss penalizing log^2(Z) is the usual cure.
Variance meets the number format
Every variance argument above assumes real numbers. In fp16 it collides with a 5-bit exponent: the smallest normal magnitude is about 6e-5 and the largest is 65504. Gradients routinely live at 1e-7 or below, which is not small — it is zero in fp16. The classic symptom is a run that trains but plateaus early, its smallest-gradient tensors silently frozen.
Loss scaling is the fix: multiply the loss by S (say 2^15) before backward, so the whole gradient distribution shifts up into representable range, then divide by S before clipping and the optimizer step. Dynamic loss scaling automates the search — halve S and skip the step whenever an inf or NaN appears, double it after a few thousand clean steps. Master weights and optimizer moments stay in fp32 regardless, because a 1e-3 relative update added to a fp16 weight rounds away entirely.
Why bf16 changed the calculus
bf16 keeps fp32’s 8-bit exponent and spends the savings on mantissa instead: it can represent roughly 1e-38 to 3e38, the same dynamic range as fp32, with about 3 decimal digits of precision. That single change makes the entire loss-scaling apparatus unnecessary. Gradients no longer underflow, activations no longer overflow, and the skipped-step machinery disappears — which is why bf16 became the default.
What you trade is precision, and the trade relocates the risk rather than deleting it. With ~8 bits of mantissa, adding a small update to a large weight can round to a no-op, and long summations lose accuracy — so accumulate matmuls and reductions in fp32, keep fp32 master weights, and compute the softmax and the loss in fp32. bf16 removes the range failures that produce NaNs; it does not remove the precision failures that produce a run which quietly learns less than it should.