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

arXiv: https://arxiv.org/abs/2601.21093 · HTML: https://mlsanity.com/paper/2601.21093/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 28 Jan 2026.

## Built on by (most cited first, 3 of 3)

- [Statistical Inference on Gradient Flows](https://mlsanity.com/paper/2606.01257/) (2026)
- [Decoupled Descent: Exact Test Error Tracking Via Approximate Message Passing](https://mlsanity.com/paper/2604.27883/) (2026)
- [Homogenization of ℓ2-Adversarial Training in High-Dimensions: Exact Dynamics under Stochastic Gradient Descent](https://mlsanity.com/paper/2607.00207/) (2026)

## Abstract (arXiv)

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 kappa asymp nα where α ∈ [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 α ∈ [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 → ∞.