# Time-bomb loss spikes: Adam drifts the LM head, BF16 pulls the trigger

> Satyajit Ghana — Head of Engineering @ Inkers Technology
> canonical: https://ai.satyajit.si/articles/lm-head-loss-spikes
> date: 2026-11-09
> tags: training, optimization, low-precision, numerics

A pretraining run that goes wrong in the first thousand steps is cheap. One that looks healthy for most of its budget and then starts climbing, and never comes back down, is the expensive kind. [The Role of Quantization in Loss Spikes](https://arxiv.org/abs/2610.38277) (Zaiqian Chen, Aditya Somasundaram, Sophie Li, Sean McLeish, John Kirchenbauer, Tom Goldstein and Micah Goldblum; Columbia, MBZUAI, Maryland and the Vector Institute; arXiv, v2 of 2 November 2026) calls these **time-bomb spikes** and gives them a cause in two parts. Adam makes the mean logit of the LM head drift. BF16 then cannot represent the differences between large logits.

Goldblum's thread on X put the puzzle in one line: "Adding the same constant to every logit changes neither the softmax distribution nor the cross entropy gradients in exact arithmetic. So why should this mean logit drift cause training to fail?" The answer the paper gives is finite precision, and the fix it proposes is small: "swap AdamW for momentum SGD only on the LM head, and keep AdamW everywhere else."

The code is public under MIT, a JAX and Flax training stack forked from Martin Marek's:

<RepoCard repo="tingtang2/quantization-loss-spikes" />

The paper trains on TPUs and B200s, which I don't have. What I can do on a CPU is isolate the LM head, make it drift, round it the way a BF16 GEMM does, and swap the optimizer. So the evidence here comes in three labels. **Measured**: I ran it, in NumPy 2.5.3 on a shared 4-core cloud VM, with BF16 simulated by rounding fp32 bits to nearest even. **Reported**: the paper's figures and the repository's configs, read and not re-run. **Reasoned**: my arithmetic on the other two.

## Two kinds of spike

The paper separates two failure shapes (reported). **Fast spikes** jump and recover within a few hundred steps. In their runs these mostly come from RMSNorm layers without a learnable gain: across five matched seed pairs at 162M parameters, adding the gain suppresses them along with the logit swings that come with them. Skipping the batch that triggered one only postpones the next, and masking as little as 0.2% of the highest-loss tokens hides them but slows convergence.

**Time-bomb spikes** survive the learnable gains. Two 404M-parameter models trained for 96B tokens, about 240 tokens per parameter or roughly 12 times the Chinchilla budget, follow the same loss curve for most of training. Then the higher-learning-rate run stops improving and its evaluation loss rises monotonically while its output logit norm accelerates (reported, Figure 3). Longer horizons make it worse: at 28M parameters, runs of 4B, 8B, 12B and 16B tokens show progressively larger late logit growth (reported, Figure 17). That matters because training far past compute-optimal is now the normal case.

<Figure src="https://media.thesatyajit.com/ai/articles/lm-head-loss-spikes/fig17.png?v=86a321aa81ca48e9" alt="Three panels against training tokens for 28M-parameter models trained to 4B, 8B, 12B and 16B tokens. Left, eval loss: all four fall to about 4.1; the 12B run turns up near 10B tokens and the 16B run shoots past 5 near 15B. Middle, output logit norm on a log scale, rising from about 5e5 to 1e7 for the longer runs and to about 3e8 for the 16B run. Right, output logit mean: near zero for most of training, then falling to about -650 for the 12B run and past -1000 for the 16B run." caption="The longer the horizon, the further the mean logit falls and the later, harder the failure: the shorter runs end before the bomb goes off. 28M-parameter models at peak learning rate 0.1 (paper, Figure 17)." />

The right panel is worth a second look: the real 28M model's mean logit reaches about $-1{,}000$ at 16B tokens. That is the same scale my toy drifts to, and the scale at which BF16 neighbours are 4 to 8 apart (reported figure; the comparison is mine).

The experiments use QK-Norm everywhere, so attention-logit growth is ruled out as a cause. What is left is the LM head.

## The model, with shapes

Every number in this section is read from the repository's Hydra configs (`configs/model/*.yaml`, `configs/experiment/*.yaml`, `configs/base.yaml`) at commit `76153f0`, and the arithmetic on them is mine (reasoned).

The architecture is a modernised GPT-2: rotary positions, GELU, RMSNorm, QK-Norm, untied input and output embeddings, no biases. The GPT-2 tokenizer gives $V = 50{,}257$. Four sizes appear:

| config | width $D$ | layers | context | head $E$ shape | head parameters |
|---|---:|---:|---:|---|---:|
| `lm28m` | 256 | 3 | 512 | 50,257 × 256 | 12.9M |
| `lm162m` | 768 | 12 | 1,024 | 50,257 × 768 | 38.6M |
| `lm404m` | 1,024 | 24 | 1,024 | 50,257 × 1,024 | 51.5M |
| `lm1p4b` | 2,048 | 24 | 2,048 | 50,257 × 2,048 | 102.9M |

At 28M, where the precision experiments run, the head is close to half the model: three blocks at $12D^2$ come to about 2.4M and the input embedding to another 12.9M. Per position, the final RMSNorm produces $h \in \mathbb{R}^{256}$ and the head computes $z = Eh \in \mathbb{R}^{50{,}257}$. In the `bf16_mixed` config that GEMM takes BF16 inputs and returns BF16 logits (`lm_head_dtype: bfloat16`, `lm_head_param_dtype: float32`), and the loss casts them to fp32 only afterwards. By the time anything runs in fp32, the logits have already been rounded once.

The time-bomb recipe at 28M (`timebomb_28m.yaml`): 131,072 tokens per step, 16B tokens, so about 122,000 steps (the 120k-step axis on the paper's figures); peak learning rate 0.1, 5% warmup, cosine to zero; AdamW with $\beta_1 = 0.9$, $\beta_2 = 0.95$, $\epsilon = 10^{-8}$; global-norm clipping at 1. The paper's setup says weight decay is $10^{-4}$ "unless otherwise specified"; the fast-spike configs set that, but both time-bomb configs set `weight_decay: 0.0`, which is what the argument below needs for the drift to be exactly invisible. One step's logits at 28M are about 6.6 billion numbers, 13 GB in BF16, which is why moving the head to FP32 is not free.

## The gradient that sums to zero

Write the LM head as an output embedding matrix $E \in \mathbb{R}^{V \times d}$ acting on the final hidden state $h$, so the logits are $z = Eh$ and $p = \operatorname{softmax}(z)$. The cross-entropy gradient with respect to $E$ is an outer product:

$$
\frac{\partial \mathcal{L}}{\partial E} = (p - y)\,h^\top,
\qquad
\sum_{i=1}^{V} \frac{\partial \mathcal{L}}{\partial E_{ij}} = h_j \sum_{i=1}^{V} (p_i - y_i) = 0 .
$$

Both $p$ and the one-hot $y$ sum to one, so every column of the gradient sums to zero over the vocabulary. Any optimizer that applies a fixed linear map to that gradient, the same for every row, keeps the zero sum. Plain SGD does. So does SGD with momentum, whose buffer is a running sum of zero-sum gradients. Under either, the mean output embedding $\mu = \frac{1}{V}\sum_i e_i$ never moves.

Adam does not. Its first moment is still zero-sum, but it divides each coordinate by its own $\sqrt{v}$. Rows for rare tokens see small, mostly one-signed gradients (they are pushed down a little every time they are not the target) and so have a small $v$. Dividing by it turns that trickle into a full-size step. Frequent tokens get large, two-signed gradients and a large $v$. The weighted rows no longer cancel, so the mean embedding drifts, and with it the mean logit $\mu^\top h$. This part is not new: the paper credits Biś et al. (2021) and Stollenwerk and Stollenwerk (2025) for it.

I wanted to see the size of that imbalance in a real step, so I built a toy (measured). The trunk is frozen: 1,024 contexts, each a 64-wide RMS-normalised hidden state that shares one mean direction with the others, which is the anisotropy real final hidden states have. Targets come from a 2,048-token vocabulary with Zipfian frequencies and a context-dependent tilt. Only $E$ trains, at batch 128 with gradient clipping at 1, like the paper's runs. The widget below is one Adam step at step 2,000, split by token-frequency rank. Each bar is how far that bin of rows moves its logits for an average context.

<ZeroSumLedger />

The momentum buffer sums to $-5 \times 10^{-8}$ over the vocabulary, which is fp32 noise. The Adam direction sums to 5,373 against an absolute sum of 9,595, so 56% of it is net push in one direction, and 87% of that net push comes from the 1,792 rows ranked past 256 (measured). At a learning rate of 0.01 that is a change of $-0.026$ in the mean logit per step. Over 40,000 steps the FP32 run below drifted to $-1{,}002$, which is the same rate (measured).

### Adam's step, term by term

For one coordinate $(i, j)$ of $E$, with gradient $g_{ij} = (p_i - y_i) h_j$ averaged over the batch, Adam keeps

$$
m_{ij} \leftarrow \beta_1 m_{ij} + (1 - \beta_1) g_{ij},
\qquad
v_{ij} \leftarrow \beta_2 v_{ij} + (1 - \beta_2) g_{ij}^2,
\qquad
\Delta E_{ij} = -\eta \, \frac{\hat m_{ij}}{\sqrt{\hat v_{ij}} + \epsilon}.
$$

Take a rare token $i$ that is almost never the target in a batch. Then $y_i = 0$ at every position and $g_{ij}$ is the batch average of $p_i h_j$: a small number whose sign follows the mean hidden state's $j$-th coordinate, batch after batch. A gradient with the same sign and roughly the same size every step has $|\hat m| \approx \sqrt{\hat v}$, so the ratio is close to $\pm 1$ whatever $p_i$ is. A token at probability $10^{-6}$ moves as fast as one at $10^{-2}$. Every such row steps by about $\eta$ per coordinate, all lowering its logit.

A frequent token is the target often enough that $g_{ij}$ flips sign between batches. Its $\hat m$ averages the flips toward zero while $\hat v$ keeps their squares, so the ratio stays well below one. Frequent rows move slowly and both ways; rare rows move fast and one way. Averaged over the vocabulary the rare rows win, which is the 87% in the ledger above (reasoned, matching the measured ledger).

Two predictions follow, tested below. Halving $\eta$ should roughly halve the drift. And the culprit is the per-row divisor, not the momentum. If every row in a column shares one second moment $s_j$, the update is $\Delta E_{ij} = -\eta\, \hat m_{ij} / s_j$, and the column sums to $-\frac{\eta}{s_j} \sum_i \hat m_{ij} = 0$ again. That is the idea behind Stollenwerk and Stollenwerk's coupled Adam, which the paper cites and does not run.

## Why the drift should not matter

Here is the puzzle. Subtract the mean embedding from every row, $E' = E - \mathbf{1}\mu^\top$. Every logit drops by the same scalar $\mu^\top h$, so $p$ is unchanged, so the gradient with respect to $E$ is unchanged. The gradient into the hidden state is unchanged too:

$$
\frac{\partial \mathcal{L}}{\partial h} = E'^\top (p - y) = E^\top (p - y) - \mu \sum_{i} (p_i - y_i) = E^\top (p - y).
$$

In exact arithmetic, without weight decay, the mean component is invisible to both passes (reasoned, following the paper's Section 3.4). A run with a drifting mean and a run without one should follow the same trajectory. My FP32 toy agrees: the Adam run drifted to a mean logit of $-1{,}002$, and its final loss of 5.038 matches, to within 0.001, the 5.037 of a run that removes the mean after every step, which never drifts at all (measured; that run is the μ-centering one below). So whatever breaks, it is not the drift itself.

## BF16 has no room for differences between big numbers

BF16 keeps 8 bits of significand. Between $2^e$ and $2^{e+1}$ the gap between neighbouring values is $2^{e-7}$: about 0.008 between 1 and 2, and 8 between 1,024 and 2,048. The paper's worked example (reported, Appendix C.1): logits $(0, 1)$ give probabilities of about $(0.27, 0.73)$. Shift both by $-1025$ and both round to $-1024$, so the softmax returns $(0.5, 0.5)$. I checked it in NumPy and got exactly that (measured). Subtracting the maximum before the exponential, which every framework does, can't help: the difference was lost when the GEMM wrote its output.

The relevant quantity is therefore the mean relative to the spread, $r = |\mu_z| / \sigma_z$. While $r$ is small, rounding is small next to the differences that set the probabilities. Once the mean dominates, it isn't. Try it:

<RoundingLab />

To get numbers at a realistic vocabulary size I shifted 64 rows of 50,257 Gaussian logits with standard deviation 2 and measured the KL divergence from the exact softmax after rounding (measured):

| mean shift | BF16 RMS error | BF16 KL | FP16 KL |
|---:|---:|---:|---:|
| 0 | 0.0033 | 2.67e-5 | 4.12e-7 |
| 64 | 0.114 | 2.78e-3 | 4.35e-5 |
| 256 | 0.456 | 4.38e-2 | 6.99e-4 |
| 512 | 0.910 | 0.167 | 2.78e-3 |
| 1,024 | 1.545 | 0.597 | 1.11e-2 |
| 2,048 | 1.930 | 2.23 | 4.38e-2 |

FP16 at a shift of 512 has the same KL as BF16 at 64: three more significand bits buy a factor of eight in headroom, and the table shows it to three digits. By a shift of 1,024 the BF16 RMS error is 1.5 logits on a spread of 2, and the distribution is mostly noise (measured).

### And FP8?

The paper does not test FP8, so I extended the static test (measured). Here the drift lives in the weights, as in training: a centred 50,257 × 128 head plus a shared mean row sized to put the mean logit at the first column's value, applied to 32 RMS-normalised hidden states, logit spread 2. "E4M3" and "E5M2" quantise both GEMM inputs to FP8 with one per-tensor scale each and accumulate in fp32; "BF16 GEMM" rounds inputs and output to BF16; "E4M3 centred" quantises $E - \mathbf{1}\mu^\top$ and adds the mean logit back in fp32.

| mean logit | BF16 GEMM, RMS / KL | E4M3, RMS / KL | E5M2, RMS / KL | E4M3 centred, RMS / KL |
|---:|---:|---:|---:|---:|
| 0 | 0.006 / 3.8e-5 | 0.074 / 2.9e-3 | 0.150 / 1.2e-2 | 0.074 / 2.9e-3 |
| −16 | 0.030 / 1.9e-4 | 0.130 / 4.2e-3 | 0.248 / 1.7e-2 | 0.074 / 2.9e-3 |
| −64 | 0.113 / 3.8e-3 | 0.439 / 3.7e-2 | 0.823 / 0.146 | 0.074 / 2.9e-3 |
| −256 | 0.433 / 8.1e-2 | 1.683 / 0.53 | 3.135 / 1.47 | 0.074 / 2.9e-3 |
| −1,024 | 1.780 / 1.00 | 5.917 / 5.32 | 11.27 / 10.4 | 0.074 / 2.9e-3 |

E4M3 at zero drift already has an error BF16 reaches only between a mean of $-16$ and $-64$, and at $-64$, which the toy's Adam runs pass within about 3,000 steps, its KL is ten times BF16's. An FP8 head would have far less room. And the damage is entirely the mean: centring before quantising holds E4M3 at its zero-drift error at every shift, because otherwise one large shared component sets the per-tensor scale and leaves the per-token differences fewer levels. That is one more reason to keep the head out of FP8 (table measured, advice reasoned).

## The paper's evidence

The cleanest test of the precision story is to hold everything fixed and change only the LM head's precision. The paper does that at 28M parameters, training 16B tokens (about 570 tokens per parameter, roughly 28 times Chinchilla) at peak learning rate 0.1, in three settings: **FP32** (head weights and GEMM in fp32), **BF16 mixed** (fp32 master weights, BF16 GEMM) and **BF16 true** (head weights and optimizer state also stored in BF16). The rest of the network stays in BF16 mixed throughout.

<Figure src="https://media.thesatyajit.com/ai/articles/lm-head-loss-spikes/fig1.png?v=ecdfdde222e8630c" alt="Three panels over 120k steps. Left, eval loss: BF16 true diverges first, near 50k steps; BF16 mixed diverges near 100k; FP32 and BF16 mixed with SGD-M stay near 4.05. Middle, output logit norm on a log scale: every AdamW baseline grows, BF16 true fastest to above 1e8, FP32 to about 4e6; SGD-M stays flat near 2e5. Right, RMS of the BF16 minus FP32 logits on a log scale: the baselines rise from about 1e-2 to above 1, while SGD-M stays flat near 6e-3." caption="Lower LM-head precision brings the time bomb earlier; FP32 survives although its logits grow too, and SGD-M on the head keeps both the logit norm and the quantization error flat. 28M-parameter models at peak learning rate 0.1 (paper, Figure 1)." />

The ordering is the one the precision story predicts. BF16 true fails first, BF16 mixed later, FP32 not at all over the horizon, even though its logit norm grows too (reported). The error metric on the right is the one I copied for the toy: recompute the head in FP32 from the same hidden states and take $\mathrm{RMS}(z_{\mathrm{bf16}} - z_{\mathrm{fp32}})$. In the failing runs it grows by orders of magnitude before the loss turns.

<Figure src="https://media.thesatyajit.com/ai/articles/lm-head-loss-spikes/fig6.png?v=64f9b29f4480dca9" alt="Three panels over 120k steps with five seeds per precision. Left, eval loss: BF16 true runs fail earliest, BF16 mixed later, FP32 stays stable. Middle, RMS logit quantization error on a log scale rising to about 1 for both BF16 settings. Right, KL divergence between FP32 and BF16 predictive distributions rising by several orders of magnitude for BF16 runs." caption="Across five seeds, the order of failure matches the order of quantization error in both the logits and the predictive distribution. 28M-parameter models at peak learning rate 0.1 (paper, Figure 6)." />

Appendix C.2 extends it to the learning-rate sweep: lower precision triggers time bombs at smaller learning rates, and at the largest rate tested every BF16 run fails while FP32 stays stable (reported, Figure 19).

## I reproduced it in a toy

Same toy as above, 40,000 steps per run, one seed, a constant learning rate after a 200-step warmup. "Mixed" rounds the GEMM's inputs and its output logits to BF16 and keeps fp32 master weights; "true" also stores $E$ and the optimizer state in BF16. The softmax and the backward run in float64 on the rounded logits, so the only precision loss is where the paper locates it. The irreducible loss of the task is 4.403 nats.

<TimeBombRuns />

What the runs show (all measured):

- **FP32, Adam.** Mean logit $-1{,}002$ at the end, best loss 4.992, final 5.038. A large drift and no damage.
- **BF16 mixed, Adam.** The same drift, $-983$, and a loss that climbs from its best of 5.006 at step 2,000 to 5.958. The RMS error between the BF16 and FP32 logits grows from 0.0005 to 1.72, and the KL from $10^{-7}$ to 0.90. Evaluated through an FP32 head, the same weights score 5.121. Most of the damage is in the forward pass, and some of it has leaked into the weights through gradients computed from wrong probabilities.
- **Lower learning rate.** At 0.003 instead of 0.01, the drift is $-425$ by the end and the loss rises from 4.820 to 4.987. Slower drift, later and smaller damage, as the paper's sweep has it.
- **SGD-M on the head, BF16 mixed.** Momentum 0.99, head learning rate 0.03 from a short sweep. Mean logit zero to the printed precision, RMS error flat at 0.004 to 0.006, final loss 4.752.
- **μ-centering, BF16 mixed.** Adam, with the mean embedding subtracted after every step. Error flat at 0.011, final loss 5.037, which is the FP32 run's loss.
- **z-loss, BF16 mixed**, $\lambda = 10^{-4}$ as in the paper. It slows the drift to $-584$ but does not stop it in 40,000 steps, and the loss still climbs to 5.574.

Two things in the toy disagree with the paper, and I'd rather say so than tune them away. First, z-loss: at the paper's $\lambda$ it only delays the failure here. At ten times that, $\lambda = 10^{-3}$, the drift stops near $-97$ and the RMS error holds at about 0.16, but the loss settles at 5.214, worse than the FP32 run's 5.038 (measured). The right strength is presumably set by how fast the drift is, which in my toy is fast. Second, BF16 true fails differently. Its drift stalls at $-120$ because, once a head weight is large, BF16 storage rounds away most of each Adam update; the loss rises early, to 5.314 by step 8,000, from lost updates rather than from logit rounding. The paper's BF16 true runs fail first too, but its figures show the quantization error leading. An SGD-M head under BF16 true is not perfectly zero-sum either, since rounding each stored weight is not: its mean logit crept to $+10$ and its error to 0.018, with the loss unchanged at 4.752 (measured).

The toy's SGD-M runs also reach a lower loss than any Adam run (4.752 against 4.992 for FP32 Adam). I would not read anything into that. My trunk is frozen and the head is a linear softmax regression, where SGD is well suited; the paper finds SGD-M "competitive", not better, on real models.

And one finding about my own set-up. I first gave the toy a learnable RMSNorm gain, trained with Adam. The gain absorbed the drift. In a full 40,000-step BF16 mixed run at learning rate 0.01, its norm shrank from 8.0 to about 1.6 while the mean embedding's norm grew to 393, and the mean logit stayed between $-6$ and $-11$. The RMS error reached only 0.027 and the loss ended at 5.107, with no climb (measured). That is why the runs above freeze the trunk. In a real network the hidden-state scale is set by much more than one gain vector, and the paper's runs, which do have learnable gains, still drift.

## The fix: SGD with momentum on the head only

The proposal is to train $E$ with SGD with momentum and everything else with AdamW. Momentum keeps the zero sum, so the mean never drifts, and no explicit centering is needed. In the repository it is one config, `configs/method/sgdm.yaml`:

```yaml
opt:
  lm_head_optimizer:
    type: sgd_momentum
    peak_lr: 0.3
    momentum: 0.99
    nesterov: false
```

`training/optimizers.py` builds it as `optax.sgd` masked onto the output embedding, chained with the AdamW for the rest of the model. The head's learning rate follows the trunk's schedule, scaled to its own peak. The paper tunes that peak separately and raises momentum from 0.9 to 0.99 (reported). For a fair comparison it gives the AdamW head the same tuning budget: tuning the head learning rate alone helps a little, raising the head's $\beta_1$ to 0.99 closes much of the gap, but every AdamW variant, with or without μ-centering or z-loss, still diverges or time-bombs at the largest trunk learning rate, where SGD-M stays stable (reported, Figure 5).

<Figure src="https://media.thesatyajit.com/ai/articles/lm-head-loss-spikes/fig7.png?v=1f810dff1ee31913" alt="Three panels over 120k steps with five seeds per method, all in BF16 mixed. Left, eval loss: the baseline median rises late, with individual seeds spiking; mu-centering, z-loss and SGD-M stay near 4.05. Middle, RMS logit quantization error: the baseline rises to about 1, z-loss to about 1e-1, mu-centering to about 4e-2, SGD-M stays flat near 6e-3. Right, KL divergence: the baseline climbs to about 1e-1 while the mitigations stay near 1e-4 to 1e-3." caption="Three mitigations that act on different things, the loss, the embeddings and the optimizer, all keep the quantization error bounded, and all prevent the time bomb. 28M-parameter models in BF16 mixed at peak learning rate 0.1, five seeds (paper, Figure 7)." />

The 404M learning-rate sweep, 96B tokens each, is where SGD-M earns its place:

<Figure src="https://media.thesatyajit.com/ai/articles/lm-head-loss-spikes/fig4.png?v=9201c3e497023082" alt="Two panels. Left, final eval loss against trunk peak learning rate from 3e-4 to 3e-1 for four methods: the AdamW baseline ranges 2.87 to 2.95 and time-bombs at 1e-2, 1e-1 and 3e-1; z-loss sits around 2.87 to 2.91; mu-centering around 2.86 to 2.89; LM head SGD-M lowest at about 2.81 to 2.83 until 3e-1, where it rises to about 2.87. Right, output logit norm over 700k steps on a log scale: the baseline grows from about 5e5 to above 1e7, z-loss holds near 1e6, mu-centering and SGD-M stay near 2e5." caption="Across trunk learning rates, SGD-M on the head gives the lowest final loss at every rate, and with μ-centering holds the logit norm down; the AdamW baseline time-bombs at three of seven rates. 404M-parameter models, 96B tokens (paper, Figure 4)." />

The table behind the fast-spike version of that sweep (reported, Table 2, final eval loss, † for a run that spiked and did not recover) makes the momentum point concrete:

| 404M, peak learning rate | 0.10 | 0.145 | 0.21 | 0.30 |
|---|---:|---:|---:|---:|
| AdamW | 3.295 | 3.248 | 11.006† | 28.241† |
| μ-centering | 3.172 | 3.221 | 5.144† | 7.115† |
| z-loss | 3.161 | 3.213 | 3.289 | 6.588† |
| SGD-M head, β = 0.9 | 3.335 | 3.304 | 3.268 | 3.244 |
| SGD-M head, β = 0.99 | 3.186 | 3.180 | 3.169 | 3.167 |

At the lowest rate μ-centering and z-loss edge out SGD-M; at the high end they fail and SGD-M does not. β = 0.9 is stable but costs 0.08 to 0.15 nats. The paper credits the momentum's $1/(1-\beta)$ amplification, and my ablation below finds that in the toy this is the whole story.

That figure carries the paper's broader claim: z-loss, μ-centering and SGD-M work for the same reason, because each keeps the head in a range where BF16 can still tell logits apart. Their Figure 7 also shows the ordering my toy shows, with SGD-M's error lowest and flattest. At 1.4B parameters, single seed, SGD-M on the head reaches a lower final training loss than AdamW, μ-centering and z-loss, and a Muon trunk with an SGD-M head reaches the lowest of all (reported, Figure 15).

The paper is also careful about a neighbouring explanation. [Barron, Gnanapragasam and Tschannen's fixed cross-entropy](/articles/softmax-fix) repairs a different rounding bug, in the loss rather than in the head's GEMM. Swapped in here, it delays but does not prevent time bombs under BF16 true; under BF16 mixed no fixed-softmax run fails within 16B tokens, but its logit norm and quantization error keep growing well above the z-loss runs (reported, Appendix C.4). It treats a symptom of large logits; it doesn't stop them getting large.

## What the code does that the paper does not say

I read the training path at commit `76153f0`, about 5,400 lines across `training/`, `model.py` and `diagnostics/`. Most of it matches the paper. What the PDF does not tell you:

- **SGD-M's real step size.** The head uses `optax.sgd(momentum=0.99)`, which keeps a heavy-ball trace $b \leftarrow g + 0.99\,b$, not an exponential average. A steady gradient therefore gets a step of $\eta \cdot g / (1 - 0.99) = 100\,\eta g$. At the config's head peak of 0.3, that is $30 g$ per step. Without Adam's normalisation a rare row's gradient is tiny, which is why the head's peak is three times the trunk's 0.1 and why the paper's appendix finds the 100× amplification of 0.99 matters.
- **No weight decay on an SGD-M head.** `optax.sgd` has none, while an AdamW head would get the run's decay. Moot in the time-bomb configs (decay zero); in the fast-spike configs ($10^{-4}$) the SGD-M head is the only undecayed part. The `adamw_b1` head is built without the decay mask the other branches use.
- **Clipping is global.** `optax.clip_by_global_norm(1.0)` is chained in front of the split optimizer, so the head's gradient and the trunk's share one norm budget. A head with a large gradient shrinks the trunk's step too.
- **μ-centering is exact only in fp32.** `center_output_embeddings` computes the mean in fp32, casts it back to the parameter dtype, then subtracts. Under BF16 true both the mean and the subtraction are rounded, so centering leaves a residual. The residual is a rounding error, not a drift, consistent with the paper's BF16 true μ-centering runs holding up.
- **z-loss strength is hard-coded.** `loss_fn_z_loss` sets `lam = 1e-4` inside the function; `configs/method/z_loss.yaml` only switches it on. The z-loss strength question my toy raised needs a source edit to test.
- **The error metric is clean.** `diagnostics/quantization.py` takes the activation the head actually consumed (already BF16), promotes it and the FP32 master weights to fp32, and recomputes the GEMM at `Precision.HIGHEST`. So $\mathrm{RMS}(z_{\mathrm{bf16}} - z_{\mathrm{fp32}})$ measures the head's own rounding, not the trunk's. It also logs `cross_entropy_difference_native_minus_fp32`, the native loss minus the loss through the FP32 shadow head: the same "gap" metric I use in the experiments below.
- **Knobs the paper does not report.** The loop has an `lm_head_mean_amplification` hook that, from a chosen step, adds `coefficient` times the current mean row to every row after each update (and turns μ-centering off when it starts). That is a direct causal test of the drift story, and no result from it is in the paper. Also undocumented: `lm_head_freeze`, `out_ln_scale_freeze`, a `final_hidden_mean_centering` that centres $h$ instead of $E$, "oblique" head optimizers with a target RMS, AdamC, a $\beta_2$ cosine anneal and an `lm_head_input_grad_fp32` switch.

The amplification hook multiplies whatever mean exists, so on an SGD-M head it does nothing. To run the causal test on a head that never drifts, I injected the drift additively instead.

## Three more experiments

For these I rebuilt the toy from its description, since the first script did not survive: same vocabulary, widths, batch, clipping, warmup and 40,000 steps, NumPy 2.5.3 on CPU. The rebuild is close, not identical: irreducible loss 5.026 nats rather than 4.403, FP32 Adam drift $-910$ rather than $-1{,}002$. Compare numbers within this section. Every run is measured, one seed unless I say otherwise.

### Inject the drift by hand

The paper's evidence for drift, then rounding error, then damage is that the three move together and that precision changes the outcome. A sharper test: take a head that does not drift, add drift and nothing else. The head here trains with SGD-M (momentum 0.99, learning rate 0.03), which keeps the mean at zero to print precision. From step 10,000, after every update, I add the same small vector to every row of $E$, pointed along the mean hidden direction. In exact arithmetic that changes nothing; at the ×1 rate it moves the mean logit about $-0.025$ a step, the Adam runs' own rate. The metric is the **gap**: eval loss through the run's own forward pass minus eval loss of the same weights through an FP32 head, logged every 100 steps. Both terms see the same weights, so optimizer noise cancels.

<DriftInjection />

What happened (all measured):

- **FP32, injected drift.** Nothing. The mean reaches $-737$, the final loss is 5.418 against 5.417 without injection, and the gap never exceeds $4 \times 10^{-6}$.
- **BF16 mixed, injected drift.** The same weights, the same drift, and a time bomb. At the ×1 rate the loss rises from 5.42 to 6.07. The gap passes 0.01 at step 13,000, at a mean of $-74$, and 0.05 at step 16,200, at $-152$. A head the paper's fix had made safe is unsafe again as soon as its mean moves: precision, not the optimizer, turns drift into damage.
- **The rate changes the timing and nothing else.** At half, double and four times the rate the damage arrives at the same mean logit. At a mean of $-100$ the gap is 0.011, 0.014 and 0.017 for ×0.5, ×1 and ×4; at $-300$ it is 0.143, 0.157 and 0.147. Only the number of steps to get there changes: the gap passes 0.05 about 11,900 steps after the drift starts at ×0.5, 6,200 at ×1 and 2,000 at ×4. Switch the widget's x-axis to the mean logit and the curves collapse onto one.

That collapse answers a question I left open in the first version of this article: when does growing error turn into damage? In the toy, at powers of two. Read the ×1 run's error against its mean logit and it is a staircase: about 0.15 for means between 64 and 128, 0.30 between 128 and 256, 0.60 between 256 and 512, 1.2 between 512 and 1,024. Each plateau is BF16's spacing at that magnitude ($2^{e-7}$: 0.5, 1, 2, 4) divided by $\sqrt{12}$, the RMS of uniform rounding error, to two digits: 0.144, 0.289, 0.577, 1.155. The loss gap climbs the same stairs, roughly fourfold per step (about 0.01, 0.04, 0.15, 0.55), because the KL grows with the square of the spacing. The time bomb in this toy is a series of detonations, one per power of two, each about four times worse than the last. Where logits spread over several binades at once, as in a real model, I would expect the stairs to smear into a ramp; the Adam runs below, with wider logits, already look smoother (reasoned).

### How much warning you get

The paper shows the quantization error growing by orders of magnitude before the loss turns. I wanted the number of steps and its spread across seeds, so I ran Adam on the head in BF16 mixed at two learning rates, three seeds each, plus one cosine-decay run, and recorded when each signal crossed a threshold (measured):

| run | RMS err > 0.1 | gap > 0.01 | gap > 0.05 | final loss | FP32 head on same weights |
|---|---:|---:|---:|---:|---:|
| lr 0.01, seed 0 | 2,500 (−61) | 3,500 (−84) | 7,500 (−174) | 6.188 | 5.654 |
| lr 0.01, seed 1 | 2,500 (−61) | 4,500 (−106) | 7,000 (−163) | 6.195 | 5.629 |
| lr 0.01, seed 2 | 2,500 (−61) | 3,500 (−83) | 10,000 (−231) | 6.237 | 5.622 |
| lr 0.005, seed 0 | 4,000 (−61) | 5,000 (−75) | 11,000 (−159) | 6.017 | 5.482 |
| lr 0.005, seed 1 | 4,000 (−61) | 5,000 (−75) | 16,500 (−237) | 5.969 | 5.486 |
| lr 0.005, seed 2 | 4,000 (−61) | 6,000 (−89) | 14,500 (−209) | 6.004 | 5.501 |
| lr 0.01, cosine to 10% | 2,500 (−61) | 5,500 (−128) | 9,500 (−214) | 5.953 | 5.408 |
| lr 0.01, FP16 mixed | 21,500 (−491) | 24,500 (−559) | never | 5.656 | 5.644 |

Steps at which each signal first crossed its threshold, with the mean logit at that moment in brackets, evaluated every 500 steps. "FP32 head" is the final weights evaluated through an FP32 forward pass.

Three things stand out. The first warning is not statistical: the RMS error crosses 0.1 at a mean logit of $-61$ in every BF16 run, whatever the seed, learning rate or schedule, because it is a property of the BF16 grid, not of the training. The first damage, a gap of 0.01 nats, follows 1,000 to 3,000 steps later (mean $-75$ to $-128$), and a gap of 0.05 comes 4,500 to 12,500 steps after the warning.

Second, the final weights are not ruined. Evaluated through an FP32 head, every BF16 run sits within 0.03 of the matching FP32 run (5.644 at learning rate 0.01, 5.473 at 0.005). The damage is almost all in the forward pass, so here you could recover by centring the head, without rewinding. Whether that holds in a real network, whose trunk has trained against wrong probabilities for thousands of steps, the paper does not test.

Third, the loss alone is a poor alarm. A detector in the spirit of the paper's (three consecutive evaluations 0.05 above the running best) fires on the FP32 runs too, because the toy's eval loss wanders that much. The gap is zero to four decimals in every FP32 run. The repository logs it as `cross_entropy_difference_native_minus_fp32`; it is what I would alert on.

The widget plots what the paper's Figure 18 plots: the error divided by the spread of the logits, against $r = |\mu_z| / \sigma_z$. The dashed curves are one-shot rounding of Gaussian logits at the same $r$, computed separately, not fitted.

<ErrorVsRatio />

Every run, Adam or injected, BF16 or FP16, lands on its format's rounding curve once $r$ passes about 10 (measured). At $r \approx 90$ the Adam runs sit at a normalised error of 0.125, and one-shot rounding of Gaussian logits at the same $r$ gives about 0.13; the injected ×1 run at $r = 389$ sits at 0.645, on the same curve. Below $r \approx 10$ the runs sit above it, because there the rounding of the GEMM's inputs, which the curve leaves out, is the larger term. So the paper's Figure 18 relationship is not a trend to fit: above small $r$ it is the BF16 grid, and a run's error is predictable from its mean and spread alone. FP16 sits a factor of eight lower all along, which is why neither FP16 run's loss moved more than 0.02 from its FP32 twin.

### What in the optimizer matters

The paper's tuning study (Figure 5 and Table 2) makes three claims about the head's optimizer: SGD-M needs momentum 0.99 rather than 0.9; raising AdamW's $\beta_1$ to 0.99 closes much of the gap; and the cause is Adam's per-coordinate scaling. The toy can test the shape of each, all in BF16 mixed:

| head optimizer | $\eta$ | effective step $\eta/(1-\beta)$ | final mean logit | final loss | FP32 head on same weights |
|---|---:|---:|---:|---:|---:|
| Adam ($\beta_1$ 0.9, $\beta_2$ 0.95) | 0.01 | n/a | −898 | 6.188 | 5.654 |
| Adam | 0.003 | n/a | −387 | 5.553 | 5.426 |
| Adam, $\beta_1$ 0.99 | 0.01 / 0.003 / 0.001 | n/a | −6,898 / −1,828 / −305 | 48.4 / 30.2 / 14.4 | diverged |
| Adam, $\beta_1$ 0.99, $\beta_2$ 0.999 | 0.01 | n/a | −21 | 5.409 | 5.410 |
| coupled Adam (one $v$ per column) | 0.01 | n/a | 0.0 | 5.616 | 5.615 |
| Adam + μ-centering | 0.01 | n/a | 0.0 | 5.645 | 5.644 |
| SGD-M, β 0.9 | 0.03 | 0.3 | 0.0 | 5.104 | 5.104 |
| SGD-M, β 0.99 | 0.01 | 1 | 0.0 | 5.159 | 5.159 |
| SGD-M, β 0.99 | 0.03 | 3 | 0.0 | 5.417 | 5.417 |
| SGD-M, β 0.9 | 0.3 | 3 | 0.0 | 5.482 | 5.483 |
| SGD-M, β 0.99 | 0.1 | 10 | 0.0 | 6.570 | 6.571 |

All BF16 mixed, 40,000 steps, one seed, measured. For reference, FP32 Adam at 0.01 ends at 5.644 and coupled Adam in FP32 at 5.615.

What the table says, in order of how much I trust it:

1. **The per-row divisor is the cause.** Coupled Adam keeps Adam's momentum, its bias correction and its normalisation, and changes one thing: every row in a column shares a second moment. The mean stays at zero to print precision and the BF16 run ends where the FP32 one does (5.616 against 5.615). This is the term-by-term prediction, and the experiment the paper did not run: you can keep an adaptive optimizer on the head if it stops adapting per token.
2. **For SGD-M, the momentum is a learning-rate multiplier and nothing else.** Matched on $\eta/(1-\beta)$, momentum 0.9 and 0.99 land within 0.07 of each other (5.482 against 5.417 at an effective step of 3), and the best run of all is β 0.9 at an effective step of 0.3. Matched on $\eta$, 0.99 looks far better than 0.9, which is how Table 2 compares them. The paper's own explanation is the $1/(1-\beta)$ factor; in the toy, that factor is the whole effect. A real head's gradient is less stationary, but I would check the effective step before concluding that 0.99 is special.
3. **Adam with $\beta_1 = 0.99$ and $\beta_2 = 0.95$ diverged at every learning rate I tried**, down to 0.001, and in FP32 too (a 400-step probe trained to 5.61 by step 200, then the mean logit fell from $-3$ to $-42$ in the next 200 steps). This is not a precision effect. When $\beta_1 > \sqrt{\beta_2}$ (here $0.99 > 0.975$), a gradient that is shrinking leaves the first moment remembering its old size while the second moment forgets it, and $\hat m / \sqrt{\hat v}$ grows by about $0.99 / 0.975$, 1.5%, per step. A rare token's gradient shrinks precisely because its logit is being pushed down, so the push feeds itself. With $\beta_2 = 0.999$ the same $\beta_1$ is stable and nearly stops the drift ($-21$ by the end). The paper reports that pair working on its head, with a tuned learning rate; my only explanation is that a real head's gradients do not decay as cleanly as a frozen-trunk toy's. I would not use it without a check.
4. **Learning rate.** Halving it from 0.01 to 0.005 cut the drift to 63% ($-565$), not 50%, and 0.003 to 43% ($-387$). Lower rates slow the drift but less than proportionally, so they buy fewer steps than you might hope (measured; the exponent is a property of this toy).

## What each fix costs

The paper's practical comparison (reported, Section 3.4), with my notes on scale (reasoned):

- **FP32 LM head.** The most direct fix and the most expensive. The head is one of the largest GEMMs in the model, with an output dimension equal to the vocabulary, and FP32 runs at a fraction of BF16 throughput. FP32 logits also double the size of what is often the largest activation tensor.
- **SGD-M on the head.** Cheaper than AdamW: one moment buffer instead of two for a $V \times d$ matrix. It needs its own learning rate, and under BF16 true the paper sees it converge to a noticeably higher loss than the other mitigations from early in training (reported, Appendix C.3).
- **z-loss.** One auxiliary term, already in many recipes. It penalises $\log^2 Z$, which acts on the logits, not on the optimizer's drift.
- **μ-centering.** A reduction over the vocabulary rows after each step. In my toy it was the cleanest: same loss as FP32, error flat.
- **Coupled Adam on the head** (Stollenwerk and Stollenwerk's proposal, not tested in the paper). Less memory than AdamW, since $v$ shrinks from $V \times d$ to $d$. In my toy it removed the drift and landed at 5.616 against FP32 Adam's 5.644. If I wanted to keep an adaptive optimizer on the head, it is what I would try first (reasoned from the measured ablation).

Gemma 2's logit soft-capping is mentioned in the related work and not tested.

## Where this sits

The pieces are older than the paper, and it says so. Biś et al. (2021) documented a common shift in Transformer embeddings; Stollenwerk and Stollenwerk (2025) tied output-embedding drift to Adam's preconditioner and proposed coupled Adam; Stollenwerk et al. (2026) proposed μ-centering. Kunstner et al. (2024) argued that Adam beats gradient descent on language models largely through heavy-tailed class imbalance: the same per-row adaptivity that causes the drift, and the tension SGD-M pays for with $\beta = 0.99$. z-loss comes from PaLM (Chowdhery et al., 2023); Wortsman et al. (2024) studied it against output-logit divergence with the elevated-learning-rate small models this paper borrows.

What is new is the link from a harmless drift to a failure through the forward GEMM's rounding, and the evidence that all three known mitigations work by bounding it. It joins a cluster of 2026 papers on rounding that breaks softmax invariants. [Barron, Gnanapragasam and Tschannen](/articles/softmax-fix) and Liu et al. found the cross-entropy gradient rounding to zero at high confidence, behind grokking's slingshot spikes. [Delta-Matching](/articles/delta-matching) covers the attention case, where BF16 in the fused backward breaks the softmax gradient's row sum. This paper is the forward-pass, LM-head member, and its Appendix C.4 shows the members are distinct.

On this site, the [Open-1B](/articles/open-1b-auditable-training) training config carries z-loss at $10^{-4}$, the paper's value.

## What I could not check

- **The paper's training runs.** 28M to 1.4B parameters on FineWeb, on TPUs, B200s, H100s and A100s. Every loss curve on this page that is not in a "measured" widget is reported. I did not run the JAX repository: no JAX here, the smallest config is 16B tokens, and the sandbox does not execute cloned third-party training code. I read it instead.
- **The repository's own causal knob.** `lm_head_mean_amplification` would test the drift story on a real network. My injection experiment is the toy analogue; the real one is a one-line config change for anyone with a TPU.
- **Fast spikes.** The RMSNorm-gain result, the token masking and the batch skipping are all reported.
- **Whether this transfers.** The paper's own future-work paragraph says the spikes are induced by "higher than normal learning rates" and "may not fully transfer to spikes that arise in more natural settings in larger models". My FP8 table is a static rounding test, not training: it shows how much less room an FP8 head would have, not that one would fail.
- **The toy's limits.** A frozen trunk, a linear softmax regression and one seed per configuration except where the lead-time table says three. The rebuilt toy reproduces the first one's behaviour but not its exact numbers. The orderings and the threshold means are what I would trust; the absolute losses are a property of a made-up task.
- **The trigger in a real network.** In the toy the damage arrives in steps, one per power of two in the mean logit. Whether a real model, with a trunk that keeps adapting to wrong probabilities, turns those steps into a self-reinforcing climb, as the paper's late curves suggest, is not something a frozen-trunk toy can show.

## What I would do with this

If you train past compute-optimal in BF16, log three numbers you probably don't log today: the mean output logit, $\mathrm{RMS}(z_{\mathrm{bf16}} - z_{\mathrm{fp32}})$, and the loss through the BF16 head minus the loss through an FP32 copy of it, all on a held-out batch. The last two cost one extra FP32 GEMM per logging step. In both the paper's runs and mine the error climbs long before the loss does, and in mine the loss gap gave thousands of steps of warning without the false alarms the loss itself raises. Watch the mean against powers of two: in BF16, each one it crosses doubles the rounding error. If the mean is drifting, remove it: μ-centering is one subtraction per step, and SGD-M on the head stops the drift at its source. Raising the head's precision to FP32 also works, at a throughput cost.

For the optimizer side of the same story, see [Muon](/articles/muon-optimizer), which the paper pairs with an SGD-M head at 1.4B. For another rounding error that breaks a sum-to-zero invariant, see [Delta-Matching](/articles/delta-matching), where the softmax backward's row sum goes wrong in FP8 attention.