A Scalable Measure of Loss Landscape Curvature for Analyzing the Training Dynamics of LLMs

TL;DR

This work addresses the need for scalable metrics of loss-landscape curvature in large language models by introducing critical sharpness $λ_c = 2/η_c$, estimable with fewer than 10 forward passes along the update direction. It provides a theoretical link between $λ_c$ and Hessian sharpness under a quadratic loss and demonstrates that $λ_c$ captures progressive sharpening and Edge of Stability across pre-training and mid-training, including up to 7B-parameter models. The authors further introduce relative critical sharpness $λ_c^{1→2}$ to study cross-landscape effects and data-mixing strategies, revealing a sweet spot in pre-training data fraction that balances retention of pre-trained capabilities with downstream adaptation. Collectively, the results offer a practical, scalable toolkit for diagnosing curvature dynamics and guiding data composition in large-scale model training, with implications for optimization stability and transfer performance.

Abstract

Understanding the curvature evolution of the loss landscape is fundamental to analyzing the training dynamics of neural networks. The most commonly studied measure, Hessian sharpness ($λ_{\max}^H$) -- the largest eigenvalue of the loss Hessian -- determines local training stability and interacts with the learning rate throughout training. Despite its significance in analyzing training dynamics, direct measurement of Hessian sharpness remains prohibitive for Large Language Models (LLMs) due to high computational cost. We analyze $\textit{critical sharpness}$ ($λ_c$), a computationally efficient measure requiring fewer than $10$ forward passes given the update direction $Δ\mathbfθ$. Critically, this measure captures well-documented Hessian sharpness phenomena, including progressive sharpening and Edge of Stability. Using this measure, we provide the first demonstration of these sharpness phenomena at scale, up to $7$B parameters, spanning both pre-training and mid-training of OLMo-2 models. We further introduce $\textit{relative critical sharpness}$ ($λ_c^{1\to 2}$), which quantifies the curvature of one loss landscape while optimizing another, to analyze the transition from pre-training to fine-tuning and guide data mixing strategies. Critical sharpness provides practitioners with a practical tool for diagnosing curvature dynamics and informing data composition choices at scale. More broadly, our work shows that scalable curvature measures can provide actionable insights for large-scale training.

Figures (11)

• Figure 1: Hessian sharpness exhibits progressive sharpening and Edge of Stability (EoS) under constant learning rate. The dashed lines corresponding to the learning rate mark the EoS threshold.
• Figure 2: Comparison of different sharpness measures on an illustrative landscape featuring a sharp valley direction and a flat river directionwen2025understanding. Hessian sharpness quantifies the curvature along the sharpest direction of the landscape (the valley). Directional sharpness measures the quadratic curvature along $\Delta \bm{\theta}$, while critical sharpness, its empirical counterpart, quantifies how far one can step along $\Delta \bm{\theta}$ before the loss increases.
• Figure 3: Comparison of different sharpness measures for MLPs trained on CIFAR-10 image classification task using SGD with learning rate $\eta$. Both critical and directional sharpness exhibit progressive sharpening and Edge of Stability, albeit some deviations from Hessian sharpness. The dashed line denotes the Edge of Stability threshold, given by $2/\eta$.
• Figure 4: Dynamics of pre-conditioned, directional, and critical sharpness during GPT-style Transformer training on FineWebEdu with AdamW using a Warmup-Stable-Decay (WSD) schedule. Critical sharpness tracks pre-conditioned sharpness throughout training, making it an effective proxy. The colored dashed lines denote the theoretical learning rate threshold $(2+2\beta_1)/(1-\beta_1)\eta$ and the two black vertical lines mark the end of warmup and stable phases.
• Figure 5: Critical Sharpness of OLMo-2 $7$B exhibits progressive sharpening throughout pre-training and mid-training. In the mid-training figure, the band around the mean trend shows the deviation across the three runs.

Theorems & Definitions (4)

• Definition 2.1: Critical learning rate $\eta_c$ and Sharpness $\lambda_c$
• Definition 2.2: Directional Sharpness $\lambda_{\text{dir}}$
• Definition 4.1: Relative Critical learning rate $\eta_c^{1 \to 2}$ and Sharpness $\lambda_c^{1 \to 2}$
• Lemma D.1: Stability Condition for Difference Equations