The Role of Quantization in Loss Spikes

Zaiqian Chen; Aditya Somasundaram; Sophie Li; Sean McLeish; John Kirchenbauer; Tom Goldstein; Micah Goldblum

cs.LGOct 28, 2026

The Role of Quantization in Loss Spikes

Authors: Zaiqian Chen, Aditya Somasundaram, Sophie Li, Sean McLeish, John Kirchenbauer, Tom Goldstein, Micah Goldblum

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 $0.1$ for ${\sim}28\times$ 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, $\mu$ -centering, $z$ -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, $\mu$ -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 $\beta_{1}$ 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 $\beta_{1}=0.99$ closes much of the gap to LM head SGD-M. Adding $\mu$ -centering or $z$ -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 $0.1$ , 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 $0.1$ , 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

Appendix

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

610.00

12.73

. ␣ denotes a leading space.

Figure 10: A single-token context swap collapses the logit distribution at a spike. Full-vocabulary logits (

V=50{,}260

) 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

17\times

, from

121.87

6.87

, 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.

\mu

-centering and

z

-loss diverge at some

\eta\in[10^{-1},3\times 10^{-1}]

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.

162M
	Peak learning rate $\eta$ (final eval loss $\downarrow$ )
Method	$10^{-1}$	$1.2{\cdot}10^{-1}$	$1.45{\cdot}10^{-1}$	$1.75{\cdot}10^{-1}$	$2.1{\cdot}10^{-1}$	$2.5{\cdot}10^{-1}$	$3{\cdot}10^{-1}$
AdamW	3.642	3.679	3.733	3.733	$12.552^{\dagger}$	$23.223^{\dagger}$	$10.845^{\dagger}$
$\mu$ -centering	3.637	3.656	3.674	3.701	$4.094^{\dagger}$	$6.063^{\dagger}$	$5.950^{\dagger}$
$z$ -loss	3.636	3.655	3.690	3.718	$5.594^{\dagger}$	$5.964^{\dagger}$	$7.495^{\dagger}$
row-centering	3.617	3.631	3.633	3.718	3.697	3.793	3.863

Appendix

Table 2: Learning rate sweep across the spike-prone regime (

\eta\in[10^{-1},\,3{\times}10^{-1}]

) for 162M and 404M models. Final evaluation loss (

\downarrow

lower is better) is reported.

\dagger

marks runs whose loss spiked during training and did not recover.

	162M		404M
Method	Eval Loss $\downarrow$	Spike Score $\downarrow$	Eval Loss $\downarrow$	Spike Score $\downarrow$
Baseline (AdamW)	$3.702\pm 0.014$	$0.088\pm 0.029$	$3.916\pm 0.278$	$0.085\pm 0.029$
$z$ -loss	$3.644\pm 0.006$	$0.064\pm 0.027$	$3.214\pm 0.008$	$0.106\pm 0.040$
$\mu$ -centering	$3.646\pm 0.009$	$0.066\pm 0.043$	$3.208\pm 0.004$	$0.064\pm 0.002$
Row-centering	$\mathbf{3.621\pm 0.016}$	$0.038\pm 0.007$	$4.858\pm 1.438$	$0.056\pm 0.011$
SGD-M, $\beta=0.9$	$3.855\pm 0.003$	$0.038\pm 0.004$	$3.301\pm 0.003$	$0.060\pm 0.002$

Appendix

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

\eta=0.15

for 162M and 404M models. Values are mean

\pm

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

\eta=0.07

(single seed). Replacing AdamW with SGD-M (

\beta=0.99

) on the LM head reaches a lower final training loss than AdamW,

\mu

-centering, and

z

-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

3\times 10^{-4}

3\times 10^{-1}

. 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,

\mathrm{RMS}(z_{\mathrm{bf16}}-z_{\mathrm{fp32}})/\mathrm{std}(z_{\mathrm{fp32}})

, increases with the relative logit mean,

|\mathrm{mean}(z)|/\mathrm{std}(z)

. 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

0.01

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 $0.1$ , 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). $\mu$ -centering, $z$ -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.

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 353 papers

Sep 10, 2026

Why Does Post-Training Quantization Work?

Post-training quantization compresses large language models (LLMs) by storing their weights at reduced precision, and each quantized weight introduces an error into the hidden states. Naively, these errors should accumulate with depth and corrupt next-token prediction; randomly initialized models accumulate these discrepancies rapidly, whereas quantized pretrained models accumulate much less hidden-state error and largely maintain downstream task performance, even though they were never trained with quantization noise. This raises the question we address: why does post-training quantization work? Comparing full-precision and quantized forward passes, we identify two mechanisms that characterize pretrained quantization robustness. First, the error a layer newly introduces tends to oppose the error it inherits from the layer's input. The two cancel partially such that the discrepancy between full-precision and quantized passes grows slowly. This counteracting residual interaction develops during pretraining. Our quantitative analysis identifies it as a major factor slowing hidden-error growth. Second, LM-head geometry preferentially preserves the scores and probabilities of high-ranked tokens, which typically represent the model's most confident predictions. Together, these mechanisms explain why quantization error that passes through numerous layers can still produce only small output changes, and we verify the findings across models and quantization settings.

Yuxiang Chen, Michael Beyer, Jun Zhu +1

Dept. of Comp. Sci. and Tech., Institute for AI, BNRist Center, THBI Lab, Tsinghua-Bosch Joint ML Center, Tsinghua University · College of AI, Tsinghua University · Bosch AI Research, Renningen, Germany

LLM Quantization Quantization Robustness Post-Training Quantization Neural Network Quantization

Ask PDF

May 11, 2026

Pretraining large language models with MXFP4 on Native FP4 Hardware

Why does full-pipeline FP4 training of large language models often diverge, even when forward activations and activation gradients remain stable? We address this question through a controlled study of MXFP4 quantization in transformer training, progressively enabling FP4 across forward propagation (Fprop), activation gradients (Dgrad), and weight gradients (Wgrad) while holding all other factors fixed. In full pretraining of Llama 3.1-8B on the C4 dataset, we observe that quantizing Wgrad is the primary driver of convergence degradation, whereas FP4 in Fprop and Dgrad alone introduces only modest additional token requirements. To interpret this behavior, we evaluate both structured and stochastic interventions under a controlled experimental setting. We find that stochastic rounding and randomized Hadamard rotations fail to stabilize training once Wgrad is quantized, whereas deterministic Hadamard rotations consistently restore stable optimization. These results suggest that FP4 training instability is driven by structured micro-scaling errors along sensitive gradient paths, rather than by insufficient stochasticity. We run experiments with native MXFP4 support on AMD Instinct MI355X GPUs, enabling controlled investigation of these effects without reliance on software emulation.

Musa Cim, Poovaiah Palangappa, Miro Hodak +3

1The Pennsylvania State University · 2Advanced Micro Devices, Inc.

Language Model Pretraining LLM Quantization Walsh-Hadamard Transform

Ask PDF

Jun 1, 2026

Massive Spikes in LLMs are Bias Vectors: Mechanistic Uncovering and Spike-Free Quantization

Massive activation spikes in Large Language Models (LLMs) severely degrade quantization by stretching dynamic ranges. While prior hypotheses characterize these as high-level scalar biases, we argue that they are merely the scalar intermediates of rigid, structural vector biases in the spike-carrying tokens. We show that these tokens converge to constant vectors after normalization that drive the attention sink and value-state drain mechanisms. We geometrically substantiate this by analyzing the coordination of projection weights:

W_K

contrastively amplifies the vector,

W_Q

aligns semantic tokens toward it, and

W_V

projects it into the spectral null-space. Furthermore, we reveal that the model actively preserves these structural biases against Rotary Positional Embedding (RoPE) perturbations by localizing them in "zones of rotational stability" utilizing low-frequency bands and coherent channel pairs. Leveraging this, we propose INSERTQUANT, a post-training quantization (PTQ) framework that clamps spikes and restores their function via pre-computed template vectors. This renders activations strictly spike-free, enabling robust low-bit quantization with high fidelity. INSERTQUANT achieves parity with state-of-the-art per-tensor quantization methods on LLMs and uniquely generalizes beyond text to other modalities such as ViTs.

Yung-Chin Chen, Chung Peng Lee, Ze-Wei Liou +1

Princeton University, NJ, USA · EnCharge AI, CA, USA

LLM Quantization Activation Outliers Activation Quantization Attention Sinks Post-Training Quantization

Ask PDF

Abstract

LLM Training

Explore similar work

Why Does Post-Training Quantization Work?

Pretraining large language models with MXFP4 on Native FP4 Hardware

Massive Spikes in LLMs are Bias Vectors: Mechanistic Uncovering and Spike-Free Quantization