Hessian Spectral Analysis at Foundation Model Scale

TL;DR

This work makes faithful Hessian analysis tractable at foundation-model scale by introducing shard-preserving Hessian-vector products under Fully Sharded Data Parallelism, enabling scalable spectral and inverse-Hessian computations. By integrating shard-local finite-difference HvPs with stochastic Lanczos quadrature and CG-based solvers, the authors demonstrate near-linear scaling and modest overhead relative to first-order training while producing Hessian spectral densities for models up to 100B parameters. The study reveals that widely used block-diagonal curvature approximations can incur large errors and poor directional alignment, highlighting cross-layer coupling that such approximations miss. Overall, the approach unlocks principled curvature-based analyses for large-scale models, with practical implications for optimization, monitoring, data governance, and model compression, while showing that storing Lanczos vectors in bf16 suffices for accurate spectra with float32 arithmetic for reductions.

Abstract

Accurate Hessian spectra of foundation models have remained out of reach, leading most prior work to rely on small models or strong structural approximations. We show that faithful spectral analysis of the true Hessian is tractable at frontier scale. Using shard-local finite-difference Hessian vector products compatible with Fully Sharded Data Parallelism, we perform stochastic Lanczos quadrature on open-source language models with up to 100B parameters, producing the first large-scale spectral density estimates beyond the sub-10B regime. We characterize the numerical behavior of this pipeline, including finite-difference bias, floating-point noise amplification, and their effect on Krylov stability in fp32 and bf16, and derive practical operating regimes that are validated empirically. We further provide end-to-end runtime and memory scaling laws, showing that full-operator spectral probing incurs only a modest constant-factor overhead over first-order training. Crucially, direct access to the Hessian reveals that widely used block-diagonal curvature approximations can fail catastrophically, exhibiting order-one relative error and poor directional alignment even in mid-scale LLMs. Together, our results demonstrate that foundation-model Hessian spectra are both computable and qualitatively misrepresented by prevailing approximations, opening the door to principled curvature-based analysis at scale.

Figures (8)

• Figure 1: Local averaging suppresses high-frequency curvature without altering large-scale geometry. (a) Finite differences act as a local averaging operator on oscillatory signals. (b) A globally convex surface with high-frequency curvature. (c) Best achievable convergence under grid-searched learning rates for GD, momentum, and Adam for surface (b). (d) Adaptive Nesterov acceleration using pointwise versus finite-difference averaged curvature on surface (b).
• Figure 2: Spectral Density of OpenAI $120$Bn parameter open weight model on wikitext-$2$, sequence length $64$, $\epsilon=10^{-4}$.
• Figure 3: Pairwise structure for $\varepsilon = 10^{-5} \ll \varepsilon_{\mathrm{opt}}$. Matrix spectra resembles that of a random matrix with no clear outliar seperation.
• Figure 4: Pairwise structure for $\varepsilon = 10^{-2} \gg \varepsilon_{\mathrm{opt}}$. Matrix spectra resembles that of a random matrix with no clear outliar seperatioon.
• Figure 5: Pairwise structure for $\varepsilon = 10^{-4} \approx \varepsilon_{\mathrm{opt}}$.

Theorems & Definitions (9)

• Theorem 3.1: Error of gradient finite-difference HvPs
• Remark 3.2: Practical step sizes
• Theorem 3.3
• Theorem 4.1: Error of Lanczos with finite-difference Hessian vector products
• Remark 4.2
• proof
• proof
• Theorem 1.1
• proof