Tensor Programs II: Neural Tangent Kernel for Any Architecture

By Greg Yang. 2020

176 citations recorded by the source corpus.

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 this https URL.

Subfields: Artificial Intelligence, Computational Mathematics, Computer Vision and Pattern Recognition, Optimization & Training, Scaling & Learning Theory

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. J. Edward Hu
2. Greg Yang
3. Sam Schoenholz
4. Soufiane Hayou
5. Jeffrey Pennington
6. Arnaud Doucet
7. Yuanzhi Li
8. Judith Rousseau

In Computational Mathematics

1. J. Edward Hu

In Computer Vision and Pattern Recognition

1. J. Edward Hu
2. Greg Yang
3. Sam Schoenholz
4. Jeffrey Pennington
5. Lechao Xiao
6. Yuanzhi Li
7. Etai Littwin
8. Shotaro Akaho

In Optimization & Training

1. J. Edward Hu
2. Greg Yang

In Scaling & Learning Theory

1. J. Edward Hu
2. Greg Yang
3. Sam Schoenholz
4. Jeffrey Pennington
5. Soufiane Hayou
6. Lechao Xiao
7. Arnaud Doucet
8. Judith Rousseau

Citations in the corpus

• Feature Learning in Infinite-Width Neural Networks (2020, methodology): “General Neural Architectures More generally, µP can be defined easily for any neural architecture whose forward pass can be written down as a Tensor Program (e.g. ResNet or Transformer; see (Yang, 2019a) for explicit programs).”
• Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks (2020, unknown): “Our results are directly applicable to in fi nite-width limit of neural networks that admit a kernel description (including feedforward, convolutional and recurrent neural networks) 13,55 – 59 , and explain their inductive bias towards simple functions 35,51,60 – 64 .”
• Tensor Programs III: Neural Matrix Laws (2020, methodology): “As with the NNGP in Appendix B, this infinite-width NTK can be computed in a straightforward way using Theorem 2.10, following the examples of Yang [51].”
• LLM Web Dynamics: Tracing Model Collapse in a Network of LLMs (2025, unknown): “The discrepancy arises because a B -component GMM is restricted to a finite parametrization, while a Transformer-based LLM occupies an effectively infinite-dimensional func-7 tion space (Pérez et al., 2019; Yang, 2020).”