High-Dimensional Limit of Stochastic Gradient Flow via Dynamical Mean-Field Theory
Sota Nishiyama ⋅ Masaaki Imaizumi
Abstract
We study the high-dimensional dynamics of multi-pass stochastic gradient descent (SGD) with small batch sizes. Existing asymptotic descriptions of SGD primarily cover online learning, large-batch regimes, or linear models. We address this gap by analyzing a continuous-time stochastic differential equation, called \emph{stochastic gradient flow} (SGF), which approximates small-batch multi-pass SGD. In the proportional limit $n,d\to\infty$ with $n/d\to\delta$, we derive a closed dynamical mean-field theory (DMFT) system and prove that it characterizes the asymptotic empirical distribution of the SGF parameters. The framework applies to nonlinear models including generalized linear models and two-layer neural networks, and recovers several existing high-dimensional descriptions, including noiseless gradient flow, online SGD, and high-dimensional linear regression.
Chat is not available.
Successful Page Load