The Role of Quantization in Loss Spikes

Summary (Overview)

• Identifies "time-bomb spikes": A particularly costly form of training instability in language model pretraining characterized by delayed, terminal loss divergences that emerge only after prolonged training, wasting substantial compute.
• Reveals finite precision as a hidden driver: Quantization error in the LM head grows by orders of magnitude before each failure, and the precision of LM head computation directly controls when time-bomb spikes occur.
• Proposes LM head SGD-M mitigation: Replacing AdamW with SGD with momentum exclusively on the LM head suppresses logit growth and improves stability over long training horizons.
• Unifies existing mitigations: z-loss, µ-centering, and LM head SGD-M—despite intervening at different points—all work by suppressing quantization error.
• Demonstrates precision-stability relationship: Lowering LM head precision (BF16 true) accelerates time-bomb spikes, while computing the LM head in FP32 keeps runs stable over the studied horizons.

Introduction and Theoretical Foundation

Background and Motivation

Pretraining Transformer-based language models at scale is notoriously difficult due to objective instability, with loss spikes causing wasted compute and derailed training runs. Large-scale training efforts have relied on heuristic recovery procedures—reverting to earlier checkpoints, skipping data near spikes, and adjusting learning rates—but these treat symptoms rather than underlying causes.

A common class of loss spikes is associated with growth and divergence of output logits in the language modeling (LM) head. Chowdhery et al. (2023) originally identified this behavior and introduced z-loss as a mitigation, which has since become a standard safeguard in many training pipelines.

Theoretical Foundation: The Puzzle of Logit Drift

Prior work (Stollenwerk & Stollenwerk, 2025) traced logit growth to Adam's adaptive preconditioning, which induces persistent drift in the mean output embedding of the LM head. However, the authors identify a critical puzzle: this drift alone cannot substantively affect the forward or backward pass under exact arithmetic.

The key theoretical insight is that shifting all logits by a constant leaves the softmax distribution unchanged:

$$p_i = \mathrm{softmax}(\mathbf{z})_i = \frac{\exp(z_i)}{\sum_{j=1}^{V} \exp(z_j)}, \quad i = 1, \ldots, V$$

The cross-entropy gradient with respect to the LM head has a zero-sum structure along the vocabulary axis:

$$\frac{\partial \mathcal{L}_{CE}}{\partial \mathbf{E}} = (\mathbf{p} - \mathbf{y})\mathbf{h}^{\intercal}, \qquad \sum_{i=1}^{V} \frac{\partial \mathcal{L}_{CE}}{\partial E_{i,j}} = h_j \underbrace{\sum_{i=1}^{V} (p_i - y_i)}_{=0} = 0$$

This zero-sum property means that in exact arithmetic, mean logit drift is harmless—it changes neither predictions nor gradients. The paper resolves this paradox by showing that finite precision computation provides the mechanism connecting harmless drift to catastrophic training failure.

Methodology

Experimental Setup

• Models: 28M–1.4B parameter models trained at elevated learning rates to reproduce loss spikes
• Architecture: Modernized GPT-2 with rotary positional embeddings, GELU activations, RMSNorm, untied embeddings, no biases
• Framework: JAX and Flax, trained on TPUs (v6e/v4) and NVIDIA B200 GPUs
• Optimizer: AdamW (ε=10⁻⁸, β₁=0.9, β₂=0.95) with linear warmup, cosine decay, gradient clipping at global norm 1
• Data: FineWeb with GPT-2 tokenizer
• Precision: BF16 mixed precision by default (FP32 master weights, BF16 runtime casts)

Key Experimental Configurations

1. Fast spike experiments (Section 2.4): 162M models with parameterless vs. parameterized RMSNorm
2. Time-bomb spike experiments (Section 3): 404M models trained for 96B tokens (~12× Chinchilla-optimal), learning rates {3×10⁻⁴, 10⁻³, 3×10⁻³, 10⁻², 3×10⁻², 10⁻¹, 3×10⁻¹}
3. Precision experiments (Section 3.4): 28M models trained for 16B tokens (~28× Chinchilla-optimal), varying LM head precision across FP32, BF16 mixed, and BF16 true

Quantization Error Measurement

At each logging step, the authors recompute LM head logits in FP32 from the same hidden states used by the native forward pass, yielding FP32 reference logits $\mathbf{z}_{fp32}$. They measure:

• Logit quantization error: RMS distance $\mathrm{RMS}(\mathbf{z}_{bf16} - \mathbf{z}_{fp32})$
• Predictive-distribution error: KL divergence $D_{KL}(p_{fp32} \| p_{bf16})$

Empirical Validation / Results

1. Fast Spikes Depend on Normalization Type

Parameterized RMSNorm gains substantially reduce fast loss spikes compared to parameterless norms. Across five matched seed pairs, learnable gains stabilized training and prevented fast spikes.

2. Time-Bomb Spikes Emerge During Overtraining

Two 404M parameter models followed nearly identical loss trajectories through most of training, but the higher learning rate run eventually stopped improving and underwent sustained, monotonic loss increase—a time-bomb. Higher learning rates and longer training horizons amplify logit growth.

3. LM Head SGD-M Prevents Time-Bomb Spikes

Replacing AdamW with SGD with momentum (momentum 0.99) exclusively on the LM head:

• Suppresses logit norm growth
• Maintains competitive final evaluation loss across learning rate sweeps
• Remains stable even at the largest trunk learning rates where all AdamW variants diverge

Tuning AdamW's head-specific learning rate alone provided only modest gains; increasing β₁ to 0.99 closed much of the optimization gap but did not consistently prevent time-bomb instability.

4. Finite Precision Accelerates Time-Bomb Spikes

Figure 6 shows that reduced precision settings develop substantial numerical error and time-bomb first, with quantization error growing by orders of magnitude before each failure.

5. All Mitigations Suppress Quantization Error

The unmitigated baseline exhibited time-bomb loss spikes accompanied by significant rise in quantization error. In contrast, µ-centering, z-loss, and LM head SGD-M all:

• Prevented time-bomb failure
• Kept both error metrics (logit quantization error and KL divergence) bounded across five seeds

6. Theoretical Analysis: Why Mean Drift Matters in Finite Precision

In exact arithmetic, centering the output embeddings (removing the vocabulary-mean component) changes neither predictions nor gradients:

$$\frac{\partial \mathcal{L}_{CE}}{\partial \mathbf{h}} = \mathbf{E}^{\intercal}(\mathbf{p} - \mathbf{y}) = \mathbf{E}^{\intercal}(\mathbf{p} - \mathbf{y}) - \boldsymbol{\mu} \sum_{i=1}^{V}(p_i - y_i) = \mathbf{E}^{\intercal}(\mathbf{p} - \mathbf{y})$$

However, in finite precision, a large logit mean implies an increase in the absolute numerical scale of individual logit values, reducing the precision with which differences between logits (which determine the predictive distribution) can be represented.

Theoretical and Practical Implications

Theoretical Implications

1. Resolves the drift paradox: The paper explains why logit drift—harmless in exact arithmetic—destabilizes training in practice: it amplifies finite precision error.

2. Unified mechanism for mitigations: z-loss (loss objective), µ-centering (output embeddings), and LM head SGD-M (optimizer) intervene at different points but share a common mechanism: suppressing quantization error growth.

3. Precision as a control variable: The precision of LM head computation directly controls when time-bomb spikes occur, establishing finite precision error as a primary driver of this instability class.

Practical Implications

• LM head SGD-M is a cheap, effective fix: Replaces AdamW only for output embeddings, stores one fewer optimizer moment buffer (reducing memory), and retains BF16 mixed precision with little overhead.
• FP32 LM head is most direct but expensive: FP32 GEMMs run at a fraction of BF16 throughput on modern accelerators, and materializing FP32 logits doubles memory for the largest activation tensor.
• Mitigations avoid precision costs: z-loss, µ-centering, and SGD-M keep the LM head in a regime where reduced precision suffices without paying the cost of higher precision.

Conclusion

The paper demonstrates that finite precision error is a hidden driver of time-bomb spikes in language model pretraining. As the logit mean grows from Adam's preconditioner-induced drift, quantization error in the LM head grows by orders of magnitude before each failure. The precision of LM head computation directly controls when failures occur.

The proposed mitigation—replacing AdamW with SGD with momentum on just the LM head—prevents mean drift at its source and stabilizes training across learning rates where unmitigated runs diverge. Existing mitigations (z-loss, µ-centering) can be understood as different ways of keeping the LM head within a regime where finite precision suffices.

Future Directions

• Testing whether findings transfer to spikes arising in more natural settings in larger models
• Building specialized optimizers for the LM head
• Evaluating mitigations in even lower precision formats such as FP8