Characterizing Optimizer-Dependent Training Dynamics Through Hessian Eigenvector Displacement and Localization

Technical Summary: Characterizing Optimizer-Dependent Training Dynamics Through Hessian Eigenvector Displacement and Localization

Problem Statement
Understanding neural network training remains challenging due to the high-dimensional, non-convex nature of the loss landscape. While the Hessian matrix (the matrix of second derivatives of the loss with respect to parameters) is a standard tool for analyzing training dynamics, generalization, and sharpness, most prior work focuses on Hessian eigenvalues (curvature magnitude) or the stability of the top-$k$ eigenspace as a whole. A critical gap exists in understanding the evolution of individual leading Hessian eigenvectors, which encode the specific directions of curvature in parameter space. It is unclear whether these eigenvectors stabilize, drift, or reorganize during mini-batch training, and whether different optimizers induce qualitatively distinct geometric trajectories through the loss landscape.

Methodology
The authors investigate the time-dependent evolution of the leading Hessian eigenvectors, $\{v^{(i)}(t)\}$, during the training of multilayer perceptrons (MLPs) on MNIST and FashionMNIST datasets. The study compares three optimizers: Stochastic Gradient Descent (SGD), Adam, and Sharpness-Aware Minimization (SAM).

To quantify eigenvector dynamics, the paper employs two complementary metrics:

1. Eigenvector Displacement: Inspired by analyses of glassy systems, this metric measures the squared Euclidean distance between an eigenvector at a starting time $t_w$ and its state after a waiting time $t$. This acts as a two-time correlation function to distinguish between stationary dynamics and "aging" behavior (where the rate of change depends on the training age). A random null model, based on the expected Hessian under random parameter initialization, serves as a baseline for complete decorrelation.
2. Inverse Participation Ratio (IPR): This metric quantifies the localization of eigenvector mass. An IPR value near $1/N$ (where $N$ is the number of parameters) indicates delocalization (uniform spread), while values near $O(1)$ indicate localization (concentration on a small subset of parameters). In this context, IPR measures how sensitivity to curvature is distributed across the network's weights.

The experiments utilize a 3-layer fully-connected MLP trained with mini-batches. The Hessian is evaluated on the full dataset at intervals, and eigenvectors are computed via Lanczos iteration using Hessian-vector products to avoid forming the full matrix.

Key Results
The study reveals clear, optimizer-dependent behaviors in both displacement and localization dynamics:

• Displacement and Stability:

  • SGD: Exhibits "aging" behavior where the leading eigenvector ($v^{(1)}$) initially changes rapidly but progressively stabilizes. The displacement saturates below the random baseline, indicating that SGD leads to progressively more stable leading curvature directions.
  • Adam: Shows substantially stronger reorganization. The eigenvector displacement flattens near the random baseline, suggesting near-complete decorrelation and persistent shifting of the leading curvature directions throughout training.
  • SAM: Displays intermediate behavior with slower but persistent drift compared to SGD.
  • These dynamics are robust across variations in batch size and learning rate, except in the full-batch regime.

• Localization (IPR):

  • Adam: Induces a strong localization phenomenon where the leading eigenvector becomes concentrated on a small subset of parameters (IPR $\approx 0.1$). This suggests that under Adam, a small subset of weights disproportionately contributes to the leading curvature directions.
  • SGD: Shows a delocalizing trend, with IPR values decreasing toward the uniform baseline ($1/N$), indicating that curvature sensitivity spreads more evenly across parameters.
  • SAM: Exhibits slightly higher localization than SGD but maintains low variance across runs.

• Alignment: The paper notes that while gradients generally align with top eigenvectors, Adam's update alignment remains low (under 0.1), likely due to its momentum parameter, whereas SGD and SAM updates align more closely with the gradient.

Significance and Claims
The paper claims that tracking Hessian eigenvector dynamics captures key differences in optimizer behavior that are not visible through spectral analysis alone. The primary significance lies in demonstrating that:

1. Optimizer Choice Dictates Geometry: Different optimizers do not merely traverse the same landscape at different speeds; they induce qualitatively distinct forms of landscape exploration. SGD stabilizes curvature directions, while Adam continuously reorganizes them.
2. Localization as a Stable Statistic: Despite the mixing of eigenvectors under Adam, the localization metric (IPR) remains a stable statistic throughout training. This suggests that even as the specific directions of curvature rotate, the nature of the curvature (whether it is driven by a few parameters or distributed) remains consistent for a given optimizer.
3. Implications for Analysis: These findings suggest that eigenvector instabilities may facilitate optimization by allowing the exploration of new regions of the loss landscape. Furthermore, the observed localization under Adam may have practical implications for targeted pruning and uncertainty quantification, as it identifies specific subsets of weights that dominate the most convex directions.

The authors conclude that these dynamics are robust to hyperparameter variations (within mini-batch regimes) and that future work could leverage these observations to improve optimizer dynamics and exploit parameter redundancy.