A Defense of the Quadratic Model

2026-07-27 · Source: stat.ML updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Mathematics & Computational Sciences · Depth: Expert, extended

Summary

The quadratic model, a simplified optimization theory model, demonstrates surprising predictive power in Large Language Model (LLM) pretraining. Researchers stress-tested this model using a 150M parameter transformer trained on 3B FineWeb tokens. They found that Taylor expanding the model and loss function at intermediate checkpoints accurately predicts optimization dynamics for windows up to 10% of total training. Analysis of these local quadratic problems revealed structured Hessian spectra, with a head dominated by the 8192-token vocabulary unembedding layer and a universal power-law tail unaffected by batch size or Adam preconditioner. Furthermore, LLM optimization typically occurs at a stochastic edge of stability for small batch sizes and a deterministic edge for large batch sizes.

Key takeaway

For AI Scientists and Machine Learning Engineers optimizing LLM pretraining, understanding the local quadratic model's predictive power is crucial. You should consider using local Taylor expansions to analyze short-term training dynamics and diagnose stability issues. This approach allows for more principled adjustments to learning rates and batch sizes, especially when operating near the observed stochastic or deterministic edges of stability, potentially improving training efficiency and convergence.

Key insights

The quadratic model accurately predicts LLM optimization dynamics locally, offering a tractable proxy for pretraining analysis.

Principles

• Local quadratic expansions predict LLM dynamics for up to 10% of training.
• Hessian spectra exhibit a vocab-size-dependent head and universal power-law tail.
• LLM training operates at a stochastic or deterministic edge of stability.

Method

Taylor expand the LLM and loss function at training checkpoints, then run preconditioned SGD on the resulting quadratic objective to predict dynamics and analyze stability.

In practice

• Use local quadratic approximations for short-term optimization trajectory prediction.
• Analyze Hessian spectra to understand layer-specific curvature contributions.
• Adjust learning rates and batch sizes based on observed stability boundaries.

Topics

• LLM Optimization
• Quadratic Models
• Hessian Spectrum
• Edge of Stability
• Taylor Expansion
• Adam Optimizer

Code references

• alexandrumeterez/QuadraticModel

Best for: Research Scientist, AI Scientist, Machine Learning Engineer

Related on AIssential

• A Defense of the Quadratic Model — The quadratic model offers a surprisingly predictive and theoretically tractable proxy for understanding LLM pretraining optimization dynamics.
• PC Layer: Polynomial Weight Preconditioning for Improving LLM Pre-Training — Polynomial preconditioning reshapes weight singular-value spectra to improve LLM training stability and efficiency without inference overhead.
• Learning Perturbations to Extrapolate Your LLM — Learnable continuous perturbations in embedding space significantly improve LLM out-of-domain extrapolation.
• Spectra: Rethinking Optimizers for LLMs Under Spectral Anisotropy — LLM gradients exhibit a "spike-tail" anisotropy where dominant linguistic structures suppress long-tail semantic learning.
• Local Stability and Gaussian Smoothing of Quantized Neural Networks — Gaussian averaging offers a smooth surrogate for quantized neural networks, enhancing stability analysis and enabling smooth surrogate gradients.
• A Quantitative Experimental Repeated Measures Study of Training Dynamics in a Small Llama Style Language Model Under a Compute-Aware Token Budget — Compute-constrained LLM training trajectories reveal non-monotonic degradation and instability, challenging endpoint-only evaluation.

See Counsel's argued verdicts on the open AI decisions leaders are weighing →

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by stat.ML updates on arXiv.org.