A Defense of the Quadratic Model

# A Defense of the Quadratic Model
Source: [https://arxiv.org/html/2607.21716](https://arxiv.org/html/2607.21716)
Pranav Ajit NairKempner Institute at Harvard UniversityDepen MorwaniKempner Institute at Harvard UniversityCengiz PehlevanKempner Institute at Harvard University Sham KakadeKempner Institute at Harvard UniversityAlex DamianKempner Institute at Harvard UniversityMIT

###### Abstract

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\.

## 1Introduction

††footnotetext:†\\dagger: Work done while at Harvard\.Large Language Model \(LLM\) optimization has quickly advanced through a feedback loop between fast experimentation and theoretical intuition\. For pretraining, many design choices such as learning rate scheduling, batch\-size scaling, and optimizer are, at least informally, motivated by traditional optimization theory\(Polyak,[1964](https://arxiv.org/html/2607.21716#bib.bib6); Zhanget al\.,[2024b](https://arxiv.org/html/2607.21716#bib.bib7); Polyak and Juditsky,[1992](https://arxiv.org/html/2607.21716#bib.bib8); Meterezet al\.,[2025](https://arxiv.org/html/2607.21716#bib.bib9); Kidambiet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib10)\)\. Such models generally act as a useful guide for reasoning about training dynamics and scaling in a tractable way\. However, it is unclear which of these models are actually descriptive of the dynamics of large\-scale pretraining\.

There are several failure modes that give rise to this gap between theoretical proxies and empirics\. On one hand, models can be overly idealized: they can capture a mathematical phenomenon, but under assumptions that are difficult to justify for modern neural networks\. However, if the model does not incorporate enough structure about the specific optimization problem \(e\.g\. worst case bounds over the set of all smooth convex problems\), it can lead to predictions that are overly pessimistic and do not lead to actionable insights\. Finally, if a model is not overly idealized and incorporates problem\-specific structure, it is often theoretically intractable\.

Arguably the simplest setting in which optimization theory can give sharp problem\-specific predictions is the quadratic model\. The advantage of this setting is that the training dynamics can be understood mode\-by\-mode across the Hessian spectrum, which enables instance\-dependent analyses that are tight for a given Hessian spectrum\. This makes it a useful sandbox for studying stability, learning rate scheduling, batch size scaling, acceleration, and preconditioning, and developing principled algorithms that transfer to neural network optimization\(Meterezet al\.,[2025](https://arxiv.org/html/2607.21716#bib.bib9),[2026](https://arxiv.org/html/2607.21716#bib.bib12); Zouet al\.,[2021](https://arxiv.org/html/2607.21716#bib.bib16); Wuet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib50); Kidambiet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib10); Jainet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib47); Morwaniet al\.,[2026](https://arxiv.org/html/2607.21716#bib.bib13); Zhanget al\.,[2024b](https://arxiv.org/html/2607.21716#bib.bib7),[2019](https://arxiv.org/html/2607.21716#bib.bib75); Paquette and Paquette,[2021](https://arxiv.org/html/2607.21716#bib.bib76); Zhanget al\.,[2024a](https://arxiv.org/html/2607.21716#bib.bib77); Wuet al\.,[2025](https://arxiv.org/html/2607.21716#bib.bib78); Ferbachet al\.,[2026](https://arxiv.org/html/2607.21716#bib.bib79); Zhanget al\.,[2026b](https://arxiv.org/html/2607.21716#bib.bib80); Varreet al\.,[2021](https://arxiv.org/html/2607.21716#bib.bib81); Varre and Flammarion,[2022](https://arxiv.org/html/2607.21716#bib.bib82); Bordelon and Pehlevan,[2021](https://arxiv.org/html/2607.21716#bib.bib83)\)\.

These works largely use the quadratic model as a useful theoretically tractable sandbox for testing algorithms and ideas, and take a leap of faith when applying them to neural networks\. In this work, we aim to tighten this connection and stress test the quadratic model by asking:

*To what extent does the quadratic model actually capture the training dynamics of LLMs?*

In[Section˜3](https://arxiv.org/html/2607.21716#S3), we test this by training a 150M transformer on 3B tokens from FineWeb and save intermediate checkpoints every 10% of training\. We then Taylor expand the model and the loss around each checkpoint and train on this proxy loss for 10% of the total tokens, and compare this to the original optimization trajectory\. We do this for two proxies\. The first,prox, linearizes the model but maintains the cross\-entropy objective, which is equivalent to multi\-class logistic regression under the NTK\(Jacotet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib60); Leeet al\.,[2019](https://arxiv.org/html/2607.21716#bib.bib61); Chizatet al\.,[2019](https://arxiv.org/html/2607.21716#bib.bib62)\)\. The second,quad, Taylor expands this objective to second order so that the loss function is a convex quadratic\. Our results indicate that these approximations can hold for up to 10% of training, and the quality of the approximation is significantly better near the end of training than at the start\.

Motivated by the accuracy of these Taylor expansions, we then turn to describing the structure of these local quadratic problems\. In[Section˜4](https://arxiv.org/html/2607.21716#S4), we estimate the spectra of the Hessian and the Gauss\-Newton matrix along the training trajectory, both with and without the Adam preconditioner\. To do so, we use extremely deep Lanczos quadrature probes with rigorous error bands and we are able to resolve the spectrum over six orders of magnitude\. We find that the spectrum often splits into a head, whose size is determined by the vocab size and is dominated by the unembedding layer, and a power\-law tail that is*universal*: identical across batch sizes and unaffected by the Adam preconditioner\. Based on these curves, we estimate the source and capacity exponents and find a capacity exponent ofα≈1\\alpha\\approx 1in the resolvable tail, and a source exponent ofβ<1\\beta<1, which contradicts the summability assumption common in literature\.

Finally, in[Section˜5](https://arxiv.org/html/2607.21716#S5), we analyze training stability by estimating how close each saved checkpoint is to the*edge of stability*\(Défossez and Bach,[2015](https://arxiv.org/html/2607.21716#bib.bib46); Jainet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib47); Maet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib48); Wuet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib50); Ma and Ying,[2021](https://arxiv.org/html/2607.21716#bib.bib51); Wuet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib53); Velikanovet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib52); Mulayoff and Michaeli,[2024](https://arxiv.org/html/2607.21716#bib.bib49); Cohenet al\.,[2021](https://arxiv.org/html/2607.21716#bib.bib4); Jastrzębskiet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib54); Jastrzebskiet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib55); Cohenet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib5); Lewkowyczet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib57); Ahnet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib58); Aroraet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib56); Damianet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib59); Cohenet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib11)\)\. This is a regime in which the linearized dynamics \(preconditioned SGD\) are incredibly sensitive to the learning rate and batch size, and any increase in the learning rate or decrease in the batch size could cause training to diverge for the linearized model\. We empirically test this by running preconditioned SGD from various checkpoints with a grid of learning rate and batch size multipliers to test which hyperparameters would cause SGD to blow up, and we compare these results to the theoretical predictions made by linear stability analysis\. We observe that at small to medium batch sizes, training occurs in a noise dominated stochastic edge of stability regime, while at large batch sizes, training operates at the deterministic edge of stability as described inCohenet al\.\([2021](https://arxiv.org/html/2607.21716#bib.bib4)\)\. We therefore conclude that the majority of LLM training occurs at the edge of stability, regardless of batch size\.

![Refer to caption](https://arxiv.org/html/2607.21716v1/x1.png)Figure 1:Validation loss curves for 150M models \(black\) trained with cosine for1×1\\timesChinchilla tokens, across batch sizes, at optimal learning rates\. \(Top\) We show agreement between the LLM loss curve when compared to training the linearized model around checkpoints recorded every10%10\\%of training \(blue\)\. \(Bottom\) We further quadratic Taylor expand the loss and train the quadratic model \(orange\)\. The losses are still in agreement for late training times, but diverge at early times, as opposed to the case where we only linearize the model\. Note that the orange and blue curves are trained on the expansions, but we evaluate them on the true model loss\. We measure how long the approximation holds for from each checkpoint in[Table˜1](https://arxiv.org/html/2607.21716#S3.T1)\. We provide further details in Section[3](https://arxiv.org/html/2607.21716#S3), and we also show a similar plot for a constant learning rate with EMA in Figure[12](https://arxiv.org/html/2607.21716#A2.F12)\(Appendix[B](https://arxiv.org/html/2607.21716#A2)\)\.
## 2Related Work

##### Linearized Dynamics\.

The Neural Tangent Kernel \(NTK\)\(Jacotet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib60); Leeet al\.,[2019](https://arxiv.org/html/2607.21716#bib.bib61)\)approximates a network through its first\-order Taylor expansion in parameter space, an approximation which is accurate whenever the parameter displacement is relatively small\. This regime is also known as "lazy" training\(Chizatet al\.,[2019](https://arxiv.org/html/2607.21716#bib.bib62)\)\. However, in practice, the parameters and the NTK evolve rapidly as features form\(Fortet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib63); Atanasovet al\.,[2021](https://arxiv.org/html/2607.21716#bib.bib64); Long,[2021](https://arxiv.org/html/2607.21716#bib.bib65); Vyaset al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib66); Ortiz\-Jiménezet al\.,[2021](https://arxiv.org/html/2607.21716#bib.bib67)\), rendering the NTK dynamics inaccurate for practical learning rates and widths\. Several other works have shown that using the "after kernel", meaning the linearized model at the end of training, is a more faithful approximation for interpretability and fine\-tuning\(Parket al\.,[2023](https://arxiv.org/html/2607.21716#bib.bib68); Long,[2021](https://arxiv.org/html/2607.21716#bib.bib65); Malladiet al\.,[2023](https://arxiv.org/html/2607.21716#bib.bib69)\)\. A complementary line of work studies higher\-order Taylor expansions\(Bai and Lee,[2019](https://arxiv.org/html/2607.21716#bib.bib71); Baiet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib70)\)as a better approximation for the true dynamics\. Relative to this literature, we linearize locally throughout the trajectory rather than at initialization, and we measure the agreement window as a fraction of the total training budget\.

##### Hessian Spectra\.

Neural network Hessian spectra have been previously studied in order to understand the optimization loss landscape\. While for small models it is possible to compute Hessian eigenvalues exactly\(Sagunet al\.,[2016](https://arxiv.org/html/2607.21716#bib.bib27),[2017](https://arxiv.org/html/2607.21716#bib.bib28)\), either by calculating the Hessian in closed form or through a relatively small number of power iterations, this quickly breaks down when parameter counts are in the millions\. A more refined algorithm, which is also the basis of[Section˜4\.2](https://arxiv.org/html/2607.21716#S4.SS2), is the Stochastic Lanczos Quadrature \(SLQ\) algorithm\(Lanczos,[1950](https://arxiv.org/html/2607.21716#bib.bib17); Golub and Meurant,[2009](https://arxiv.org/html/2607.21716#bib.bib22)\)\. Since its introduction, numerous works have studied this algorithm, providing accuracy bounds on the spectral estimates given by the quadrature algorithm\(Ubaruet al\.,[2017](https://arxiv.org/html/2607.21716#bib.bib24); Adamset al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib25); Chenet al\.,[2021](https://arxiv.org/html/2607.21716#bib.bib21),[2022](https://arxiv.org/html/2607.21716#bib.bib26); Linet al\.,[2016](https://arxiv.org/html/2607.21716#bib.bib23)\)\. In neural networks, the first applications of SLQ were in vision models, where several papers reported that Hessians of trained classification tend to separate into a bulk and outlier eigenvalues, roughly equal to the number of classes\(Papyan,[2018](https://arxiv.org/html/2607.21716#bib.bib29),[2019](https://arxiv.org/html/2607.21716#bib.bib30),[2020](https://arxiv.org/html/2607.21716#bib.bib31); Ghorbaniet al\.,[2019](https://arxiv.org/html/2607.21716#bib.bib33)\), an effect related to "neural collapse"\(Papyanet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib32)\)\. Several open\-source toolkits implement SLQ for neural network spectra, including the Tensorflow/JAX accompanyingGhorbaniet al\.\([2019](https://arxiv.org/html/2607.21716#bib.bib33)\), PyHessian\(Yaoet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib34)\), GradVis\(Chatzimichailidiset al\.,[2019](https://arxiv.org/html/2607.21716#bib.bib35)\), and the Deep Curvature suite\(Granziolet al\.,[2019](https://arxiv.org/html/2607.21716#bib.bib36)\)\. More recently, empirical spectral measurements have been used to derive optimization prescriptions in neural network training\.Granziolet al\.\([2022](https://arxiv.org/html/2607.21716#bib.bib37)\)use SLQ to derive an on\-the\-fly hyperparameter learner from the Hessian spectra during training\. In transformers,Zhanget al\.\([2024c](https://arxiv.org/html/2607.21716#bib.bib38)\)found that the Hessian spectrum is very heterogenous across parameter blocks, and further used these findings to design Adam\-mini, a more memory efficient version of Adam\(Zhanget al\.,[2025](https://arxiv.org/html/2607.21716#bib.bib39)\)\. To the best of our knowledge, the first application of SLQ at up to 10B parameter scale for LLMs areGranziol \([2025](https://arxiv.org/html/2607.21716#bib.bib40)\); Granziol and Juarev \([2026](https://arxiv.org/html/2607.21716#bib.bib41)\), running SLQ at a depth of around10−2010\-20probes, which gives a poor resolution in the spectrum tails\. Finally, there is a wide body of literature established on the connections between the neural tangent kernel \(NTK\) eigenvalues, and the Gauss\-Newton\(Nociet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib42); Lauditiet al\.,[2026](https://arxiv.org/html/2607.21716#bib.bib43); Jianget al\.,[2026](https://arxiv.org/html/2607.21716#bib.bib44); Kalraet al\.,[2025](https://arxiv.org/html/2607.21716#bib.bib45)\)\.

##### Linear Stability\.

The dynamics of SGD on a quadratic are a stochastic linear dynamical system\. In the least squares setting with batch size11,Défossez and Bach \([2015](https://arxiv.org/html/2607.21716#bib.bib46)\)showed that the expected loss diverges when the step size exceeds a critical threshold determined by the fourth moment of the data\.Jainet al\.\([2018](https://arxiv.org/html/2607.21716#bib.bib47)\); Maet al\.\([2018](https://arxiv.org/html/2607.21716#bib.bib48)\)later extended this to larger batch sizes and showed that the maximal stable learning rate grows linearly in the batch size up to a critical batch size\. Motivated by the fact that at a minimizer, the loss landscape is locally convex,Wuet al\.\([2018](https://arxiv.org/html/2607.21716#bib.bib50)\); Ma and Ying \([2021](https://arxiv.org/html/2607.21716#bib.bib51)\); Wuet al\.\([2022](https://arxiv.org/html/2607.21716#bib.bib53)\); Velikanovet al\.\([2022](https://arxiv.org/html/2607.21716#bib.bib52)\); Mulayoff and Michaeli \([2024](https://arxiv.org/html/2607.21716#bib.bib49)\)applied this stability analysis to describe which minima can be attractors for SGD\. Linear stability in neural networks has also been extensively studied in the full batch setting under the name*edge of stability*\.Cohenet al\.\([2021](https://arxiv.org/html/2607.21716#bib.bib4)\)showed that the largest Hessian eigenvalue typically increases during GD training until it reaches the critical stability threshold2/η2/\\etaafter which it oscillates for the remainder of training around this value\. This extends empirical observations made byJastrzębskiet al\.\([2018](https://arxiv.org/html/2607.21716#bib.bib54)\); Jastrzebskiet al\.\([2020](https://arxiv.org/html/2607.21716#bib.bib55)\)\.Cohenet al\.\([2022](https://arxiv.org/html/2607.21716#bib.bib5)\)also studied the edge of stability phenomenon for adaptive optimizers like Adam and showed that the right notion of local stability is given by the eigenvalues of the preconditioned Hessian\. Several mechanisms have been proposed for how training is sustained at the edge of stability including the catapult effect for large initial learning rates\(Lewkowyczet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib57)\)a self\-stabilization effect resulting from the cubic term in the loss Taylor expansion\(Damianet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib59)\), and progressive sharpening driven by normalization\(Aroraet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib56)\)\. More recently,Cohenet al\.\([2024](https://arxiv.org/html/2607.21716#bib.bib11)\)have proposed an ODE capturing the dynamics of GD and adaptive optimizers at the edge of stability which empirically captures the true training dynamics\. In the case of noise, a separate line of work has analyzed the dynamics of SGD and momentum at the edge of stability\(Andreyev and Beneventano,[2024](https://arxiv.org/html/2607.21716#bib.bib74); Andreyevet al\.,[2026](https://arxiv.org/html/2607.21716#bib.bib73)\)\.

Most closely related to our stability experiments in[Section˜5](https://arxiv.org/html/2607.21716#S5), concurrent work byCaiet al\.\([2026](https://arxiv.org/html/2607.21716#bib.bib99)\)studies the edge of stability phenomenon in LLMs at scales up to 1\.7B parameters\. They probe the edge of stability by perturbing the learning rate at a fixed batch size, and observe that large batch training occurs at the edge of stability while training well below the critical batch size does not\. However, because their probes only perturb the learning rate, they cannot distinguish instability caused by discretization from instability caused by stochasticity\. Jointly perturbing\(η,B\)\(\\eta,B\)allows us to separate the two, and we find that at small batch sizes training occurs at a stochastic rather than deterministic edge of stability\. Our setting also differs in that we use the linearized model and the corresponding Gauss\-Newton matrices as our convex quadratic proxy, whileCaiet al\.\([2026](https://arxiv.org/html/2607.21716#bib.bib99)\)instead add a proximal L2 penalty to the local problem which shifts the negative eigenvalues up to zero\.

##### Prescriptive Optimization Theory\.

Lastly, we acknowledge two complementary lines of work striving to connect optimization theory with training practice, providing theoretically grounded prescriptions for hyperparameter selection in pretraining\. The first shows that convergence bounds from convex optimization theory can be used to derive practical prescriptions for hyperparameter scaling in LLMs\(Schaippet al\.,[2025](https://arxiv.org/html/2607.21716#bib.bib84); Hägeleet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib85); Shulginet al\.,[2026](https://arxiv.org/html/2607.21716#bib.bib86); Defazioet al\.,[2023](https://arxiv.org/html/2607.21716#bib.bib93); Defazio,[2026](https://arxiv.org/html/2607.21716#bib.bib94); Defazioet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib95)\)\. The second models the discrete optimizer dynamics through stochastic differential equations, yielding scaling rules for hyperparameters across batch sizes and principled comparisons between adaptive methods\(Liet al\.,[2017](https://arxiv.org/html/2607.21716#bib.bib87); Malladiet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib88); Compagnoniet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib89),[2026](https://arxiv.org/html/2607.21716#bib.bib90),[2023](https://arxiv.org/html/2607.21716#bib.bib91); Zhanget al\.,[2026a](https://arxiv.org/html/2607.21716#bib.bib92)\)\.

## 3Quadratics are Predictive for Pretraining

Taylor’s theorem is a central tool for studying smooth functions, as it allows complicated nonconvex objectives to be approximated locally by polynomials, which are often far easier to analyze\. Studying arbitrary nonconvex objectives directly is often too broad to yield sharp predictions for real world training: worst\-case analyses have to account for pathological functions that may be irrelevant to the training runs observed in practice\. Taylor’s theorem provides a natural way to make the problem local\. Around a point on the training trajectory, a complicated objective can be approximated by a quadratic model whose behavior is governed by the local gradient and curvature\. This turns the question from one about arbitrary nonconvex optimization into a concrete empirical test: do local quadratic expansions during pretraining accurately predict the complete training dynamics?

### 3\.1Local models

Throughout the manuscript, we consider three related losses, suppressing the dependence on the minibatch for notational simplicity\. Letℓ\\ellbe cross entropy loss,θ\\thetabe the parameters, andfθf\_\{\\theta\}be the model on a single sample so the loss isℒ​\(θ\)=ℓ​\(fθ\)\\mathcal\{L\}\(\\theta\)=\\ell\(f\_\{\\theta\}\)\. Our first simplified loss,prox, linearizesfθf\_\{\\theta\}around a reference pointθref\\theta\_\{\\mathrm\{ref\}\}but maintains the cross\-entropy objective:

ℒprox​\(θ\):=ℓ​\(fθlin\)​where​fθlin=fθref\+∇θfθref​\(θ−θref\)\.\\mathcal\{L\}\_\{\\mathrm\{prox\}\}\(\\theta\):=\\ell\(f^\{\\mathrm\{lin\}\}\_\{\\theta\}\)\\mbox\{\\quad where\\quad\}f^\{\\mathrm\{lin\}\}\_\{\\theta\}=f\_\{\\theta\_\{\\mathrm\{ref\}\}\}\+\\nabla\_\{\\theta\}f\_\{\\theta\_\{\\mathrm\{ref\}\}\}\(\\theta\-\\theta\_\{\\mathrm\{ref\}\}\)\.This simplification, also known as the NTK approximation\(Jacotet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib60); Chizatet al\.,[2019](https://arxiv.org/html/2607.21716#bib.bib62)\)or the GN\-proximal loss\(Burke,[1985](https://arxiv.org/html/2607.21716#bib.bib1); Drusvyatskiy,[2017](https://arxiv.org/html/2607.21716#bib.bib2); Abreuet al\.,[2025](https://arxiv.org/html/2607.21716#bib.bib3)\), reduces the non\-convex dynamics to solving a convex multi\-class logistic regression problem where the samples are given by the network Jacobians atθref\\theta\_\{\\mathrm\{ref\}\}\.

We can further simplify by Taylor expanding the proximal loss to second order to getquad:

ℒquad​\(θ\):=ℒ​\(θref\)\+∇ℒ​\(θref\)​\(θ−θref\)\+12​\(θ−θref\)T​G​\(θ−θref\)\\mathcal\{L\}\_\{\\mathrm\{quad\}\}\(\\theta\):=\\mathcal\{L\}\(\\theta\_\{\\mathrm\{ref\}\}\)\+\\nabla\\mathcal\{L\}\(\\theta\_\{\\mathrm\{ref\}\}\)\(\\theta\-\\theta\_\{\\mathrm\{ref\}\}\)\+\\tfrac\{1\}\{2\}\(\\theta\-\\theta\_\{\\mathrm\{ref\}\}\)^\{T\}G\(\\theta\-\\theta\_\{\\mathrm\{ref\}\}\)whereG=∇2ℒprox​\(θref\)G=\\nabla^\{2\}\\mathcal\{L\}\_\{\\mathrm\{prox\}\}\(\\theta\_\{\\mathrm\{ref\}\}\)is the Gauss\-Newton matrix atθref\\theta\_\{\\mathrm\{ref\}\}\.

Note thatℒquad​\(⋅\)\\mathcal\{L\}\_\{\\mathrm\{quad\}\}\(\\cdot\)is not the usual second order Taylor expansion, as we use the Gauss\-Newton matrixG=∇2ℒproxG=\\nabla^\{2\}\\mathcal\{L\}\_\{\\mathrm\{prox\}\}, rather than the HessianH=∇2ℒH=\\nabla^\{2\}\\mathcal\{L\}\. This corresponds to keeping the positive semidefinite \(PSD\) term in the common Hessian decomposition:

H​\(θ\)=G​\(θ\)⏟⪰0\+∇fℓ​\(fθ\)T​∇θ2fθ\.\\displaystyle H\(\\theta\)\\quad=\\quad\\underbrace\{G\(\\theta\)\}\_\{\\succeq 0\}\\quad\+\\quad\\cancel\{\\nabla\_\{f\}\\ell\(f\_\{\\theta\}\)^\{T\}\\nabla^\{2\}\_\{\\theta\}f\_\{\\theta\}\}\.Keeping only the Gauss–Newton term ensures that the quadratic objective is convex: its curvature matrix is positive semidefinite, so the local proxy cannot become unbounded below along negative\-curvature directions\. Thus, our quadratic model should be viewed as the second\-order approximation induced by the linearized network, rather than the full Taylor expansion of the original model\. We compare the Hessian and Gauss–Newton spectra in Section[4\.2](https://arxiv.org/html/2607.21716#S4.SS2), and leave a more detailed study of the dropped term to future work\.

### 3\.2Definitions and Experimental Methodology

We train a 150M \(non\-embedding\) parameter model — 168M with embedding and unembedding matrices — on 3B tokens from FineWeb at batch sizesB=1,64,1024B=1,64,1024, sequence length10241024, at the optimally tuned learning rates for each batch size \(see[Figure˜8](https://arxiv.org/html/2607.21716#A1.F8)for the full sweep\)\. We train our models using Adam with no weight decay, warmup of10%10\\%,ϵ=10−8\\epsilon=10^\{\-8\},β1=0\.9\\beta\_\{1\}=0\.9and we dynamically setβ2\\beta\_\{2\}at steptttomax⁡\(0\.95,1−10\.01​t\)\\max\(0\.95,1\-\\frac\{1\}\{0\.01t\}\), so the second moment bufferν\\nuaverages over the last 1% of tokens\. For evaluation, we also track exponential moving averages \(EMA\) of both the parameters and the second moment buffer with parameterγt=max⁡\(0,1−10\.04​t\)\\gamma\_\{t\}=\\max\(0,1\-\\frac\{1\}\{0\.04t\}\)which averages over the last 4% of tokens\. For cosine, we decay to10%10\\%of the peak learning rate\.

For the architecture, we use a standard OLMO\-style transformer with embedding dimension10241024, depth1212,1616heads\. We also use the CompleteP\(Deyet al\.,[2026](https://arxiv.org/html/2607.21716#bib.bib19)\)parameterization to set the per\-coordinate learning rates for Adam\. To ensure our experiments are not dominated by embedding layers, we set our vocab size to81928192by training our own byte\-pair encoding \(BPE\) tokenizers\. This ensures that the embedding and unembedding layers take up 10% of the total parameters, which matches the ratio for many larger models\(Liuet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib72)\)\.111Using the standard GPT\-2 vocab size of 50257 would cause the embedding and unembedding layers to take up 40% of the total parameters, which is extremely unrealistic at larger scale\.We note that with CompleteP multipliers, the learning rates appear much larger than those standard in the literature\. This is because the intermediate layers in the model have a multiplier of1/D1/Dover the base learning rate, whereD=1024D=1024\. Thusη=1\\eta=1with CompleteP approximately corresponds toη=0\.001\\eta=0\.001under standard parameterization\.

### 3\.3Pretraining as a Chain of Local Quadratic Problems

Table 1:Percentage of the full training token budget for which the linearized loss predicts the true loss decrease from the starting checkpoint to within10%10\\%, i\.e\.,\|\(ℒ^t−ℒt\)/\(ℒref−ℒt\)\|≤0\.1\\left\|\(\\hat\{\\mathcal\{L\}\}\_\{t\}\-\\mathcal\{L\}\_\{t\}\)/\(\\mathcal\{L\}\_\{\\mathrm\{ref\}\}\-\\mathcal\{L\}\_\{t\}\)\\right\|\\leq 0\.1\. Each continuation spans10%10\\%of training, so the maximum entry is1010\.For each pretraining run, we save checkpointsθk\\theta\_\{k\}at every10%10\\%of the total token budget\. Starting from eachθk\\theta\_\{k\}, we continue training for the next10%10\\%of the original budget, but replace the loss with its local expansion aroundθref=θk\\theta\_\{\\mathrm\{ref\}\}=\\theta\_\{k\}: the optimizer is run onℒprox\\mathcal\{L\}\_\{\\mathrm\{prox\}\}for theproxexperiments and onℒquad\\mathcal\{L\}\_\{\\mathrm\{quad\}\}for thequadexperiments\. To isolate the effect of the loss approximation, everything else is kept identical to the original run: Adam is initialized from the optimizer state at checkpointkk, we follow the original learning rate schedule, and we maintain the same4%4\\%parameter EMA used for validation\. After this continuation, we evaluate the resulting EMA parameters on the*true*nonlinear model and report the full validation lossℒ​\(⋅\)\\mathcal\{L\}\(\\cdot\)\.

We show our results in[Figure˜1](https://arxiv.org/html/2607.21716#S1.F1)for cosine decay, and in[Figure˜12](https://arxiv.org/html/2607.21716#A2.F12)\([Appendix˜B](https://arxiv.org/html/2607.21716#A2)\) for a constant learning rate with EMA\. While early time linearizations quickly diverge from the loss, middle to late time expansions match the ground truth for roughly55–10%10\\%of the total training budget\. In general,proxmatches better thanquadat early times as well, across all three batch sizes, andquadonly begins to match after60%60\\%of training\. We leave a further study of why this is the case to future work\.

## 4The Hessian of LLMs

In[Section˜3](https://arxiv.org/html/2607.21716#S3), we showed that quadratic Taylor expansions are surprisingly predictive over short timescales\. This is incredibly powerful, because we have a broad toolbox for studying and understanding optimization on quadratics\. Importantly, these are not*arbitrary*quadratic objectives, and instead possess special structure inherited from the local loss landscape\. In this section, we will shed light on this structure by estimating properties of the population Hessian, which is arguably the most important quantity for quadratic optimization\.

Because Adam is a preconditioned method, we distinguish between the*raw*HessianHHand the*preconditioned*HessianHP:=P−1​HH\_\{P\}:=P^\{\-1\}HwhereP:=diag​\[ν\+ϵ\]P:=\\mathrm\{diag\}\[\\sqrt\{\\nu\}\+\\epsilon\]is the Adam preconditioner\.222We also bake in the CompleteP\(Deyet al\.,[2026](https://arxiv.org/html/2607.21716#bib.bib19)\)multipliers into both the raw Hessian and preconditioned Hessian so that the contribution of different layers to the Hessian spectrum remains stable across scales\.Note that the preconditioned Hessian has real eigenvalues asHPH\_\{P\}is similar toP−1/2​H​P−1/2P^\{\-1/2\}HP^\{\-1/2\}, which is symmetric\. The raw Hessian governs the behavior of SGD, while the preconditioned Hessian governs the behavior of preconditioned methods like Adam\(Cohenet al\.,[2022](https://arxiv.org/html/2607.21716#bib.bib5)\)\.

We also distinguish between the Hessian of the linearized loss, the Gauss\-Newton matrixGG, and the Hessian of the original loss, the loss HessianHH\. We will useG,HG,Hto refer to the raw Gauss\-Newton and Hessian matrices, and useGP,HPG\_\{P\},H\_\{P\}to refer to the preconditioned versions\.

### 4\.1Methodology

Algorithm 1Lanczos quadrature with reorthogonalization1:Symmetric matrix\-vector product

v↦H​vv\\mapsto Hv, probe

v0v\_\{0\}, depth

mm
2:

q1←v0/‖v0‖q\_\{1\}\\leftarrow v\_\{0\}/\\norm\{v\_\{0\}\}
3:for

j=1,…,mj=1,\\ldots,mdo

4:

z←H​qjz\\leftarrow Hq\_\{j\}
5:

αj←qj⊤​z\\alpha\_\{j\}\\leftarrow q\_\{j\}^\{\\top\}z
6:for

i=1,…,ji=1,\\ldots,jdo

7:

z←z−qi​\(qi⊤​z\)z\\leftarrow z\-q\_\{i\}\(q\_\{i\}^\{\\top\}z\)
8:endfor

9:

βj←‖z‖\\beta\_\{j\}\\leftarrow\\norm\{z\}
10:

qj\+1←z/βjq\_\{j\+1\}\\leftarrow z/\\beta\_\{j\}
11:endfor

12:

Tm←tridiag⁡\(β1,…,βm−1;α1,…,αm\)T\_\{m\}\\leftarrow\\operatorname\{tridiag\}\(\\beta\_\{1\},\\ldots,\\beta\_\{m\-1\};\\alpha\_\{1\},\\ldots,\\alpha\_\{m\}\)
13:Diagonalize

Tm=U​diag​\(ϑ\)​U⊤T\_\{m\}=U\\mathrm\{diag\}\(\\vartheta\)U^\{\\top\}
14:

ω←U1,:⊙2\\omega\\leftarrow U\_\{1,:\}^\{\\odot 2\},

V←Q​UV\\leftarrow QUwhere

Q=\[q1,…,qm\]Q=\[q\_\{1\},\\ldots,q\_\{m\}\]
15:returnRitz values, weights and vectors

\{\(ϑi,ωi,V:,i\)\}i=1m\\\{\(\\vartheta\_\{i\},\\omega\_\{i\},V\_\{:,i\}\)\\\}\_\{i=1\}^\{m\}

Algorithm 2Gauss–Radau quadrature1:Tridiagonal Lanczos matrix

TmT\_\{m\}, final residual

βm\\beta\_\{m\}, dimension

nn, target value

λ∉spec​\(Tm\)\\lambda\\not\\in\\mathrm\{spec\}\(T\_\{m\}\)
2:

γ←λ\+βm2​em⊤​\(Tm−λ​I\)−1​em\\gamma\\leftarrow\\lambda\+\\beta\_\{m\}^\{2\}e\_\{m\}^\{\\top\}\(T\_\{m\}\-\\lambda I\)^\{\-1\}e\_\{m\}
3:

T←\(Tmβm​emβm​em⊤γ\)T\\leftarrow\\begin\{pmatrix\}T\_\{m\}&\\beta\_\{m\}e\_\{m\}\\\\ \\beta\_\{m\}e\_\{m\}^\{\\top\}&\\gamma\\end\{pmatrix\}
4:

T=U​diag⁡\(ϑ\)​U⊤T=U\\operatorname\{diag\}\(\\vartheta\)U^\{\\top\}
5:

wk←n​U1​k2w\_\{k\}\\leftarrow nU\_\{1k\}^\{2\}for

k=1,…,m\+1k=1,\\ldots,m\+1
6:

i−​\(λ\)←∑k:ϑk\>λwki\_\{\-\}\(\\lambda\)\\leftarrow\\sum\_\{k:\\vartheta\_\{k\}\>\\lambda\}w\_\{k\}
7:

i\+​\(λ\)←i−​\(λ\)\+wji\_\{\+\}\(\\lambda\)\\leftarrow i\_\{\-\}\(\\lambda\)\+w\_\{j\}where

ϑj=λ\\vartheta\_\{j\}=\\lambda
8:returninterval

\[i−​\(λ\),i\+​\(λ\)\]\[i\_\{\-\}\(\\lambda\),i\_\{\+\}\(\\lambda\)\]

Naively, the Hessian spectrum is impossible to compute because a model withppparameters has a HessianH∈ℝp×pH\\in\\mathbb\{R\}^\{p\\times p\}, which is far too large to store or eigen\-decompose\. However, we can efficiently compute Hessian\-vector productsv→H​vv\\to Hvwhich allows us to sketch the Hessian and estimate its spectrum using Lanczos quadrature\. For details, we refer the reader toGolub and Meurant \([2009](https://arxiv.org/html/2607.21716#bib.bib22)\), but we describe the basic idea here\. Given a starting vectorv0v\_\{0\}and access to matrix\-vector productsv→H​vv\\to Hv, Lanczos sketches the distributionμ​\(v0\)\\mu\(v\_\{0\}\)whereμ​\(v\)\\mu\(v\)is defined by:

μ​\(v\):=∑i=1p\(v⋅ui\)2​δλi\\displaystyle\\mu\(v\):=\\sum\_\{i=1\}^\{p\}\(v\\cdot u\_\{i\}\)^\{2\}\\delta\_\{\\lambda\_\{i\}\}and\{\(λi,ui\)\}\\\{\(\\lambda\_\{i\},u\_\{i\}\)\\\}are the eigenvalues and eigenvectors ofHH\. This can be interpreted as a reweighted empirical spectral distribution in which each eigenvalueλi\\lambda\_\{i\}is given probability\(v⋅ui\)2\(v\\cdot u\_\{i\}\)^\{2\}\.

UsingmmHessian vector products, the Lanczos recurrence returnsmmRitz valuesϑ1,…,ϑm\\vartheta\_\{1\},\\ldots,\\vartheta\_\{m\}andmmweightsω1,…,ωm\\omega\_\{1\},\\ldots,\\omega\_\{m\}which sum to11, and produces the estimate:

μ​\(v0\)≈∑i=1mωi​δϑi\.\\displaystyle\\mu\(v\_\{0\}\)\\approx\\sum\_\{i=1\}^\{m\}\\omega\_\{i\}\\delta\_\{\\vartheta\_\{i\}\}\.This approximation is guaranteed to match up to the first2​m−12m\-1moments\. An imprecise but useful interpretation is that each Ritz valueϑi\\vartheta\_\{i\}is a representative for a group of eigenvalues\{λj\}j∈ℐ\\\{\\lambda\_\{j\}\\\}\_\{j\\in\\mathcal\{I\}\}, and the corresponding weightωi≈∑j∈ℐ\(v0⋅ui\)2\\omega\_\{i\}\\approx\\sum\_\{j\\in\\mathcal\{I\}\}\(v\_\{0\}\\cdot u\_\{i\}\)^\{2\}is the total weight of this group of eigenvalues\.

In this section we use the Lanczos recurrence to inspect three properties of the Hessian:

1. 1\.Spectral Decay:estimate the mapi→λii\\to\\lambda\_\{i\}to look for power laws in the spectrum
2. 2\.Eigenvector Distribution:estimate the mapi→‖Πℓ​ui‖2i\\to\\norm\{\\Pi\_\{\\ell\}u\_\{i\}\}^\{2\}whereΠℓ\\Pi\_\{\\ell\}denotes the orthogonal projection onto layers of typeℓ\\ell\(embed,unembed,attn\.q/k/v,mlp\.up/down\)
3. 3\.Gradient\-Hessian Alignment:estimate the mapi→\(∇ℒ​\(θ\)⋅ui\)2i\\to\(\\nabla\\mathcal\{L\}\(\\theta\)\\cdot u\_\{i\}\)^\{2\}

##### Experimental Details

For \(1\) and \(2\) we use Lanczos probes with randomv0∼Sd−1v\_\{0\}\\sim S^\{d\-1\}so that𝔼v0​μ​\(v0\)\\mathbb\{E\}\_\{v\_\{0\}\}\\mu\(v\_\{0\}\)is the empirical spectral distribution ofHH\. For \(3\) we run an additional gradient\-start probe withv0=∇ℒ​\(θ\)v\_\{0\}=\\nabla\\mathcal\{L\}\(\\theta\)\. For each of the checkpointsB=1,64,1024B=1,64,1024andT=10%,50%,100%T=10\\%,50\\%,100\\%we run these Lanczos probes on the preconditioned Gauss NewtonGPG\_\{P\}, the preconditioned HessianHPH\_\{P\}, the raw Gauss NewtonGG, and the raw HessianHH\. All Lanczos probes were run with a probe depth ofm=1200m=1200and stored the full Lanczos basis with two\-pass modified Gram–Schmidt reorthogonalization in fp32 precision\. The matrix\-vector products were computed over 10M tokens\.

##### Accuracy and Ablations

These large scale Lanczos runs require a lot of memory and compute, so we ran ablations to test which parameters were most important:

- •Probe count:We empirically found that a single random probe was sufficient to resolve the spectrumi→λii\\to\\lambda\_\{i\}to high accuracy, and given a fixed matrix\-vector product budget, the primary bottleneck is always the Lanczos depthmm, so we ran one random probe and one gradient probe for each \(checkpoint, matrix\) combination\. This is justified by the fact that deep in the tails, the weight of each Ritz value is standing in for the average weight of a large block of eigenvalues, which creates an averaging effect even with a single probe\.
- •Precision:We found that the Lanczos three term recurrence quickly lost accuracy around depth100100regardless of precision, so we elected to store the full Lanczos basis and do multi\-pass reorthogonalization for all runs\. We found that storing the basis in fp16 led to a drop in accuracy, but there was no drop in accuracy between fp32 and fp64, so we ran in fp32\. The basis for a single checkpoint was around 750GB in memory which we sharded across 16 H100 GPUs\.
- •Sample Size:We compared the spectrum at 10M tokens and 100M tokens and found no significant difference, even deep in the tails, so we elected to compute the matrix\-vector products using 10M tokens for compute efficiency \(Figure[17](https://arxiv.org/html/2607.21716#A2.F17)\)\.

We also tested our methodology against synthetic spectra for which the ground truth is known\. These comparisons can be found in[Figure˜18](https://arxiv.org/html/2607.21716#A2.F18)\. We found extremely strong agreement between the SLQ estimates and the ground truth spectra, even deep into the tails\.

### 4\.2The Hessian Spectrum

![Refer to caption](https://arxiv.org/html/2607.21716v1/x2.png)Figure 2:We estimate the spectrum of the Gauss Newton matrix, the Hessian, and their preconditioned versions at different checkpoints\. The solid lines denote the best estimate for the spectrum, and the shaded bands represent error bands from finite depth Lanczos probes \(m=1200m=1200\)\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x3.png)

\(a\)Top eigenvalues atB=1B=1![Refer to caption](https://arxiv.org/html/2607.21716v1/x4.png)

\(b\)Negative Hessian eigenvalues
![Refer to caption](https://arxiv.org/html/2607.21716v1/x5.png)\(c\)Spectra atT=50%,100%T=50\\%,100\\%
![Refer to caption](https://arxiv.org/html/2607.21716v1/x6.png)\(d\)Final spectra aligned ati=106i=10^\{6\}

Figure 3:Analyses of the Hessian spectra for cosine decay\. See[Figure˜13](https://arxiv.org/html/2607.21716#A2.F13)for the constant schedule\.We plot the spectra for the runs using cosine decay in[Figure˜2](https://arxiv.org/html/2607.21716#S4.F2)\. We begin with two important points about plotting methodology and expected accuracy before interpreting the experimental results\.

##### Plotting Methodology:

Lanczos naturally estimates the mapi→λii\\to\\lambda\_\{i\}\. Specifically, the Lanczos recurrence guarantees that ifi⋆​\(λ\)i^\{\\star\}\(\\lambda\)is the index ofλ\\lambdain the empirical spectral densityμ​\(v0\)\\mu\(v\_\{0\}\)theni⋆​\(ϑj\)∈\[ij−1,ij\]i^\{\\star\}\(\\vartheta\_\{j\}\)\\in\[i\_\{j\-1\},i\_\{j\}\]whereij:=p​∑k=1jωki\_\{j\}:=p\\sum\_\{k=1\}^\{j\}\\omega\_\{k\}are the partial sums of the Ritz weights\. This guarantee can be extended to intermediateλ≠ϑj\\lambda\\neq\\vartheta\_\{j\}using Gauss\-Radau bounds \([Algorithm˜2](https://arxiv.org/html/2607.21716#alg2)\)\. These guarantee that for anyλ\\lambda,i⋆​\(λ\)∈\[i−​\(λ\),i\+​\(λ\)\]i^\{\\star\}\(\\lambda\)\\in\[i\_\{\-\}\(\\lambda\),i\_\{\+\}\(\\lambda\)\]\. We can therefore create a uniform grid in log\-space forλ\\lambda, computei±​\(λ\)i\_\{\\pm\}\(\\lambda\), and shade the resulting band\. We also plot a solid best fit curve for visualization purposes\.

##### Theoretical Error Estimates:

A simple heuristic is that Lanczos can resolve the spectrum with small multiplicative error for eigenvaluesλ≳λ1/m2\\lambda\\gtrsim\\lambda\_\{1\}/m^\{2\}\. Form=1200m=1200, this means we can resolve eigenvalues up to 1\.5M times smaller than the top eigenvalue, which is consistent with the Gauss\-Radau error bands in our experiments\. We empirically found that the error bands for the Hessian are larger than for the Gauss Newton, which we attribute to the fact that Lanczos has to resolve both the positive and negative extremal eigenvalues before resolving the tail nearλ=0\\lambda=0, which effectively halves its effective probe depthmm\.

#### 4\.2\.1Experimental Results

We observed remarkable structure in these spectra, and highlight a few key points:

##### Head\-Tail Split:

In most of the Hessian/Gauss Newton spectra, especially at smaller batch sizes, there is a sharp transition aroundi=8192i=8192, which is our model’s vocab size\. In[Figure˜4](https://arxiv.org/html/2607.21716#S4.F4), we show that this is not a coincidence as the topVVeigenvectors are dominated by theD×VD\\times Vunembedding layer\. We believe this may be related to neural collapse\(Papyan,[2018](https://arxiv.org/html/2607.21716#bib.bib29),[2019](https://arxiv.org/html/2607.21716#bib.bib30); Papyanet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib32); Papyan,[2020](https://arxiv.org/html/2607.21716#bib.bib31)\), in which each tokenvvhas a unique unembedding vectorhv∈SD−1​\(D\)h\_\{v\}\\in S^\{D\-1\}\(\\sqrt\{D\}\), and the Hessian of the last layer is well approximated by the rankVVmatrix:∑v∈\[V\]pv​\(hv​evT\)⊗2\\sum\_\{v\\in\[V\]\}p\_\{v\}\(h\_\{v\}e\_\{v\}^\{T\}\)^\{\\otimes 2\}wherepvp\_\{v\}are the marginal token probabilities\. To test this, we plottedpv\\sqrt\{p\_\{v\}\}against the topVVeigenvalues of the preconditioned spectrum333The square root is to account for Adam’s preconditioner at small batch sizesand found a surprisingly good fit\. For example, both the true spectrum and the token probabilitiespvp\_\{v\}have 3 nearly equal outliers \(corresponding to the tokensthe,,, and\.\)\.

##### Spectrum Evolution:

We observe that both the preconditioned and raw Hessian and Gauss Newton matrices do not change significantly between the50%50\\%and100%100\\%checkpoints\. We fit the capacity exponent as well in[Figure˜3](https://arxiv.org/html/2607.21716#S4.F3)and see that in the tail we get roughlyα≈0\.96\\alpha\\approx 0\.96for Gauss\-Newton andα≈0\.65\\alpha\\approx 0\.65for Hessian\. Note that as shown in Figure[3](https://arxiv.org/html/2607.21716#S4.F3), this exponent, when fitted on the tail, does not depend on batch size or whether we take into account the preconditioner\.

##### Gauss\-Newton Hessian agreement:

We observe that at small and medium batch sizes, the spectra of theGPG\_\{P\}andHPH\_\{P\}agree best in the head, while at large batches they agree best in the tails\. Similarly, the spectra of the raw matricesG,HG,Halways agree in the head and agree best in the tail forB=1024B=1024\.

##### Tail Universality:

Perhaps our most surprising finding is that regardless of the batch size or preconditioner, the tail of the spectrum appears to be*universal*, meaning that the power law is identical and can only be shifted by an absolute constant\. To visualize this, we aligned the spectra of both the raw and preconditioned Gauss Newton and Hessian matrices ati=106i=10^\{6\}\. Especially for the Gauss Newton, the tail fori∈\[105,107\]i\\in\[10^\{5\},10^\{7\}\]appears to match for all batch sizes and regardless of whether the matrix is preconditioned, implying the Adam preconditioner does not affect the tail\.

##### Negative Eigenvalues:

Unsurprisingly, both the preconditioned Hessian and raw Hessian have negative eigenvalues, shown in[Figure˜3](https://arxiv.org/html/2607.21716#S4.F3)\. Perhaps more surprising is that the negative eigenvalues of the preconditioned HessianHPH\_\{P\}are as large as the positive ones atB=64B=64and significantly larger atB=1024B=1024\. In[Figures˜10](https://arxiv.org/html/2607.21716#A1.F10)and[14](https://arxiv.org/html/2607.21716#A2.F14)we can see that the corresponding eigenvectors are almost entirely localized tomlp\.up, the MLP layer that reads from the residual stream\. We hypothesize that these are related to dead neurons which actuate the rare token effect \(see[Section˜5\.1](https://arxiv.org/html/2607.21716#S5.SS1.SSS0.Px1)\)\.

### 4\.3The Hessian Eigenvectors

![Refer to caption](https://arxiv.org/html/2607.21716v1/x7.png)Figure 4:An estimate of the mass of each eigenvector in each parameter group using SLQ\.Because we maintain the basisVVin[Algorithm˜1](https://arxiv.org/html/2607.21716#alg1), we can also compute the Ritz vectorsu1,…,umu\_\{1\},\\ldots,u\_\{m\}associated to the Ritz valuesϑ1,…,ϑm\\vartheta\_\{1\},\\ldots,\\vartheta\_\{m\}\. This lets us estimate which parts of the network are contributing to which parts of the Hessian spectrum\. Explicitly, given a unit vectorvvand a parameter blockℐ\\mathcal\{I\}, we define the mass ofvvinℐ\\mathcal\{I\}to be∑i∈ℐvi2\\sum\_\{i\\in\\mathcal\{I\}\}v\_\{i\}^\{2\}\. Note that if we partition the parameters into disjoint subsets\{I1,…,Ik\}\\\{I\_\{1\},\\ldots,I\_\{k\}\\\}then these masses must sum to11\. We plot the results in[Figure˜4](https://arxiv.org/html/2607.21716#S4.F4)\.

We highlight three main observations:

##### The head is dominated byunembed:

In[Section˜4\.2](https://arxiv.org/html/2607.21716#S4.SS2), we observed that forB=1,64B=1,64, there was a clear split between the head and tail of the spectrum aroundi=Vi=V\. In[Figure˜4](https://arxiv.org/html/2607.21716#S4.F4), we can see that the corresponding eigenvectors live largely in the unembedding layer, which we connected to neural collapse\(Papyanet al\.,[2020](https://arxiv.org/html/2607.21716#bib.bib32)\)in[Section˜4\.2](https://arxiv.org/html/2607.21716#S4.SS2)\. This effect disappears forB=1024B=1024\. We attribute this to the fact that under neural collapse the Hessian of the unembedding layer is adapted to the standard coordinate basis, so full batch Adam is able to completely precondition it away\(Cohenet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib11)\), while small batch Adam only applies a square root\.

##### The tail is spread uniformly across the layers:

After the head, the remaining eigenvectors appear to be spread evenly across the different types of layers, and the eigenvector mass is approximately proportional to the number of parameters in each layer\.

##### Large Hessian outliers are dominated bymlp\.up:

We observed in[Section˜4\.2](https://arxiv.org/html/2607.21716#S4.SS2)that atB=1024B=1024, the Hessian has extreme outlier eigenvalues, both positive and negative\. In[Figure˜4](https://arxiv.org/html/2607.21716#S4.F4)\(for a more complete plot across batch sizes and checkpoints, we refer the reader to[Figure˜10](https://arxiv.org/html/2607.21716#A1.F10)for cosine and[Figure˜14](https://arxiv.org/html/2607.21716#A2.F14)for constant\) we can see that the corresponding eigenvectors live entirely inmlp\.up\. Our hypothesis is that nearly dead neurons inmlp\.upcause the Adam second moment vectorν\\nuto collapse to0, so thatdiag​\(ν−1/2\)​H\\mathrm\{diag\}\(\\nu^\{\-1/2\}\)Hdiverges in those directions\. See[Section˜5\.1](https://arxiv.org/html/2607.21716#S5.SS1.SSS0.Px1)for additional discussion\.

### 4\.4Source Conditions \(Hessian\-gradient alignment\)

![Refer to caption](https://arxiv.org/html/2607.21716v1/x8.png)Figure 5:An estimate of the source condition⟨ui,∇L​\(θ\)⟩2/λi\\expectationvalue\{u\_\{i\},\\nabla L\(\\theta\)\}^\{2\}/\\lambda\_\{i\}using SLQ at the final 100% checkpoint\.In addition to a random Lanczos probe, we also run Lanczos starting from the gradient, which captures the distribution of the gradient across the spectrum\. In analyses of linear regression, this is often called the*source*condition and is parameterized by the exponentβ\\beta:

⟨ui,∇ℒ​\(θ\)⟩2/λi∝i−β\\displaystyle\\expectationvalue\{u\_\{i\},\\nabla\\mathcal\{L\}\(\\theta\)\}^\{2\}/\\lambda\_\{i\}\\propto i^\{\-\\beta\}\(Source Condition\)λi∝i−α\\displaystyle\\lambda\_\{i\}\\propto i^\{\-\\alpha\}\(Capacity Condition\)where\{\(λi,ui\)\}\\\{\(\\lambda\_\{i\},u\_\{i\}\)\\\}are the eigenvalues and eigenvectors of the Hessian\. This can also be interpreted as the amount of excess loss contained in theii\-th eigenvector because for a quadratic objective:

ℒ​\(θ\)−ℒ⋆=∑i=1p⟨ui,∇ℒ​\(θ\)⟩22​λi\.\\displaystyle\\mathcal\{L\}\(\\theta\)\-\\mathcal\{L\}^\{\\star\}=\\sum\_\{i=1\}^\{p\}\\frac\{\\expectationvalue\{u\_\{i\},\\nabla\\mathcal\{L\}\(\\theta\)\}^\{2\}\}\{2\\lambda\_\{i\}\}\.\(1\)This easily generalizes to preconditioned methods\. If we define the symmetrized preconditioned Hessian to beP−1/2​H​P−1/2P^\{\-1/2\}HP^\{\-1/2\}and replace the gradient with the semi\-preconditioned gradientg←P−1/2​gg\\leftarrow P^\{\-1/2\}g, then the same decomposition holds\. In[Figure˜5](https://arxiv.org/html/2607.21716#S4.F5), we estimate this source condition for both the preconditioned Gauss\-Newton matrix and the preconditioned Hessian\. We note that the map that Lanczos naturally estimates is the cumulative sumi→∑j<i\(g⋅uj\)2i\\to\\sum\_\{j<i\}\(g\\cdot u\_\{j\}\)^\{2\}, so to estimatei→\(g⋅uj\)2i\\to\(g\\cdot u\_\{j\}\)^\{2\}, we have to differentiate this map which is numerically unstable\. We therefore first average the value over a multiplicative window of\[i/2,2​i\]\[i/2,2i\]to smooth it before differentiating\.

The primary observation is that in the part of the tail we can resolve, the source condition for the Gauss Newton matrix is consistently aroundβ∈\[0\.1,0\.2\]\\beta\\in\[0\.1,0\.2\]while the source condition for the Hessian is closer toβ∈\[0\.3,0\.6\]\\beta\\in\[0\.3,0\.6\]\. This contradicts the common assumptionβ\>1\\beta\>1in the literature on power laws, which guarantees[eq\.˜1](https://arxiv.org/html/2607.21716#S4.E1)is summable asp→∞p\\to\\inftyand, in the standard bias\-variance decomposition, gives rise to power laws in the bias which depend on the source and capacity exponents\. Instead, whenβ<1\\beta<1, as we measure here, the problem actually behaves like it is finite dimensional\(Paquetteet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib96), Figure 2\)\. Whenβ<1\\beta<1, there is a critical timeT⋆=pαT^\{\\star\}=p^\{\\alpha\}, whereppis the number of parameters, after which the bias decays exponentially\. Meanwhile the variance term hits a noise floor without averaging and decays likeO​\(1/T\)O\(1/T\)with averaging\(Polyak and Juditsky,[1992](https://arxiv.org/html/2607.21716#bib.bib8); Dieuleveutet al\.,[2018](https://arxiv.org/html/2607.21716#bib.bib98)\)\. We empirically verify this in[Figure˜9](https://arxiv.org/html/2607.21716#A1.F9)by fitting a power law of the formℒt=ℒ⋆\+A​T−γ\\mathcal\{L\}\_\{t\}=\\mathcal\{L\}^\{\\star\}\+AT^\{\-\\gamma\}to all of our optimally tuned runs whereℒ⋆≈2\.502\\mathcal\{L\}^\{\\star\}\\approx 2\.502is shared\. Nearly all runs haveγ≈1\\gamma\\approx 1, which matches the variance dominated regime predicted by theory\.

## 5Edge of Stability

Cohenet al\.\([2021](https://arxiv.org/html/2607.21716#bib.bib4)\)demonstrated that full\-batch optimization typically occurs at the*edge of stability*\(EOS\), meaning that increasing the learning rate to\(1\+ϵ\)​η\(1\+\\epsilon\)\\etawould cause the optimizer to diverge on the local quadratic Taylor expansion, while decreasing it to\(1−ϵ\)​η\(1\-\\epsilon\)\\etawould cause the optimizer to converge\. For gradient descent, this is captured by conditionη≈2/λ1​\(H\)\\eta\\approx 2/\\lambda\_\{1\}\(H\)\. This was later generalized to adaptive optimizers, including Adam, inCohenet al\.\([2022](https://arxiv.org/html/2607.21716#bib.bib5)\)where they showed that if you freeze the preconditioner, adaptive optimizers also typically operate at the edge of stability\.

These empirical results primarily hold in the full\-batch \(or at least very large batch\) setting\. However, a similar phenomenon occurs at smaller batch sizes\. In particular, there is a*stochastic edge of stability*regime in which decreasing your batch size can cause you to diverge on the local quadratic Taylor expansion\. In this section, we empirically test whether optimization in LLMs occurs at either the deterministic or stochastic edge of stability, and we compare these results with the theoretical predictions made by linear stability analysis\.

### 5\.1Methodology

To test whether optimization occurs at the edge of stability, we reuse a similar experimental setup as in Section[3\.2](https://arxiv.org/html/2607.21716#S3.SS2)\. We Taylor expand the model usingquad, around the EMA iterateθ¯\\bar\{\\theta\}, so that the objective on each minibatch is quadratic\. We then freeze Adam’s second moment bufferν\\nuto the EMAν¯\\bar\{\\nu\}\. We then run preconditioned SGD with this fixed preconditioner for10%10\\%of the total token budget, or until the loss diverges\. To test how close a checkpoint is to the edge of stability, we run this experiment over a5×55\\times 5grid of learning rate and batch size multipliers:\{1/4,1/2,1,2,4\}\\\{1/4,1/2,1,2,4\\\}\.444For batch size11we could not use a batch size multiplier less than11so we generated a5×35\\times 3grid\.We also repeat this experiment over1010random seeds, where each random seed shuffles the sequences in the dataset\. We say a run diverges if its loss hits100100, and we plot the fraction of runs which diverge as a function of the learning rate and batch size multipliers in Figure[7](https://arxiv.org/html/2607.21716#S5.F7)\. While we arbitrarily cap the loss divergence thresholds at100100,[Figure˜7](https://arxiv.org/html/2607.21716#S5.F7)is incredibly stable for any cutoff between1010and10510^\{5\}\.

##### Extreme Outliers

An important detail, which results from freezing the preconditioner, is that thequadtraining is now susceptible to divergence due to extreme outliers in the dynamics, which we observed were connected to rare tokens or dead neurons\(Voitaet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib20)\)\. We believe an illustrative example is to look at the dynamics of the embedding matrixEE\. In every minibatch, if tokeniidoes not occur then its corresponding embeddingEiE\_\{i\}gets no gradient\. Because the second moment bufferν\\nuis an EMA of these zero gradients, if the token isn’t seen for roughly11−β2\\tfrac\{1\}\{1\-\\beta\_\{2\}\}steps, then those entries ofν\\nubecome exponentially small\. If this token later reoccurs, then Adam will divide its gradient byν\+ϵ≈0\\sqrt\{\\nu\}\+\\epsilon\\approx 0, which can cause a spike\. While this phenomenon is clearest for the embedding matrix, it can also occur throughout the model, e\.g\. from neurons which are rarely activated\(Voitaet al\.,[2024](https://arxiv.org/html/2607.21716#bib.bib20)\)\. To circumvent this issue, we estimate the Frobenius norm of the preconditioned Gauss\-Newton matrix for each sequence‖P−1​Gi‖F2\\\|P^\{\-1\}G\_\{i\}\\\|\_\{F\}^\{2\}and filter out the top1%1\\%outliers\. Without this correction, nearly all hyperparameters diverged after a single outsized spike\.

Before we discuss the experimental results, we will briefly discuss the predictions that theory makes for stability of SGD on quadratics so that we can better interpret the experimental results and validate the predictive power of the theory\.

### 5\.2Theoretical Predictions for Stability on Quadratics

![Refer to caption](https://arxiv.org/html/2607.21716v1/figures/stable_region.png)Figure 6:A toy simulation of the linear stability on a 1D regression problem whereHi∼i\.i\.d\.Poisson​\(Bcrit\)H\_\{i\}\\stackrel\{\{\\scriptstyle i\.i\.d\.\}\}\{\{\\sim\}\}\\mathrm\{Poisson\}\(B\_\{\\text\{crit\}\}\)\. Note the MSE stable region is monotonic inη,B,τ\\eta,B,\\tauwhile the Lyapunov stable region is not\. Our EOS experiments determine where on this cartoon training occurs by perturbingη,B\\eta,B\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x9.png)Figure 7:Heatmap showing instability for models trained with cosine decay as a function of\(η,B\)\(\\eta,B\)across batch sizeB=1B=1,6464and10241024, at10%10\\%,50%50\\%and100%100\\%of training\. Cell color indicates how many seeds diverged out of a total of1010seeds per run\. We provide the analogous figure for constant learning rate in Figure[16](https://arxiv.org/html/2607.21716#A2.F16)\(Appendix[B](https://arxiv.org/html/2607.21716#A2)\)\.The dynamics of SGD on a quadratic are a linear dynamical system\. In particular, ifHℬH\_\{\\mathcal\{B\}\}denotes the Hessian on a minibatchℬ\\mathcal\{B\}then we can rewrite the SGD update as:

θ−θ⋆\\displaystyle\\theta\-\\theta^\{\\star\}←θ−θ⋆−η​∇ℒℬ​\(θ\)\\displaystyle\\leftarrow\\theta\-\\theta^\{\\star\}\-\\eta\\nabla\\mathcal\{L\}\_\{\\mathcal\{B\}\}\(\\theta\)SGD update\(2\)=θ−θ⋆−η​\[∇ℒℬ​\(θ\)−∇ℒℬ​\(θ⋆\)\]−η​∇ℒℬ​\(θ⋆\)\\displaystyle=\\theta\-\\theta^\{\\star\}\-\\eta\[\\nabla\\mathcal\{L\}\_\{\\mathcal\{B\}\}\(\\theta\)\-\\nabla\\mathcal\{L\}\_\{\\mathcal\{B\}\}\(\\theta^\{\\star\}\)\]\-\\eta\\nabla\\mathcal\{L\}\_\{\\mathcal\{B\}\}\(\\theta^\{\\star\}\)add and subtract=\(I−η​Hℬ\)⏟stochastic contraction​\(θ−θ⋆\)−η​∇ℒℬ​\(θ⋆\)⏟additive noise\.\\displaystyle=\\underbrace\{\(I\-\\eta H\_\{\\mathcal\{B\}\}\)\}\_\{\\mathclap\{\\text\{stochastic contraction\}\}\}\(\\theta\-\\theta^\{\\star\}\)\-\\underbrace\{\\eta\\nabla\\mathcal\{L\}\_\{\\mathcal\{B\}\}\(\\theta^\{\\star\}\)\}\_\{\\text\{additive noise\}\}\.gradients are affineIn other words, at every step we multiply the current displacement toθ⋆\\theta^\{\\star\}by a random matrixI−η​HℬI\-\\eta H\_\{\\mathcal\{B\}\}then add a random noise vector independent ofθ\\theta\. For linear dynamical systems of this form, there are two relevant notions of stability: mean square stability and Lyapunov stability\.

###### Definition 1\.

A linear dynamical systemθt\\theta\_\{t\}ismean squarestable iflimt→∞𝔼​\[‖θt‖2\]<∞\\lim\_\{t\\to\\infty\}\\mathbb\{E\}\[\\\|\\theta\_\{t\}\\\|^\{2\}\]<\\infty, andLyapunovstable iflimt→∞𝔼​log⁡‖θt‖<∞\\lim\_\{t\\to\\infty\}\\mathbb\{E\}\\log\\\|\\theta\_\{t\}\\\|<\\infty\.

Note that by Jensen’s inequality, mean square stability implies Lyapunov stability:

𝔼​\[log⁡‖θ‖\]=12​𝔼​\[log⁡‖θ‖2\]≤12​log⁡𝔼​‖θ‖2\.\\mathbb\{E\}\[\\log\\\|\\theta\\\|\]=\\tfrac\{1\}\{2\}\\mathbb\{E\}\[\\log\\\|\\theta\\\|^\{2\}\]\\leq\\tfrac\{1\}\{2\}\\log\\mathbb\{E\}\\\|\\theta\\\|^\{2\}\.For most processes, mean\-square stability and Lyapunov stability nearly coincide\. However, mean square stability is significantly more sensitive to heavy tails\. As an illustrative example, letη=B=1\\eta=B=1and run SGD on the loss functionsℒ1=12​θ2\\mathcal\{L\}\_\{1\}=\\tfrac\{1\}\{2\}\\theta^\{2\},ℒ2=32​θ2\\mathcal\{L\}\_\{2\}=\\tfrac\{3\}\{2\}\\theta^\{2\}\. Then at every step you multiplyθ\\thetaby either1−η=01\-\\eta=0or1−3​η=−21\-3\\eta=\-2, each with probability1/21/2\. The marginal at timettis0with probability1−2−t1\-2^\{\-t\}and\(−2\)t\(\-2\)^\{t\}with probability2−t2^\{\-t\}\. This process is mean square unstable because𝔼​\[θ2\]=2t→∞\\mathbb\{E\}\[\\theta^\{2\}\]=2^\{t\}\\to\\infty, but it is Lyapunov stable, which reflects thatθt→0\\theta\_\{t\}\\to 0almost surely\. Mean\-square stability captures wild fluctuations, while Lyapunov stability captures the “typical” behavior of the process\. We will begin by briefly discussing mean\-square stability before discussing Lyapunov stability, as it is more mathematically tractable and provides better intuition\.

##### Mean Square Stability

There is a simple closed form for mean\-square stability\. If

Σ=𝔼​\[\(θ−θ⋆\)​\(θ−θ⋆\)⊤\]\\Sigma=\\mathbb\{E\}\[\(\\theta\-\\theta^\{\\star\}\)\(\\theta\-\\theta^\{\\star\}\)^\{\\top\}\]then from[eq\.˜2](https://arxiv.org/html/2607.21716#S5.E2),Σ\\Sigmais updated by

Σ←𝔼ℬ​\[\(I−η​Hℬ\)​Σ​\(I−η​Hℬ\)\]\+η2​C​where​C=Covℬ​\[∇ℒℬ​\(θ⋆\)\]\.\\displaystyle\\Sigma\\leftarrow\\mathbb\{E\}\_\{\\mathcal\{B\}\}\[\(I\-\\eta H\_\{\\mathcal\{B\}\}\)\\Sigma\(I\-\\eta H\_\{\\mathcal\{B\}\}\)\]\+\\eta^\{2\}C\\mbox\{\\quad where\\quad\}C=\\mathrm\{Cov\}\_\{\\mathcal\{B\}\}\[\\nabla\\mathcal\{L\}\_\{\\mathcal\{B\}\}\(\\theta^\{\\star\}\)\]\.This is a deterministic linear system forΣ\\Sigma, soΣ\\Sigmawill either converge or diverge depending on the eigenvalues of the linear operatorΦη,B\\Phi\_\{\\eta,B\}:

Φη,B​\[Σ\]:=𝔼ℬ​\[\(I−η​Hℬ\)​Σ​\(I−η​Hℬ\)\]\.\\displaystyle\\Phi\_\{\\eta,B\}\[\\Sigma\]:=\\mathbb\{E\}\_\{\\mathcal\{B\}\}\[\(I\-\\eta H\_\{\\mathcal\{B\}\}\)\\Sigma\(I\-\\eta H\_\{\\mathcal\{B\}\}\)\]\.Stability of theΣ\\Sigmarecurrence requiresλ1​\(Φη,B\)≤1\\lambda\_\{1\}\(\\Phi\_\{\\eta,B\}\)\\leq 1\. Because𝔼​‖θt‖2=tr​Σt\\mathbb\{E\}\\\|\\theta\_\{t\}\\\|^\{2\}=\\mathrm\{tr\}~\\Sigma\_\{t\}, this is also the condition for mean square stability\. Note that the learning rateη\\etaenters directly, and the batch sizeBBenters implicitly through the sampling of the minibatchℬ\\mathcal\{B\}\. With some effort, it is possible to make these dependencies more explicit\. For the remainder of this section, we will useτ:=η/B\\tau:=\\eta/Bto denote the “temperature” or learning rate to batch size ratio of the process\.

###### Proposition 1\(Mulayoff and Michaeli,[2024](https://arxiv.org/html/2607.21716#bib.bib49), Theorem 5 & Proposition 6\)\.

Letτ:=η/B\\tau:=\\eta/Bbe the temperature and let𝒮MSE=\{\(η,τ\)\|λ1​\(Φη,B\)<1\}\\mathcal\{S\}\_\{\\text\{MSE\}\}=\\\{\(\\eta,\\tau\)~\|~\\lambda\_\{1\}\(\\Phi\_\{\\eta,B\}\)<1\\\}be the set of mean square stable hyperparameters for SGD\. Then the stable set𝒮MSE\\mathcal\{S\}\_\{\\text\{MSE\}\}satisfies the following properties:

- •𝒮MSE\\mathcal\{S\}\_\{\\text\{MSE\}\}is monotonic inη\\etaandτ\\tau, i\.e\. if\(η,τ\)∈𝒮MSE\(\\eta,\\tau\)\\in\\mathcal\{S\}\_\{\\text\{MSE\}\}then\(c1​η,c2​τ\)∈𝒮MSE\(c\_\{1\}\\eta,c\_\{2\}\\tau\)\\in\\mathcal\{S\}\_\{\\text\{MSE\}\}for anyc1,c2≤1c\_\{1\},c\_\{2\}\\leq 1\.
- •There is a maximum learning rateηcrit\\eta\_\{\\text\{crit\}\}and maximum temperatureτcrit\\tau\_\{\\text\{crit\}\}, i\.e\.𝒮MSE⊆\[0,ηcrit\)×\[0,τcrit\)\\mathcal\{S\}\_\{\\text\{MSE\}\}\\subseteq\[0,\\eta\_\{\\text\{crit\}\}\)\\times\[0,\\tau\_\{\\text\{crit\}\}\)\.ηcrit\\eta\_\{\\text\{crit\}\}is given by2/λ1​\(H\)2/\\lambda\_\{1\}\(H\), as for GD, whileτcrit\\tau\_\{\\text\{crit\}\}depends on the homogeneity of the minibatch Hessians\.

The effect of these observations is that inη,τ\\eta,\\tauspace, the set of hyperparameters for which SGD is mean\-square stable is a rounded rectangle with edgesη<ηcrit\\eta<\\eta\_\{\\text\{crit\}\}andτ<τcrit\\tau<\\tau\_\{\\text\{crit\}\}\. Transforming this intoη,B\\eta,Bspace, this becomes a wedge shape\. This is reflected in the toy simulation in[Figure˜6](https://arxiv.org/html/2607.21716#S5.F6)\.

##### Lyapunov Stability

Lyapunov stability is significantly less structured than mean square stability\. For example,[Figure˜6](https://arxiv.org/html/2607.21716#S5.F6)shows a linear regression problem where\(η,B\)\(\\eta,B\)is Lyapunov unstable but\(η,2​B\)\(\\eta,2B\)is stable, so Lyapunov stability does not satisfy a monotonicity property\. Furthermore, the set of Lyapunov stable hyperparameters is often not even connected\. This distinction turns out to be relevant for LLMs\. In our experiments, we did observe settings in which\(η,B\)\(\\eta,B\)is stable but\(η/2,B/2\)\(\\eta/2,B/2\)is not stable, which violates monotonicity\. This is inconsistent with mean\-square stability and suggests that there are regimes during training in which SGD is Lyapunov stable but not mean square stable\.

### 5\.3Experimental Results

Our experimental results demonstrate that the majority of training occurs within a factor of22of the edge of stability, and the shape of the plots is incredibly consistent with the theoretical predictions in[Section˜5\.2](https://arxiv.org/html/2607.21716#S5.SS2)\. Figure[7](https://arxiv.org/html/2607.21716#S5.F7)shows that forB=1B=1andB=64B=64, the stability boundary follows the diagonal line, whereas atB=1024B=1024there is a flat cutoff, consistent with the illustration from[Figure˜6](https://arxiv.org/html/2607.21716#S5.F6)\. Moreover, while the\(η,B\)=\(1,1\)\(\\eta,B\)=\(1,1\)cell is not always in the stability region, the\(0\.5,1\)\(0\.5,1\)cell is always stable, indicating that every model is running at most a factor of22over the stability boundary, which is typical even for full\-batch gradient descent\(Cohenet al\.,[2021](https://arxiv.org/html/2607.21716#bib.bib4)\)and matches the conclusions ofCaiet al\.\([2026](https://arxiv.org/html/2607.21716#bib.bib99)\)\.

Our results also contextualize the finding ofCaiet al\.\([2026](https://arxiv.org/html/2607.21716#bib.bib99)\)that the edge of stability disappears at small batch sizes\. One confounder is that their probes run for a fixed number of steps \(≤50\\leq 50\), rather than a fixed number of tokens, so at small batch sizes the simulations process very few tokens and can fail to detect instabilities\. Our results suggest that after normalizing by tokens, rather than steps, LLMs operate within a factor of22of the stochastic edge of stability at small batch sizes\.

## 6Discussion and Limitations

##### Quadratics as a Pretraining Sandbox\.

While LLM loss landscapes are notgloballyquadratic, the optimizer appears to enter regions that are well approximated by a prox\-linear or Gauss\-Newton Taylor expansion for a nontrivial fraction of the training budget \(up to 10%\)\. This may partially explain why intuition and algorithms for the quadratic model seem to extend to LLM pretraining\.

##### Structure of the Local Quadratic\.

When the local dynamics are captured by thequadTaylor expansion, we can understand these dynamics through the spectrum of the Hessian/Gauss\-Newton matrix\. We estimated these spectra using deep Lanczos quadrature probes and found that the local quadratic models have extremely structured spectra, with power\-law like behavior in both the head and the tail\. We also find that the tails are*universal*meaning that their shape does not vary with training time or batch size and is not affected by the Adam preconditioner\. The SLQ probes also allow us to localize different parts of the spectrum to different layers in the network\.

##### Local Stability\.

Motivated by the observation that full\-batch training typically occurs at the edge of stability\(Cohenet al\.,[2021](https://arxiv.org/html/2607.21716#bib.bib4),[2022](https://arxiv.org/html/2607.21716#bib.bib5)\)and by theoretical analyses of linear stability of SGD\(Mulayoff and Michaeli,[2024](https://arxiv.org/html/2607.21716#bib.bib49); Défossez and Bach,[2015](https://arxiv.org/html/2607.21716#bib.bib46)\), we empirically test whether LLM pretraining operates near a \(stochastic\) stability boundary\. We first study the set of locally stable hyperparameters\(η,B\)\(\\eta,B\), and show how the shape of its boundary determines whether the instability is caused by discretization or stochasticity\. We then locate where training occurs within this set by perturbing\(η,B\)\(\\eta,B\)and checking for instability\. At small batch sizes, training typically occurs at a stochastic edge of stability, while at large batch sizes it typically occurs at a deterministic edge of stability\.

##### Progressive Sharpening and Curvature Regularization\.

It has been widely observed that the stochasticity in SGD/Adam drives optimization towards flatter regions of the loss landscapeKeskaret al\.\([2017](https://arxiv.org/html/2607.21716#bib.bib100)\); Jastrzębskiet al\.\([2018](https://arxiv.org/html/2607.21716#bib.bib54)\)\. This curvature regularization effect cannot be captured by the quadratic model since the curvature is by definition fixed\. Therefore, the fact that the quadratic model does seem to capture the training dynamics near the end of the training suggests that this curvature regularization effect is relatively weak by the end of training\. However, we do observe evidence of progressive sharpening for our runs with cosine decay\. As the learning rate decays, the runs leave the edge of stability and the curvature once again begins to grow, which matches the observations inCaiet al\.\([2026](https://arxiv.org/html/2607.21716#bib.bib99)\)\.

##### Limitations\.

Our experiments focus on a fixed model scale, architecture, and optimizer\. While the agreement between the local and ground truth model is remarkable, it is an important future direction to explore how the model size and other optimizers — particularly second order optimizers — influence the spectrum exponents and loss agreement\. Moreover, Figure[1](https://arxiv.org/html/2607.21716#S1.F1)shows that early on in training, both the quadratic and the prox\-linear approximations are much less accurate than at mid\-to\-late time, suggesting that higher\-order terms in the loss play an important role in the early phase of pretraining\. Finally, while the late time agreement for cosine is remarkable, constant with EMA tracks the original very poorly \(Figure[12](https://arxiv.org/html/2607.21716#A2.F12)\)\. We speculate that this is related to the edge of stability behavior of the constant runs during training and we leave a further study to future work\.

### 6\.1Conclusion

Overall, our findings support a simple but powerful mental model: LLM pretraining can be locally understood as optimization over a sequence of structured quadratic problems\. Instead of asking for a single global theory of a highly nonconvex loss landscape, we can ask how the local quadratic model changes along the trajectory, how the optimizer interacts with its spectrum, and how we can leverage these insights to design more efficient training pipelines\. In this sense, the quadratic model is not merely a toy problem, but an experimentally validated proxy for the optimization dynamics of pretraining, and a useful foundation for building theory that applies to real world models\.

### Acknowledgements and Funding

The authors would like to thank Max Shad and Bala Desinghu for their help with the cluster\. AM, DM, PN acknowledge the support of a Kempner Institute Graduate Research Fellowship\. CP is supported by an NSF CAREER Award \(IIS\-2239780\), DARPA grants DIAL\-FP\-038 and AIQ\-HR00112520041, the Simons Collaboration on the Physics of Learning and Neural Computation, and the William F\. Milton Fund from Harvard University\. The authors acknowledge that this work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence\. SK and DM acknowledge support from the Office of Naval Research under award N0001422\-1\-2377 and the National Science Foundation Grant under award \#IIS 2229881\. DM is also supported by a Simons Investigator Fellowship, NSF grant DMS\-2134157, DARPA grant W911NF2010021, and DOE grant DE\-SC0022199\.

## References

- N\. Abreu, N\. Vyas, S\. Kakade, and D\. Morwani \(2025\)The potential of second\-order optimization for llms: a study with full gauss\-newton\.arXiv preprint arXiv:2510\.09378\.Cited by:[§3\.1](https://arxiv.org/html/2607.21716#S3.SS1.p1.7)\.
- R\. P\. Adams, J\. Pennington, M\. J\. Johnson, J\. Smith, Y\. Ovadia, B\. Patton, and J\. Saunderson \(2018\)Estimating the spectral density of large implicit matrices\.arXiv preprint arXiv:1802\.03451\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- K\. Ahn, J\. Zhang, and S\. Sra \(2022\)Understanding the unstable convergence of gradient descent\.InInternational conference on machine learning,pp\. 247–257\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1)\.
- A\. Andreyev, A\. Ananthkumar, M\. Walden, T\. Poggio, and P\. Beneventano \(2026\)Momentum further constrains sharpness at the edge of stochastic stability\.arXiv preprint arXiv:2604\.14108\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- A\. Andreyev and P\. Beneventano \(2024\)Edge of stochastic stability: revisiting the edge of stability for sgd\.arXiv preprint arXiv:2412\.20553\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- S\. Arora, Z\. Li, and A\. Panigrahi \(2022\)Understanding gradient descent on the edge of stability in deep learning\.InInternational Conference on Machine Learning,pp\. 948–1024\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- A\. Atanasov, B\. Bordelon, and C\. Pehlevan \(2021\)Neural networks as kernel learners: the silent alignment effect\.arXiv preprint arXiv:2111\.00034\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- Y\. Bai, B\. Krause, H\. Wang, C\. Xiong, and R\. Socher \(2020\)Taylorized training: towards better approximation of neural network training at finite width\.arXiv preprint arXiv:2002\.04010\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- Y\. Bai and J\. D\. Lee \(2019\)Beyond linearization: on quadratic and higher\-order approximation of wide neural networks\.arXiv preprint arXiv:1910\.01619\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- B\. Bordelon and C\. Pehlevan \(2021\)Learning curves for sgd on structured features\.arXiv preprint arXiv:2106\.02713\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- J\. V\. Burke \(1985\)Descent methods for composite nondifferentiable optimization problems\.Mathematical Programming33\(3\),pp\. 260–279\.Cited by:[§3\.1](https://arxiv.org/html/2607.21716#S3.SS1.p1.7)\.
- Y\. Cai, H\. Huang, H\. Wen, D\. Liu, Y\. Ma, and K\. Lyu \(2026\)Does LLM pre\-training typically occur at the edge of stability?\.InWorkshop on Scientific Methods for Understanding Deep Learning,External Links:[Link](https://openreview.net/forum?id=QSb05IuPsy)Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p2.1),[§5\.3](https://arxiv.org/html/2607.21716#S5.SS3.p1.7),[§5\.3](https://arxiv.org/html/2607.21716#S5.SS3.p2.2),[§6](https://arxiv.org/html/2607.21716#S6.SS0.SSS0.Px4.p1.1)\.
- A\. Chatzimichailidis, J\. Keuper, F\. Pfreundt, and N\. R\. Gauger \(2019\)Gradvis: visualization and second order analysis of optimization surfaces during the training of deep neural networks\.In2019 IEEE/ACM Workshop on Machine Learning in High Performance Computing Environments \(MLHPC\),pp\. 66–74\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- T\. Chen, T\. Trogdon, and S\. Ubaru \(2021\)Analysis of stochastic lanczos quadrature for spectrum approximation\.InInternational Conference on Machine Learning,pp\. 1728–1739\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- T\. Chen, T\. Trogdon, and S\. Ubaru \(2022\)Randomized matrix\-free quadrature for spectrum and spectral sum approximation\.arXiv preprint arXiv:2204\.01941,pp\. 55\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- L\. Chizat, E\. Oyallon, and F\. Bach \(2019\)On lazy training in differentiable programming\.Advances in neural information processing systems32\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p5.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2607.21716#S3.SS1.p1.7)\.
- J\. M\. Cohen, A\. Damian, A\. Talwalkar, J\. Z\. Kolter, and J\. D\. Lee \(2024\)Understanding optimization in deep learning with central flows\.arXiv preprint arXiv:2410\.24206\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2),[§4\.3](https://arxiv.org/html/2607.21716#S4.SS3.SSS0.Px1.p1.3)\.
- J\. M\. Cohen, B\. Ghorbani, S\. Krishnan, N\. Agarwal, S\. Medapati, M\. Badura, D\. Suo, D\. Cardoze, Z\. Nado, G\. E\. Dahl,et al\.\(2022\)Adaptive gradient methods at the edge of stability\.arXiv preprint arXiv:2207\.14484\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2),[§4](https://arxiv.org/html/2607.21716#S4.p2.5),[§5](https://arxiv.org/html/2607.21716#S5.p1.3),[§6](https://arxiv.org/html/2607.21716#S6.SS0.SSS0.Px3.p1.2)\.
- J\. M\. Cohen, S\. Kaur, Y\. Li, J\. Z\. Kolter, and A\. Talwalkar \(2021\)Gradient descent on neural networks typically occurs at the edge of stability\.arXiv preprint arXiv:2103\.00065\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2),[§5\.3](https://arxiv.org/html/2607.21716#S5.SS3.p1.7),[§5](https://arxiv.org/html/2607.21716#S5.p1.3),[§6](https://arxiv.org/html/2607.21716#S6.SS0.SSS0.Px3.p1.2)\.
- E\. M\. Compagnoni, L\. Biggio, A\. Orvieto, F\. N\. Proske, H\. Kersting, and A\. Lucchi \(2023\)An sde for modeling sam: theory and insights\.InInternational Conference on Machine Learning,pp\. 25209–25253\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- E\. M\. Compagnoni, R\. Islamov, F\. N\. Proske, A\. Lucchi, A\. Orvieto, and E\. Gorbunov \(2026\)On the interaction of batch noise, adaptivity, and compression, under\(L0,L1\)\(L\_\{0\},L\_\{1\}\)\-smoothness: an sde approach\.External Links:2506\.00181,[Link](https://arxiv.org/abs/2506.00181)Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- E\. M\. Compagnoni, T\. Liu, R\. Islamov, F\. N\. Proske, A\. Orvieto, and A\. Lucchi \(2024\)Adaptive methods through the lens of sdes: theoretical insights on the role of noise\.arXiv preprint arXiv:2411\.15958\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Damian, E\. Nichani, and J\. D\. Lee \(2022\)Self\-stabilization: the implicit bias of gradient descent at the edge of stability\.arXiv preprint arXiv:2209\.15594\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- A\. Defazio, A\. Cutkosky, H\. Mehta, and K\. Mishchenko \(2023\)Optimal linear decay learning rate schedules and further refinements\.arXiv preprint arXiv:2310\.07831\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Defazio, X\. Yang, H\. Mehta, K\. Mishchenko, A\. Khaled, and A\. Cutkosky \(2024\)The road less scheduled\.Advances in Neural Information Processing Systems37,pp\. 9974–10007\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Defazio \(2026\)Schedulefree\+: scaling learning\-rate\-free & schedule\-free learning to large language models\.arXiv preprint arXiv:2605\.19095\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Défossez and F\. Bach \(2015\)Averaged least\-mean\-squares: bias\-variance trade\-offs and optimal sampling distributions\.InArtificial Intelligence and Statistics,pp\. 205–213\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2),[§6](https://arxiv.org/html/2607.21716#S6.SS0.SSS0.Px3.p1.2)\.
- N\. Dey, B\. Zhang, L\. Noci, M\. Li, B\. Bordelon, S\. Bergsma, C\. Pehlevan, B\. Hanin, and J\. Hestness \(2026\)Don’t be lazy: completep enables compute\-efficient deep transformers\.Advances in Neural Information Processing Systems38,pp\. 137707–137739\.Cited by:[§3\.2](https://arxiv.org/html/2607.21716#S3.SS2.p2.8),[footnote 2](https://arxiv.org/html/2607.21716#footnote2)\.
- A\. Dieuleveut, A\. Durmus, and F\. Bach \(2018\)Bridging the gap between constant step size stochastic gradient descent and markov chains\.External Links:1707\.06386,[Link](https://arxiv.org/abs/1707.06386)Cited by:[§4\.4](https://arxiv.org/html/2607.21716#S4.SS4.p2.12)\.
- D\. Drusvyatskiy \(2017\)The proximal point method revisited\.arXiv preprint arXiv:1712\.06038\.Cited by:[§3\.1](https://arxiv.org/html/2607.21716#S3.SS1.p1.7)\.
- D\. Ferbach, K\. Everett, G\. Gidel, E\. Paquette, and C\. Paquette \(2026\)Dimension\-adapted momentum outscales sgd\.Advances in Neural Information Processing Systems38,pp\. 112780–112977\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- S\. Fort, G\. K\. Dziugaite, M\. Paul, S\. Kharaghani, D\. M\. Roy, and S\. Ganguli \(2020\)Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel\.Advances in Neural Information Processing Systems33,pp\. 5850–5861\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- B\. Ghorbani, S\. Krishnan, and Y\. Xiao \(2019\)An investigation into neural net optimization via hessian eigenvalue density\.InInternational Conference on Machine Learning,pp\. 2232–2241\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- G\. H\. Golub and G\. Meurant \(2009\)Matrices, moments and quadrature with applications\.Princeton University Press\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1),[§4\.1](https://arxiv.org/html/2607.21716#S4.SS1.p1.7)\.
- D\. Granziol and K\. Juarev \(2026\)Hessian spectral analysis at foundation model scale\.arXiv preprint arXiv:2602\.00816\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- D\. Granziol, X\. Wan, and T\. Garipov \(2019\)Deep curvature suite\.arXiv preprint arXiv:1912\.09656\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- D\. Granziol, S\. Zohren, and S\. Roberts \(2022\)Learning rates as a function of batch size: a random matrix theory approach to neural network training\.Journal of Machine Learning Research23\(173\),pp\. 1–65\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- D\. Granziol \(2025\)HessFormer: hessians at foundation scale\.arXiv preprint arXiv:2505\.11564\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- A\. Hägele, E\. Bakouch, A\. Kosson, L\. B\. Allal, L\. Von Werra, and M\. Jaggi \(2024\)Scaling laws and compute\-optimal training beyond fixed training durations\.Advances in Neural Information Processing Systems37,pp\. 76232–76264\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Jacot, F\. Gabriel, and C\. Hongler \(2018\)Neural tangent kernel: convergence and generalization in neural networks\.Advances in neural information processing systems31\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p5.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2607.21716#S3.SS1.p1.7)\.
- P\. Jain, S\. M\. Kakade, R\. Kidambi, P\. Netrapalli, and A\. Sidford \(2018\)Parallelizing stochastic gradient descent for least squares regression: mini\-batching, averaging, and model misspecification\.Journal of machine learning research18\(223\),pp\. 1–42\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1),[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- S\. Jastrzębski, Z\. Kenton, N\. Ballas, A\. Fischer, Y\. Bengio, and A\. Storkey \(2018\)On the relation between the sharpest directions of dnn loss and the sgd step length\.arXiv preprint arXiv:1807\.05031\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2),[§6](https://arxiv.org/html/2607.21716#S6.SS0.SSS0.Px4.p1.1)\.
- S\. Jastrzebski, M\. Szymczak, S\. Fort, D\. Arpit, J\. Tabor, K\. Cho, and K\. Geras \(2020\)The break\-even point on optimization trajectories of deep neural networks\.arXiv preprint arXiv:2002\.09572\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- K\. Jiang, J\. Cohen, and Y\. Li \(2026\)Understanding the evolution of the neural tangent kernel at the edge of stability\.Advances in Neural Information Processing Systems38,pp\. 31996–32036\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- D\. S\. Kalra, T\. He, and M\. Barkeshli \(2025\)Universal sharpness dynamics in neural network training: fixed point analysis, edge of stability, and route to chaos\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 55966–56000\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- N\. S\. Keskar, D\. Mudigere, J\. Nocedal, M\. Smelyanskiy, and P\. T\. P\. Tang \(2017\)On large\-batch training for deep learning: generalization gap and sharp minima\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=H1oyRlYgg)Cited by:[§6](https://arxiv.org/html/2607.21716#S6.SS0.SSS0.Px4.p1.1)\.
- R\. Kidambi, P\. Netrapalli, P\. Jain, and S\. Kakade \(2018\)On the insufficiency of existing momentum schemes for stochastic optimization\.In2018 Information Theory and Applications Workshop \(ITA\),pp\. 1–9\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p1.1),[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- C\. Lanczos \(1950\)An iteration method for the solution of the eigenvalue problem of linear differential and integral operators\.Journal of research of the National Bureau of Standards45\(4\),pp\. 255–282\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- C\. Lauditi, C\. Pehlevan, and B\. Bordelon \(2026\)Spectral dynamics in deep networks: feature learning, outlier escape, and learning rate transfer\.arXiv preprint arXiv:2605\.07870\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- J\. Lee, L\. Xiao, S\. Schoenholz, Y\. Bahri, R\. Novak, J\. Sohl\-Dickstein, and J\. Pennington \(2019\)Wide neural networks of any depth evolve as linear models under gradient descent\.Advances in neural information processing systems32\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p5.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- A\. Lewkowycz, Y\. Bahri, E\. Dyer, J\. Sohl\-Dickstein, and G\. Gur\-Ari \(2020\)The large learning rate phase of deep learning: the catapult mechanism\.arXiv preprint arXiv:2003\.02218\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- Q\. Li, C\. Tai,et al\.\(2017\)Stochastic modified equations and adaptive stochastic gradient algorithms\.InInternational Conference on Machine Learning,pp\. 2101–2110\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- L\. Lin, Y\. Saad, and C\. Yang \(2016\)Approximating spectral densities of large matrices\.SIAM review58\(1\),pp\. 34–65\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- A\. Liu, B\. Feng, B\. Xue, B\. Wang, B\. Wu, C\. Lu, C\. Zhao, C\. Deng, C\. Zhang, C\. Ruan,et al\.\(2024\)Deepseek\-v3 technical report\.arXiv preprint arXiv:2412\.19437\.Cited by:[§3\.2](https://arxiv.org/html/2607.21716#S3.SS2.p2.8)\.
- P\. M\. Long \(2021\)Properties of the after kernel\.arXiv preprint arXiv:2105\.10585\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- C\. Ma and L\. Ying \(2021\)On linear stability of sgd and input\-smoothness of neural networks\.Advances in Neural Information Processing Systems34,pp\. 16805–16817\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- S\. Ma, R\. Bassily, and M\. Belkin \(2018\)The power of interpolation: understanding the effectiveness of sgd in modern over\-parametrized learning\.InInternational Conference on Machine Learning,pp\. 3325–3334\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- S\. Malladi, K\. Lyu, A\. Panigrahi, and S\. Arora \(2022\)On the sdes and scaling rules for adaptive gradient algorithms\.Advances in Neural Information Processing Systems35,pp\. 7697–7711\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- S\. Malladi, A\. Wettig, D\. Yu, D\. Chen, and S\. Arora \(2023\)A kernel\-based view of language model fine\-tuning\.InInternational Conference on Machine Learning,pp\. 23610–23641\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- A\. Meterez, D\. Morwani, J\. Wu, C\. Oncescu, C\. Pehlevan, and S\. Kakade \(2025\)Seesaw: accelerating training by balancing learning rate and batch size scheduling\.arXiv preprint arXiv:2510\.14717\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p1.1),[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- A\. Meterez, P\. A\. Nair, D\. Morwani, C\. Pehlevan, and S\. Kakade \(2026\)Anytime pretraining: horizon\-free learning\-rate schedules with weight averaging\.arXiv preprint arXiv:2602\.03702\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- D\. Morwani, A\. Meterez, P\. Nair, and S\. Kakade \(2026\)Compute efficiency and serial runtime tradeoffs for stochastic momentum methods\.arXiv preprint arXiv:2606\.19179\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- R\. Mulayoff and T\. Michaeli \(2024\)Exact mean square linear stability analysis for sgd\.InThe Thirty Seventh Annual Conference on Learning Theory,pp\. 3915–3969\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2),[§6](https://arxiv.org/html/2607.21716#S6.SS0.SSS0.Px3.p1.2),[Proposition 1](https://arxiv.org/html/2607.21716#Thmproposition1)\.
- L\. Noci, A\. Meterez, T\. Hofmann, and A\. Orvieto \(2024\)Why do learning rates transfer? reconciling optimization and scaling limits for deep learning\.arXiv preprint arXiv:2402\.17457\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- G\. Ortiz\-Jiménez, S\. Moosavi\-Dezfooli, and P\. Frossard \(2021\)What can linearized neural networks actually say about generalization?\.Advances in Neural Information Processing Systems34,pp\. 8998–9010\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- V\. Papyan, X\. Han, and D\. L\. Donoho \(2020\)Prevalence of neural collapse during the terminal phase of deep learning training\.Proceedings of the National Academy of Sciences117\(40\),pp\. 24652–24663\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1),[§4\.2\.1](https://arxiv.org/html/2607.21716#S4.SS2.SSS1.Px1.p1.11),[§4\.3](https://arxiv.org/html/2607.21716#S4.SS3.SSS0.Px1.p1.3)\.
- V\. Papyan \(2018\)The full spectrum of deepnet hessians at scale: dynamics with sgd training and sample size\.arXiv preprint arXiv:1811\.07062\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1),[§4\.2\.1](https://arxiv.org/html/2607.21716#S4.SS2.SSS1.Px1.p1.11)\.
- V\. Papyan \(2019\)Measurements of three\-level hierarchical structure in the outliers in the spectrum of deepnet hessians\.arXiv preprint arXiv:1901\.08244\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1),[§4\.2\.1](https://arxiv.org/html/2607.21716#S4.SS2.SSS1.Px1.p1.11)\.
- V\. Papyan \(2020\)Traces of class/cross\-class structure pervade deep learning spectra\.Journal of Machine Learning Research21\(252\),pp\. 1–64\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1),[§4\.2\.1](https://arxiv.org/html/2607.21716#S4.SS2.SSS1.Px1.p1.11)\.
- C\. Paquette and E\. Paquette \(2021\)Dynamics of stochastic momentum methods on large\-scale, quadratic models\.Advances in Neural Information Processing Systems34,pp\. 9229–9240\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- E\. Paquette, C\. Paquette, L\. Xiao, and J\. Pennington \(2024\)4\+ 3 phases of compute\-optimal neural scaling laws\.Advances in Neural Information Processing Systems37,pp\. 16459–16537\.Cited by:[§4\.4](https://arxiv.org/html/2607.21716#S4.SS4.p2.12)\.
- S\. M\. Park, K\. Georgiev, A\. Ilyas, G\. Leclerc, and A\. Madry \(2023\)Trak: attributing model behavior at scale\.arXiv preprint arXiv:2303\.14186\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- B\. T\. Polyak and A\. B\. Juditsky \(1992\)Acceleration of stochastic approximation by averaging\.SIAM journal on control and optimization30\(4\),pp\. 838–855\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p1.1),[§4\.4](https://arxiv.org/html/2607.21716#S4.SS4.p2.12)\.
- B\. T\. Polyak \(1964\)Some methods of speeding up the convergence of iteration methods\.Ussr computational mathematics and mathematical physics4\(5\),pp\. 1–17\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p1.1)\.
- L\. Sagun, L\. Bottou, and Y\. LeCun \(2016\)Eigenvalues of the hessian in deep learning: singularity and beyond\.arXiv preprint arXiv:1611\.07476\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- L\. Sagun, U\. Evci, V\. U\. Guney, Y\. Dauphin, and L\. Bottou \(2017\)Empirical analysis of the hessian of over\-parametrized neural networks\.arXiv preprint arXiv:1706\.04454\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- F\. Schaipp, A\. Hägele, A\. Taylor, U\. Simsekli, and F\. Bach \(2025\)The surprising agreement between convex optimization theory and learning\-rate scheduling for large model training\.arXiv preprint arXiv:2501\.18965\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- E\. Shulgin, D\. von Rütte, T\. H\. Zhang, N\. Ajroldi, B\. Schölkopf, and A\. Orvieto \(2026\)Deriving hyperparameter scaling laws via modern optimization theory\.arXiv preprint arXiv:2603\.15958\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- S\. Ubaru, J\. Chen, and Y\. Saad \(2017\)Fast estimation of tr\(f\(a\)\) via stochastic lanczos quadrature\.SIAM Journal on Matrix Analysis and Applications38\(4\),pp\. 1075–1099\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- A\. Varre and N\. Flammarion \(2022\)Accelerated sgd for non\-strongly\-convex least squares\.InConference on Learning Theory,pp\. 2062–2126\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- A\. V\. Varre, L\. Pillaud\-Vivien, and N\. Flammarion \(2021\)Last iterate convergence of sgd for least\-squares in the interpolation regime\.\.Advances in Neural Information Processing Systems34,pp\. 21581–21591\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- M\. Velikanov, D\. Kuznedelev, and D\. Yarotsky \(2022\)A view of mini\-batch sgd via generating functions: conditions of convergence, phase transitions, benefit from negative momenta\.arXiv preprint arXiv:2206\.11124\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- E\. Voita, J\. Ferrando, and C\. Nalmpantis \(2024\)Neurons in large language models: dead, n\-gram, positional\.InFindings of the Association for Computational Linguistics: ACL 2024,pp\. 1288–1301\.Cited by:[§5\.1](https://arxiv.org/html/2607.21716#S5.SS1.SSS0.Px1.p1.9)\.
- N\. Vyas, Y\. Bansal, and P\. Nakkiran \(2022\)Limitations of the ntk for understanding generalization in deep learning\.arXiv preprint arXiv:2206\.10012\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Wu, P\. L\. Bartlett, S\. M\. Kakade, J\. D\. Lee, and B\. Yu \(2025\)Risk comparisons in linear regression: implicit regularization dominates explicit regularization\.arXiv preprint arXiv:2509\.17251\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- L\. Wu, C\. Ma,et al\.\(2018\)How sgd selects the global minima in over\-parameterized learning: a dynamical stability perspective\.Advances in Neural Information Processing Systems31\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1),[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- L\. Wu, M\. Wang, and W\. Su \(2022\)The alignment property of sgd noise and how it helps select flat minima: a stability analysis\.Advances in Neural Information Processing Systems35,pp\. 4680–4693\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p7.1),[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px3.p1.2)\.
- Z\. Yao, A\. Gholami, K\. Keutzer, and M\. W\. Mahoney \(2020\)Pyhessian: neural networks through the lens of the hessian\.In2020 IEEE international conference on big data \(Big data\),pp\. 581–590\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- C\. Zhang, R\. Islamov, E\. M\. Compagnoni, J\. Pang, A\. Lucchi, and A\. Orvieto \(2026a\)Beyond a single explanation of the adam–sgd gap\.arXiv preprint arXiv:2606\.14259\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px4.p1.1)\.
- D\. Zhang, X\. Tang, Y\. Liang, and D\. Zou \(2026b\)Scaling laws for precision in high\-dimensional linear regression\.arXiv preprint arXiv:2602\.19241\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- G\. Zhang, L\. Li, Z\. Nado, J\. Martens, S\. Sachdeva, G\. Dahl, C\. Shallue, and R\. B\. Grosse \(2019\)Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model\.Advances in neural information processing systems32\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- H\. Zhang, Y\. Liu, Q\. Chen, and C\. Fang \(2024a\)The optimality of \(accelerated\) sgd for high\-dimensional quadratic optimization\.arXiv preprint arXiv:2409\.09745\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- H\. Zhang, D\. Morwani, N\. Vyas, J\. Wu, D\. Zou, U\. Ghai, D\. Foster, and S\. Kakade \(2024b\)How does critical batch size scale in pre\-training?\.arXiv preprint arXiv:2410\.21676\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p1.1),[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.
- Y\. Zhang, C\. Chen, T\. Ding, Z\. Li, R\. Sun, and Z\. Luo \(2024c\)Why transformers need adam: a hessian perspective\.Advances in neural information processing systems37,pp\. 131786–131823\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- Y\. Zhang, C\. Chen, Z\. Li, T\. Ding, C\. Wu, D\. D\. Kingma, Y\. Ye, Z\. Luo, and R\. Sun \(2025\)Adam\-mini: use fewer learning rates to gain more\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 28033–28063\.Cited by:[§2](https://arxiv.org/html/2607.21716#S2.SS0.SSS0.Px2.p1.1)\.
- D\. Zou, J\. Wu, V\. Braverman, Q\. Gu, and S\. Kakade \(2021\)Benign overfitting of constant\-stepsize sgd for linear regression\.InConference on Learning Theory,pp\. 4633–4635\.Cited by:[§1](https://arxiv.org/html/2607.21716#S1.p3.1)\.

## Appendix AAdditional Plots

![Refer to caption](https://arxiv.org/html/2607.21716v1/x10.png)Figure 8:\(η,B\)\(\\eta,B\)pretraining sweep for the constant with EMA and cosine trained models\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x11.png)Figure 9:Joint power law fit for the loss as a function of the tokens of the formL=L⋆\+A​T−γL=L\_\{\\star\}\+AT^\{\-\\gamma\}\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x12.png)Figure 10:A full sweep of the eigenvector distribution at the end of training for cosine decay\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x13.png)Figure 11:An extended version of[Figure˜5](https://arxiv.org/html/2607.21716#S4.F5)with additional checkpoints\. This is for the set of runs using cosine decay after warmup\. See[Figure˜15](https://arxiv.org/html/2607.21716#A2.F15)for the constant schedule runs\.
## Appendix BConstant Learning Rate with EMA

All of the plots in the main paper were run with 10% warmup followed by cosine decay to 10% of the maximal learning rate\. We also repeated the experiments with a constant learning rate schedule, which remains fixed after warmup\.

![Refer to caption](https://arxiv.org/html/2607.21716v1/x14.png)Figure 12:Validation loss curves for 150M models \(black\) trained with constant learning with EMA for1×1\\timesChinchilla tokens, across batch sizes, at optimal learning rates, under the same setup as Figure[1](https://arxiv.org/html/2607.21716#S1.F1)\. The quadratic expansions are much more unstable, and we leave a study of this to future work\. We provide further details in Section[3](https://arxiv.org/html/2607.21716#S3)\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x15.png)Figure 13:The spectra using a constant learning rate schedule after warmup, rather than cosine decay\. This is otherwise equivalent to[Figure˜3](https://arxiv.org/html/2607.21716#S4.F3)\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x16.png)Figure 14:A full sweep of the eigenvector distribution at the end of training for the constant learning rate schedule\. This is otherwise equivalent to[Figures˜4](https://arxiv.org/html/2607.21716#S4.F4)and[10](https://arxiv.org/html/2607.21716#A1.F10)\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x17.png)Figure 15:The source condition using a constant learning rate schedule after warmup, rather than cosine decay\. This is otherwise equivalent to[Figures˜5](https://arxiv.org/html/2607.21716#S4.F5)and[11](https://arxiv.org/html/2607.21716#A1.F11)\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x18.png)Figure 16:Heatmap showing instability for models trained with constant learning rate with EMA under the same setup as Figure[1](https://arxiv.org/html/2607.21716#S1.F1)\. As in[Figure˜1](https://arxiv.org/html/2607.21716#S1.F1), the smaller batch sizes operate near the stochastic edge of stability\. However, because the stable region in theB=1024B=1024plots is diagonal then vertical, it operates at both the deterministic and stochastic edge of stability\. This is the case whenever\(η,B\)\(\\eta,B\)is stable but both\(2​η,2​B\)\(2\\eta,2B\)and\(η,B/2\)\(\\eta,B/2\)are unstable\. This is distinct from our cosine schedule experiments for whichB=1024B=1024operated entirely at the deterministic edge of stability\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x19.png)Figure 17:Comparison of the preconditioned Gauss Newton spectrum atT=100%T=100\\%computed over 10M and 100M tokens, both at sequence length10241024\. Each plot shows a different pretraining batch size, which is distinct from the number of samples used to estimate the Hessian\. The spectra largely agree, even deep in the tail, so our remaining experiments estimate the Hessian over 10M tokens\.![Refer to caption](https://arxiv.org/html/2607.21716v1/x20.png)Figure 18:Comparing the ground truth \(dashed black\) with the SLQ estimate \(orange\) for various spectra\. For the ground truth, we tested two power law spectraλi∝i−α\\lambda\_\{i\}\\propto i^\{\-\\alpha\}along with the estimated spectra forP−1​GP^\{\-1\}GandP−1​HP^\{\-1\}Hto check consistency with an additional SLQ probe\. The light orange bands show the Gauss\-Radau error bands from[Algorithm˜2](https://arxiv.org/html/2607.21716#alg2)while the light blue bands show a 99% confidence interval for the error that comes from sampling a single random probe\. The bottom row shows that the relative errorlog⁡\(λ^i/λi\)\\log\(\\hat\{\\lambda\}\_\{i\}/\\lambda\_\{i\}\)remains small deep in the spectrum\.