# Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models

> Training spikes in LLM pretraining are reproduced by quadratic models with heavy-tailed inputs, where stochastic gradient noise interacting with power-law eigenvalue distributions breaks concentration of measure.

- **Source:** [arXiv](https://arxiv.org/abs/2610.13524)
- **Published:** 2026-10-17
- **Permalink:** https://picx.dev/p/DLWo5o
- **Whiteboard:** https://picx.dev/p/DLWo5o/image

## Summary

## 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):

| Setting | $\alpha = 0.3$ | $\alpha = 0.9$ |
|---------|---------------|---------------|
| SGD | Diverges without spikes | **Shows spikes before divergence** |
| Adam | Diverges without spikes | **Shows spikes before divergence** |

**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):

| Experiment | Data size | Time |
|-----------|-----------|------|
| Train | 25.5B tokens | 40 hrs (TPUv4×32) |
| Eval $\delta_{t,s}^{w}$ | 336M tokens | 1.1 hrs |
| Eval $\mathbb{E}(\mathbf{h}_t\mathbf{h}_t^{\top})$ | 336M tokens | 8 min |
| $\mathbb{E}(\nabla\mathbf{W}_t)$ & $\mathbb{E}(\|\nabla_{\mathbf{B}}\mathbf{W}_t\|^2)$ | 336M tokens | 15 min |
| Tr$(\mathbb{E}(\nabla^2\mathbf{W}_t))$ | 336M tokens | 1.7 hrs |
| **Simulation** | **1500 steps** | **5 sec (CPU)** |

### 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."

---

_Markdown view of https://picx.dev/p/DLWo5o, served by PicX — AI-generated visual whiteboard summaries of research papers._