cs.LGOct 28, 2026
SaveThe Role of Quantization in Loss Spikes
Organizations: Columbia University · Mohamed bin Zayed University of Artificial Intelligence · University of Maryland · Vector Institute
Abstract
Loss spikes in language model pretraining can waste substantial compute and derail otherwise healthy training runs. A primary driver of these instabilities is the growth and divergence of output logits in the final language modeling (LM) head. In this work, we study a particularly costly form of this instability associated with output logit growth: delayed, terminal divergences that emerge only after prolonged training, which we call time-bomb spikes. We find that these instabilities are preceded by rapidly growing quantization error in the LM head, which eventually culminates in a loss spike. Crucially, reducing numerical precision accelerates time-bomb spikes, while increasing precision improves stability. We then propose a simple yet effective mitigation that replaces AdamW on the LM head with SGD with momentum, suppressing logit growth and improving stability over long training horizons. Finally, we find that existing mitigation techniques from the literature (along with our own) all reduce quantization error, suggesting a shared numerical mechanism behind their effectiveness. Together, our results identify finite precision error as a hidden driver of time-bomb spikes and show that suppressing this error is key to stable long-horizon pretraining.
Figures & tables
Figure 1: Reduced precision accelerates time-bomb spikes, while LM head SGD-M prevents them. Lowering LM head precision brings the failure earlier, with BF16 true diverging first, BF16 mixed later, and FP32 remaining stable throughout training ( left ). Output logit norm grows under every AdamW baseline, including FP32 ( middle ), but it destabilizes training only in reduced precision, where LM head quantization error grows by orders of magnitude before each failure ( right ). Quantization error is measured as the RMS difference between native BF16 logits and FP32 logits computed from the same hidden states. Replacing AdamW with SGD-M on the LM head keeps both logit norm and quantization error nearly flat, preventing the time bomb even in BF16 mixed precision. Results are for 28M parameter models trained at peak learning rate for the Chinchilla optimal token horizon, using AdamW throughout (baselines) or SGD with momentum on the LM head (SGD-M).
Figure 2: Learnable gain parameters in normalization layers stabilize training and can prevent fast loss spikes. Comparing a 162M spiking run with the same seed and configuration only differing in the inclusion of learnable RMSNorm parameters, the parameterized norm configuration exhibits no fast spikes, output logit swings, or large logit norm growth.
Figure 3: Long horizon training can result in time-bomb instability. Two 404M parameter models train similarly for most of pretraining, but the higher LR run exhibits delayed, unrecoverable loss increases accompanied by accelerating output logit growth.
Figure 4: SGD-M is competitive in both optimization quality and control of LM head growth. Across trunk learning rates, -centering, -loss, and LM head SGD-M improve stability relative to the AdamW baseline, while differing in both final evaluation loss (left) and the degree of LM head growth suppression (right). Across high learning rates, -centering and LM head SGD-M strongly suppress output logit norm growth, whereas the baseline develops the largest late growth and corresponding time-bomb spikes (right). Faint curves show individual learning rates; solid curves and shaded regions show the median and interquartile range across learning rates.
Figure 5: Increasing LM head is more effective than tuning its learning rate alone. Left: Final evaluation loss across trunk learning rates. Head learning rate tuning yields modest gains, while closes much of the gap to LM head SGD-M. Adding -centering or -loss improves some settings, but all AdamW variants diverge or time-bomb at the largest trunk learning rate, whereas SGD-M remains stable. Dashed lines denote tuned head learning rates. Right: output logit norm trajectories, showing how each method constrains LM head growth.
Figure 6: Reduced precision in the LM head accelerates time-bomb instability and increases quantization error as logits grow. For 28M parameter models trained at peak LR , BF16 true runs fail earliest, BF16 mixed runs deteriorate later, and FP32 remains stable across five seeds (left). The same ordering appears in the quantization error between BF16 native and FP32 reference logits (middle) and in the KL divergence between the corresponding predictive distributions (right), indicating that loss spikes coincide with increasingly large finite precision error. Solid lines show the median across seeds and faint lines show individual runs.
Figure 7: Methods that prevent time-bomb failure suppress quantization error. For 28M-parameter models trained in BF16 mixed at peak LR , the unmitigated baseline develops time-bomb spikes (left) together with orders-of-magnitude growth in logit quantization error (middle) and predictive distribution quantization error (right). All mitigation methods prevent the time bomb and keep both numerical-error metrics bounded across five seeds. Solid lines show medians and faint lines show individual runs. The same effect is more pronounced under BF16 true training; the unmitigated baseline fails catastrophically after significant growth in numerical error ( Figure 20 ).
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Learning rates can control the type of training instability elicited . Example training runs at 1.4B scale trained for 10B tokens. We focus on the learning rates similar to the middle one in order to elicit fast loss spikes to analyze.
Figure 9: Skipping a small fraction of high-loss tokens eliminates loss spikes but hurts convergence. A 162M model trained at the Chinchilla-optimal token budget with peak learning rate 0.15. One run is the baseline; the other skips high-loss tokens identified via a rolling window threshold.
| Original context (spiking) | Swapped context (recovered) | |||
| Rank | Token | Logit | Token | Logit |
| 1 | ible | 782.05 | ␣the | 95.01 |
| 2 | ments | 669.86 | ␣in | 94.12 |
| 3 | ured | 640.67 | ␣however | 93.80 |
| 4 | ive | 603.56 | ␣until | 92.99 |
| 5 | ivism | 595.26 | ␣and | 92.96 |
Table 1: A single-token context swap collapses the spike. Top-5 predicted tokens at the Rank-1 spike event before and after replacing the final context token Command with a comma. Before the swap, the model assigns extreme logits to unrelated subword tokens; after the swap, the top predictions are ordinary function words with logits in a normal range, and the loss falls from to . ␣ denotes a leading space.
Figure 10: A single-token context swap collapses the logit distribution at a spike. Full-vocabulary logits ( ) at the Rank-1 spike event before (blue) and after (orange) replacing the final context token Command with a comma, with the target token ery held fixed. The standard deviation contracts , from to , recovering the range typical of non-spiking steps.
Figure 11: Fast loss spikes coincide with sharp swings in the output logit mean. A spiking 162M run compared with a non-spiking run using the same configuration and differing only in random seed.
Figure 12: Skipping high-loss batches defers rather than eliminates fast spike instability. Following Karpathy (2024) , we skip any batch whose loss exceeds 7 standard deviations from a rolling 1000-step window and revert the optimizer state. Reverting to either the previous or second-to-previous optimizer state does not prevent subsequent fluctuations in output logit mean characteristic of spike events.
Figure 13: Parameterized RMSNorms effectively mitigate fast loss spikes. The five parameterized runs are highly consistent and overlap closely, indicating low seed-to-seed variability relative to the parameter-less runs.
Figure 14: LM head SGD-M is the only intervention stable across the spike-prone learning rate regime at both model sizes. Final evaluation loss versus peak learning rate for 162M and 404M models. -centering and -loss diverge at some at both sizes, and row-centering, while stable at 162M, diverges inconsistently at 404M and is omitted there. LM head SGD-M remains stable throughout and achieves the lowest loss at the high end.
| Peak learning rate (final eval loss ) | |||||||
| Method | |||||||
| 162M | |||||||
| AdamW | 3.642 | 3.679 | 3.733 | 3.733 | |||
| -centering | 3.637 | 3.656 | 3.674 | 3.701 | |||
| -loss | 3.636 | 3.655 | 3.690 | 3.718 | |||
| row-centering | 3.617 | 3.631 | 3.633 | 3.718 | 3.697 | 3.793 | 3.863 |
Table 2: Learning rate sweep across the spike-prone regime ( ) for 162M and 404M models. Final evaluation loss ( lower is better) is reported. marks runs whose loss spiked during training and did not recover.
| 162M | 404M | |||
|---|---|---|---|---|
| Method | Eval Loss | Spike Score | Eval Loss | Spike Score |
| Baseline (AdamW) | ||||
| -loss | ||||
| -centering | ||||
| Row-centering | ||||
| SGD-M, | ||||
Table 3: Tuned SGD-M on the LM head suppresses spikes and remains competitive on evaluation loss. Final evaluation loss and training-loss spike score (defined in Section A.6 ) at peak learning rate for 162M and 404M models. Values are mean standard error across 5 seeds; bold marks the lowest evaluation loss for each model scale.
Figure 15: LM head SGD-M also stabilizes training at 1.4B scale and pairs well with Muon. Training loss curves for 1.4B models at peak learning rate (single seed). Replacing AdamW with SGD-M ( ) on the LM head reaches a lower final training loss than AdamW, -centering, and -loss, and combining a Muon backbone with LM head SGD-M reaches the lowest loss of any configuration shown.
Figure 16: Time-bomb instability emerges across the long-horizon learning rate sweep alongside accelerating output logit growth. Evaluation loss (left), output logit norm (middle), and output logit mean (right) for 404M-parameter models trained for 96B tokens across peak learning rates from to . Runs train similarly for much of pretraining, while higher-learning rate settings exhibit substantially greater late stage logit norm growth and, in several cases, deterioration in evaluation loss.
Figure 17: Longer training horizons are associated with increasingly severe late stage LM head growth. Evaluation loss (left), output logit norm (middle), and output logit mean (right) for 28M parameter models with parameterized RMSNorm, and peak LR of 0.1 trained with horizons of 4B, 8B, 12B, and 16B tokens. Extending training exposes progressively larger output logit norm growth and late stage loss deterioration, consistent with time-bomb instability becoming more prominent during prolonged training.
Figure 18: Normalized logit error grows as the logit mean becomes large relative to variation across logits. Across BF16 true and BF16 mixed baseline runs, the normalized RMS error between native BF16 logits and FP32 reference logits, , increases with the relative logit mean, . Colors indicate precision, markers indicate peak learning rate, and lines trace individual training trajectories.
Figure 19: Lower precision triggers time-bomb spikes at smaller learning rates. Fraction of seeds exhibiting a time-bomb instability across peak learning rates and numerical precision settings. As precision decreases from FP32 to BF16 mixed to BF16 true, instability emerges at progressively smaller learning rates. A time-bomb is defined as three consecutive post-midpoint evaluations at least above the running-best evaluation loss.
Figure 20: Methods that prevent time-bomb spikes also suppress quantization error. For 28M-parameter models trained in BF16 true at peak learning rate , the unmitigated baseline develops severe unrecoverable loss spikes (left) together with orders-of-magnitude growth in logit quantization error (middle) and predictive distribution quantization error (right). -centering, -loss, and LM head SGD-M prevent the time bomb and keep both numerical-error metrics bounded across five seeds. Solid lines show medians and faint lines show individual runs.
More in this topic
LLM Training
Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.
See all 353 papers