Training-loss spikes reproduced by a quadratic model, and predicted from a checkpoint
Synopsis
Training-loss spikes have been hard to reproduce under controlled conditions, and their cause is not fully understood. Researchers show that simple quadratic models with heavy-tailed, power-law inputs reproduce the spikes, then align per-layer training dynamics of deep networks with such models and give a lightweight simulation algorithm. In a 1.3B-parameter Transformer language model case study, where an AdamW hyperparameter was changed to induce spikes, the simulation from a step-21k checkpoint predicts bumps for the Spiky1 and Spiky2 configurations while Normal stays comparatively smooth; a schedule derived from observed metrics improved performance in that case.
Interpretation
Training-loss spikes can appear in simple quadratic models without nonlinear feature learning, which changes the impression that spikes belong only to the non-convex dynamics of deep networks. Spikes had often been attributed to feature learning and non-convexity in deep networks, and quadratic models were considered incapable of showing them. Across optimizers, learning rates and power-law exponents in power-law random feature models, spikes appear before divergence under both SGD and Adam only when the data are heavy-tailed and the learning rate is near the divergence threshold.
A layer-wise quadratic loss treats each network layer as a standalone linear unit and uses a second-order Taylor expansion to describe that layer's contribution to validation loss. Earlier analyses of stochastic optimization mostly stayed with full-batch gradient descent or worst-case bounds, without a per-layer treatment of stochastic gradient noise. Under this expansion, each layer's change in validation loss becomes a quadratic function of its parameters whose gradient matches the validation-loss gradient at the training point, separating that layer's evolution from the rest of the network.
The simulation calls no backward pass and instead synthesizes pseudo-gradients from each step's mean and second moment, keeping the cost small. Measuring the full Hessian tensor directly is infeasible at billion-parameter scale, and the layer-wise simulation reduces what is needed to one forward pass of statistics over a validation set. On a 1.3B-parameter, 24-block Transformer decoder, simulating six feed-forward layers with two random seeds each from step 21k to 22.5k shows clear bumps for Spiky1 and Spiky2 while Normal stays comparatively smooth; the simulation takes about 5 seconds on CPU, against about 8 minutes for a full-layer evaluation.
Perspective
The result addresses researchers and engineers studying the stability of billion-scale language-model pre-training. The demonstrated scope is one 1.3B-parameter, 24-block Transformer decoder on C4, with simulation from a step-21k checkpoint covering up to step 22.5k and only six feed-forward layers with two random seeds each. The schedule derived from observed metrics improved validation loss in that case, as a proof of concept for the theory.
The simulation does not always reproduce every spike statistic: on power-law random feature models it tends to underestimate spike size at small power-law exponents and clearly overestimate it at large ones. The layer-wise quadratic model also cannot explain why first-order terms fall faster than expected or why second-order terms jump rather than return, leaving cross-layer correlations and Taylor expansions beyond second order open. In the case study, simulated spike positions vary with random seed and do not correspond point by point to the actual training trajectory.