Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models
Ran Tian, Courtney Paquette
cs.LG
2026-10-09
Heavy-tailed inputs and near-divergence learning rates reproduce loss spikes even in convex quadratic models; a lightweight simulation predicts them from a checkpoint.
Pre-training large language models occasionally produces training spikes: short, sharp surges in the loss that settle back down. PaLM-scale runs have documented them, and practitioners treat them as pathological. The surge may signal unstable training, suboptimal hyperparameters, or a worse final model, and violent swings are risky under the low-precision arithmetic these runs use.
The obstacle is reproduction. Figure 1 gives a clean contrast: ordinary loss fluctuations follow the data shuffle. Keep the shuffle fixed, vary the init seed or minor architecture details, and the shape of the curve barely changes. Spikes break that rule. Two runs differing only in init seed spike at different steps. Scrubbing data outliers only partially helps for the same reason: random spikes persist in validation loss, so a single data point is usually not the cause. Prior theory mostly orbited sharpness, the largest Hessian eigenvalue, and no toy model's loss curve ever matched real spikes.
The core claim: ordinary fluctuations are driven by data shuffling alone, because high dimensionality cancels model-originated randomness; when the data distribution is heavy-tailed, that cancellation breaks and spikes appear.
The vehicle is the Power-Law Random Feature (PLRF) model: a frozen random matrix maps inputs to features, linear regression with quadratic loss follows, and input coordinates are scaled by i^{-α}, with α measuring heavy-tailedness. No depth, convex loss, no non-convex feature learning. Yet at α=0.9 it produces spikes that look like the real thing. Two conditions must hold together: heavy tails (large α) and a learning rate close to divergence. At α=0.3, raising the learning rate just yields smooth divergence.
The intuition, for SGD: split the update into the deterministic full-batch descent and a random-walk term r that accumulates gradient noise. In high dimensions r's norm concentrates while its direction stays random, so the validation loss, a quadratic function of r, collapses onto a deterministic trajectory. That concentration of measure requires the eigenvalues of the quadratic not to be too heavy-tailed. If a few eigenvalues dominate, the loss surges the moment r points along those directions: a spike. The learning rate must be large enough for r to accumulate, and just above that lies divergence; the SGD divergence threshold has an exact formula (Eq. 2.2). Appendix A.6 shows concentration holds only for α≤0.5, and natural language sits right at the edge: viewing the input embedding layer as a linear map from one-hot vectors, its eigenvalue distribution is exactly the Zipf distribution of unigram probabilities, corresponding to α=0.5.
Connecting this to real networks takes two steps. A layer is defined as a standalone linear unit whose input and output gradient do not depend on its own weights (most DNNs decompose this way; shared weights such as an embedding reused as the output head are the exception). A second-order Taylor expansion of the validation loss at step t then gives each layer a quadratic loss, and Theorem 3.3 shows per-example gradients can be approximated as coming from a Random Feature model, via an NTK-style latent space of extremely high dimension. Prediction 3.5 gives an explicit formula for the second moment of gradients; on the Transformer its empirical error stays mostly within 5% while both sides vary over orders of magnitude across layers and steps.
The simulation (Algorithm 3.1) then samples Gaussian pseudo-gradients matching the mean and second moment at each step and feeds them to the optimizer. Startup needs four quantities from one pass over a validation set: the eigendecomposition of the activation correlation matrix, the average gradient, the per-batch gradient norm, and the Hessian trace, with setup costing about two TPU hours in total. After that no DNN backward pass is ever run; 1,500 steps for one layer take 5 seconds on CPU.
PLRF toy runs (D=1024, V=20000, batch 8):
| Setting | Behavior |
| α=0.3, learning rate raised (SGD/Adam) | smooth divergence, no spikes |
| α=0.9, learning rate near divergence | both SGD and Adam spike before diverging |
The Transformer case study is the centerpiece. A 1.3B decoder (24 blocks) trained on C4 to the Chinchilla-optimal token count, 25.5B tokens over 40 hours on TPUv4×32, with AdamW hyperparameters inherited from a tuned configuration; β2 was then changed from 0.95 to 0.999 to induce spikes. Spiky1 and Spiky2 differ only in init seed. Simulating from the 21k checkpoint to 22.5k: both Spiky configurations show clear bumps across several FFN layers (blocks 0/4/8 around step 22k for Spiky1, blocks 0/12 around 21.5k for Spiky2), while the Normal simulation stays comparatively smooth. Simulated spike positions vary with the seed too. In actual training, only Spiky2 spiked in that window; the simulation delivers a verdict on whether a configuration can spike, not the timing.
Supporting measurements back the mechanism. The eigenvalue distribution of the activation correlation matrix is significantly more heavy-tailed under spiky configurations. The equivalent batch factor of the 1.3B model is about 1: per-example losses in DNNs amplify gradient noise and cancel the damping from a large batch. The takeaway is blunt: gradient noise cannot be ignored when training a Transformer, λmax is only distantly relevant, and the full eigenvalue spectrum is what matters. The authors also derive a β2 schedule from observed metrics that empirically improves performance in their case, framed as proof of concept rather than a replacement for existing stabilization methods.
This is the first study to explicitly tie training spikes to heavy-tailed data distributions, with a convex toy model whose curves match the real phenomenon (the code runs on Colab in two minutes). Three practical payoffs. A warning tool: from any checkpoint, a few CPU seconds estimate whether a configuration will spike, a cheap pre-check before committing to continued training. A monitoring signal: the power-law exponent of the activation correlation spectrum works as a health indicator. A corrected intuition: large batches do not imply low gradient noise, since the effective noise level here matches batch size 1, and the folk knowledge that β2 near 1 invites spikes now has a mechanistic account. To be clear, this is a mechanism-and-tool paper: no new stabilization method, and the practical recipe is incremental.
The authors list several themselves. The layer decomposition excludes shared weights (an embedding reused as the output head), left to future work. The layer-wise quadratic model fails to explain the first-order Taylor terms: they run about 20× the scale of the second-order terms and fall faster than a quadratic model predicts, with spikes appearing as bumps, jumps, or dives that quadratic models cannot produce at that scale; cross-layer correlations and higher-order expansions are needed. The simulation also omits weight decay.
The assumption chain deserves scrutiny. The latent distribution is treated as non-drifting and i.i.d. Gaussian over short horizons (Assumption 3.4), the output-gradient covariance is collapsed to isotropic (Assumption 3.6), and the simulation tracks an average slice of the weight matrix rather than the whole thing. The case study covers a single 1.3B scale and a single β2-triggered path; whether spikes induced by data or architecture changes share the mechanism is untested. How β2 reshapes the Hessian spectrum is, in the authors' own words, a curious mystery.