A mean field view of the landscape of two-layer neural networks

By Song Mei, Andrea Montanari, Phan-Minh Nguyen. 2018

1,021 citations recorded by the source corpus.

Significance Multilayer neural networks have proven extremely successful in a variety of tasks, from image classification to robotics. However, the reasons for this practical success and its precise domain of applicability are unknown. Learning a neural network from data requires solving a complex optimization problem with millions of variables. This is done by stochastic gradient descent (SGD) algorithms. We study the case of two-layer networks and derive a compact description of the SGD dynamics in terms of a limiting partial differential equation. Among other consequences, this shows that SGD dynamics does not become more complex when the network size increases. Multilayer neural networks are among the most powerful models in machine learning, yet the fundamental reasons for this success defy mathematical understanding. Learning a neural network requires optimizing a nonconvex high-dimensional objective (risk function), a problem that is usually attacked using stochastic gradient descent (SGD). Does SGD converge to a global optimum of the risk or only to a local optimum? In the former case, does this happen because local minima are absent or because SGD somehow avoids them? In the latter, why do local minima reached by SGD have good generalization properties? In this paper, we consider a simple case, namely two-layer neural networks, and prove that—in a suitable scaling limit—SGD dynamics is captured by a certain nonlinear partial differential equation (PDE) that we call distributional dynamics (DD). We then consider several specific examples and show how DD can be used to prove convergence of SGD to networks with nearly ideal generalization error. This description allows for “averaging out” some of the complexities of the landscape of neural networks and can be used to prove a general convergence result for noisy SGD.

Subfields: Artificial Intelligence, Statistical and Nonlinear Physics, Statistics and Probability, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Statistical and Nonlinear Physics

In Statistics and Probability

In Model Architectures

Citations in the corpus

Wide neural networks of any depth evolve as linear models under gradient descent (2019, background): “Mei et al. (2018); Chizat & Bach (2018a); Rotskoff & Vanden-Eijnden (2018); Sirignano & Spiliopoulos (2018) analyze the mean ﬁeld SGD dynamics of training neural networks in the large-width limit.”
Scaling description of generalization with number of parameters in deep learning (2019, neutral): “This result explains why the generalization performance converges as N → ∞, a result previously obtained for single hidden layer neural networks using a different approach [32, 33, 34, 35].”
Neural Tangent Kernel: Convergence and Generalization in Neural Networks (2018, background): “of training in the large-width limit for shallow networks has seen recent progress as well [11].”
The large learning rate phase of deep learning: the catapult mechanism (2020, background): “…limit of infinite width (Daniely, 2017; Jacot et al., 2018; Lee et al., 2019; Du et al., 2019; Zou et al., 2018; Allen-Zhu et al., 2019; Li & Liang, 2018; Chizat et al., 2019; Mei et al., 2018; Rotskoff & Vanden-Eijnden, 2018; Sirignano & Spiliopoulos, 2018; Woodworth et al., 2019; Naveh et al.).”
Feature Learning in Infinite-Width Neural Networks (2020, methodology): “In summary, For all θ ∈ R, ft(ξ) stays fixed for all t and ξ if we set al ← al + θ, bl ← bl − θ, c← c− 2θ. (13) This insight in particular implies MFP is a special case of MUP in the case of 1-hidden-layer MLPs.”
When do neural networks outperform kernel methods? (2020, background): “Further even in the infinite width limit, it is known that two-layers neural networks can actually capture a richer class of functions than the associated RKHS, provided SGD training is scaled differently from the lazy regime [23, 7, 28, 30, 8].”
Generalization Guarantees for Neural Networks via Harnessing the Low-rank Structure of the Jacobian (2019, methodology): “[18, 54, 42, 53, 49] relate the distribution of the network weights to Wasserstein gradient flows using mean field analysis.”
Continual Learning in the Teacher-Student Setup: Impact of Task Similarity (2021, background): “The complementary “mean-field” limit of finite input dimension and an infinite number of hidden neurons was analysed (Mei et al., 2018; Chizat & Bach, 2018; Sirignano & Spiliopoulos, 2020; Rotskoff & Vanden-Eijnden, 2018).”
Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks (2019, methodology): “…ﬁeld analysis to establish that for inﬁnitely wide neural networks, the empirical distribution of the neural network parameters can be described as a Wasserstein gradient ﬂow (Mei et al., 2018; Chizat & Bach, 2018a; Sirignano & Spiliopoulos, 2018; Rotskoﬀ & Vanden-Eijnden, 2018; Wei et al., 2018).”
On Lazy Training in Differentiable Programming (2018, background): “In contrast, the critical scaling α(m) = 1/m, allows to converge as m→∞ to a non degenerate dynamic described by a partial differential equation and referred to as the mean-field limit [24, 7, 29, 33].”
On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport (2018, background): “Several independent works [24, 28, 32] have studied the many-particle limit of training a neural network with a single large hidden layer and a quadratic loss R.”
On the Convergence Rate of Training Recurrent Neural Networks (2018, methodology): “There is huge literature on using the mean-field theory to study neural networks [11, 15, 36, 41, 45–47, 49, 54, 66–68].”
Mean Field Analysis of Neural Networks: A Law of Large Numbers (2018, background): “In [38], it is assumed that the gradient of the neural network is a priori globally Lipschitz and bounded and under this assumption a similar PDE result as well as certain rates of convergence are established.”
Kernel and Deep Regimes in Overparametrized Models (2019, neutral): “[18] Song Mei, Andrea Montanari, and Phan-Minh Nguyen.”