Quantitative Propagation of Chaos for SGD in Wide Neural Networks

Abstract

In this paper, we investigate the limiting behavior of a continuous-time counterpart of the Stochastic Gradient Descent (SGD) algorithm applied to two-layer overparameterized neural networks, as the number or neurons (ie, the size of the hidden layer) $N\to +\infty$. Following a probabilistic approach, we show 'propagation of chaos' for the particle system defined by this continuous-time dynamics under different scenarios, indicating that the statistical interaction between the particles asymptotically vanishes. In particular, we establish quantitative convergence with respect to $N$ of any particle to a solution of a mean-field McKean-Vlasov equation in the metric space endowed with the Wasserstein distance. In comparison to previous works on the subject, we consider settings in which the sequence of stepsizes in SGD can potentially depend on the number of neurons and the iterations. We then identify two regimes under which different mean-field limits are obtained, one of them corresponding to an implicitly regularized version of the minimization problem at hand. We perform various experiments on real datasets to validate our theoretical results, assessing the existence of these two regimes on classification problems and illustrating our convergence results.

References in corpus (13)

• The Loss Surfaces of Multilayer Networks
• Visualizing the Loss Landscape of Neural Nets
• A Mean Field View of the Landscape of Two-Layers Neural Networks
• Breaking the Curse of Dimensionality with Convex Neural Networks
• The statistics of critical points of Gaussian fields on large-dimensional spaces
• On the Convergence Rate of Training Recurrent Neural Networks
• Perspective: Energy Landscapes for Machine Learning
• On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport
• Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit
• On the saddle point problem for non-convex optimization
• Explaining Landscape Connectivity of Low-cost Solutions for Multilayer Nets
• Mean-Field Neural ODEs via Relaxed Optimal Control
• Markov Transitions and the Propagation of Chaos

Cited by in corpus (8)

• Propagation of chaos: a review of models, methods and applications. II. Applications
• The Heavy-Tail Phenomenon in SGD
• Heavy Tails in SGD and Compressibility of Overparametrized Neural Networks
• Law of large numbers and central limit theorem for wide two-layer neural networks: the mini-batch and noisy case
• Wide stochastic networks: Gaussian limit and PAC-Bayesian training
• Implicit Compressibility of Overparametrized Neural Networks Trained with Heavy-Tailed SGD
• Painless step size adaptation for SGD
• Law of Large Numbers for Bayesian two-layer Neural Network trained with Variational Inference