[Paper Review] Tensor Programs II: Neural Tangent Kernel for Any Architecture

Greg Yang|arXiv (Cornell University)|Jun 25, 2020

Stochastic Gradient Optimization Techniques59 references53 citations

TL;DR

This work proves that the neural tangent kernel (NTK) converges deterministically for any neural network architecture in the infinite-width limit and clarifies when gradient independence assumptions are valid, with practical tensor-program-based methods and reference implementations.

ABSTRACT

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 https://github.com/thegregyang/NTK4A.

Motivation & Objective

• Demonstrate that randomly initialized networks of any architecture have a deterministic NTK limit as widths grow without bound.
• Identify conditions under which gradient independence assumptions yield correct NTK limits (Simple GIA Check).
• Show failures of gradient independence when the Simple GIA Check does not hold and provide correct NTK computation guidance.
• Present a friendly exposition of NTK limits via tensor programs for wide networks.
• Provide reference implementations of infinite-width NTKs for recurrent networks, transformers, and batch normalization.

Proposed method

• Apply the tensor programs framework to analyze wide neural networks and derive NTK limits.
• Formalize the gradient independence assumption (GIA) and establish the Simple GIA Check as a condition for correct NTK calculations.
• Prove convergence of the NTK to a deterministic limit for arbitrary architectures as widths tend to infinity.
• Derive and present NTK limits for recurrent neural networks, transformers, and batch normalization.
• Offer practical guidance and reference implementations for computing infinite-width NTKs.

Experimental results

Research questions

• RQ1 Does the NTK converge to a deterministic limit for any neural network architecture in the infinite-width limit?
• RQ2 Under what conditions is the gradient independence assumption valid for NTK calculations (Simple GIA Check)?
• RQ3 How can one compute the NTK limit for specific architectures such as RNNs, transformers, and batch normalization?
• RQ4 What are the failure modes of GIA, and how do they affect NTK results?
• RQ5 Can tensor-programs provide a friendly, actionable framework for understanding NTK limits in wide networks?

Key findings

• The NTK converges to a deterministic limit as network widths tend to infinity for any architecture.
• A Simple GIA Check identifies a condition under which GIA-based NTK calculations are correct.
• When the Simple GIA Check fails, GIA can yield incorrect NTK results.
• The paper presents a friendly exposition of NTK results using tensor programs.
• Reference implementations are provided for infinite-width NTKs of RNNs, transformers, and batch normalization.