Researchers link LM-head quantization to loss spikes after billions of tokens

• Billions of apparently healthy training tokens can precede terminal loss divergence driven by LM-head quantization error, according to the researchers' announcement.
• A targeted optimizer change prevented the failures in the reported experiments: SGD with momentum on the LM head, with AdamW retained elsewhere.
• Broader applicability remains open: a follow-up calls for validation in other settings.

AdamW can push the vocabulary logits together toward extreme values, even though softmax ignores a common offset in exact arithmetic, according to the paper. The released experiment configurations can switch only the final normalization and LM head to FP32 while leaving the transformer trunk in BF16.

The researchers call these delayed, terminal divergences “time-bomb spikes.” In a long-horizon experiment, they trained 404M-parameter models for 96B tokens, approximately 240 tokens per parameter and 12 times the Chinchilla-optimal token budget.

Two runs with peak learning rates of 1e-3 and 1e-2 followed similar loss trajectories through much of training. The higher-learning-rate run eventually developed sustained increases in evaluation loss alongside accelerating output-logit growth.

Delayed loss divergence and the optimizer intervention

AdamW's adaptive scaling creates the failure chain described in the paper's analysis:

1. Mean embedding drift. Cross-entropy gradients sum to zero across the vocabulary axis. Adam's coordinate-wise rescaling breaks that property, allowing the mean output embedding to drift.
2. Growing logit offsets. The embedding drift moves the mean output logit toward extreme values. A common offset would leave softmax probabilities unchanged in exact arithmetic.
3. Lost numerical resolution. BF16's spacing between representable values grows with magnitude. Large offsets reduce the resolution available for differences between logits, which determine the probabilities.

In the 28M-parameter precision ablation, BF16 computation with BF16-stored head weights failed first; BF16 computation with FP32-stored weights failed later. FP32 remained stable throughout the studied horizon despite continued logit growth, while quantization error grew by orders of magnitude before the reduced-precision failures.

The fix is pleasingly surgical: the researchers replace AdamW with SGD with momentum exclusively on the LM head.

• Transformer trunk: retains AdamW.
• LM head: uses SGD-M, with momentum 0.99 in the reported comparisons.
• Head learning rate: is tuned independently of the trunk learning rate.

SGD-M preserves the gradient's zero-sum structure without Adam's coordinate-wise adaptive rescaling. In the experiments, it kept logit norm and quantization error nearly flat and prevented the delayed divergence.

The paper also compares two existing interventions:

• z-loss: penalizes growth of the log-partition function.
• μ-centering: subtracts the vocabulary-axis mean from output embeddings after each optimizer step.

Both also bounded quantization error, suggesting a shared numerical mechanism behind otherwise different mitigations.

The study's models span 28M to 1.4B parameters and use FineWeb, a GPT-2 tokenizer, untied embeddings, and QK-Norm. The delayed-spike experiments use learnable RMSNorm gains, no weight decay, and a batch size of 131,072 tokens; the authors report that smaller batches brought onset forward and weight decay delayed it.

Micah Goldblum, a coauthor of the paper, described the result as holding “in our experiments” and expressed interest in broader validation.

For a separate comparison the team had not investigated, he said it would require models with the same FP32 training loss.

Fast, transient spikes receive a separate treatment in the paper: adding learnable RMSNorm gains substantially reduced them across five matched seed pairs. Delayed time-bomb instability persisted with those gains enabled.

The released repository includes training code, optimizer configurations, and diagnostic tools. Its experiment guide exposes three LM-head precision modes:

• bf16_true: BF16 computation and stored weights.
• bf16_mixed: BF16 computation with FP32-stored weights.
• fp32_lm_head: FP32 final normalization and LM head, with a BF16 trunk.

Setting log_lm_head_quantization_error_metrics=true computes native-precision and FP32-shadow logits from the same hidden states, then logs their differences under lm_head_quantization/. The repository also supplies separate sweep scripts for precision, mitigation methods, and long-horizon learning rates.