Does LLM Pre-Training Typically Occur at the Edge of Stability?

Yuhang Cai; Haodong Wen; Haofeng Huang; Deyi Liu; Yiyuan Ma; Kaifeng Lyu

cs.LGOct 27, 2026

Does LLM Pre-Training Typically Occur at the Edge of Stability?

Authors: Yuhang Cai, Haodong Wen, Haofeng Huang, Deyi Liu, Yiyuan Ma, Kaifeng Lyu

Organizations: UC Berkeley · Tsinghua University · ByteDance Seed

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.

Figures & tables

Panel of Figure 1 — Figure 1: LLM pretraining on FineWeb-10B exhibits two distinct failure modes of the quadratic approximation. (a) EoC phenomenon : regardless of the LR choice, the training dynamics under quadratic approximation diverge due to persistent negative curvature, while the true dynamics remain stable. (b) EoS phenomenon : after compensating for negative curvature, the quadratic loss remains as stable as the true loss for $0.8\times$ and $1.0\times$ the training LR, but increasing the LR slightly to $1.2\times$ causes the quadratic loss to diverge. This shows that the current training LR is close to a critical stability threshold. See Sections 5.1 and 5.2 for details.

Panel of Figure 2 — Figure 2: EoS and EoC appear across model sizes. At most checkpoints, the bottom heatmaps show a sharp transition from bounded displacement to explosion near LR factor $1.0$ , indicating that the training LR is at the critical threshold of EoS. In the middle row, $-\lambda_{\min}$ has a smaller magnitude than $\lambda_{\max}$ , corresponding to saddleness below the $0.1$ reference line in the top plots and consistent with EoC. Each heatmap column is a checkpoint, and the vertical axis gives the simulated LR divided by the training LR. See Section 6.2 for the results and Section B.1 for training and detection settings.

Panel of Figure 3 — Figure 3: Training may leave the EoS regime after LR decay. In (a) and (b), the transition from bounded displacement to explosion moves to larger LR factors during cosine decay and after the abrupt multistep LR drop, respectively. After the larger LR drop in (c), the local simulations remain bounded at the training LR for the rest of the observed run, where training enters the stable regime and does not come back to EoS. See Sections 6.4 , B.4 and B.5 for further discussion.

Panel of Figure 4 — Figure 4: The EoS batch size (EBZ) and the critical batch size (CBZ) show a positive correlation across datasets. (a) Larger EBZ generally corresponds to larger CBZ, and the dashed line is a linear fit. (b) The critical LR factor $q=\widehat{\eta}_{\star}/\eta$ , averaged over checkpoints, decreases as batch size increases. Stars mark the interpolated crossings of $q=1.1$ , which define EBZ. See Sections 6.3 and B.3 for the experimental protocol and threshold estimation.

Panel of Figure 5 — Figure 5: Training may go beyond EoS and enter an extremely unstable regime with large LR. (a) As the peak LR increases, the critical LR factor $q=\widehat{\eta}_{\star}/\eta$ , averaged over the tested checkpoints, falls from near the reference line at $1.0$ to well below it. Values near $1.0$ indicate EoS, while smaller values mean that the training LR exceeds the estimated critical LR. (b) These large peak LRs also yield higher final evaluation loss. See Sections 6.2 and 6.4 for further discussion.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Mini-batch Lanczos estimates are stable for the top eigenpair but unstable for the bottom eigenpair. We run Lanczos on ten independently sampled mini-batches at the same checkpoint for the Adam-preconditioned Hessian. The maximum eigenvectors have high pairwise absolute cosine and similar eigenvalues across batches. In contrast, the minimum eigenvectors are nearly orthogonal across batches and the corresponding $\lambda_{\min}$ values vary substantially, showing that per-batch Lanczos estimates are dominated by batch-specific negative-curvature directions.

Model	#Layers	Hidden Dim	#Heads	Head Dim	Params
GPT-extra-small	4	768	12	64	43M
GPT-small	12	768	12	64	104M
GPT-medium	24	1,024	16	64	344M
GPT-large	36	1,280	20	64	787M

Appendix

Table 1: Model architectures used in our experiments.

Panel of Figure 7 — Figure 7: EoS and EoC appears on model trained with SGD. We train ResNet20 and ResNet32 with MSE loss and SGD, using batch size $1024$ and weight decay $5\times 10^{-4}$ for $50$ epochs. Panels follow Figure 2 , using the raw Hessian and epoch on the horizontal axis.

Panel of Figure 8 — Figure 8: EoS and EoC measurements under different LRs. We train GPT-extra-small (43M) on FineWeb-10B using AdamW with the WSD schedule and vary the LR. All runs use 400 linear warmup steps, 7600 stable steps, and 2000 linear decay steps. We identify $0.004$ as the optimal peak LR via a separate sweep.

Panel of Figure 9 — Figure 9: EoS emerges only above a batch size threshold. We train GPT-extra-small (43M) on FineWeb using AdamW with the WSD schedule at global batch sizes $256$ , $512$ , and $1{,}024$ with the same number of training tokens. Starting from LR $0.002$ , weight decay $0.1$ , and $\beta_{2}=0.95$ at batch size $512$ , we scale the LR, the weight decay, and $\beta_{2}$ with the square-root rule of Section B.3 . Each panel follows the reporting protocol of Figure 2 . (a) At batch size $256$ , the transition is near factor $1.0$ only at the first checkpoints, and from step $3{,}600$ on no LR factor in the grid explodes. (b, c) At $512$ and $1{,}024$ , the transition stays close to factor $1.0$ for most of training.

Panel of Figure 10 — Figure 10: Effects of different compensation coefficients. We train GPT-extra-small (43M) with AdamW on FineWeb-10B, run EoS detection with different compensation coefficients $\lambda_{\mathsf{P}}$ , and report the loss at local step 25. (a) $\lambda_{\mathsf{P}}=0$ (no compensation). (b) $\lambda_{\mathsf{P}}\approx 408$ (estimated). (c) $\lambda_{\mathsf{P}}=600$ .

Panel of Figure 11 — Figure 11: Ablations on the multistep LR drop. We train GPT-extra-small (43M) and vary the peak LR and the drop ratio of the multistep schedule. (a)–(c) The peak LR is fixed at $0.002$ and the drop ratio varies over $0.6$ , $0.45$ , and $0.3$ , with the drop at step 4000. (d)–(f) The drop ratio is fixed at $0.7$ and the peak LR varies over $0.008$ , $0.004$ , and $0.0005$ , with the drop at step 3000.

Panel of Figure 12 — Figure 12: Ablations for the linear decay phase. We train GPT-extra-small (43M) and vary the decay length and decay ratio of the linear decay phase. (a)–(c) The decay ratio is fixed at $0.1$ , with decay lengths of 3000, 5000, and 8000 steps. (d)–(f) The decay length is fixed at 6000 steps, with decay ratios of $0.2$ , $0.4$ , and $0.6$ .

Figure 13: Training loss of GPT-large. We compare training loss versus steps at the empirically optimal peak LR (

0.0025

) and the largest tested peak LR (

0.04

). Both runs are trained for 20 tokens per parameter.

Figure 14: GPT-small training stays near EoS up to peak LR $0.01$ and goes beyond EoS at larger peak LRs. Up to peak LR $0.01$ , the critical LR factor $q=\widehat{\eta}_{\star}/\eta$ , the smallest LR factor whose simulation explodes, lies between $0.9$ and $1.2$ at $51$ of the $60$ checkpoints, and $35$ of these are at EoS ( $0.9\leq q\leq 1.1$ ). Over the same peak LRs, $44$ of the $60$ checkpoints satisfy the EoC criterion. For peak LRs from $0.02$ on, the training LR exceeds the estimated critical LR ( $q<0.9$ ) at $19$ of the $20$ checkpoints, and at $5$ of the $10$ checkpoints of peak LR $0.04$ no LR factor in the grid stays bounded. In these two runs, the rows below LR factor $0.6$ , the $-\lambda_{\min}$ curve, and the saddleness use the compensation coefficient $\lambda_{\mathsf{P}}$ estimated at LR factor $0.1$ , while the rows from $0.6$ up use the one estimated at $0.6$ . Because the two coefficients differ, a bounded factor $0.6$ sits above exploding lower factors at one checkpoint of peak LR $0.02$ . Each panel is one peak LR, each heatmap column is one of ten checkpoints in the stable phase, and the black line marks the largest LR factor that stays bounded. Saddleness above $0.7$ runs off the top of the axis, and blank heatmap cells were not simulated. See Figure 2 for the panel format and Section B.6 for the setup.

Figure 15: GPT-medium training stays near EoS up to peak LR $5\times 10^{-3}$ and goes beyond EoS at larger peak LRs. Up to peak LR $5\times 10^{-3}$ , the critical LR factor $q=\widehat{\eta}_{\star}/\eta$ , the smallest LR factor whose simulation explodes, lies between $0.9$ and $1.2$ at $43$ of the $50$ checkpoints, and $31$ of these are at EoS ( $0.9\leq q\leq 1.1$ ). Over the same peak LRs, $38$ of the $50$ checkpoints satisfy the EoC criterion. For peak LRs from $0.01$ on, the training LR exceeds the estimated critical LR ( $q<0.9$ ) at $26$ of the $30$ checkpoints, and at $7$ of the $10$ checkpoints of peak LR $0.04$ no LR factor in the grid stays bounded. In these three runs, the rows below LR factor $0.6$ , the $-\lambda_{\min}$ curve, and the saddleness use the compensation coefficient $\lambda_{\mathsf{P}}$ estimated at LR factor $0.1$ , while the rows from $0.6$ up use the one estimated at $0.6$ . Each panel is one peak LR, each heatmap column is one of ten checkpoints in the stable phase, and the black line marks the largest LR factor that stays bounded. Saddleness above $0.7$ runs off the top of the axis, and blank heatmap cells were not simulated. See Figure 2 for the panel format and Section B.6 for the setup.

Figure 16: GPT-large training stays near EoS up to peak LR $1.25\times 10^{-3}$ and goes beyond EoS at larger peak LRs. Up to peak LR $1.25\times 10^{-3}$ , the critical LR factor $q=\widehat{\eta}_{\star}/\eta$ , the smallest LR factor whose simulation explodes, lies between $0.9$ and $1.2$ at $29$ of the $30$ checkpoints, and $19$ of these are at EoS ( $0.9\leq q\leq 1.1$ ). Over the same peak LRs, $26$ of the $30$ checkpoints satisfy the EoC criterion. For peak LRs from $2.5\times 10^{-3}$ on, the training LR exceeds the estimated critical LR ( $q<0.9$ ) at $44$ of the $50$ checkpoints, and at every checkpoint of peak LR $0.04$ no LR factor in the grid stays bounded. In these five runs, the rows below LR factor $0.6$ , the $-\lambda_{\min}$ curve, and the saddleness use the compensation coefficient $\lambda_{\mathsf{P}}$ estimated at LR factor $0.1$ , while the rows from $0.6$ up use the one estimated at $0.6$ . Because the two coefficients differ, a bounded factor $0.6$ sits above exploding lower factors at two checkpoints of peak LR $2.5\times 10^{-3}$ and one of $0.01$ . Each panel is one peak LR, each heatmap column is one of ten checkpoints in the stable phase, and the black line marks the largest LR factor that stays bounded. Saddleness above $0.7$ runs off the top of the axis, and blank heatmap cells were not simulated. See Figure 2 for the panel format and Section B.6 for the setup.

Figure 17: Summary of a representative training run. We show training dynamics for GPT-extra-small trained on FineWeb-10B using AdamW with LR $0.004$ , weight decay $0.1$ , and the WSD LR schedule.

Edge of Stability

Edge of Stability studies how large-step training approaches curvature-dependent stability limits and how the resulting dynamics affect learning.

See all 37 papers

A Defense of the Quadratic Model
Jul 23, 2026Alexandru Meterez, Pranav Ajit Nair, Depen Morwani +3Language Model Pretraining Edge of Stability
1Kempner Institute at Harvard University · †: Work done while at Harvard. · 2MIT
Product-Stability: Provable Convergence for Gradient Descent on the Edge of Stability
Apr 3, 2026Eric GanDeep Learning Optimization Edge of Stability
Independent Researcher.
Muon Sublates the Edge of Stability in LLM Pretraining
Sep 28, 2026Yanzhe Chen, Qifang Zhao, Xiaoxiao Xu +1Stochastic Optimization Language Model Pretraining
School of Mathematical Sciences, Shanghai Jiao Tong University, China · Alibaba Inc., Hangzhou, China · School of Mathematical Sciences, Institute of Natural Sciences and MOE-LSC, Shanghai Jiao Tong University, China

Jul 23, 2026J/K move · Enter open · S save

A Defense of the Quadratic Model

Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan+2

1Kempner Institute at Harvard University · †: Work done while at Harvard. · 2MIT

Language Model Pretraining Edge of Stability Second-Order Optimization Neural Network Training Dynamics

Open paper AskPDF arXiv

Due to the complexity of neural network loss landscapes, optimization theory is forced to rely on idealized models, and there is generally a tradeoff between how theoretically tractable the model is, and how accurately it describes the true optimization dynamics. In this work, we stress test the simplest possible model of optimization -- the quadratic model -- and show that it can be surprisingly predictive in an LLM setting with 150M parameters and 3B training tokens. Specifically, we show that Taylor expanding the model and the loss function at intermediate checkpoints through training can accurately predict the optimization dynamics over windows that can last up to 10% of training. Having established this agreement, we then turn to analyzing the structure of these local quadratic optimization problems through two lenses: the Hessian spectrum and local stability. Using Lanczos quadrature with extremely deep probes, we are able to estimate the Hessian spectrum deep into the tail, and we find a surprising amount of structure in both the eigenvalues and eigenvectors, which depends on the batch size, preconditioner, and training time. We also empirically test local linear stability at intermediate checkpoints and compare it to theoretical predictions to demonstrate that optimization in LLMs typically occurs at a stochastic edge of stability, whose nature is also determined by batch size. Our results indicate the quadratic model may be a theoretically tractable proxy for pretraining optimization dynamics.

Abstract

Edge of Stability

Explore similar work

A Defense of the Quadratic Model

Product-Stability: Provable Convergence for Gradient Descent on the Edge of Stability

Muon Sublates the Edge of Stability in LLM Pretraining