Universality of high-dimensional scaling limits of stochastic gradient descent
This addresses the theoretical understanding of SGD behavior in high-dimensional machine learning models, such as neural networks, by establishing universality conditions, though it is incremental in extending prior Gaussian-specific results to broader distributions.
The paper demonstrates that the ordinary differential equation (ODE) limits describing the evolution of summary statistics under stochastic gradient descent (SGD) in high-dimensional settings are universal across data distributions with matching first two moments, provided initialization and ground truth vectors are coordinate-delocalized, but shows non-universality in cases of coordinate-aligned initialization or stochastic fluctuations.
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.