Tensor Programs IIb: Architectural Universality of Neural Tangent Kernel Training Dynamics

Greg Yang,Etai Littwin
DOI: https://doi.org/10.48550/arXiv.2105.03703
2021-05-08
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.
Machine Learning,Neural and Evolutionary Computing,Probability
What problem does this paper attempt to address?