A Defense of the Quadratic Model

TL;DR

Meterez et al. stress-test the simplest possible model of optimization — a local quadratic Taylor expansion of the loss and a linear Taylor expansion of the model — as a predictive tool for LLM pretraining. At 150M parameters and 3B training tokens, expanding at intermediate checkpoints and continuing optimization on the linearized/quadratic surrogate reproduces the real training loss over windows lasting up to ~10% of the run. Beyond the empirical agreement, they use deep Lanczos quadrature (probe depth m=1200) to resolve the Hessian spectrum across six orders of magnitude, find a structured “vocab-sized head + power-law tail” that persists across batch sizes and is largely invariant to the Adam preconditioner, and empirically test local linear stability — showing LLM training typically occurs at a stochastic edge of stability at small batch and a deterministic edge of stability at large batch. The core claim: a local-quadratic proxy is a theoretically tractable and empirically defensible description of late-training optimization dynamics.

Key claims

Continuing training from an intermediate checkpoint using (a) the linearized model with cross-entropy loss and (b) a quadratic Taylor expansion of the full loss both match the real run’s loss trajectory over windows of ~10% of training, with agreement improving as training progresses [§2, §3, thread 2/8].
Agreement is worse early in training and worse under a constant LR schedule than under cosine — the quadratic proxy is late-time, schedule-sensitive rather than universally valid [§3, thread 6/8].
The Hessian spectrum, resolved by Stochastic Lanczos Quadrature with unusually deep probes (m=1200), has a bimodal structure: a head of roughly vocab-size eigenvalues (interpreted as neural collapse) and a power-law tail [§4, thread 3-4/8].
The power-law tail is essentially universal across batch sizes and unaffected by the Adam preconditioner, suggesting it reflects a data/architecture property rather than an optimizer artifact [§4, thread 4/8].
Local linear stability at intermediate checkpoints — freezing the preconditioner and sweeping (η, B) from the checkpoint until divergence — shows LLM training sits at a stochastic edge of stability for small batch sizes and a deterministic edge of stability for large batch sizes, matching linear-stability theory predictions [§5, thread 5/8].
The nature of the edge (stochastic vs. deterministic) is determined by batch size, tying the sharpness of the operating point to the noise-vs-drift balance of stochastic optimization [§5].

Method

The setup is a 150M-parameter decoder LLM trained on 3B tokens with checkpoints saved every 10% of training. At each checkpoint the authors freeze the parameters θ_t, form the linearized model f_lin(x; θ) = f(x; θ_t) + ∇_θ f(x; θ_t) · (θ - θ_t) and the quadratic loss L_quad(θ) = L(θ_t) + g_t · (θ - θ_t) + ½ (θ - θ_t)^T H_t (θ - θ_t) where g_t, H_t are the mini-batch gradient and Hessian at θ_t. They continue training on each surrogate using the same optimizer + schedule and compare loss trajectories to the real run.

Hessian-spectrum estimation uses Stochastic Lanczos Quadrature with a Krylov subspace of dimension m=1200 — an order of magnitude deeper than typical prior work, which is what enables resolving eigenvalue mass across six orders of magnitude and cleanly separating the vocab-size head from the power-law tail. Spectra are computed at each checkpoint under multiple batch sizes and with/without Adam preconditioning applied.

Local stability is probed by fixing θ_t and its Adam preconditioner state, running preconditioned SGD from that state on a (learning-rate, batch-size) grid, and marking which combinations diverge. Comparison to linear-stability theory predictions (deterministic EoS at η · λ_max ≈ 2 for large batch, stochastic EoS at a batch-dependent threshold below that for small batch) yields the small-B-stochastic / large-B-deterministic split.

Results

Agreement window. Late in training (near cosine decay tail), quadratic + linearized surrogates match the real loss trajectory for continuation windows lasting up to ~10% of the total training run [§3, thread 2/8].
Schedule sensitivity. Constant-LR runs exhibit worse quadratic-model agreement than cosine, and early-training checkpoints exhibit worse agreement than late-training ones [§3, thread 6/8].
Hessian spectrum: head + tail. The eigenvalue distribution shows a distinct head of roughly |V| eigenvalues (vocab-size, consistent with neural-collapse geometry of the output layer) and a smooth power-law tail extending across six orders of magnitude [§4].
Preconditioner invariance. The power-law tail exponent is essentially unchanged whether the raw Hessian or the Adam-preconditioned Hessian P^{-1/2} H P^{-1/2} is spectrally analyzed — Adam’s diagonal preconditioner rescales but does not restructure the bulk of the tail [§4, Fig. Hessian spectrum panels].
Batch-size invariance of the tail. The same power-law tail appears across a sweep of batch sizes, even though the head shifts [§4].
Stability boundary. From late-training checkpoints, the divergence curve in (η, B) space lies on the stochastic-EoS boundary for small B and rotates onto the deterministic-EoS boundary η · λ_max = 2 for large B, matching the theoretical prediction of linear-stability analysis [§5].

Why it’s interesting

This paper is the natural companion to Generalization at the Edge of Stability on the wiki: both take the Cohen-et-al. Edge-of-Stability regime as a given empirical fact about how LLM optimization actually behaves, both deploy Stochastic Lanczos Quadrature with deep probes to resolve the Hessian spectrum, and both derive structure that pointwise sharpness measures miss — but they answer complementary questions. Generalization at the Edge of Stability asks what generalization theory still applies inside the chaotic EoS regime (Sharpness Dimension over a random pullback attractor); this paper asks whether local optimization dynamics inside that regime remain tractable, and finds the answer is yes — a quadratic Taylor expansion reproduces the real trajectory for ~10% of training. Together they form a mini-cluster: “EoS is empirically inevitable in LLM training; here is what remains predictable inside it.”

For the [[training-stability-at-scale]] fork on the wiki, this is a data-side rather than architecture-side or optimizer-side contribution. The architecture line (A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training Gated Attention, mHC: Manifold-Constrained Hyper-Connections Birkhoff-polytope residual mixing) locates stability in structural fixes; the optimizer line (Controlled LLM Training on Spectral Sphere SSO, Aurora: A Leverage-Aware Optimizer for Rectangular Matrices Aurora) locates it in update geometry. This paper says: whatever knobs you turn, the underlying object being optimized is well-approximated by a structured quadratic with vocab-size head + universal power-law tail, and training operates on the stability boundary of that object. The batch-size-determines-which-edge finding also gives a mechanistic hook into Hyperparameter scaling laws — the Fitting/Transfer/Enforce three-way fork all implicitly picks a point on that stability boundary, and the deterministic/stochastic split is a cleaner language for what happens as B varies than raw LR sweeps.