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

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 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.

• The paper introduces critical sharpness as a scalable proxy for loss curvature that accurately tracks training dynamics in large language models.
• It employs an efficient line search method requiring only 5–6 forward passes to estimate the critical learning rate, bypassing expensive Hessian computations.
• It also defines relative critical sharpness to inform data mixing strategies and mitigate catastrophic forgetting in multi-task learning scenarios.

Scalable Curvature Measures for Loss Landscape Analysis in LLM Training Dynamics

Introduction

This work addresses a critical challenge in deep learning: the quantification and analysis of loss landscape curvature at scale, particularly in the context of LLMs. The curvature of the loss landscape—often analyzed via the Hessian of the loss with respect to model parameters—has been shown to dictate training stability, optimization behavior, and, to some extent, generalization. Traditional sharpness measures, primarily the largest Hessian eigenvalue ($\lambda_{\max}^H$), are computationally prohibitive for LLMs due to the scale and the ubiquity of training primitives such as fused kernels with unavailable second-derivative implementations. This paper introduces and systematically studies critical sharpness ($\lambda_c$), a scalable, optimizer-aware curvature measure, and its extensions, thereby enabling actionable insights into curvature-driven training dynamics and data composition strategies at the scale of billions of parameters (2601.16979).

Critical Sharpness: Definition and Estimation

The critical sharpness $\lambda_c$ is defined as $2/\eta_c$, where $\eta_c$ is the smallest learning rate that would increase the loss in the model’s current update direction. Unlike Hessian-based measures, $\lambda_c$ is evaluated via an efficient line search requiring only a handful of forward passes, making it well-suited for distributed and large-batch LLM training.

Figure 1: Comparison of sharpness measures in an illustrative landscape with both sharp and flat directions. Hessian sharpness is maximal along the steepest direction; critical sharpness instead quantifies the maximal safe step along the update direction.

The critical sharpness measure thus directly connects geometric curvature to stable optimization: it estimates the "natural length scale" of the landscape along update directions used by practical optimizers, includes optimizer effects (such as preconditioning in Adam), and is robust to implementation details that defeat explicit Hessian computation.

Empirically, the estimation procedure—composed of exponential and binary search along the update direction—is shown to reliably converge to $\eta_c$ in approximately $5$–$6$ forward passes, requiring neither gradients nor Hessian-vector products, and thus scales to models where Hessian sharpness is intractable.

Theoretical Foundations: Relationship to Hessian Sharpness

Under a quadratic approximation to the loss, critical sharpness is intimately related to both directional sharpness (curvature along the update direction) and Hessian sharpness. Specifically, for gradient descent, directional sharpness $\lambda_{\mathrm{dir}}$ is a weighted sum of Hessian eigenvalues, weighted by the alignment of the gradient with each eigendirection:

$$\lambda_{\mathrm{dir}} = \frac{\sum_i c_i^2 \lambda_i^H}{\sum_i c_i^2} \leq \lambda_{\max}^H$$

where $c_i$ is the projection of the gradient onto the $i$th Hessian eigenvector. Equality is achieved when the gradient perfectly aligns with the sharpest direction. For adaptive optimizers, a similar weighted relationship holds for the preconditioned Hessian. In non-convex or non-quadratic landscapes, critical sharpness becomes a robust empirical proxy that tracks the maximal stable learning rate along realistic update directions.

Figure 2: Progressive sharpening and the onset of Edge-of-Stability in batch gradient descent, as observed by multiple sharpness measures for large and small batch sizes.

Empirical Evaluation: Progressive Sharpening and Edge-of-Stability at Scale

A central investigation concerns whether well-documented sharpness phenomena—including progressive sharpening (the increase of sharpness over training) and Edge of Stability (EoS, the regime where sharpness fluctuates around a critical threshold)—manifest at LLM scale and are captured by critical sharpness.

Across a variety of architectures (e.g., GPT-style Transformers) and optimizers (e.g., AdamW), critical sharpness faithfully tracks the progression of preconditioned Hessian sharpness, mirroring the dynamics predicted by theory and previously observed for small models. In particular, with realistic learning rate schedules (warmup followed by decay), critical sharpness reflects the lowering and subsequent oscillation around the stability threshold, verifies progressive sharpening, and identifies instability regions.

Figure 3: Relative critical sharpness for Dolmino mix subsets during OLMo-2 mid-training, with variation bands denoting batch-wise variability.

Via checkpoint analysis of OLMo-2 models at $7$B scale, the paper provides evidence for continued progressive sharpening throughout both pre-training and mid-training, as quantified by critical sharpness.

Relative Critical Sharpness and Data Mixing

A striking and novel extension is the introduction of relative critical sharpness $\lambda_c^{1 \to 2}$, quantifying how updates from optimizing one loss (e.g., a finetuning objective) affect the curvature landscape of another (e.g., pretraining loss). This is formalized as the maximum allowable learning rate before the pretraining loss increases when stepping along the update direction derived from the finetuning task.

By measuring $\lambda_c^{1 \to 2}$ across mixes of pretraining and downstream finetuning data, the analysis reveals the impact of data mixing on catastrophic forgetting and downstream task trade-offs. Key empirical findings include:

• When the training mix is dominated by finetuning data, pretraining loss landscapes become highly sensitive, resulting in high critical sharpness (i.e., a narrow, sharp basin), and rapid forgetting occurs.
• As the proportion of pretraining data in the mix increases, the critical sharpness for pretraining tasks drops and the stable basin expands, but excessive emphasis can limit specialization.
• There exists a "sweet spot" in the pretraining data ratio where critical sharpness curves for distinct tasks intersect, maximizing multi-task learning rates without instability.

Performance heatmaps as a function of data mix ratio and learning rate confirm the predictive power of critical sharpness: optimal GSM8K (math) improvements occur outside the pretraining basin, while general capabilities (MMLU) are preserved only within the basin, aligning with curvature-based predictions.

Implications and Future Directions

The computational efficiency and diagnostic reliability of critical sharpness and its relative variant have clear practical implications:

• Training Diagnostics: Practitioners can monitor sharpness and EoS onset at scale, diagnosing instabilities or suboptimal learning rate schedules in LLM training.
• Data Composition: The relative critical sharpness measure informs optimal data mixing strategies for continual learning and finetuning, avoiding exhaustive ablation studies or the need for downstream evaluations at every mixing ratio.
• Forgetting Mitigation: The analytical framework enables systematic mitigation of catastrophic forgetting in LLMs, supporting tailored balancing of specialization and retention.
• Curvature-Aware Optimization: The alignment between update directions and sharpest landscape directions, as made explicit in this framework, motivates future methods for adaptive learning rate schedules, curvature-aware optimization, and generalization control.

Theoretically, these findings generalize previously local observations of sharpness and EoS to the modern large-scale regime, and suggest that optimizer-aware curvature measures are fundamentally more robust and practical than Hessian-based metrics in nonconvex, high-dimensional landscapes.

Conclusion

This paper establishes critical sharpness as a scalable, empirically robust proxy for loss landscape curvature in large-scale deep learning. Through rigorous theoretical analysis, systematic empirical validation, and practical demonstrations in real-world LLM pretraining and finetuning regimes, it enables actionable understanding of optimization dynamics, data composition, and catastrophic forgetting. The introduction of relative critical sharpness represents a versatile analytical tool for continual and multi-task learning contexts, pointing towards future developments in curvature-aware training algorithms and robust, scalable model diagnostics.