# Weighted Neural Tangent Kernel: A Generalized and Improved Network-Induced
Kernel

Lei Tan  Shutong Wu  Xiaolin Huang

###### Abstract

The Neural Tangent Kernel (NTK) has recently attracted intense study, as it
describes the evolution of an over-parameterized Neural Network (NN) trained
by gradient descent. However, it is now well-known that gradient descent is
not always a good optimizer for NNs, which can partially explain the
unsatisfactory practical performance of the NTK regression estimator. In this
paper, we introduce the Weighted Neural Tangent Kernel (WNTK), a generalized
and improved tool, which can capture an over-parameterized NN’s training
dynamics under different optimizers. Theoretically, in the infinite-width
limit, we prove: i) the stability of the WNTK at initialization and during
training, and ii) the equivalence between the WNTK regression estimator and
the corresponding NN estimator with different learning rates on different
parameters. With the proposed weight update algorithm, both empirical and
analytical WNTKs outperform the corresponding NTKs in numerical experiments.

Machine Learning, ICML

††Shanghai Jiao Tong University, Shanghai, China. Correspondence to: Lei Tan
$<$<EMAIL_ADDRESS>$>$, Shutong Wu $<$<EMAIL_ADDRESS>$>$, Xiaolin
Huang $<$<EMAIL_ADDRESS>$>$.

##  1 Introduction

Neural Networks (NNs) are increasingly gaining popularity due to their
powerful flexibility in model fitting and their impressive performance in many
applications. However, owing to their nested and nonlinear architectures, the
theoretical study on NNs is relatively lacking. In recent years, the Neural
Tangent Kernel (NTK) has been theoretically proved to be an excellent tool to
describe a NN in the infinite-width limit using the over-parameterization
property (Jacot et al., 2018). An over-parameterized NN trained by gradient
descent is equivalent to the corresponding NTK regression estimator (Arora et
al., 2019a). Some regularization methods of NNs can also be transferred into
regularization terms of the NTK regression estimator (Hu et al., 2019).

Although the NTK has built up a promising bridge between theoretical analysis
and practice learning algorithms, the performance gap is still obvious, as
shown by Arora et al. (2019a): Accuracy of CNTKs (NTKs that correspond to
Convolutional Neural Networks) is always about 5% lower than that of the
corresponding Convolutional Neural Networks (CNNs). Many modifications have
been proposed to fulfill this gap (Hanin & Nica, 2019; Chen et al., 2020).
Nevertheless, those modifications generally depart from the initial idea that
the NTK is a tool to simulate the NN’s training phase dynamics.

In this paper, we generalize the original NTK to improve its classification
performance while keeping a close link to NN training. The original NTK
describes the dynamics of the training phase of an over-parameterized NN using
gradient descent as the optimizer. However, for an under-parameterized NN, it
has been well-known that gradient descent is not always a good optimizer.
Instead, there have been many efficient algorithms, e.g., SGD with Momentum
(Sutskever et al., 2013; Cutkosky & Mehta, 2020), AdaGrad (Zeiler, 2012), Adam
(Zhang, 2018). Those modifications multiply certain adjustment matrices on the
gradient descent direction, which should be correspondingly reflected in the
NTK.

Our idea is to introduce weight terms on the NTK, resulting in the Weighted
Neural Tangent Kernel (WNTK). Consider a fully connected network with input
dimension $d_{0}\in{\mathbb{N}^{+}}$ and width $d_{l}\in{\mathbb{N}^{+}}$ at
the $l$-th layer. Let $f^{(l)}$ denote the output of the $l$-th layer of the
NN and let all activation functions of different layers be a coordinate-wise
one $\sigma$. $P$ is the number of parameters and $\theta\in{\mathbb{R}^{P}}$
represents all the parameters. For each input sample $x\in\mathbb{R}^{d_{0}}$,
the recurrence relation for this NN architecture is defined as

$f^{(l)}(\theta,x)=W^{(l)}\sqrt{\frac{1}{d_{l-1}}}\sigma(f^{(l-1)}(\theta,x)),$
(1)

where $W^{(l)}\in{\mathbb{R}}^{d_{l}\times d_{l-1}}$ is the weight matrix
between the $(l-1)$-th and the $l$-th layer. We formally give the WNTK:

$\mathcal{A}(x_{i},x_{j})=\sum_{p=1}^{P}a_{p}(\partial_{\theta_{p}}f^{(L)}(\theta,x_{i}))^{\top}\partial_{\theta_{p}}f^{(L)}(\theta,x_{j}),$
(2)

where $a_{p}\in{\mathbb{R}}$. Intuitively, the introduced weights could model
the adjustment on gradient descent and thus improve the classification
performance. Simultaneously, the WNTK reserves two theoretical properties: i)
stability at initialization and during training; ii) equivalence between a
WNTK regression estimator and the corresponding NN estimator. These two
properties are crucial for linking NNs and NTKs. However, unlike NTKs locked
to fixed descent direction, WNTKs can describe different optimizers.

There are other ways to understand the weight terms.

1\. Consider a NN under gradient descent. Theoretical works have shown that
techniques, like Dropout and Freezing Layer, can accelerate its convergence,
enhance its generalization ability and improve its performance (Srivastava et
al., 2014; Brock et al., 2017). The essence of such techniques is to set
different learning rates on different parameters, which is equivalent to
applying adjusted gradient descent. Therefore, weight terms in the WNTK can
also be interpreted as learning rates on parameters. A theoretical explanation
on this link between weight terms and learning rates could be found in Section
4.

2\. From the view of kernel learning, the NTK could be regarded as an equal-
weighted sum of multiple kernels, each of which describes the gradient on each
layers, i.e.,

$\displaystyle\ \ \Theta^{(L)}_{l}(x_{i},x_{j})$ (3) $\displaystyle=$
$\displaystyle\sum_{\theta_{p}\ \mathrm{in\ layer}\
l}(\partial_{\theta_{p}}f^{(L)}(\theta,x_{i}))^{\top}\partial_{\theta_{p}}f^{(L)}(\theta,x_{j})$

where $x_{i},x_{j}$ are two input samples in the training set. Many
discussions on multiple kernel learning (Sonnenburg et al., 2006; Gönen &
Alpaydın, 2011) imply that giving different weights, just like the WNTK

$\mathcal{A}(x_{i},x_{j})=\sum_{l=1}^{L}a^{L}_{l}\Theta^{(L)}_{l}(x_{i},x_{j}),$
(4)

could improve the performance. For these layer-wise weights, we establish an
update algorithm to find a good WNTK, see Algorithm 1. One can also element-
wisely adjust the weights, which falls into the idea of non-parametric kernel
learning (Zhuang et al., 2011; Liu et al., 2020) that learns the element
values.

To summarize, this paper proposes the WNTK, analyzes its properties, designs a
learning algorithm, and evaluates the performance of both empirical and
analytic kernels. The main contributions are listed below:

1\. We introduce the Weighted Neural Tangent Kernel, a generalized and
improved tool to simulate over-parameterized NN dynamics. Compared to the NTK,
the WNTK could describe the dynamics of the NN trained by different
optimizers.

2\. In the infinite-width limit, we prove that the WNTK is stable both at
initialization (Theorem 1) and during the corresponding NN training phase
(Theorem 2) under adjusted gradient descent.

3\. In the infinite-width limit, we prove the equivalence between a WNTK
regression estimator and the corresponding NN estimator under adjusted
gradient descent (Theorem 4).

4\. We design a weight update algorithm (Algorithm 1) and verify in the
experiment that the WNTK does have significant improvement from the NTK.

The rest of the paper is organized as follows. Section 2 summarizes the
research background and the related work. Section 3 reviews the derivation of
the NTK and introduces the notations. Section 4 reveals the connection between
the WNTK and NN training dynamics. Section 5 establishes a weight update
algorithm for learning a good WNTK. Section 6 presents experimental results.

##  2 Related Work

Initially, the correspondence between wide NNs and kernel machines was noted
by Neal’s pioneering work (2012), which was followed by a series of papers
(Williams, 1997; Lee et al., 2017; Hazan & Jaakkola, 2015; Matthews et al.,
2018; Novak et al., 2018; Garriga-Alonso et al., 2018). Meanwhile, over-
parameterized NNs showed stunning performance and strong generalization
ability in practice, which gave rise to many discussions on the nature of
over-parameterization (Arora et al., 2019b; Allen-Zhu et al., 2018; Chizat et
al., 2018; Li & Liang, 2018; Ghorbani et al., 2019a; Arora et al., 2018). It
has been shown that, under the infinite width assumption, the evolution of a
wide NN of any depth possesses the same dynamics as a linearized model (Lee et
al., 2019). Based on all these discussions, several different NN-induced
kernels were proposed as tools to explain the behavior of NNs (Daniely et al.,
2016; Shankar et al., 2020; Anselmi et al., 2015), where belong the NTK (Jacot
et al., 2018) and its derivations (Arora et al., 2019a; Li et al., 2019; Du et
al., 2019; Alemohammad et al., 2020). The most remarkable property of the NTK
is that the NTK predictor is equivalent to the corresponding NN predictor
trained by gradient descent. As the inner connection between kernel machines
and NNs is progressively being revealed (Daniely et al., 2016; Belkin et al.,
2018), the NTK is considered as a bridge between NNs and traditional kernel
machines.

Although Du et al. (2018) prove that gradient descent optimizes over-
parameterized wide NNs, the performance of the NTK in practice is relatively
unsatisfactory compared to NNs (Samarin et al., 2020; Ghorbani et al., 2019b;
Wei et al., 2019). The NTK can outperform the corresponding NN only on small
datasets (Arora et al., 2019c). On large datasets, there exists a noticeable
performance gap between them (Arora et al., 2019a), which triggers many works
on improving the NTK’s performance, like adding label-awareness (Chen et al.,
2020), conducting finite-width correction (Hanin & Nica, 2019), and creating
neural tangent hierarchy (Huang & Yau, 2020).

The main idea of this paper is to generalize the NTK to discuss different
optimizers rather than gradient descent. It has been intensively discussed
that Adam (Zhang, 2018), AdaGrad (Zeiler, 2012), SGD with Momentum (Sutskever
et al., 2013; Cutkosky & Mehta, 2020), Etc, have better performance than
gradient descent, at least on under-parameterized NNs.

##  3 Recap of NTK

Let us focus on the NN architecture described by (1). Consider an $L$-layer
NN. The corresponding NTK is defined as follows:

$\displaystyle\Theta(x_{i},x_{j})$
$\displaystyle=\Theta^{(L)}(\theta)(x_{i},x_{j})$ (5)
$\displaystyle=\sum_{p=1}^{P}(\partial_{\theta_{p}}f^{(L)}(\theta,x_{i}))^{\top}\partial_{\theta_{p}}f^{(L)}(\theta,x_{j}),$

The two most significant works on the NTK are: i) It converges in probability
to a closed-form kernel at initialization; ii) It stays constant during
training if the NN is sufficiently wide (Jacot et al., 2018). By default,
$\Theta$ and $\mathcal{A}$ refer to the NTK and the WNTK on the training set,
and $\dot{\sigma}$ represents the derivative of $\sigma$. The following
formula shows the stable closed-form kernel induced by an $L$-layer dense
network

$\displaystyle\Theta\to\Theta^{(L)}_{\infty}\otimes Id_{d_{L}},$ (6)

where $\otimes$ stands for the Kronecker product, $Id_{d_{L}}$ denotes the
identity matrix of size $d_{L}$, and $\Theta^{(L)}_{\infty}$ is defined
recursively by

$\displaystyle\Theta^{(1)}_{\infty}(x_{i},x_{j})=$ $\displaystyle\
\Sigma^{(1)}(x_{i},x_{j}),$ $\displaystyle\Theta^{(l)}_{\infty}(x_{i},x_{j})=$
$\displaystyle\
\Theta^{(l-1)}_{\infty}(x_{i},x_{j})\dot{\Sigma}^{(l)}(x_{i},x_{j})$
$\displaystyle\ +\Sigma^{(l)}(x_{i},x_{j}),$
$\displaystyle\Sigma^{(1)}(x_{i},x_{j})=$ $\displaystyle\
\frac{1}{d_{0}}x^{\top}_{i}x_{j},$ $\displaystyle\Sigma^{(l)}(x_{i},x_{j})=$
$\displaystyle\
\mathbb{E}_{f\sim\mathcal{N}(0,\Sigma^{(l-1)})}[\sigma(f(x_{i}))\sigma(f(x_{j}))],$
$\displaystyle\dot{\Sigma}^{(l)}(x_{i},x_{j})=$ $\displaystyle\
\mathbb{E}_{f\sim\mathcal{N}(0,\Sigma^{(l-1)})}[\dot{\sigma}(f(x_{i}))\dot{\sigma}(f(x_{j}))].$

Under the assumption that $d_{1}=...=d_{L-1}=d$ and $d_{L}=1$, the following
formula reflects the stability of the NTK during the training phase of an
over-parameterized NN, or equivalently, the convergence between the NTKs
$\Theta_{0}$ and $\Theta_{t}$ at moment $0$ and $t$:

$\displaystyle\forall t,||\Theta_{0}-\Theta_{t}||_{F}\to 0,as\ d\to\infty.$
(7)

Under the infinite width assumption, the fully trained NN estimator
$f^{\star}$ can be proved to be equivalent to a NTK regression estimator
(Jacot et al., 2018), i.e.,

$\displaystyle
f^{\star}(\theta,x)\to\Theta(x,\mathcal{X})^{\top}(\Theta(\mathcal{X},\mathcal{X}))^{-1}\mathcal{Y},$
(8)

where $(\mathcal{X},\mathcal{Y})$ stands for the training set of size $n$ and
$x$ is a sample.

##  4 Theory: WNTK and Neural Network Training dynamics

Intuitively, the WNTK is a generalization of the NTK and in this section we
will show that this generalization enables us to describe the training process
under adjusted gradient descent using different optimizers. For the NN
architecture described in (1), we reshape the formula of the WNTK (2) as the
following,

$\displaystyle\mathcal{A}(x_{i},x_{j})=(\partial_{\theta}f^{(L)}(\theta,x_{i}))^{\top}(a\odot\partial_{\theta}f^{(L)}(\theta,x_{j})).$
(9)

If we assume $d_{L}=1$, we have,

$\displaystyle\mathcal{A}=\mathcal{A}(\mathcal{X},\mathcal{X})=\nabla_{\theta}f(\mathcal{X})(a\odot\nabla_{\theta}f(\mathcal{X})^{\top}),$
(10)

where $a=[a_{1},a_{2},...,a_{P}]^{\top}$, $a_{p}\in{\mathbb{R}}$,
$\nabla_{\theta}f(\mathcal{X})\in{\mathbb{R}}^{n\times P}$ represents the
gradient of $f$ on $\theta$ with respect to the training set $\mathcal{X}$,
and $\odot$ refers to the penetrating face product.

###  4.1 Stability Analysis of WNTK

In this subsection, we theoretically show that the WNTK captures the behavior
of an over-parameterized NN fully-trained by a non-equally weighted optimizer.
In the infinite-width limit, the WNTK, of which the corresponding NN is
randomly initialized, converges in probability to a deterministic limit.

###### Theorem 1

(Stability at Initialization). Consider a network of depth L at
initialization, with a Lipschitz nonlinearity $\sigma$. Suppose that, for
every parameter $\theta_{p}$ in the $l$-th layer, its weight term in the WNTK
follows a layer-wise distribution $\rho_{l}$. In the limit as the layer width
$d_{1},...,d_{L-1}$ tends to infinity sequentially, the layer-wise WNTK
converges in probability to a deterministic limiting kernel

$\displaystyle\mathcal{A}\to\mathcal{A}^{(L)}_{\infty}\otimes Id_{d_{L}}.$
(11)

With $\mu_{l}$ referring to the expectation of all element-wise weights which
follow the layer-wise distribution $\rho_{l}$, the scalar kernel
$\mathcal{A}^{(L)}_{\infty}:{\mathbb{R}}^{d_{0}}\times{\mathbb{R}}^{d_{0}}\to\mathbb{R}$
is defined recursively by

$\displaystyle\mathcal{A}^{(1)}_{\infty}(x_{i},x_{j})=$ $\displaystyle\
\mu_{1}\Sigma^{(1)}(x_{i},x_{j}),$
$\displaystyle\mathcal{A}^{(l)}_{\infty}(x_{i},x_{j})=$ $\displaystyle\
\mathcal{A}^{(l-1)}_{\infty}(x_{i},x_{j})\dot{\Sigma}^{(l)}(x_{i},x_{j})$
$\displaystyle+\mu_{l}\Sigma^{(l)}(x_{i},x_{j}),$

where

$\displaystyle\Sigma^{(1)}(x_{i},x_{j})$
$\displaystyle=\frac{1}{d_{0}}x^{\top}_{i}x_{j},$
$\displaystyle\Sigma^{(l)}(x_{i},x_{j})$
$\displaystyle=\mathbb{E}_{f\sim\mathcal{N}(0,\Sigma^{(l-1)})}[\sigma(f(x_{i}))\sigma(f(x_{j}))],$
$\displaystyle\dot{\Sigma}^{(l)}(x_{i},x_{j})$
$\displaystyle=\mathbb{E}_{f\sim\mathcal{N}(0,\Sigma^{(l-1)})}[\dot{\sigma}(f(x_{i}))\dot{\sigma}(f(x_{j}))].$

Proof Sketch. The demonstration is inspired by Jacot et al. (2018) and is by
induction. Take $\mathcal{A}^{(l)}$ the WNTK of an $l$-layer fully-connected
NN. For the first hidden layer, it can be easily demonstrated that
$\mathcal{A}^{(1)}\to\mathcal{A}^{(1)}_{\infty}\otimes Id_{d_{1}}$. For layer
$l>1$, we divide $\mathcal{A}^{(l)}$ into two parts: the sum of derivatives on
parameters in the first $(l-1)$ hidden layers and the sum of derivatives on
parameters in the $l$-th layer. In the first case, as
$\mathcal{A}^{(l-1)}\to\mathcal{A}^{(l-1)}_{\infty}\otimes Id_{d_{l-1}}$ is
assumed to be correct by induction hypothesis,
$\mathcal{A}^{(l)}_{\mathrm{first}\ (l-1)\
\mathrm{layers}}\to\mathcal{A}^{(l-1)}_{\infty}(x_{i},x_{j})\dot{\Sigma}^{(l)}(x_{i},x_{j})$
can be derived. In the second case, it can be easily demonstrated that
$\mathcal{A}^{(l)}_{l\mathrm{-th}\
\mathrm{layer}}\to\mu_{l}\Sigma^{(l)}(x_{i},x_{j})$. Thus, we have proved the
whole theorem. A detailed proof can be found in Appendix.

Note that (6) and (11) represent one of the key properties possessed by the
NTK and the WNTK, respectively. Similar to $\Theta^{(L)}_{\infty}$,
$\mathcal{A}^{(L)}_{\infty}$ only depends on the choice of activation
function, the depth of the NN, and the variance of the parameters at
initialization.

Next, we study the stability of the WNTK of a standardly parameterized NN when
applying full-batch gradient descent under Assumptions 1-4 (see below). The
demonstration is built on the recent work by Lee et al. (2019), which proves
the stability of the NTK during the training phase of an over-parameterized
NN.

Standard Parameterization: A standard parameterized NN is formulated as

$\displaystyle\begin{cases}h^{l+1}=x^{l}W^{l+1}\\\
x^{l+1}=\sigma(h^{l+1})\end{cases}$

with $x^{l}$ being the output of the $l$-th layer, $x^{l+1}$ the output of the
$(l+1)$-th layer, and each parameter $w^{l}_{i,j}$ in the connection matrices
$W^{l}$ initialized from i.i.d. Gaussians
$\mathcal{N}(0,\frac{\sigma^{2}_{w}}{d})$. In the following, we assume
$d_{L}=1$.

Assumptions 1-4:

  * •

The widths of the hidden layers are identical, i.e. $d_{1}=...=d_{L-1}=d$.

  * •

The analytic WNTK is full-rank. And we set $\eta_{\mathrm{critical}}$
$=2(\lambda_{\min}+\lambda_{\max})^{-1}>0$. $\lambda_{\min}$ and
$\lambda_{\max}$ are the smallest and the largest eigenvalues of the WNTK.
$\lambda_{\max}\geq\lambda_{\min}>0$.

  * •

The training set $(\mathcal{X},\mathcal{Y})$ is contained in some compact sets
and $x_{i}\neq x_{j}$ for all $x_{i},x_{j}\in\mathcal{X}$.

  * •

The activation function $\sigma$ satisfies: $|\sigma(0)|$,
$||\dot{\sigma}||_{\infty}$, $\sup_{x_{i}\neq
x_{j}}|\dot{\sigma}(x_{i})-\dot{\sigma}(x_{j})|/|x_{i}-x_{j}|<\infty$.

Assumption 2 is satisfied when all inputs are limited to the unit sphere. See
the proof of the positive-definiteness of the NTK in Jacot et al. (2018). A
similar conclusion can be derived for the WNTK.

###### Theorem 2

(Stability during Training). Under Assumptions 1-4 and suppose that an L-layer
neural network $f$ with random initialization is trained by gradient decent
with different learning rate $\eta a_{p}$ on different parameter $\theta_{p}$,
where $\eta=\frac{\eta_{0}}{d}$. For $\delta_{0}>0$ and
$\eta_{0}<\eta_{\mathrm{critical}}$, there exist $R_{0}>0$, $N\in\mathbb{N}$,
and $K>1$ such that for every $d\geq N$, we have the following inequality with
probability at least $(1-\delta_{0})$:

$\displaystyle||\mathcal{A}_{0}-\mathcal{A}_{t}||_{F}\leq\frac{6K^{3}R_{0}a_{\max}^{2}}{\lambda_{\min}}d^{-\frac{1}{2}},\forall
t,$ (12)

where $\mathcal{A}_{0}$ and $\mathcal{A}_{t}$ are the WNTK at moment $0$ and
$t$ under standard parameterization, and $a_{\max}=\max\\{a_{i}|i=1,...,P\\}$.

Proof Sketch. We begin with the definition of $\mathcal{A}$ (10). By applying
local lipschitzness properties of the Jacobean (Lemma 1 in Appendix), we find
that a sufficient condition of
$||\mathcal{A}_{0}-\mathcal{A}_{t}||_{F}\leq\frac{6K^{3}R_{0}a_{\max}^{2}}{\lambda_{\min}}d^{-\frac{1}{2}}$
is
$||\theta_{0}-\theta_{t}||_{2}\leq\frac{3KR_{0}a_{\max}}{\lambda_{\min}}d^{-\frac{1}{2}}$.
Secondly, we deduce by induction that the latter inequality can be attained if
$||f(\theta_{t},\mathcal{X})-\mathcal{Y}||_{2}\leq(1-\frac{\eta_{0}\lambda_{\min}}{3})^{\top}R_{0}$.
Thirdly, we prove the correctness of the last inequality under Assumptions
1-4. A detailed proof can be found in Appendix.

For the WNTK, (12) shows the stability of the WNTK during training, playing
similarly the rule of (7) in the NTK analysis. Adjustment on descent direction
does not affect the stability of the WNTK. So, we find no dependence between
the stability of the WNTK and the optimizer of an over-parameterized NN. The
weight terms in the WNTK can characterize different optimizers.

###  4.2 Equivalence between WNTK Regression Estimator and NN Estimator using
Different Optimizer

In this subsection, we will show that a standard parameterized wide NN
estimator under adjusted gradient descent is equivalent to its related WNTK
regression estimator. Here we consider sum of squares as the loss function,
i.e., $(x_{i},y_{i})\in(\mathcal{X},\mathcal{Y})$,

$\displaystyle\mathcal{L}(f(\theta,\mathcal{X}))=\frac{1}{2}\sum_{i=1}^{n}(f(\theta,x_{i})-y_{i})^{2}.$

First, under adjusted gradient descent, an over-parameterized NN still drops
into the lazy regime, and thus can be approximated by a linearized network
defined as follows,

$\displaystyle
f^{\mathrm{lin}}_{t}(x)=f(\theta_{0},x)+\nabla_{\theta}f(x)|_{\theta=\theta_{0}}(\theta_{t}-\theta_{0}).$
(13)

$f^{\mathrm{lin}}_{t}(x)$ and $\theta_{t}$ denote the value of
$f^{\mathrm{lin}}(x)$ and $\theta$ at moment $t$, respectively. The dynamics
of adjusted gradient flow using this linearized model are governed by

$\dot{\theta}_{t}=-\eta(a\odot\nabla_{\theta}f_{0}(\mathcal{X})^{\top})\nabla_{f^{\mathrm{lin}}_{t}(\mathcal{X})}\mathcal{L},$
(14)
$\dot{f}^{\mathrm{lin}}_{t}(x)=-\eta\mathcal{A}_{0}(x,\mathcal{X})\nabla_{f^{\mathrm{lin}}_{t}(\mathcal{X})}\mathcal{L}.$
(15)

###### Theorem 3

(WNTK Lazy Regime). Under Assumptions 1-4, applying gradient descent with
different learning rate $\eta a_{p}$ on different parameter $\theta_{p}$, for
every $x\in{\mathbb{R}}^{d_{0}}$ satisfying $||x||_{2}\leq 1$, for all $t\geq
0$, as $d\to\infty$, with probability arbitrarily close to 1 over random
initialization, we have

$\displaystyle||f(\theta_{t},\mathcal{X})-f^{\mathrm{lin}}_{t}(\mathcal{X})||_{2}\to
0,$ (16) $\displaystyle||f(\theta_{t},x)-f^{\mathrm{lin}}_{t}(x)||_{2}\to 0.$

The proof of this theorem follows a similar proof in Lee et al. (2019). The
details could be found in Appendix.

###### Theorem 4

(Equivalence between WNTK Regression and Neural Network). Fix a set of
learning rate $\eta a>0$. Consider gradient descent of an infinitely wide
neural network with initialization $\theta(0)$. If $\eta$ is bounded by a
known value $\eta_{\mathrm{c}}$ and the neural network is defined to have a
very small initial output, the fully trained NN estimator $f^{\star}$ is
equivalent to the WNTK ridge regression estimator with regularization term
going to zero:

$\displaystyle
f^{\star}(\theta,x)=\mathcal{A}(x,\mathcal{X})^{\top}(\mathcal{A}(\mathcal{X},\mathcal{X}))^{-1}\mathcal{Y},$
(17)

where $\mathcal{A}(\mathcal{X},\mathcal{X})\in\mathbb{R}^{n\times n}$ stand
for the WNTK on the training set and
$\mathcal{A}(x,\mathcal{X})=[\mathcal{A}(x,x_{1}),...,\mathcal{A}(x,x_{n})]^{\top}\in\mathbb{R}^{n}$.

Proof. Let us first review this NN’s adjusted gradient descent dynamics. At
iteration $k$, adjusted gradient descent on $\theta$ with learning rate $a$
iterates as follows:

$\displaystyle\theta(k+1)=\theta(k)-\eta(a\odot\nabla_{\theta}\mathcal{L}(\theta(k))),a\in{\mathbb{R}}^{P}.$
(18)

Under the infinite-width assumption, according to the property of lazy regime
(16), as $d\to\infty$:

$\displaystyle\mathcal{L}(f(\theta,\mathcal{X}))=$ $\displaystyle\
\frac{1}{2}\sum_{i=1}^{n}(f(\theta(0),x_{i})+$ (19) $\displaystyle\
\left<\nabla_{\theta}f(\theta(0),x_{i}),\theta-\theta(0)\right>-y_{i})^{2}.$

We substitute (19) into (18) and we get

$\displaystyle\ \theta(k+1)$ $\displaystyle=$ $\displaystyle\
\theta(k)-\eta(a\odot(\sum_{i=1}^{n}(f(\theta(k),x_{i})-y_{i})\nabla_{\theta}f(\theta(0),x_{i})))$
$\displaystyle=$ $\displaystyle\
\theta(k)-\eta(a\odot(\sum_{i=1}^{n}(f(\theta(0),x_{i})+$ $\displaystyle\
\left<\nabla_{\theta}f(\theta(0),x_{i}),\theta(k)-\theta(0)\right>-y_{i})(\nabla_{\theta}f(\theta(0),x_{i}))))$
$\displaystyle=$ $\displaystyle\ \theta(k)-\eta(a\odot U)L-\eta(a\odot
U)U^{T}(\theta(k)-\theta(0)),$

where

$\displaystyle L=\ (l_{0}(0),l_{1}(0),...,l_{n}(0))^{T}\in{\mathbb{R}}^{n},$
$\displaystyle l_{i}(0)=\ f(\theta(0),x_{i})-y_{i}\in{\mathbb{R}},$
$\displaystyle U=\
(\nabla_{\theta}f_{t}(\mathcal{X})|_{\theta=\theta_{0}})^{\top}\in{\mathbb{R}}^{P\times
n}.$

Let $c(k)$ $=\theta(k)$ $+(\eta(a\odot U)U^{T})^{-1}$ $(\eta(a\odot
U)L-\eta(a\odot U)U^{T}\theta(0))$. We find the following iterative relation

$\displaystyle c(k+1)=(Id-\eta(a\odot U)U^{T})c(k).$

Set $\eta_{\mathrm{c}}=\frac{1}{|(a\odot U)U^{T}|}$. According to (10), we can
deduce, as $k\to\infty$,

$\displaystyle\theta(k)=$ $\displaystyle\ \theta(0)-(Id-(Id-\eta(a\odot
U)U^{T})^{k})$ $\displaystyle\ (\eta(a\odot U)U^{T})^{-1}\eta(a\odot U)L$
$\displaystyle=$ $\displaystyle\ \theta(0)-(\eta(a\odot
U)U^{T})^{-1}\eta(a\odot U)L$ $\displaystyle=$ $\displaystyle\
\theta(0)-(a\odot U)(\mathcal{A})^{-1}L.$

As we assume that the neural network is defined to have a very small initial
output, the predictor of this wide neural network is:

$\displaystyle f^{\star}(\theta,x)=$ $\displaystyle\
f^{lin}_{\infty}(x)=\left<\nabla_{\theta}f(\theta(0),x),\theta-\theta(0)\right>$
$\displaystyle=$ $\displaystyle\
-(\nabla_{\theta}f_{t}(x)|_{\theta=\theta_{0}})^{\top}(a\odot
U)(\mathcal{A})^{-1}L$ $\displaystyle=$ $\displaystyle\
\mathcal{A}(x,\mathcal{X})^{\top}(\mathcal{A})^{-1}\mathcal{Y}.$

$\hfill\blacksquare$

Note that for an empirical WNTK of a finite-width NN, we can no longer ignore
the initial output and the prediction takes the following form:

$\displaystyle f^{\star}(\theta,x)\simeq$ $\displaystyle\
\mathcal{A}(x,\mathcal{X})^{\top}(\mathcal{A})^{-1}\mathcal{Y}+(f(\theta_{0},x)$
(20) $\displaystyle\
-\mathcal{A}(x,\mathcal{X})^{\top}(\mathcal{A})^{-1}f(\theta_{0},\mathcal{X})).$

Moreover, regularization in the WNTK ridge regression is equivalent to
regularization in over-parameterized NNs. Such properties have already been
discussed on the NTK (Wei et al., 2019).

Below is an example, where WNTK simulates the training process of a NN under
gradient descent with momentum. Take $\beta\in[0,1]$. At iteration $k$, the
gradient descent with momentum is given by the following,

$\displaystyle\Delta\theta=$
$\displaystyle-\eta\nabla_{\theta}f(\theta(k),\mathcal{X})$
$\displaystyle-\eta\beta(\nabla_{\theta}f(\theta(k),\mathcal{X})-\nabla_{\theta}f(\theta(k-1),\mathcal{X})).$

Suppose that this NN has only one output at the final layer. Choose the
following weights in WNTK,

$\displaystyle
a_{p}=1+\beta(\nabla_{\theta_{p}}f(\theta(0),\mathcal{X}))^{-1}H_{p}(\theta(0)),$

where $H_{p}(\theta(0),\mathcal{X})$ refers to the Hessian Matrix at
$\theta(0)$ with respect to $\theta_{p}$, namely, the derivative gradient
information at the initial point $\theta(0)$. Analytically, it can be
approximated by the second derivative of $f$. Practically, it can be obtained
by training the neural network for one epoch using full-batch gradient
descent. In this case, the WNTK represents the same dynamics as gradient
descent with momentum:

$\displaystyle\Delta\theta_{p}$
$\displaystyle=-\eta(a\odot\nabla_{\theta_{p}}f(\theta(0),\mathcal{X}))$
$\displaystyle=-\eta(\nabla_{\theta_{p}}f(\theta(0),\mathcal{X})+\beta
H_{p}(\theta(0),\mathcal{X})).$

Gradient descent with momentum and its derives has been revealed to be useful
and effective on over-parameterized NN (Liu & Belkin, 2018). The above example
shows that such improvements on optimizer can be equivalently transferred to
WNTK regression estimators.

##  5 Algorithm: Weight Update

In the previous sections, we show the equivalence between the WNTK regression
estimator and the corresponding over-parameterized NN estimator under the
adjusted descent direction. Because of the outperformance of other optimizers
over gradient descent, the WNTK is promising to have better performance than
the NTK if we can find a good descent direction, i.e., good weights. One
problem here is that WNTK, also NTK, are certainly overfitting on the training
set. To achieve better generalization performance, we have tried cross-
validation, kernel alignment, regularization rules, Etc, and find that the
following update rule is the most suitable.

Gradient Descent on Layer-wise Weight. One way to find the optimal layer-wise
weights is to perform gradient descent on weights. The original training data
can be divided into a training set
$(\mathcal{X}_{\mathcal{T}},\mathcal{Y}_{\mathcal{T}})$ and a validation set
$(\mathcal{X}_{\mathcal{V}},\mathcal{Y}_{\mathcal{V}})$. We use
$\mathcal{X}_{\mathcal{T}}$ to generate the WNTK estimator and use its
prediction on $\mathcal{X}_{\mathcal{V}}$ to update its layer-wise weights.
The weights could be updated layer-wisely or element-wisely. Consider an
$L$-layer NN and $d_{L}=1$. With the definition of layer-wise WNTK (3) and
(4), the layer-wise gradient of a weight term $a^{L}_{i}$ at the $l$-th layer
is given as follows,

$\displaystyle\ \frac{\partial f^{\star}}{\partial
a^{L}_{i}}(\theta,\mathcal{X}_{\mathcal{V}})$ $\displaystyle=\ $
$\displaystyle\frac{\partial\mathcal{A}_{\mathcal{V}}^{\top}}{\partial
a^{L}_{i}}\mathcal{A}^{-1}\mathcal{Y}_{\mathcal{T}}+\mathcal{A}_{\mathcal{V}}^{\top}\frac{\partial\mathcal{A}^{-1}}{\partial
a^{L}_{i}}\mathcal{Y}_{\mathcal{T}}$ $\displaystyle=\ $
$\displaystyle({\Theta_{\mathcal{V}}}^{(L)}_{l})^{\top}\mathcal{A}^{-1}\mathcal{Y}_{\mathcal{T}}-\mathcal{A}_{\mathcal{V}}^{\top}\mathcal{A}^{-1}\Theta^{(L)}_{l}\mathcal{A}^{-1}\mathcal{Y}_{\mathcal{T}},$

where

$\displaystyle{\Theta_{\mathcal{V}}}^{(L)}_{l}$
$\displaystyle=\sum_{\theta_{p}\ \mathrm{in\ layer}\
l}\partial_{\theta_{p}}f^{(L)}(\theta,\mathcal{X}_{\mathcal{T}})^{\top}\partial_{\theta_{p}}f^{(L)}(\theta,\mathcal{X}_{\mathcal{V}})$

and $\mathcal{A}_{\mathcal{V}}$ represents the WNTK between
$\mathcal{X}_{\mathcal{T}}$ and $\mathcal{X}_{\mathcal{V}}$. The whole
algorithm could be found in Algorithm 1. Notice that the above closed-form
formula enables us to avoid calculating the gradient through the Auto-Gradient
method in packages like Tensorflow or Pytorch, which is of high cost. Instead,
by storing the layer-wise kernels $\Theta^{(L)}_{l}$ and $\mathcal{A}^{-1}$ in
memory, the calculation of gradients of layer-wise weights becomes much
lighter. The time cost of this algorithm is mainly from the matrix inversion
operation $\mathcal{A}^{-1}$, which is of $\mathcal{O}(n^{3})$. So the time
complexity of each gradient update iteration is $\mathcal{O}(Ln^{3})$. Its
space complexity is $\mathcal{O}(Ln^{2})$, which mainly comes from the storage
of $\mathcal{A}^{-1}$ and $\Theta^{(L)}_{l}$, $l=1,..,L$.

With little modification, the above update is applicable to train element-wise
weights. In our experiment, we choose layer-wise update strategy, since it can
better control complexity. Moreover, with different layer-wise weights, the
WNTK could explain how layer-wise information is integrated and why one needs
different learning rates for training NNs, which has been reported in many
practical papers (Takase et al., 2018; Liu et al., 2019).

To investigate the meaning of layer-wise weights, we consider a fully
connected NN at infinite width and focus on two successive layers. Reformulate
the layer-wise kernel (3) as the following,

$\Theta^{(L)}_{l}(x_{i},x_{j})=\Sigma^{(l)}(x_{i},x_{j})\prod_{l^{\prime}=l+1}^{L}\dot{\Sigma}^{(l^{\prime})}(x_{i},x_{j}).$
(21)

The activation function is supposed to be $ReLU$ and we have

$\displaystyle\
\mathbb{E}_{(u,v)\sim\mathcal{N}(0,D\Lambda^{(l)}D)}[\sigma(u)\sigma(v)]$
$\displaystyle=$ $\displaystyle\
\frac{\lambda(\pi-\arccos(\lambda))+\sqrt{1-\lambda^{2}}}{2\pi}c_{1}c_{2},$
$\displaystyle\
\mathbb{E}_{(u,v)\sim\mathcal{N}(0,D\Lambda^{(l)}D)}[\dot{\sigma}(u)\dot{\sigma}(v)]$
$\displaystyle=$ $\displaystyle\ \frac{\pi-\arccos(\lambda)}{2\pi},$

where $|\lambda|<1$,

$\displaystyle D=\begin{bmatrix}c_{1}&0\\\ 0&c_{2}\end{bmatrix}\mathrm{,\ and\
}\Lambda^{(l)}=\begin{bmatrix}1&\lambda\\\ \lambda&1\end{bmatrix}$

under the following constraint,

$\displaystyle
D\Lambda^{(l)}D=\begin{bmatrix}\Sigma^{(l-1)}(x_{i},x_{i})&\Sigma^{(l-1)}(x_{i},x_{j})\\\
\Sigma^{(l-1)}(x_{j},x_{i})&\Sigma^{(l-1)}(x_{j},x_{j})\end{bmatrix}.$

Note that the analytic kernel $\mathcal{A}^{(L)}_{\infty}$ in (11) and the
WNTK are well determined, once we fix the activation function, the depth of
the NN, and the variance of the parameters at initialization. Take a look at
the difference between $\Theta^{(L)}_{l}(x_{i},x_{j})$ and
$\Theta^{(L)}_{l-1}(x_{i},x_{j})$, which is equal to the difference between
two WNTKs of the same NN with different layer-wise weights
$\begin{cases}a_{i}=0,\ i\neq l\\\ a_{l}=1\end{cases}$ and
$\begin{cases}a_{i}=0,\ i\neq l+1\\\ a_{l+1}=1\end{cases}$.
By applying (21), this difference takes the following form:

$\displaystyle\ \Theta^{(L)}_{l-1}(x_{i},x_{j})-\Theta^{(L)}_{l}(x_{i},x_{j})$
$\displaystyle=$ $\displaystyle\
(\Sigma^{(l-1)}(x_{i},x_{j})\dot{\Sigma}^{(l)}(x_{i},x_{j})-\Sigma^{(l)}(x_{i},x_{j}))C$
$\displaystyle=$ $\displaystyle\ \frac{\sqrt{1-\lambda^{2}}c_{1}c_{2}}{2\pi}C$
$\displaystyle=$ $\displaystyle\
\sqrt{\Sigma^{(l-1)}(x,x)\Sigma^{(l-1)}(x^{\prime},x^{\prime})-\Sigma^{(l-1)}(x_{i},x_{j})^{2}}C,$

where $C=\prod_{l^{\prime}=l+1}^{L}\dot{\Sigma}^{(l^{\prime})}(x_{i},x_{j})$
is a constant and is not related to the $(l-1)$-th either the $l$-th layer’s
output. It is apparent that, compared to $\Theta^{(L)}_{l}(x_{i},x_{j})$,
$\Theta^{(L)}_{l-1}(x_{i},x_{j})$ possesses additional information related to
the similarity of inputs at the $(l-1)$-th layer. Hence, each layer represents
similar but not identical information. So, the layer-wise weight term, by
nature, should be non-equally weighted. It is reasonable to expect that a
network’s optimal kernel has different weights on different layer-wise
kernels.

Algorithm 1 Gradient Descent on Layer-wise Weight

Input: training data $\\{(\mathcal{X},\mathcal{Y})\\}$, initial layer-wise
weight $a^{L}=[a^{L}_{1},a^{L}_{2},...,a^{L}_{L}]^{\top}$, train-validation
ratio $r$, learning rate $\eta_{w}$.

repeat

Split $\\{(\mathcal{X},\mathcal{Y})\\}$ randomly into to a training set
$(\mathcal{X}_{\mathcal{T}},\mathcal{Y}_{\mathcal{T}})$ and a validation set
$(\mathcal{X}_{\mathcal{V}},\mathcal{Y}_{\mathcal{V}})$ with respect to the
train-validation ratio $r$. Generate the WNTK $\mathcal{A}$ and the estimator
$f^{\star}$ from $\mathcal{X}_{\mathcal{T}}$ and $a^{L}$.

for $i=1$ to $L$ do

$\Delta_{i}=2(f^{\star}(\theta,\mathcal{X}_{\mathcal{V}})-\mathcal{Y}_{\mathcal{V}})^{\top}\frac{\partial
f^{\star}}{\partial a^{L}_{i}}(\theta,\mathcal{X}_{\mathcal{V}})$

Update $a^{L}_{i}=a^{L}_{i}-\eta_{w}\times\Delta_{i}$.

end for

until Stopping criteria is reached.

##  6 Experiment

This section provides empirical support that the proposed WNTK, also the
Weighted Convolutional Neural Tangent Kernel (WCNTK), does improve the
classification performance. We will first evaluate the performance of
empirical kernels on CIFAR-10, which is an image task. Then we compare the
classification performances between the analytic WNTK and NTK on 10 UCI
datasets, which are all non-image. The empirical kernel matrices and NNs are
calculated and trained using the Tensorflow package, with a NVIDIA RTX 2080Ti
GPU. All the codes could be found from
https://github.com/ASTAugustin/ICML_WNTK.

Table 1: Binary-classification accuracy on CIFAR-10 dataset.
dataset | Size | NTK_pur | WNTK_pur | NTK_ini | WNTK_ini | NN
---|---|---|---|---|---|---
Air/Auto | 100 |  71.33$\pm$1.68 |  76.45$\pm$1.81 |  77.33$\pm$2.22 | 79.55$\pm$3.30 |  78.30$\pm$2.03
| 1000 |  84.22$\pm$0.91 |  85.34$\pm$0.87 |  85.75$\pm$0.87 |  86.98$\pm$1.01 | 87.82$\pm$0.55
| 10000 |  88.76$\pm$0.77 |  89.66$\pm$0.65 |  92.70$\pm$0.80 |  93.61$\pm 0.63$ | 94.29$\pm$0.54
Cat/Bird | 100 |  59.96$\pm$2.98 |  64.37$\pm$3.29 |  62.11$\pm$3.57 | 65.41$\pm$2.37 |  56.87$\pm$2.22
| 1000 |  72.22$\pm$1.62 | 73.67$\pm$1.06 |  69.78$\pm$1.74 |  71.14$\pm$1.36 |  67.38$\pm$6.09
| 10000 |  76.15$\pm$0.35 |  76.67$\pm$0.58 |  78.00$\pm$1.81 | 79.15$\pm$1.62 |  77.46$\pm$0.90
Cat/Deer | 100 |  61.58$\pm$2.07 | 66.47$\pm$2.04 |  63.06$\pm$7.41 |  64.06$\pm$7.48 |  60.03$\pm$7.95
| 1000 |  70.94$\pm$1.49 | 71.36$\pm$1.61 |  70.50$\pm$1.05 |  71.22$\pm$1.96 |  70.46$\pm$1.29
| 10000 |  77.30$\pm$0.85 |  77.73$\pm$0.69 |  79.96$\pm$0.94 | 81.42$\pm$0.83 |  78.76$\pm$0.94
Deer/Dog | 100 |  61.18$\pm$1.79 | 65.12$\pm$1.39 |  59.63$\pm$4.08 |  62.92$\pm$2.77 |  61.31$\pm$4.29
| 1000 |  68.91$\pm$1.24 |  70.04$\pm$1.09 |  70.52$\pm$2.15 | 71.78$\pm$3.55 |  71.56$\pm$1.38
| 10000 |  76.33$\pm$0.71 |  77.66$\pm$0.64 |  82.61$\pm$1.00 |  84.46$\pm$0.96 | 84.85$\pm$0.69

###  6.1 Experiment of Empirical Kernels on CIFAR-10

We empirically compare the CNTK and the WCNTK on binary-classification on
CIFAR-10 dataset, a widely used image dataset and contains 60,000 $32\times
32$ colored images in 10 different classes. For each class, there are 5,000
training samples and 1,000 test samples. We test the performance of Vanilla
CNN and the corresponding kernels (CNTK and WCNTK) on different subsets of
CIFAR-10. Samples of different sizes ($10^{2}$, $10^{3}$, and $10^{4}$) are
randomly selected for training, and we use the test set for evaluation.

In this experiment, a five-layer Vanilla CNN, which consists of four
convolutional layers and a fully connected layer, is used. The number of
channels is 64, and the activation function is $tanh$. A fully connected layer
with one neuron is used as the output. The neural network output is multiplied
by a small constant in order that the second term of final prediction of the
empirical CNTK and CWNTK (20) becomes negligible without changing the descent
direction. For WNTK, layer-wise weights are learned by Algorithm 1 with the
following setting: initialized layer-wise weights $a^{L}_{1}=0.25$,
$a^{L}_{i}=1\;(\forall i>1)$, learning rate $\eta_{w}=0.01$, and train-
validation ratio $r=0.2$. The regularization term in kernel regression is 1.

Remark. The selection of $a^{L}_{1}=0.25$ and $a^{L}_{i}=1\;(\forall i>1)$ is
motivated by the fact that when applying Algorithm 1, the first layer’s weight
usually drops significantly to nearly 0. To speed up the convergence, we
heuristically give a small initial value.

Each experiment is repeated ten times, and we report the mean, together with
the standard deviation, in Table 1, where (W)NTK_pur stands for (W)CNTK of CNN
with randomly initialized parameters and (W)NTK_ini for (W)CNTK of CNN with
pretrained parameters. From Table 1, we could observe that:

  * •

The WNTK indeed significantly improves the performance of the NTK, by
averagely $2.14\%$ for random initialization and $1.65\%$ for pretrained
parameters.

  * •

Pretraining is helpful for both empirical NTKs and WNTKs, with which, both
empirical NTKs and WNTKs can get a presentable improvement about $1\%$ to
$2\%$.

  * •

Empirical WNTKs could be competitive to or even better than the used NNs,
especially if the pretraining is used and when the data size is not too large.
The used NNs are not wide enough and the weights of WNTKs can also be assigned
element-wisely. These possibilities of amelioration make it promising to
further improve the performance of the WNTK.

###  6.2 Experiment of Analytic Kernels on UCI

Next we consider tasks on UCI Machine Learning Repository. The analytic NTK
and WNTK are evaluated on these non-image datasets, compared with a wide NN,
simulating the corresponding infinite-width NN. The selected task name,
together with the size (S.), the number of features (F.), and the number of
classes (F.) are shown in Table 2. Different to the previous experiments, UCI
data are non-image and we test the analytic WNTK of infinite-width NN, instead
of an empirical one.

Following the same setting of (11), the WNTK in this experiment is
analytically generated from a five-layer fully connected NN with $ReLU$
activation under the infinite-width assumption. In pre-processing, all
training data are standardized. The initial NN parameters are randomly given
with unit parameter variance. The regularization term in kernel regression is
set to be $0.1$. For comparison, we also train wide five-layer fully connected
NNs, where the number of the channels is set to 1024 for both the input layer
and every hidden layer, in order to simulate NNs of infinite width.

In Table 2, we report the accuracy calculated by 5-fold cross-validation. The
WNTK improves accuracy from the NTK by $2\%$ to $3\%$. Such improvement
reflects the fact that, under the infinite width assumption, the analytical
WNTK is also better in generalization ability than the analytical NTK. Note
that both WNTKs and NTKs are much better than fully connected NNs on small
datasets. Such advantage of NN-induced kernels when facing small datasets has
also been demonstrated by Arora et al. (2019c).

Table 2: Classification accuracy on 10 UCI datasets. dataset | s./f./c. | NTK | WNTK | NN
---|---|---|---|---
tissue | (106,9,6) | 69.23 | 72.12 | 28.52
echo | (131,10,2) | 81.82 | 84.09 | 75.78
planning | (182,12,2) | 65.55 | 71.11 | 60.20
cancer | (198,33,2) | 81.63 | 84.69 | 76.81
heart | (294,12,2) | 83.56 | 85.62 | 78.41
survival | (306,3,2) | 70.72 | 73.68 | 66.82
tumor | (330,17,15) | 49.09 | 51.83 | 15.85
ilpd | (583,9,2) | 71.23 | 73.29 | 66.65
credit | (690,14,2) | 60.61 | 64.83 | 80.80
contrac | (1473,9,3) | 48.37 | 52.51 | 39.99

##  7 Conclusion

Motivated by the fact that the Neural Tangent Kernel is restricted to describe
NN dynamics with gradient descent, this paper proposes the Weighted Neural
Tangent Kernel as a generalized and improved version of the NTK that can
simulate wide NN dynamics with different optimizers. With rigorous discussion,
we prove that the WNTK is also stable at initialization and during training. A
sufficiently wide net fully trained by adjusted gradient descent is proved to
be a linearized model and equivalent to the WNTK regression predictor. Hence,
the WNTK possesses all properties of the NTK, and we believe that the WNTK is
a more powerful tool to understand the correspondence between wide NNs and
kernels. With developing a weight learning algorithm, we also verify the
WNTK’s superiority over the NTK.

## References

  * Alemohammad et al. (2020) Alemohammad, S., Wang, Z., Balestriero, R., and Baraniuk, R.  The recurrent neural tangent kernel.  _arXiv preprint arXiv:2006.10246_ , 2020.
  * Allen-Zhu et al. (2018) Allen-Zhu, Z., Li, Y., and Liang, Y.  Learning and generalization in overparameterized neural networks, going beyond two layers.  _arXiv preprint arXiv:1811.04918_ , 2018.
  * Anselmi et al. (2015) Anselmi, F., Rosasco, L., Tan, C., and Poggio, T.  Deep convolutional networks are hierarchical kernel machines.  _arXiv preprint arXiv:1508.01084_ , 2015.
  * Arora et al. (2018) Arora, S., Cohen, N., and Hazan, E.  On the optimization of deep networks: Implicit acceleration by overparameterization.  In _International Conference on Machine Learning_ , pp. 244–253. PMLR, 2018.
  * Arora et al. (2019a) Arora, S., Du, S., Hu, W., Li, Z., Salakhutdinov, R., and Wang, R.  On exact computation with an infinitely wide neural net.  _ArXiv_ , abs/1904.11955, 2019a.
  * Arora et al. (2019b) Arora, S., Du, S., Hu, W., Li, Z., and Wang, R.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In _International Conference on Machine Learning_ , pp. 322–332. PMLR, 2019b.
  * Arora et al. (2019c) Arora, S., Du, S. S., Li, Z., Salakhutdinov, R., Wang, R., and Yu, D.  Harnessing the power of infinitely wide deep nets on small-data tasks.  _arXiv preprint arXiv:1910.01663_ , 2019c.
  * Belkin et al. (2018) Belkin, M., Ma, S., and Mandal, S.  To understand deep learning we need to understand kernel learning.  In _International Conference on Machine Learning_ , pp. 541–549. PMLR, 2018.
  * Brock et al. (2017) Brock, A., Lim, T., Ritchie, J. M., and Weston, N.  Freezeout: Accelerate training by progressively freezing layers.  _arXiv preprint arXiv:1706.04983_ , 2017.
  * Chen et al. (2020) Chen, S., He, H., and Su, W. J.  Label-aware neural tangent kernel: Toward better generalization and local elasticity.  _arXiv preprint arXiv:2010.11775_ , 2020.
  * Chizat et al. (2018) Chizat, L., Oyallon, E., and Bach, F.  On lazy training in differentiable programming.  _arXiv preprint arXiv:1812.07956_ , 2018.
  * Cutkosky & Mehta (2020) Cutkosky, A. and Mehta, H.  Momentum improves normalized sgd.  In _International Conference on Machine Learning_ , pp. 2260–2268. PMLR, 2020.
  * Daniely et al. (2016) Daniely, A., Frostig, R., and Singer, Y.  Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity.  _arXiv preprint arXiv:1602.05897_ , 2016.
  * Dragomir (2003) Dragomir, S. S.  _Some Gronwall type inequalities and applications_.  Nova Science, 2003.
  * Du et al. (2018) Du, S. S., Zhai, X., Poczos, B., and Singh, A.  Gradient descent provably optimizes over-parameterized neural networks.  _arXiv preprint arXiv:1810.02054_ , 2018.
  * Du et al. (2019) Du, S. S., Hou, K., Póczos, B., Salakhutdinov, R., Wang, R., and Xu, K.  Graph neural tangent kernel: Fusing graph neural networks with graph kernels.  _arXiv preprint arXiv:1905.13192_ , 2019.
  * Garriga-Alonso et al. (2018) Garriga-Alonso, A., Rasmussen, C. E., and Aitchison, L.  Deep convolutional networks as shallow gaussian processes.  _arXiv preprint arXiv:1808.05587_ , 2018.
  * Ghorbani et al. (2019a) Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A.  Limitations of lazy training of two-layers neural networks.  _arXiv preprint arXiv:1906.08899_ , 2019a.
  * Ghorbani et al. (2019b) Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A.  Linearized two-layers neural networks in high dimension.  _arXiv preprint arXiv:1904.12191_ , 2019b.
  * Gönen & Alpaydın (2011) Gönen, M. and Alpaydın, E.  Multiple kernel learning algorithms.  _The Journal of Machine Learning Research_ , 12:2211–2268, 2011.
  * Hanin & Nica (2019) Hanin, B. and Nica, M.  Finite depth and width corrections to the neural tangent kernel.  _arXiv preprint arXiv:1909.05989_ , 2019.
  * Hazan & Jaakkola (2015) Hazan, T. and Jaakkola, T.  Steps toward deep kernel methods from infinite neural networks.  _arXiv preprint arXiv:1508.05133_ , 2015.
  * Hu et al. (2019) Hu, W., Li, Z., and Yu, D.  Understanding generalization of deep neural networks trained with noisy labels.  _ArXiv_ , abs/1905.11368, 2019.
  * Huang & Yau (2020) Huang, J. and Yau, H.-T.  Dynamics of deep neural networks and neural tangent hierarchy.  In _International Conference on Machine Learning_ , pp. 4542–4551. PMLR, 2020.
  * Jacot et al. (2018) Jacot, A., Gabriel, F., and Hongler, C.  Neural tangent kernel: Convergence and generalization in neural networks.  In _Advances in neural information processing systems_ , pp. 8571–8580, 2018.
  * Lee et al. (2017) Lee, J., Bahri, Y., Novak, R., Schoenholz, S. S., Pennington, J., and Sohl-Dickstein, J.  Deep neural networks as gaussian processes.  _arXiv preprint arXiv:1711.00165_ , 2017.
  * Lee et al. (2019) Lee, J., Xiao, L., Schoenholz, S., Bahri, Y., Novak, R., Sohl-Dickstein, J., and Pennington, J.  Wide neural networks of any depth evolve as linear models under gradient descent.  In _Advances in neural information processing systems_ , pp. 8572–8583, 2019.
  * Li & Liang (2018) Li, Y. and Liang, Y.  Learning overparameterized neural networks via stochastic gradient descent on structured data.  _arXiv preprint arXiv:1808.01204_ , 2018.
  * Li et al. (2019) Li, Z., Wang, R., Yu, D., Du, S. S., Hu, W., Salakhutdinov, R., and Arora, S.  Enhanced convolutional neural tangent kernels.  _arXiv preprint arXiv:1911.00809_ , 2019.
  * Liu & Belkin (2018) Liu, C. and Belkin, M.  Accelerating sgd with momentum for over-parameterized learning.  _arXiv preprint arXiv:1810.13395_ , 2018.
  * Liu et al. (2020) Liu, F., Huang, X., Gong, C., Yang, J., and Li, L.  Learning data-adaptive non-parametric kernels.  _Journal of Machine Learning Research_ , 21(208):1–39, 2020.
  * Liu et al. (2019) Liu, L., Jiang, H., He, P., Chen, W., Liu, X., Gao, J., and Han, J.  On the variance of the adaptive learning rate and beyond.  _arXiv preprint arXiv:1908.03265_ , 2019.
  * Matthews et al. (2018) Matthews, A. G. d. G., Rowland, M., Hron, J., Turner, R. E., and Ghahramani, Z.  Gaussian process behaviour in wide deep neural networks.  _arXiv preprint arXiv:1804.11271_ , 2018.
  * Neal (2012) Neal, R. M.  _Bayesian learning for neural networks_ , volume 118.  Springer Science & Business Media, 2012.
  * Novak et al. (2018) Novak, R., Xiao, L., Lee, J., Bahri, Y., Yang, G., Hron, J., Abolafia, D. A., Pennington, J., and Sohl-Dickstein, J.  Bayesian deep convolutional networks with many channels are gaussian processes.  _arXiv preprint arXiv:1810.05148_ , 2018.
  * Samarin et al. (2020) Samarin, M., Roth, V., and Belius, D.  On the empirical neural tangent kernel of standard finite-width convolutional neural network architectures.  _arXiv preprint arXiv:2006.13645_ , 2020.
  * Shankar et al. (2020) Shankar, V., Fang, A., Guo, W., Fridovich-Keil, S., Ragan-Kelley, J., Schmidt, L., and Recht, B.  Neural kernels without tangents.  In _International Conference on Machine Learning_ , pp. 8614–8623. PMLR, 2020.
  * Sonnenburg et al. (2006) Sonnenburg, S., Rätsch, G., Schäfer, C., and Schölkopf, B.  Large scale multiple kernel learning.  _The Journal of Machine Learning Research_ , 7:1531–1565, 2006.
  * Srivastava et al. (2014) Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R.  Dropout: a simple way to prevent neural networks from overfitting.  _The journal of machine learning research_ , 15(1):1929–1958, 2014.
  * Sutskever et al. (2013) Sutskever, I., Martens, J., Dahl, G., and Hinton, G.  On the importance of initialization and momentum in deep learning.  In _International conference on machine learning_ , pp. 1139–1147. PMLR, 2013.
  * Takase et al. (2018) Takase, T., Oyama, S., and Kurihara, M.  Effective neural network training with adaptive learning rate based on training loss.  _Neural Networks_ , 101:68–78, 2018.
  * Wei et al. (2019) Wei, C., Lee, J. D., Liu, Q., and Ma, T.  Regularization matters: Generalization and optimization of neural nets vs their induced kernel.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Williams (1997) Williams, C. K.  Computing with infinite networks.  _Advances in neural information processing systems_ , pp. 295–301, 1997.
  * Zeiler (2012) Zeiler, M. D.  Adadelta: an adaptive learning rate method.  _arXiv preprint arXiv:1212.5701_ , 2012.
  * Zhang (2018) Zhang, Z.  Improved adam optimizer for deep neural networks.  In _2018 IEEE/ACM 26th International Symposium on Quality of Service (IWQoS)_ , pp. 1–2. IEEE, 2018.
  * Zhuang et al. (2011) Zhuang, J., Tsang, I. W., and Hoi, S. C.  A family of simple non-parametric kernel learning algorithms.  _The Journal of Machine Learning Research_ , 12:1313–1347, 2011.

##  Appendix A Appendix

###  A.1 Notations

  * •

$\otimes$ stands for the Kronecker product, and $\odot$ stands for the
penetrating face product.

  * •

Without specific definition, $\dot{\sigma}$ represents the derivative of
$\sigma$, and $\delta$ represents the Dirac delta function.
$\delta_{kk^{\prime}}=0$ except for the case $k=k^{\prime}$, $\delta_{kk}=1$.

  * •

$||.||_{op}$ is the Operator Norm, and $||.||_{F}$ is the Frobenius Norm.

###  A.2 Additional Experiment Settings

In the CIFAR-10 experiment, pretraining is considered a trick to select a good
initialization of $a$. It refers to training the NN with a small learning rate
to a low accuracy threshold on the validation set before generating the
empirical kernels. The idea of pretraining comes from the observation that the
(W)NTK’s performance is significantly improved if the corresponding NN is
trained for some epochs. To ensure that the improvement of the NTK is not
attributed to its corresponding NN, pretraining only attains a relatively low
accuracy. For training sets of size $\leq 1,000$, the accuracy threshold is
set to 60%, and for others, the accuracy threshold is set to 70%, where the NN
is usually trained for only several epochs. In this experiment, the Adam
optimizer is used for NN training. Each NN’s accuracy is calculated using the
parameters with the highest accuracy on the validation set trained in 200
epochs. Due to efficiency considerations, for all experiments, we do not use
training tricks like batch normalization, dropout, data augmentation, Etc.

We conduct four binary classification experiments on several datasets of
CIFAR-10. These datasets are chosen because their NTK accuracy is not too
high. The utilization of datasets with high NTK accuracy may influence the
WNTK’s performance improvement, as both methods are already close to the upper
limit of classification accuracy.

In our experiment, the WCNTK is calculated empirically. To understand it
better, we can find a detailed explanation of the analytic WCNTK in Arora et
al. (2019a).

In the UCI experiment, the five-layer architecture inevitably reduces the
performance of NNs. Such an architecture is chosen to better reflect the
effect of weight terms in kernel estimators, as generally, kernel machines are
more often used and have better performance on small datasets, compared to
NNs.

###  A.3 Proof of Theorem 1.

Theorem 1 (Stability at Initialization). Consider a network of depth L at
initialization, with a Lipschitz nonlinearity $\sigma$. Suppose that, for
every parameter $\theta_{p}$ in the $l$-th layer, its weight term in the WNTK
follows a layer-wise distribution $\rho_{l}$. In the limit as the layer width
$d_{1},...,d_{L-1}$ tends to infinity sequentially, the layer-wise WNTK
converges in probability to a deterministic limiting kernel

$\displaystyle\mathcal{A}\to\mathcal{A}^{(L)}_{\infty}\otimes Id_{d_{L}}.$

With $\mu_{l}$ referring to the expectation of all element-wise weights which
follow the layer-wise distribution $\rho_{l}$, the scalar kernel
$\mathcal{A}^{(L)}_{\infty}:{\mathbb{R}}^{d_{0}}\times{\mathbb{R}}^{d_{0}}\to\mathbb{R}$
is defined recursively by

$\displaystyle\mathcal{A}^{(1)}_{\infty}(x_{i},x_{j})=$ $\displaystyle\
\mu_{1}\Sigma^{(1)}(x_{i},x_{j}),$
$\displaystyle\mathcal{A}^{(l)}_{\infty}(x_{i},x_{j})=$ $\displaystyle\
\mathcal{A}^{(l-1)}_{\infty}(x_{i},x_{j})\dot{\Sigma}^{(l)}(x_{i},x_{j})$
$\displaystyle+\mu_{l}\Sigma^{(l)}(x_{i},x_{j}),$

where

$\displaystyle\Sigma^{(1)}(x_{i},x_{j})$
$\displaystyle=\frac{1}{d_{0}}x^{\top}_{i}x_{j},$
$\displaystyle\Sigma^{(l)}(x_{i},x_{j})$
$\displaystyle=\mathbb{E}_{f\sim\mathcal{N}(0,\Sigma^{(l-1)})}[\sigma(f(x_{i}))\sigma(f(x_{j}))],$
$\displaystyle\dot{\Sigma}^{(l)}(x_{i},x_{j})$
$\displaystyle=\mathbb{E}_{f\sim\mathcal{N}(0,\Sigma^{(l-1)})}[\dot{\sigma}(f(x_{i}))\dot{\sigma}(f(x_{j}))].$

Proof. The proof is by induction. In the case of depth $=1$,

$\displaystyle\ \mathcal{A}^{(1)}_{kk^{\prime}}(x_{i},x_{j})$ $\displaystyle=$
$\displaystyle\ \sum_{\mathrm{\theta_{p}\ in\ layer\
1}}a_{p}[\partial_{\theta_{p}}f^{(1)}_{k}(\theta,x_{i})\times\partial_{\theta_{p}}f^{(1)}_{k^{\prime}}(\theta,x_{j})]\delta_{kk^{\prime}}$
$\displaystyle\to$ $\displaystyle\
\mu_{1}\frac{1}{d_{0}}x^{T}_{i}x_{j}\delta_{kk^{\prime}}$ $\displaystyle\to$
$\displaystyle\ \mu_{1}\Sigma^{(1)}(x_{i},x_{j})\delta_{kk^{\prime}}$
$\displaystyle=$ $\displaystyle\
\mathcal{A}^{(1)}_{\infty}(x_{i},x_{j})\delta_{kk^{\prime}}.$

So,

$\displaystyle\mathcal{A}^{(1)}\to\ \mathcal{A}^{(1)}_{\infty}\otimes
Id_{d_{1}}.$

In the case of depth $=l+1$, by the induction hypothesis, we assume that

$\displaystyle\mathcal{A}^{(l)}\to\mathcal{A}^{(l)}_{\infty}\otimes
Id_{d_{l}},$

or, equivalently, if we reduce the dimension of the last formula,

$\displaystyle\mathcal{A}_{kk^{\prime}}^{(l)}(x_{i},x_{j})\to\mathcal{A}^{(l)}_{\infty}(x_{i},x_{j})\delta_{kk^{\prime}}.$

In the following, we will demonstrate
$\mathcal{A}^{(l+1)}\to\mathcal{A}^{(l+1)}_{\infty}\otimes Id_{d_{l+1}}$.

$\displaystyle\mathcal{A}_{kk^{\prime}}^{(l+1)}(x_{i},x_{j})$ $\displaystyle=$
$\displaystyle\
(\partial_{\theta}f_{k}^{(l+1)}(\theta,x_{i}))^{T}(a\odot\partial_{\theta}f_{k^{\prime}}^{(l+1)}(\theta,x_{j}))\delta_{kk^{\prime}}.$

Let’s divide the sum of all parameters into two parts: the sum of parameters
in the first $l$ layer and the sum of parameters in the $(l+1)$ layer. In the
first case, we have $\partial_{\theta}f^{(l+1)}_{k}(\theta,x_{i})$
$=\frac{1}{\sqrt{d_{l}}}\sum^{d_{l}}_{m=1}\partial_{\theta}(f_{k}^{(l)}(\theta,x_{i}))\dot{\sigma}(f_{k}^{(l)}(\theta,x_{i}))W^{(l)}_{mk}$.

So, as $d_{l}\to\infty$, we can make the following approximation by the law of
large number,

$\displaystyle(\partial_{\theta}f_{k}^{(l+1)}(\theta,x_{i}))^{T}(a\odot\partial_{\theta}f_{k^{\prime}}^{(l+1)}(\theta,x_{j}))$
$\displaystyle\to$ $\displaystyle\
\dot{\Sigma}^{(l+1)}(x_{i},x_{j})\sum_{m=1}^{d_{l}}\mathcal{A}^{(l)}_{kk^{\prime}}(x_{i},x_{j})W^{(l)}_{mk}W^{(l)}_{mk^{\prime}}$
$\displaystyle\to$ $\displaystyle\
\mathcal{A}^{(l)}_{\infty}(x_{i},x_{j})\dot{\Sigma}^{(l+1)}(x_{i},x_{j})\delta_{kk^{\prime}}.$

In the second case, just similar to the case of depth $=1$, we can easily find
out,

$\displaystyle(\partial_{\theta}f_{k}^{(l+1)}(\theta,x_{i}))^{T}(a\odot\partial_{\theta}f_{k^{\prime}}^{(l+1)}(\theta,x_{j}))$
$\displaystyle\to$ $\displaystyle\ \mu_{l}\Sigma^{(l+1)}(x_{i},x_{j}).$

Now, it is easy to conclude that,

$\displaystyle\mathcal{A}_{kk^{\prime}}^{(l+1)}(x_{i},x_{j})$
$\displaystyle\to$ $\displaystyle\
(\mathcal{A}^{(l)}_{\infty}(x_{i},x_{j})\dot{\Sigma}^{(l+1)}(x_{i},x_{j})+\mu_{l}\Sigma^{(l+1)}(x_{i},x_{j}))\delta_{kk^{\prime}}$
$\displaystyle\to$ $\displaystyle\
\mathcal{A}^{(l+1)}_{\infty}(x_{i},x_{j})\delta_{kk^{\prime}},$

or, equivalently,

$\displaystyle\mathcal{A}^{(l+1)}\to\mathcal{A}^{(l+1)}_{\infty}\otimes
Id_{d_{l+1}}.$

$\hfill\blacksquare$

###  A.4 Proof of Theorem 2.

Assumptions [1-4]:

  * •

The widths of the hidden layers are identical, i.e. $d_{1}=...=d_{L-1}=d$.

  * •

The analytic WNTK is full-rank. And we set $\eta_{\mathrm{critical}}$
$=2(\lambda_{\min}+\lambda_{\max})^{-1}>0$. $\lambda_{\min}$ and
$\lambda_{\max}$ are the smallest and the largest eigenvalues of the WNTK.
$\lambda_{\max}\geq\lambda_{\min}>0$.

  * •

The training set $(\mathcal{X},\mathcal{Y})$ is contained in some compact sets
and $x_{i}\neq x_{j}$ for all $x_{i},x_{j}\in\mathcal{X}$.

  * •

The activation function $\sigma$ satisfies: $|\sigma(0)|$,
$||\dot{\sigma}||_{\infty}$, $\sup_{x_{i}\neq
x_{j}}|\dot{\sigma}(x_{i})-\dot{\sigma}(x_{j})|/|x_{i}-x_{j}|<\infty$.

###### Lemma 1

(Local Lipschitzness of the Jacobean). There is a $K>0$ such that for every
$C>0$, with high probability over random initialization (w.h.p.o.r.i.) the
following holds:

$\displaystyle\begin{cases}d^{-\frac{1}{2}}\frac{||J(\theta)-J(\theta^{\prime})||_{F}}{||\theta-\theta^{\prime}||_{2}}\leq
K\\\ d^{-\frac{1}{2}}||J(\theta)||_{F}\leq
K\end{cases},\forall\theta,\theta^{\prime}\in
B(\theta_{0},Cd^{-\frac{1}{2}}),$

where

$\displaystyle
B(\theta_{0},R):=\\{\theta:||\theta-\theta_{0}||_{2}<R\\}\textit{\ and \
}J(\theta)\in\mathbb{R}^{n\times P}.$

Note that $d_{L}$ is assumed to be 1 in the following. The proof of Lemma 1
can be found in Lee’s work (2019). The proof of Theorem 2 relies on the Local
Lipschitzness of the Jacobean.

Theorem 2 (Stability during Training). Under Assumptions 1-4 and suppose that
an L-layer neural network $f$ with random initialization is trained by
gradient decent with different learning rate $\eta a_{p}$ on different
parameter $\theta_{p}$, where $\eta=\frac{\eta_{0}}{d}$. For $\delta_{0}>0$
and $\eta_{0}<\eta_{\mathrm{critical}}$, there exist $R_{0}>0$,
$N\in\mathbb{N}$, and $K>1$ such that for every $d\geq N$, we have the
following inequality with probability at least $(1-\delta_{0})$:

$\displaystyle||\mathcal{A}_{0}-\mathcal{A}_{t}||_{F}\leq\frac{6K^{3}R_{0}a_{\max}^{2}}{\lambda_{\min}}d^{-\frac{1}{2}},\forall
t,$

where $\mathcal{A}_{0}$ and $\mathcal{A}_{t}$ are the WNTK at moment $0$ and
$t$ under standard parameterization, and $a_{\max}=\max\\{a_{i}|i=1,...,P\\}$.

Proof. First, take $K$ and $D_{0}$ big enough that the Local Lipschitzness of
the Jacobean holds with the probability $(1-\delta_{0}/3)$, $\forall d>D_{0}$.
Let’s assume that $\forall t,\theta_{t}\in B(\theta_{0},Cd^{-\frac{1}{2}})$
with $C=\frac{3KR_{0}a_{\max}}{\lambda_{\min}}$. According to the definition
of WNTK and Lemma 1, we have

$\displaystyle\ ||\mathcal{A}_{0}-\mathcal{A}_{t}||_{F}$ $\displaystyle=$
$\displaystyle\ \frac{1}{d}||J(\theta_{0})(a\odot
J(\theta_{0})^{T})-J(\theta_{t})(a\odot J(\theta_{t})^{T})||_{F}$
$\displaystyle\leq$ $\displaystyle\
\frac{a_{\max}}{d}(||J(\theta_{0})||_{op}||J(\theta_{0})^{T}-J(\theta_{t})^{T}||_{F}$
$\displaystyle\
+||J(\theta_{t})||_{F}||J(\theta_{t})^{T}-J(\theta_{0})^{T}||_{op})$
$\displaystyle\leq$ $\displaystyle\
a_{\max}2K^{2}||\theta_{0}-\theta_{t}||_{2}$ $\displaystyle\leq$
$\displaystyle\
\frac{6K^{3}R_{0}a_{\max}^{2}}{\lambda_{\min}}d^{-\frac{1}{2}}.$

So, our target is now to prove by induction the following inequality with a
known small probability:

$\displaystyle||\theta_{t}-\theta_{0}||_{2}\leq\frac{3KR_{0}a_{\max}}{\lambda_{\min}}d^{-\frac{1}{2}},\forall
t.$

At initialization, by applying Lemma 1, We define
$g(\theta)=f(\theta,\mathcal{X})-\mathcal{Y}$ and we have

$\displaystyle||\theta_{1}-\theta_{0}||_{2}\leq$ $\displaystyle\ \eta
a_{\max}||J(\theta_{0})||_{op}||g(\theta_{0})||_{2}$ $\displaystyle\leq$
$\displaystyle\ \eta_{0}Ka_{\max}||g(\theta_{0})||_{2}d^{-\frac{1}{2}}.$

Since $f$ converges in distribution to a mean zero Gaussian, there exists
$R_{0}$ and $D_{1}$, such that for every $d>D_{1}$, with probability at least
$(1-\delta/3)$ over random initialization,

$\displaystyle||g(\theta_{0})||_{2}<R_{0}.$

According to Assumption 2, we have
$\eta_{0}\leq\frac{2}{\lambda_{\min}+\lambda_{\max}}$. Hence,
$\eta_{0}\leq\frac{3}{\lambda_{\min}}$ and we can prove that,

$\displaystyle||\theta_{1}-\theta_{0}||_{2}\leq\frac{3KR_{0}a_{\max}}{\lambda_{\min}}d^{-\frac{1}{2}}.$

From now on, we will prove the part of iteration. We assume that $\forall
t^{\prime}\leq
t,||\theta_{t^{\prime}}-\theta_{0}||_{2}\leq\frac{3KR_{0}a_{\max}}{\lambda_{\min}}d^{-\frac{1}{2}}$.
We have

$\displaystyle||\theta_{t+1}-\theta_{0}||_{2}$
$\displaystyle\leq\sum_{t^{\prime}=0}^{t}||\theta_{t^{\prime}+1}-\theta_{t^{\prime}}||_{2}$
$\displaystyle\leq\eta
a_{\max}\sum_{t^{\prime}=0}^{t}||J(\theta_{t}^{\prime})||_{op}||g(\theta_{t^{\prime}})||_{2}$
$\displaystyle\leq\eta_{0}Ka_{\max}d^{-\frac{1}{2}}\sum_{t^{\prime}=0}^{t}||g(\theta_{t^{\prime}})||_{2}.$

Suppose we have $\forall t^{\prime}\leq
t,||g(\theta_{t^{\prime}})||_{2}\leq(1-\frac{\eta_{0}\lambda_{\min}}{3})^{t^{\prime}}R_{0}$.
It can be deduced that
$\sum_{t^{\prime}=0}^{t}||g(\theta_{t^{\prime}})||_{2}\leq\frac{3}{\eta_{0}\lambda_{min}}R_{0}$,
and
$||\theta_{t+1}-\theta_{0}||_{2}\leq\frac{3KR_{0}a_{\max}}{\lambda_{\min}}d^{-\frac{1}{2}}$.
So, in the rest of this proof, we will demonstrate the existence of the upper
bound of $g(\theta_{t^{\prime}}),\forall t^{\prime}\leq t$. We will still
finish this proof by induction.

At initialization, it has been discussed above that $g(\theta_{0})\leq R_{0}$.
For the part of iteration, we assume that we already have $\forall
t^{\star}\leq
t^{\prime},||g(\theta_{t^{\star}})||_{2}\leq(1-\frac{\eta_{0}\lambda_{\min}}{3})^{t^{\star}}R_{0}$.
So,

$\displaystyle||g(\theta_{t^{\prime}+1})||_{2}=$ $\displaystyle\
||g(\theta_{t^{\prime}+1})-g(\theta_{t^{\prime}})+g(\theta_{t^{\prime}})||_{2}$
$\displaystyle=$ $\displaystyle\
||J(\theta_{t^{\prime}})(\theta_{t^{\prime}+1}-\theta_{t^{\prime}})+g(\theta_{t^{\prime}})||_{2}$
$\displaystyle=$ $\displaystyle\ ||-\eta
J(\theta^{\prime}_{t^{\prime}})(a\odot
J(\theta_{t^{\prime}})^{T})g(\theta_{t^{\prime}})+g(\theta_{t^{\prime}})||_{2}$
$\displaystyle\leq$ $\displaystyle\ ||1-\eta
J(\theta^{\prime}_{t^{\prime}})(a\odot
J(\theta_{t^{\prime}})^{T})||_{op}||g(\theta_{t^{\prime}})||_{2}$
$\displaystyle\leq$ $\displaystyle\ ||1-\eta
J(\theta^{\prime}_{t^{\prime}})(a\odot J(\theta_{t^{\prime}})^{T})||_{op}$
$\displaystyle\ (1-\frac{\eta_{0}\lambda_{\min}}{3})^{t^{\prime}}R_{0}.$

where where $\theta^{\prime}_{t^{\prime}}$ is some linear interpolation
between $\theta_{t^{\prime}}$ and $\theta_{t^{\prime}+1}$. It remains to show
with a known small probability

$\displaystyle||1-\eta J(\theta^{\prime}_{t^{\prime}})(a\odot
J(\theta_{t^{\prime}})^{T})||_{op}\leq 1-\frac{\eta_{0}\lambda_{min}}{3}.$

We have,

$\displaystyle\ ||1-\eta J(\theta^{\prime}_{t^{\prime}})(a\odot
J(\theta_{t^{\prime}})^{T})||_{op}$ $\displaystyle\leq$ $\displaystyle\
||1-\eta_{0}\mathcal{A}_{\infty}^{(L)}||_{op}+\eta_{0}||\mathcal{A}_{\infty}^{(L)}-\mathcal{A}_{0}||_{op}$
$\displaystyle\ +\eta||J(\theta_{0})(a\odot
J(\theta_{0})^{T})-J(\theta^{\prime}_{t^{\prime}})(a\odot
J(\theta_{t^{\prime}})^{T})||_{op}$ $\displaystyle\leq$ $\displaystyle\
||1-\eta_{0}\mathcal{A}_{\infty}^{(L)}||_{op}+\eta_{0}||\mathcal{A}_{\infty}^{(L)}-\mathcal{A}_{0}||_{op}$
$\displaystyle\
+a_{\max}\eta_{0}K^{2}(||\theta_{t^{\prime}}-\theta_{0}||_{2}+||\theta^{\prime}_{t^{\prime}}-\theta_{0}||_{2}).$

Assumption 2 implies

$||1-\eta_{0}\mathcal{A}_{\infty}^{(L)}||_{op}\leq 1-\eta_{0}\lambda_{\min}.$
(22)

As we have proved the convergence of $\mathcal{A}_{0}$ in Theorem 1, one can
find $D_{2}$ independent of $t^{\prime}$, such that the following event holds
with probability $(1-\delta_{0}/3)$,

$||\mathcal{A}_{\infty}^{(L)}-\mathcal{A}_{0}||_{op}\leq\frac{\lambda_{\min}}{3}.$
(23)

With (22), (23) and Lemma 1, we have

$\displaystyle\ ||1-\eta J(\theta^{\prime}_{t^{\prime}})(a\odot
J(\theta_{t^{\prime}})^{T})||_{op}$ $\displaystyle\leq$ $\displaystyle\
1-\eta_{0}\lambda_{\min}+\frac{\eta_{0}\lambda_{\min}}{3}+2a_{\max}^{2}\eta_{0}K^{2}\frac{3KR_{0}}{\sqrt{d}\lambda_{\min}}$
$\displaystyle\leq$ $\displaystyle\ 1-\frac{\eta_{0}\lambda_{\min}}{3},$

under the constraint that,

$\displaystyle
d\geq(\frac{18K^{3}R_{0}a_{\max}^{2}}{\lambda_{\min}^{2}})^{2}=D_{3}.$

Therefore, we only need to set $N=\max\\{D_{0},D_{1},D_{2},D_{3}\\}$ and we
can finally conclude that,

$\displaystyle||\mathcal{A}_{t}-\mathcal{A}_{0}||_{F}\leq\frac{6K^{3}R_{0}a_{\max}^{2}}{\lambda_{\min}}d^{-\frac{1}{2}}=\mathcal{O}(d^{-\frac{1}{2}}),\forall
t,$

with probability $p\geq(1-\delta_{0}/3)^{3}\geq(1-\delta_{0})$.
$\hfill\blacksquare$

###  A.5 Proof of Theorem 3.

Theorem 3 (WNTK Lazy Regime). Under Assumptions 1-4, applying gradient descent
with different learning rate $\eta a_{p}$ on different parameter $\theta_{p}$,
for every $x\in{\mathbb{R}}^{d_{0}}$ satisfying $||x||_{2}\leq 1$, for all
$t\geq 0$, as $d\to\infty$, with probability arbitrarily close to 1 over
random initialization, we have

$\displaystyle||f(\theta_{t},\mathcal{X})-f^{\mathrm{lin}}_{t}(\mathcal{X})||_{2}\to
0,$ $\displaystyle||f(\theta_{t},x)-f^{\mathrm{lin}}_{t}(x)||_{2}\to 0.$

Proof. This proof follows a similar proof in Lee et al. (2019). Via continuous
time gradient descent, the evolution of the parameters and the output function
of the neural network in this theorem can be written as follows.

$\displaystyle\dot{\theta}_{t}$
$\displaystyle=-\eta(a\odot\nabla_{\theta}f(\theta_{t},\mathcal{X}))^{T}\nabla_{f(\theta_{t},\mathcal{X})}\mathcal{L},$
$\displaystyle\dot{f}(\theta_{t},x)$
$\displaystyle=-\eta\mathcal{A}_{t}(x,\mathcal{X})\nabla_{f(\theta_{t},\mathcal{X})}\mathcal{L}.$

Correspondingly, the dynamics of gradient flow using the linearized model are
governed by,

$\displaystyle\dot{\theta}_{t}$
$\displaystyle=-\eta(a\odot\nabla_{\theta}f_{0}(\mathcal{X}))^{T}\nabla_{f^{\mathrm{lin}}_{t}(\mathcal{X})}\mathcal{L},$
$\displaystyle\dot{f}^{\mathrm{lin}}_{t}(x)$
$\displaystyle=-\eta\mathcal{A}_{0}(x,\mathcal{X})\nabla_{f^{\mathrm{lin}}_{t}(\mathcal{X})}\mathcal{L}.$

We define $g^{\mathrm{lin}}(t)=f^{\mathrm{lin}}_{t}(\mathcal{X})-\mathcal{Y}$
and $g(t)=f(\theta_{t},\mathcal{X})-\mathcal{Y}$. We begin our proof by the
following equation:

$\displaystyle\frac{d}{dt}(\exp(\eta\mathcal{A}_{0}t)(g^{\mathrm{lin}}(t)-g(t))$
$\displaystyle=$
$\displaystyle\eta(\exp(\eta\mathcal{A}_{0}t)(\mathcal{A}_{t}-\mathcal{A}_{0})g(t)).$

Integrating both sides and using the fact $g^{\mathrm{lin}}(0)=g(0)$,

$\displaystyle(g^{\mathrm{lin}}(t)-g(t))=$
$\displaystyle-\int_{0}^{t}\eta(\exp(\eta\mathcal{A}_{0}(s-t))(\mathcal{A}_{s}-\mathcal{A}_{0})(g^{\mathrm{lin}}(s)-g(s)))ds$
$\displaystyle+\int_{0}^{t}\eta(\exp(\eta\mathcal{A}_{0}(s-t))(\mathcal{A}_{s}-\mathcal{A}_{0})g^{\mathrm{lin}}(s))ds.$

Taking the norm gives

$\displaystyle||g^{\mathrm{lin}}(t)-g(t)||_{2}\leq$
$\displaystyle\eta\int_{0}^{t}||\exp^{\eta\mathcal{A}_{0}(s-t)}||_{op}||\mathcal{A}_{s}-\mathcal{A}_{0}||_{op}||g^{\mathrm{lin}}(s)-g(s)||_{2}ds$
$\displaystyle+\eta\int_{0}^{t}||\exp^{\eta\mathcal{A}_{0}(s-t)}||_{op}||\mathcal{A}_{s}-\mathcal{A}_{0}||_{op}||g^{\mathrm{lin}}(s)||_{2}ds.$

Let $\lambda_{0}>0$ be the smallest eigenvalue of $\mathcal{A}_{0}$.

$\displaystyle||g^{\mathrm{lin}}(t)-g(t)||_{2}\leq$
$\displaystyle\eta\int_{0}^{t}\exp^{\eta\lambda_{0}(s-t)}||\mathcal{A}_{s}-\mathcal{A}_{0}||_{op}||g^{\mathrm{lin}}(s)-g(s)||_{2}ds$
$\displaystyle+\eta\int_{0}^{t}\exp^{\eta\lambda_{0}(s-t)}||\mathcal{A}_{s}-\mathcal{A}_{0}||_{op}||g^{\mathrm{lin}}(s)||_{2}ds.$

We conduct the following variable replacement. Let

$\displaystyle u(t)$
$\displaystyle=\exp^{\eta\lambda_{0}t}||g^{\mathrm{lin}}(s)-g(s)||_{2},$
$\displaystyle\alpha(t)$
$\displaystyle=\eta\int_{0}^{t}\exp^{\eta\lambda_{0}(s)}||\mathcal{A}_{s}-\mathcal{A}_{0}||_{op}||g^{\mathrm{lin}}(s)||_{2}ds,$
$\displaystyle\beta(t)$
$\displaystyle=\eta||\mathcal{A}_{t}-\mathcal{A}_{0}||_{op}.$

The above inequality can be rewritten as

$\displaystyle u(t)\leq\alpha(t)+\int_{0}^{t}\beta{s}u(s)ds.$

By applying an integral form of the Gronwall’s inequality (Dragomir, 2003), we
can get:

$\displaystyle u(t)\leq\alpha(t)\exp(\int_{0}^{t}\beta{s}ds).$

According to the dynamics of the linearized network, we have

$\displaystyle||g^{\mathrm{lin}}(t)-g(t)||_{2}\leq\exp^{-\lambda_{0}\eta
t}||g^{\mathrm{lin}}(0)||_{2}.$

Then

$\displaystyle||g^{\mathrm{lin}}(t)-g(t)||_{2}\leq$ (24)
$\displaystyle\exp^{-\lambda_{0}\eta
t}||g^{\mathrm{lin}}(0)||_{2}\int_{0}^{t}\beta(s)ds\exp(\int_{0}^{t}\beta(s)ds).$

Let $\sigma_{t}=\sup_{0\leq s\leq t}||\mathcal{A}_{s}-\mathcal{A}_{0}||_{op}$.
According to Theorem 2,

$\displaystyle\forall t,\sigma_{t}=O(d^{-\frac{1}{2}})\to 0.$ (25)

Thus there exists $P(t)$ a polynomial, such that

$\displaystyle||g^{\mathrm{lin}}(t)-g(t)||_{2}\leq(\eta
t\sigma_{t}\exp^{-\lambda\eta t+\sigma_{t}\eta
t})||g^{\mathrm{lin}}(0)||_{2}P(t).$

Therefore,

$\displaystyle||f(\theta_{t},\mathcal{X})-f^{\mathrm{lin}}_{t}(\mathcal{X})||_{2}$
$\displaystyle=$ $\displaystyle||g(t)-g^{\mathrm{lin}}(t)||_{2}\leq
O(d^{-\frac{1}{2}})\to 0.$

Similarly, we can prove the convergence on a test sample. Let $y$ be the true
label and We define $g^{\mathrm{lin}}(t,x)=f^{\mathrm{lin}}_{t}(x)-y$ and
$g(t,x)=f(\theta_{t},x)-y$. We have

$\displaystyle\frac{d}{dt}(g^{\mathrm{lin}}(t,x)-g(t,x))$ $\displaystyle=$
$\displaystyle-\eta(\mathcal{A}_{0}(x,\mathcal{X})-\mathcal{A}_{t}(x,\mathcal{X}))g^{\mathrm{lin}}(t)$
$\displaystyle+\eta\mathcal{A}_{t}(x,\mathcal{X})(g(t)-g^{\mathrm{lin}}(t)).$

We conduct the same integration over $[0,t]$ and take the norm.

$\displaystyle||g^{\mathrm{lin}}(t,x)-g(t,x)||_{2}$ $\displaystyle\leq$
$\displaystyle\eta\int_{0}^{t}\exp^{-\eta\lambda_{0}s}||\mathcal{A}_{s}(x,\mathcal{X})-\mathcal{A}_{0}(x,\mathcal{X})||_{2}ds$
$\displaystyle+\eta\int_{0}^{t}(||\mathcal{A}_{0}(x,\mathcal{X})||_{2}+||\mathcal{A}_{s}(x,\mathcal{X})-\mathcal{A}_{0}(x,\mathcal{X})||_{2})$
$\displaystyle||g(s)-g^{\mathrm{lin}}(s)||_{2}ds.$

Lemma 1 can imply, for all $t$,

$\displaystyle||\mathcal{A}_{s}(x,\mathcal{X})-\mathcal{A}_{0}(x,\mathcal{X})||_{2}=O(d^{-\frac{1}{2}}).$

This implies that the first term in the right side of the inequality
$=O(d^{-\frac{1}{2}})$. In terms of the the second term, we can find its
upper-bound by exactly the same way as in the proof on the training set. By
applying (24) and (25), we can prove that the second term
$\leq||\mathcal{A}_{0}(x,\mathcal{X})||_{2}\int_{0}^{t}(\eta
s\sigma_{s}\exp^{-\eta\lambda_{0}s+\sigma_{s}\eta
s})||g^{\mathrm{lin}}(0)||_{2}dtP(t)=O(d^{-\frac{1}{2}})$. So, we end the
whole proof by

$\displaystyle\ ||f(\theta_{t},x)-f^{\mathrm{lin}}_{t}(x)||_{2}$
$\displaystyle=$ $\displaystyle\ ||g(t,x)-g^{\mathrm{lin}}(t,x)||_{2}\leq
O(d^{-\frac{1}{2}})\to 0.$

$\hfill\blacksquare$