A Defense of the Quadratic Model

2026-07-23

AI Summary

This work addresses the challenge of characterizing the highly complex loss landscape in large language model (LLM) pretraining, where existing theories struggle to balance analytical tractability with accurate dynamic prediction. By performing Taylor expansions of both the model and loss function at mid-training, the authors construct a local quadratic approximation and combine it with Lanczos quadrature and Hessian spectral estimation. For the first time, they validate this approach on a 150M-parameter LLM trained on 3B tokens, demonstrating predictive accuracy over a training window spanning 10% of total steps. Their analysis reveals that the quadratic model faithfully captures optimization trajectories, that the tail structure of the Hessian spectrum is strongly influenced by batch size, preconditioning, and training stage, and that optimization typically resides in a stochastic edge-of-stability regime dictated by batch size—uncovering a deep connection between local stability and hyperparameter choice.

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.

Problem

• quadratic model
• optimization dynamics
• loss landscape
• large language models
• Hessian spectrum

Innovation

• quadratic model
• Hessian spectrum
• stochastic edge of stability
• Lanczos quadrature
• optimization dynamics