PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ran Tian, Courtney Paquette

arXiv:2610.13524v1cs.LG

TL;DR

Training spikes in LLM pre-training are empirically observed, but their causes remain incompletely understood and their sensitivity to initialization and data shuffling complicates study. The paper uses power-law random-feature quadratic models and layer-wise simulations to investigate the phenomenon. Its simulations suggest a role for second-order gradient noise in some DNN layers, and the authors report that an appropriate β2 schedule empirically improves performance.

  • Problem

    Training spikes are loss surges whose causes remain incompletely understood, and their sensitivity to data shuffling and initialization complicates their study.

    Source: §1 · p. 1

  • Method

    The paper uses power-law random-feature quadratic models and a layer-wise simulation algorithm to model DNN training dynamics.

    Source: §1 · p. 2, §3.3 · p. 8

  • Results

    The simulations suggest that DNN training spikes likely originate from second-order gradient noise within some layers, and an appropriate β2 schedule empirically improves performance.

    Source: §1 · p. 3, §5 · p. 15

  • Takeaways & Limitations

    The authors conclude that quadratic models can at least partially reveal the underlying mechanism of empirical training spikes.

    Source: §6 · p. 16

    PLRF cannot capture the changing Hessian during DNN training, and the exact mechanism by which β2 affects training dynamics, especially Hessian eigenvalue distributions, remains unresolved.

    Source: §5 · pp. 14, 16

Abstract

from arXiv · show

Intermittent ``spikes'' in loss, during the pre-training of Large Language Models (LLMs), remains a mysterious phenomenon both theoretically and empirically. In this paper, we show that simple quadratic models with heavy-tailed, power-law distributed inputs can replicate these spikes. We introduce a theoretical framework and simulation algorithm to model per-layer training dynamics of Deep Neural Networks (DNNs), using such quadratic models. In a case study on an Adam hyperparameter configuration (where $\beta_2$ is set close to $1$) known to induce spikes in Transformer LMs, our simulation successfully predicts bumps indicating spiky behavior from a given checkpoint. Furthermore, we derive an appropriate $\beta_2$ schedule from observed metrics, which empirically improves performance.

1. Introduction.

The paper addresses training spikes whose causes are difficult to isolate, and proposes power-law random-feature models and layer-wise simulations to study them. It reports that heavy-tailed inputs can produce spikes in simple quadratic models and that simulated second-order gradient noise may help explain spikes in DNN layers.

1. Introduction.

  • Training spikes are sharp, temporary loss surges whose causes remain incompletely understood and whose occurrence can be difficult to reproduce reliably.They may signal instability or sub-optimal hyperparameters, and fluctuating metrics can also trigger numerical issues under low-precision arithmetic.Source: §1 · p. 1
  • Unlike ordinary loss fluctuations, spike timing can vary with model initialization even when data shuffling is held fixed, so a single data outlier is not usually the cause.The introduction notes that removing data outliers may only partially address spikes because random validation-loss surges can remain.Source: §1 · p. 1, Figure 1 · p. 2
  • PLRF models are used to demonstrate the proposed mechanism, and a PLRF-like algorithm simulates individual DNN-layer training dynamics.The paper presents these quadratic models as a lightweight way to study elusive DNN phenomena.Source: §1 · pp. 2–3
  • Heavy-tailed inputs let quadratic PLRF models exhibit training spikes, with larger α a necessary condition for spikes in the reported comparison.The introduction identifies this as the first explicit link, to the authors’ knowledge, between training spikes and heavy-tailed data distributions.Source: §1 · p. 2
  • The simulation framework predicts that DNN training spikes likely originate from second-order gradient noise within some layers.Its approach extends Taylor-based prediction of future dynamics to stochastic optimizers such as SGD and Adam, using a second-order expansion to account for gradient noise.Source: §1 · pp. 2–3
  • In a Transformer-LM case study, the simulation predicts spiky behavior from a checkpoint, and the authors report that an appropriate β2 schedule empirically improves performance.The case study changes Adam’s β2 to induce spikes; the schedule is proposed after analysis of training metrics.Source: §1 · p. 3

2. Background.

The background contrasts smooth behavior expected from full-batch quadratic optimization with spikes that can arise under stochastic optimization and heavy-tailed inputs. It introduces PLRF models and explains how spikes can coexist with a smooth expected validation loss.

2. Background.

  • A PLRF model uses power-law-scaled Gaussian inputs, with α as the power-law exponent; its features feed linear regression with quadratic loss.The paper uses Ax to approximate DNN hidden vectors, with D as the DNN model dimension and V as a very large latent-space dimension.Source: §2 · p. 3
  • SGD dynamics for PLRF can be solved analytically, enabling numerical predictions of compute-optimal scaling; the study also reports that large α can cause training spikes.The background contrasts this stochastic setting with the straightforward full-batch quadratic divergence threshold.Source: §2 · p. 3
  • For α = 0.9, SGD and Adam show spikes before divergence, whereas for α = 0.3 the validation loss diverges without spikes; the reported conditions are heavy-tailed data and a large learning rate near divergence.The comparison uses D = 1024, V = 20000, and batch size B = 8.Source: §2 · p. 3
  • For PLRF with α > 0.5, expected validation loss remains smooth, but convergence to that expectation can no longer be proved and individual runs can spike even at massive scales.This qualifies the usual Concentration of SGD result, which describes a smooth deterministic curve when D and V are sufficiently large.Source: §2 · p. 4
  • With stochastic gradient noise, quadratic models can spike when inputs are heavy-tailed; third-order terms are not necessary, unlike the cited full-batch Edge-of-Stability account.For full-batch gradient descent on quadratic loss, divergence occurs above 2/λmax, and the cited stabilization behavior near that threshold involves DNN third-order terms.Source: §2 · p. 5

3. Theory.

The section develops a layer-wise quadratic-loss framework and random-feature approximation for modeling DNN layer dynamics. It then uses synthetic stochastic gradients to simulate those dynamics and explain when spikes arise.

3.1. Layer-wise Quadratic Loss.

  • A second-order Taylor expansion defines a layer-wise quadratic loss for evaluating one layer's contribution while other parameters are fixed, although training updates remain unfrozen.The resulting loss is quadratic in the selected layer's parameters, and its gradient is matched to the validation-loss gradient at the checkpoint.Source: §3.1 · pp. 5–6

3.2. Per-Layer Approximation by Random Feature Models.

  • A high-dimensional random-feature approximation represents layer activations and output gradients, with gradient-norm error mostly within 5% in a Transformer LM despite an infeasible latent space.Under Assumption 3.4, the latent features are modeled using transformed i.i.d. Gaussian inputs.Source: §3.2 · pp. 6–7

3.3. Simulation Algorithm.

  • The algorithm matches Hessian trace, initializes eigenmode-wise gradient norms, and synthesizes batch gradients; on PLRF, it reproduces training dynamics including spikes.It simulates a D-dimensional average slice and approximates output-gradient covariance by a scalar operator when most eigenvalues are close.Source: §3.3 · p. 8

3.4. A Sketch of why Training Spikes Occur.

  • Under heavy-tailed spectra, accumulated stochastic noise can align with large-eigenvalue directions, producing spikes when the learning rate is slightly below divergence.At high dimension, ordinary fluctuations concentrate around a deterministic trajectory; the spike mechanism appears when that concentration breaks down.Source: §3.4 · p. 10

4. Simulation of Transformer LM Training.

Checkpoint-based simulations show bumps in both spiky configurations while the Normal simulation is smoother, but the layer-wise quadratic account does not capture all measured dynamics. The empirical comparisons point to limits in explaining total validation-loss spikes with layer-wise quadratic terms alone.

4. Simulation of Transformer LM Training

  • Both Spiky1 and Spiky2 simulations show clear bumps, while Normal is comparatively smooth; actual training spikes only for Spiky2 from steps 21k to 22k.The simulation bumps also appear at different positions across random seeds.Source: §4 · pp. 10–11, Figure 3 · p. 12
  • Measured second-order terms jump in Spiky2 alongside its training spike, with jump sizes matching simulation fluctuations in some layers but not all.The measured terms jump rather than return to their previous level as total-loss spikes do.Source: §4 · p. 11
  • Layer-wise quadratic models do not fully explain total validation loss; cross-layer correlations or Taylor terms beyond second order may be needed.The measured second-order terms show jumps rather than the total loss’s spike pattern.Source: §4 · p. 11
  • The first-order terms decrease much faster than a layer-wise quadratic model predicts, which the authors associate with DNN nonlinearity and feature learning.The authors also suggest that DNN nonlinearity produces larger noise than the quadratic model predicts.Source: §4 · p. 11
  • Spiky2 first-order terms show bumps, jumps, or dives across blocks, behaviors that layer-wise quadratic models cannot reproduce at this scale.Figure 4 uses a shared y-axis range of −4 to 1 across the three configurations.Source: §4 · p. 11, Figure 4 · p. 13

5. Discussion and Related Works.

The discussion connects training spikes to stochastic optimization and heavy-tailed spectra while emphasizing that quadratic layer-wise models only approximate DNN dynamics. It presents monitoring metrics and a β2 schedule as theory-motivated proof-of-concept directions, not replacements for existing methods.

5. Discussion and Related Works

  • PLRF models produce visually spike-like loss trajectories without log-scaled axes, though the authors stress their examples’ aesthetic similarity to observed spikes.The supplementary code runs on Google Colab within two minutes.Source: §5 · p. 14
  • Transformer gradient noise is comparable to a Random Feature model with batch-size ≈1 because point-wise DNN training loss can offset the effect of large batches.The authors conclude that gradient noise cannot be ignored in Transformer LM training.Source: §5 · p. 14
  • For stochastic optimization, the distribution of all Hessian eigenvalues matters more than sharpness λmax alone, and the gradient-noise variance can be calculated for Random Feature models.This contrasts with analyses that assume a constant noise variance.Source: §5 · pp. 14, 16
  • Layer-wise quadratic models cannot capture the evolving DNN Hessian; the authors also report that their layer-wise shortcut reduces computation but leaves the β2–spectrum mechanism open.They report a slightly decreasing power-law exponent for Normal training despite an increasing Hessian trace.Source: §5 · pp. 14–15
  • The authors connect language’s Zipf-like token frequencies to power-law eigenvalues of the input correlation matrix, with α = 0.5 under the PLRF definition.They suggest that scaling statistics of language universals may underlie neural scaling laws.Source: §5 · p. 15
  • The paper proposes a β2 schedule that improves performance in its case study, while noting that performance and spikes depend on many factors and the mechanism remains unclear.The authors also describe the recipe as proof of concept rather than a replacement for previous methods.Source: §5 · pp. 15–16
  • The authors propose monitoring the power-law exponent of forward-activation correlation eigenvalues and using their simulation algorithm to test optimization techniques.The explored β2 schedule is presented as a proof of concept, not a replacement for prior methods.Source: §5 · p. 16

6. Conclusion.

The paper presents theory, an algorithm, and experiments aligning PLRF models with layer-wise DNN training dynamics. The authors argue that quadratic models partially reveal the mechanism behind empirical training spikes.

6. Conclusion

  • The paper aligns PLRF models with layer-wise DNN training dynamics through theory, an algorithm, and experiments.The authors see the similarity between simulated and empirical spikes as evidence that quadratic models partially reveal the underlying mechanism.Source: §6 · p. 16
  • The authors’ conclusion is qualified: quadratic models reveal the mechanism only partially, based on the apparent similarity between simulations and empirical spikes.Source: §6 · p. 16
  • The conclusion rests on the apparent similarity between PLRF simulations and empirical training spikes.Source: §6 · p. 16

AI Use Statement.

The authors report that generative AI was not used to develop the theory, claims, experiments, or methods. They used it during writing for rephrasing and related-paper summaries, then checked and manually rewrote the outputs.

AI Use Statement

  • Generative AI was not used to develop theory, formulate mathematical claims, design experiments, or implement methods.Source: AI Use Statement · p. 16
  • During writing, AI was consulted to rephrase draft paragraphs and summarize related papers; the authors checked and manually rewrote the outputs.Source: AI Use Statement · p. 16
  • The authors take responsibility for the paper’s final text, claims, proofs, and code.Source: AI Use Statement · p. 16

Reproducibility Statement.

The statement reports reproducibility materials and Transformer experiment settings, then records the hyperparameter dependence of the reported dynamics and conventions used in the theory appendix.

Reproducibility Statement.

  • The supplementary materials include Python code for training spikes and simulation, and the Transformer LM experiments use a decoder-only Transformer.Hyperparameters are listed in § C.1.Source: Reproducibility Statement · p. 16
  • The presented training dynamics depend on hyperparameters, especially the learning-rate schedule and weight decay, while the observed phenomena are expected to generalize across modeling choices.The text directs readers to Appendix A for mathematical details.Source: Reproducibility Statement · p. 17

Appendix A. Theory and Proof.

  • The appendix specifies vector, matrix, trainable-tensor, function-call, and gradient notation, and assumes finite training and validation datasets drawn from the same distribution.Input data points are denoted by x and may have any internal structure.Source: Appendix A · p. 24

A.1. Proof of Theorem 3.3.

The proof constructs a universal linear representation for Taylor-expanded DNN dynamics by concatenating training examples. This lets it express future activations and outputs as linear maps of an augmented input.

A.1. Proof of Theorem 3.3.

  • Concatenating the training examples’ Jacobians and augmenting the input with current activations yields a shared matrix A for the first-order expansion of future activations.The construction addresses the mismatch between point-wise loss expansions and the Random Feature model’s universal matrix.Source: §A.1 · p. 24
  • The proof similarly expands future outputs by combining the activation change with the weight change, represented through an augmented input and padded weight difference.This gives an affine representation in the future weights.Source: §A.1 · p. 24
  • For the loss-gradient construction, the proof concatenates data-dependent matrices across training examples and augments the input with the current gradient to obtain an affine function of future weights.The Jacobians and related matrices used in the construction do not depend on the forecast horizon s.Source: §A.1 · p. 25

A.2. Proof of Prediction 3.5.

The proof uses Gaussian input coordinates and Wick’s formula to evaluate the relevant matrix expectation. Under a sufficiently general orthonormal basis, off-diagonal projection terms vanish in the high-dimensional limit.

A.2. Proof of Prediction 3.5.

  • With Gaussian coordinates, Wick’s formula reduces the fourth-moment calculation to three index pairings; one pairing gives the identity contribution.The identity follows from the orthonormal columns of V.Source: §A.2 · pp. 25–26
  • The other pairings yield rank-one projection terms whose average magnitude is about 1/D, making them negligible relative to the identity when D is sufficiently large.This argument assumes V is sufficiently general; a Haar-measure assumption gives an overwhelming-probability bound.Source: §A.2 · p. 26
  • Although the lemma’s proof uses properties specific to Gaussian inputs, the authors expect its high-dimensional conclusion to extend asymptotically to other mean-zero, unit-variance random vectors.Source: §A.2 · p. 26

A.3. Equivalent Batch Factor.

These appendices develop an equivalent batch factor and the simulation details, then derive a learning-rate threshold and identify when high-dimensional concentration holds. The results rely on stated approximations and dimension conditions.

A.3. Equivalent Batch Factor.

  • The mini-batch gradient-noise calculation requires B ≪ D, under which increasing batch size decreases gradient noise.The derivation attributes the condition to cross-interactions within a batch.Source: §A.3 · p. 27
  • For DNNs, the coefficient in the quadratic-loss approximation varies across data points, so the simulation estimates its average to match the layer loss Hessian trace and absorbs a larger-than-one multiple into the equivalent batch factor.The factor is intended to coincide with batch size for Random Feature models; the DNN approximation remains rough.Source: §A.3 · pp. 27–28

A.4. Detailed Derivation of Algorithm 3.1.

  • The simulation initializes eigenmode magnitudes to match the gradient RMS, draws synthetic per-batch gradients from Gaussian distributions with predicted means and variances, then applies the optimizer.The sign of each initialized coordinate is irrelevant because the simulation depends on its absolute value.Source: §A.4 · pp. 28–29

A.5. Derivation of Learning-rate Threshold.

  • The learning-rate threshold is derived by assuming gradient variance remains constant at the threshold; convergence requires γ to stay below it.The learning rate should scale in proportion to 1/Σ_i λ_i for the D → ∞ limit to be well-defined.Source: §A.5 · pp. 29–30, §A.6 · p. 31

A.6. Concentration Effect Only Holds for α ≤0.5.

  • For power-law eigenvalues proportional to i^-2α and i.i.d. mode variables, random fluctuations cancel in the high-dimensional limit if and only if α ≤ 0.5.The argument applies the law of large numbers when the relevant weighted variables have finite expectation and variance.Source: §A.6 · p. 31

Appendix B. PLRF Simulation Experiments.

The appendix tests PLRF simulation against spiky trajectories and applies layer-wise measurements to a Transformer LM. It reports qualitative spike prediction, heavy-tailed activation spectra, and a near-unit equivalent batch factor.

Appendix B. PLRF Simulation Experiments.

  • The simulation matches average loss flows and qualitatively predicts spikiness, but underestimates spike size at B=1 and overestimates it at B=64.The authors attribute the mismatch to non-Gaussian per-example gradients and a second-moment estimate that is asymptotically accurate only for small batches.Source: Appendix B · pp. 31–32

C.1. Model and Training Settings.

  • The k=2 Transformer setup uses 24 blocks, 32 attention heads, 2048-token sequences, AdamW with β1=0.9 and β2=0.95, and about 48,500 training steps.The model-component counts and optimizer settings are given for the case-study configuration.Source: §C.1 · p. 33
  • Per-example and per-batch gradient squared norms are distinct quantities because the batch-gradient square is taken after summing gradients.The paper notes there is no obvious relation between their dataset averages.Source: §C.1 · p. 34

C.2. Power-Law Distribution of Forward Activations.

  • Spiky training has more heavy-tailed activation eigenvalue spectra, with larger fitted α and more unstable power-law statistics than Normal training.The fitting procedure estimates α by likelihood and the truncation threshold by minimizing the KS statistic.Source: §C.2 · p. 35
  • The Spiky–Normal spectral gap persists in lower layers, while upper-layer distributions become more similar and appear to converge near the end of training.The reported examples include block 0 for lower layers and blocks 12 and 20 for upper layers.Source: §C.2 · p. 35

C.3. Empirical Verification of Prediction 3.5.

  • The empirical relative error of Equation (3.4) is below 0.1 and mostly within 5%, despite gradient norms varying by multiple magnitudes during training.The ratio plot is used to verify the formula for per-example gradient norms.Source: §C.3 · pp. 35–36

C.4. Is the “Batch-size” Large or Small?.

  • Token count versus sequence count affects the theoretical batch-size interpretation, while the observed small B indicates a typical stochastic-optimization regime.The paper notes that its analysis is more accurate when batch-size is small relative to hidden dimension D.Source: §C.4 · p. 36
  • The measured equivalent batch factor B stays close to 1; its first factor is slightly below tokens per batch, while the third factor makes B small.The second factor is close to 1, as predicted.Source: §C.4 · p. 36

C.5. The Appropriate β2.

The section derives a metric-guided β2 schedule from the learning rate and measured curvature and gradient metrics, then tests it against fixed settings. The schedule improves validation loss, while a constant setting at its final β2 is spiky.

C.5. The Appropriate β2.

  • The recipe estimates normalized Hessian trace T and per-batch gradient scale N from validation data alongside Adam’s learning-rate schedule η_t.The proposal is motivated by matching the self-correlation times of gradient noise and Adam’s running average of gradient squares.Source: §C.5 · p. 37
  • The approximation of N by the mean of Adam’s running gradient-square state works better in Normal than Spiky training, where large β2 can prolong surges.The authors suggest prolonged surges cause the running state to overestimate the average.Source: §C.5 · p. 37
  • To capture N's comparatively stable and consistent behavior across all layers for both Normal and Spiky, the authors manually fit a schedule with initial value 0.018, final value 0.008, and exponent 3.The schedule is shown as pink dashed curves in the Fourth Row of Figure 12.Source: §C.5 · p. 37
  • The proposed β2 schedule can improve validation-loss performance and slightly outperform β2 = (1−init value), although it ends at β2 = (1−final value); training at the latter setting is spiky, while the former is not.Source: §C.5 · p. 38

Related papers

Loading 2610.13524v1…