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

2021/05/08 by Greg Yang, Etai Littwin, Yang, Greg +1 · 8 citations
Computer Science · Mathematics · #Advanced Neural Network Applications #FOS: Computer and information sciences #FOS: Mathematics #Machine Learning (cs.LG) #Neural and Evolutionary Computing (cs.NE) #Probability (math.PR) #Stochastic Gradient Optimization Techniques #Tensor decomposition and applications #cs.LG #cs.NE #math.PR

paper · pdf · doi:10.48550/arxiv.2105.03703

ICML 2021

arxiv created 2021/05/08 · openalex publication_date 2021/05/08 · arxiv updated 2021/05/11 · openalex created_date 2021/05/24 · openalex updated_date 2026/07/28

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.

Citations

• Deep Convolutional Networks on Graph-Structured Data
• Gradient-based learning applied to document recognition
• A Convergence Theory for Deep Learning via Over-Parameterization
• Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes
• Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation

Cited by

• Data Value in the Age of Scaling: Understanding LLM Scaling Dynamics Under Real-Synthetic Data Mixtures
• When Bias Pretends to Be Truth: How Spurious Correlations Undermine Hallucination Detection in LLMs
• Quantitative convergence of trained single layer neural networks to Gaussian processes
• Theoretical Foundations of Representation Learning using Unlabeled Data: Statistics and Optimization
• Evolution of Concepts in Language Model Pre-Training
• Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs
• Querying Kernel Methods Suffices for Reconstructing their Training Data
• Scaling Diffusion Transformers Efficiently via μP

Related

• Tensor Programs II: Neural Tangent Kernel for Any Architecture
• Tensor Programs IVb: Adaptive Optimization in the Infinite-Width Limit
• Neural Tangent Kernel: A Survey
• Mathematical Foundations of Neural Tangents and Infinite-Width Networks
• Neural Tangent Kernel: Convergence and Generalization in Neural Networks