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
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× and 1.0× the training LR, but increasing the LR slightly to 1.2× 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.
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, −λmin has a smaller magnitude than λ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.
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.
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=η⋆/η , 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.
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=η⋆/η , 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 λ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.
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×10−4 for 50 epochs. Panels follow Figure 2 , using the raw Hessian and epoch on the horizontal axis.
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.
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 β2=0.95 at batch size 512 , we scale the LR, the weight decay, and β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.
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 λP , and report the loss at local step 25. (a) λP=0 (no compensation). (b) λP≈408 (estimated). (c) λP=600 .
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.
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=η⋆/η , 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≤q≤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 −λmin curve, and the saddleness use the compensation coefficient λ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×10−3 and goes beyond EoS at larger peak LRs. Up to peak LR 5×10−3 , the critical LR factor q=η⋆/η , 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≤q≤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 −λmin curve, and the saddleness use the compensation coefficient λ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×10−3 and goes beyond EoS at larger peak LRs. Up to peak LR 1.25×10−3 , the critical LR factor q=η⋆/η , 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≤q≤1.1 ). Over the same peak LRs, 26 of the 30 checkpoints satisfy the EoC criterion. For peak LRs from 2.5×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 −λmin curve, and the saddleness use the compensation coefficient λ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×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.
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.
Alexandru Meterez, Pranav Ajit Nair, Depen Morwani +3
1Kempner Institute at Harvard University · †: Work done while at Harvard. · 2MIT
Empirically, modern deep learning training often occurs at the Edge of Stability (EoS), where the sharpness of the loss exceeds the threshold below which classical convergence analysis applies. Despite recent progress, existing theoretical explanations of EoS either rely on restrictive assumptions or focus on specific squared-loss-type objectives. In this work, we introduce and study a structural property of loss functions that we term product-stability. We show that for losses with product-stable minima, gradient descent applied to objectives of the form (x,y)↦l(xy) can provably converge to the local minimum even when training in the EoS regime. This framework substantially generalizes prior results and applies to a broad class of losses, including binary cross entropy. Using bifurcation diagrams, we characterize the resulting training dynamics, explain the emergence of stable oscillations, and precisely quantify the sharpness at convergence. Together, our results offer a principled explanation for stable EoS training for a wider class of loss functions.
Muon is increasingly used for language-model pretraining, yet its large-step dynamics are not captured by the classical edge-of-stability (EoS) picture of gradient descent (GD). In GD, loss neutrality, equal-magnitude update reversal, and marginal stability meet at a single learning-rate-dependent edge. We show that Muon breaks this coupling. For stochastic no-momentum Muon, we derive a coherence-corrected conditional loss-neutral boundary 2ρb/η, while temporal alignment follows a separate geometry. Controlled experiments show that loss balance and temporal alignment respond differently to learning rate and batch size. Across our language model experiments, the 130M Llama-like LLM runs exhibit loss-boundary tracking with weak negative alignment, whereas the studied 1B LLM configuration shows stronger partial cancellation; in both settings, directions remain far from coherent reversal while training continues to improve. These results support a split EoS picture for Muon: a stochastic loss-neutral edge survives, but it is not accompanied by a universal temporal-direction signature. The source code for reproducing the experiments can be found in https://github.com/cyzebra/Muon-Sublates-the-Edge-of-Stability-in-LLM-Pretraining
Yanzhe Chen, Qifang Zhao, Xiaoxiao Xu +1
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