Tensor Programs IIb: Architectural Universality of Neural Tangent Kernel Training Dynamics
arXiv:2105.03703
Abstract
Yang (2020a) recently showed that the Neural Tangent Kernel (NTK) at initialization has an infinite-width limit for a large class of architectures including modern staples such as ResNet and Transformers. However, their analysis does not apply to training. Here, we show the same neural networks (in the so-called NTK parametrization) during training follow a kernel gradient descent dynamics in function space, where the kernel is the infinite-width NTK. This completes the proof of the *architectural universality* of NTK behavior. To achieve this result, we apply the Tensor Programs technique: Write the entire SGD dynamics inside a Tensor Program and analyze it via the Master Theorem. To facilitate this proof, we develop a graphical notation for Tensor Programs.
arXiv commentICML 2021
References in corpus (14)
- 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
- On Exact Computation with an Infinitely Wide Neural Net
- 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
- Tensor Programs II: Neural Tangent Kernel for Any Architecture
- Steps Toward Deep Kernel Methods from Infinite Neural Networks
- Tensor Programs I: Wide Feedforward or Recurrent Neural Networks of Any Architecture are Gaussian Processes
- Infinite attention: NNGP and NTK for deep attention networks
- On Random Kernels of Residual Architectures