[Paper Review] Neural Tangent Kernel: Convergence and Generalization in Neural Networks

Arthur Paul Jacot, Franck Gabriel+1 more|arXiv (Cornell University)|Jun 20, 2018

Neural Networks and Applications1 references1,508 citations

TL;DR

The paper introduces the Neural Tangent Kernel (NTK), proves its deterministic limiting behavior in infinite-width nets, and shows training dynamics under gradient descent reduce to kernel gradient descent in function space.

ABSTRACT

At initialization, artificial neural networks (ANNs) are equivalent to Gaussian processes in the infinite-width limit, thus connecting them to kernel methods. We prove that the evolution of an ANN during training can also be described by a kernel: during gradient descent on the parameters of an ANN, the network function $f_θ$ (which maps input vectors to output vectors) follows the kernel gradient of the functional cost (which is convex, in contrast to the parameter cost) w.r.t. a new kernel: the Neural Tangent Kernel (NTK). This kernel is central to describe the generalization features of ANNs. While the NTK is random at initialization and varies during training, in the infinite-width limit it converges to an explicit limiting kernel and it stays constant during training. This makes it possible to study the training of ANNs in function space instead of parameter space. Convergence of the training can then be related to the positive-definiteness of the limiting NTK. We prove the positive-definiteness of the limiting NTK when the data is supported on the sphere and the non-linearity is non-polynomial. We then focus on the setting of least-squares regression and show that in the infinite-width limit, the network function $f_θ$ follows a linear differential equation during training. The convergence is fastest along the largest kernel principal components of the input data with respect to the NTK, hence suggesting a theoretical motivation for early stopping. Finally we study the NTK numerically, observe its behavior for wide networks, and compare it to the infinite-width limit.

Motivation & Objective

• Understand how training dynamics of fully connected neural networks behave in the infinite-width limit.
• Show that the network function follows kernel gradient descent with respect to a limiting NTK.
• Establish conditions under which the limiting NTK is positive definite and ensures convergence.
• Connect neural network training to kernel methods and explain implications for generalization and early stopping.

Proposed method

• Define fully connected networks with a Lipschitz nonlinearity and random Gaussian initialization.
• Introduce the Neural Tangent Kernel (NTK) and prove its convergence to a deterministic limit as layer widths go to infinity.
• Show that NTK remains constant during training in the infinite-width limit.
• Derive a linear differential equation governing the network function under least-squares loss.
• Analyze convergence via kernel gradient descent and kernel principal components of the data.
• Provide numerical experiments comparing finite-width networks to the infinite-width theory.

Experimental results

Research questions

• RQ1 Does the NTK converge to a deterministic limiting kernel in the infinite-width limit?
• RQ2 Does the NTK stay constant during training, enabling a function-space description of training dynamics?
• RQ3 Under what conditions is the limiting NTK positive definite and thus guarantees convergence?
• RQ4 How does training under least-squares loss relate to kernel principal components and early stopping?
• RQ5 How well do finite-width networks approximate the infinite-width NTK theory in practice?

Key findings

• The network function at initialization converges to a Gaussian process, and the NTK converges in probability to a deterministic limit as widths go to infinity.
• During training, the NTK stays asymptotically constant, allowing gradient descent to be described as kernel gradient descent in function space.
• For non-polynomial Lipschitz nonlinearity and data on the sphere, the limiting NTK is positive definite for depth at least two, ensuring convergence.
• Under least-squares regression, the training dynamics reduce to a linear differential equation governed by the NTK, with convergence fastest along the largest kernel principal components.
• Numerical experiments show wide networks approach the infinite-width limit and that the NTK remains close to its limit during training.