Universality of high-dimensional scaling limits of stochastic gradient descent

TL;DR

This paper demonstrates that the high-dimensional limits of stochastic gradient descent (SGD) are universal across a broad class of data distributions, provided certain conditions are met, and explores conditions where universality does not hold.

Contribution

The paper proves the universality of the ODE limits for SGD in high dimensions across various mixture distributions, extending previous Gaussian-specific results.

Findings

ODE limits are universal for mixtures with matching first two moments.

Universality fails when initialization is coordinate aligned.

Fluctuation limits are not universal across different distributions.

Abstract

We consider statistical tasks in high dimensions whose loss depends on the data only through its projection into a fixed-dimensional subspace spanned by the parameter vectors and certain ground truth vectors. This includes classifying mixture distributions with cross-entropy loss with one and two-layer networks, and learning single and multi-index models with one and two-layer networks. When the data is drawn from an isotropic Gaussian mixture distribution, it is known that the evolution of a finite family of summary statistics under stochastic gradient descent converges to an autonomous ordinary differential equation (ODE), as the dimension and sample size go to $\infty$ and the step size goes to $0$ commensurately. Our main result is that these ODE limits are universal in that this limit is the same whenever the data is drawn from mixtures of arbitrary product distributions whose first two moments match the corresponding Gaussian distribution, provided the initialization and ground truth vectors are coordinate-delocalized. We complement this by proving two corresponding non-universality results. We provide a simple example where the ODE limits are non-universal if the initialization is coordinate aligned. We also show that the stochastic differential equation limits arising as fluctuations of the summary statistics around their ODE's fixed points are not universal.