Neural tangent kernel

An Artificial Neural Network (ANN) with scalar output consists in a family of functions ${\displaystyle f\left(\cdot ,\theta \right):\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} }$ parametrized by a vector of parameters ${\displaystyle \theta \in \mathbb {R} ^{P}}$.

The Neural Tangent Kernel (NTK) is a kernel ${\displaystyle \Theta$ :\mathbb {R} ^{n_{\mathrm {in} }}\times \mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} } defined by

$${\displaystyle \Theta \left(x,y;\theta \right)=\sum _{p=1}^{P}\partial _{\theta _{p}}f\left(x;\theta \right)\partial _{\theta _{p}}f\left(y;\theta \right).}$$

In the language of kernel methods, the NTK ${\displaystyle \Theta }$ is the kernel associated with the feature map ${\displaystyle \left(x\mapsto \partial _{\theta _{p}}f\left(x;\theta \right)\right)_{p=1,\ldots ,P}}$.

An ANN with vector output of size ${\displaystyle n_{\mathrm {out} }}$ consists in a family of functions ${\displaystyle f\left(\cdot$ ;\theta \right):\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} ^{n_{\mathrm {out} }}} parametrized by a vector of parameters ${\displaystyle \theta \in \mathbb {R} ^{P}}$.

In this case, the Neural Tangent Kernel ${\displaystyle \Theta$ :\mathbb {R} ^{n_{\mathrm {in} }}\times \mathbb {R} ^{n_{\mathrm {in} }}\to {\mathcal {M}}_{n_{\mathrm {out} }}\left(\mathbb {R} \right)} is a matrix-valued kernel, with values in the space of ${\displaystyle n_{\mathrm {out} }\times n_{\mathrm {out} }}$ matrices, defined by

$${\displaystyle \Theta _{k,l}\left(x,y;\theta \right)=\sum _{p=1}^{P}\partial _{\theta _{p}}f_{k}\left(x;\theta \right)\partial _{\theta _{p}}f_{l}\left(y;\theta \right).}$$

For a dataset ${\displaystyle \left(x_{i}\right)_{i=1,\ldots ,n}\subset \mathbb {R} ^{n_{\mathrm {in} }}}$ with scalar labels ${\displaystyle \left(z_{i}\right)_{i=1,\ldots ,n}\subset \mathbb {R} }$ and a loss function ${\displaystyle c:\mathbb {R} \times \mathbb {R} \to \mathbb {R} }$, the associated empirical loss, defined on functions ${\displaystyle f:\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} }$, is given by

$${\displaystyle {\mathcal {C}}\left(f\right)=\sum _{i=1}^{n}c\left(f\left(x_{i}\right),z_{i}\right).}$$

When training the ANN ${\displaystyle f\left(\cdot$ ;\theta \right):\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} } is trained to fit the dataset (i.e. minimize ${\displaystyle {\mathcal {C}}}$) via continuous-time gradient descent, the parameters ${\displaystyle \left(\theta \left(t\right)\right)_{t\geq 0}}$ evolve through the ordinary differential equation:

$${\displaystyle \partial _{t}\theta \left(t\right)=-\nabla {\mathcal {C}}\left(f\left(\cdot$$ ;\theta \right)\right).}

During training the ANN output function follows an evolution differential equation given in terms of the NTK:

${\displaystyle \partial _{t}f\left(x;\theta \left(t\right)\right)=-\sum _{i=1}^{n}\Theta \left(x,x_{i};\theta \right)\partial _{w}c\left(w,z_{i}\right){\Big |}_{w=f\left(x_{i};\theta \left(t\right)\right)}.}$

This equation shows how the NTK drives the dynamics of ${\displaystyle f\left(\cdot$ ;\theta \left(t\right)\right)} in the space of functions ${\displaystyle \mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} }$ during training.

For a dataset ${\displaystyle \left(x_{i}\right)_{i=1,\ldots ,n}\subset \mathbb {R} ^{n_{\mathrm {in} }}}$ with vector labels ${\displaystyle \left(z_{i}\right)_{i=1,\ldots ,n}\subset \mathbb {R} ^{n_{\mathrm {out} }}}$ and a loss function ${\displaystyle c:\mathbb {R} ^{n_{\mathrm {out} }}\times \mathbb {R} ^{n_{\mathrm {out} }}\to \mathbb {R} }$, the corresponding empirical loss on functions ${\displaystyle f:\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} ^{n_{\mathrm {out} }}}$ is defined by

$${\displaystyle {\mathcal {C}}\left(f\right)=\sum _{i=1}^{n}c\left(f\left(x_{i}\right),z_{i}\right).}$$

The training of ${\displaystyle f_{\theta \left(t\right)}}$ through continuous-time gradient descent yields the following evolution in function space driven by the NTK:

$${\displaystyle \partial _{t}f_{k}\left(x;\theta \left(t\right)\right)=-\sum _{i=1}^{n}\sum _{l=1}^{n_{\mathrm {out} }}\Theta _{k,l}\left(x,x_{i};\theta \right)\partial _{w_{l}}c\left(\left(w_{1},\ldots ,w_{n_{\mathrm {out} }}\right),z_{i}\right){\Big |}_{w=f\left(x_{i};\theta \left(t\right)\right)}.}$$

The NTK ${\displaystyle \Theta \left(x,x_{i};\theta \right)}$ represents the influence of the loss gradient ${\displaystyle \partial _{w}c\left(w,z_{i}\right){\big |}_{w=f\left(x_{i};\theta \right)}}$ with respect to example ${\displaystyle i}$ on the evolution of ANN output ${\displaystyle f\left(x;\theta \right)}$ through a gradient descent step: in the scalar case, this reads

$${\displaystyle f\left(x;\theta \left(t+\epsilon \right)\right)-f\left(x;\theta \left(t\right)\right)\approx \epsilon \sum _{i=1}^{n}\Theta \left(x,x_{i};\theta \left(t\right)\right)\partial _{w}c\left(w,z_{i}\right){\big |}_{w=f\left(x_{i};\theta \right)}.}$$

In particular, each data point ${\displaystyle x_{i}}$ influences the evolution of the output ${\displaystyle f\left(x;\theta \right)}$ for each ${\displaystyle x}$ throughout the training, in a way that is captured by the NTK ${\displaystyle \Theta \left(x,x_{i};\theta \right)}$.

Consider an ANN with fully-connected layers ${\displaystyle \ell =0,\ldots ,L}$ of widths ${\displaystyle n_{0}=n_{\mathrm {in} },n_{1},\ldots ,n_{L}=n_{\mathrm {out} }}$, so that ${\displaystyle f\left(\cdot$ ;\theta \right)=R_{L-1}\circ \cdots \circ R_{0}} , where ${\displaystyle R_{\ell }=\sigma \circ A_{\ell }}$ is the composition of an affine transformation ${\displaystyle A_{i}}$ with the pointwise application of a nonlinearity ${\displaystyle \sigma$ :\mathbb {R} \to \mathbb {R} } , where ${\displaystyle \theta }$ parametrizes the maps ${\displaystyle A_{0},\ldots ,A_{L-1}}$. The parameters ${\displaystyle \theta \in \mathbb {R} ^{P}}$ are initialized randomly, in an independent identically distributed way.

The scale of the NTK as the widths grow is affected by the exact parametrization of the ${\displaystyle A_{i}}$'s and by the initialization of the parameters. This motivates the so-called NTK parametrization ${\displaystyle A_{\ell }\left(x\right)={\frac {1}{\sqrt {n_{\ell }}}}W^{\left(\ell \right)}x+b^{\left(\ell \right)}}$. This parametrization ensures that if the parameters ${\displaystyle \theta \in \mathbb {R} ^{P}}$ are initialized as standard normal variables, the NTK has a finite nontrivial limit. In the large-width limit, the NTK converges to a deterministic (non-random) limit ${\displaystyle \Theta _{\infty }}$, which stays constant in time.

The NTK ${\displaystyle \Theta _{\infty }}$ is explicitly given by ${\displaystyle \Theta _{\infty }=\Theta ^{\left(L\right)}}$, where ${\displaystyle \Theta ^{\left(L\right)}}$ is determined by the set of recursive equations:

${\displaystyle {\begin{aligned}\Theta ^{\left(1\right)}\left(x,y\right)&=\Sigma ^{\left(1\right)}\left(x,y\right),\\\Sigma ^{\left(1\right)}\left(x,y\right)&={\frac {1}{n_{\mathrm {in} }}}x^{T}y+1,\\\Theta ^{\left(\ell +1\right)}\left(x,y\right)&=\Theta ^{\left(\ell \right)}\left(x,y\right){\dot {\Sigma }}^{\left(\ell +1\right)}\left(x,y\right)+\Sigma ^{\left(\ell +1\right)}\left(x,y\right),\\\Sigma ^{\left(\ell +1\right)}\left(x,y\right)&=L_{\Sigma ^{\left(\ell \right)}}^{\sigma }\left(x,y\right),\\{\dot {\Sigma }}^{\left(\ell +1\right)}\left(x,y\right)&=L_{\Sigma ^{\left(\ell \right)}}^{\dot {\sigma }},\end{aligned}}}$

where ${\displaystyle L_{K}^{f}}$ denotes the kernel defined in terms of the Gaussian expectation:

${\displaystyle L_{K}^{f}\left(x,y\right)=\mathbb {E} _{\left(X,Y\right)\sim {\mathcal {N}}\left(0,{\begin{pmatrix}K\left(x,x\right)&K\left(x,y\right)\\K\left(y,x\right)&K\left(y,y\right)\end{pmatrix}}\right)}\left[f\left(X\right)f\left(Y\right)\right].}$

In this formula the kernels ${\displaystyle \Sigma ^{\left(\ell \right)}}$ are the so-called activation kernels[12][13][14] of the ANN.

The NTK describes the evolution of neural networks under gradient descent in function space. Dual to this perspective is an understanding of how neural networks evolve in parameter space, since the NTK is defined in terms of the gradient of the ANN's outputs with respect to its parameters. In the infinite width limit, the connection between these two perspectives becomes especially interesting. The NTK remaining constant throughout training at large widths co-occurs with the ANN being well described throughout training by its first order Taylor expansion around its parameters at initialization:[9]

${\displaystyle f\left(x;\theta (t)\right)=f\left(x;\theta (0)\right)+\nabla _{\theta }f\left(x;\theta (0)\right)\left(\theta (t)-\theta (0)\right)+{\mathcal {O}}\left(\min \left(n_{1}\dots n_{L-1}\right)^{-{\frac {1}{2}}}\right).}$

The NTK can be studied for various ANN architectures,[10] in particular Convolutional Neural Networks (CNNs),[15] Recurrent Neural Networks (RNNs), Transformer Neural Networks.[16] In such settings, the large-width limit corresponds to letting the number of parameters grow, while keeping the number of layers fixed: for CNNs, this amounts to letting the number of channels grow.

For a convex loss functional ${\displaystyle {\mathcal {C}}}$ with a global minimum, if the NTK remains positive-definite during training, the loss of the ANN ${\displaystyle {\mathcal {C}}\left(f\left(\cdot$ ;\theta \left(t\right)\right)\right)} converges to that minimum as ${\displaystyle t\to \infty }$. This positive-definiteness property has been shown in a number of cases, yielding the first proofs that large-width ANNs converge to global minima during training.[1][7][17]

The NTK gives a rigorous connection between the inference performed by infinite-width ANNs and that performed by kernel methods: when the loss function is the least-squares loss, the inference performed by an ANN is in expectation equal to the kernel ridge regression (with zero ridge) with respect to the NTK ${\displaystyle \Theta _{\infty }}$. This suggests that the performance of large ANNs in the NTK parametrization can be replicated by kernel methods for suitably chosen kernels.[1][10]