Curvature-Aligned Probing for Local Loss-Landscape Stabilization

Nikita Kiselev; Andrey Grabovoy

cs.LGApr 16, 2026

Curvature-Aligned Probing for Local Loss-Landscape Stabilization

Authors: Nikita Kiselev, Andrey Grabovoy

Organizations: Department of Intelligent Systems Moscow Institute of Physics and Technology

Abstract

Local loss-landscape stabilization under sample growth is typically measured either pointwise or through isotropic averaging in the full parameter space. Despite practical value, both choices probe directions that contribute little to the dominant local deformation of strongly anisotropic neural landscapes. We recast stabilization as an observational problem and introduce a unified family of criteria parameterized by an aggregation order and a probing distribution; within this family we propose a curvature-aligned criterion $Δ_2^{(D)}$ that probes the loss increment field in the top- $D$ eigenspace of the empirical Hessian near a trained solution. Solely from a local quadratic model, we prove that $Δ_2^{(D)}$ preserves the $O(k^{-2})$ mean-squared rate of the full-space criterion while replacing ambient-dimension curvature dependence with dependence on the subspace dimension $D$ ; a corollary gives a closed-form spectral expression and a proposition identifies the top- $D$ eigenspace as extremal within the eigenspace-aligned family. We also derive scalable estimators based on Hessian-vector products, subspace Monte Carlo, and a closed-form Gaussian-moment proxy. On a decoder-only transformer, a curvature-aligned probe occupying a tiny fraction of parameter space already reproduces the full-space mean-squared signal to within numerical noise throughout the validated local regime, and the closed-form estimator is orders of magnitude faster than direct Monte Carlo after subspace construction.

Explore similar work

Jun 15, 2026cs.LG

Taming Curvature: Architecture Warm-Up for Stable Transformer Training

Training billion-parameter Transformers is often brittle, with transient loss spikes and divergence that waste compute. Even though the recently developed Edge of Stability (EoS) theory provides a powerful tool to understand and control the stability of optimization methods via the (preconditioned) curvature, these curvature-controlling methods are not popular in large-scale Transformer training due to the complexity of curvature estimation. To this end, we first introduce a fast online estimator of the largest (preconditioned) Hessian eigenvalue (i.e., curvature) based on a warm-started variant for power iteration with Hessian-vector products. We show theoretically, and verify empirically, that the proposed method makes per-iteration curvature tracking feasible at billion parameter scale while being more accurate. Using this tool, we find that training instabilities coincide with surges in preconditioned curvature and that curvature grows with depth. Motivated by these observations, we propose architecture warm-up: progressively growing network depth to carefully control the preconditioned Hessian and stabilize training. Experiments on large Transformers validate that our approach enables efficient curvature tracking and reduces instabilities compared to existing state-of-the-art stabilization techniques without slowing down convergence.

Sameera Ramasinghe, Ajanthan Thalaiyasingam, Hadi Mohaghegh Dolatabadi +6

Transformer Variant Curvature Prediction Quality Hessian Eigenvalue Low-Precision Training Warm Starts Newton

Ask PDF

May 22, 2026cs.LG

Spectral Asymptotics of Neural Network Loss Landscapes: An Exact Decomposition of the Curvature Exponent

The curvature exponent

α

h_k \propto σ_k^α

-- governing how Hessian eigenvalues scale with gradient singular values -- varies systematically across layer types (

α\approx 2

for convolutions,

\approx 1

for transformer attention,

< 1

for MLP up-projections). Why? We prove the Spectral Alignment Decomposition:

α= 2 + d\logΦ_k / d\logσ_k

, where

Φ_k

measures alignment between Kronecker factor eigenbases and gradient singular directions. This reduces "why does

α

vary?" to a geometric question we answer for LayerNorm, residual connections, and softmax heads. The decomposition implies a spectral transfer identity

s = αγ

linking curvature exponent, effective gradient rank-decay

γ

, and Hessian decay exponent

s

. The identity is algebraic; its empirical content is that

α

and

γ

, fit on independent data (HVPs vs. SVD), recover

s

to ~2% median error across 93 layers, five architectures, and three datasets -- with no free parameters. A zeta-function bound on participation ratio shows curvature concentrates onto effectively one direction per layer. As a proof of concept, we derive the architecture-adaptive preconditioner

T(σ;α)

and show that Spectral Newton -- implementing

T

in the gradient singular basis -- outperforms AdamW on vision benchmarks where

α\approx 2

Anherutowa Calvo

Hessian Eigenvalue Hessian Matrix Local Curvature Convex Loss Singular Value Decomposition Newton+1

Ask PDF

May 26, 2026cs.LG

Worker Disagreement Reveals Sharp Directions in Local SGD

Deep neural network training often exhibits highly anisotropic loss geometry, where a few sharp dominant Hessian directions coexist with a large flatter bulk. Gradients tend to align disproportionately with these dominant directions, although stable progress often requires movement through flatter bulk directions. Estimating the dominant subspace is therefore useful but costly with direct Hessian-based methods. We show that standard Local SGD exposes this geometry through worker disagreement. We theoretically show that the worker-average gap covariance is shaped by stochastic-gradient noise and Hessian curvature, causing workers to disagree along sharp, curvature-sensitive directions. Thus, worker-average gaps provide a cheap Hessian-free estimator of the dominant subspace. Experiments on MLPs, CNNs, and Transformers show that subspaces formed by worker-average gaps capture a substantial fraction of the gradient component lying in the dominant Hessian eigenspace.

Tolga Dimlioglu, Kristi Topollai, Anna Choromanska

Hessian Eigenvalue Stochastic Gradient Descent Covariance Estimation Anisotropy Directionality Disagreement

Ask PDF