Summary (Overview)

• This paper introduces Power-Law Random Feature (PLRF) models—simple quadratic models with heavy-tailed, power-law distributed inputs—that can replicate the mysterious "training spikes" observed during LLM pre-training.
• The authors develop a theoretical framework and simulation algorithm that aligns each DNN layer with a Random Feature model, enabling lightweight simulation of per-layer training dynamics.
• Key finding: the power-law exponent $\alpha$ of the data distribution is critical; a larger $\alpha$ (heavy-tailed data) is a necessary condition for training spikes to occur.
• The paper demonstrates that stochastic gradient noise (captured by 2nd-order Taylor expansion) is the driving mechanism behind spikes, not 3rd-order terms as previously believed for full-batch training.
• In a case study on a 1.3B parameter Transformer LM, the simulation successfully predicts spiky behavior from a given checkpoint, and the authors derive a $\beta_2$ schedule for Adam that avoids spikes and empirically improves performance.

---

Introduction and Theoretical Foundation

Background and Motivation

Training spikes—temporary, sharp surges in loss during LLM pre-training—are a well-documented empirical phenomenon (Chowdhery et al., 2023; Molybog et al., 2023). These spikes are considered pathological because they may indicate:

• Training instability
• Sub-optimal hyperparameters
• Reduced final model predictive power
• Numerical issues in low-precision arithmetic

A crucial observation is that training spikes are sensitive to both data shuffling AND model initialization, unlike normal loss fluctuations which remain consistent under identical data shuffles regardless of initialization (Figure 1). This extra randomness complicates the study and mitigation of training spikes.

Key Theoretical Foundation: The Random Feature Model

The paper builds on the Random Feature model with quadratic loss:

$$\ell[\mathbf{x}; \mathbf{w}] \stackrel{\text{def}}{=} \frac{c}{2}\left(\mathbf{w}^{\top}\mathbf{A}\mathbf{x} - b[\mathbf{x}]\right)^{2}.\tag{2.1}$$

Here, $\mathbf{w}$ is the $D$-dimensional trainable weights, $\mathbf{A}$ is a fixed $D \times V$ random matrix, $b[\mathbf{x}]$ is a scalar label function, and $c$ is a constant.

The Power-Law Random Feature (PLRF) model assumes inputs $\mathbf{x} = (i^{-\alpha}x_i)_{1 \leq i \leq V}$ where $x_i \sim \mathcal{N}(0,1)$ i.i.d., and $\alpha$ is the power-law exponent of the input data.

The Divergence Threshold

For SGD with small batch size ($B \ll \sqrt{D}$), the learning-rate threshold $\gamma$ of divergence is given by:

$$\frac{1}{B}\sum_{i} \frac{\gamma\lambda_{i}}{2 - \gamma\lambda_{i}} = 1,\tag{2.2}$$

where $\{\lambda_i\}$ are the eigenvalues of the Hessian. In particular:

$$\gamma < \min\left(\frac{2B}{\sum_i\lambda_i}, \frac{2}{\lambda_{\max}}\right).$$

Key insight: Unlike full-batch gradient descent (where the threshold is simply $2/\lambda_{max}$), stochastic optimization has a different divergence threshold due to gradient noise—this is where spikes originate.

---

Methodology

Layer-wise Quadratic Loss Framework

The authors define a layer (Definition 3.2) as a subset $\mathbf{W} \subset \mathcal{W}$ of model parameters such that:

1. $\mathbf{W}$ acts as matrix multiplication: $y = \mathbf{W}h$
2. The forward pass to input $h$ does not depend on $\mathbf{W}$
3. The backward pass to output gradient $\frac{\partial\ell}{\partial y}$ does not depend on $\mathbf{W}$

For such a layer, the per-example gradient is:

$$\frac{\partial\ell}{\partial\mathbf{W}}[\mathbf{x};\mathcal{W}] = \mathbf{h}[\mathbf{x};\mathcal{W}^{\mathrm{F}}]\frac{\partial\ell}{\partial\mathbf{y}}[\mathbf{x};\mathbf{y},\mathcal{W}^{\mathrm{B}}].\tag{3.2}$$

The layer-wise quadratic loss is defined via 2nd-order Taylor expansion:

$$L[\mathbf{w}] \stackrel{\text{def}}{=} \frac{1}{2}(\mathbf{w} - \mathbf{w}^{*})^{\top}\mathbb{E}_{\mathbf{x}\sim\mathcal{D}}\left(\frac{\partial^{2}\ell}{\partial\mathbf{w}^{\top}\partial\mathbf{w}}[\mathbf{x};\mathcal{W}_t]\right)(\mathbf{w} - \mathbf{w}^{*}).\tag{3.1}$$

Key Prediction (Prediction 3.5)

Under Assumption 3.4 (i.i.d. Gaussian inputs after linear transformation), the expected gradient covariance satisfies:

$$\mathbb{E}(\mathbf{g}_i\mathbf{g}_j^{\top}) \approx \begin{cases} \lambda_i \mathbb{E}(\nabla\mathbf{y}_{t+s}\nabla\mathbf{y}_{t+s}^{\top}) & \text{if } i = j \\ \mathbf{0} & \text{if } i \neq j \end{cases}\tag{3.3}$$

where $\mathbf{g}_i, \mathbf{g}_j$ are the $i,j$-th column vectors of $(\nabla\mathbf{W}_{t+s})\mathbf{U}$, and $\mathbf{U}$ comes from the eigen-decomposition $\mathbb{E}(\mathbf{h}_t\mathbf{h}_t^{\top}) = \mathbf{U}\text{Diag}(\boldsymbol{\lambda})\mathbf{U}^{\top}$.

This leads to:

$$\mathbb{E}(\|\nabla\mathbf{W}_{t+s}\|^2) \approx \sum_{i=1}^{D}\lambda_i \mathbb{E}(\|\nabla\mathbf{y}_{t+s}\|^2).\tag{3.4}$$

Simulation Algorithm (Algorithm 3.1)

The simulation algorithm synthesizes pseudo-gradients by sampling Gaussians to match the average and 2nd moment of real gradients:

1. Inputs required: eigenvalues/eigenvectors of $\mathbb{E}(\mathbf{h}_t\mathbf{h}_t^{\top})$, average gradient $\mathbb{E}(\nabla W_t)$, mean squared-norm of per-batch gradients $\mathbb{E}(\|\nabla_{\mathcal{B}}W_t\|^2)$, and Hessian trace $\text{Tr}(\mathbb{E}(\nabla^2W_t))$
2. Key steps:
   • Set constant $c$ to match Hessian trace
   • Compute equivalent batch factor $\mathfrak{B}$
   • Initialize $z = c(\mathbf{w} - \mathbf{w}^{*})$
   • At each step: sample stochastic noise $n_k \sim \mathcal{N}\left(0, \lambda_k|\sum_i\lambda_iz_i^2/\mathfrak{B} - \lambda_kz_k^2|\right)$, synthesize gradient, call optimizer

Theoretical Explanation of Spikes

The authors decompose SGD updates into:

$$\overline{\mathbf{w}}_{t+1} - \mathbf{w}^{*} = (\mathbf{I} - \eta_t\mathbf{H})(\overline{\mathbf{w}}_t - \mathbf{w}^{*})\tag{3.8}$$ $$\mathbf{r}_{t+1} = (\mathbf{I} - \eta_t\mathbf{H})\mathbf{r}_t - \eta_t\mathbf{n}_t.$$

Mechanism: The stochastic noise $\mathbf{r}_t$ accumulates as a random walk. When the eigenvalues of the quadratic loss function are heavy-tailed, the validation loss suddenly surges when $\mathbf{r}_t$ randomly points toward directions of large eigenvalues—a "spike" occurs. Concentration of measure breaks down precisely when the eigenvalue distribution is too heavy-tailed.

---

Empirical Validation / Results

PLRF Model Experiments

The paper demonstrates training spikes in PLRF models (Figure 2):

Conditions for training spikes in PLRF:

1. Heavy-tailed data (large $\alpha$)
2. Near divergence (large learning-rate)

Transformer LM Case Study

The authors trained a 1.3B parameter Transformer decoder (24 blocks) on the C4 dataset with AdamW, changing $\beta_2$ from 0.95 (Normal) to 0.999 (Spiky1, Spiky2).

Compute requirements (Table 1):

Simulation Results

The simulation from step 21k to 22.5k showed:

• Clear bumps for Spiky1 and Spiky2 configurations (e.g., block_{0,4,8} at step ~22k for Spiky1, block_{0,12} at step ~21.5k for Spiky2)
• Smooth trajectory for the Normal configuration
• Spiky positions differ between random seeds, demonstrating the random nature of training spikes

The empirical error of equation (3.4) was mostly within 5%, even though both sides simultaneously vary multiple magnitudes across different layers and training steps.

Key Observations on Limitations

1. 2nd-order terms exhibit "jumps" rather than "spikes" (not resuming to previous level), suggesting correlations across layers and higher-order terms matter
2. 1st-order terms decrease much faster than expected from a layer-wise quadratic model, indicating the DNN's non-linear nature accelerates feature learning
3. The equivalent batch factor for Transformer LM training was found to be $\approx 1$, meaning gradient noise cannot be ignored despite large batch sizes

---

Theoretical and Practical Implications

Theoretical Implications

1. Training spikes can occur in convex quadratic models—contrary to previous belief that 3rd-order terms are necessary (Cohen et al., 2021; Damian et al., 2023; Cohen et al., 2025)
2. Concentration of measure breaks down for heavy-tailed data distributions ($\alpha > 0.5$), even at massive scales
3. Sharpness $\lambda_{max}$ is only distantly relevant to stochastic optimization dynamics; the full eigenvalue distribution matters
4. Natural language Zipf's Law ($\alpha = 0.5$ for token unigram probabilities) connects directly to the power-law framework

Practical Implications

1. $\beta_2$ schedule for Adam: The authors derive an appropriate increasing $\beta_2$ schedule that avoids spikes and empirically improves performance (§ C.5)
2. Monitoring metrics: The power-law exponent of eigenvalue distributions of forward activation correlation matrices serves as an indicator of healthy training
3. Lightweight simulation: The algorithm runs in 5 seconds on CPU per layer, requiring only one pass over a validation set—enabling rapid testing of optimization techniques

Connection to Related Work

The paper connects to several research streams:

• Edge-of-Stability (Cohen et al., 2021, 2022): stochastic optimization introduces noise that changes the dynamics
• Gradient noise analysis (Mandt et al., 2016; Li et al., 2019): SDE approaches, but with explicit variance calculation instead of assuming constant noise
• Double Descent (Adlam and Pennington, 2020): similar distinction between data-oriented and model-oriented randomness
• "Stochastic edge-of-stability" (Agarwala and Pennington, 2025): differences in learning-rate thresholds between full-batch and SGD

---

Conclusion

The paper demonstrates that simple quadratic models with heavy-tailed inputs can replicate training spikes, challenging the assumption that spikes require complex non-convex dynamics. The key mechanism is stochastic gradient noise interacting with heavy-tailed eigenvalue distributions, causing concentration of measure to break down.

Main takeaways:

1. PLRF models provide a minimal, reproducible setting for studying training spikes (code runs in 2 minutes on Google Colab)
2. The simulation algorithm enables lightweight prediction of spiky behavior from a given checkpoint
3. The $\beta_2$ schedule recipe improves training stability and performance

Future directions:

• Understanding how correlations across layers and Taylor expansions beyond 2nd order contribute to the full validation loss dynamics
• Investigating why the 1st-order terms decrease faster than quadratic models predict
• Exploring how optimization hyperparameters like $\beta_2$ affect the eigenvalue distribution of the Hessian
• Moving from worst-case bounds to average-case analysis of stochastic optimization with heavy-tailed inputs

Final remark: "Given the apparent similarity between empirical training spikes and our simulations, we believe quadratic models at least partially reveals the essence of underlying mechanism."