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

summary

The gist

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

In short

The study investigates intermittent loss spikes during Large Language Model pre-training by using simple quadratic models with heavy-tailed inputs, called Power-Law Random Feature (PLRF) models. The authors developed a simulation framework to model layer dynamics and found that these spikes occur under specific conditions related to data distribution and learning rates. They also derived an optimal $\beta_2$ schedule for Adam optimization that empirically improved performance.

Key concepts

Power-Law Random Feature (PLRF) Models: These are simple linear regression models with a quadratic loss function used to approximate the dynamics of individual layers in a Deep Neural Network. They are employed because they can replicate the training spikes observed in complex DNNs when trained on heavy-tailed data distributions.
Heavy-Tailed Data (Large $\alpha$): This refers to data where extreme values occur more frequently than expected, characterized by a large power-law exponent $\alpha$. The paper shows that for training spikes to appear in PLRF models, the data must possess this heavy-tailed characteristic.
$\beta_2$ Schedule (Adam Optimization): $\beta_2$ is a hyperparameter in the Adam optimization algorithm that controls how quickly the second moment estimate adapts. The authors derived an appropriate schedule for $\beta_2$ based on observed training metrics, which empirically led to better validation loss performance by avoiding detrimental spikes.

Terminology used across episodes

Reinforcement Learning (RL) — explained in 25 episodes
Sample Complexity — explained in 20 episodes
Retrieval-Augmented Generation (RAG) — explained in 19 episodes
Federated Learning — explained in 17 episodes
Flow Matching — explained in 17 episodes
Quantum Fisher Information (QFI) — explained in 16 episodes

This episode discusses

The paper

Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models · Read on arXiv

Ran Tian, Courtney Paquette

Google DeepMind

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.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models".

Jane: The gist The paper shows that simple quadratic models with heavy-tailed,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we’re looking at this paper, "Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models." It tackles that mystery about intermittent loss spikes in LLM pre-training.

Jane: The core idea they present is that simple quadratic models with heavy-tailed, power-law distributed inputs can actually replicate these intermittent “spikes” in loss, which is what they call training spikes.

Lu: They set up a theoretical framework and a simulation algorithm to model the per-layer training dynamics of Deep Neural Networks using those quadratic models.

Meng: The authors use this framework to explore why these spikes happen, moving away from just looking at the overall network behavior during training.

Lalam: It’s about modeling what happens within each layer, which helps them pinpoint where the instability is coming from.

The paper's summary: Tom: They summarize their contributions by first developing this framework that aligns DNN layers with simple Random Feature models to enable lightweight simulations of training dynamics.

Jane: That’s significant because it allows them to scientifically study these elusive phenomena like training spikes in DNNs, which is something they say is surprising because even convex quadratic models show similar behavior.

Lu: They are showing that spiky behavior isn't just tied to the complex, non-convex nature of feature learning in DNNs, as some research suggests, but it can happen even with simpler PLRF models.

Meng: They conduct a case study on an Adam hyperparameter configuration where β2 is set close to one in a Transformer LM, and their simulation successfully predicts the spiky behavior for a given model checkpoint <ref:2610.13524#pg1,a case study on an Adam hyperparameter configuration where β2 is set>.

Lalam: Beyond just predicting the spikes, they derive an appropriate β2 schedule from observed metrics, which they claim empirically improves performance when analyzing several training metrics in detail <ref:2610.13524#pg1,derive an appropriate β2 schedule from observed metrics, which>.

The paper's improvements: Tom: The paper suggests several practical improvements. First, they can predict "bumps" indicating spiky behavior from a given model checkpoint using their simulation algorithm and Power-Law Random Feature models, as shown in Section four and Figure three <ref:2610.13524#pg2>.

Jane: They also show that the power-law exponent alpha is critical; a larger alpha is a necessary condition for these spikes to occur when comparing PLRF training trajectories across different optimizers or learning rates.

Lu: They point out two specific conditions for training spikes in PLRF: one, heavy-tailed data, meaning a large alpha, and two, near divergence when comparing the validation loss with different learning rates.

Meng: They observed that when alpha is zero point three and the learning rate increases, the validation loss diverges without showing any spikes; but when alpha is zero point nine and you compare different optimizers, both exhibit training spikes before divergence.

Lalam: They also derive a specific β2 schedule, which they claim empirically improves performance in terms of validation loss when analyzing several metrics in detail <ref:2610.13524#pg1>.

Conclusion: Tom: So to wrap up the "Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models," the paper suggests that quadratic models at least partially reveal the underlying mechanism behind these spikes.

Jane: They urge the community to move away from analyzing full-batch gradient descent and worst-case bounds, moving toward studying stochastic optimization and average-case behavior of DNN training with heavy-tailed data.

Lu: The proposed β2 schedule is presented as a recipe that can indeed avoid spikes in the simulation and empirically improve performance <ref:2610.13524#pg1>.

Meng: It’s important to note that the paper shows the mechanism, but it doesn't solve every training issue; it points us toward better optimization recipes rather than just a single fix.

Lalam: The authors also flag a limitation: they are using quadratic models, which is an approximation, so it only reveals part of the underlying mechanism rather than explaining everything completely.

Tom: So that’s the idea behind this paper on training spikes, using layer-wise random feature models to simulate and predict them. We’ll be looking at how these simulation techniques compare to other methods in our next discussion on arXiv.

More episodes

← Home