Tensor Programs IIb: Architectural Universality of Neural Tangent Kernel Training Dynamics
This foundational result extends NTK theory to training dynamics, impacting all of ML/AI by providing a universal framework for analyzing neural network optimization.
The paper proves that neural networks in the NTK parametrization follow kernel gradient descent dynamics with the infinite-width NTK during training, establishing architectural universality for NTK behavior across architectures like ResNet and Transformers.
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.