Connecting Parameter Magnitudes and Hessian Eigenspaces at Scale using Sketched Methods

TL;DR

The paper addresses the link between early crystallization of magnitude-based parameter masks and the top-$k$ Hessian eigenspaces in deep networks. It introduces a Grassmannian-based framework to compare the spans of $k$-sparse masks and top-$k$ Hessian eigenvectors, with overlap chosen as the interpretable similarity measure. To enable large-scale analysis, the authors develop SEIGH, a matrix-free sketched-SVD method that computes Hessian eigendecompositions with $O(Dk)$ memory and can yield over $10^3$ eigenpairs for models with more than $10^7$ parameters. Empirical results show that the overlap between magnitude masks and top Hessian eigenspaces is consistently above random chance and grows with model size, suggesting that large parameters tend to align with directions of larger loss curvature. This framework provides a scalable way to analyze Hessians in DL and offers a new perspective on the structure of neural networks beyond pruning alone.

Abstract

Recently, it has been observed that when training a deep neural net with SGD, the majority of the loss landscape's curvature quickly concentrates in a tiny *top* eigenspace of the loss Hessian, which remains largely stable thereafter. Independently, it has been shown that successful magnitude pruning masks for deep neural nets emerge early in training and remain stable thereafter. In this work, we study these two phenomena jointly and show that they are connected: We develop a methodology to measure the similarity between arbitrary parameter masks and Hessian eigenspaces via Grassmannian metrics. We identify *overlap* as the most useful such metric due to its interpretability and stability. To compute *overlap*, we develop a matrix-free algorithm based on sketched SVDs that allows us to compute over 1000 Hessian eigenpairs for nets with over 10M parameters --an unprecedented scale by several orders of magnitude. Our experiments reveal an *overlap* between magnitude parameter masks and top Hessian eigenspaces consistently higher than chance-level, and that this effect gets accentuated for larger network sizes. This result indicates that *top Hessian eigenvectors tend to be concentrated around larger parameters*, or equivalently, that *larger parameters tend to align with directions of larger loss curvature*. Our work provides a methodology to approximate and analyze deep learning Hessians at scale, as well as a novel insight on the structure of their eigenspace.

Figures (19)

• Figure 1: Overlap between top-$k$ parameter magnitude masks and top-$k$ Hessian eigenspaces is consistently and substantially above random chance.(Left) Measurements for a $7030.0$-parameter network trained on $16\!\times\!16$ downsampled MNIST until convergence (see \ref{['sec:exp_tinymnist']} and \ref{['fig:collapse']}): Exact $\emph{overlap}$ between top-$k$ parameters and eigenvectors with $k\!=\!350.0$ ( ), approximate $\emph{overlap}$ via sketched eigendecomposition ( ), and chance-level baseline ( ) (see \ref{['sec:metrics_comparison']}). Note how $\emph{overlap}$ is larger than chance, and sketched $\emph{overlap}$ is a good approximation. (Right) Lines show the ratio sketched $\emph{overlap}$ vs. chance-level baseline for a model with $>\!11.0$M parameters trained on ImageNet (see \ref{['sec:exp_largescale']}), at three different points during training, and as a function of $k$. Note how $\emph{overlap}$ is always higher than baseline, up to a factor of 1000.
• Figure 2: Behaviour of different Grassmannian metrics for random pairs of matrices and masks.(Left) As a function of $D$, for different sparsity ratios $\rho\!\coloneqq\!k/D$. (Right) As a function of $\rho$, for different ambient dimensions $D$. Lines show the median metric for $50$ random pairs, the (almost imperceptible) shaded regions span the $5$-$95$ percentiles (see \ref{['fig:synth_r_BB']} for broader distributions). Each row shows a selected Grassmannian metric (see \ref{['sec:metrics_grassmann']} for more metrics): $\mathrm{dist}_{\overline{c,2}}$ is a representative example of a collapsing metric, being $\approx\!0.0$ almost everywhere. The $\emph{overlap}$ metric is non-collapsing, and its expectation equals $\rho$.
• Figure 3: Overlap between parameter magnitude masks and top Hessian eigenspaces as a function of $k$ at three different points during the training process. Shown is the sketched $\emph{overlap}$ between the mask of the largest $k$ parameters by magnitude and the top-$k$ Hessian eigenspace. Consistently across the training process and for all problems and values of $k$, we observe an $\emph{overlap}$ substantially larger than the random chance baseline ( ) introduced in \ref{['sec:metrics_comparison']}.
• Figure 4: (Left)Overlap between parameter magnitude masks and top Hessian eigenspaces as a function of training step. Shown is the sketched $\emph{overlap}$ between masks of $k$-largest parameters by magnitude and top-$k$ Hessian eigenspaces, for different values of $k$ at different training steps. The observed $\emph{overlap}$ raises early on, and is consistently and substantially above random chance baseline ( , introduced in \ref{['sec:metrics_comparison']}). (Right)Factor by which the measured sketched $\emph{overlap}$ is greater than the random baseline. For most of the measured training process and choices of $k$, the observed $\emph{overlap}$ is at least $10\times$ larger than random chance baseline. This multiple factor over the random baseline is largest for small $k$ and increases with network size, surpassing $10^3$.
• Figure 5: Training evolution of the model used to report \ref{['fig:perturb_train', 'fig:perturb_test', 'fig:rand_mask_baseline']}. It corresponds to the $7030.0$-parameter model trained on $16\!\times\!16$ downsampled MNIST, as described in \ref{['sec:exp_tinymnist']}.

Theorems & Definitions (2)

• Lemma A.1
• proof