LLM pre-training dynamics often occur near critical learning stability points
Does LLM Pre-Training Typically Occur at the Edge of Stability?
Machine Learning
Summary
Training large language models (LLMs) involves adjusting many small parts to make them work better. Usually, researchers use a simple math tool called quadratic approximation to understand how these adjustments behave. The authors found that this simple tool sometimes doesn’t match what actually happens during training, especially with big models like LLMs. They identified two important behaviors during training—one that always happens and another that mostly happens depending on learning settings—which helps explain when these math tools fail and when they are useful.
What this means in practice
- •For machine learning engineers: Design better training schedules by understanding when model dynamics reach critical stability points during LLM pre-training.
- •For data center operations teams: Adjust batch sizes and learning rates during large model training to maintain stable resource use and improve training reliability.
Authors
Yuhang Cai, Haodong Wen, Haofeng Huang, Deyi Liu, Yiyuan Ma, Kaifeng Lyu
Abstract
Quadratic approximations provide a common lens for studying neural network optimization, but recent evidence challenges their predictive validity. In full-batch gradient descent with learning rate (LR) $η$, Cohen et al. (2021) observed the Edge of Stability (EoS), where the largest Hessian eigenvalue concentrates near $2/η$, seemingly at odds with classical stability conditions predicted by quadratic approximations. In this work, we revisit the fidelity of quadratic approximations as a model of neural network training dynamics, with a particular focus on their failure modes in LLM training. We first highlight a distinct failure mode of the quadratic approximation, where persistent negative curvature causes the quadratic dynamics to diverge regardless of the LR, even though the real dynamics remain stable. We term this phenomenon the Edge of Convexity (EoC). After removing the effect of this LR-independent instability, we then propose a generalized definition of Edge of Stability (EoS) for stochastic training with adaptive optimizers, based on whether increasing the LR induces significant divergence in the quadratic dynamics while leaving the real dynamics largely unchanged. Across different LLM pretraining runs with various model sizes up to 787M, we find that (1) EoC is always observed across LLM pretraining, and (2) EoS is also prevalent but not universal: it disappears when the LR becomes sufficiently small (e.g., after decay) or when the batch size falls below a critical threshold that has a positive correlation with the critical batch size. Together, these findings characterize when and how quadratic approximations fail and provide a foundation for future work on understanding training dynamics.