High-dimensional SGD aligns with emerging outlier eigenspaces
Gerard Ben Arous Reza Gheissari Jiaoyang Huang Aukosh Jagannath
We rigorously study the joint evolution of training dynamics via stochastic
gradient descent (SGD) and the spectra of empirical Hessian and gradient
matrices. We prove that in two canonical classification tasks for multi-class
high-dimensional mixtures and either 1 or 2-layer neural networks, the SGD
trajectory rapidly aligns with emerging low-rank outlier eigenspaces of the
Hessian and gradient matrices. Moreover, in multi-layer settings this alignment
occurs per layer, with the final layer's outlier eigenspace evolving over the
course of training, and exhibiting rank deficiency when the SGD converges to
sub-optimal classifiers. This establishes some of the rich predictions that
have arisen from extensive numerical studies in the last decade about the
spectra of Hessian and information matrices over the course of training in
overparametrized networks.