Spectral alignment for high-dimensional SGD, Aukosh Jagannath

Over the last decade, a body of rich predictions has been made about the spectra of empirical Hessian and information matrices over the course of training (via SGD) in overparametrized networks. I'll present a recent work, in collaboration with G. Ben Arous (NYU Courant), R.Ghessari (Northwestern U.), and J. Huang (U. Penn), where we rigorously establish some of these predictions. 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.