High‐dimensional limit theorems for SGD: Effective dynamics and critical scaling
-
Published:2023-10-04
Issue:3
Volume:77
Page:2030-2080
-
ISSN:0010-3640
-
Container-title:Communications on Pure and Applied Mathematics
-
language:en
-
Short-container-title:Comm Pure Appl Math
Author:
Arous Gérard Ben1,
Gheissari Reza2,
Jagannath Aukosh34
Affiliation:
1. Courant Institute New York University New York USA
2. Department of Mathematics Northwestern University Evanston USA
3. Department of Statistics and Actuarial Science, Cheriton School of Computer Science University of Waterloo Waterloo Canada
4. Department of Applied Mathematics Cheriton School of Computer Science, University of Waterloo Waterloo Canada
Abstract
AbstractWe study the scaling limits of stochastic gradient descent (SGD) with constant step‐size in the high‐dimensional regime. We prove limit theorems for the trajectories of summary statistics (i.e., finite‐dimensional functions) of SGD as the dimension goes to infinity. Our approach allows one to choose the summary statistics that are tracked, the initialization, and the step‐size. It yields both ballistic (ODE) and diffusive (SDE) limits, with the limit depending dramatically on the former choices. We show a critical scaling regime for the step‐size, below which the effective ballistic dynamics matches gradient flow for the population loss, but at which, a new correction term appears which changes the phase diagram. About the fixed points of this effective dynamics, the corresponding diffusive limits can be quite complex and even degenerate. We demonstrate our approach on popular examples including estimation for spiked matrix and tensor models and classification via two‐layer networks for binary and XOR‐type Gaussian mixture models. These examples exhibit surprising phenomena including multimodal timescales to convergence as well as convergence to sub‐optimal solutions with probability bounded away from zero from random (e.g., Gaussian) initializations. At the same time, we demonstrate the benefit of overparametrization by showing that the latter probability goes to zero as the second layer width grows.
Funder
National Science Foundation
Adolph C. and Mary Sprague Miller Institute for Basic Research in Science, University of California Berkeley
Natural Sciences and Engineering Research Council of Canada
Canada Research Chairs
Subject
Applied Mathematics,General Mathematics
Reference78 articles.
1. K.Ahn S.Bubeck S.Chewi Y. T.Lee F.Suarez andY.Zhang Learning threshold neurons via the “edge of stability” 2022.
2. A.Anastasiou K.Balasubramanian andM. A.Erdogdu Normal approximation for stochastic gradient descent via non‐asymptotic rates of martingale CLT Conference on Learning Theory PMLR 2019pp.115–137.
3. D.Araújo R. I.Oliveira andD.Yukimura A mean‐field limit for certain deep neural networks arXiv preprint arXiv:1906.00193 (2019).
4. S.Arora Z.Li andA.Panigrahi Understanding gradient descent on the edge of stability in deep learning K.Chaudhuri S.Jegelka L.Song C.Szepesvari G.Niu andS.Sabato (eds.) Proceedings of the 39th International Conference on Machine Learning Proceedings of Machine Learning Research vol.162 PMLR 2022 pp.948–1024.https://proceedings.mlr.press/v162/arora22a.html.
Cited by
1 articles.
订阅此论文施引文献
订阅此论文施引文献,注册后可以免费订阅5篇论文的施引文献,订阅后可以查看论文全部施引文献
1. High-dimensional limit of one-pass SGD on least squares;Electronic Communications in Probability;2024-01-01