High-dimensional learning dynamics of multi-pass Stochastic Gradient Descent in multi-index models

Zhou Fan; Leda Wang

stat.MLJan 28, 2026

High-dimensional learning dynamics of multi-pass Stochastic Gradient Descent in multi-index models

Authors: Zhou Fan, Leda Wang

Organizations: Department of Statistics and Data Science, Yale University

Abstract

The learning dynamics of Stochastic Gradient Descent (SGD) are believed to play an important role in influencing the generalization properties of trained models. While the precise dynamics of how online (one-pass) SGD optimizes the population risk are now understood in a broad range of high-dimensional problems, a comparable understanding of multi-pass SGD for empirical risk minimization is still largely lacking. In this work, we develop an asymptotically exact characterization of the dynamics of multi-pass SGD for a non-convex multi-index model, in the high-dimensional regime where sample size $n$ and data dimension $d$ increase proportionally. We consider any sub-linear batch size $κ\asymp n^α$ with $α\in [0,1)$ and a commensurate ``critical'' scaling of the learning rate, and show that the coordinate-wise dynamics of SGD admit a novel continuous-time scaling limit, given by a system of dynamical mean-field equations that are driven by a scalar Poisson jump process representing the asymptotic limit of SGD sampling noise. This scaling limit connects several previously disparate theories around SGD, namely (i) the classical interpretation of SGD as an approximation to gradient flow on the empirical risk, (ii) the characterization of online SGD via an autonomous ODE in its overlap statistics, and (iii) the coincidence between multi-pass SGD dynamics and diffusion approximations in least-squares problems. We illustrate how these perspectives are recovered by specializing our result to small learning rates, large sample sizes, and linear models, respectively. Outside of these specializations, the scaling limit of SGD is distinct both from diffusion approximations and gradient flow, and may enable SGD to converge to local optima of the non-convex empirical risk with different statistical properties.

Figures & tables

Figure 1. Comparison of SGD and SME dynamics in a linear model. Parameters: $n=8000$ , $d=10000$ , $\eta=0.8$ , $\lambda=0.1$ . The dynamics of the squared norm $\|\bm{\theta}_{t}\|^{2}/d$ , overlap $\langle\bm{\theta}_{t},\bm{\theta}^{*}\rangle/d$ , and training loss agree closely for SGD and SME, whereas those of the empirical CDF $n^{-1}\sum_{i=1}^{n}\mathbf{1}\{|\mathbf{x}_{i}^{\top}\bm{\theta}|\leq 1\}$ show a discrepancy.

Figure 2. Divergence of SGD and SME dynamics in a model with Huber loss and $\tanh$ activation. Parameters: $n=8000$ , $d=10000$ , $\eta=3.0$ , $\lambda=0.1$ . Solid/dotted curves depict the dynamics of SGD/SME and their DMFT predictions over 4 training epochs, exhibiting a divergence that is predicted by DMFT. Dashed curves for training times $t>4$ depict the subsequent dynamics of gradient flow initialized from the last SGD/SME iterate, in which the norm and overlap statistics re-converge.

Figure 3. Divergence of SGD and SME dynamics in a model with Huber loss and $\sin$ activation. Parameters are the same as Figure 2 . In this example, the difference between SGD and SME persists through the later gradient flow training, suggesting that the dynamics reach local optima with distinct statistical properties.

Figure 4. Effect of learning rate on SGD vs. SME divergence. Overlap dynamics are plotted in units of rescaled time $\tau=\eta t$ , for various learning rates $\eta\in\{0.5,1.25,2.5\}$ . In addition to the single-sample case, a fixed batch size $\kappa=n^{1/3}$ is included for comparison; the strong alignment between these two simulations provides compelling support for our theory. At small learning rates ( $\eta=0.5$ ), the SME diffusion provides an accurate approximation of SGD, and their gap widens as $\eta$ increases.

Figure 5. Convergence to one-pass SGD in the large sample limit. Overlap dynamics are plotted in units of rescaled time $\tau=\gamma t$ , fixing $d=4000$ and varying $\gamma=n/d\in\{0.5,1.0,5.0,20.0\}$ . As the training sample size $\gamma$ increases, multi-pass SGD dynamics (solid lines) and their DMFT predictions (dotted lines) both converge to the simulated overlap statistics of one-pass SGD (black dashed line).

Dynamical Mean-Field Theory

See all 8 papers

Interactive Learning of Single-Index Models via Stochastic Gradient Descent
Feb 19, 2026Nived Rajaraman, Yanjun HanSingle-Index Models Generalized Linear Bandits
Microsoft Research NYC · Courant Institute of Mathematical Sciences and Center for Data Science, New York University
Exact Dynamics of Multi-class Stochastic Gradient Descent
Oct 15, 2025Elizabeth Collins-Woodfin, Inbar SeroussiLogistic Regression High-Dimensional Asymptotics
High-dimensional Limit of SGD for Diagonal Linear Networks
May 16, 2026Begoña García Malaxechebarría, Courtney Paquette, Maryam Fazel +1Stochastic Optimization Convergence Neural Network Training Dynamics
University of Washington · McGill University · Google DeepMind +2

Feb 19, 2026J/K move · Enter open · S save

Interactive Learning of Single-Index Models via Stochastic Gradient Descent

Nived Rajaraman, Yanjun Han

Microsoft Research NYC · Courant Institute of Mathematical Sciences and Center for Data Science, New York University

Single-Index Models Generalized Linear Bandits Stochastic Gradient Descent Statistical Learning Theory

Open paper AskPDF arXiv

Stochastic gradient descent (SGD) is a cornerstone algorithm for high-dimensional optimization, renowned for its empirical successes. Recent theoretical advances have provided a deep understanding of how SGD enables feature learning in high-dimensional nonlinear models, most notably the single-index model with i.i.d. data. In this work, we study the sequential learning problem for single-index models, also known as generalized linear bandits or ridge bandits, where SGD is a simple and natural solution, yet its learning dynamics remain largely unexplored. We show that, similar to the optimal interactive learner, SGD undergoes a distinct burn-in'' phase before entering the learning'' phase in this setting. Moreover, with an appropriately chosen learning rate schedule, a single SGD procedure simultaneously achieves near-optimal (or best-known) sample complexity and regret guarantees across both phases, for a broad class of link functions. Our results demonstrate that SGD remains highly competitive for learning single-index models under adaptive data.

Abstract

Dynamical Mean-Field Theory

Explore similar work

Interactive Learning of Single-Index Models via Stochastic Gradient Descent

Exact Dynamics of Multi-class Stochastic Gradient Descent

High-dimensional Limit of SGD for Diagonal Linear Networks