Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation

By Greg Yang. 2019

313 citations recorded by the source corpus.

Several recent trends in machine learning theory and practice, from the design of state-of-the-art Gaussian Process to the convergence analysis of deep neural nets (DNNs) under stochastic gradient descent (SGD), have found it fruitful to study wide random neural networks. Central to these approaches are certain scaling limits of such networks. We unify these results by introducing a notion of a straightline \emph{tensor program} that can express most neural network computations, and we characterize its scaling limit when its tensors are large and randomized. From our framework follows (1) the convergence of random neural networks to Gaussian processes for architectures such as recurrent neural networks, convolutional neural networks, residual networks, attention, and any combination thereof, with or without batch normalization; (2) conditions under which the \emph{gradient independence assumption} -- that weights in backpropagation can be assumed to be independent from weights in the forward pass -- leads to correct computation of gradient dynamics, and corrections when it does not; (3) the convergence of the Neural Tangent Kernel, a recently proposed kernel used to predict training dynamics of neural networks under gradient descent, at initialization for all architectures in (1) without batch normalization. Mathematically, our framework is general enough to rederive classical random matrix results such as the semicircle and the Marchenko-Pastur laws, as well as recent results in neural network Jacobian singular values. We hope our work opens a way toward design of even stronger Gaussian Processes, initialization schemes to avoid gradient explosion/vanishing, and deeper understanding of SGD dynamics in modern architectures.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Scaling & Learning Theory

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 Computer Vision and Pattern Recognition

In Scaling & Learning Theory

Citations in the corpus

Wide neural networks of any depth evolve as linear models under gradient descent (2019, background): “Note that the convergence of the empirical NTK in probability is proved rigorously in [37].”
Feature Learning in Infinite-Width Neural Networks (2020, methodology): “5556 55It is completely OK to define an initial scalar using randomness from other parts of the program, as long as this scalar converges almost surely to a deterministic limit 56We cannot define it using a Moment instruction because, intuitively, the mechanism of this convergence is through CLT, no”
Feature Purification: How Adversarial Training Performs Robust Deep Learning (2020, neutral): “[4, 7, 8, 13, 14, 26–29, 38, 45, 50, 57, 61, 104, 115]).”
Learning curves for continual learning in neural networks: Self-knowledge transfer and forgetting (2021, methodology): “Then, the model update on the the class 1 at the n -th task, i.e., f ) in the inﬁnite-width limit (Jacot et al., 2018; Yang, 2019).”
On Exact Computation with an Infinitely Wide Neural Net (2019, background): “Theorem 3.1 improves upon previous results [Jacot et al., 2018, Yang, 2019] that establish convergence to the NTK at initialization in the following sense: 1.”
Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers (2018, background): “” There is also a line of work studying the relationship between neural networks and NTKs (neural tangent kernels) [3, 4, 7, 27, 31, 51].”
Neural Tangents: Fast and Easy Infinite Neural Networks in Python (2019, background): “The results for individual architectures have subsequently been generalized, and it was shown that a GP correspondence holds for a general class of networks that can be mapped to so-called tensor programs in (Yang, 2019).”
A Mean Field Theory of Batch Normalization (2019, methodology): “8 α-RELU Recall that α-ReLUs (Yang & Schoenholz (2017)) are, roughly speaking, the αth power of ReLU.”
Enhanced Convolutional Neural Tangent Kernels (2019, background): “Recently, NTK s induced by various neural network architectures are derived and shown to achieve strong empirical performance [Arora et al., 2019, Yang, 2019, Du et al., 2019a].”
Finite Depth and Width Corrections to the Neural Tangent Kernel (2019, background): “A number of articles [3, 8, 15, 22] have followed up on the original NTK work [14].”
Disentangling trainability and generalization in deep learning (2019, methodology): “…index i is used because as the number of channels approach inﬁnity, this kernel will also converge in distribution to a deterministic kernel Θ ( l +1) (Yang, 2019), which can also be computed recursively in a similar manner to the NTK for fully-connected networks as (Yang, 2019; Arora et al.,…”
Disentangling feature and lazy training in deep neural networks (2019, background): “…are initialized as i.i.d. random variables with zero mean and a ﬁxed variance of order h − 1 / 2 , the output function f ( x ) is a Gaussian random processes with some covariance that can be computed (Neal 1996, Williams 1997, Lee et al 2018, Matthews et al 2018, Novak et al 2019, Yang 2019).”
Kernel and Rich Regimes in Overparametrized Models (2019, background): “Of particular interest is the tangent kernel at initialization, Kw(0) [15, 22, 23].”
Disentangling feature and lazy learning in deep neural networks: an empirical study (2019, background): “The output function f ( w, x ) is a Gaussian random processes with some covariance that can be computed [Neal, 1996, Williams, 1997, Lee et al., 2018, de G. Matthews et al., 2018, Novak et al., 2019, Yang, 2019].”
Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks (2019, background): “Their result was later improved by Yang [36] such that the widths of different layers can go to infinity simultaneously.”
Backward Feature Correction: How Deep Learning Performs Deep Learning (2020, background): “2 With huge non-convexity obstacles, existing theoretical works, to the best of our knowledge, have only been able to demonstrate that multi-layer neural networks can efficiently perform learning tasks that are already known solvable by non-hierarchical methods, such as kernel methods [3, 6, 7, 9, 1”