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

2026-02-06

AI Summary

This work addresses the absence of an analytical framework for characterizing the dynamics of small-batch, multi-pass stochastic gradient descent (SGD) in high-dimensional nonlinear models. The authors propose the Stochastic Gradient Flow (SGF) model, which extends dynamical mean-field theory (DMFT) to stochastic settings for the first time, yielding a low-dimensional system of continuous-time equations that precisely describe the asymptotic distribution of parameters in the high-dimensional limit. This approach unifies existing theories for online SGD and high-dimensional linear regression, and applies broadly to generalized linear models and two-layer neural networks. Moreover, the framework recovers several established SGD dynamics as special cases through closed-form equations, thereby providing a principled and general theoretical foundation for analyzing SGD in complex, high-dimensional learning scenarios.

Technology Category

• Machine Learning: Learning with Manifolds
• Reasoning under Uncertainty: Stochastic Optimization
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Web Mining and Content Analysis: Models for Web evolution
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• high-dimensional limit
• stochastic gradient descent
• nonlinear models
• stochastic gradient flow
• dynamical mean-field theory

Innovation

• Stochastic Gradient Flow
• Dynamical Mean-Field Theory
• High-Dimensional Asymptotics
• Multi-pass SGD
• Nonlinear Models