High-Dimensional Limit of Stochastic Gradient Flow via Dynamical Mean-Field Theory

TL;DR

This paper develops a dynamical mean-field theory framework to analyze the high-dimensional asymptotic behavior of stochastic gradient flow, providing a unifying perspective on SGD dynamics in complex models.

Contribution

It introduces a novel DMFT-based analytical approach to characterize the high-dimensional limit of stochastic gradient flow for nonlinear models, extending existing theories.

Findings

Derives low-dimensional equations describing SGF in high dimensions

Recovers existing high-dimensional SGD descriptions as special cases

Provides a unified theoretical framework for analyzing SGD dynamics

Abstract

Modern machine learning models are typically trained via multi-pass stochastic gradient descent (SGD) with small batch sizes, and understanding their dynamics in high dimensions is of great interest. However, an analytical framework for describing the high-dimensional asymptotic behavior of multi-pass SGD with small batch sizes for nonlinear models is currently missing. In this study, we address this gap by analyzing the high-dimensional dynamics of a stochastic differential equation called a \emph{stochastic gradient flow} (SGF), which approximates multi-pass SGD in this regime. In the limit where the number of data samples $n$ and the dimension $d$ grow proportionally, we derive a closed system of low-dimensional and continuous-time equations and prove that it characterizes the asymptotic distribution of the SGF parameters. Our theory is based on the dynamical mean-field theory (DMFT) and is applicable to a wide range of models encompassing generalized linear models and two-layer neural networks. We further show that the resulting DMFT equations recover several existing high-dimensional descriptions of SGD dynamics as special cases, thereby providing a unifying perspective on prior frameworks such as online SGD and high-dimensional linear regression. Our proof builds on the existing DMFT technique for gradient flow and extends it to handle the stochasticity in SGF using tools from stochastic calculus.