Foundation

Some Math behind Neural Tangent Kernel

Neural networks are widely recognized as over-parameterized and can often fit data easily, achieving near-zero training loss while still generalizing reasonably well to a test dataset. Despite the fact that these parameters are randomly initialized, the optimization process can reliably produce comparably strong results. This behavior also persists when the number of model parameters exceeds the number of training data points. The neural tangent kernel (NTK) (Jacot et al. 2018) is a kernel-based framework for describing how neural networks evolve during gradient-descent training. It provides key insight into why sufficiently wide neural networks can consistently converge to a global minimum when trained to minimize an empirical loss. In this post, we take a detailed look at the motivation and definition of the NTK, and we also cover the proof of deterministic convergence across different random initializations for neural networks in the infinite-width limit by characterizing the NTK in that regime.

2022 · 17 min read · Curated and presented by Arthur Sedek

Neural networks are well known to be over-parameterized. In practice, they can often fit the training data to near-zero loss while still generalizing reasonably well on the test dataset. Even though these parameters are initialized randomly, the optimization procedure can reliably produce comparably strong solutions. This behavior persists even when the number of model parameters exceeds the number of training data points.

Neural tangent kernel (NTK) (Jacot et al. 2018) is a kernel-based framework for explaining the evolution of neural networks trained via gradient descent. It offers key insights into why sufficiently wide neural networks can consistently converge to a global minimum when minimizing an empirical loss. In this post, we take a deep dive into the motivation and definition of the NTK, and we walk through the proof that, for infinite-width networks, convergence is deterministic across different random initializations by characterizing the NTK in that regime.

🤓 Unlike my earlier posts, this one focuses primarily on a small set of foundational papers rather than providing broad coverage of the full literature. Many interesting works appeared after NTK that modify or extend the theory for understanding NN learning dynamics, but they are out of scope here. The aim is to present all the mathematics behind NTK in a clear, easy-to-follow format, so the post is highly math-intensive. If you spot any errors, please let me know, and I will be glad to correct them promptly. Thanks in advance!

Basics

This section reviews several elementary concepts that are essential for understanding the neural tangent kernel. Feel free to skip it.

Vector-to-vector Derivative

Given an input vector $\mathbf{x} \in \mathbb{R}^n$ (as a column vector) and a function $f: \mathbb{R}^n \to \mathbb{R}^m$, the derivative of $f$ with respective to $\mathbf{x}$ is a $m\times n$ matrix, also known as the Jacobian matrix:

$ J = \frac{\partial f}{\partial \mathbf{x}} = \begin{bmatrix} \frac{\partial f_1}{\partial x_1} & \dots &\frac{\partial f_1}{\partial x_n} \\ \vdots & & \\ \frac{\partial f_m}{\partial x_1} & \dots &\frac{\partial f_m}{\partial x_n} \\ \end{bmatrix} \in \mathbb{R}^{m \times n} $

Throughout the post, integer subscripts refer to a single entry in a vector or matrix. That is, $x_i$ denotes the $i$-th element of the vector $\mathbf{x}$, and $f_i(.)$ is the $i$-th entry of the function output.

The gradient of a vector with respect to a vector is defined as $\nabla_\mathbf{x} f = J^\top \in \mathbb{R}^{n \times m}$, and this formulation also applies when $m=1$ (that is, for a scalar output).

Differential Equations

Differential equations characterize relationships between one or more functions and their derivatives. Two main classes are commonly distinguished.

(1) ODE (Ordinary differential equation) involves an unknown function of a single random variable. ODEs are the primary type of differential equation used in this post. A general ODE can be written as $(x, y, \frac{dy}{dx}, \dots, \frac{d^ny}{dx^n}) = 0$.
(2) PDE (Partial differential equation) involves unknown multivariable functions and their partial derivatives.

Let us review the simplest setting and its solution method. Separation of variables (Fourier method) can be applied when all terms involving one variable can be moved to one side of the equation, while the remaining terms are moved to the other side. For example:

$ \begin{aligned} \text{Given }a\text{ is a constant scalar:}\quad\frac{dy}{dx} &= ay \\ \text{Move same variables to the same side:}\quad\frac{dy}{y} &= adx \\ \text{Put integral on both sides:}\quad\int \frac{dy}{y} &= \int adx \\ \ln (y) &= ax + C' \\ \text{Finally}\quad y &= e^{ax + C'} = C e^{ax} \end{aligned} $

Central Limit Theorem

Given a collection of i.i.d. random variables, $x_1, \dots, x_N$ with mean $\mu$ and variance $\sigma^2$, the Central Limit Theorem (CTL) states that the expectation would be Gaussian distributed when $N$ becomes really large.

$ \bar{x} = \frac{1}{N}\sum_{i=1}^N x_i \sim \mathcal{N}(\mu, \frac{\sigma^2}{n})\quad\text{when }N \to \infty $

CTL also extends to multidimensional vectors. In that case, rather than a single scale $\sigma^2$, we compute the covariance matrix of the random variable $\Sigma$.

Taylor Expansion

The Taylor expansion expresses a function as an infinite sum of terms, where each term is defined using derivatives of the function. The Taylor expansion of a function $f(x)$ around $x=a$ can be written as: $ f(x) = f(a) + \sum_{k=1}^\infty \frac{1}{k!} (x - a)^k\nabla^k_xf(x)\vert_{x=a} $ where $\nabla^k$ denotes the $k$-th derivative.

The first-order Taylor expansion is commonly used as a linear approximation of the function value:

$ f(x) \approx f(a) + (x - a)\nabla_x f(x)\vert_{x=a} $

Kernel & Kernel Methods

A kernel is, in essence, a similarity function between two data points, $K: \mathcal{X} \times \mathcal{X} \to \mathbb{R}$. It quantifies how sensitive the prediction for one data sample is to the prediction for the other, or equivalently, how similar the two data points are. A kernel should be symmetric, $K(x, x’) = K(x’, x)$.

Depending on the structure of the problem, some kernels can be decomposed into two feature maps, one for each data point, so that the kernel value is the inner product of these features: $K(x, x’) = \langle \varphi(x), \varphi(x’) \rangle$.

Kernel methods are non-parametric, instance-based machine learning algorithms. Suppose we know the labels of the training samples $\{x^{(i)}, y^{(i)}\}$. Then the label for a new input $x$ is predicted as a weighted sum $\sum_{i} K(x^{(i)}, x)y^{(i)}$.

Gaussian Processes

A Gaussian process (GP) is a non-parametric approach that models a multivariate Gaussian probability distribution over a collection of random variables. A GP places a prior over functions and then updates the posterior over functions based on observed data.

Given a set of data points $\{x^{(1)}, \dots, x^{(N)}\}$, a GP assumes they follow a joint multivariate Gaussian distribution specified by a mean $\mu(x)$ and a covariance matrix $\Sigma(x)$. The entry at location $(i,j)$ in the covariance matrix $\Sigma(x)$ is determined by a kernel $\Sigma_{i,j} = K(x^{(i)}, x^{(j)})$, also called a covariance function. The essential idea is that if two data points are considered similar by the kernel, their function outputs should also be close. Making predictions for unobserved points with a GP is equivalent to sampling from this distribution, using the conditional distribution of unknown points given the observed ones.

See this post for a high-quality, highly visual tutorial on Gaussian Processes.

Notation

Consider a fully connected neural network with parameters $\theta$, $f(.;\theta): \mathbb{R}^{n_0} \to \mathbb{R}^{n_L}$. Layers are indexed from 0 (input) to $L$ (output). Each layer contains $n_0, \dots, n_L$ neurons, including an input of size $n_0$ and an output of size $n_L$. The network has $P = \sum_{l=0}^{L-1} (n_l + 1) n_{l+1}$ parameters in total, and therefore $\theta \in \mathbb{R}^P$.

The training dataset contains $N$ data points, $\mathcal{D}=\{\mathbf{x}^{(i)}, y^{(i)}\}_{i=1}^N$. We denote all inputs by $\mathcal{X}=\{\mathbf{x}^{(i)}\}_{i=1}^N$ and all labels by $\mathcal{Y}=\{y^{(i)}\}_{i=1}^N$.

Next, we examine the forward-pass computation in each layer in detail. For $l=0, \dots, L-1$, each layer $l$ defines an affine transformation $A^{(l)}$ with weight matrix $\mathbf{w}^{(l)} \in \mathbb{R}^{n_{l} \times n_{l+1}}$ and bias term $\mathbf{b}^{(l)} \in \mathbb{R}^{n_{l+1}}$, followed by a pointwise nonlinearity $\sigma(.)$ that is Lipschitz continuous.

$ \begin{aligned} A^{(0)} &= \mathbf{x} \\ \tilde{A}^{(l+1)}(\mathbf{x}) &= \frac{1}{\sqrt{n_l}} {\mathbf{w}^{(l)}}^\top A^{(l)} + \beta\mathbf{b}^{(l)}\quad\in\mathbb{R}^{n_{l+1}} & \text{; pre-activations}\\ A^{(l+1)}(\mathbf{x}) &= \sigma(\tilde{A}^{(l+1)}(\mathbf{x}))\quad\in\mathbb{R}^{n_{l+1}} & \text{; post-activations} \end{aligned} $

Note that the NTK parameterization rescales the weights by $1/\sqrt{n_l}$ in the transformation to avoid divergence in the infinite-width limit. The constant scalar $\beta \geq 0$ controls the magnitude of the bias terms.

In the analysis below, all network parameters are initialized i.i.d. from a Gaussian distribution $\mathcal{N}(0, 1)$.

Neural Tangent Kernel

Neural tangent kernel (NTK) (Jacot et al. 2018) is a central tool for understanding neural network training under gradient descent. At a high level, it captures how updating model parameters using one data sample changes the predictions on other samples.

We develop the intuition for NTK step by step.

The empirical loss $\mathcal{L}: \mathbb{R}^P \to \mathbb{R}_+$ minimized during training is defined using a per-sample cost function $\ell: \mathbb{R}^{n_0} \times \mathbb{R}^{n_L} \to \mathbb{R}_+$:

$ \mathcal{L}(\theta) =\frac{1}{N} \sum_{i=1}^N \ell(f(\mathbf{x}^{(i)}; \theta), y^{(i)}) $

Then, by the chain rule, the gradient of the loss is:

$ \nabla_\theta \mathcal{L}(\theta)= \frac{1}{N} \sum_{i=1}^N \underbrace{\nabla_\theta f(\mathbf{x}^{(i)}; \theta)}_{\text{size }P \times n_L} \underbrace{\nabla_f \ell(f, y^{(i)})}_{\text{size } n_L \times 1} $

To track how the network parameters $\theta$ evolve over time, observe that each gradient descent step introduces a small increment with an infinitesimal step size. When the update step is sufficiently small, it can be approximated by a time derivative:

$ \frac{d\theta}{d t} = - \nabla_\theta\mathcal{L}(\theta) = -\frac{1}{N} \sum_{i=1}^N \nabla_\theta f(\mathbf{x}^{(i)}; \theta) \nabla_f \ell(f, y^{(i)}) $

Again applying the chain rule, the network output evolves according to:

$ \frac{df(\mathbf{x};\theta)}{dt} = \frac{df(\mathbf{x};\theta)}{d\theta}\frac{d\theta}{dt} = -\frac{1}{N} \sum_{i=1}^N \color{blue}{\underbrace{\nabla_\theta f(\mathbf{x};\theta)^\top \nabla_\theta f(\mathbf{x}^{(i)}; \theta)}_\text{Neural tangent kernel}} \color{black}{\nabla_f \ell(f, y^{(i)})} $

This reveals the Neural Tangent Kernel (NTK), defined by the blue term in the expression above, $K: \mathbb{R}^{n_0}\times\mathbb{R}^{n_0} \to \mathbb{R}^{n_L \times n_L}$:

$ K(\mathbf{x}, \mathbf{x}'; \theta) = \nabla_\theta f(\mathbf{x};\theta)^\top \nabla_\theta f(\mathbf{x}'; \theta) $

where each entry of the output matrix at location $(m, n), 1 \leq m, n \leq n_L$ is:

$ K_{m,n}(\mathbf{x}, \mathbf{x}'; \theta) = \sum_{p=1}^P \frac{\partial f_m(\mathbf{x};\theta)}{\partial \theta_p} \frac{\partial f_n(\mathbf{x}';\theta)}{\partial \theta_p} $

The corresponding “feature map” for an input $\mathbf{x}$ is $\varphi(\mathbf{x}) = \nabla_\theta f(\mathbf{x};\theta)$.

Infinite Width Networks

To explain why individual gradient descent steps behave so similarly across different random initializations, several pioneering theoretical works begin with the infinite-width setting. Below, we examine a detailed NTK-based proof showing how infinite-width networks are guaranteed to converge to a global minimum when trained to minimize an empirical loss.

Connection with Gaussian Processes

Deep neural networks have a close connection with Gaussian processes (Neal 1994). The output functions of a $L$-layer network, $f_i(\mathbf{x}; \theta)$ for $i=1, \dots, n_L$, are i.i.d. centered Gaussian processes with covariance $\Sigma^{(L)}$, defined recursively by:

$ \begin{aligned} \Sigma^{(1)}(\mathbf{x}, \mathbf{x}') &= \frac{1}{n_0}\mathbf{x}^\top{\mathbf{x}'} + \beta^2 \\ \lambda^{(l+1)}(\mathbf{x}, \mathbf{x}') &= \begin{bmatrix} \Sigma^{(l)}(\mathbf{x}, \mathbf{x}) & \Sigma^{(l)}(\mathbf{x}, \mathbf{x}') \\ \Sigma^{(l)}(\mathbf{x}', \mathbf{x}) & \Sigma^{(l)}(\mathbf{x}', \mathbf{x}') \end{bmatrix} \\ \Sigma^{(l+1)}(\mathbf{x}, \mathbf{x}') &= \mathbb{E}_{f \sim \mathcal{N}(0, \lambda^{(l)})}[\sigma(f(\mathbf{x})) \sigma(f(\mathbf{x}'))] + \beta^2 \end{aligned} $

Lee & Bahri et al. (2018) proved this result by mathematical induction:

(1) Start with $L=1$, where there is no nonlinearity and the input undergoes only a single affine transformation:

$ \begin{aligned} f(\mathbf{x};\theta) = \tilde{A}^{(1)}(\mathbf{x}) &= \frac{1}{\sqrt{n_0}}{\mathbf{w}^{(0)}}^\top\mathbf{x} + \beta\mathbf{b}^{(0)} \\ \text{where }\tilde{A}_m^{(1)}(\mathbf{x}) &= \frac{1}{\sqrt{n_0}}\sum_{i=1}^{n_0} w^{(0)}_{im}x_i + \beta b^{(0)}_m\quad \text{for }1 \leq m \leq n_1 \end{aligned} $

Because weights and biases are initialized i.i.d., all output dimensions of this network ${\tilde{A}^{(1)}_1(\mathbf{x}), \dots, \tilde{A}^{(1)}_{n_1}(\mathbf{x})}$ are also i.i.d. For different inputs, the $m$-th network outputs $\tilde{A}^{(1)}_m(.)$ have a joint multivariate Gaussian distribution, which is equivalent to a Gaussian process with the following covariance function (with mean $\mu_w=\mu_b=0$ and variance $\sigma^2_w = \sigma^2_b=1$):

$ \begin{aligned} \Sigma^{(1)}(\mathbf{x}, \mathbf{x}') &= \mathbb{E}[\tilde{A}_m^{(1)}(\mathbf{x})\tilde{A}_m^{(1)}(\mathbf{x}')] \\ &= \mathbb{E}\Big[\Big( \frac{1}{\sqrt{n_0}}\sum_{i=1}^{n_0} w^{(0)}_{i,m}x_i + \beta b^{(0)}_m \Big) \Big( \frac{1}{\sqrt{n_0}}\sum_{i=1}^{n_0} w^{(0)}_{i,m}x'_i + \beta b^{(0)}_m \Big)\Big] \\ &= \frac{1}{n_0} \sigma^2_w \sum_{i=1}^{n_0} \sum_{j=1}^{n_0} x_i{x'}_j + \frac{\beta \mu_b}{\sqrt{n_0}} \sum_{i=1}^{n_0} w_{im}(x_i + x'_i) + \sigma^2_b \beta^2 \\ &= \frac{1}{n_0}\mathbf{x}^\top{\mathbf{x}'} + \beta^2 \end{aligned} $

(2) For the induction step, assume the proposition holds for $L=l$, that is, a $l$-layer network. Under this assumption, $\tilde{A}^{(l)}_m(.)$ is a Gaussian process with covariance $\Sigma^{(l)}$, and $\{\tilde{A}^{(l)}_i\}_{i=1}^{n_l}$ are i.i.d.

We then prove the proposition for $L=l+1$. The outputs are computed as:

$ \begin{aligned} f(\mathbf{x};\theta) = \tilde{A}^{(l+1)}(\mathbf{x}) &= \frac{1}{\sqrt{n_l}}{\mathbf{w}^{(l)}}^\top \sigma(\tilde{A}^{(l)}(\mathbf{x})) + \beta\mathbf{b}^{(l)} \\ \text{where }\tilde{A}^{(l+1)}_m(\mathbf{x}) &= \frac{1}{\sqrt{n_l}}\sum_{i=1}^{n_l} w^{(l)}_{im}\sigma(\tilde{A}^{(l)}_i(\mathbf{x})) + \beta b^{(l)}_m \quad \text{for }1 \leq m \leq n_{l+1} \end{aligned} $

We can infer that the expected value of the sum of contributions from the previous hidden layer is zero:

$ \begin{aligned} \mathbb{E}[w^{(l)}_{im}\sigma(\tilde{A}^{(l)}_i(\mathbf{x}))] &= \mathbb{E}[w^{(l)}_{im}]\mathbb{E}[\sigma(\tilde{A}^{(l)}_i(\mathbf{x}))] = \mu_w \mathbb{E}[\sigma(\tilde{A}^{(l)}_i(\mathbf{x}))] = 0 \\ \mathbb{E}[\big(w^{(l)}_{im}\sigma(\tilde{A}^{(l)}_i(\mathbf{x}))\big)^2] &= \mathbb{E}[{w^{(l)}_{im}}^2]\mathbb{E}[\sigma(\tilde{A}^{(l)}_i(\mathbf{x}))^2] = \sigma_w^2 \Sigma^{(l)}(\mathbf{x}, \mathbf{x}) = \Sigma^{(l)}(\mathbf{x}, \mathbf{x}) \end{aligned} $

Because $\{\tilde{A}^{(l)}_i(\mathbf{x})\}_{i=1}^{n_l}$ are i.i.d., the central limit theorem implies that as the hidden layer becomes infinitely wide $n_l \to \infty$, $\tilde{A}^{(l+1)}_m(\mathbf{x})$ is Gaussian distributed with variance $\beta^2 + \text{Var}(\tilde{A}_i^{(l)}(\mathbf{x}))$. Note that ${\tilde{A}^{(l+1)}_1(\mathbf{x}), \dots, \tilde{A}^{(l+1)}_{n_{l+1}}(\mathbf{x})}$ remain i.i.d.

Therefore, $\tilde{A}^{(l+1)}_m(.)$ is equivalent to a Gaussian process with covariance function:

$ \begin{aligned} \Sigma^{(l+1)}(\mathbf{x}, \mathbf{x}') &= \mathbb{E}[\tilde{A}^{(l+1)}_m(\mathbf{x})\tilde{A}^{(l+1)}_m(\mathbf{x}')] \\ &= \frac{1}{n_l} \sigma\big(\tilde{A}^{(l)}_i(\mathbf{x})\big)^\top \sigma\big(\tilde{A}^{(l)}_i(\mathbf{x}')\big) + \beta^2 \quad\text{;similar to how we get }\Sigma^{(1)} \end{aligned} $

When $n_l \to \infty$, the central limit theorem gives:

$ \Sigma^{(l+1)}(\mathbf{x}, \mathbf{x}') \to \mathbb{E}_{f \sim \mathcal{N}(0, \Lambda^{(l)})}[\sigma(f(\mathbf{x}))^\top \sigma(f(\mathbf{x}'))] + \beta^2 $

The Gaussian process obtained via this construction is known as the Neural Network Gaussian Process (NNGP) (Lee & Bahri et al. (2018)).

Deterministic Neural Tangent Kernel

We are now ready to address the most important proposition from the NTK paper:

When $n_1, \dots, n_L \to \infty$ (an infinite-width network), the NTK converges such that:

(1) it is deterministic at initialization, meaning the kernel is independent of the specific initialization values and is determined only by the architecture, and
(2) it remains constant throughout training.

The proof again uses mathematical induction:

(1) First, $K^{(0)} = 0$ always holds. When $L=1$, we can write the NTK explicitly. It is deterministic and does not depend on network initialization. Because there is no hidden layer, there is no width that can be taken to infinity.

$ \begin{aligned} f(\mathbf{x};\theta) &= \tilde{A}^{(1)}(\mathbf{x}) = \frac{1}{\sqrt{n_0}} {\mathbf{w}^{(0)}}^\top\mathbf{x} + \beta\mathbf{b}^{(0)} \\ K^{(1)}(\mathbf{x}, \mathbf{x}';\theta) &= \Big(\frac{\partial f(\mathbf{x}';\theta)}{\partial \mathbf{w}^{(0)}}\Big)^\top \frac{\partial f(\mathbf{x};\theta)}{\partial \mathbf{w}^{(0)}} + \Big(\frac{\partial f(\mathbf{x}';\theta)}{\partial \mathbf{b}^{(0)}}\Big)^\top \frac{\partial f(\mathbf{x};\theta)}{\partial \mathbf{b}^{(0)}} \\ &= \frac{1}{n_0} \mathbf{x}^\top{\mathbf{x}'} + \beta^2 = \Sigma^{(1)}(\mathbf{x}, \mathbf{x}') \end{aligned} $

(2) Next, for $L=l$, assume that a $l$-layer network with $\tilde{P}$ parameters in total, $\tilde{\theta} = (\mathbf{w}^{(0)}, \dots, \mathbf{w}^{(l-1)}, \mathbf{b}^{(0)}, \dots, \mathbf{b}^{(l-1)}) \in \mathbb{R}^\tilde{P}$, has an NTK that converges to a deterministic limit as $n_1, \dots, n_{l-1} \to \infty$.

$ K^{(l)}(\mathbf{x}, \mathbf{x}';\tilde{\theta}) = \nabla_{\tilde{\theta}} \tilde{A}^{(l)}(\mathbf{x})^\top \nabla_{\tilde{\theta}} \tilde{A}^{(l)}(\mathbf{x}') \to K^{(l)}_{\infty}(\mathbf{x}, \mathbf{x}') $

Note that $K_\infty^{(l)}$ does not depend on $\theta$.

Now consider $L=l+1$. Relative to a $l$-layer network, a $(l+1)$-layer network includes an additional weight matrix $\mathbf{w}^{(l)}$ and bias term $\mathbf{b}^{(l)}$, so the full parameter set additionally contains $\theta = (\tilde{\theta}, \mathbf{w}^{(l)}, \mathbf{b}^{(l)})$.

The output function of this $(l+1)$-layer network is:

$ f(\mathbf{x};\theta) = \tilde{A}^{(l+1)}(\mathbf{x};\theta) = \frac{1}{\sqrt{n_l}} {\mathbf{w}^{(l)}}^\top \sigma\big(\tilde{A}^{(l)}(\mathbf{x})\big) + \beta \mathbf{b}^{(l)} $

We also have its derivatives with respect to the different parameter blocks. For brevity, denote $\tilde{A}^{(l)} = \tilde{A}^{(l)}(\mathbf{x})$ in the following expression:

$ \begin{aligned} \nabla_{\color{blue}{\mathbf{w}^{(l)}}} f(\mathbf{x};\theta) &= \color{blue}{ \frac{1}{\sqrt{n_l}} \sigma\big(\tilde{A}^{(l)}\big)^\top } \color{black}{\quad \in \mathbb{R}^{1 \times n_l}} \\ \nabla_{\color{green}{\mathbf{b}^{(l)}}} f(\mathbf{x};\theta) &= \color{green}{ \beta } \\ \nabla_{\color{red}{\tilde{\theta}}} f(\mathbf{x};\theta) &= \frac{1}{\sqrt{n_l}} \nabla_\tilde{\theta}\sigma(\tilde{A}^{(l)}) \mathbf{w}^{(l)} \\ &= \color{red}{ \frac{1}{\sqrt{n_l}} \begin{bmatrix} \dot{\sigma}(\tilde{A}_1^{(l)})\frac{\partial \tilde{A}_1^{(l)}}{\partial \tilde{\theta}_1} & \dots & \dot{\sigma}(\tilde{A}_{n_l}^{(l)})\frac{\partial \tilde{A}_{n_l}^{(l)}}{\partial \tilde{\theta}_1} \\ \vdots \\ \dot{\sigma}(\tilde{A}_1^{(l)})\frac{\partial \tilde{A}_1^{(l)}}{\partial \tilde{\theta}_\tilde{P}} & \dots & \dot{\sigma}(\tilde{A}_{n_l}^{(l)})\frac{\partial \tilde{A}_{n_l}^{(l)}}{\partial \tilde{\theta}_\tilde{P}}\\ \end{bmatrix} \mathbf{w}^{(l)} \color{black}{\quad \in \mathbb{R}^{\tilde{P} \times n_{l+1}}} } \end{aligned} $

where $\dot{\sigma}$ is the derivative of $\sigma$, and the entry at location $(p, m), 1 \leq p \leq \tilde{P}, 1 \leq m \leq n_{l+1}$ in the matrix $\nabla_{\tilde{\theta}} f(\mathbf{x};\theta)$ can be written as:

$ \frac{\partial f_m(\mathbf{x};\theta)}{\partial \tilde{\theta}_p} = \sum_{i=1}^{n_l} w^{(l)}_{im} \dot{\sigma}\big(\tilde{A}_i^{(l)} \big) \nabla_{\tilde{\theta}_p} \tilde{A}_i^{(l)} $

The NTK for this $(l+1)$-layer network is defined accordingly:

$ \begin{aligned} & K^{(l+1)}(\mathbf{x}, \mathbf{x}'; \theta) \\ =& \nabla_{\theta} f(\mathbf{x};\theta)^\top \nabla_{\theta} f(\mathbf{x};\theta) \\ =& \color{blue}{\nabla_{\mathbf{w}^{(l)}} f(\mathbf{x};\theta)^\top \nabla_{\mathbf{w}^{(l)}} f(\mathbf{x};\theta)} + \color{green}{\nabla_{\mathbf{b}^{(l)}} f(\mathbf{x};\theta)^\top \nabla_{\mathbf{b}^{(l)}} f(\mathbf{x};\theta)} + \color{red}{\nabla_{\tilde{\theta}} f(\mathbf{x};\theta)^\top \nabla_{\tilde{\theta}} f(\mathbf{x};\theta)} \\ =& \frac{1}{n_l} \Big[ \color{blue}{\sigma(\tilde{A}^{(l)})\sigma(\tilde{A}^{(l)})^\top} + \color{green}{\beta^2} \\ &+ \color{red}{ {\mathbf{w}^{(l)}}^\top \begin{bmatrix} \dot{\sigma}(\tilde{A}_1^{(l)})\dot{\sigma}(\tilde{A}_1^{(l)})\sum_{p=1}^\tilde{P} \frac{\partial \tilde{A}_1^{(l)}}{\partial \tilde{\theta}_p}\frac{\partial \tilde{A}_1^{(l)}}{\partial \tilde{\theta}_p} & \dots & \dot{\sigma}(\tilde{A}_1^{(l)})\dot{\sigma}(\tilde{A}_{n_l}^{(l)})\sum_{p=1}^\tilde{P} \frac{\partial \tilde{A}_1^{(l)}}{\partial \tilde{\theta}_p}\frac{\partial \tilde{A}_{n_l}^{(l)}}{\partial \tilde{\theta}_p} \\ \vdots \\ \dot{\sigma}(\tilde{A}_{n_l}^{(l)})\dot{\sigma}(\tilde{A}_1^{(l)})\sum_{p=1}^\tilde{P} \frac{\partial \tilde{A}_{n_l}^{(l)}}{\partial \tilde{\theta}_p}\frac{\partial \tilde{A}_1^{(l)}}{\partial \tilde{\theta}_p} & \dots & \dot{\sigma}(\tilde{A}_{n_l}^{(l)})\dot{\sigma}(\tilde{A}_{n_l}^{(l)})\sum_{p=1}^\tilde{P} \frac{\partial \tilde{A}_{n_l}^{(l)}}{\partial \tilde{\theta}_p}\frac{\partial \tilde{A}_{n_l}^{(l)}}{\partial \tilde{\theta}_p} \\ \end{bmatrix} \mathbf{w}^{(l)} } \color{black}{\Big]} \\ =& \frac{1}{n_l} \Big[ \color{blue}{\sigma(\tilde{A}^{(l)})\sigma(\tilde{A}^{(l)})^\top} + \color{green}{\beta^2} \\ &+ \color{red}{ {\mathbf{w}^{(l)}}^\top \begin{bmatrix} \dot{\sigma}(\tilde{A}_1^{(l)})\dot{\sigma}(\tilde{A}_1^{(l)})K^{(l)}_{11} & \dots & \dot{\sigma}(\tilde{A}_1^{(l)})\dot{\sigma}(\tilde{A}_{n_l}^{(l)})K^{(l)}_{1n_l} \\ \vdots \\ \dot{\sigma}(\tilde{A}_{n_l}^{(l)})\dot{\sigma}(\tilde{A}_1^{(l)})K^{(l)}_{n_l1} & \dots & \dot{\sigma}(\tilde{A}_{n_l}^{(l)})\dot{\sigma}(\tilde{A}_{n_l}^{(l)})K^{(l)}_{n_ln_l} \\ \end{bmatrix} \mathbf{w}^{(l)} } \color{black}{\Big]} \end{aligned} $

where the individual entry at location $(m, n), 1 \leq m, n \leq n_{l+1}$ of the matrix $K^{(l+1)}$ can be written as:

$ \begin{aligned} K^{(l+1)}_{mn} =& \frac{1}{n_l}\Big[ \color{blue}{\sigma(\tilde{A}_m^{(l)})\sigma(\tilde{A}_n^{(l)})} + \color{green}{\beta^2} + \color{red}{ \sum_{i=1}^{n_l} \sum_{j=1}^{n_l} w^{(l)}_{im} w^{(l)}_{in} \dot{\sigma}(\tilde{A}_i^{(l)}) \dot{\sigma}(\tilde{A}_{j}^{(l)}) K_{ij}^{(l)} } \Big] \end{aligned} $

When $n_l \to \infty$, the blue and green terms have a limit (see the proof in the previous section):

$ \frac{1}{n_l}\sigma(\tilde{A}^{(l)})\sigma(\tilde{A}^{(l)}) + \beta^2\to \Sigma^{(l+1)} $

and the red term has the limit:

$ \sum_{i=1}^{n_l} \sum_{j=1}^{n_l} w^{(l)}_{im} w^{(l)}_{in} \dot{\sigma}(\tilde{A}_i^{(l)}) \dot{\sigma}(\tilde{A}_{j}^{(l)}) K_{ij}^{(l)} \to \sum_{i=1}^{n_l} \sum_{j=1}^{n_l} w^{(l)}_{im} w^{(l)}_{in} \dot{\sigma}(\tilde{A}_i^{(l)}) \dot{\sigma}(\tilde{A}_{j}^{(l)}) K_{\infty,ij}^{(l)} $

Later, Arora et al. (2019) provided a proof under a weaker limiting condition. Their result does not require every hidden layer to be infinitely wide, but instead assumes only that the minimum width is sufficiently large.

Linearized Models

From the previous section, and by the chain rule for derivatives, we have established that the gradient update on the output of an infinite-width network takes the following form. For brevity, we omit the inputs in the analysis below:

$ \begin{aligned} \frac{df(\theta)}{dt} &= -\eta\nabla_\theta f(\theta)^\top \nabla_\theta f(\theta) \nabla_f \mathcal{L} & \\ &= -\eta\nabla_\theta f(\theta)^\top \nabla_\theta f(\theta) \nabla_f \mathcal{L} & \\ &= -\eta K(\theta) \nabla_f \mathcal{L} \\ &= \color{cyan}{-\eta K_\infty \nabla_f \mathcal{L}} & \text{; for infinite width network}\\ \end{aligned} $

To follow how $\theta$ evolves over time, we treat it as a function of the time step $t$. Using a Taylor expansion, the learning dynamics of the network can be reduced to:

$ f(\theta(t)) \approx f^\text{lin}(\theta(t)) = f(\theta(0)) + \underbrace{\nabla_\theta f(\theta(0))}_{\text{formally }\nabla_\theta f(\mathbf{x}; \theta) \vert_{\theta=\theta(0)}} (\theta(t) - \theta(0)) $

This formulation is commonly known as the linearized model, since $\theta(0)$, $f(\theta(0))$, and $\nabla_\theta f(\theta(0))$ are all constants. Assuming the incremental time step $t$ is extremely small and that the parameters are updated via gradient descent:

$ \begin{aligned} \theta(t) - \theta(0) &= - \eta \nabla_\theta \mathcal{L}(\theta) = - \eta \nabla_\theta f(\theta)^\top \nabla_f \mathcal{L} \\ f^\text{lin}(\theta(t)) - f(\theta(0)) &= - \eta\nabla_\theta f(\theta(0))^\top \nabla_\theta f(\mathcal{X};\theta(0)) \nabla_f \mathcal{L} \\ \frac{df(\theta(t))}{dt} &= - \eta K(\theta(0)) \nabla_f \mathcal{L} \\ \frac{df(\theta(t))}{dt} &= \color{cyan}{- \eta K_\infty \nabla_f \mathcal{L}} & \text{; for infinite width network}\\ \end{aligned} $

We ultimately recover the same learning dynamics. This implies that, in the infinite-width limit, a neural network can be substantially simplified and described by the linearized model above (Lee & Xiao, et al. 2019).

In the special case where the empirical loss is the MSE loss, $\nabla_\theta \mathcal{L}(\theta) = f(\mathcal{X}; \theta) - \mathcal{Y}$, the network dynamics reduce to a linear ODE that admits a closed-form solution:

$ \begin{aligned} \frac{df(\theta)}{dt} =& -\eta K_\infty (f(\theta) - \mathcal{Y}) & \\ \frac{dg(\theta)}{dt} =& -\eta K_\infty g(\theta) & \text{; let }g(\theta)=f(\theta) - \mathcal{Y} \\ \int \frac{dg(\theta)}{g(\theta)} =& -\eta \int K_\infty dt & \\ g(\theta) &= C e^{-\eta K_\infty t} & \end{aligned} $

When $t=0$, we obtain $C=f(\theta(0)) - \mathcal{Y}$, and therefore:

$ f(\theta) = (f(\theta(0)) - \mathcal{Y})e^{-\eta K_\infty t} + \mathcal{Y} \\ = f(\theta(0))e^{-K_\infty t} + (I - e^{-\eta K_\infty t})\mathcal{Y} $

Lazy Training

It has been observed that when a neural network is highly over-parameterized, it can fit the data with the training loss rapidly converging to zero, while the network parameters themselves change very little. This phenomenon is referred to as lazy training. Put differently, even when the loss $\mathcal{L}$ decreases substantially, the change in the network differential $f$ (also known as the Jacobian matrix) remains very small.

Let $\theta(0)$ denote the initial network parameters, and let $\theta(T)$ denote the final parameters after the loss has been minimized to zero. The parameter-space displacement can be approximated using a first-order Taylor expansion:

$ \begin{aligned} \hat{y} = f(\theta(T)) &\approx f(\theta(0)) + \nabla_\theta f(\theta(0)) (\theta(T) - \theta(0)) \\ \text{Thus }\Delta \theta &= \theta(T) - \theta(0) \approx \frac{\|\hat{y} - f(\theta(0))\|}{\| \nabla_\theta f(\theta(0)) \|} \end{aligned} $

Continuing with the first-order Taylor approximation, we can also track how the differential of $f$ changes:

$ \begin{aligned} \nabla_\theta f(\theta(T)) &\approx \nabla_\theta f(\theta(0)) + \nabla^2_\theta f(\theta(0)) \Delta\theta \\ &= \nabla_\theta f(\theta(0)) + \nabla^2_\theta f(\theta(0)) \frac{\|\hat{y} - f(\mathbf{x};\theta(0))\|}{\| \nabla_\theta f(\theta(0)) \|} \\ \text{Thus }\Delta\big(\nabla_\theta f\big) &= \nabla_\theta f(\theta(T)) - \nabla_\theta f(\theta(0)) = \|\hat{y} - f(\mathbf{x};\theta(0))\| \frac{\nabla^2_\theta f(\theta(0))}{\| \nabla_\theta f(\theta(0)) \|} \end{aligned} $

Let $\kappa(\theta)$ represent the relative change in the differential of $f$ with respect to the change in parameter space:

$ \kappa(\theta = \frac{\Delta\big(\nabla_\theta f\big)}{\| \nabla_\theta f(\theta(0)) \|} = \|\hat{y} - f(\theta(0))\| \frac{\nabla^2_\theta f(\theta(0))}{\| \nabla_\theta f(\theta(0)) \|^2} $

Chizat et al. (2019) proved, for a two-layer neural network, that $\mathbb{E}[\kappa(\theta_0)] \to 0$ (that is, entering the lazy regime) when the number of hidden neurons $\to \infty$. For additional discussion of linearized models and lazy training, see this post.

Citation

Cited as:

Weng, Lilian. (Sep 2022). Some math behind neural tangent kernel. Lil’Log. https://lilianweng.github.io/posts/2022-09-08-ntk/.

@article{weng2022ntk,
  title   = "Some Math behind Neural Tangent Kernel",
  author  = "Weng, Lilian",
  journal = "Lil'Log",
  year    = "2022",
  month   = "Sep",
  url     = "https://lilianweng.github.io/posts/2022-09-08-ntk/"
}