# An Infinite-Width Analysis on the Jacobian-Regularised Training of a Neural
Network

Taeyoung Kim
School of Computing, KAIST
Daejeon, South Korea
<EMAIL_ADDRESS>&Hongseok Yang
School of Computing and Kim Jaechul Graduate School of AI, KAIST
Daejeon, South Korea
<EMAIL_ADDRESS>
###### Abstract

The recent theoretical analysis of deep neural networks in their infinite-
width limits has deepened our understanding of initialisation, feature
learning, and training of those networks, and brought new practical techniques
for finding appropriate hyperparameters, learning network weights, and
performing inference. In this paper, we broaden this line of research by
showing that this infinite-width analysis can be extended to the Jacobian of a
deep neural network. We show that a multilayer perceptron (MLP) and its
Jacobian at initialisation jointly converge to a Gaussian process (GP) as the
widths of the MLP’s hidden layers go to infinity and characterise this GP. We
also prove that in the infinite-width limit, the evolution of the MLP under
the so-called robust training (i.e., training with a regulariser on the
Jacobian) is described by a linear first-order ordinary differential equation
that is determined by a variant of the Neural Tangent Kernel. We
experimentally show the relevance of our theoretical claims to wide finite
networks, and empirically analyse the properties of kernel regression solution
to obtain an insight into Jacobian regularisation.

##  1 INTRODUCTION

The recent theoretical analysis of deep neural networks in their infinite-
width limits has substantially deepened our understanding of initialisation
[Neal, 1996, Lee et al., 2018, Matthews et al., 2018, Yang, 2019, Peluchetti
et al., 2020, Lee et al., 2022], feature learning [Yang and Hu, 2021],
hyperparameter turning [Yang et al., 2021], and training of those networks
[Jacot et al., 2018, Lee et al., 2019, Yang and Littwin, 2021]. Seemingly
impossible questions, such as whether the gradient descent achieves the zero
training error and where it converges eventually, are answered [Jacot et al.,
2018], and new techniques for scaling hyperparameters [Yang and Hu, 2021] or
finding appropriate hyperparameters [Yang et al., 2021] have been developed.
Also, the tools used in the analysis, such as the Neural Tangent Kernel, turn
out to be useful for analysing pruning and other optimisations for deep neural
networks [Liu and Zenke, 2020].

Our goal is to extend this infinite-width analysis to the Jacobian of a deep
neural network. By Jacobian, we mean the input-output Jacobian of a network’s
output. This Jacobian has information about the smoothness of the network’s
output and has been used to measure the robustness of a trained network
against noise [Peck et al., 2017] or to learn a network that achieves high
accuracy on a test dataset even when examples in the dataset are corrupted
with noise. The Jacobian of a neural network also features in the work on
network testing, verification, and adversarial attack [Goodfellow et al.,
2015, Zhang et al., 2019, Wang et al., 2021].

We show that at initialisation, a multilayer perceptron (MLP) and its Jacobian
jointly converge to a zero-mean Gaussian process (GP) as the widths of the
MLP’s hidden layers go to infinity, and characterise this GP by describing its
kernel inductively. Our result can be used to compute analytically the
posterior of the MLP (viewed as a random function) conditioned on not just
usual input-output pairs $(x^{(1)},y^{(1)}),\ldots,(x^{(N)},y^{(N)})$ in the
training set, but also derivatives at some inputs, in the infinite-width
setting.

We also analyse the training dynamics of an MLP under the so-called _robust-
training_ objective, which contains a regulariser on the Jacobians of the MLP
at each training input, in addition to the standard loss over input-output
pairs in the training set. As in the case of the standard training without the
Jacobian regulariser, we show that the training can be characterised by a
kernel that is defined in terms of the derivatives of the MLP output and its
Jacobian with respect to the MLP parameters. We call this kernel Jacobian
Neural Tangent Kernel (JNTK) due to its similarity with the standard Neural
Tangent Kernel (NTK) and show that it becomes independent of the MLP
parameters at initialisation, as the widths of the MLP go to infinity. Thus,
the JNTK at initialisation is deterministic in the infinite-width limit. Then,
we show the JNTK of the infinitely wide MLP stays constant during the robust
training. We identify a linear first-order ordinary differential equation
(ODE) that characterises the evolution of this infinitely wide MLP during
robust training and describe the analytic solution of the ODE. We
experimentally confirm that the conclusions of our theoretical analysis apply
to finite MLPs with large but realistic widths. We also include empirical
analysis of the ODE solution obtained and empirically check when our
assumption is satisfied.

The rest of the paper is organised as follows. In Section 2, we describe the
MLPs and robust-learning objective used in the paper. We then present our
theoretical results for initialisation in Section 3 and for training dynamics
in Section 4, and describe preliminary experimental observations in Section 5.
We conclude our paper with related works in Section 6. The missing proofs and
additional experiments can be found in the supplementary material.

##  2 SETUP

We use the following convention and notation. For any $n\in\mathbb{N}$, $[n]$
means the set $\\{1,\ldots,n\\}$. We write $x^{(1:N)}$ to denote the set
$\\{x^{(1)},\ldots,x^{(N)}\\}$. We count the indices of vectors and matrices
from $0$, not $1$. For instance, for $\alpha,\beta\in[d_{0}]$, a vector
$x\in\mathbb{R}^{d_{0}}$, we write $x_{\alpha-1}$ for the $\alpha$-th
component of $x$, and for the matrix $M$ of size $d_{0}\times d_{0}$, we write
$M_{\alpha-1,\beta-1}$ for the $(\alpha,\beta)$-th entry of $x$. Note that we
often use matrix of size $(1+d_{0})\times(1+d_{0})$, where $M_{\alpha,\beta}$
denotes the $(\alpha+1,\beta+1)$-th entry of $M$ for $\alpha,\beta\in[d_{0}]$.

We use multilayer perceptrons (MLPs)
$f:\mathbb{R}^{d_{0}}\to\mathbb{R}^{d_{L+1}}$, which are defined as follows:
for all inputs $x\in\mathbb{R}^{d_{0}}$ and all layers $l=2,\ldots,L$,

$\displaystyle h^{(1)}(x)\in\mathbb{R}^{d_{1}},$ $\displaystyle
h^{(1)}(x)\coloneqq\phi(W^{(1)}x),$ $\displaystyle
h^{(l)}(x)\in\mathbb{R}^{d_{l}},$ $\displaystyle
h^{(l)}(x)\coloneqq\phi\bigg{(}\frac{1}{\sqrt{d_{l-1}}}W^{(l)}h^{(l-1)}(x)\bigg{)},$
$\displaystyle f(x)\in\mathbb{R}^{d_{L+1}},$ $\displaystyle
f(x)\coloneqq\frac{\kappa}{\sqrt{d_{L}}}W^{(L+1)}h^{(L)}(x).$

Here the weights $W^{(l)}\in\mathbb{R}^{d_{l}\times d_{l-1}}$ are initialised
with independent samples from $\mathcal{N}(0,1)$. Note that due to this
randomness of $W^{(l)}$, the MLP $f$ is a random function at initialisation.
Following [Arora et al., 2019], we introduce a fixed non-trainable $\kappa>0$,
which controls the scale of the output of $f$ and thus the randomness (or
variance) of $f$ at initialisation. Suppressing the initial randomness is
known to improve the performance of a neural network in practice [Hu et al.,
2020], and simplify the theoretical analysis of the network.

To simplify the presentation, we do not include the bias terms in our
definition. We also assume that $d\coloneqq d_{1}=\ldots=d_{L}$ and
$d_{L+1}=1$. We also include some assumptions on the activation function
$\phi$ as

######  Assumption 1 (Activation Assumption).

(1) All of $\phi$ and its first and second derivatives $\dot{\phi}$ and
$\ddot{\phi}$ are Lipschitz continuous with the Lipschitz coefficient
$M_{1},M_{2},M_{3}$, respectively.
(2) The activation $\phi$ is normalised such that
$\mathop{\mathbb{E}}_{z\sim\mathcal{N}(0,1)}\left[\phi(z)^{2}\right]=1$.

We denote the set of all trainable parameters by
$\theta\coloneqq\\{W_{ij}^{(l)}\\}_{l,i,j}$, and write $\theta_{t}$ for $t\geq
0$ to refer to the parameters trained by the gradient flow until time $t$. We
sometimes make the dependency on the width $d$ and the parameters $\theta$
explicit in the MLP, and write $f_{d}$, or $f_{d,\theta}$.

Throughout the paper, we study the training of an MLP $f_{d}$ that regularises
the Jacobian of $f_{d}$. Formally, this means that the training optimises the
following _robust-training_ objective for a given training set
$\mathcal{D}=\\{(x^{(1)},y^{(1)}),\ldots,(x^{(N)},y^{(N)})\\}$:

$\displaystyle\mathcal{L}(\theta)\coloneqq\frac{1}{2N}\sum_{i=1}^{N}\Big{(}$
$\displaystyle(f_{d,\theta}(x^{(i)})-y^{(i)})^{2}\lambda\sum_{\alpha\in[d_{0}]}\big{(}J(f_{d,\theta})(x^{(i)})_{\alpha-1}\big{)}^{2}\Big{)}$
(1)

where $J(f):\mathbb{R}^{d_{0}}\to\mathbb{R}^{d_{0}}$ is the Jacobian of $f$
defined as $J(f)(x)_{\alpha-1}=\frac{\partial f(x)}{\partial x_{\alpha-1}}$,
and $0\leq\lambda\leq 1$ is a hyperparameter determining the importance of the
Jacobian regulariser. By including the Jacobian regulariser, this objective
encourages $f_{d,\theta}$ to change little at each training input $x^{(i)}$
when the input is perturbed, so that when a new input $x$ is close to the
input $x^{(i)}$ of a training example, the learnt predictor tends to return an
output close to $y^{(i)}$ of the example. This encouragement has often been
used to learn a robust predictor in the literature [Hoffman et al., 2019].

To establish our results, we need assumptions for the training dataset
$\mathcal{D}$ and a future input $x^{*}\in\mathbb{R}^{d_{0}}$.

######  Assumption 2 (Dataset Assumption).

(1) The inputs have the unit norm: $\|x^{(i)}\|=1$ for all $i\in[N]$, and
$\|x^{*}\|=1$. (2) The outputs are bounded: $|y^{(i)}|\leq 1$ for all
$i\in[N]$.

##  3 INFINITE-WIDTH LIMIT AT INITIALISATION

We start by describing our results on the infinite-width limit of the Jacobian
of an MLP $f_{d,\theta}$ at _initialisation_. Recall that due to the random
initialisation of the parameters $\theta$, both the MLP and its Jacobian are
random functions at initialisation. It is well-known that as $d$ tends to
infinity, $f_{d,\theta}$ at initialisation converges to a zero-mean Gaussian
process [Lee et al., 2018, Matthews et al., 2018]. The kernel of this GP is
commonly called the NNGP kernel and has an inductive characterisation over the
depth of the MLP. What concerns us here is to answer whether these results
extend to the Jacobian of $f_{d,\theta}$. Does the Jacobian also converge to a
GP? If so, how is the limiting GP of the Jacobian related to the limiting GP
of the network output? Also, in that case, can we characterise the kernel of
the limiting GP of the Jacobian?

The next theorem summarises our answers to these questions.

######  Theorem 3 (GP Convergence at Initialisation).

Suppose that Assumption 2 holds. As $d$ goes to infinity, the function
$x\longmapsto(f_{d}(x),J(f_{d})(x)_{0},\ldots,J(f_{d})(x)_{d_{0}-1})^{\intercal}$
from $\mathbb{R}^{d_{0}}$ to $\mathbb{R}^{d_{0}+1}$ converges weakly in finite
marginal 111A random function $f$ converges weakly in finite marginal, if
every finite evaluation $\\{f(x^{(i)})\\}_{i=1}^{N}$ converges weakly for any
$N<\infty$. to a zero-mean GP with the kernel $\kappa^{2}\Sigma^{(L)}$, which
we call _Jacobian NNGP kernel_ , where
$\Sigma^{(L)}:\mathbb{R}^{d_{0}}\times\mathbb{R}^{d_{0}}\to\mathbb{R}^{(d_{0}+1)\times(d_{0}+1)}$
is defined inductively as follows: for all inputs
$x,x^{\prime}\in\mathbb{R}^{d_{0}}$, layers $l\in[L]$, and indices
$\alpha,\beta\in[d_{0}]$,

$\displaystyle\begin{aligned} \Sigma^{(0)}(x,x^{\prime})_{00}\coloneqq\langle
x,x^{\prime}\rangle,&\quad\Sigma^{(0)}(x,x^{\prime})_{0\beta}\coloneqq
x_{\beta-1},\\\ \Sigma^{(0)}(x,x^{\prime})_{\alpha 0}\coloneqq
x_{\alpha-1}^{\prime},&\quad\Sigma^{(0)}(x,x^{\prime})_{\alpha\beta}\coloneqq\mathbb{I}_{[{\alpha=\beta}]},\end{aligned}$
$\displaystyle\Sigma^{(l)}(x,x^{\prime})_{00}\coloneqq\mathop{\mathbb{E}}_{g}[\phi(g(x)_{0})\cdot\phi(g(x^{\prime})_{0})],$
$\displaystyle\Sigma^{(l)}(x,x^{\prime})_{\alpha\beta}\coloneqq\mathop{\mathbb{E}}_{g}[\dot{\phi}(g(x)_{0})g(x)_{\alpha}\cdot\dot{\phi}(g(x^{\prime})_{0})g(x^{\prime})_{\beta}],$
$\displaystyle\Sigma^{(l)}(x,x^{\prime})_{0\beta}\coloneqq\mathop{\mathbb{E}}_{g}[\phi(g(x)_{0})\cdot\dot{\phi}(g(x^{\prime})_{0})g(x^{\prime})_{\beta}],$
$\displaystyle\Sigma^{(l)}(x,x^{\prime})_{\alpha
0}\coloneqq\mathop{\mathbb{E}}_{g}[\dot{\phi}(g(x)_{0})g(x)_{\alpha}\cdot\phi(g(x^{\prime})_{0})],$

where $\mathbb{I}_{[{-}]}$ is the indicator function, the subscript
${[-]}_{\alpha}$ denotes the $(\alpha+1)$-th component of a vector, and the
random variable $g$ in the expectations is a zero-mean GP with the kernel
$\Sigma^{(l-1)}$.

The $(1,1)$-th entry of the kernel $\kappa^{2}\Sigma^{(L)}(x,x^{\prime})$,
denoted by $\kappa^{2}\Sigma^{(L)}(x,x^{\prime})_{00}$, specifies the
covariance between the outputs of the limiting MLP at $x$ and $x^{\prime}$. It
is precisely the standard NNPG kernel. The other entries describe how the
components of the Jacobian of the limiting MLP are related between themselves
and also with the MLP output.

We prove this theorem using the tensor-program framework [Yang, 2019]. Our
proof translates our convergence question into the one on a tensor program and
uses the so-called Master theorem for tensor programs to derive an answer to
the question. See the supplementary material for the details.

The entries of the Jacobian NNGP kernel can be derived from the $(1,1)$-th
entry $\kappa^{2}\Sigma^{(L)}(x,x^{\prime})_{00}$ via differentiation:

######  Theorem 4.

For the kernel $\Sigma^{(L)}$ in Theorem 3, the following equalities hold for
all $\alpha,\beta\in[d_{0}]$:

$\displaystyle\Sigma^{(L)}(x,x^{\prime})_{0\beta}=\frac{\partial\Sigma^{(L)}(x,x^{\prime})_{00}}{\partial
x^{\prime}_{\beta-1}},$ $\displaystyle\Sigma^{(L)}(x,x^{\prime})_{\alpha
0}=\frac{\partial\Sigma^{(L)}(x,x^{\prime})_{00}}{\partial x_{\alpha-1}},$
$\displaystyle\Sigma^{(L)}(x,x^{\prime})_{\alpha\beta}=\frac{\partial^{2}\Sigma^{(L)}(x,x^{\prime})_{00}}{\partial
x_{\alpha-1}\partial x^{\prime}_{\beta-1}}.$

Thus, the standard result on GPs (Section 9.4 of [Rasmussen and Williams,
2005]) implies that the realisation of the limiting GP defines a
differentiable function in the first component of its output with probability
one, and the Jacobian of this function is also a GP with a kernel induced from
$\Sigma^{(L+1)}(x,x^{\prime})_{00}$ via differentiation.

##  4 JACOBIAN NEURAL TANGENT KERNEL

We next analyse the training dynamics of an MLP under the robust-training
objective in (1), which regularises the Jacobian of the MLP. As in the case of
the standard objective without the Jacobian regulariser, the training dynamics
of a _finite_ MLP is described by a kernel induced by the MLP. The next
definition describes this kernel.

######  Definition 5 (Finite Jacobian NTK).

The _finite Jacobian Neural Tangent Kernel (finite JNTK)_ of a finite MLP
$f_{d}$ with parameters $\theta$ is a function
$\Theta_{d,\theta}:\mathbb{R}^{d_{0}}\times\mathbb{R}^{d_{0}}\to\mathbb{R}^{(d_{0}+1)\times(d_{0}+1)}$
defined as follows: for all $x,x^{\prime}\in\mathbb{R}^{d_{0}}$ and
$\alpha,\beta\in[d_{0}]$,

$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{00}\coloneqq\left\langle\frac{\partial
f_{d}(x)}{\partial\theta},\,\frac{\partial
f_{d}(x^{\prime})}{\partial\theta}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{\alpha\beta}\coloneqq\left\langle\frac{\partial
J(f_{d})(x)_{\alpha-1}}{\partial\theta},\,\frac{\partial
J(f_{d})(x^{\prime})_{\beta-1}}{\partial\theta}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{0\beta}\coloneqq\left\langle\frac{\partial
f_{d}(x)}{\partial\theta},\,\frac{\partial
J(f_{d})(x^{\prime})_{\beta-1}}{\partial\theta}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{\alpha
0}\coloneqq\left\langle\frac{\partial
J(f_{d})(x)_{\alpha-1}}{\partial\theta},\,\frac{\partial
f_{d}(x^{\prime})}{\partial\theta}\right\rangle.$

The finite JNTK includes the standard NTK of an MLP in the $(1,1)$-th entry,
as expected. In addition, it computes the relationship between the gradients
of the $\alpha$ and $\beta$-th components of the Jacobian of the MLP, and also
between the gradient of a component of the Jacobian and the gradient of the
MLP itself.

The next lemma shows that the finite JNTK determines the robust training with
the objective in (1).

######  Lemma 6.

Assume that the parameters of the MLP $f_{d}$ evolve by the continuous version
of the gradient update formalised by the ODE
$\frac{d\theta_{t}}{dt}=-\frac{\partial\mathcal{L}(\theta_{t})}{\partial\theta_{t}}$.
Then for $\alpha\in[d_{0}]$,

$\displaystyle\frac{df_{d,\theta_{t}}(x)}{dt}\
{}={-}\frac{1}{N}\sum_{i=1}^{N}\left\langle\Psi(f_{d,\theta_{t}},x^{(i)},y^{(i)}),\Theta_{d,\theta_{t}}(x,x^{(i)})_{0}\right\rangle,$
(2) $\displaystyle\frac{dJ(f_{d,\theta_{t}})(x)_{\alpha-1}}{dt}\
{}={-}\frac{1}{N}\sum_{i=1}^{N}\left\langle\Psi(f_{d,\theta_{t}},x^{(i)},y^{(i)}),\Theta_{d,\theta_{t}}(x,x^{(i)})_{\alpha-1}\right\rangle,$
(3)

where $\Theta_{d,\theta_{t}}(x,x^{(i)})_{0}$ is the $(d_{0}+1)$-dimensional
vector
$(\Theta_{d,\theta_{t}}(x,x^{(i)})_{00},\ldots,\Theta_{d,\theta_{t}}(x,x^{(i)})_{0d_{0}})^{\intercal}$
obtained from the first row of the kernel output
$\Theta_{d,\theta_{t}}(x,x^{(i)})$,
$\Theta_{d,\theta_{t}}(x,x^{(i)})_{\alpha}$ is defined similarly but from the
$(\alpha+1)$-th row as $(\Theta_{d,\theta_{t}}(x,x^{(i)})_{\alpha
0},\ldots,\Theta_{d,\theta_{t}}(x,x^{(i)})_{\alpha d_{0}})^{\intercal}$, and
$\Psi(f_{d,\theta_{t}},x^{(i)},y^{(i)})$ is the below vector in
$\mathbb{R}^{d_{0}+1}$ constructed as:

$\displaystyle\bigg{(}$ $\displaystyle
f_{d,\theta_{t}}(x^{(i)})-y^{(i)},J(f_{d,\theta_{t}})(x^{(i)})_{0},\ldots,J(f_{d,\theta_{t}})(x^{(i)})_{d_{0}-1}\bigg{)}^{\intercal}.$

Each summand in the ODEs in (2) and (3) is the inner product of two vectors of
size $d_{0}+1$, where only the second vector depends on the input $x$ and this
dependency is via the kernel $\Theta_{d,\theta_{t}}$. A similar form appears
in the analysis of the standard training via the usual NTK. There it is also
proved that if a loss function has a subgradient [Chen et al., 2021], as the
widths of the MLP involved increase, the equation simplifies greatly: the
kernel used in the input-dependent vector stops depending on the initial
values of the parameters (i.e., $\theta_{0}$) and stays constant over time
$t$, and so, the ODE on the MLP becomes a simple linear ODE that can be solved
analytically. We prove that a similar simplification is possible for the
robust training with Jacobian regulariser in the infinite-width limit.

First, we show that the finite JNTK at initialisation becomes deterministic
(i.e., it does not depend on the initial parameter values $\theta_{0}$) in the
infinite-width limit.

######  Theorem 7 (Convergence of Finite JNTK at Initialisation).

Suppose that Assumption 2 holds. As $d$ goes to infinity,
$\Theta_{d,\theta_{0}}$ almost surely converges to a function
$\kappa^{2}\Theta$ in finite marginal, called _limiting JNTK_ , that does not
depend on $\theta_{0}$. Here $\Theta$ is a (deterministic) function of type
$\mathbb{R}^{d_{0}}\times\mathbb{R}^{d_{0}}\to\mathbb{R}^{(d_{0}+1)\times(d_{0}+1)}$,
and has the following form: for $l\in[L]$,
$x,x^{\prime}\in\mathbb{R}^{d_{0}}$, and $\alpha,\beta\in[d_{0}]$,

$\displaystyle\Theta(x,x^{\prime})_{00}$
$\displaystyle\coloneqq\sum_{l=0}^{L}\big{(}\Sigma^{(l)}(x,x^{\prime})_{00}\times\Delta^{(l)}(x,x^{\prime})\big{)},$
$\displaystyle\Theta(x,x^{\prime})_{0\beta}$
$\displaystyle\coloneqq\frac{\partial\Theta(x,x^{\prime})_{00}}{\partial
x_{\beta}^{\prime}},$ $\displaystyle\Theta(x,x^{\prime})_{\alpha 0}$
$\displaystyle\coloneqq\frac{\partial\Theta(x,x^{\prime})_{00}}{\partial
x_{\alpha}},$ $\displaystyle\Theta(x,x^{\prime})_{\alpha\beta}$
$\displaystyle\coloneqq\frac{\partial^{2}\Theta(x,x^{\prime})_{00}}{\partial
x_{\alpha}\partial x_{\beta}^{\prime}},$

where $\Sigma^{(l)}$ is the kernel defined inductively in Theorem 3, and
$\Delta^{(l)}$ is defined with a zero-mean GP $g^{(u)}$ with the kernel
$\Sigma^{(u)}(x,x^{\prime})_{00}$ as follows:
$\Delta^{(l)}(x,x^{\prime})\coloneqq\prod_{u=l}^{L-1}\mathop{\mathbb{E}}_{g^{(u)}}[\dot{\phi}(g^{(u)}(x))\dot{\phi}(g^{(u)}(x^{\prime}))]$.
The explicit form of JNTK can be found in the supplementary material, Section
H.

Next, we show that the finite JNTK stays constant during training in the
infinite-width limit, under the following assumption:

######  Assumption 8 (Full Rank of the Limiting JNTK).

The minimum eigenvalue of the $N(d_{0}+1)\times N(d_{0}+1)$ matrix
$\Theta(x^{(1:N)},x^{(1:N)})$ is greater than $0$.

Note that the assumption is stated with $\Theta$, which is independent of the
random initialisation of parameters $\theta$ and, thus, deterministic. In
practice, we can test this assumption computationally, as we did in our
experiments shown later.

The following result is a counterpart of the standard result on the change of
the NTK [Jacot et al., 2018, Lee et al., 2019] during training without our
Jacobian regulariser. It lets us solve the ODEs (2) and (3) and derive their
analytic solutions in the infinite-width limit.

######  Theorem 9 (Constancy of Finite JNTK during Training).

Suppose that Assumptions 2 and 8 hold. Let $\theta_{t}$ be the parameter of
the MLP at time $t$, trained with gradient flow on the dataset
$\mathcal{D}=(x^{(1:N)},y^{(1:N)})$. Then, there exists a function $F$, not
depending on the width and the parameters of the MLP, such that for all
$\epsilon,\delta\in(0,1)$, if the MLP width satisfies $d\geq
F(x^{(1:N)},L,\epsilon,\delta)$, then with probability at least $1-\delta$
over the randomness of the MLP initialisation, the finite JNTK stays near the
corresponding limiting JNTK during training: for all $t\geq 0$ and all
$i,j\in[N]$,

$\left\|(1/\kappa^{2})\Theta_{d,\theta_{t}}(x^{(i)},x^{(j)})-\Theta(x^{(i)},x^{(j)})\right\|_{F}\leq\epsilon.$

We prove the theorem 9 in the supplementary material, Section K. The structure
of the proof is similar to that of the standard results. We first show that if
the parameters at time $t$ stay close to their initialisation, then the finite
JNTK also does. Then, we prove that the gradient flow does not move the
parameters far away from their initialisation. Unlike in the proof of the
standard results, when we upper-bound the movement of the parameter matrix
$W^{(l)}_{t}$ at each layer $l$ and time $t$ from its initialisation
$W^{(l)}_{0}$, we consider all of $1$, $2$, and $\infty$ matrix norms; the
standard proof considers only the $2$ norm. This more-precise bookkeeping is
needed because of the Jacobian regulariser in our setting. Now, Theorem 9
makes the ODEs in (2) and (3) the linear first-order ODEs, which allow us to
derive their analytic solutions at $t=\infty$ for a given initial condition.

We describe the solution of the ODE in (2) at time $t=\infty$. Recall that
$\lambda$ is the coefficient of the gradient regulariser in the robust-
training objective (Equation (1)). Let
$\Theta_{\lambda}:\mathbb{R}^{d_{0}}\times\mathbb{R}^{d_{0}}\to\mathbb{R}^{(d_{0}+1)\times(d_{0}+1)}$
be the following function: for all $x,x^{\prime}\in\mathbb{R}^{d_{0}}$,
$\Theta_{\lambda}(x,x^{\prime})\coloneqq\Lambda\Theta(x,x^{\prime})\Lambda$,
where $\Lambda$ is the ${(d_{0}+1)\times(d_{0}+1)}$ diagonal matrix with
$\Lambda_{00}=1$ and $\Lambda_{ii}=\sqrt{\lambda}$ for all $i\neq 0$. Using
$\Theta_{\lambda}$, we define key players in our theorem on the ODE solution.
The first is the matrix
$\Theta_{\lambda}(x^{(1:N)},x^{(1:N)})\in\mathbb{R}^{N(d_{0}+1)\times
N(d_{0}+1)}$ that is constructed by including every
$\Theta_{\lambda}(x^{(i)},x^{(j)})\in\mathbb{R}^{(d_{0}+1)\times(d_{0}+1)}$ as
the $(i,j)$-th submatrix. The next players are the vectors
$\mathbf{y}^{(1:N)}\in\mathbb{R}^{N(d_{0}+1)}$ and
$\Theta_{\lambda}(x,x^{(1:N)})_{0}\in\mathbb{R}^{N(d_{0}+1)}$ defined by
stacking the below vectors: for all $i\in[N]$ and $x\in\mathbb{R}^{d_{0}}$,

$\displaystyle\mathbf{y}^{(i)}$
$\displaystyle\coloneqq(y_{i},0,\ldots,0)^{\intercal}\in\mathbb{R}^{1+d_{0}},$
$\displaystyle\Theta_{\lambda}(x,x^{(i)})_{0}$
$\displaystyle\coloneqq(\Theta_{\lambda}(x^{*},x^{(i)})_{00},\ldots,\Theta_{\lambda}(x^{*},x^{(i)})_{0d_{0}})^{\intercal}.$

The last player is the following real-valued function $f_{\mathrm{ntk}}$ on
$\mathbb{R}^{d_{0}}$, which solves the infinite-width version of the ODE in
(2) at time $t=\infty$: for $x\in\mathbb{R}^{d_{0}}$,

$\displaystyle
f_{\mathrm{ntk}}(x)={\Theta_{\lambda}(x,x^{(1:N)})_{0}^{\intercal}\Theta_{\lambda}(x^{(1:N)},x^{(1:N)})^{-1}\mathbf{y}^{(1:N)}.}$
(4)

Using what we have defined, we present our result on the solution of the ODE
in (2) at time $t=\infty$.

######  Theorem 10 (MLPs Learnt by Robust Training).

Assume that Assumptions 2 and 8 hold. Then, there exist functions $F$ and $G$
such that for all $\epsilon,\delta>0$ and $x^{*}\in\mathbb{R}^{d_{0}}$ with
$\|x^{*}\|=1$, and also for all $d$ and $\kappa$, if $d\geq
F(x^{(1:N)},x^{*},L,\delta,\epsilon)$ and $\kappa\leq
G(x^{(1:N)},x^{*},L,\delta,\epsilon)$, then the following statement holds with
probability at least $1-\delta$. If $f_{d,\theta_{\infty}}(x^{*})$ is the
solution of the ODE in (2) at time $t=\infty$, we have
$|f_{d,\theta_{\infty}}(x^{*})-f_{\mathrm{ntk}}(x^{*})|\leq\epsilon$ and the
training loss converges to $0$ exponentially with respect to $t$.

This theorem shows that in the infinite-width limit, the fully-trained MLP
converges to $f_{\mathrm{ntk}}$, which has the form of the standard kernel
regressor.

The functions $F,G$ are introduced to implicitly argue the bounds. The first
function $F$, which lower bounds the width of network $d$ requires the network
to be large enough to enable the infinite width behaviours. The second
function $G$ upper bounds the scale of the output of the network with
$\kappa$, so that the initial randomness is suppressed, therefore the training
loss at $t=0$ can be bounded.

We also extend the previous that a similar result holds for gradient descent
with sufficiently small learning rate $\eta$, omitting the convergence towards
kernel regression solution.

######  Theorem 11.

Assume that Assumptions 2 and 8 hold. Then, there exist functions $F,G$ and
$H$ such that for all $\delta>0$ and $d,\kappa,\eta$, if $d\geq
F(x^{(1:N)},L,\delta)$, $\kappa\leq G(x^{(1:N)},L,\delta)$, and $\eta\leq
H(\lambda,x^{(1:N)},L)$, then the training loss of gradient descent with
learning rate $\eta$ converges to $0$ exponentially with respect to $t$ with
probability at least $1-\delta$.

Here we include one more implicit bound $H$, which requires the learning rate
to be small enough to guarantee global convergence.

##  5 EXPERIMENTS

We experimentally checked our results using finite MLPs with varying widths
and depths and simple datasets, analysed the robustness of the kernel
regression solution obtained, and empirically investigated the validity of
Assumption 8. In this section, we report our findings from those experiments.
We repeated the experiments 10 times and plotted their results with 95%
confidence regions.

###  5.1 Validation Of Main Results

We tested our three results: (i) an MLP and its Jacobian jointly converge to a
GP (Theorem 3); (ii) the finite JNTK of an MLP converges to a deterministic
kernel at initialisation (Theorem 7); (iii) the finite JNTK of an MLP stays
constant throughout robust training with high probability, which grows to $1$
as the MLP gets wider (Theorem 9).

For the first two convergence results (Theorems 3 and 7), we used the simple
synthetic dataset in $\mathbb{R}^{4}$ of size 256, which approximates the
optimal $\epsilon$-net of the $S^{3}$ where $y^{(i)}$ is chosen arbitrarily
since it does not matter for these results. In these tests, we tried MLPs with
all the width-depth combinations in $\\{2^{i}\\}_{i=6}^{13}\times\\{1,2,3\\}$.
This dataset is constructed to cover the possible angles between two inputs,
and the standard basis. For the third result regarding the evolution of a JNTK
during robust training (Theorem 9), we used the Algerian forest fire dataset
[Abid and Izeboudjen, 2020] from the UCI Machine Learning repository, which
contains $224$ data points with input dimension $11$. We fixed the depth of
the network as 11 which was the smallest depth that satisfies Assumption 8
with GeLU activation, for varying width $\\{2^{i}\\}_{i=6}^{13}$. In all of
these tests, we set $\kappa$ to $0.1$, the value that matters only for the
third test and made the training of the test achieve reasonable performance.
More details on the experiments can be found in Section A.

Figure 1: $\max$-norm distance between the Jacobian NNGP kernel and the
estimate of its finite counterpart. The $x$-axis represents the width of an
MLP, and the $y$-axis the $\max$-norm distance.

Figure 2: $\max$-norm distance between the $(1/\kappa^{2})$-scaled finite JNTK
$(1/\kappa^{2})\Theta_{d,\theta_{0}}$ and the similarly-scaled limiting JNTK
$\Theta$ at initialisation. The $x$-axis represents the width of an MLP, and
the $y$-axis the $\max$-norm distance.

For the validation of Theorem 3, we focused on the convergence to
$\Sigma^{(L)}$, the Jacobian NNGP kernel scaled down by $\kappa^{2}$. This
does not directly validate the theorem since it does not verify the
Gaussianity of the output. But this is an effective proxy and the convergence
to $\Sigma^{(L)}$ is critical in the proof of Theorem 7. We estimated the
finite-width counterpart of the Jacobian NNGP kernel of an MLP $f_{d,\theta}$
written simply as $f$ below: for all $i,j\in[N]$,

$\widehat{\Sigma}^{(L)}(x^{(i)},x^{(j)})=\widehat{\mathrm{Cov}}\left(\begin{matrix}f(x^{(i)})\\\
J(f)(x^{(i)})\end{matrix},\begin{matrix}f(x^{(j)})\\\
J(f)(x^{(j)})\end{matrix}\right)$

where $\widehat{\mathrm{Cov}}$ means a Monte-Carlo estimate of covariance with
one million samples. Then, we measured the $\max$-norm distance between the
Gram matrix of the $(\alpha,\beta)$-th component of the Jacobian NNGP kernel
and that of its finite estimated counterpart with $1/\kappa^{2}$ scaling, for
all $\alpha,\beta\in\\{0,1,2,3,4\\}$. We denote the measured quantity by
$\Delta\Sigma_{\alpha\beta}^{(L)}$, which is formally defined as shown below:

$\displaystyle\max_{i,j\in[64]}\left|\frac{1}{\kappa^{2}}\widehat{\Sigma}^{(L)}(x^{(i)},x^{(j)})_{\alpha\beta}-\Sigma^{(L)}(x^{(i)},x^{(j)})_{\alpha\beta}\right|.$

Figure 2 visualises $\Delta\Sigma_{\alpha\beta}^{(L)}$ for four choices of
$(\alpha,\beta)$ representing four types of covariance: between the outputs
$(00)$, between the Jacobians with respect to a fixed input coordinate $(22)$,
between the output and the Jacobian $(02)$, and between the Jacobians with
respect to different input coordinates $(12)$. As the width of the network
increases, we can see that $\Delta\Sigma_{\alpha\beta}^{(L)}$ decreases. (We
do not expect that even as we increase the width indefinitely, the graphs
converge to $0$ because of the Monte-Carlo errors in the computation of
$\widehat{\Sigma}^{(L)}$.) This validates our convergence result.

The validation of the next Theorem 7 similarly focused on the convergence to
$\Theta$, the limiting JNTK scaled down by $\kappa^{2}$. We computed the
$\max$-norm distance between the Gram matrix of the $(\alpha,\beta)$-th
component of the $(1/\kappa^{2})$-scaled finite JNTK and that of the limiting
JNTK for $\alpha,\beta\in\\{0,1,2,3,4\\}$:

$\displaystyle\max_{i,j\in[64]}\left|\frac{1}{\kappa^{2}}\Theta_{d,\theta_{0}}(x^{(i)},x^{(j)})_{\alpha\beta}-\Theta(x^{(i)},x^{(j)})_{\alpha\beta}\right|.$

Figure 2 shows the results of these computations. We can see that as the
network width increases, the $(1/\kappa^{2})$-scaled finite NTK converges to
$\Theta$ at initialisation, which confirms Theorem 7. Note that
$\Delta\Theta_{22}$ is much larger than $\Delta\Theta_{00}$’s. This is because
while both the former and the latter are defined as sums of terms, the former
has four times more summands than the latter.

Finally, we validated the claim of Theorem 9 that the JNTK stays constant
during robust training. We used a binary classification dataset but regarded
it as a regression dataset with targets in $\\{+1,-1\\}$, and trained the
network under the regression loss with Jacobian regularisation coefficient
$\lambda=0.01$. We mimicked the gradient flow in the theorem with gradient
descent with learning rate $1$, and measured the following $\max$-norm
distance $\Delta\Theta_{\alpha\beta,t}$ at the training steps
$t=2^{0},2^{1},\ldots,2^{11}$: for $\alpha,\beta\in\\{0,1,\ldots,11\\}$,

$\max_{i,j\in[N]}\left|\frac{1}{\kappa^{2}}\Theta_{d,\theta_{t}}(x^{(i)},x^{(j)})_{\alpha\beta}-\Theta(x^{(i)},x^{(j)})_{\alpha\beta}\right|.$

Figure 3: $\max$-norm distance between the $(1/\kappa^{2})$-scaled finite JNTK
$(1/\kappa^{2})\Theta_{d,\theta_{t}}$ at the training step $t$, and the
similarly-scaled limiting JNTK $\Theta$ at initialisation. The x-axis
represents the width of an MLP, and the y-axis represents $\max$-norm
distance.

Figure 3 shows the measured $\max$-norm distance for the four cases of
$(\alpha,\beta)\in\\{(2,2),(0,2),(0,0),(1,2)\\}$. As the network gets wider,
the distance gets smaller among all training steps $t$ which validates our
claim. Note that as the training step $t$ increases, the distance also
increases, but the increment becomes smaller, which indicates that the JNTK
changes less as the training progresses.

###  5.2 Analysis Of Kernel Regression

As stated in Theorem 10, in the regime of large width and small learning rate,
the neural network under robust training converges to the solution
$f_{\mathrm{ntk}}$ of a kernel-regression problem, as the training time goes
to infinity. Following [Tsilivis and Kempe, 2022], we analysed the precision
and robustness of this kernel-regression solution by looking at the
eigenvalues and eigenvectors of the matrix
$\Theta_{\lambda}(x^{(1:N)},x^{(1:N)})$ in the solution. We used an 11-layer
neural network with the GeLU activation, where the depth is chosen to be the
smallest among those that satisfy the Assumption 8 in our experiments. For
each $i\in[N(d_{0}+1)]$, let
$f_{i}(x)\coloneqq\Theta_{\lambda}(x,x^{(1:N)})_{0}^{\intercal}\lambda_{i}^{-1}v_{i}v_{i}^{\intercal}\mathbf{y}^{(1:N)}$,
where $\lambda_{i}$ and $v_{i}$ are the $i$-th eigenvalue (ordered from the
largest) and the corresponding eigenvector of the matrix
$\Theta_{\lambda}(x^{(1:N)},x^{(1:N)})$. Call $f_{i}$ eigenfeature. Then, by
Equation 4, $f_{\mathrm{ntk}}(x)=\sum_{i\in[N(d_{0}+1)]}f_{i}(x)$. We
interpret this as $f_{\mathrm{ntk}}$ being a linear combination of
eigenfeatures $f_{i}$, and analyse the precision and robustness of these
eigenfeatures.

Figure 4: Pairplot of test accuracy, test accuracy after 0.01 perturbation,
test accuracy after 0.1 perturbation of eigenfeatures of robust training. The
colours denote the coefficient of Jacobian regularisation.

Figure 5: Pairplot of test accuracy, test accuracy after 0.01 perturbation,
test accuracy after 0.1 perturbation of eigenfeatures of standard training.
The ellipses correspond to 4$\sigma$ confidence region of multivariate
Gaussian distribution fitting points, to show the correlation clearly.

Figure 5 shows the pairplot of the following three measures on these
eigenfeatures $f_{i}$: test accuracy of $f_{i}$, test accuracy of $f_{i}$
after the $0.01$ perturbation of the test inputs, and test accuracy of $f_{i}$
with respect to the $0.1$ perturbation. The perturbation was obtained by
changing the inputs via one-step projected gradient descent on the analytic
solution without Jacobian regularisation, and $0.01$ and $0.1$ refer to
learning rates. The 0.1 perturbation was enough to change the labels of some
inputs, while for the 0.01 perturbation, there exists a classifier that
classifies all the inputs perfectly even after this perturbation. In the plot,
we denote eigenfeatures by their eigenvalues and use the colours to represent
the Jacobian-regularisation coefficient from $\lambda=1$ (Purple) to
$\lambda=10^{-3}$ (Red). The result on standard training in Figure 5 shows a
distinct result, where it shows less or no correlation between accuracy and
robustness.

While the empirical analysis of NTK [Tsilivis and Kempe, 2022] found, in the
setting of standard training without Jacobian regularisation, eigenfeatures
that are highly accurate but less robust, or highly robust but less accurate,
we could not find such eigenfeatures in our setting. We observed that adding
Jacobian regularisation promotes alignment between accuracy and robustness so
that achieving high test accuracy implies satisfying robustness.

Figure 6: Smallest eigenvalue of network with varying depth with GeLU or erf
activation, for both JNTK and standard NTK. The red cross denotes depths that
JNTK have the positive smallest eigenvalue.

###  5.3 Empirical Analysis Of Assumption 8

Assumption 8 is a standard assumption in the NTK analysis; it is known that
the standard NTK $\Theta(x,x^{\prime})_{00}$ has its smallest eigenvalue
bounded away from zero if no inputs are parallel and the activation function
is smooth [Du et al., 2019a]. However, for our JNTK, it is not known to be
true theoretically. Our empirical analysis suggests that whether the
assumption holds or not depends on activation function, depth, and even the
dataset itself.

To check when Assumption 8 holds, we empirically analysed the smallest
eigenvalue of the JNTK Gram matrix for different depths and activation
functions. The results are summarised in Figure 6. The blue-dot curves show
the smallest eigenvalues of the NTK Gram matrices, which are the matrices
built by only using $\Theta(x,x^{\prime})_{00}$ terms. Note that the plot is
always positive. The green curves show the smallest eigenvalues of the JNTK
Gram matrices. Note that the smallest eigenvalues in the JNTK case are much
smaller than those in the NTK case. Also, to be greater than $0$, it requires
deeper networks, for example, $11$ in the GeLU-activation case and $6$ in the
erf-activation case. We conjecture that this requirement on depth comes from
the expanded size of our Gram matrix (that is, $(N(d_{0}+1))^{2}$ instead of
$N^{2}$), and also from Jacobian regularisation’s forcing the network to be
flat around training examples and thereby reducing its flexibility. This
conjecture is consistent with the finding of [Poole et al., 2016] that the
complexity of the network increases exponentially as the depth increases.

##  6 RELATED WORKS AND CONCLUSION

Jacobian regularisation was introduced in several forms including [Hoffman et
al., 2019, Jakubovitz and Giryes, 2018], with different functions being
differentiated. While being simple, it has been a strong baseline for defence
against adversarial attacks [Jakubovitz and Giryes, 2018], with several
extensions. The reason for the success of Jacobian regularisation has not been
well understood theoretically. Only recently the implicit bias of the finite
difference version of Jacobian regularisation was studied for linear models
[Karakida et al., 2023], which favours an L1 regularisation solution.

In this paper, we provided a theoretical analysis of robust training with
Jacobian regularisation from the perspective of infinite-width limits. We
characterised the limiting NNGP and NTK kernels of feed-forward neural
networks trained under Jacobian regularisation and showed that with high
probability, such networks converge to solutions of certain kernel regression
problems, as the widths of those networks are large enough. We used these
theoretical results to analyse the precision and robustness of the networks
trained with Jacobian regularisation, and also experimentally demonstrated the
validity of these results.

## References

  * [Abid and Izeboudjen, 2020] Abid, F. and Izeboudjen, N. (2020).  Predicting forest fire in algeria using data mining techniques: Case study of the decision tree algorithm.  In Ezziyyani, M., editor, Advanced Intelligent Systems for Sustainable Development (AI2SD’2019), pages 363–370, Cham. Springer International Publishing.
  * [Arora et al., 2019] Arora, S., Du, S. S., Hu, W., Li, Z., Salakhutdinov, R. R., and Wang, R. (2019).  On exact computation with an infinitely wide neural net.  In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
  * [Chen et al., 2021] Chen, Y., Huang, W., Nguyen, L., and Weng, T.-W. (2021).  On the equivalence between neural network and support vector machine.  In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems, volume 34, pages 23478–23490. Curran Associates, Inc.
  * [Du et al., 2019a] Du, S., Lee, J., Li, H., Wang, L., and Zhai, X. (2019a).  Gradient descent finds global minima of deep neural networks.  In Chaudhuri, K. and Salakhutdinov, R., editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1675–1685. PMLR.
  * [Du et al., 2019b] Du, S., Lee, J., Li, H., Wang, L., and Zhai, X. (2019b).  Gradient descent finds global minima of deep neural networks.  In Chaudhuri, K. and Salakhutdinov, R., editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1675–1685. PMLR.
  * [Goodfellow et al., 2015] Goodfellow, I. J., Shlens, J., and Szegedy, C. (2015).  Explaining and harnessing adversarial examples.  In Bengio, Y. and LeCun, Y., editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
  * [Hoffman et al., 2019] Hoffman, J., Roberts, D. A., and Yaida, S. (2019).  Robust learning with jacobian regularization.
  * [Hu et al., 2020] Hu, W., Li, Z., and Yu, D. (2020).  Simple and effective regularization methods for training on noisily labeled data with generalization guarantee.  In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
  * [Jacot et al., 2018] Jacot, A., Gabriel, F., and Hongler, C. (2018).  Neural tangent kernel: Convergence and generalization in neural networks.  In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.
  * [Jakubovitz and Giryes, 2018] Jakubovitz, D. and Giryes, R. (2018).  Improving dnn robustness to adversarial attacks using jacobian regularization.  In Ferrari, V., Hebert, M., Sminchisescu, C., and Weiss, Y., editors, Computer Vision – ECCV 2018, pages 525–541, Cham. Springer International Publishing.
  * [Karakida et al., 2023] Karakida, R., Takase, T., Hayase, T., and Osawa, K. (2023).  Understanding gradient regularization in deep learning: Efficient finite-difference computation and implicit bias.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J., editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 15809–15827. PMLR.
  * [Lee et al., 2022] Lee, H., Yun, E., Yang, H., and Lee, J. (2022).  Scale mixtures of neural network gaussian processes.  In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net.
  * [Lee et al., 2018] Lee, J., Bahri, Y., Novak, R., Schoenholz, S. S., Pennington, J., and Sohl-Dickstein, J. (2018).  Deep neural networks as gaussian processes.  In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings. OpenReview.net.
  * [Lee et al., 2019] Lee, J., Xiao, L., Schoenholz, S., Bahri, Y., Novak, R., Sohl-Dickstein, J., and Pennington, J. (2019).  Wide neural networks of any depth evolve as linear models under gradient descent.  In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
  * [Liu and Zenke, 2020] Liu, T. and Zenke, F. (2020).  Finding trainable sparse networks through neural tangent transfer.  In Proceedings of the 37th International Conference on Machine Learning (ICML’20), pages 6336–6347.
  * [Lohweg, 2013] Lohweg, V. (2013).  banknote authentication.  UCI Machine Learning Repository.  DOI: https://doi.org/10.24432/C55P57.
  * [Matthews et al., 2018] Matthews, A. G., Hron, J., Rowland, M., Turner, R. E., and Ghahramani, Z. (2018).  Gaussian process behaviour in wide deep neural networks.  In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings. OpenReview.net.
  * [Neal, 1996] Neal, R. M. (1996).  Priors for infinite networks.  In Bayesian Learning for Neural Networks, pages 29–53. Springer New York.
  * [Nguyen et al., 2021] Nguyen, Q., Mondelli, M., and Montufar, G. F. (2021).  Tight bounds on the smallest eigenvalue of the neural tangent kernel for deep relu networks.  In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8119–8129. PMLR.
  * [Peck et al., 2017] Peck, J., Roels, J., Goossens, B., and Saeys, Y. (2017).  Lower bounds on the robustness to adversarial perturbations.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
  * [Peluchetti et al., 2020] Peluchetti, S., Favaro, S., and Fortini, S. (2020).  Stable behaviour of infinitely wide deep neural networks.  In Chiappa, S. and Calandra, R., editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 1137–1146. PMLR.
  * [Poole et al., 2016] Poole, B., Lahiri, S., Raghu, M., Sohl-Dickstein, J., and Ganguli, S. (2016).  Exponential expressivity in deep neural networks through transient chaos.  In Lee, D., Sugiyama, M., Luxburg, U., Guyon, I., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc.
  * [Rasmussen and Williams, 2005] Rasmussen, C. E. and Williams, C. K. I. (2005).  Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning).  The MIT Press.
  * [Sejnowski and Gorman, 2017] Sejnowski, T. and Gorman, R. (2017).  Connectionist Bench (Sonar, Mines vs. Rocks).  UCI Machine Learning Repository.  DOI: https://doi.org/10.24432/C5T01Q.
  * [Tsilivis and Kempe, 2022] Tsilivis, N. and Kempe, J. (2022).  What can the neural tangent kernel tell us about adversarial robustness?  In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A., editors, Advances in Neural Information Processing Systems, volume 35, pages 18116–18130. Curran Associates, Inc.
  * [Vershynin, 2018] Vershynin, R. (2018).  High-Dimensional Probability: An Introduction with Applications in Data Science.  Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press.
  * [Wang et al., 2021] Wang, J., Chen, J., Sun, Y., Ma, X., Wang, D., Sun, J., and Cheng, P. (2021).  RobOT: Robustness-Oriented Testing for Deep Learning Systems, page 300–311.  IEEE Press.
  * [Yang, 2019] Yang, G. (2019).  Wide feedforward or recurrent neural networks of any architecture are gaussian processes.  In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
  * [Yang, 2020] Yang, G. (2020).  Tensor programs ii: Neural tangent kernel for any architecture.
  * [Yang et al., 2021] Yang, G., Hu, E., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J. (2021).  Tuning large neural networks via zero-shot hyperparameter transfer.  In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems, volume 34, pages 17084–17097. Curran Associates, Inc.
  * [Yang and Hu, 2021] Yang, G. and Hu, E. J. (2021).  Tensor programs iv: Feature learning in infinite-width neural networks.  In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11727–11737. PMLR.
  * [Yang and Littwin, 2021] Yang, G. and Littwin, E. (2021).  Tensor programs iib: Architectural universality of neural tangent kernel training dynamics.  In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11762–11772. PMLR.
  * [Zhang and Wei, 2022] Zhang, H. and Wei, H. (2022).  Sharper sub-weibull concentrations.  Mathematics, 10(13).
  * [Zhang et al., 2019] Zhang, H., Zhang, P., and Hsieh, C. (2019).  Recurjac: An efficient recursive algorithm for bounding jacobian matrix of neural networks and its applications.  In The Thirty-Third AAAI Conference on Artificial Intelligence, AAAI 2019, Honolulu, Hawaii, USA, January 27 - February 1, 2019, pages 5757–5764. AAAI Press.

###### Contents

  1. 1 INTRODUCTION
  2. 2 SETUP
  3. 3 INFINITE-WIDTH LIMIT AT INITIALISATION
  4. 4 JACOBIAN NEURAL TANGENT KERNEL
  5. 5 EXPERIMENTS
    1. 5.1 Validation Of Main Results
    2. 5.2 Analysis Of Kernel Regression
    3. 5.3 Empirical Analysis Of Assumption 8
  6. 6 RELATED WORKS AND CONCLUSION
  7. A EXPERIMENT DETAIL
  8. B ADDITIONAL EXPERIMENTS
  9. C ADDITIONAL NOTATIONS
  10. D REVIEW OF THE TENSOR-PROGRAM FRAMEWORK
  11. E PROOF OF THEOREM 3
  12. F PROOF OF THEOREM 4
  13. G REMARK ON ASSUMPTION 8
  14. H PROOF OF THEOREM 7
  15. I FINITE JNTK AFTER PERTURBATION
    1. I.1 Bounds at initialisation
    2. I.2 Bounds after perturbation
    3. I.3 Proof of Theorem 23
  16. J PROOF OF THEOREM 9
  17. K PROOF OF THEOREM 10
  18. L PROOF OF THEOREM 11

##  Appendix A EXPERIMENT DETAIL

For the GP convergence and convergence of finite JNTK at initialisation, we
utilised a synthetic dataset generated using the Fibonacci Lattice algorithm
extended to the hypersphere, which approximates the 0.5 $\epsilon$-net of the
$S^{3}$. Regarding the evolution of the finite JNTK during robust training, we
set the values of $\lambda$ to 0.01, $\kappa$ to 0.1, and the learning rate to
1. These specific values were selected to ensure that the test accuracy
approached 100% by the end of 2048 epochs.

All the natural datasets first have their entries scaled to be contained in
[-1, 1], then normalised to have a unit norm to match our assumption. We
include two additional datasets in the supplementary materials, the banknote
authentication dataset [Lohweg, 2013] which is 4-dimensional data with 1372
data points, and the connectionist bench dataset [Sejnowski and Gorman, 2017]
which is 60-dimensional data with 208 data points. The banknote authentication
dataset [Lohweg, 2013] contains several data points that are nearly parallel,
so the standard NTK kernel becomes nearly singular, making the comparison
impossible. To mitigate this, we filtered out some data points that are nearly
parallel. In specific, for GeLU activation, we removed data points $x^{(i)}$
such that $|\langle x^{(i)},x^{(j)}\rangle|>0.99$ for some $j<i$, and 0.995
for erf activation, each resulting 186, 320 data points respectively.

For our experiments, we utilised 3 NVIDIA RTX 6000 GPUs to validate the
convergence of the GP over three days. Additionally, we employed 3 NVIDIA RTX
A5000 GPUs to verify the constancy of finite JNTK during training for one day.
All other experiments were performed on CPUs, each completed in under an hour.

##  Appendix B ADDITIONAL EXPERIMENTS

In this section, we describe additional experiments on the minimum eigenvalue
assumption and also on the kernel regression solution analysis. We repeat the
same experiments in Figures 4 and 6 but with different datasets [Lohweg, 2013,
Sejnowski and Gorman, 2017].

Figure B7: The minimum eigenvalues of the standard NTK kernel and the JNTK
kernel under the GeLU and erf activation functions on the banknote
authentication dataset. The x-axis is the depth of the network, and the y-axis
is the minimum eigenvalue. The x-marks indicate that the nonzero minimum
eigenvalue assumption is satisfied, and the o-marks indicate that the
assumption is violated. We do not mark the points for the standard NTK kernel
because the assumption is always satisfied for the standard NTK kernel. Figure
B8: The minimum eigenvalues of the standard NTK kernel and the JNTK kernel
under the GeLU and erf activation functions on the connectionist bench
dataset. The x-axis is the depth of the network, and the y-axis is the minimum
eigenvalue. The x-marks indicate that the nonzero minimum eigenvalue
assumption is satisfied, and the o-marks indicate that the assumption is
violated. We do not mark the points for the standard NTK kernel because the
assumption is always satisfied for the standard NTK kernel.

Figures B7 and B8 show the minimum eigenvalues of the standard NTK kernel and
the JNTK kernel under the GeLU and erf activation functions on the banknote
authentication dataset and the connectionist bench dataset, respectively. Both
plots show a similar tendency as Figure 6 in the main paper. Note that in the
connectionist bench dataset, the assumption is satisfied for lower-depth
networks, being invalidated only for shallow GeLU networks. This can be
explained by the fact that the connectionist bench dataset has higher
dimensions than the banknote authentication and forest fire datasets.

We give the kernel regression analysis results for varying datasets and
activations in Figures B9, B10, B11, B12, and B13. We use the smallest depth
network that satisfies the minimum eigenvalue assumption for each dataset and
activation, which can be found in Figures B7 and B8, and also in Figure 6 in
the main text.

Figure B9: Kernel regression analysis on forest fire dataset, and erf
activation. The left figure is the result without Jacobian regularisation, and
the right figure is the result with Jacobian regularisation.

Figure B10: Kernel regression analysis on the banknote authentication dataset
under the GeLU activation. The left figure is the result without Jacobian
regularisation, and the right figure is the result with Jacobian
regularisation.

Figure B11: Kernel regression analysis on the banknote authentication dataset
under the erf activation. The left figure is the result without Jacobian
regularisation, and the right figure is the result with Jacobian
regularisation.

Figure B12: Kernel regression analysis on the connectionist bench dataset
under the GeLU activation. The left figure is the result without Jacobian
regularisation, and the right figure is the result with Jacobian
regularisation.

Figure B13: Kernel regression analysis on the connectionist bench dataset
under the erf activation. The left figure is the result without Jacobian
regularisation, and the right figure is the result with Jacobian
regularisation.

##  Appendix C ADDITIONAL NOTATIONS

We introduce symbols $h^{(0)}(x)$ and $g^{(l)}(x)$ for $l=1,\ldots,L+1$ and
$x\in\mathbb{R}^{d_{0}}$, and let them denote the input $x$ and the pre-
activation values at layers $l$. Thus, the MLP
$f:\mathbb{R}^{d_{0}}\to\mathbb{R}^{d_{L+1}}$ is defined as follows: for all
inputs $x\in\mathbb{R}^{d_{0}}$ and all layers $l=2,\ldots,L$,

$\displaystyle h^{(0)}(x)\in\mathbb{R}^{d_{0}},$ $\displaystyle
h^{(0)}(x)\coloneqq x,$ $\displaystyle g^{(1)}(x)\in\mathbb{R}^{d_{1}}$
$\displaystyle g^{(1)}(x)\coloneqq W^{(1)}x,$ $\displaystyle
h^{(1)}(x)\in\mathbb{R}^{d_{1}},$ $\displaystyle
h^{(1)}(x)\coloneqq\phi(g^{(1)}(x)),$ $\displaystyle
g^{(l)}(x)\in\mathbb{R}^{d_{l}},$ $\displaystyle
g^{(l)}(x)\coloneqq\frac{1}{\sqrt{d}}W^{(l)}h^{(l-1)}(x),$ $\displaystyle
h^{(l)}(x)\in\mathbb{R}^{d_{l}},$ $\displaystyle
h^{(l)}(x)\coloneqq\phi(g^{(l)}(x)),$ $\displaystyle
g^{(L+1)}(x)\in\mathbb{R}^{d_{L+1}},$ $\displaystyle
g^{(L+1)}(x)\coloneqq\frac{1}{\sqrt{d}}W^{(L+1)}h^{(L)}(x),$ $\displaystyle
f(x)\in\mathbb{R}^{d_{L+1}},$ $\displaystyle f(x)\coloneqq\kappa
g^{(L+1)}(x).$

Recall our assumption that $d=d_{1}=\ldots=d_{L}$ and $d_{L+1}=1$.

We use the following constant $C_{l}$, especially for simplifying the
multipliers in the neural network:

$C_{l}\coloneqq\begin{cases}1&l=0,\\\ d&l\geq 1.\end{cases}$

We write $G_{C}(n)$ for the sum of the geometric series
$C^{0},C^{1},\ldots,C^{n}$, i.e.,

$G_{C}(n)\coloneqq\sum_{i=0}^{n}C^{i}.$

In the supplementary materials, we consider norms on random variables, such as
the sub-Gaussian norm $\|\ \cdot\ \|_{SG}$ and the sub-exponential norm $\|\
\cdot\ \|_{SE}$.

For the sums $\sum_{i=a}^{b}$ and products $\prod_{i=a}^{b}$, we use empty sum
/ product convention when $a>b$,

$\sum_{i=a}^{b}({}\cdots{})=0,\qquad\prod_{i=a}^{b}({}\cdots{})=1.$

##  Appendix D REVIEW OF THE TENSOR-PROGRAM FRAMEWORK

We quickly review the tensor-program framework by Greg Yang [Yang, 2019, Yang,
2020]. To simplify the presentation, our review describes only a simplified
version of the framework where all the hidden layers have the same width $n$.
For the full version where the hidden layers have different widths (but these
widths are sent to infinity with a fixed ratio), see the original papers of
the framework [Yang, 2019, Yang, 2020].

_Tensor programs_ are a particular type of straight-line programs that
represent computations involving neural networks, such as forward computation
and backpropagation. Each tensor program consists of three parts, namely,
initialisation, main body, and output, each of which we will explain next.

The initialisation part declares $\mathbb{R}^{d}$-valued input variables
$x_{1},\ldots,x_{p}$, each being initialised with i.i.d. Gaussian entries, and
also $d\times d$ matrices $M_{1},\ldots,M_{q}$ where each matrix is
initialised with i.i.d. Gaussian entries and independent with all the other
matrices and the input variables. That is, for $i\in[p]$, $i^{\prime}\in[q]$,
and $k,k^{\prime}\in[d]-1$,

$\displaystyle x_{i}$ $\displaystyle\in\mathbb{R}^{d},$
$\displaystyle(x_{i})_{k}$
$\displaystyle\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathcal{N}(0,\sigma_{i}^{2}),$
$\displaystyle M_{i^{\prime}}$ $\displaystyle\in\mathbb{R}^{d\times d},$
$\displaystyle(M_{i^{\prime}})_{k,k^{\prime}}$
$\displaystyle\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathcal{N}\left(0,\frac{\tau_{i^{\prime}}^{2}}{d}\right).$

The input variables can be correlated. This correlation is described by a
covariance matrix $(C_{i,i^{\prime}})_{i,i^{\prime}\in[p]}$ that is assumed
given: for $i,i^{\prime}\in[p]$ and $k,k^{\prime}\in[d]-1$,

$\mathrm{Cov}((x_{i})_{k},(x_{i}^{\prime})_{k}^{\prime})=\begin{cases}C_{i,i^{\prime}},&\text{if
}k=k^{\prime};\\\ 0,&\text{otherwise}.\end{cases}$

The main body of a tensor program is a sequence of two types of assignments
such that no input variables are assigned and the other variables are assigned
once. The two types of assignments are as follows:

MatMul

$h=Mx$ or $h=M^{\intercal}x$ where $M$ is one of the matrices declared in the
initialisation part, and $x$ is a variable assigned by the next NonLin-type
assignment before the current MatMul-type assignment.

NonLin

$x=\phi(h_{1},\ldots,h_{m})$ where $h_{1},\ldots,h_{m}$ are input variables or
those assigned by the first MatMul-type assignment before the current
assignment $x=\phi(h_{1},\ldots,h_{m})$, the
$\phi:\mathbb{R}^{m}\to\mathbb{R}$ is a function, and this function is applied
to the vectors $h_{1},\ldots,h_{m}$ pointwise in this assignment.

The last output part of a tensor program has the form:

$\displaystyle\frac{1}{d}\sum_{k=0}^{d-1}\psi((h_{1})_{k},\ldots,(h_{m})_{k})\text{
or
}\left(\frac{1}{\sqrt{d}}v^{\intercal}h_{1},\ldots,\frac{1}{\sqrt{d}}v^{\intercal}h_{m}\right)$

where $\psi$ is a function of type $\mathbb{R}^{m}\to\mathbb{R}$,
$h_{1},\ldots,h_{m}$ are variables, and $v$ is an input variable.

Every variable $v$ (which can be $x$ or $h$) in a tensor program has an
associated real-valued random variable $Z^{v}$, which is defined inductively.
For input variables $x_{1},\ldots,x_{p}$, we have real-valued random variables
$Z^{x_{1}},\ldots,Z^{x_{p}}\in\mathbb{R}$ that are jointly distributed by the
zero-mean multivariate Gaussian with the following covariance:

$\mathrm{Cov}(Z^{x_{i}},Z^{x_{i^{\prime}}})=C_{i,i^{\prime}}.$

For assigned variables, their random variables are defined as follows.

MatMul

For the assignment $h=Mx$ or $h=M^{\intercal}x$, the corresponding $Z^{h}$ is
the zero-mean Gaussian random variable that is independent with
$Z^{h^{\prime}}$ previously-defined by induction whenever the matrix
$M^{\prime}$ in $h^{\prime}=M^{\prime}x^{\prime}$ (including $M^{\intercal}$)
is not the same as $M$, and is correlated with $Z^{h^{\prime\prime}}$ with
$h^{\prime\prime}=Mx^{\prime\prime}$ for the same $M$ as follows:

$\mathrm{Cov}(Z^{h},Z^{h^{\prime\prime}})=\sigma_{M}^{2}\mathop{\mathbb{E}}[Z^{x}Z^{x^{\prime\prime}}]$

where $\sigma_{M}^{2}$ is the variance of the entries of the matrix $M$.

NonLin

For the assignment $x=\phi(h_{1},\ldots,h_{m})$, the corresponding random
variable $Z^{x}$ is defined by the following equation:

$Z^{x}=\phi(Z^{h_{1}},\ldots,Z^{h_{m}}).$

Note that $h_{1},\ldots,h_{m}$ are assigned before the assignment for $x$ so
that the random variables $Z^{h_{1}},\ldots,Z^{h_{m}}$ are defined by the
inductive construction.

The main reason for using the tensor-program framework is the master theorem
which states that as $n$ goes to infinity, the $k$-th components of variables
in a tensor program jointly converge the random variables that we have just
defined inductively.

######  Theorem 12 (Master Theorem).

If a tensor program uses only polynomially-bounded nonlinear functions $\phi$
in NonLin (i.e., there are some $C,c_{1},c_{2}>0$ such that for any
$u\in\mathbb{R}^{m}$, $\phi(u)\leq C\lVert u\rVert^{c_{1}}+c_{2}$) and it
satisfies the so-called BP-likeness (see [Yang, 2019, Yang, 2020] for the
definition), then for all polymonially-bounded $\psi$, we have the following
almost-sure convergence:

$\frac{1}{d}\sum_{k=0}^{d-1}\psi((h_{1})_{k},\ldots,(h_{m})_{k})\stackrel{{\scriptstyle
a.s.}}{{\to}}\mathop{\mathbb{E}}[\psi(Z^{h_{1}},\ldots,Z^{h_{m}})],$

as $d$ tends to infinity.

Using this Master theorem and Proposition G.4 from [Yang, 2019], we can show
that the output distribution converges to a Gaussian random variable.

######  Corollary 13 (GP Convergence).

Assume that a tensor program uses only polynomially-bounded nonlinear
functions $\phi$ in NonLin and it satisfies the so-called BP-likeness. Also,
assume that the output of the tensor program is

$\left(\frac{1}{\sqrt{d}}v^{\intercal}h_{1},\ldots,\frac{1}{\sqrt{d}}v^{\intercal}h_{m}\right)$

for the input random vector $v$ such that its entries are initialised with
$\mathcal{N}(0,\sigma_{v}^{2})$, the vector is not used anywhere else in the
program, and it is independent of all other input variables in the program.
Then, as $n$ tends to infinity, the output of the program converges in
distribution to the following Gaussian distribution:

$\left(\frac{1}{\sqrt{d}}v^{\intercal}h_{1},\ldots,\frac{1}{\sqrt{d}}v^{\intercal}h_{k}\right)\stackrel{{\scriptstyle
dist.}}{{\to}}\mathcal{N}(0,K)$

where

$K_{ij}=\sigma_{v}^{2}\mathop{\mathbb{E}}[Z^{h_{i}}Z^{h_{j}}].$

##  Appendix E PROOF OF THEOREM 3

We express the computation of the network’s outputs and Jacobians on all the
training inputs as a tensor program and use the Master theorem 13 to show that
these outputs and Jacobians jointly converge to Gaussian random variables.

We first define the input vectors and matrices. We have the following
$(1+d_{0})N$-many input vectors, which correspond to the values of the
preactivations on the training inputs and the Jacobians of these
preactivations:

$g^{(1)}_{1},\;\ldots,\;g^{(1)}_{N},\;(J_{1}g^{(1)}_{1}),\;\ldots,\;(J_{d_{0}}g^{(1)}_{N})\in\mathbb{R}^{d}$

These input vectors are initialised as follows: for $k\in[d]-1$, $i,j\in[N]$,
and $\alpha,\beta\in[d_{0}]$,

$\displaystyle\mathrm{Var}\left((g_{i}^{(1)})_{k}\right)$ $\displaystyle=1,$
$\displaystyle\mathrm{Var}\left((J_{\alpha}g_{i}^{(1)})_{k}\right)$
$\displaystyle=1,$
$\displaystyle\mathrm{Cov}\left((g_{i}^{(1)})_{k},(g_{j}^{(1)})_{k}\right)$
$\displaystyle=\left\langle x^{(i)},x^{(j)}\right\rangle,$
$\displaystyle\mathrm{Cov}\left((J_{\alpha}g_{i}^{(1)})_{k},(g_{j}^{(1)})_{k}\right)$
$\displaystyle=\left\langle e_{\alpha-1},x^{(j)}\right\rangle,$
$\displaystyle\mathrm{Cov}\left((g_{i}^{(1)})_{k},(J_{\beta}g_{j}^{(1)})_{k}\right)$
$\displaystyle=\left\langle x^{(i)},e_{\beta-1}\right\rangle,$
$\displaystyle\mathrm{Cov}\left((J_{\alpha}g_{i}^{(1)})_{k},(J_{\beta}g_{j}^{(1)})_{k}\right)$
$\displaystyle=\left\langle e_{\alpha-1},e_{\beta-1}\right\rangle,$

where $e_{\alpha-1}$ is the unit vector with its $\alpha$-th component being 1
and the others being 0. Note that this initialisation corresponds to having a
random matrix $W^{(1)}$ intialised with i.i.d. samples from
$\mathcal{N}(0,1)$, and setting $g^{(1)}_{i}$ to $W^{(1)}x^{(i)}$ and
$J_{\alpha}g_{i}^{(1)}$ to $\frac{\partial W^{(1)}x}{\partial
x_{\alpha-1}}|_{x=x^{(i)}}$.

We also have an input vector $V^{(L+1)}\in\mathbb{R}^{d}$ whose entries are
set to the i.i.d. samples from $\mathcal{N}(0,\kappa^{2})$. The vector is
independent of the previously defined input vectors.

For the input matrices, we have $V^{(l)}$ for each layer $l=2,\ldots,L$. The
entries of $V^{(l)}$ are initialised with i.i.d. samples from
$\mathcal{N}(0,1/d)$.

Next, we define the body of the program that computes the network’s outputs on
all the training inputs as well as the corresponding Jacobians. This body uses
the following variables:

$\displaystyle
h_{1}^{(1)},\;\ldots,\;h_{N}^{(1)},\;\ldots,\;h_{1}^{(L)},\;\ldots,\;h_{N}^{(L)}\in\mathbb{R}^{d},$
$\displaystyle
g_{1}^{(2)},\;\ldots,\;g_{N}^{(2)},\;\ldots,\;g_{1}^{(L)},\;\ldots,\;g_{N}^{(L)}\in\mathbb{R}^{d},$
$\displaystyle(J_{1}h_{1}^{(1)}),\;\ldots,\;(J_{d_{0}}h_{N}^{(1)}),\;\ldots,\;(J_{1}h_{1}^{(L)}),\;\ldots,\;(J_{d_{0}}h_{N}^{(L)})\in\mathbb{R}^{d},$
$\displaystyle(J_{1}g_{1}^{(2)}),\;\ldots,\;(J_{d_{0}}g_{N}^{(2)}),\;\ldots,\;(J_{1}g_{1}^{(L)}),\;\ldots,\;(J_{d_{0}}g_{N}^{(L)})\in\mathbb{R}^{d}.$

It is defined as follows:

$\displaystyle h_{1}^{(1)}=\phi(g_{1}^{(1)});$ $\displaystyle\ldots$
$\displaystyle h_{N}^{(1)}=\phi(g_{N}^{(1)});$ $\displaystyle
g_{1}^{(2)}=V^{(2)}h_{1}^{(1)};$ $\displaystyle\ldots$ $\displaystyle
g_{N}^{(2)}=V^{(2)}h_{N}^{(1)};$ $\displaystyle
h_{1}^{(2)}=\phi(g_{1}^{(2)});$ $\displaystyle\ldots$ $\displaystyle
h_{N}^{(2)}=\phi(g_{N}^{(2)});$ $\displaystyle\vdots$ $\displaystyle\vdots$
$\displaystyle g_{1}^{(L)}=V^{(L)}h_{1}^{(L-1)};$ $\displaystyle\ldots$
$\displaystyle g_{N}^{(L)}=V^{(L)}h_{N}^{(L-1)};$ $\displaystyle
h_{1}^{(L)}=\phi(g_{1}^{(L)});$ $\displaystyle\ldots$ $\displaystyle
h_{N}^{(L)}=\phi(g_{N}^{(L)});$
$\displaystyle(J_{1}h_{1}^{(1)})=\dot{\phi}(g_{1}^{(1)})\odot(J_{1}g_{1}^{(1)});$
$\displaystyle\ldots$
$\displaystyle(J_{d_{0}}h_{N}^{(1)})=\dot{\phi}(g_{N}^{(1)})\odot(J_{d_{0}}g_{N}^{(1)});$
$\displaystyle(J_{1}g_{1}^{(2)})=V^{(2)}(J_{1}h_{1}^{(1)});$
$\displaystyle\ldots$
$\displaystyle(J_{d_{0}}g_{N}^{(2)})=V^{(2)}(J_{d_{0}}h_{N}^{(1)});$
$\displaystyle(J_{1}h_{1}^{(2)})=\dot{\phi}(g_{1}^{(2)})\odot(J_{1}g_{1}^{(2)});$
$\displaystyle\ldots$
$\displaystyle(J_{d_{0}}h_{N}^{(2)})=\dot{\phi}(g_{N}^{(2)})\odot(J_{d_{0}}g_{N}^{(2)});$
$\displaystyle\vdots$ $\displaystyle\vdots$
$\displaystyle(J_{1}g_{1}^{(L)})=V^{(L)}(J_{1}h_{1}^{(L-1)});$
$\displaystyle\ldots$
$\displaystyle(J_{d_{0}}g_{N}^{(L)})=V^{(L)}(J_{d_{0}}h_{N}^{(L-1)});$
$\displaystyle(J_{1}h_{1}^{(L)})=\dot{\phi}(g_{1}^{(L)})\odot(J_{1}g_{1}^{(L)});$
$\displaystyle\ldots$
$\displaystyle(J_{d_{0}}h_{N}^{(L)})=\dot{\phi}(g_{N}^{(L)})\odot(J_{d_{0}}g_{N}^{(L)});$

Finally, we make the program output the following $(1+d_{0})N$-dimensional
vector:

$\displaystyle\Bigg{\langle}\frac{1}{\sqrt{d}}(V^{(L+1)})^{\intercal}h_{1}^{(L)},\;\ldots,\;\frac{1}{\sqrt{d}}(V^{(L+1)})^{\intercal}h_{N}^{(L)},\;$
$\displaystyle\qquad\frac{1}{\sqrt{d}}(V^{(L+1)})^{\intercal}(J_{1}h_{1}^{(L)}),\;\ldots,\;\frac{1}{\sqrt{d}}(V^{(L+1)})^{\intercal}(J_{1}h_{N}^{(L)}),\;$
$\displaystyle\qquad\ldots$
$\displaystyle\qquad\frac{1}{\sqrt{d}}(V^{(L+1)})^{\intercal}(J_{d_{0}}h_{1}^{(L)}),\;\ldots,\;\frac{1}{\sqrt{d}}(V^{(L+1)})^{\intercal}(J_{d_{0}}h_{N}^{(L)})\Bigg{\rangle}.$

Now Corollary 13 implies that these random variables jointly converge in
distribution to a zero-mean multivariate Gaussian distribution as $d$ tends to
infinity. Furthermore, the Master theorem (i.e., Theorem 12) gives the
following inductive formula for computing the expectation in the definition of
the covariance matrix in the corollary:

$\displaystyle\mathrm{Cov}\left(Z^{g_{i}^{(l)}},Z^{g_{j}^{(l)}}\right)$
$\displaystyle=\mathop{\mathbb{E}}\left[\phi\left(Z^{g_{i}^{(l-1)}}\right)\phi\left(Z^{g_{j}^{(l-1)}}\right)\right],$
$\displaystyle\mathrm{Cov}\left(Z^{J_{\alpha}g_{i}^{(l)}},Z^{g_{j}^{(l)}}\right)$
$\displaystyle=\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l-1)}}\right)Z^{J_{\alpha}g_{i}^{(l-1)}}\phi\left(Z^{g_{j}^{(l-1)}}\right)\right],$
$\displaystyle\mathrm{Cov}\left(Z^{g_{i}^{(l)}},Z^{J_{\beta}g_{j}^{(l)}}\right)$
$\displaystyle=\mathop{\mathbb{E}}\left[\phi\left(Z^{g_{i}^{(l-1)}}\right)\dot{\phi}\left(Z^{g_{j}^{(l-1)}}\right)Z^{J_{\beta}g_{j}^{(l-1)}}\right],$
$\displaystyle\mathrm{Cov}\left(Z^{J_{\alpha}g_{i}^{(l)}},Z^{J_{\beta}g_{j}^{(l)}}\right)$
$\displaystyle=\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l-1)}}\right)Z^{J_{\alpha}g_{i}^{(l-1)}}\dot{\phi}\left(Z^{g_{j}^{(l-1)}}\right)Z^{J_{\beta}g_{j}^{(l-1)}}\right].$

By definition, the random vectors
$(Z^{g_{i}^{(1)}},Z^{J_{1}g_{i}^{(1)}},\ldots,Z^{J_{d_{0}}g_{i}^{(l)}})\in\mathbb{R}^{1+d_{0}}$
for $i\in[N]$ have the covariance given by $\Sigma^{(0)}(x^{(i)},x^{(j)})$ in
the statement of Theorem 3. Note that the inductive definition that we have
just given is identical to the inductive definition of
$\Sigma^{(l-1)}(x^{(i)},x^{(j)})$ in Theorem 3. Thus, the random vectors
$(Z^{g_{i}^{(l)}},Z^{J_{1}g_{i}^{(l)}},\ldots,Z^{J_{d_{0}}g_{i}^{(l)}})\in\mathbb{R}^{1+d_{0}}$
for $i\in[N]$ have the covariance given by $\Sigma^{(l-1)}(x^{(i)},x^{(j)})$.
As a result, the covariance of the limiting distribution of the network output
and its Jacobian is given by

$\displaystyle\mathrm{Cov}\left(f(x^{(i)}),f(x^{(j)})\right)$
$\displaystyle=\kappa^{2}\Sigma^{(L)}(x^{(i)},x^{(j)})_{00},$
$\displaystyle\mathrm{Cov}\left(J(f)(x^{(i)})_{\alpha-1},f(x^{(j)})\right)$
$\displaystyle=\kappa^{2}\Sigma^{(L)}(x^{(i)},x^{(j)})_{\alpha 0},$
$\displaystyle\mathrm{Cov}\left(f(x^{(i)}),J(f)(x^{(j)})_{\beta-1}\right)$
$\displaystyle=\kappa^{2}\Sigma^{(L)}(x^{(i)},x^{(j)})_{0\beta},$
$\displaystyle\mathrm{Cov}\left(J(f)(x^{(i)})_{\alpha-1},J(f)(x^{(j)})_{\beta-1}\right)$
$\displaystyle=\kappa^{2}\Sigma^{(L)}(x^{(i)},x^{(j)})_{\alpha\beta}.$

This proves Theorem 3.

We use this theorem to show that if we take $\kappa$ small enough and $d$
large enough, both outputs and Jacobians at initialisation are close to zero.

######  Remark 14.

By the definition of $\Sigma^{L}$ and the assumption on the activiation
function $\phi$, for all $x\in\mathbb{R}^{d_{0}}$ with $\|x\|=1$, the diagonal
entries of $\kappa^{2}\Sigma^{(L)}(x,x)$ are at most
$\kappa^{2}(M_{1}^{2L}+1)$. Thus, using the tail bound for the standard normal
distribution and also union bound, we can show that for all $\epsilon>0$ and
$x\in\mathbb{R}^{d_{0}}$ with $\|x\|=1$, if
$g\sim\mathcal{GP}(0,\kappa^{2}\Sigma^{(L)})$, then

$P\left(\left\|g(x)\right\|_{\infty}>\epsilon\right)\leq
2(1+d_{0})\exp\left(-\frac{\epsilon^{2}}{2\kappa^{2}(M_{1}^{2L}+1)}\right).$

Since every training input $x^{(i)}$ satisfies $\|x^{(i)}\|=1$ by assumption,
we can instantiate the above bound on those inputs and get, by union bound,
that

$P\left(\forall
i\in[N].\,\left\|g(x^{(i)})\right\|_{\infty}\leq\epsilon\right)\geq
1-2N(1+d_{0})\exp\left(-\frac{\epsilon^{2}}{2\kappa^{2}(M^{2L}_{1}+1)}\right).$
(5)

By Theorem 3, there exists a function $F_{1}$ that if

$d\geq F_{1}\left(x^{(1:N)},\,L,\,\delta,\
\frac{\epsilon}{\kappa\sqrt{M^{2L}_{1}+1}}\right),$

then

$\left|P\left(\forall
i\in[N],\alpha\in[d_{0}].\,\left(f_{d}(x^{(i)}),J(f_{d})(x^{(i)})_{\alpha-1}\in[-\epsilon,\epsilon]\right)\right)-P\left(\left\|g(x^{(i)})\right\|_{\infty}\leq\epsilon\right)\right|\leq\frac{\delta}{2}.$
(6)

From the bounds in (5) and (6) it follows that for all $\delta\in(0,1)$, if
$\kappa$ is sufficiently small and $d$ is sufficiently large, then with
probability at least $1-\delta$,

$f_{d}(x^{(i)}),\,J(f_{d})(x^{(i)})_{\alpha-1}\in[-\epsilon,\epsilon]$

for all $i\in[N]$ and $\alpha\in[d_{0}]$.

##  Appendix F PROOF OF THEOREM 4

We start with one general result that allows us to move the derivative from
the outside of an expectation to the inside when the expectation is taken over
a GP.

######  Theorem 15.

Let $K:\mathbb{R}^{d_{0}}\times\mathbb{R}^{d_{0}}\to\mathbb{R}$ be a symmetric
kernel. Consider functions $\psi:\mathbb{R}\to\mathbb{R}$ and
$\varphi:\mathbb{R}^{2}\to\mathbb{R}$ such that

  * •

$\psi$ and $\varphi$ are polynomially bounded; and

  * •

there exists a polynomially-bounded function
$\dot{\psi}:\mathbb{R}\to\mathbb{R}$ that satisfies
$\int_{0}^{x}\dot{\psi}(t)dt=\psi(x)-\psi(0)$ for all $x\in\mathbb{R}$.

Let $x,x^{\prime}\in\mathbb{R}^{d_{0}}$, $\alpha,\beta\in[d_{0}]$, and

$\displaystyle\Gamma$
$\displaystyle=\begin{pmatrix}K(x,x)&K(x,x^{\prime})&\frac{\partial
K(x,y^{\prime})}{\partial
y^{\prime}_{\alpha-1}}\Big{|}_{y^{\prime}=x}&\frac{\partial
K(x,y^{\prime})}{\partial
y^{\prime}_{\beta-1}}\Big{|}_{y^{\prime}=x^{\prime}}\\\
K(x^{\prime},x)&K(x^{\prime},x^{\prime})&\frac{\partial
K(x^{\prime},y^{\prime})}{\partial
y^{\prime}_{\alpha-1}}\Big{|}_{y^{\prime}=x}&\frac{\partial
K(x^{\prime},y^{\prime})}{\partial
y^{\prime}_{\beta-1}}\Big{|}_{y^{\prime}=x^{\prime}}\\\ \frac{\partial
K(y,x)}{\partial y_{\alpha-1}}\Big{|}_{y=x}&\frac{\partial
K(y,x^{\prime})}{\partial
y_{\alpha-1}}\Big{|}_{y=x}&\frac{\partial^{2}K(y,y^{\prime})}{\partial
y_{\alpha-1}\partial
y^{\prime}_{\alpha-1}}\Big{|}_{(y,y^{\prime})=(x,x)}&\frac{\partial^{2}K(y,y^{\prime})}{\partial
y_{\alpha-1}\partial
y^{\prime}_{\beta-1}}\Big{|}_{(y,y^{\prime})=(x,x^{\prime})}\\\ \frac{\partial
K(y,x)}{\partial y_{\beta-1}}\Big{|}_{y=x^{\prime}}&\frac{\partial
K(y,x^{\prime})}{\partial
y_{\beta-1}}\Big{|}_{y=x^{\prime}}&\frac{\partial^{2}K(y,y^{\prime})}{\partial
y_{\beta-1}\partial
y^{\prime}_{\alpha-1}}\Big{|}_{(y,y^{\prime})=(x^{\prime},x)}&\frac{\partial^{2}K(y,y^{\prime})}{\partial
y_{\beta-1}\partial
y^{\prime}_{\beta-1}}\Big{|}_{(y,y^{\prime})=(x^{\prime},x^{\prime})}\end{pmatrix}.$

Then,

$\frac{\partial}{\partial
x^{\prime}_{\beta-1}}\mathop{\mathbb{E}}_{z\sim\mathcal{N}(0,\Gamma)}\left[\psi(z_{1})\cdot\varphi(z_{0},z_{2})\right]=\mathop{\mathbb{E}}_{z\sim\mathcal{N}(0,\Gamma)}\left[\dot{\psi}(z_{1})\cdot
z_{3}\cdot\varphi(z_{0},z_{2})\right].$ (7)

###### Proof.

Fix $x,x^{\prime}\in\mathbb{R}^{d_{0}}$ and $\alpha,\beta\in[d_{0}]$. Let
$\Gamma$ be the matrix in the statement of the theorem, and
$\Sigma=\Gamma_{-3,-3}\in\mathbb{R}^{3\times 3}$ the submatrix of $\Gamma$
without the last column and row. Let $I_{3}$ and $I_{4}$ be the $3\times 3$
and $4\times 4$ identity matrices, respectively.

We will first introduce a function $H(u)$ which gives the reparameterisation
of $\Sigma$ for $z_{1}$, defined as

$\displaystyle a_{0}$
$\displaystyle=\frac{\Sigma_{0,1}}{\sqrt{\Sigma_{0,0}}},$ $\displaystyle
a_{2}$
$\displaystyle=\frac{\Sigma_{1,2}-\Sigma_{0,1}\Sigma_{0,2}/\Sigma_{0,0}}{\sqrt{\Sigma_{2,2}-\Sigma_{0,2}^{2}/\Sigma_{0,0}}},$
$\displaystyle a_{1}$ $\displaystyle=\sqrt{\Sigma_{1,1}-a_{0}^{2}-a_{1}^{2}},$
$\displaystyle H(u)$ $\displaystyle=a_{0}u_{0}+a_{1}u_{1}+a_{2}u_{2}.$

This gives the reparameterisation as

$\left(z_{0},z_{1},z_{2}\right)\stackrel{{\scriptstyle
dist}}{{=}}\left(\sqrt{\Sigma_{0,0}}u_{0},H(u),\frac{\Sigma_{0,2}}{\sqrt{\Sigma_{0,0}}}u_{0}+\sqrt{\Sigma_{2,2}-\frac{\Sigma_{0,2}^{2}}{\Sigma_{0,0}}}u_{2}\right)$

where $z\sim\mathcal{N}(0,\Sigma)$ and $u\sim\mathcal{N}(0,I_{3})$.

Applying this reparameterisation to our expectations allows us to remove the
dependency of distribution on the inputs,

$\displaystyle\mathop{\mathbb{E}}_{z\sim\mathcal{N}(0,\Gamma)}\left[\psi(z_{1})\varphi(z_{0},z_{2})\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}_{z\sim\mathcal{N}(0,\Sigma)}\left[\psi(z_{1})\varphi(z_{0},z_{2})\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}_{u\sim\mathcal{N}(0,I_{3})}\left[\psi(H(u))\varphi\left(\sqrt{\Sigma_{0,0}}u_{0},\frac{\Sigma_{0,2}}{\sqrt{\Sigma_{0,0}}}u_{0}+\sqrt{\Sigma_{2,2}-\frac{\Sigma_{0,2}^{2}}{\Sigma_{0,0}}}u_{2}\right)\right].$

Taking the derivative w.r.t. $x_{\beta-1}^{\prime}$, gives

$\displaystyle\frac{\partial}{\partial
x_{\beta-1}^{\prime}}\mathop{\mathbb{E}}_{z\sim\mathcal{N}(0,\Gamma)}\left[\psi(z_{1})\varphi(z_{0},z_{2})\right]$
$\displaystyle={}\frac{\partial}{\partial
x_{\beta-1}^{\prime}}\mathop{\mathbb{E}}_{u\sim\mathcal{N}(0,I_{3})}\left[\psi(H(u))\varphi\left(\sqrt{\Sigma_{0,0}}u_{0},\frac{\Sigma_{0,2}}{\sqrt{\Sigma_{0,0}}}u_{0}+\sqrt{\Sigma_{2,2}-\frac{\Sigma_{0,2}^{2}}{\Sigma_{0,0}}}u_{2}\right)\right]$
$\displaystyle={}\mathop{\mathbb{E}}_{u\sim\mathcal{N}(0,I_{3})}\left[\frac{\partial}{\partial
x_{\beta-1}^{\prime}}\left(\psi(H(u))\varphi\left(\sqrt{\Sigma_{0,0}}u_{0},\frac{\Sigma_{0,2}}{\sqrt{\Sigma_{0,0}}}u_{0}+\sqrt{\Sigma_{2,2}-\frac{\Sigma_{0,2}^{2}}{\Sigma_{0,0}}}u_{2}\right)\right)\right]$
$\displaystyle={}\mathop{\mathbb{E}}_{u\sim\mathcal{N}(0,I_{4})}\left[H^{\prime}(u)\dot{\varphi}(H(u))\varphi\left(\sqrt{\Sigma_{0,0}}u_{0},\frac{\Sigma_{0,2}}{\sqrt{\Sigma_{0,0}}}u_{0}+\sqrt{\Sigma_{2,2}-\frac{\Sigma_{0,2}^{2}}{\Sigma_{0,0}}}u_{2}\right)\right],$

using the fact that $\Sigma_{0,0},\Sigma_{0,2},\Sigma_{2,2}$ do not involve
$x^{\prime}$, and $H^{\prime}(u)$ is defined by differentiating $H(u)$ by
$x_{\beta-1}^{\prime}$ and adding additional $b_{3}u_{3}$, giving

$\displaystyle H^{\prime}(u)$
$\displaystyle=b_{0}u_{0}+b_{1}u_{1}+b_{2}u_{2}+b_{3}u_{3},$ $\displaystyle
b_{0}$
$\displaystyle=\frac{1}{\sqrt{\Sigma_{0,0}}}\frac{\partial\Sigma_{0,1}}{\partial
x_{\beta-1}^{\prime}},$ $\displaystyle b_{2}$
$\displaystyle=\frac{1}{\sqrt{\Sigma_{2,2}-\Sigma_{0,2}^{2}/\Sigma_{0,0}}}\left(\frac{\partial\Sigma_{1,2}}{\partial
x_{\beta-1}^{\prime}}-\frac{\Sigma_{0,2}}{\Sigma_{0,0}}\frac{\partial\Sigma_{0,1}}{\partial
x_{\beta-1}^{\prime}}\right),$ $\displaystyle b_{1}$
$\displaystyle=\frac{1}{2a_{2}}\left(\frac{\partial\Sigma_{1,1}}{\partial
x_{\beta-1}^{\prime}}-2a_{0}b_{0}-2a_{2}b_{2}\right),$ $\displaystyle b_{3}$
$\displaystyle=\sqrt{\Gamma_{3,3}-b_{0}^{2}-b_{1}^{2}-b_{2}^{2}}.$

Now if $u\sim\mathcal{N}(0,I_{4})$, the 4-dimensional random vector

$\left(\sqrt{\Sigma_{0,0}u_{0}},H(u),\frac{\Sigma_{0,2}}{\sqrt{\Sigma_{0,0}}}u_{0}+\sqrt{\Sigma_{2,2}-\frac{\Sigma_{0,2}^{2}}{\Sigma_{0,0}}}u_{2},H^{\prime}(u)\right)$

has the distribution $\mathcal{N}(0,\Gamma)$. This can be verified by
computing the covariance,

$\displaystyle\mathrm{Cov}(\sqrt{\Sigma_{0,0}}u_{0},H^{\prime}(u))$
$\displaystyle=\frac{\Sigma_{0,1}}{\partial x_{\beta-1}^{\prime}},$
$\displaystyle\mathrm{Cov}(H(u),H^{\prime}(u))$
$\displaystyle=\frac{1}{2}\frac{\partial\Sigma_{1,1}}{\partial
x_{\beta-1}^{\prime}}=\frac{\partial K(x^{\prime},y^{\prime})}{\partial
y^{\prime}_{\beta-1}}\Big{|}_{y^{\prime}=x^{\prime}},$
$\displaystyle\mathrm{Cov}\left(\frac{\Sigma_{0,2}}{\sqrt{\Sigma_{0,0}}}u_{0}+\sqrt{\Sigma_{2,2}-\frac{\Sigma_{0,2}^{2}}{\Sigma_{0,0}}}u_{2},H^{\prime}(u)\right)$
$\displaystyle=\frac{\partial\Sigma_{1,2}}{\partial x_{\beta-1}^{\prime}},$
$\displaystyle\mathrm{Var}(H^{\prime}(u))$
$\displaystyle=\frac{\partial^{2}K(y,y^{\prime})}{\partial y_{\beta-1}\partial
y^{\prime}_{\beta-1}}\Big{|}_{(y,y^{\prime})=(x^{\prime},x^{\prime})},$

where the second equality comes from the symmetry of $K$:

$\frac{\partial}{\partial
x_{\beta-1}^{\prime}}\Sigma_{1,1}=\left(\frac{\partial}{\partial
y_{\beta-1}}K(y,x^{\prime})\big{|}_{y=x^{\prime}}+\frac{\partial}{\partial
y_{\beta-1}^{\prime}}K(x^{\prime},y^{\prime})\big{|}_{y^{\prime}=x^{\prime}}\right)=2\frac{\partial}{\partial
y_{\beta-1}^{\prime}}K(x^{\prime},y^{\prime})\big{|}_{y^{\prime}=x^{\prime}}.$

The last equality uses the independence between $u_{3}$ and
$u_{0},u_{1},u_{2}$ and the zero mean property of $u_{3}$.

From what we have just shown, the final result follows:

$\displaystyle\frac{\partial}{\partial
x^{\prime}_{\beta-1}}\mathop{\mathbb{E}}_{z\sim\mathcal{N}(0,\Sigma)}\left[\psi(z_{1})\cdot\varphi(z_{0},z_{2})\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}_{u\sim\mathcal{N}(0,I_{4})}\left[\dot{\psi}(H(u))H^{\prime}(u)\varphi\left(\sqrt{\Sigma_{1,1}u_{0}},\frac{\Sigma_{0,2}}{\sqrt{\Sigma_{0,0}}}u_{0}+\sqrt{\Sigma_{2,2}-\frac{\Sigma_{0,2}^{2}}{\Sigma_{0,0}}}u_{2}\right)\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}_{z\sim\mathcal{N}(0,\Gamma)}\left[\dot{\psi}(z_{1})\cdot
z_{3}\cdot\varphi(z_{0},z_{2})\right].$

∎

Using Theorem 15, we prove Theorem 4. Pick $\alpha,\beta\in[d_{0}]$, and
$l\in[L]\cup\\{0\\}$. We have to show that for all
$x,x^{\prime}\in\mathbb{R}^{d_{0}}$,

$\displaystyle\Sigma^{(l)}(x,x^{\prime})_{\alpha 0}$
$\displaystyle=\frac{\partial}{\partial
x_{\alpha-1}}\Sigma^{(l)}(x,x^{\prime})_{00},$ (8)
$\displaystyle\Sigma^{(l)}(x,x^{\prime})_{0\beta}$
$\displaystyle=\frac{\partial}{\partial
x_{\beta-1}^{\prime}}\Sigma^{(l)}(x,x^{\prime})_{00},$ (9)
$\displaystyle\Sigma^{(l)}(x,x^{\prime})_{\alpha\beta}$
$\displaystyle=\frac{\partial^{2}}{\partial x_{\alpha-1}\partial
x_{\beta-1}^{\prime}}\Sigma^{(l)}(x,x^{\prime})_{00}.$ (10)

The proof is given by induction on $l$. When $l=0$, all the three desired
equations hold for all $x,x^{\prime}\in\mathbb{R}^{d_{0}}$, since
differentiation of $\langle x,x^{\prime}\rangle$ gives LHS. For the induction
step, assume that the three equations (8), (9), and (10) hold for all
$x,x^{\prime}\in\mathbb{R}^{d_{0}}$ for the $l$-th layer. We need to show that
the equations also hold for all $x,x^{\prime}\in\mathbb{R}^{d_{0}}$ for the
$(l+1)$-th layer. First, we prove that the equation (9) holds for all
$x,x^{\prime}\in\mathbb{R}^{d_{0}}$; the equation (8) can be proved similarly.
We define the $4\times 4$ matrix $\Gamma$ as in Theorem 15, with the kernel
$(x,x^{\prime})\mapsto\Sigma^{(l)}(x,x^{\prime})_{00}$, where following
equality can be proven with the induction hypothesis,

$\Gamma=\begin{pmatrix}\Sigma^{(l)}(x,x)_{00}&\Sigma^{(l)}(x,x^{\prime})_{00}&\Sigma^{(l)}(x,x)_{0\alpha}&\Sigma^{(l)}(x,x^{\prime})_{0\beta}\\\
\Sigma^{(l)}(x,x^{\prime})_{00}&\Sigma^{(l)}(x^{\prime},x^{\prime})_{00}&\Sigma^{(l)}(x,x^{\prime})_{\alpha
0}&\Sigma^{(l)}(x^{\prime},x^{\prime})_{0\beta}\\\
\Sigma^{(l)}(x,x)_{0\alpha}&\Sigma^{(l)}(x,x^{\prime})_{\alpha
0}&\Sigma^{(l)}(x,x)_{\alpha\alpha}&\Sigma^{(l)}(x,x^{\prime})_{\alpha\beta}\\\
\Sigma^{(l)}(x,x^{\prime})_{0\beta}&\Sigma^{(l)}(x^{\prime},x^{\prime})_{0\beta}&\Sigma^{(l)}(x,x^{\prime})_{\alpha\beta}&\Sigma^{(l)}(x^{\prime},x^{\prime})_{\beta\beta}\end{pmatrix}.$

Using this equality and Theorem 15 with $\psi(z_{1})=\phi(z_{1})$ and
$\varphi(z_{0},z_{2})=\phi(z_{0})$,

$\displaystyle\Sigma^{(l+1)}(x,x^{\prime})_{0\beta}$
$\displaystyle=\mathop{\mathbb{E}}_{g}\left[\phi(g(x))\dot{\phi}(g(x^{\prime}))J(g)(x^{\prime})_{\beta-1}\right]$
$\displaystyle=\mathop{\mathbb{E}}_{u\sim\mathcal{N}(0,\Gamma)}\left[\phi(u_{0})\dot{\phi}(u_{1})u_{3}\right]$
$\displaystyle=\frac{\partial}{\partial
x_{\beta-1}^{\prime}}\mathop{\mathbb{E}}_{u\sim\mathcal{N}(0,\Gamma)}\left[\phi(u_{0})\phi(u_{1})\right]$
$\displaystyle=\frac{\partial}{\partial
x_{\beta-1}^{\prime}}\mathop{\mathbb{E}}_{g}\left[\phi(g(x))\phi(g(x^{\prime}))\right]$
$\displaystyle=\frac{\partial}{\partial
x_{\beta-1}^{\prime}}\Sigma^{(l+1)}(x,x^{\prime})_{00}$

We can prove the equation (10) with a similar method. This time, we define the
matrix $\Xi$ which is a matrix defined by permutation of the rows and columns
of $\Gamma$ by the bijection that maps $(1,2,3,4)$ to $(2,1,4,3)$, so that the
following holds:

$(u_{0},u_{1},u_{2},u_{3})\stackrel{{\scriptstyle
dist}}{{=}}(w_{1},w_{0},w_{3},w_{2})$

where $u\sim\mathcal{N}(0,\Gamma)$ and $w\sim\mathcal{N}(0,\Xi)$. Then with
Theorem 15 with the choice of $\psi(z_{1})=\phi(z_{1})$ and
$\varphi(z_{0},z_{2})=\dot{\phi}(z_{0})z_{2}$ and the second equality (9) just
proven, the third equality can be shown as

$\displaystyle\Sigma^{(l+1)}(x,x^{\prime})_{\alpha\beta}$
$\displaystyle=\mathop{\mathbb{E}}_{g}\left[\dot{\phi}(g(x))\dot{g}(x)\dot{\phi}(g(x^{\prime}))\dot{g}(x^{\prime})\right]$
$\displaystyle=\mathop{\mathbb{E}}_{w\sim\mathcal{N}(0,\Xi)}\left[\dot{\phi}(w_{1})w_{3}\dot{\phi}(w_{0})w_{2}\right]$
$\displaystyle=\frac{\partial}{\partial
x_{\alpha-1}}\mathop{\mathbb{E}}_{w\sim\mathcal{N}(0,\Xi)}\left[\phi(w_{1})\dot{\phi}(w_{0})w_{2}\right]$
$\displaystyle=\frac{\partial}{\partial
x_{\alpha-1}}\mathop{\mathbb{E}}_{g}\left[\phi(g(x))\dot{\phi}(g(x^{\prime}))\dot{g}(x^{\prime})\right]$
$\displaystyle=\frac{\partial}{\partial
x_{\alpha-1}}\Sigma^{(l+1)}(x,x^{\prime})_{0\beta}$
$\displaystyle=\frac{\partial^{2}}{\partial x_{\alpha-1}\partial
x_{\beta-1}^{\prime}}\Sigma^{(l+1)}(x,x^{\prime})_{\alpha\beta}.$

##  Appendix G REMARK ON ASSUMPTION 8

In this section, we illustrate the difficulty of finding a sufficient
condition for Assumption 8 by proving that one of the standard conditions for
the positive definiteness of the NTK kernel does not work for the JNTK kernel.
We provide a counterexample for the condition. To simplify the presentation,
we do not impose the assumptions in Assumption 1. The result for the
normalised case can be obtained by scaling every matrix by $\frac{1}{3}$.

We first review two results that guarantee the positive definiteness of the
NTK kernel.

######  Theorem 16 (Proposition F.1 of [Du et al., 2019b]).

Suppose activation is analytic and non-polynomial. Then, the standard NTK
kernel

$\Theta(x^{(1:N)},x^{(1:N)})_{00}$

has positive eigenvalues as long as no two data satisfy $x^{(i)}=cx^{(j)}$ for
some $c\in\mathbb{R}$.

######  Theorem 17 (Adapted from Theorem 3.2 of [Nguyen et al., 2021]).

Suppose the Hermite expansion of activation

$\sigma(x)=\frac{1}{\sqrt{2\pi}}\sum_{r=0}^{\infty}\mu_{r}(\sigma)H_{r}(x)\exp(-x^{2}/2)$

has infinitely many nonzero $\mu_{r}(\sigma)$. Then, the standard NTK kernel

$\Theta(x^{(1:N)},x^{(1:N)})_{00}$

has positive eigenvalues as long as no two data satisfy $x^{(i)}=cx^{(j)}$ for
some $c\in\mathbb{R}$.

###### Proof.

We use the following inequality stated in the proof of Theorem 3.2 of [Nguyen
et al., 2021]:

$\lambda_{\min}(\Theta(x^{(1:N)},x^{(1:N)})_{00})\geq\mu_{r}(\sigma)^{2}\frac{\min_{i\in[N]}\|x^{(i)}\|^{2r}-(N-1)\max_{i\neq
j}|\langle x^{(i)},x^{(j)}\rangle|^{r}}{\max_{i\in[N]}\|x^{(i)}\|^{2(r-1)}}.$

Then, we can choose $r$ large enough such that $\mu_{r}(\sigma)\neq 0$ and
$\max_{i\neq j}|\langle x^{(i)},x^{(j)}\rangle|^{r}<\frac{1}{N-1}$, which
guarantees the positive definiteness of the NTK kernel. ∎

Our counterexample is related to Theorem 17. Consider neural networks with the
activation function $\sigma(x)=x^{2}$, which is a polynomial function. This
activation satisfies the condition of Theorem 17, which requires that the
Hermite coefficients of the activation be non-zero for infinitely many $r$. It
is because otherwise, the summation would include $\exp(-x^{2}/2)$ term, which
is not a polynomial.

######  Theorem 18.

For any input dimension $d_{0}$ and square activation $\sigma(x)=x^{2}$, there
exists dataset $\\{x^{(i)}\\}_{i=1}^{N}$ such that

$\lambda_{\min}\left(\Theta(x^{(1:N)},x^{(1:N)})\right)=0$

but

$\lambda_{\min}\left(\Theta(x^{(1:N)},x^{(1:N)})_{00}\right)>0$

for shallow neural networks, i.e., $L=1$.

###### Proof.

We first give the analytic form of the JNTK kernel in this case, which can be
computed via Isserlis’ theorem.

$\displaystyle\Theta(x,y)$ $\displaystyle=\Sigma(x,y)+\dot{\Sigma}(x,y),$
$\displaystyle\Sigma(x,y)_{00}$
$\displaystyle=\mathbb{E}_{w\sim\mathcal{N}(0,I_{d})}\left[\langle
w,x\rangle^{2}\langle w,y\rangle^{2}\right]$
$\displaystyle=2\mathrm{Cov}(\langle w,x\rangle,\langle
w,y\rangle)^{2}+\mathrm{Var}(\langle w,x\rangle)\mathrm{Var}(\langle
w,y\rangle)$ $\displaystyle=2\langle x,y\rangle^{2}+\|x\|^{2}\|y\|^{2},$
$\displaystyle\Sigma(x,y)_{\alpha 0}$ $\displaystyle=\frac{\partial}{\partial
x_{\alpha-1}}\Sigma(x,y)_{00}$ $\displaystyle=4\langle x,y\rangle
y_{\alpha-1}+2x_{\alpha-1}\|y\|^{2},$ $\displaystyle\Sigma(x,y)_{0\beta}$
$\displaystyle=\frac{\partial}{\partial y_{\beta-1}}\Sigma(x,y)_{00}$
$\displaystyle=4\langle x,y\rangle x_{\beta-1}+2y_{\beta-1}\|x\|^{2},$
$\displaystyle\Sigma(x,y)_{\alpha\beta}$
$\displaystyle=\frac{\partial^{2}}{\partial x_{\alpha-1}\partial
y_{\beta-1}}\Sigma(x,y)_{00}$
$\displaystyle=4x_{\beta-1}y_{\alpha-1}+4x_{\alpha-1}y_{\beta-1}+4\langle
x,y\rangle\mathbb{I}_{\alpha=\beta},$ $\displaystyle\dot{\Sigma}(x,y)_{00}$
$\displaystyle=\mathbb{E}_{w\sim\mathcal{N}(0,I_{d})}\left[4\langle
w,x\rangle\langle w,y\rangle\right]$ $\displaystyle=4\langle x,y\rangle,$
$\displaystyle\dot{\Sigma}(x,y)_{\alpha 0}$
$\displaystyle=\frac{\partial}{\partial x_{\alpha-1}}\dot{\Sigma}(x,y)_{00}$
$\displaystyle=4y_{\alpha-1},$ $\displaystyle\dot{\Sigma}(x,y)_{0\beta}$
$\displaystyle=\frac{\partial}{\partial y_{\beta-1}}\dot{\Sigma}(x,y)_{00}$
$\displaystyle=4x_{\beta-1},$ $\displaystyle\dot{\Sigma}(x,y)_{\alpha\beta}$
$\displaystyle=\frac{\partial^{2}}{\partial x_{\alpha-1}\partial
y_{\beta-1}}\dot{\Sigma}(x,y)_{00}$
$\displaystyle=4\mathbb{I}_{\alpha=\beta}.$

In matrix form, we have

$\displaystyle\Sigma(x,y)$ $\displaystyle=\begin{pmatrix}2\langle
x,y\rangle^{2}+\|x\|^{2}\|y\|^{2}&4\langle x,y\rangle
x^{\intercal}+2\|x\|^{2}y^{\intercal}\\\ 4\langle x,y\rangle
y+2\|y\|^{2}x&4xy^{\intercal}+4yx^{\intercal}+4\langle x,y\rangle
I_{d}\end{pmatrix},$ $\displaystyle\dot{\Sigma}(x,y)$
$\displaystyle=\begin{pmatrix}4\langle x,y\rangle&4x^{\intercal}\\\
4y&4I_{d}\end{pmatrix}.$

We will show the result by showing

$\mathrm{rank}\ \Theta(x^{(1:N)},x^{(1:N)})\leq\mathrm{rank}\
\Sigma(x^{(1:N)},x^{(1:N)})+\mathrm{rank}\
\dot{\Sigma}(x^{(1:N)},x^{(1:N)})<N(d_{0}+1).$

We can bound the rank of individual matrices as

$\displaystyle\mathrm{rank}\ \Sigma(x^{(1:N)},x^{(1:N)})$ $\displaystyle\leq
N(d_{0}+1)-2N+1,$ $\displaystyle\mathrm{rank}\
\dot{\Sigma}(x^{(1:N)},x^{(1:N)})$ $\displaystyle\leq d_{0},$

for some $x^{(1:N)}$ which will be proven in the next two lemmas that follow.
So, setting $N>(d_{0}+1)/2$ shows that $\Theta(x^{(1:N)},x^{(1:N)})$ is
singular, therefore $\lambda_{\min}(\Theta(x^{(1:N)},x^{(1:N)}))=0$. ∎

######  Lemma 19.

For any $x^{(1:N)}$ such that no $x^{(i)}$ is zero vector, we have at least
$2N-1$ linearly independent vectors that

$\Sigma(x^{(1:N)},x^{(1:N)})v=0,$

so that $\mathrm{rank}\ \Sigma(x^{(1:N)},x^{(1:N)})\leq N(d_{0}+1)-N$.

###### Proof.

Let’s first define the following vectors:

$v_{i}=(\underbrace{0,\ldots,0}_{(i-1)\cdot(d_{0}+1)},-2,x_{0}^{(i)},\ldots,x_{d_{0}-1}^{(i)},\underbrace{0,\ldots,0}_{(N-i)\cdot(d_{0}+1)})^{\intercal}$

for $i=1,\ldots,N$. Then these vectors are linearly independent, and for
$i\in[N]$,

$\displaystyle\left(\Sigma(x^{(1:N)},x^{(1:N)})v_{j}\right)_{i(d_{0}+1)}=$
$\displaystyle-4\langle
x^{(i)},x^{(j)}\rangle^{2}-2\left\|x^{(i)}\right\|^{2}\left\|x^{(i)},x^{(j)}\right\|^{2}$
$\displaystyle+4\langle
x^{(i)},x^{(j)}\rangle\left(x^{(i)}\right)^{\intercal}x^{(j)}+2\left\|x^{(i)}\right\|^{2}\left(x^{(j)}\right)^{\intercal}x^{(j)},$
$\displaystyle\left(\Sigma(x^{(1:N)},x^{(1:N)})v_{j}\right)_{i(d_{0}+1)+\alpha}$
$\displaystyle=-8\langle x^{(i)},x^{(j)}\rangle
x_{\alpha}^{(j)}-4\left\|x^{(j)}\right\|^{2}x_{\alpha}^{(i)}$
$\displaystyle+4x_{\alpha}^{(i)}\left(x^{(j)}\right)^{\intercal}x^{(j)}+4x_{\alpha}^{(j)}\left(x^{(i)}\right)^{\intercal}x^{(j)}+4\langle
x^{(i)},x^{(j)}\rangle x_{\alpha}^{(j)},$

showing that they all lies in the null space of $\Sigma(x^{(1:N)},x^{(1:N)})$.

We then define the following vectors:

$w_{i}=(0,-x_{0}^{(i)},\ldots,-x_{d_{0}-1}^{(i)}\underbrace{0,\ldots,0}_{(i-2)\cdot(d_{0}+1)},0,x_{0}^{(1)},\ldots,x_{d_{0}-1}^{(1)},\underbrace{0,\ldots,0}_{(N-i)\cdot(d_{0}+1)})^{\intercal}$

for $i=2,\ldots,N$.

Similarly for $i\in[N]$,

$\displaystyle\left(\Sigma(x^{(1:N)},x^{(1:N)})w_{j}\right)_{i(d_{0}+1)}=$
$\displaystyle-4\langle
x^{(i)},x^{(1)}\rangle\left(x^{(i)}\right)^{\intercal}x^{(j)}-2\left\|x^{(i)}\right\|^{2}\left(x^{(1)}\right)^{\intercal}x^{(j)}$
$\displaystyle+4\langle
x^{(i)},x^{(j)}\rangle\left(x^{(i)}\right)^{\intercal}x^{(1)}+2\left\|x^{(i)}\right\|^{2}\left(x^{(j)}\right)^{\intercal}x^{(1)},$
$\displaystyle\left(\Sigma(x^{(1:N)},x^{(1:N)})w_{j}\right)_{i(d_{0}+1)+\alpha}=$
$\displaystyle-4x_{\alpha}^{(i)}\left(x^{(1)}\right)^{\intercal}x^{(j)}-4x_{\alpha}^{(1)}\left(x^{(i)}\right)^{\intercal}x^{(j)}-4\langle
x^{(i)},x^{(1)}\rangle x_{\alpha}^{(j)}$
$\displaystyle+4x_{\alpha}^{(i)}\left(x^{(j)}\right)^{\intercal}x^{(1)}+4x_{\alpha}^{(j)}\left(x^{(i)}\right)^{\intercal}x^{(1)}+4\langle
x^{(i)},x^{(j)}\rangle x_{\alpha}^{(1)}$

showing that these vectors also lie in the null space of
$\Sigma(x^{(1:N)},x^{(1:N)})$.

It is clear that all vectors $\\{v_{i}\\}_{i=1}^{N}\cup\\{w_{j}\\}_{j=2}^{N}$
are linearly independent. These show that the null space is at least $N+N-1$
dimension, therefore the rank of

$\Sigma(x^{(1:N)},x^{(1:N)})\leq N(d_{0}+1)-2N+1.$

∎

######  Lemma 20.

For any $x^{(1:N)}$, we have

$\mathrm{rank}\ \dot{\Sigma}(x^{(1:N)},x^{(1:N)})\leq d_{0}.$

###### Proof.

Instead of giving the eigendecomposition, we show that there exists $d_{0}$
vectors that

$\dot{\Sigma}(x^{(1:N)},x^{(1:N)})=\sum_{\alpha=0}^{d_{0}-1}\lambda_{\alpha}v_{\alpha}v_{\alpha}^{\intercal}$

for $\lambda_{\alpha}\geq 0$ and $v_{\alpha}\in\mathbb{R}^{N(d_{0}+1)}$.

We define these vectors as

$\displaystyle v_{\alpha}$
$\displaystyle=(x_{\alpha}^{(1)},\underbrace{0,\ldots,0}_{\alpha},1,\underbrace{0,\ldots,0}_{d_{0}-\alpha-1},x_{\alpha}^{(2)},\underbrace{0,\ldots,0}_{\alpha},1,\underbrace{0,\ldots,0}_{d_{0}-\alpha-1},\ldots,x_{\alpha}^{(N)},\underbrace{0,\ldots,0}_{\alpha},1,\underbrace{0,\ldots,0}_{d_{0}-\alpha-1})^{\intercal}$

Then, we have

$\displaystyle\left(v_{\gamma}v_{\gamma}^{\intercal}\right)_{i(d_{0}+1),j(d_{0}+1)}$
$\displaystyle=x_{\gamma}^{(i)}x_{\gamma}^{(j)},$
$\displaystyle\left(v_{\gamma}v_{\gamma}^{\intercal}\right)_{i(d_{0}+1)+\alpha,j(d_{0}+1)}$
$\displaystyle=x_{\gamma}^{(j)}\mathbb{I}_{\alpha=\gamma},$
$\displaystyle\left(v_{\gamma}v_{\gamma}^{\intercal}\right)_{i(d_{0}+1),j(d_{0}+1)+\beta}$
$\displaystyle=x_{\gamma}^{(i)}\mathbb{I}_{\beta=\gamma},$
$\displaystyle\left(v_{\gamma}v_{\gamma}^{\intercal}\right)_{i(d_{0}+1)+\alpha,j(d_{0}+1)+\beta}$
$\displaystyle=\mathbb{I}_{\alpha=\gamma}\mathbb{I}_{\beta=\gamma}$

which shows that taking $\lambda_{\gamma}=4$ proves the identity. ∎

######  Remark 21.

We note that the bound of Lemma 19 is not tight. If we computationally test
this rank condition, the rank of $\Sigma(x^{(1:N)},x^{(1:N)})$ follows

$\mathrm{rank}\
\Sigma(x^{(1:N)},x^{(1:N)})=\sum_{i=0}^{N-1}\max(d_{0}-i,0)=\begin{cases}\frac{N(2d_{0}-N+1)}{2}&\text{if
}d_{0}>N\\\ \frac{d_{0}(d_{0}+1)}{2}&\text{if }d_{0}\leq N\end{cases}$

for non-parallel inputs $x^{(1:N)}$.

Conversely, the bound of Lemma 20 is tight, i.e.,

$\mathrm{rank}\ \dot{\Sigma}(x^{(1:N)},x^{(1:N)})=d_{0}$

for non-parallel inputs $x^{(1:N)}$.

##  Appendix H PROOF OF THEOREM 7

Before giving the proof, we first state the complete form of $\Theta$ in
Theorem 7, which is the $1/\kappa^{2}$-scaled version of the limiting JNTK:

$\displaystyle\Theta(x,x^{\prime})_{00}$
$\displaystyle\coloneqq\sum_{l=0}^{L}\left(\Sigma^{(l)}(x,x^{\prime})_{00}\>\prod_{u=l}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}\right),$
$\displaystyle\Theta(x,x^{\prime})_{\alpha 0}$
$\displaystyle\coloneqq\sum_{l=0}^{L}\Bigg{(}\Sigma^{(l)}(x,x^{\prime})_{\alpha
0}\>\prod_{u=l}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}$
$\displaystyle+\Sigma^{(l)}(x,x^{\prime})_{00}\>\sum_{u_{\alpha}=l}^{L-1}\Gamma^{(u_{\alpha})}(x,x^{\prime})_{\alpha
0}\>\prod_{\begin{subarray}{c}u=l\\\ u\neq
u_{\alpha}\end{subarray}}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}\Bigg{)}$
$\displaystyle\Theta(x,x^{\prime})_{0\beta}$
$\displaystyle\coloneqq\sum_{l=0}^{L}\Bigg{(}\Sigma^{(l)}(x,x^{\prime})_{0\beta}\>\prod_{u=l}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}$
$\displaystyle+\Sigma^{(l)}(x,x^{\prime})_{00}\>\sum_{u_{\beta}=l}^{L-1}\Gamma^{(u_{\beta})}(x,x^{\prime})_{0\beta}\>\prod_{\begin{subarray}{c}u=l\\\
u\neq u_{\beta}\end{subarray}}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}\Bigg{)}$
$\displaystyle\Theta(x,x^{\prime})_{\alpha\beta}$
$\displaystyle\coloneqq\sum_{l=0}^{L}\Bigg{(}\Sigma^{(l)}(x,x^{\prime})_{\alpha\beta}\>\prod_{u=l}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}$
$\displaystyle+\Sigma^{(l)}(x,x^{\prime})_{\alpha
0}\>\sum_{u_{\beta}=l}^{L-1}\Gamma^{(u_{\beta})}(x,x^{\prime})_{0\beta}\>\prod_{\begin{subarray}{c}u=l\\\
u\neq u_{\beta}\end{subarray}}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}$
$\displaystyle+\Sigma^{(l)}(x,x^{\prime})_{0\beta}\>\sum_{u_{\alpha}=l}^{L-1}\Gamma^{(u_{\alpha})}(x,x^{\prime})_{\alpha
0}\>\prod_{\begin{subarray}{c}u=l\\\ u\neq
u_{\alpha}\end{subarray}}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}$
$\displaystyle+\Sigma^{(l)}(x,x^{\prime})_{00}\>\sum_{u_{\alpha\beta}=l}^{L-1}\Gamma^{(u_{\alpha\beta})}(x,x^{\prime})_{\alpha\beta}\>\prod_{\begin{subarray}{c}u=l\\\
u\neq u_{\alpha\beta}\end{subarray}}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}$
$\displaystyle+\Sigma^{(l)}(x,x^{\prime})_{00}\>\sum_{\begin{subarray}{c}u_{\alpha},u_{\beta}=l\\\
u_{\alpha}\neq
u_{\beta}\end{subarray}}^{L-1}\Gamma^{(u_{\alpha})}(x,x^{\prime})_{\alpha
0}\>\Gamma^{(u_{\beta})}(x,x^{\prime})_{0\beta}\>\prod_{\begin{subarray}{c}u=l\\\
u\neq u_{\alpha}\\\ u\neq
u_{\beta}\end{subarray}}^{L-1}\Gamma^{(u)}(x,x^{\prime})_{00}\Bigg{)},$

where
$\Gamma^{(u)}:\mathbb{R}^{d_{0}}\times\mathbb{R}^{d_{0}}\to\mathbb{R}^{(1+d_{0})\times(1+d_{0})}$
is defined as follows: for $g^{(u)}\sim\mathcal{GP}(0,\Sigma^{(u)})$,

$\displaystyle\Gamma^{(u)}(x,x^{\prime})_{00}$
$\displaystyle\coloneqq\mathop{\mathbb{E}}_{g^{(u)}}\left[\dot{\phi}(g^{(u)}(x)_{0})\dot{\phi}(g^{(u)}(x^{\prime})_{0})\right],$
$\displaystyle\Gamma^{(u)}(x,x^{\prime})_{\alpha 0}$
$\displaystyle\coloneqq\mathop{\mathbb{E}}_{g^{(u)}}\left[\ddot{\phi}(g^{(u)}(x)_{0})\dot{\phi}(g^{(u)}(x^{\prime})_{0})g^{(u)}(x)_{\alpha}\right],$
$\displaystyle\Gamma^{(u)}(x,x^{\prime})_{0\beta}$
$\displaystyle\coloneqq\mathop{\mathbb{E}}_{g^{(u)}}\left[\dot{\phi}(g^{(u)}(x)_{0})\ddot{\phi}(g^{(u)}(x^{\prime})_{0})g^{(u)}(x^{\prime})_{\beta}\right],$
$\displaystyle\Gamma^{(u)}(x,x^{\prime})_{\alpha\beta}$
$\displaystyle\coloneqq\mathop{\mathbb{E}}_{g^{(u)}}\left[\ddot{\phi}(g^{(u)}(x)_{0})\ddot{\phi}(g^{(u)}(x^{\prime})_{0})g^{(u)}(x)_{\alpha}g^{(u)}(x^{\prime})_{\beta}\right].$

We prove Theorem 7 by using the proof strategy for Theorem 3 in Section E
again. That is, we express the computation of backpropagated gradient as a
tensor program, and apply Theorem 12 to show the claimed convergence in
Theorem 7.

Recall the definition of the finite JNTK: for $\alpha,\beta\in[d_{0}]$ and
$x,x^{\prime}\in\mathbb{R}^{d_{0}}$,

$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{00}$
$\displaystyle=\left\langle\frac{\partial f(x)}{\partial\theta},\frac{\partial
f(x^{\prime})}{\partial\theta}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{\alpha 0}$
$\displaystyle=\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial\theta},\frac{\partial
f(x^{\prime})}{\partial\theta}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{0\beta}$
$\displaystyle=\left\langle\frac{\partial f(x)}{\partial\theta},\frac{\partial
J(f)(x^{\prime})_{\beta-1}}{\partial\theta}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{\alpha\beta}$
$\displaystyle=\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial\theta},\frac{\partial
J(f)(x^{\prime})_{\beta-1}}{\partial\theta}\right\rangle.$

We rephrase this finite JNTK such that the contribution of each layer is
explicit:

$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{00}$
$\displaystyle=\sum_{l=1}^{L+1}\left\langle\frac{\partial f(x)}{\partial
W^{(l)}},\frac{\partial f(x^{\prime})}{\partial W^{(l)}}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{\alpha 0}$
$\displaystyle=\sum_{l=1}^{L+1}\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial W^{(l)}},\frac{\partial f(x^{\prime})}{\partial
W^{(l)}}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{0\beta}$
$\displaystyle=\sum_{l=1}^{L+1}\left\langle\frac{\partial f(x)}{\partial
W^{(l)}},\frac{\partial J(f)(x^{\prime})_{\beta-1}}{\partial
W^{(l)}}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{\alpha\beta}$
$\displaystyle=\sum_{l=1}^{L+1}\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial W^{(l)}},\frac{\partial
J(f)(x^{\prime})_{\beta-1}}{\partial W^{(l)}}\right\rangle.$ (11)

By the definitions of the MLP $f$ and its components $g^{(l)}$ and $h^{(l)}$
in Section C, for every $l=2,\ldots,L+1$, we have

$\displaystyle\frac{\partial f(x)}{\partial W^{(l)}}$
$\displaystyle=\frac{1}{\sqrt{d}}\frac{\partial f(x)}{\partial
g^{(l)}(x)}\left(h^{(l-1)}(x)\right)^{\intercal},$
$\displaystyle\frac{\partial J(f)(x)_{\alpha-1}}{\partial W^{(l)}}$
$\displaystyle=\frac{1}{\sqrt{d}}\frac{\partial J(f)(x)_{\alpha-1}}{\partial
g^{(l)}(x)}\left(h^{(l-1)}(x)\right)^{\intercal}+\frac{1}{\sqrt{d}}\frac{\partial
J(f)(x)_{\alpha-1}}{\partial(\partial g^{(l)}(x)/\partial
x_{\alpha-1})}\left(\frac{\partial h^{(l-1)}(x)}{\partial
x_{\alpha-1}}\right)^{\intercal}.$

The derivatives with respect to $W^{(1)}$ have similar forms shown below:

$\displaystyle\frac{\partial f(x)}{\partial W^{(1)}}$
$\displaystyle=\frac{\partial f(x)}{\partial
g^{(1)}(x)}\left(x\right)^{\intercal},$ $\displaystyle\frac{\partial
J(f)(x)}{\partial W^{(1)}}$ $\displaystyle=\frac{\partial
J(f)(x)_{\alpha-1}}{\partial
g^{(1)}(x)}\left(x\right)^{\intercal}+\frac{\partial
J(f)(x)_{\alpha-1}}{\partial(\partial g^{(1)}(x)/\partial
x_{\alpha-1})}\left(e_{\alpha-1}\right)^{\intercal}.$

We plug in these characterisations of the derivatives in the layer-wise
description of the finite JNTK in (11), and simplify the results using
$\langle v_{1}w_{1}^{\intercal},v_{2}w_{2}^{\intercal}\rangle=\langle
v_{1},v_{2}\rangle\cdot\langle w_{1},w_{2}\rangle$. Recall that $C_{l}=d$ if
$l=2,\ldots,L+1$ and $C_{1}=1$.

$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{00}$ (12)
$\displaystyle=\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle
h^{(l-1)}(x),h^{(l-1)}(x^{\prime})\right\rangle\left\langle\frac{\partial
f(x)}{\partial g^{(l)}(x)},\frac{\partial f(x^{\prime})}{\partial
g^{(l)}(x^{\prime})}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{\alpha 0}$ (13)
$\displaystyle=\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle
h^{(l-1)}(x),h^{(l-1)}(x^{\prime})\right\rangle\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial g^{(l)}(x)},\frac{\partial
f(x^{\prime})}{\partial g^{(l)}(x^{\prime})}\right\rangle$
$\displaystyle+\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle\frac{\partial
h^{(l-1)}(x)}{\partial
x_{\alpha-1}},h^{(l-1)}(x^{\prime})\right\rangle\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial(\partial g^{(l)}(x)/\partial
x_{\alpha-1})},\frac{\partial f(x^{\prime})}{\partial
g^{(l)}(x^{\prime})}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{0\beta}$ (14)
$\displaystyle=\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle
h^{(l-1)}(x),h^{(l-1)}(x^{\prime})\right\rangle\left\langle\frac{\partial
f(x)_{\alpha-1}}{\partial g^{(l)}(x)},\frac{\partial
J(f)(x^{\prime})_{\beta-1}}{\partial g^{(l)}(x^{\prime})}\right\rangle$
$\displaystyle+\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle
h^{(l-1)}(x),\frac{\partial h^{(l-1)}(x^{\prime})}{\partial
x_{\beta-1}^{\prime}}\right\rangle\left\langle\frac{\partial f(x)}{\partial
g^{(l)}(x)},\frac{\partial J(f)(x^{\prime})_{\beta-1}}{\partial(\partial
g^{(l)}(x^{\prime})/\partial x_{\beta-1}^{\prime})}\right\rangle,$
$\displaystyle\Theta_{d,\theta}(x,x^{\prime})_{\alpha\beta}$ (15)
$\displaystyle=\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle
h^{(l-1)}(x),h^{(l-1)}(x^{\prime})\right\rangle\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial g^{(l)}(x)},\frac{\partial
J(f)(x^{\prime})_{\beta-1}}{\partial g^{(l)}(x^{\prime})}\right\rangle$
$\displaystyle+\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle\frac{\partial
h^{(l-1)}(x)}{\partial
x_{\alpha-1}},h^{(l-1)}(x^{\prime})\right\rangle\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial(\partial g^{(l)}(x)/\partial
x_{\alpha-1})},\frac{\partial J(f)(x^{\prime})_{\beta-1}}{\partial
g^{(l)}(x^{\prime})}\right\rangle$
$\displaystyle+\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle
h^{(l-1)}(x),\frac{\partial h^{(l-1)}(x^{\prime})}{\partial
x_{\beta-1}^{\prime}}\right\rangle\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial g^{(l)}(x)},\frac{\partial
J(f)(x^{\prime})_{\beta-1}}{\partial(\partial g^{(l)}(x^{\prime})/\partial
x_{\beta-1}^{\prime})}\right\rangle$
$\displaystyle+\sum_{l=1}^{L+1}\frac{1}{C_{l-1}}\left\langle\frac{\partial
h^{(l-1)}(x)}{\partial x_{\alpha-1}},\frac{\partial
h^{(l-1)}(x^{\prime})}{\partial
x_{\beta-1}^{\prime}}\right\rangle\left\langle\frac{\partial
J(f)(x)_{\alpha-1}}{\partial(\partial g^{(l)}(x)/\partial
x_{\alpha-1})},\frac{\partial J(f)(x^{\prime})_{\beta-1}}{\partial(\partial
g^{(l)}(x^{\prime})/\partial x_{\beta-1}^{\prime})}\right\rangle.$ (16)

So it is enough to show that all the inner products here converge to
constants. Then, we have the convergence of the finite JNTK at initialisation.
The tensor program in Section E already computes both components of the first
inner product of each summand. Applying the Master theorem to these inner
products normalised by $C_{l-1}$ gives the almost-sure convergence of these
normalised inner products to
$\Sigma^{(l-1)}(x,x^{\prime})_{00},\Sigma^{(l-1)}(x,x^{\prime})_{\alpha
0},\Sigma^{(l-1)}(x,x^{\prime})_{0\beta},\Sigma^{(l-1)}(x,x^{\prime})_{\alpha\beta}$.
Now it remains to show the convergence of the second inner product of each
summand. We do this by expressing the backpropagated gradients in a tensor
program.

We first include a new random vector for the initialisation,
$V^{(L+1)}\in\mathbb{R}^{d}$, which is independent of all other vectors, and
its entries are i.i.d. $\mathcal{N}(0,\kappa^{2})$. The body of this program
first includes the variables that appeared in Section E, and additionally the
following variables:

$\displaystyle
Dh_{1}^{(1)},\ldots,Dh_{N}^{(1)},\ldots,Dh_{1}^{(L)},\ldots,Dh_{N}^{(L)}\in\mathbb{R}^{d},$
$\displaystyle
Dg_{1}^{(1)},\ldots,Dg_{N}^{(1)},\ldots,Dg_{1}^{(L)},\ldots,Dg_{N}^{(L)}\in\mathbb{R}^{d},$
$\displaystyle
D_{\alpha}J_{\alpha}h_{1}^{(1)},\ldots,D_{\alpha}J_{\alpha}h_{N}^{(1)},\ldots,D_{\alpha}J_{\alpha}h_{1}^{(L)},\ldots,D_{\alpha}J_{\alpha}h_{N}^{(L)}\in\mathbb{R}^{d},$
$\displaystyle
D_{\alpha}J_{\alpha}g_{1}^{(1)},\ldots,D_{\alpha}J_{\alpha}g_{N}^{(1)},\ldots,D_{\alpha}J_{\alpha}g_{1}^{(L)},\ldots,D_{\alpha}J_{\alpha}g_{N}^{(L)}\in\mathbb{R}^{d},$
$\displaystyle
D_{\alpha}h_{1}^{(1)},\ldots,D_{\alpha}h_{N}^{(1)},\ldots,D_{\alpha}h_{1}^{(L)},\ldots,D_{\alpha}h_{N}^{(L)}\in\mathbb{R}^{d},$
$\displaystyle
D_{\alpha}g_{1}^{(1)},\ldots,D_{\alpha}g_{N}^{(1)},\ldots,D_{\alpha}g_{1}^{(L)},\ldots,D_{\alpha}g_{N}^{(L)}\in\mathbb{R}^{d}.$

These variables are defined as follows:

$\displaystyle Dh_{1}^{(L)}$ $\displaystyle\coloneqq V^{(L+1)},$
$\displaystyle\cdots$ $\displaystyle Dh_{N}^{(L)}$ $\displaystyle\coloneqq
V^{(L+1)},$ $\displaystyle Dg_{1}^{(L)}$
$\displaystyle\coloneqq\dot{\phi}\left(g_{1}^{(L)}\right)\odot Dh_{1}^{(L)},$
$\displaystyle\cdots$ $\displaystyle Dg_{N}^{(L)}$
$\displaystyle\coloneqq\dot{\phi}\left(g_{N}^{(L)}\right)\odot Dh_{N}^{(L)},$
$\displaystyle Dh_{1}^{(L-1)}$
$\displaystyle\coloneqq\left(V^{(L)}\right)^{\intercal}Dg_{1}^{(L)},$
$\displaystyle\cdots$ $\displaystyle Dh_{N}^{(L-1)}$
$\displaystyle\coloneqq\left(V^{(L)}\right)^{\intercal}Dg_{N}^{(L)},$
$\displaystyle\vdots$ $\displaystyle\vdots$ $\displaystyle\vdots$
$\displaystyle Dg_{1}^{(1)}$
$\displaystyle\coloneqq\dot{\phi}\left(g_{1}^{(1)}\right)\odot Dh_{1}^{(1)},$
$\displaystyle\cdots$ $\displaystyle Dg_{N}^{(1)}$
$\displaystyle\coloneqq\dot{\phi}\left(g_{N}^{(1)}\right)\odot Dh_{N}^{(1)}.$

These random vectors correspond to the backpropagated gradient, but we need
the scaling as

$\displaystyle Dh_{i}^{(l)}$ $\displaystyle=\sqrt{d}\frac{\partial
f(x^{(i)})}{\partial h^{(l)}(x^{(i)})},$ $\displaystyle Dg_{i}^{(l)}$
$\displaystyle=\sqrt{d}\frac{\partial f(x^{(i)})}{\partial g^{(l)}(x^{(i)})}$

where this $\sqrt{d}$ scaling comes from the scaling at output
$\frac{1}{\sqrt{d}}\left(V^{(L+1)}\right)^{\intercal}h_{i}^{(L)}$.

Similarly, we define the backpropagation of the Jacobian gradient, defined as

$\displaystyle D_{\alpha}J_{\alpha}h_{1}^{(L)}$ $\displaystyle\coloneqq
V^{(L+1)},$ $\displaystyle\cdots$ $\displaystyle
D_{\alpha}J_{\alpha}h_{N}^{(L)}$ $\displaystyle\coloneqq V^{(L+1)},$
$\displaystyle D_{\alpha}J_{\alpha}g_{1}^{(L)}$
$\displaystyle\coloneqq\dot{\phi}\left(g_{1}^{(L)}\right)\odot
D_{\alpha}J_{\alpha}h_{1}^{(L)},$ $\displaystyle\cdots$ $\displaystyle
D_{\alpha}J_{\alpha}g_{N}^{(L)}$
$\displaystyle\coloneqq\dot{\phi}\left(g_{N}^{(L)}\right)\odot
D_{\alpha}J_{\alpha}h_{N}^{(L)},$ $\displaystyle
D_{\alpha}J_{\alpha}h_{1}^{(L-1)}$
$\displaystyle\coloneqq\left(V^{(L)}\right)^{\intercal}D_{\alpha}J_{\alpha}g_{1}^{(L)},$
$\displaystyle\cdots$ $\displaystyle D_{\alpha}J_{\alpha}h_{N}^{(L-1)}$
$\displaystyle\coloneqq\left(V^{(L)}\right)^{\intercal}D_{\alpha}J_{\alpha}g_{N}^{(L)},$
$\displaystyle\vdots$ $\displaystyle\vdots$ $\displaystyle\vdots$
$\displaystyle D_{\alpha}J_{\alpha}g_{1}^{(1)}$
$\displaystyle\coloneqq\dot{\phi}\left(g_{1}^{(1)}\right)\odot
D_{\alpha}J_{\alpha}h_{1}^{(1)},$ $\displaystyle\cdots$ $\displaystyle
D_{\alpha}J_{\alpha}g_{N}^{(1)}$
$\displaystyle\coloneqq\dot{\phi}\left(g_{N}^{(1)}\right)\odot
D_{\alpha}J_{\alpha}h_{N}^{(1)},$

and

$\displaystyle D_{\alpha}h_{1}^{(L)}$ $\displaystyle\coloneqq\mathbf{0},$
$\displaystyle\vdots$ $\displaystyle D_{\alpha}h_{N}^{(L)}$
$\displaystyle\coloneqq\mathbf{0},$ $\displaystyle D_{\alpha}g_{1}^{(L)}$
$\displaystyle\coloneqq\ddot{\phi}\left(g_{1}^{(L)}\right)\odot
J_{\alpha}g_{1}^{(L)}\odot
D_{\alpha}J_{\alpha}h_{1}^{(L)}+\dot{\phi}\left(g_{1}^{(L)}\right)\odot
D_{\alpha}h_{1}^{(L)},$ $\displaystyle\vdots$ $\displaystyle
D_{\alpha}g_{N}^{(L)}$
$\displaystyle\coloneqq\ddot{\phi}\left(g_{N}^{(L)}\right)\odot
J_{\alpha}g_{N}^{(L)}\odot
D_{\alpha}J_{\alpha}h_{N}^{(L)}+\dot{\phi}\left(g_{N}^{(L)}\right)\odot
D_{\alpha}h_{N}^{(L)},$ $\displaystyle D_{\alpha}h_{1}^{(L-1)}$
$\displaystyle\coloneqq\left(V^{(L)}\right)^{\intercal}D_{\alpha}g_{1}^{(L)},$
$\displaystyle\vdots$ $\displaystyle D_{\alpha}h_{N}^{(L-1)}$
$\displaystyle\coloneqq\left(V^{(L)}\right)^{\intercal}D_{\alpha}g_{N}^{(L)},$
$\displaystyle\vdots$ $\displaystyle D_{\alpha}g_{1}^{(1)}$
$\displaystyle\coloneqq\ddot{\phi}\left(g_{1}^{(1)}\right)\odot
J_{\alpha}g_{1}^{(1)}\odot
D_{\alpha}J_{\alpha}h_{1}^{(1)}+\dot{\phi}\left(g_{1}^{(1)}\right)\odot
D_{\alpha}h_{1}^{(1)},$ $\displaystyle\vdots$ $\displaystyle
D_{\alpha}g_{N}^{(1)}$
$\displaystyle\coloneqq\ddot{\phi}\left(g_{N}^{(1)}\right)\odot
J_{\alpha}g_{N}^{(1)}\odot
D_{\alpha}J_{\alpha}h_{N}^{(1)}+\dot{\phi}\left(g_{N}^{(1)}\right)\odot
D_{\alpha}h_{N}^{(1)}.$

Again, these random vectors correspond to the backpropagated gradient with
scaling, as

$\displaystyle D_{\alpha}J_{\alpha}g_{i}^{(l)}$
$\displaystyle=\sqrt{d}\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial\left(\frac{\partial
g^{(l)}(x^{(i)})}{\partial x_{\alpha-1}^{(i)}}\right)},$ $\displaystyle
D_{\alpha}J_{\alpha}h_{i}^{(l)}$ $\displaystyle=\sqrt{d}\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial\left(\frac{\partial
h^{(l)}(x^{(i)})}{\partial x_{\alpha-1}^{(i)}}\right)},$ $\displaystyle
D_{\alpha}g_{i}^{(l)}$ $\displaystyle=\sqrt{d}\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial g^{(l)}(x^{(i)})},$ $\displaystyle
D_{\alpha}h_{i}^{(l)}$ $\displaystyle=\sqrt{d}\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial h^{(l)}(x^{(i)})}.$

The outputs of the program whose convergence is guaranteed by Theorem 12 are
followings: for $i,j\in[N]$, $\alpha,\beta\in[d_{0}]$ and $l\in[N]$,

$\displaystyle\left\langle\frac{\partial f(x^{(i)})}{\partial
g^{(l)}(x^{(i)})},\frac{\partial f(x^{(j)})}{\partial
g^{(l)}(x^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right],$
(17) $\displaystyle\left\langle\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial g^{(l)}(x^{(i)})},\frac{\partial
f(x^{(j)})}{\partial g^{(l)}(x^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right],$
(18) $\displaystyle\left\langle\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial(\partial g^{(l)}(x^{(i)})/\partial
x_{\alpha-1}^{(i)})},\frac{\partial f(x^{(j)})}{\partial
g^{(l)}(x^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}g_{i}^{(l)}}Z^{g_{j}^{(l)}}\right],$
(19) $\displaystyle\left\langle\frac{\partial f(x^{(i)})_{\alpha-1}}{\partial
g^{(l)}(x^{(i)})},\frac{\partial J(f)(x^{(j)})_{\beta-1}}{\partial
g^{(l)}(x^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{D_{\beta}g_{j}^{(l)}}\right],$
(20) $\displaystyle\left\langle\frac{\partial f(x^{(i)})}{\partial
g^{(l)}(x^{(i)})},\frac{\partial J(f)(x^{(j)})_{\beta-1}}{\partial(\partial
g^{(l)}(x^{(j)})/\partial x_{\beta-1}^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{D_{\beta}J_{\beta}g_{j}^{(l)}}\right],$
(21) $\displaystyle\left\langle\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial g^{(l)}(x^{(i)})},\frac{\partial
J(f)(x^{(j)})_{\beta-1}}{\partial g^{(l)}(x^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{D_{\beta}g_{j}^{(l)}}\right],$
(22) $\displaystyle\left\langle\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial(\partial g^{(l)}(x^{(i)})/\partial
x_{\alpha-1}^{(i)})},\frac{\partial J(f)(x^{(j)})_{\beta-1}}{\partial
g^{(l)}(x^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}g_{i}^{(l)}}Z^{D_{\beta}g_{j}^{(l)}}\right],$
(23) $\displaystyle\left\langle\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial g^{(l)}(x^{(i)})},\frac{\partial
J(f)(x^{(j)})_{\beta-1}}{\partial(\partial g^{(l)}(x^{(j)})/\partial
x_{\beta-1}^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{D_{\beta}J_{\beta}g_{j}^{(l)}}\right],$
(24) $\displaystyle\left\langle\frac{\partial
J(f)(x^{(i)})_{\alpha-1}}{\partial(\partial g^{(l)}(x^{(i)})/\partial
x_{\alpha-1}^{(i)})},\frac{\partial J(f)(x^{(j)})_{\beta-1}}{\partial(\partial
g^{(l)}(x^{(j)})/\partial x_{\beta-1}^{(j)})}\right\rangle$
$\displaystyle\to\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}g_{i}^{(l)}}Z^{D_{\beta}J_{\beta}g_{j}^{(l)}}\right].$
(25)

Our final computation is evaluating these nine expectations, which can be done
recursively. Before proceeding further, let’s focus on two variables,
$D_{\alpha}J_{\alpha}g_{i}^{(l)}$ and $Dg_{i}^{(l)}$. Their initialisation
$D_{\alpha}J_{\alpha}h_{i}^{(L)}$ and $Dh_{i}^{(L)}$ are the same, and the
recursive definition also matches, so they are equal always. This implies that

$\displaystyle Dg_{i}^{(l)}$ $\displaystyle=D_{\alpha}J_{\alpha}g_{i}^{(l)},$
$\displaystyle Dg_{j}^{(l)}$ $\displaystyle=D_{\beta}J_{\beta}g_{j}^{(l)},$
$\displaystyle Dh_{i}^{(l)}$ $\displaystyle=D_{\alpha}J_{\alpha}h_{i}^{(l)},$
$\displaystyle Dh_{j}^{(l)}$ $\displaystyle=D_{\beta}J_{\beta}h_{j}^{(l)}.$

This implies that most of the expectations are equal, thereby we only need to
compute three expectations, (17), (18), and (22). All other expectations can
be reduced to these three expectations, by equivalence we just shown or
symmetry.

Before giving the proof of all equations, we first handle simple cases.
Equation (17) is a term that also appears in the standard NTK, and can be
computed with the following recurrence relation,

$\displaystyle\mathop{\mathbb{E}}\left[Z^{Dh_{i}^{(L)}}Z^{Dh_{j}^{(L)}}\right]$
$\displaystyle=\kappa^{2},$ (26)
$\displaystyle\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right]$
$\displaystyle=\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)Z^{Dh_{i}^{(l)}}\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)Z^{Dh_{j}^{(l)}}\right]$
$\displaystyle=\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)\right]\times\mathop{\mathbb{E}}\left[Z^{Dh_{i}^{(l)}}Z^{Dh_{j}^{(l)}}\right],$
(27)
$\displaystyle\mathop{\mathbb{E}}\left[Z^{Dh_{i}^{(l-1)}}Z^{Dh_{j}^{(l-1)}}\right]$
$\displaystyle=\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right]$
(28)

where we used the independency of $Z$-variables $Z^{g_{i}^{(l)}}$ and
$Z^{Dh_{i}^{(l)}}$, this gives

$\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right]=\Gamma^{(l)}(x^{(i)},x^{(j)})_{00}\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l+1)}}Z^{Dg_{j}^{(l+1)}}\right],$

and resolving recurrence relations shows

$\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right]=\kappa^{2}\prod_{u=l}^{L-1}\Gamma^{(u)}(x^{(i)},x^{(j)})_{00}.$

Now we can extend this approach to all the variables. The first steps for (26)
are

$\displaystyle\mathop{\mathbb{E}}\left[Z^{Dh_{i}^{(L)}}Z^{Dh_{j}^{(L)}}\right]$
$\displaystyle=\kappa^{2},$
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(L)}}Z^{Dh_{j}^{(L)}}\right]$
$\displaystyle=0,$
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(L)}}Z^{Dh_{j}^{(L)}}\right]$
$\displaystyle=\kappa^{2},$
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(L)}}Z^{D_{\beta}h_{j}^{(L)}}\right]$
$\displaystyle=0,$
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(L)}}Z^{D_{\beta}h_{j}^{(L)}}\right]$
$\displaystyle=0,$
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(L)}}Z^{D_{\beta}J_{\beta}h_{j}^{(L)}}\right]$
$\displaystyle=\kappa^{2}.$

For the recursion step for the activation application corresponding to (27),
we have

$\displaystyle\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)Z^{Dh_{i}^{(l)}}\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)Z^{Dh_{j}^{(l)}}\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)\right]\mathop{\mathbb{E}}\left[Z^{Dh_{i}^{(l)}}Z^{Dh_{j}^{(l)}}\right]$
$\displaystyle=$
$\displaystyle\Gamma^{(l)}(x^{(i)},x^{(j)})_{00}\mathop{\mathbb{E}}\left[Z^{Dh_{i}^{(l)}}Z^{Dh_{j}^{(l)}}\right],$
(29)
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}\left[\left(\ddot{\phi}\left(Z^{g_{i}^{(l)}}\right)\cdot
Z^{J_{\alpha}g_{i}^{(l)}}\cdot
Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}+\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)\cdot
Z^{D_{\alpha}h_{i}^{(l)}}\right)\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)Z^{Dh_{j}^{(l)}}\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}\left[\ddot{\phi}\left(Z^{g_{i}^{(l)}}\right)\cdot
Z^{J_{\alpha}g_{i}^{(l)}}\cdot\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)\right]\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}Z^{Dh_{j}^{(l)}}\right]$
$\displaystyle\qquad+\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)\right]\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{Dh_{j}^{(l)}}\right]$
(30) $\displaystyle=$ $\displaystyle\Gamma^{(l)}(x^{(i)},x^{(j)})_{\alpha
0}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}Z^{Dh_{j}^{(l)}}\right]+\Gamma^{(l)}(x^{(i)},x^{(j)})_{00}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{Dh_{j}^{(l)}}\right],$
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{D_{\beta}g_{j}^{(l)}}\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}\bigg{[}\left(\ddot{\phi}\left(Z^{g_{i}^{(l)}}\right)\cdot
Z^{J_{\alpha}g_{i}^{(l)}}\cdot
Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}+\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)\cdot
Z^{D_{\alpha}h_{i}^{(l)}}\right)$
$\displaystyle\qquad\times\left(\ddot{\phi}\left(Z^{g_{j}^{(l)}}\right)\cdot
Z^{J_{\beta}g_{j}^{(l)}}\cdot
Z^{D_{\beta}J_{\beta}h_{j}^{(l)}}+\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)\cdot
Z^{D_{\beta}h_{j}^{(l)}}\right)\bigg{]}$ $\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}\left[\ddot{\phi}\left(Z^{g_{i}^{(l)}}\right)Z^{J_{\alpha}g_{i}^{(l)}}\ddot{\phi}\left(Z^{g_{j}^{(l)}}\right)Z^{J_{\beta}g_{j}^{(l)}}\right]\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}J_{\beta}h_{j}^{(l)}}\right]$
$\displaystyle\quad+\mathop{\mathbb{E}}\left[\ddot{\phi}\left(Z^{g_{i}^{(l)}}\right)Z^{J_{\alpha}g_{i}^{(l)}}\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)\right]\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}h_{j}^{(l)}}\right]$
$\displaystyle\quad+\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)\ddot{\phi}\left(Z^{g_{j}^{(l)}}\right)Z^{J_{\beta}g_{j}^{(l)}}\right]\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}J_{\beta}h_{j}^{(l)}}\right]$
$\displaystyle\quad+\mathop{\mathbb{E}}\left[\dot{\phi}\left(Z^{g_{i}^{(l)}}\right)\dot{\phi}\left(Z^{g_{j}^{(l)}}\right)\right]\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}h_{j}^{(l)}}\right]$
$\displaystyle=$
$\displaystyle\Gamma^{(l)}(x^{(i)},x^{(j)})_{\alpha\beta}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}J_{\beta}h_{j}^{(l)}}\right]+\Gamma^{(l)}(x^{(i)},x^{(j)})_{\alpha
0}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}h_{j}^{(l)}}\right]$
$\displaystyle\quad+\Gamma^{(l)}(x^{(i)},x^{(j)})_{0\beta}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}J_{\beta}h_{j}^{(l)}}\right]+\Gamma^{(l)}(x^{(i)},x^{(j)})_{00}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}h_{j}^{(l)}}\right].$
(31)

Finally for the recursion step corresponding to matrix multiplication (28), we
have

$\displaystyle\mathop{\mathbb{E}}\left[Z^{Dh_{i}^{(l-1)}}Z^{Dh_{j}^{(l-1)}}\right]$
$\displaystyle=\mathop{\mathbb{E}}\left[Z^{Dg_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right],$
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l-1)}}Z^{Dh_{j}^{(l-1)}}\right]$
$\displaystyle=\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right],$
$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l-1)}}Z^{D_{\beta}h_{j}^{(l-1)}}\right]$
$\displaystyle=\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{D_{\beta}g_{j}^{(l)}}\right].$

Our final task is solving the recurrence relations,

$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{Dg_{j}^{(l)}}\right]$
$\displaystyle=$ $\displaystyle\kappa^{2}\Gamma^{(l)}(x^{(i)},x^{(j)})_{\alpha
0}\prod_{u=l+1}^{L-1}\Gamma^{(u)}(x^{(i)},x^{(j)})_{00}+\Gamma^{(l)}(x^{(i)},x^{(j)})_{00}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{Dh_{j}^{(l)}}\right]$
$\displaystyle=$
$\displaystyle\kappa^{2}\sum_{u_{\alpha}=l}^{L-1}\Gamma^{(u_{\alpha})}(x^{(i)},x^{(j)})_{\alpha
0}\prod_{\begin{subarray}{c}u=l\\\ u\neq
u_{\alpha}\end{subarray}}^{L-1}\Gamma^{(u)}(x^{(i)},x^{(j)})_{00}$

and

$\displaystyle\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l)}}Z^{D_{\beta}g_{j}^{(l)}}\right]$
$\displaystyle=\Gamma^{(l)}(x^{(i)},x^{(j)})_{\alpha\beta}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}J_{\beta}h_{j}^{(l)}}\right]+\Gamma^{(l)}(x^{(i)},x^{(j)})_{\alpha
0}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}J_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}h_{j}^{(l)}}\right]$
$\displaystyle\quad+\Gamma^{(l)}(x^{(i)},x^{(j)})_{0\beta}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}J_{\beta}h_{j}^{(l)}}\right]+\Gamma^{(l)}(x^{(i)},x^{(j)})_{00}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}h_{i}^{(l)}}Z^{D_{\beta}h_{j}^{(l)}}\right]$
$\displaystyle=\kappa^{2}\Gamma^{(l)}(x^{(i)},x^{(j)})_{\alpha\beta}\prod_{u=l+1}^{L-1}\Gamma^{(u)}(x^{(i)},x^{(j)})_{00}$
$\displaystyle\quad+\kappa^{2}\Gamma^{(l)}(x^{(i)},x^{(j)})_{\alpha
0}\sum_{u_{\beta}=l+1}^{L-1}\Gamma^{(u_{\beta})}(x^{(i)},x^{(j)})_{0\beta}\prod_{\begin{subarray}{c}u=l\\\
u\neq u_{\beta}\end{subarray}}^{L-1}\Gamma^{(u)}(x^{(i)},x^{(j)})_{00}$
$\displaystyle\quad+\kappa^{2}\Gamma^{(l)}(x^{(i)},x^{(j)})_{0\beta}\sum_{u_{\alpha}=l+1}^{L-1}\Gamma^{(u_{\alpha})}(x^{(i)},x^{(j)})_{\alpha
0}\prod_{\begin{subarray}{c}u=l\\\ u\neq
u_{\alpha}\end{subarray}}^{L-1}\Gamma^{(u)}(x^{(i)},x^{(j)})_{00}$
$\displaystyle\quad+\Gamma^{(l)}(x^{(i)},x^{(j)})_{00}\mathop{\mathbb{E}}\left[Z^{D_{\alpha}g_{i}^{(l+1)}}Z^{D_{\beta}g_{j}^{(l+1)}}\right]$
$\displaystyle=\kappa^{2}\sum_{u^{\prime}=l}^{L-1}\Gamma^{(u^{\prime})}(x^{(i)},x^{(j)})_{\alpha\beta}\prod_{\begin{subarray}{c}u=l\\\
u\neq l^{\prime}\end{subarray}}^{L-1}\Gamma^{(u)}(x^{(i)},x^{(j)})_{00}$
$\displaystyle\quad+\kappa^{2}\sum_{\begin{subarray}{c}u_{\alpha}=l\\\
u_{\beta}=l\\\ u_{\alpha}\neq
u_{\beta}\end{subarray}}\Gamma^{(u_{\alpha})}(x^{(i)},x^{(j)})_{\alpha
0}\Gamma^{(u_{\beta})}(x^{(i)},x^{(j)})_{0\beta}\prod_{\begin{subarray}{c}u=l\\\
u\neq u_{\alpha}\\\ u\neq
u_{\beta}\end{subarray}}^{L-1}\Gamma^{(u)}(x^{(i)},x^{(j)})_{00}.$

Together with some trivial computations:

$\displaystyle\left\langle h^{(0)}(x^{(i)}),h^{(0)}(x^{(j)})\right\rangle$
$\displaystyle=\langle x^{(i)},x^{(j)}\rangle,$
$\displaystyle\left\langle\frac{\partial h^{(0)}(x^{(i)})}{\partial
x_{\alpha-1}^{(i)}},h^{(0)}(x^{(j)})\right\rangle$ $\displaystyle=\langle
e_{\alpha-1},x^{(j)}\rangle,$ $\displaystyle\left\langle
h^{(0)}(x^{(i)}),\frac{\partial h^{(0)}(x^{(j)})}{\partial
x_{\beta-1}^{(j)}}\right\rangle$ $\displaystyle=\langle
x^{(i)},e_{\beta-1}\rangle,$ $\displaystyle\left\langle\frac{\partial
h^{(0)}(x^{(i)})}{\partial x_{\alpha-1}^{(i)}},\frac{\partial
h^{(0)}(x^{(j)})}{\partial x_{\beta-1}^{(j)}}\right\rangle$
$\displaystyle=\langle e_{\alpha-1},e_{\beta-1}\rangle,$
$\displaystyle\left\langle\frac{\partial f(x)}{\partial
g^{(L+1)}(x)},\frac{\partial f(x)}{\partial g^{(L+1)}(x)}\right\rangle$
$\displaystyle=\kappa^{2}$

we can show the convergence of every inner product in the layer-wise
description of the finite JNTK, which proves the desired convergence of the
finite JNTK.

To utilise this theorem, we need a version that gives a convergence rate.
Since the tensor-program framework does not give an explicit convergence rate,
we will use an implicit function to give a convergence rate.

######  Remark 22.

There exists function $F_{2}(\cdot,\cdot,\cdot,\cdot)$, so that if the width
of network $d$ is large enough,

$d\geq F_{2}\left(x^{(1:N)},L,\delta,\epsilon\right),$

then the finite JNTK and the limiting JNTK satisfy

$\left\|\Theta_{d,\theta_{0}}(x^{(i)},x^{(j)})-\kappa^{2}\Theta(x^{(i)},x^{(j)})\right\|_{F}\leq\kappa^{2}\epsilon$

for all $i,j\in[N]$, with probability at least $1-\delta$.

##  Appendix I FINITE JNTK AFTER PERTURBATION

######  Theorem 23.

Let $\theta_{0}=\\{W_{0}^{(l)}\\}_{l=1}^{L+1}$ be the weights of the $L$-layer