LIT-753: Rigorous dynamical mean field theory for stochastic gradient descent methods

Gerbelot et al. (2022) — ARXIV-2210.06591

Key takeaways

• What is proved. Closed-form equations for the exact high-dimensional asymptotics of a generic family of discrete-time first-order iterations, learning an estimator (an M-estimator, a shallow network with a finite number q of hidden units) by empirical risk minimization on Gaussian data. Samples n and dimension d go to infinity together; q stays finite. The equations coincide with a discretization of the dynamical mean-field theory (DMFT) physicists derive for gradient flow.
• What the family covers. Multi-pass SGD with mini-batches modelled as independent Bernoulli selections — so the batch must be a finite fraction of the dataset, not a constant number of examples; heavy-ball and Nesterov momentum; time-dependent updates including learning-rate schedules and thermal (Langevin) noise; any differentiable regularizer; and, through non-separable update functions, data with a generic well-conditioned covariance rather than identity.
• What is new against the prior rigorous result (Celentano et al., reference 11 in the paper): stochasticity of the algorithm and non-identity data covariance, by a different proof — iterative Gaussian conditioning rather than a mapping onto approximate message passing — which also shows explicitly how the memory kernels of the effective dynamics build up. The paper treats discrete time only, arguing that is what algorithms run.
• Numerics: for a teacher-student perceptron with ridge regularization, the cosine similarity to the teacher (which fixes the generalization error in that model) from the solved equations agrees with SGD simulations at d = 1000, across learning rates and batch fractions.
• Limits it states: the solver is stable only to relatively short times, a long-standing problem for DMFT; the data covariance cannot change during training, so a learned feature map is outside it; and extracting simple governing quantities from the equations is left open.

Standing in the anthology

A different exact limit from the ones the record holds. THEORY-009, from LIT-271 and its neighbours, is also a “mean-field” account of SGD, but of another limit: width goes to infinity and the object that moves is the distribution of neurons. Here width stays finite and it is the number of samples and the dimension that grow. LIT-454 derives SGD’s limiting dynamics as a Langevin equation, with an exact Ornstein–Uhlenbeck solution for linear regression. Neither connects to this paper’s equations, and the record has no document that uses DMFT.

It does not bear on the optimizer cluster. SOTA-121, SOTA-165 and the theory around them concern matrix preconditioners on deep networks at scale. This paper covers first-order iterations with a scalar step size and momentum on Gaussian data and shallow models; whether its non-separable update functions could express a matrix preconditioner is not something it discusses. Its relation to the record’s batch-size material (SOTA-198, LIT-058) is that batch size enters as a fraction of the dataset, which is not the regime those practices are about.

Filed as a seed: the rigorous end of SGD-dynamics theory, for when the record has a claim about high-dimensional training dynamics that it could test.

Unread — no NOTE.