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

Greg Yang|arXiv (Cornell University)|Feb 13, 2019

Gaussian Processes and Bayesian Inference78 references187 citations

TL;DR

The paper presents a unifying tensor-program framework to derive scaling limits of wide neural networks, establishing Gaussian process behavior, gradient independence conditions, and Neural Tangent Kernel convergence across standard architectures without batch norm.

ABSTRACT

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.

Motivation & Objective

• Define a unifying tensor program framework that expresses most neural network computations under weight sharing.
• Characterize the scaling limit of these programs as widths go to infinity under Glorot-style initialization.
• Derive Gaussian Process behavior for broad architectures (RNNs, CNNs, ResNets, attention, etc.).
• Analyze when the gradient independence assumption yields correct gradient dynamics and provide corrections when not.
• Prove convergence of the Neural Tangent Kernel at initialization for architectures without batch normalization.

Proposed method

• Introduce tensor programs with G-, A-, and H-vars to encode neural computations.
• Define common dimension classes (CDCs) and sampling schemes for weights and inputs.
• Prove that under wide limits, G-vars converge to Gaussians with computable means and covariances (Theorems 4.3, 5.1, 6.3).
• Derive DNN-GP correspondence for standard architectures under broad nonlinearities (Corollary 2.1).
• Derive (informal) gradient independence validity (Corollary 2.3) and corrects where necessary.
• Establish Neural Tangent Kernel convergence Kθ → K∞ for finite input sets without batch norm (Corollary 2.4).

Experimental results

Research questions

• RQ1 Under what conditions do wide neural networks with weight sharing converge to Gaussian processes across common architectures?
• RQ2 When is the gradient independence assumption valid for backpropagation, and how can the correct gradient dynamics be computed if it fails?
• RQ3 How does the Neural Tangent Kernel behave at initialization across standard architectures without batch normalization, and when does it converge to a limit kernel K∞?
• RQ4 Can the framework recover classical random-matrix results (e.g., semicircle, Marchenko-Pastur) as special cases?
• RQ5 What is the role of weight sharing (transposes) in the scaling limits of various architectures (RNNs, CNNs, residuals, attention)?

Key findings

• The DNN-GP correspondence generalizes to standard architectures and nonlinearities, producing Gaussian process limits as widths grow (Corollary 2.1).
• The gradient independence assumption yields correct backpropagation dynamics for polynomially bounded nonlinearities under certain conditions, with explicit corrections where it fails (Corollary 2.3).
• Neural Tangent Kernel converges to a limit K∞ at initialization for finite input sets across standard architectures without batch norm (Corollary 2.4).
• The tensor-program framework can rederive classical random-matrix results and relate to state-evolution-like analyses in related algorithms (e.g., AMP).
• The work provides a general method to analyze signal propagation and gradient dynamics, enabling design of initialization schemes to avoid gradient explosion/vanishing.