🤖 machine learning

A Defense of the Quadratic Model

This paper demonstrates that the simple quadratic model, when applied via Taylor expansion at intermediate checkpoints, can surprisingly accurately predict optimization dynamics in large language models, revealing structured Hessian spectra and confirming that training typically occurs at a stochastic edge of stability dependent on batch size.

Original authors: Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade, Alex Damian

Published 2026-07-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade, Alex Damian

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to navigate a massive, foggy mountain range at night. This mountain is the "loss landscape" of a giant Artificial Intelligence, a complex terrain where every step the AI takes is an attempt to find the lowest valley (the best possible performance). The problem is that the mountain is so huge and twisted that no one can see the whole map at once. For years, scientists have tried to predict how the AI moves by using simple, flat maps or idealized models. Sometimes these maps work, but often they are too simple to be useful, or too complicated to calculate. The big question is: Can we use a simple, flat map to predict the AI's journey through this chaotic mountain?

To understand this, you need to know a few things. First, AI models learn by "optimizing," which just means taking tiny steps to reduce their mistakes. Second, the "Hessian" is a mathematical way of describing how steep or curved the ground is at any specific spot. If you know the curve, you can predict if a step will take you down into a valley or up a cliff. Finally, there is a concept called the "edge of stability," which is like walking right on the edge of a cliff where the ground is so steep that one wrong step sends you tumbling, yet the AI somehow keeps walking without falling. This paper asks if the simplest possible map—a flat, quadratic (bowl-shaped) approximation—can actually predict how a giant AI behaves in the real world.

The authors of this paper decided to stress-test this simple "quadratic model" on a Large Language Model (LLM) with 150 million parameters, trained on 3 billion words. They wanted to see if, instead of trying to understand the whole crazy mountain, they could just look at a small, flat patch of ground around where the AI currently is, and use that to predict the next few steps.

They found something surprisingly cool: the simple bowl-shaped map works! When they stopped the training every 10% of the way through and created a flat, quadratic version of the problem at that exact spot, this simple model could accurately predict the AI's behavior for the next 10% of the training. It's as if, even though the mountain is a jagged mess, the AI spends so much time in specific "bowl-shaped" valleys that a simple map of a bowl is enough to tell you where it will go next. The prediction was even better near the end of training than at the beginning.

Once they proved the simple map worked, they used it to peek under the hood of the AI's learning process. They looked at the "Hessian spectrum," which is like a list of all the different slopes and curves in the landscape. They discovered that this list isn't random; it has a very specific structure. The top part of the list is dominated by the part of the AI that handles vocabulary (the "unembedding" layer), while the long, trailing part of the list follows a universal "power law." This means the shape of the tail is identical whether the AI is using a tiny batch of data or a huge one, and it doesn't care about the specific tricks (preconditioners) the optimizer uses to speed things up. It's a universal pattern hidden in the chaos.

Finally, they investigated the "edge of stability." They wanted to know if the AI is walking carefully in the middle of the valley or dancing on the very edge of a cliff. By tweaking the learning rate and the batch size (the size of the data chunks), they found that the AI is indeed dancing on the edge. When the AI uses small batches, it dances on a "stochastic edge of stability," meaning it's teetering because of the random noise in the data. When it uses huge batches, it teeters on a "deterministic edge," where the instability is caused by the math of the steps themselves. In both cases, the AI is operating right at the limit of what is stable, which suggests that this precarious balancing act is actually how these models learn best.

The paper concludes that while the AI's world is incredibly complex, we can understand it by treating it as a sequence of simple, local quadratic problems. Instead of needing a single, impossible theory for the whole mountain, we can just look at the local bowl the AI is standing in. This simple model isn't just a toy; it's a surprisingly accurate proxy for how real-world AI pretraining actually works, giving us a new, tractable way to understand and improve these massive systems.

Technical Summary: A Defense of the Quadratic Model

Problem Statement

Optimization theory for Large Language Models (LLMs) often relies on idealized models to navigate the complexity of non-convex loss landscapes. However, a significant gap exists between theoretical tractability and empirical accuracy: models that are mathematically simple often fail to describe real-world dynamics, while models that capture specific problem structures are frequently theoretically intractable. The central question addressed by this work is: To what extent does the simplest possible optimization model—the quadratic model—actually capture the training dynamics of LLMs?

Methodology

The authors conduct a rigorous "stress test" of the quadratic model using a 150M parameter transformer trained on 3 billion tokens from FineWeb. The methodology involves three primary components:

1. Local Quadratic Approximation (Section 3)

The authors evaluate the predictive power of local Taylor expansions at intermediate checkpoints (every 10% of training). They define two proxy losses:

Prox (Linearized): Linearizes the model $f_\theta$ around a reference point $\theta_{ref}$ while maintaining the original cross-entropy objective. This is equivalent to multi-class logistic regression under the Neural Tangent Kernel (NTK) regime.
Quad (Quadratic): Taylor expands the Prox loss to the second order, resulting in a convex quadratic objective. Crucially, they use the Gauss-Newton matrix ( $G$ ) rather than the full Hessian ( $H$ ) to ensure the quadratic objective remains convex (positive semidefinite), avoiding unboundedness in negative curvature directions.

They train these proxy models for 10% of the total token budget starting from each checkpoint and compare the resulting optimization trajectories against the true LLM training dynamics.

2. Hessian Spectrum Estimation (Section 4)

To analyze the structure of these local quadratic problems, the authors estimate the spectra of the Hessian and Gauss-Newton matrices (both raw and Adam-preconditioned) along the training trajectory.

Technique: They employ Stochastic Lanczos Quadrature (SLQ) with extremely deep probes ( $m=1200$ ) and rigorous error bands (Gauss-Radau bounds).
Scale: Probes are run on checkpoints at 10%, 50%, and 100% of training across batch sizes $B \in \{1, 64, 1024\}$ .
Analysis: They examine spectral decay (power laws), eigenvector distribution across network layers, and the alignment between the gradient and eigenvectors (source conditions).

3. Edge of Stability (EOS) Analysis (Section 5)

The authors empirically test whether LLM training operates at the "edge of stability."

Setup: They freeze the Adam preconditioner and run preconditioned SGD on the local quadratic proxy.
Perturbation: They vary learning rate ( $\eta$ ) and batch size ( $B$ ) multipliers on a grid to determine the stability boundary.
Metric: A run is considered divergent if the loss exceeds a threshold (100). They distinguish between Mean Square Stability (sensitive to heavy tails) and Lyapunov Stability (typical behavior).

Key Results

1. Predictive Power of Quadratic Models

Agreement Window: The local quadratic and linearized models accurately predict the true optimization dynamics for windows lasting up to 10% of the total training budget.
Temporal Accuracy: The approximation quality is significantly better in the middle-to-late stages of training compared to the early stages.
Prox vs. Quad: The linearized model (Prox) generally matches the true loss better than the quadratic model (Quad) in early training, while Quad converges to high accuracy only after 60% of training.

2. Structure of the Hessian Spectrum

Head-Tail Split: The spectrum exhibits a sharp transition around index $i \approx 8192$ (the vocabulary size). The "head" (top eigenvalues) is dominated by the unembedding layer, consistent with the "neural collapse" phenomenon.
Universal Tail: The "tail" of the spectrum (indices $> 10^5$ ) follows a power law that is universal: it is identical across different batch sizes and unaffected by the Adam preconditioner.
Exponents:
- Capacity Exponent ( $\alpha$ ): Estimated at $\alpha \approx 0.96$ for the Gauss-Newton tail and $\alpha \approx 0.65$ for the Hessian tail.
- Source Exponent ( $\beta$ ): The source condition (gradient alignment) yields $\beta \in [0.1, 0.2]$ for Gauss-Newton and $\beta \in [0.3, 0.6]$ for the Hessian. This contradicts the common literature assumption of $\beta > 1$ , suggesting the problem behaves as if it were finite-dimensional with exponential bias decay after a critical time.
Negative Eigenvalues: The preconditioned Hessian exhibits large negative eigenvalues, particularly at larger batch sizes, localized to the MLP up-projection layer. The authors hypothesize these are linked to "dead neurons" and the rare token effect.

3. Edge of Stability Dynamics

Stochastic vs. Deterministic:
- At small to medium batch sizes ( $B=1, 64$ ), training occurs at a stochastic edge of stability. The stability boundary follows a diagonal line in the $(\eta, B)$ space, indicating that instability is driven by stochasticity (noise).
- At large batch sizes ( $B=1024$ ), training operates at the deterministic edge of stability, characterized by a flat cutoff in the stability region, consistent with full-batch gradient descent behavior.
Proximity to Boundary: Across all settings, training operates within a factor of 2 of the stability boundary.

Significance and Claims

The paper argues that the quadratic model is not merely a theoretical toy but a valid, experimentally verified proxy for LLM pretraining dynamics.

Theoretical Tractability: The findings suggest that complex LLM optimization can be locally understood as a sequence of structured quadratic problems. This allows researchers to leverage the rich toolbox of quadratic optimization theory (e.g., stability analysis, learning rate scheduling) to reason about LLM training.
Revisiting Assumptions: The results challenge standard assumptions in the literature, specifically the summability of source conditions ( $\beta > 1$ ) and the universality of the Hessian tail under preconditioning.
Practical Implications: The identification of the "stochastic edge of stability" at small batch sizes provides a refined understanding of why training diverges under specific hyperparameter configurations, distinguishing between discretization errors and stochastic noise.

The authors conclude that while global loss landscapes are non-convex, the local quadratic approximation offers a powerful, tractable framework for analyzing and designing training pipelines for large-scale models.