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

The gist The paper shows that simple quadratic models with heavy-tailed, power-law distributed inputs can replicate intermittent “spikes” in loss during pre-training of Large Language Models (LLMs) by introducing a theoretical framework and simulation algorithm to model per-layer training dynamics using such models.

How it works

The authors introduce a theoretical framework that posits that normal fluctuations in loss are driven solely by data shuffling because high model dimensionality typically cancels out model-originated randomness when the data distribution is heavy-tailed, resulting in training spikes. They demonstrate this mechanism using Power-Law Random Feature (PLRF) models, which are simple linear-regression models with quadratic loss.

The simulation framework is largely inspired by Lee et al. (2019), utilizing a Taylor expansion of the DNN at training step t to predict dynamics up to a future step t + s. This approach is advanced by incorporating stochastic optimization algorithms like SGD and Adam, which are more representative of real-world practices than full-batch gradient descent.

Key Findings on PLRF Models

The simulation reveals that even simple models like PLRF can exhibit training spikes. The power-law exponent α, which quantifies the “heavy-tailedness” of the data distribution, is critical; a larger α is a necessary condition for these spikes to occur when comparing PLRF training trajectories across different optimizers, learning rates, and power-law exponents.

The findings highlight two conditions for training spikes in PLRF: (1) heavy-tailed data (large α) and (2) near divergence (large learning-rate) when comparing the validation loss of PLRF with different learning rates, they observe that when α = 0.3, as the learning-rate increases, validation loss diverges without showing any spikes. In contrast, when α = 0.9, both optimizers exhibit training spikes before divergence.

Theoretical Framework and Dynamics

The framework aligns each DNN layer with a Random Feature model using a layer-wise quadratic loss. The metric δ w t,s is defined as a quadratic function of wt+s, which can be written as “L[wt+s]+Constant”.

The simulation algorithm uses Prediction 3.5 to synthesize “pseudo gradients” by randomly sampling Gaussians to match the average and 2nd moment, at each step of the simulation. This involves tracking the D-dimensional vector z = c(w − w∗), where (wi)1≤i≤D and (w∗i)1≤i≤D are the representations of w and w∗ under the eigen basis.

Simulation Results and Insights

The simulation successfully predicts spiky behavior for a given model checkpoint in a case study on an Adam hyperparameter configuration (where β2 is set close to 1) known to induce spikes in Transformer LMs. The authors also derive an appropriate β2 schedule from observed metrics, which empirically improves performance when analyzing several training metrics in detail and offering optimization recipes such as an appropriate β2 schedule, which can avoid spikes and empirically improve performance.

Empirically verifying Prediction 3.5, the ratio HY/M, which should be close to 1 according to Equation (3.4), shows that the relative error of Equation (3.4) is less than 0.1, mostly within 5%. Furthermore, in Figure 9, the ratio HY/M is close to 1 as predicted.

Conclusion and Future Directions

The work suggests that quadratic models at least partially reveals the essence of underlying mechanism. The authors urge the community to move from analyzing full-batch gradient descent and worst-case bounds, to the more realistic and challenging settings of stochastic optimization and average-case behavior of DNN training, with heavy-tailed input data and scaling-up of model sizes. The proposed β2 schedule can indeed improve performance in terms of validation loss.

--- Page 38 ---

The paper is a scientific investigation into the underlying mechanisms causing intermittent training spikes in Deep Neural Networks (DNNs) by modeling layer-wise dynamics using simple quadratic Random Feature models and simulating training behavior under stochastic optimization.

Conclusion and Future Directions

The work suggests that quadratic models at least partially reveals the essence of underlying mechanism.

1. Improved Training Spike Prediction: The system can predict bumps indicating spiky behavior from a given checkpoint by using a simulation algorithm that models per-layer training dynamics using Power-Law Random Feature (PLRF) models, as shown in Section 4 and Figure 3.

2. Empirical Hyperparameter Optimization: The system can derive an appropriate Adam hyperparameter schedule by analyzing observed metrics, specifically proposing the recipe: set (C.3) β2 ≈ 1 − s / ηt T√N, which empirically improves performance in validation loss (Table 3).

3. Robust Training Dynamics Understanding: The system can analyze the underlying cause of training spikes by linking them to heavy-tailed data distributions, noting that a larger α is a necessary condition for these spikes to occur (Section 2) and confirming that training spikes can occur at different steps between 'Spiky1' and 'Spiky2' (Figure 1).

4. Layer-wise Noise Characterization: The system can diagnose the source of instability by analyzing the evolution of eigenvalue distributions, noting that eigenvalues from Spiky are more heavy-tailed (Figure 7), providing an indicator of healthy training dynamics.

5. Adaptive Batch-Size Heuristics: The system can suggest heuristics for tuning the batch-size to minimize gradient noise, stating that one can empirically measure the third factor and tune the batch-size to make B ≈ 1 (Section C.4).

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 β 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 β 2 schedule from observed metrics, which empirically improves performance.