High-Dimensional Scaling Limits of Online Stochastic Gradient Descent in Single-Index Models

Abstract

We analyze the scaling limits of stochastic gradient descent (SGD) with a constant step size in the high-dimensional regime in single-index models. Specifically, we prove limit theorems for the trajectories of finite-dimensional summary statistics of SGD as the dimension tends to infinity. These scaling limits enable the analysis of both ballistic dynamics, described by a system of ordinary differential equations (ODEs), and diffusive dynamics, captured by a system of stochastic differential equations (SDEs). Additionally, we analyze a critical step-size scaling regime where, below this threshold, the effective ballistic dynamics align with the gradient flow of the population loss. In contrast, a new diffusive correction term appears at the threshold due to fluctuations around the fixed points. Furthermore, we discuss nearly sharp thresholds for the number of samples required for consistent estimation, which depend solely on an intrinsic property of the activation function known as the information exponent. Our main contribution is demonstrating that if a single-index model has an information exponent greater than two, the deterministic scaling limit, corresponding to the ballistic phase, or so-called dynamical mean-field theory in statistical physics, fails to achieve consistent estimation in high-dimensional inference problems. This shows the necessity of diffusive correction terms to accurately describe the dynamics of online SGD in single-index models via SDEs such as an Ornstein-Uhlenbeck process.