Rigorous dynamical mean field theory for stochastic gradient descent methods
arXiv:2210.06591 · doi:10.1137/23M1594388
Abstract
We prove closed-form equations for the exact high-dimensional asymptotics of a family of first order gradient-based methods, learning an estimator (e.g. M-estimator, shallow neural network, ...) from observations on Gaussian data with empirical risk minimization. This includes widely used algorithms such as stochastic gradient descent (SGD) or Nesterov acceleration. The obtained equations match those resulting from the discretization of dynamical mean-field theory (DMFT) equations from statistical physics when applied to gradient flow. Our proof method allows us to give an explicit description of how memory kernels build up in the effective dynamics, and to include non-separable update functions, allowing datasets with non-identity covariance matrices. Finally, we provide numerical implementations of the equations for SGD with generic extensive batch-size and with constant learning rates.
arXiv comment40 pages, 4 figures
References in corpus (12)
- Analytical Solution of the Off-Equilibrium Dynamics of a Long Range Spin-Glass Model
- The dynamics of message passing on dense graphs, with applications to compressed sensing
- Universality in polytope phase transitions and message passing algorithms
- Solution of the dynamics of liquids in the large-dimensional limit
- Numerical implementation of dynamical mean field theory for disordered systems: application to the Lotka-Volterra model of ecosystems
- Out-of-equilibrium dynamical mean-field equations for the perceptron model
- A simple theory for the dynamics of mean-field-like models of glass-forming fluids
- The effective noise of Stochastic Gradient Descent
- Stochasticity helps to navigate rough landscapes: comparing gradient-descent-based algorithms in the phase retrieval problem
- Numerical solution of the dynamical mean field theory of infinite-dimensional equilibrium liquids
- On the Dynamics of Liquids in the Large-Dimensional Limit
- High dimensional optimization under non-convex excluded volume constraints