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

TL;DR

This work develops a rigorous, high-dimensional dynamical mean-field framework for the learning dynamics of multi-pass mini-batch SGD in isotropic, high-dimensional multi-index models. It introduces DMFT limit processes driven by a scalar Poisson jump (SGD) and a Gaussian diffusion (SME), showing that coordinate observables converge to a self-consistent limit characterized by kernels $(C_\theta,R_\theta,C_f,R_f,R_f^*,\Gamma)$ in a fixed-point system. A key result is that, for batch size scaling $\kappa\asymp n^\alpha$ with $\alpha\in[0,1)$, the SGD limit is independent of $\alpha$, while SGD, SME, and gradient flow are generally distinct except in the linear regression case where they coincide. The paper also establishes convergence of discrete-time DMFT to the continuous DMFT as the discretization $\delta\to0$, and provides detailed discretization-error bounds for both SGD and SME, along with numerical demonstrations highlighting when SGD/SME predictions align or differ in linear and nonconvex settings. Overall, the framework clarifies how stochastic gradient noise shapes high-dimensional learning dynamics and offers practical routes to simulate DMFT limits efficiently via the AMP-based discrete system.

Abstract

We study the learning dynamics of a multi-pass, mini-batch Stochastic Gradient Descent (SGD) procedure for empirical risk minimization in high-dimensional multi-index models with isotropic random data. In an asymptotic regime where the sample size $n$ and data dimension $d$ increase proportionally, for any sub-linear batch size $κ\asymp n^α$ where $α\in [0,1)$, and for a commensurate ``critical'' scaling of the learning rate, we provide an asymptotically exact characterization of the coordinate-wise dynamics of SGD. This characterization takes the form of a system of dynamical mean-field equations, driven by a scalar Poisson jump process that represents the asymptotic limit of SGD sampling noise. We develop an analogous characterization of the Stochastic Modified Equation (SME) which provides a Gaussian diffusion approximation to SGD. Our analyses imply that the limiting dynamics for SGD are the same for any batch size scaling $α\in [0,1)$, and that under a commensurate scaling of the learning rate, dynamics of SGD, SME, and gradient flow are mutually distinct, with those of SGD and SME coinciding in the special case of a linear model. We recover a known dynamical mean-field characterization of gradient flow in a limit of small learning rate, and of one-pass/online SGD in a limit of increasing sample size $n/d \to \infty$.

Figures (5)

• 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 $\|\boldsymbol{\theta}_t\|^2/d$, overlap $\langle \boldsymbol{\theta}_t, \boldsymbol{\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 \boldsymbol{\theta}| \le 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 \ref{['fig:tanh']}. 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\}$. 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).

Theorems & Definitions (40)

• Remark 2.4: Forms of the response processes
• Theorem 2.5
• Theorem 2.6
• Remark 2.7
• Remark 2.8
• Definition 3.1
• Lemma 3.2
• proof
• proof : Proof of Theorem \ref{['thm:uniqueness-existence']}(a)
• Lemma 3.3