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 κ≍nα with α∈[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 , η=0.8 , λ=0.1 . The dynamics of the squared norm ∥θt∥2/d , overlap ⟨θt,θ∗⟩/d , and training loss agree closely for SGD and SME, whereas those of the empirical CDF n−1∑i=1n1{∣xi⊤θ∣≤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 , η=3.0 , λ=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 τ=ηt , for various learning rates η∈{0.5,1.25,2.5} . In addition to the single-sample case, a fixed batch size κ=n1/3 is included for comparison; the strong alignment between these two simulations provides compelling support for our theory. At small learning rates ( η=0.5 ), the SME diffusion provides an accurate approximation of SGD, and their gap widens as η increases.
Figure 5. Convergence to one-pass SGD in the large sample limit. Overlap dynamics are plotted in units of rescaled time τ=γt , fixing d=4000 and varying γ=n/d∈{0.5,1.0,5.0,20.0} . As the training sample size γ 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).
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.
Nived Rajaraman, Yanjun Han
Microsoft Research NYC · Courant Institute of Mathematical Sciences and Center for Data Science, New York University
We develop a framework for analyzing the learning dynamics of high-dimensional problems trained using one-pass stochastic gradient descent (SGD) with data from multiple anisotropic classes. Our main theorem provides exact expressions for quantities of interest, including the risk and the overlap with the true signal, in terms of a deterministic system of ODEs, valid in the high-dimensional limit. The theorem holds for a broad class of optimization problems and extends to settings where the number of classes grows with dimension. To illustrate its utility, we investigate in detail the effect of the data's anisotropic structure on the problems of binary logistic regression and least-squares (LS) loss. We study the LS in a linear multiclass setup and derive a learning-rate threshold that depends on the average eigenvalue of the covariance matrices. In the binary logistic regression, we study three cases: isotropic covariances, data covariance matrices with a large fraction of zero eigenvalues (denoted as the zero-one model), and covariance matrices with power-law spectra. We show that a structural phase transition occurs. In particular, for the zero-one model and the power-law model with sufficiently large power, SGD aligns more closely with values of the class mean that are projected onto the ``clean directions'' (i.e., directions of smaller variance). This is supported by analytical studies and numerical simulations, which show the exact asymptotic behavior of the loss in the high-dimensional limit. The effects of data anisotropy that we demonstrate are likely to hold beyond these examples and illustrate one application of the broader theorem that we prove.
Understanding the behavior of stochastic gradient methods is a central problem in modern machine learning. Recent work has highlighted diagonal linear networks as a simplified yet expressive setting for analyzing the optimization and generalization properties of neural models. In this work, we show that in the high-dimensional regime, stochastic gradient descent on diagonal linear networks is well-approximated by continuous dynamics governed by a stochastic differential equation (SDE), which explicitly decouples the drift from the gradient noise. We further derive a deterministic partial differential equation whose solution propagates the relevant state of the iterates and characterizes the time evolution of a broad class of observable statistics, including the risk, curvature, and other metrics for optimality. Finally, we show that, under a suitable parametrization, the stochastic dynamics are globally well posed and converge exponentially fast to zero risk with high probability, yielding a fully explicit non-asymptotic description of their long-time behavior. Numerical simulations corroborate our theoretical findings.
Begoña García Malaxechebarría, Courtney Paquette, Maryam Fazel +1
University of Washington · McGill University · Google DeepMind +2