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

• URL: http://arxiv.org/abs/2601.21093v1
• Date: Wed, 28 Jan 2026 22:28:12 GMT
• Title: High-dimensional learning dynamics of multi-pass Stochastic Gradient Descent in multi-index models
• Authors: Zhou Fan, Leda Wang,
• Abstract summary: We study the learning dynamics of a multi-pass, mini-batch Gradient Descent (SGD) procedure for empirical risk minimization.<n>In an limit regime where the sample size $n$ and data dimension $d$ increase proportionally, for any sub-linear batch size $asymp n where $in [0,1)$, we provide anally exact characterization of the coordinate-wise dynamics of SGD.
• Score: 2.2129910930772
• License: http://creativecommons.org/licenses/by/4.0/
• 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$.

Related papers

• High-Dimensional Limit of Stochastic Gradient Flow via Dynamical Mean-Field Theory [6.2000582635449994]
  Modern machine learning models are typically trained via multi-pass gradient descent (SGD) with small batch sizes.<n>We analyze the high-dimensional dynamics of a differential equation called a emphstochastic gradient flow (SGF)<n>We show that the resulting DMFT equations recover several existing high-dimensional descriptions of SGD dynamics as special cases.
  arXiv  Detail & Related papers  (2026-02-06T02:37:10Z)
• Exact Dynamics of Multi-class Stochastic Gradient Descent [4.1538344141902135]
  We develop a framework for analyzing the training and learning rate dynamics on a variety of high-dimensional optimization problems trained using one-pass gradient descent (SGD)<n>We give exact expressions for a large class of functions of the limiting dynamics, including the risk and the overlap with the true signal, in terms of a deterministic solution to a system of ODEs.
  arXiv  Detail & Related papers  (2025-10-15T20:31:49Z)
• G$^2$RPO: Granular GRPO for Precise Reward in Flow Models [74.21206048155669]
  We propose a novel Granular-GRPO (G$2$RPO) framework that achieves precise and comprehensive reward assessments of sampling directions.<n>We introduce a Multi-Granularity Advantage Integration module that aggregates advantages computed at multiple diffusion scales.<n>Our G$2$RPO significantly outperforms existing flow-based GRPO baselines.
  arXiv  Detail & Related papers  (2025-10-02T12:57:12Z)
• Stochastic Differential Equations models for Least-Squares Stochastic Gradient Descent [6.3151583550712065]
  We study the dynamics of a continuous-time model of the Gradient Descent (SGD) We analyze degenerate Differential Equations (squareSDEs) that model SGD either in the case of the training loss (finite samples) or the population one (online setting)
  arXiv  Detail & Related papers  (2024-07-02T14:52:21Z)
• A Mean-Field Analysis of Neural Stochastic Gradient Descent-Ascent for Functional Minimax Optimization [90.87444114491116]
  This paper studies minimax optimization problems defined over infinite-dimensional function classes of overparametricized two-layer neural networks. We address (i) the convergence of the gradient descent-ascent algorithm and (ii) the representation learning of the neural networks. Results show that the feature representation induced by the neural networks is allowed to deviate from the initial one by the magnitude of $O(alpha-1)$, measured in terms of the Wasserstein distance.
  arXiv  Detail & Related papers  (2024-04-18T16:46:08Z)
• Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index models [10.781866671930857]
  We analyze the dynamics of streaming gradient descent (SGD) in the high-dimensional limit. We demonstrate a deterministic equivalent of SGD in the form of a system of ordinary differential equations. In addition to the deterministic equivalent, we introduce an SDE with a simplified diffusion coefficient.
  arXiv  Detail & Related papers  (2023-08-17T13:33:02Z)
• Capturing dynamical correlations using implicit neural representations [85.66456606776552]
  We develop an artificial intelligence framework which combines a neural network trained to mimic simulated data from a model Hamiltonian with automatic differentiation to recover unknown parameters from experimental data. In doing so, we illustrate the ability to build and train a differentiable model only once, which then can be applied in real-time to multi-dimensional scattering data.
  arXiv  Detail & Related papers  (2023-04-08T07:55:36Z)
• NAG-GS: Semi-Implicit, Accelerated and Robust Stochastic Optimizer [45.47667026025716]
  We propose a novel, robust and accelerated iteration that relies on two key elements. The convergence and stability of the obtained method, referred to as NAG-GS, are first studied extensively. We show that NAG-arity is competitive with state-the-art methods such as momentum SGD with weight decay and AdamW for the training of machine learning models.
  arXiv  Detail & Related papers  (2022-09-29T16:54:53Z)
• The effective noise of Stochastic Gradient Descent [9.645196221785694]
  Gradient Descent (SGD) is the workhorse algorithm of deep learning technology. We characterize the parameters of SGD and a recently-introduced variant, persistent SGD, in a neural network model. We find that noisier algorithms lead to wider decision boundaries of the corresponding constraint satisfaction problem.
  arXiv  Detail & Related papers  (2021-12-20T20:46:19Z)
• Inverting brain grey matter models with likelihood-free inference: a tool for trustable cytoarchitecture measurements [62.997667081978825]
  characterisation of the brain grey matter cytoarchitecture with quantitative sensitivity to soma density and volume remains an unsolved challenge in dMRI. We propose a new forward model, specifically a new system of equations, requiring a few relatively sparse b-shells. We then apply modern tools from Bayesian analysis known as likelihood-free inference (LFI) to invert our proposed model.
  arXiv  Detail & Related papers  (2021-11-15T09:08:27Z)
• On the Double Descent of Random Features Models Trained with SGD [78.0918823643911]
  We study properties of random features (RF) regression in high dimensions optimized by gradient descent (SGD) We derive precise non-asymptotic error bounds of RF regression under both constant and adaptive step-size SGD setting. We observe the double descent phenomenon both theoretically and empirically.
  arXiv  Detail & Related papers  (2021-10-13T17:47:39Z)
• Dynamical mean-field theory for stochastic gradient descent in Gaussian mixture classification [25.898873960635534]
  We analyze in a closed learning dynamics of gradient descent (SGD) for a single-layer neural network classifying a high-dimensional landscape. We define a prototype process for which can be extended to a continuous-dimensional gradient flow. In the full-batch limit, we recover the standard gradient flow.
  arXiv  Detail & Related papers  (2020-06-10T22:49:41Z)
• Kernel and Rich Regimes in Overparametrized Models [69.40899443842443]
  We show that gradient descent on overparametrized multilayer networks can induce rich implicit biases that are not RKHS norms. We also demonstrate this transition empirically for more complex matrix factorization models and multilayer non-linear networks.
  arXiv  Detail & Related papers  (2020-02-20T15:43:02Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.