Tensor Programs II: Neural Tangent Kernel for Any Architecture

Abstract

We prove that a randomly initialized neural network of *any architecture* has its Tangent Kernel (NTK) converge to a deterministic limit, as the network widths tend to infinity. We demonstrate how to calculate this limit. In prior literature, the heuristic study of neural network gradients often assumes every weight matrix used in forward propagation is independent from its transpose used in backpropagation (Schoenholz et al. 2017). This is known as the *gradient independence assumption (GIA)*. We identify a commonly satisfied condition, which we call *Simple GIA Check*, such that the NTK limit calculation based on GIA is correct. Conversely, when Simple GIA Check fails, we show GIA can result in wrong answers. Our material here presents the NTK results of Yang (2019a) in a friendly manner and showcases the *tensor programs* technique for understanding wide neural networks. We provide reference implementations of infinite-width NTKs of recurrent neural network, transformer, and batch normalization at https://github.com/thegregyang/NTK4A.

arXiv comment11 pages of main text. 60 pages total. August 2020: Fixed "BP-like" definition

References in corpus (30)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks
• Semi-Supervised Classification with Graph Convolutional Networks
• Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering
• Deep Convolutional Networks on Graph-Structured Data
• Neural Tangent Kernel: Convergence and Generalization in Neural Networks
• A Convergence Theory for Deep Learning via Over-Parameterization
• Deep Neural Networks as Gaussian Processes
• On Exact Computation with an Infinitely Wide Neural Net
• Gradient Descent Provably Optimizes Over-parameterized Neural Networks
• Exponential expressivity in deep neural networks through transient chaos
• Stochastic Gradient Descent Optimizes Over-parameterized Deep ReLU Networks
• Toward Deeper Understanding of Neural Networks: The Power of Initialization and a Dual View on Expressivity
• Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation
• Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes
• Graph Neural Tangent Kernel: Fusing Graph Neural Networks with Graph Kernels
• The Convergence Rate of Neural Networks for Learned Functions of Different Frequencies
• On the Convergence Rate of Training Recurrent Neural Networks
• Dynamical Isometry and a Mean Field Theory of RNNs: Gating Enables Signal Propagation in Recurrent Neural Networks
• Towards Characterizing Divergence in Deep Q-Learning
• A Mean Field Theory of Batch Normalization
• Steps Toward Deep Kernel Methods from Infinite Neural Networks
• Asymptotics of Wide Networks from Feynman Diagrams
• On the Selection of Initialization and Activation Function for Deep Neural Networks
• Dynamics of Deep Neural Networks and Neural Tangent Hierarchy
• Infinite attention: NNGP and NTK for deep attention networks
• Tensor Programs III: Neural Matrix Laws
• Neural Temporal-Difference and Q-Learning Provably Converge to Global Optima
• On the Neural Tangent Kernel of Deep Networks with Orthogonal Initialization
• The Nonlinearity Coefficient - Predicting Generalization in Deep Neural Networks

Cited by in corpus (24)

• Neural Networks and Quantum Field Theory
• Nonperturbative renormalization for the neural network-QFT correspondence
• Bayesian Deep Ensembles via the Neural Tangent Kernel
• Feature Learning in Infinite-Width Neural Networks
• Tensor Programs III: Neural Matrix Laws
• Inductive Biases and Variable Creation in Self-Attention Mechanisms
• Tensor Programs IIb: Architectural Universality of Neural Tangent Kernel Training Dynamics
• On the Neural Tangent Kernel of Deep Networks with Orthogonal Initialization
• A quantum algorithm for training wide and deep classical neural networks
• Metric Flows with Neural Networks
• Infinite Neural Network Quantum States: Entanglement and Training Dynamics
• Towards Deepening Graph Neural Networks: A GNTK-based Optimization Perspective
• Shift Invariance Can Reduce Adversarial Robustness
• The Recurrent Neural Tangent Kernel
• The Limitations of Large Width in Neural Networks: A Deep Gaussian Process Perspective
• The Nonlinearity Coefficient - A Practical Guide to Neural Architecture Design
• Implicit bias of deep linear networks in the large learning rate phase
• Kernel-Based Smoothness Analysis of Residual Networks
• The Future is Log-Gaussian: ResNets and Their Infinite-Depth-and-Width Limit at Initialization
• Analyzing Finite Neural Networks: Can We Trust Neural Tangent Kernel Theory?
• Quantifying Epistemic Uncertainty in Deep Learning
• Implicit Acceleration and Feature Learning in Infinitely Wide Neural Networks with Bottlenecks
• Enhanced Recurrent Neural Tangent Kernels for Non-Time-Series Data
• Understanding Approximate Fisher Information for Fast Convergence of Natural Gradient Descent in Wide Neural Networks