On Exact Computation with an Infinitely Wide Neural Net

Introduction

How well does a classic deep net architecture like AlexNet or VGG19 perform on a standard dataset such as CIFAR-10 when its “width”— namely, number of channels in convolutional layers, and number of nodes in fully-connected internal layers — is allowed to increase to infinity? Questions about these “infinite limits” of deep nets have naturally emerged in the ongoing effort to understand the power of deep learning. In mathematics it is often easier to study objects in the infinite limit. Furthermore, the infinite limit could conceivably make sense in deep learning, since over-parametrization seems to help optimization a lot and doesn’t hurt generalization much (Zhang et al., 2017): deep neural nets with millions of parameters work well even for datasets with $50$k training examples. So why not imagine nets whose width goes to infinity?

Allowing width to go to infinity also connects deep learning in an interesting way with other areas of machine learning. A single hidden-layer neural network with i.i.d. random parameters, in the limit of infinite width, is a function drawn from a Gaussian process (GP) (Neal, 1996). This model as well as analogous ones with multiple layers (Lee et al., 2018; Matthews et al., 2018), convolutional filters (Novak et al., 2019; Garriga-Alonso et al., 2019) and other architectures (Yang, 2019) make up the GP view of deep learning. These correspond to infinitely wide deep nets whose all parameters are chosen randomly (with careful scaling), and only the top (classification) layer is optimized.

From now on we will use weakly-trained nets to refer to nets whose layers receive random initialization and only the top layer is trained by gradient descent. We use fully-trained to refer to nets whose all parameters are trained by gradient descent. It has long been known that weakly-trained convolutional nets have reasonable performance on MNIST and CIFAR-10. Weakly-trained nets that are fully-connected instead of convolutional, can also be thought of as “multi-layer random kitchen sinks,” which also have a long history.

where $\bm{x},\bm{x}^{\prime}$ are two inputs. This kernel method makes sense when the width goes to infinity.

The objects of interest in this paper are not weakly-trained nets, but fully-trained nets. In the finite case, analysis of optimization and generalization of fully-trained nets is of course an open problem. One may also ask:

Can we understand the power of fully-trained nets whose width goes to infinity?

A priori this question doesn’t seem any easier than the finite case, and empirical evaluation seems computationally infeasible due to the infinite limit. They also do not correspond to a kernel method in any obvious way.

Recent papers suggest that neural nets whose width greatly exceeds the number of training data points can rapidly reduce training error to via gradient descent, and under some conditions, the trained net also exhibits good generalization (Du et al., 2019, 2018b; Li and Liang, 2018; Allen-Zhu et al., 2018a, b; Zou et al., 2018; Arora et al., 2019; Cao and Gu, 2019). Extra-wideness plays a crucial role in the proof: it is shown that as width increases, training causes increasingly smaller changes (in a proportionate sense) in the parameters. This raises the possibility that as one increases the width to infinity, a certain limiting behavior can emerge even in the fully-trained net. A recent paper by Jacot et al. (2018) isolated a notion implicit in the above papers, which they called the Neural Tangent Kernel (NTK). They suggested — via a proof that is slightly heuristic — that this fixed kernel characterizes the behavior of fully-connected infinite width neural networks whose layers have been trained by gradient descent. The NTK is different from the Gaussian process kernels discussed earlier, and is defined using the gradient of the output of the randomly initialized net with respect to its parameters, i.e.,

Here, the gradient $\frac{\partial f(\bm{\theta},\bm{x})}{\partial\bm{\theta}}$ appears from considering gradient descent, as will be explained in Section 3. One may also generalize the NTK to convolutional neural nets, and we call the corresponding kernel Convolutional Neural Tangent Kernel (CNTK).

Though NTK and CNTK are defined by an infinite limit, a recent paper (Lee et al., 2019) attempted to understand their properties via a finite approximation of the infinite limit kernel by Monte Carlo methods. However, as will be shown in Section B, using random features generated from practically sized nets can degrade the performance a lot. It was still open what is the full power of exact CNTK on modern datasets. This is a challenging question especially for CNTK with pooling operations, since when convolution with pooling is involved, it was believed that exact computation of kernels (for either convolutional Gaussian process kernel or CNTK) is infeasible for large datasets like CIFAR-10 (Novak et al., 2019).

Furthermore, we give a more rigorous, non-asymptotic proof that the NTK captures the behavior of a fully-trained wide neural net under weaker condition than previous proofs. We also experimentally show that the random feature methods for approximating CNTK in earlier work do not compute good approximations, which is clear from their much worse performance on CIFAR.

1 Notation

Related Work

From a Gaussian process (GP) viewpoint, the correspondence between infinite neural networks and kernel machines was first noted by Neal (1996). Follow-up work extended this correspondence to more general shallow neural networks (Williams, 1997; Roux and Bengio, 2007; Hazan and Jaakkola, 2015). More recently, this was extended to deep and convolutional neural networks (Lee et al., 2018; Matthews et al., 2018; Novak et al., 2019; Garriga-Alonso et al., 2019) and a variety of other architectures (Yang, 2019). However, these kernels, as we discussed in Section 1, represent weakly-trained nets, instead of fully-trained nets.

Beyond GPs, the connection between neural networks and kernels is also studied in the compositional kernel literature. Cho and Saul (2009) derived a closed-form kernel formula for rectified polynomial activations, which include ReLU as a special case. Daniely et al. (2016) proposed a general framework to transform a neural network to a compositional kernel and later Daniely (2017) showed for sufficiently wide neural networks, stochastic gradient descent can learn functions that lie in the corresponding reproducing kernel Hilbert space. However, the kernels studied in these works still correspond to weakly-trained neural networks.

This paper is inspired by a line of recent work on over-parameterized neural networks (Du et al., 2019, 2018b; Du and Hu, 2019; Li and Liang, 2018; Allen-Zhu et al., 2018b, a; Zou et al., 2018; Cao and Gu, 2019). These papers established that for (convolutional) neural networks with large but finite width, (stochastic) gradient descent can achieve zero training error. A key component in these papers is showing that the weight matrix at each layer is close to its initialization. This observation implies that the kernel defined in Equation (2) is still close to its initialization. Arora et al. (2019) explicitly used this observation to derive generalization bounds for two-layer over-parameterized neural networks. Chizat and Bach (2018) argued that these results in the kernel regime may be too simple to be able to explain the success of deep learning, while on the other hand, out results show that CNTK is at least able to perform well on tasks like CIFAR-10 classification. Also see the survey Fan et al. (2019) for recent advance in deep learning theory.

Jacot et al. (2018) derived the exact same kernel from kernel gradient descent. They showed that if the number of neurons per layer goes to infinity in a sequential order, then the kernel remains unchanged for a finite training time. They termed the derived kernel Neural Tangent Kernel (NTK). We follow the same naming convention and name its convolutional extension Convolutional Neural Tangent Kernel (CNTK). Later, Yang (2019) derived a formula of CNTK as well as a mechanistic way to derive NTK for different architectures. Comparing with (Yang, 2019), our CNTK formula has a more explicit convolutional structure and results in an efficient GPU-friendly computation method. Recently, Lee et al. (2019) tried to empirically verify the theory in (Jacot et al., 2018) by studying the linearization of neural nets. They observed that in the first few iterations, the linearization is close to the actual neural net. However, as will be shown in Section B, such linearization can decrease the classification accuracy by $5\%$ even on a “CIFAR-2" (airplane V.S. car) dataset. Therefore, exact kernel evaluation is important to study the power of NTK and CNTK.

Neural Tangent Kernel

Note that the above dynamics is identical to the dynamics of kernel regression under gradient flow, for which at time $t\to\infty$ the final prediction function is (assuming $\bm{u}(0)=\bm{0}$)

In Theorem 3.2, we rigorously prove that a fully-trained sufficiently wide ReLU neural network is equivalent to the kernel regression predictor (5) on any given data point.

To give the formula of NTK, we also need to define a derivative covariance:

The final NTK expression for the fully-connected neural network is

where we let $\dot{\Sigma}^{(L+1)}(\bm{x},\bm{x}^{\prime})=1$ for convenience. We refer readers to Section D for the derivation of this formula. Rigorously, for ReLU activation, we have the following theorem that gives a concrete bound on the hidden widths that is sufficient for convergence to the NTK at initialization:

The proof of Theorem 3.1 is given in Section E. Theorem 3.1 improves upon previous results (Jacot et al., 2018; Yang, 2019) that also established similar convergence in the following sense:

Previous results are asymptotic, i.e., they require the widths to go to infinity, while Theorem 3.1 gives a non-asymptotic bound on the required layer widths.

Jacot et al. (2018) required sequential limit, i.e., $d_{1},\ldots,d_{L}$ go to infinity one by one, and Yang (2019) let $d_{1},\ldots,d_{L}$ go to infinity at the same rate. On the other hand, Theorem 3.1 only requires $\min_{h\in[L]}d_{h}$ to be sufficiently large, which is the weakest notion of limit.

Equivalence between wide neural net and kernel regression with NTK.

Built on Theorem 3.1, we can further incorporate the training process and show the equivalence between a fully-trained sufficiently wide neural net and the kernel regression solution using the NTK, as described in Lemma 3.1 and the discussion after it.

Since the above solution corresponds to the linear dynamics in Equation (4) with zero initialization, in order to establish equivalence between neural network and kernel regression, we would like the initial output of the neural network to be small. Therefore, we apply a small multiplier $\kappa>0$, and let the final output of the neural network be $f_{nn}(\bm{\theta},\bm{x})=\kappa f\left(\bm{\theta},\bm{x}\right).$ We let $f_{nn}(\bm{x}_{te})=\lim_{t\rightarrow\infty}f_{nn}(\bm{\theta}(t),\bm{x}_{te})$ be the prediction of the neural network at the end of training.

The following theorem establishes the equivalence between the fully-trained wide neural network $f_{nn}$ and the kernel regression predictor $f_{ntk}$ using the NTK.

Several comments are in sequel. Theorem 3.2 is, to our knowledge, the first result that rigorously shows the equivalence between a fully-trained neural net and a kernel predictor. Comparing with (Jacot et al., 2018), our bound is non-asymptotic whereas (Jacot et al., 2018) only has an asymptotic result; furthermore, Jacot et al. (2018) required the width of every layer to go to infinity in a sequential order, while we can have the same number of neurons per layer, which is closer to practice. Comparing with recent results on over-parameterized neural nets (Arora et al., 2019; Allen-Zhu et al., 2018b, a; Du et al., 2019, 2018b; Li and Liang, 2018; Zou et al., 2018), our theorem is a more precise characterization of the learned neural network. That is, the prediction is essentially a kernel predictor. Therefore, to study the properties of these over-parameterized nets, such as their generalization power, it is sufficient to study the corresponding NTK.

While this theorem only gives guarantee for a single point, using a union bound, we can show that this guarantee holds for (exponentially many) finite testing points. Combing this with the standard analysis of hold-out validation set, we can conclude that a fully-trained wide neural net enjoys the same generalization ability as its corresponding NTK.

For the proof of Theorem 3.2, we first use a generic argument to show that the perturbation on the prediction can be reduced to the perturbation on kernel value at the initialization and during training. Theorem 3.1 guarantees a small perturbation on kernel value at initialization. For the perturbation during training, we use high level proof idea from Du et al. (2018b); Arora et al. (2019) to reduce the perturbation on the kernel value to the perturbation on the gradient of each prediction with respect to weight matrices. Then we adopt technical lemmas from Allen-Zhu et al. (2018b) to obtain bounds on the perturbation of the gradient. The proof of Theorem 3.2 is given in Section F. We remark that Jacot et al. (2018); Lee et al. (2019) provided proofs for the training part. However, both are asymptotic results and only apply to finite training time. In contrast, we give a finite-width perturbation bound and our result applies to infinite training time.

Convolutional Neural Tangent Kernel

In this section we study convolutional neural nets (CNNs) and their corresponding CNTKs. We study two architectures, vanilla CNN and CNN with global average pooling (GAP). In this section we define vanilla CNN and present its corresponding CNTK formula. The derivation of this formula is deferred to Section G. We present the definition of CNN with GAP and its CNTK in Section H.

Equation (10) shows that patch $[\bm{w}*\bm{x}]_{ij}$ depends on $[\bm{x}]_{i-\frac{q-1}{2}:i+\frac{q-1}{2},j-\frac{q-1}{2}:j+\frac{q-1}{2}}$. Our CNTK formula also relies on this dependency. For $(i,j,i^{\prime},j^{\prime})\in[{P}]\times[{Q}]\times[{P}]\times[{Q}]$, define

A vanilla CNN consisting of $L$ convolution layers and one fully-connected layer is formally defined as follows:

For $h=1,\ldots,L$, $\beta=1,\ldots,{C}^{(h)}$, the intermediate outputs are defined as

For this architecture, using the same reasoning as in Section D, we obtain the following convolutional neural tangent kernel formula. The details are provided in Section G.

We let $\bm{x},\bm{x}^{\prime}$ be two input images.

For $\alpha=1,\ldots,{C}^{(0)}$, $(i,j,i^{\prime},j^{\prime})\in[{P}]\times[{Q}]\times[{P}]\times[{Q}]$, define

For $(i,j,i^{\prime},j^{\prime})\in[{P}]\times[{Q}]\times[{P}]\times[{Q}]$, define

Note that $\bm{\Sigma}(\bm{x},\bm{x}^{\prime})$ and $\dot{\bm{\Sigma}}(\bm{x},\bm{x}^{\prime})$ share similar structures as their NTK counterparts in Equations (7) and (8). The only difference is that we have one more step, taking the trace over patches. This step represents the convolution operation in the corresponding CNN. Next, we can use a recursion to compute the CNTK:

First, we define $\bm{\Theta}^{(0)}(\bm{x},\bm{x}^{\prime})=\bm{\Sigma}^{(0)}(\bm{x},\bm{x}^{\prime})$.

For $h=1,\ldots,L-1$ and $(i,j,i^{\prime},j^{\prime})\in[{P}]\times[{Q}]\times[{P}]\times[{Q}]$, we define

For $h=L$ , we define $\bm{\Theta}^{(L)}(\bm{x},\bm{x}^{\prime})=\dot{\bm{K}}^{(L)}(\bm{x},\bm{x}^{\prime})\odot\bm{\Theta}^{(L-1)}(\bm{x},\bm{x}^{\prime})+\bm{K}^{(L)}(\bm{x},\bm{x}^{\prime}).$

In Section H we give the CNTK formula for CNNs with GAP, which is similar to vanilla CNNs. To compute the CNTK matrix corresponding to a CNN with GAP that has $L$ convolution layers and one fully-connected layer on $n$ samples, the time complexity is $O(n^{2}{P}^{2}{Q}^{2}L)$. Previous work assumed that directly computing convolutional kernel (with pooling) exactly is computationally infeasible, and thus resorted to approximations like Monte Carlo sampling (Novak et al., 2019). We are able to scale the exact CNTK computation to the full CIFAR-10 dataset and 20-layer CNN with GAP. We present our efficient computation approach in Section I.

Experiments

We evaluate the performances of CNNs and their corresponding CNTKs on the CIFAR-10 dataset. The implementation details are in Section A. We also compare the performances between CNTKs and their corresponding random feat Due to space limit, we defer these results on random features to Section B.

We test two types of architectures, vanilla CNN and CNN with global average pooling (GAP), as described in Sections 4 and H. We also test CNTKs with only 2,000 training data to see whether their performances are consistent with CNTKs and CNNs using the full training set. The results are summarized in Table 1. Notice that in Table 1, depth is the total number of layers (including both convolution layers and fully-connected layers).

Several comments are in sequel. First, CNTKs are very powerful kernels. The best kernel, 11-layer CNTK with GAP, achieves 77.43% classification accuracy on CIFAR-10. This results in a significant new benchmark for performance of a pure kernel-based method on CIFAR-10, being $10\%$ higher than methods reported in (Novak et al., 2019).

Second, we find that for both CNN and CNTK, depth can affect the classification accuracy. This observation demonstrates that depth not only matters in deep neural networks but can also affect the performance of CNTKs.

Third, the global average pooling operation can significantly increase the classification accuracy by 8% - 10% for both CNN and CNTK. Based on this finding, we expect that many techniques that improve the performance of neural networks are in some sense universal, i.e., these techniques can also benefit kernel methods.

Fourth, we find that there is still a 5% - 6% performance gap between CNTKs and CNNs. Since CNTKs exactly correspond to infinitely wide CNNs, this performance gap implies that finite width has its benefits. Therefore, it is likely that recent theoretical work on over-parameterization that operates in the NTK regime cannot fully explain the success of neural networks yet, and we believe it is an interesting open problem to characterize this gap.

Potential application in neural architecture search.

Finally, we find that performances of CNTK-V-2Ks and CNTK-GAP-2Ks are highly correlated to their CNN-V, CNTK-V, CNN-GAP and CNTK-GAP counterparts. Again we see CNTK-GAP-2Ks outperform CNTK-V-2Ks by a large margin (about $8\%$ - $9\%$). One potential application of this observation is to guide neural architecture search. We can compute the kernel on a small training data, test it on a validation set, and choose neural network architectures based on the performance of this small kernel on the validation set. We leave large scale experiments of this idea for future work.

Conclusion

By giving the first practical algorithm for computing CNTKs exactly, this paper allows investigation of the behavior of infinitely wide (hence infinitely over-parametrized) deep nets, which turns out to not be much worse than that of their finite counterparts. We also give a fully rigorous proof that a sufficiently wide net is approximately equivalent to the kernel regression predictor, thus yielding a powerful new off-the-shelf kernel. We leave it as an open problem to understand the behavior of infinitely wide nets with features such as Batch Normalization or Residual Layers. Of course, one can also hope that the analysis of infinite nets provides rigorous insight into finite ones.

Acknowledgments

We thank Jason D. Lee, Haochuan Li and Xiyu Zhai for useful discussions. S. Arora, W. Hu and Z. Li are supported by NSF, ONR, Simons Foundation, Schmidt Foundation, Mozilla Research, Amazon Research, DARPA and SRC. R. Salakhutdinov and R. Wang are supported in part by NSF IIS-1763562, Office of Naval Research grant N000141812861, and Nvidia NVAIL award. We thank Amazon Web Services for providing compute time for the experiments in this paper. This work was done while S. S. Du was a Ph.D. student at Carnegie Mellon University.

References

Appendix A Experiment Details

Due to efficiency considerations, for all experiments, we use no data augmentation. Tricks like batch normalization, dropout, weight decay, etc. are not used for proper comparison. We fix the filter $q$ to be $3$ and stride to be $1$. We use zero padding to make sure the number of patches keeps unchanged after each convolutional layer. We set the number of convolution layers to be $2$, $3$, $5$, $10$, or $20$. For both CNNs and CNTKs, we use the quadratic loss as the objective function.

and choose the class with largest value as the prediction.

The architecture of CNNs is as described in Section 4 and Section H. We set the number of the channels of the network as 1024 and $\kappa$ as $0.05$. To train CNNs, we use stochastic gradient descent (SGD) with fixed learning rate. We report the best average performance over 3 trials among the different learning rate chosen from $\{0.1,1,10\}$. The test accuracy is measured by taking average of the 10 epochs after reaching full training accuracy except the depth-3 vanilla CNN, which couldn’t attain full training accuracy within 3000 epochs for all learning rates

Our neural networks are trained using the PyTorch package, using (possibly multiple) NVIDIA Tesla V100 GPUs. We calculate the kernel values using the CuPyhttps://cupy.chainer.org. package. For time-consuming operations, we write native CUDA codes to speed up the calculation. All experiments are performed on Amazon Web Services (AWS).

Appendix B Additional Experiments on Random Features

We verify the importance of using the exact kernels instead of the approximated ones from random features (as done in [Lee et al., 2019]). The random features are generated by taking the gradient of the randomly initialized CNNs with respect to the weight matrices. For all CNNs we set the number of channels to be $128$. We compare the performances of the exact kernels and the random kernels on a CIFAR-2 dataset, i.e., the first two class in CIFAR-10. For each kernel generated by random features, we test $10$ times and report the median. The results are summarized in Table 2.

Note that even on the simple CIFAR-2 dataset, random features have much worse accuracies than exact kernels by a large margin. This experiment demonstrates the importance of using the exact kernels instead of approximated ones.

Appendix C Proof of Lemma 3.1

The parameters $\bm{\theta}$ evolve according to the differential equation

where $t\geq 0$ is a continuous time index. Under Equation (13), the evolution of the network output $f(\bm{\theta}(t),\bm{x}_{i})$ can be written as

Appendix D NTK Derivation

In this section we derive NTK for the fully-connected neural net defined in Section 3.

First we explain how the Gaussian process covariance in Equation (7) is obtained. The intuition is that $\left[\bm{f}^{(h+1)}(\bm{x})\right]_{i}=\sum_{j=1}^{d_{h}}\left[\bm{W}^{(h+1)}\right]_{i,j}\left[\bm{g}^{(h)}(\bm{x})\right]_{j}$ is a centered Gaussian process conditioned on $\bm{f}^{(h)}$ ($\forall i\in[d_{h+1}]$), with covariance

which converges to $\Sigma^{(h)}(\bm{x},\bm{x}^{\prime})$ as $d_{h}\to\infty$ given that each $\left[\bm{f}^{(h)}\right]_{j}$ is a centered Gaussian process with covariance $\Sigma^{(h-1)}$. This yields the inductive definition in Equation (7).

Recall that we need to compute the value that $\left\langle\frac{\partial f(\bm{\theta},\bm{x})}{\partial\bm{\theta}},\frac{\partial f(\bm{\theta},\bm{x}^{\prime})}{\partial\bm{\theta}}\right\rangle$ converges to at random initialization in the infinite width limit. We can write the partial derivative with respect to a particular weight matrix $\bm{W}^{(h)}$ in a compact form:

Then, for any $h\in[L+1]$, we can compute

Note that we have established in Equation (16) that

For the other factor $\left\langle\mathbf{b}^{(h)}(\bm{x}),\mathbf{b}^{(h)}(\bm{x}^{\prime})\right\rangle$, from Equation (17) we get

Although $\bm{W}^{(h+1)}$ and $\mathbf{b}_{h+1}(\bm{x})$ are dependent, the Gaussian initialization of $\bm{W}^{(h+1)}$ allows us to replace $\bm{W}^{(h+1)}$ with a fresh new sample $\widetilde{\bm{W}}^{(h+1)}$ without changing its limit: (This is made rigorous for ReLU activation in Theorem 3.1.)

Applying this approximation inductively in Equation (19), we get

Finally, since $\left\langle\frac{\partial f(\bm{\theta},\bm{x})}{\partial\bm{\theta}},\frac{\partial f(\bm{\theta},\bm{x}^{\prime})}{\partial\bm{\theta}}\right\rangle=\sum_{h=1}^{L+1}\left\langle\frac{\partial f(\bm{\theta},\bm{x})}{\partial\bm{W}^{(h)}},\frac{\partial f(\bm{\theta},\bm{x}^{\prime})}{\partial\bm{W}^{(h)}}\right\rangle$, we obtain the final NTK expression for the fully-connected neural network:

Appendix E Proof of Theorem 3.1

When $\sigma$ is k-homogeneous function, we have

Thus $t_{\sigma}(\bm{\Sigma})$ can be written as $c_{\sigma}\left(\Sigma_{11}\Sigma_{22}\right)^{\frac{k}{2}}\hat{t}({\frac{\Sigma_{12}}{\sqrt{\Sigma_{11}\Sigma_{22}}}})$,

For all activation function $\sigma$, $t_{\sigma}\left(\begin{bmatrix}1&1\\ 1&1\end{bmatrix}\right)=1$.

For all 1-homogeneous activation $\sigma$, $\hat{t}_{\sigma}(1)=1$ and $t_{\sigma}\left(\begin{bmatrix}a&a\\ a&a\end{bmatrix}\right)=a^{k}$ .

For $\sigma\left(z\right)=\max(0,z)$, $\hat{t}_{\sigma}(\rho)=\frac{\sqrt{1-\rho^{2}}+\rho\arcsin\rho}{\pi}+\frac{x}{2}$, $\hat{t}_{\dot{\sigma}}(\rho)=\frac{1}{2}+\frac{\arcsin\rho}{\pi}$ and $c_{\sigma}=c_{\dot{\sigma}}=2$.

For any $-\frac{\pi}{2}\leq y^{\prime}\leq y\leq\frac{\pi}{2}$, $\sin y-\sin y^{\prime}\geq 2\sin^{2}\frac{y-y^{\prime}}{2}$.

$\sin y\geq\frac{2y}{\pi},\ \forall y\in[0,\frac{\pi}{2}]$.

(1). From $-\frac{\pi}{2}\leq y^{\prime}\leq y^{\prime}\leq\frac{\pi}{2}$ we know $\frac{-\pi}{2}+\frac{y-y^{\prime}}{2}\leq\frac{y+y^{\prime}}{2}\leq\frac{\pi}{2}-\frac{y-y^{\prime}}{2}$, which implies that $\cos(\frac{y+y^{\prime}}{2})\geq\sin(\frac{y-y^{\prime}}{2})$. Thus,

(2). Note that $\left(\frac{\sin y}{y}\right)^{\prime}=\frac{y\cos y-\sin y}{y^{2}}=\frac{\cos y}{y^{2}}(y-\tan y)<0$, $\frac{\sin y}{y}$ is decreasing on $[0,\frac{\pi}{2}]$. Thus $\frac{\sin y}{y}\geq\frac{1}{\frac{\pi}{2}}=\frac{2}{\pi},\ \forall y\in[0,\frac{\pi}{2}]$.

(3). Let $y,y^{\prime}\in[-\frac{\pi}{2},\frac{\pi}{2}]$, such that $\sin y=z,\sin y^{\prime}=z^{\prime}$. W.l.o.g., we assume $y^{\prime}<y$, $z^{\prime}<z$. Combing (1) and (2), we have $z-z^{\prime}=\sin y-\sin y^{\prime}\geq 2\sin^{2}\frac{y-y^{\prime}}{2}\geq\frac{2(y-y^{\prime})^{2}}{\pi^{2}}$. Thus $z-z^{\prime}\leq\frac{2\epsilon^{2}}{\pi^{2}}\Longrightarrow\arcsin z-\arcsin z^{\prime}=y-y^{\prime}\leq\epsilon.$ ∎

Recall the definition in Equation (7) and (8), we have

For $\sigma\left(z\right)=\max(z,0)$, we have

Let $\bm{D}=\bm{D}(\bm{x},\bm{x}^{\prime})=\bm{D}^{(h)}(\bm{x})\bm{D}^{(h)}(\bm{x}^{\prime})$ is a 0-1 diagonal matrix. We define the following events:

$\mathcal{A}^{h}\left(\bm{x},\bm{x}^{\prime},\epsilon_{1}\right):=\left\{\left|\bm{g}^{(h)}(x^{(0)})^{\top}\bm{g}^{(h)}(\bm{x})-\bm{\Sigma}^{(h)}(\bm{x}^{(0)},\bm{x})\right|\leq\epsilon_{1}\right\}$, $\forall 0\leq h\leq L$

$\overline{\mathcal{A}}^{h}\left(\bm{x},\bm{x}^{\prime},\epsilon_{1}\right)=\mathcal{A}^{h}\left(\bm{x},\bm{x},\epsilon_{1}\right)\cap\mathcal{A}^{h}\left(\bm{x},\bm{x}^{\prime},\epsilon_{1}\right)\cap\mathcal{A}^{h}\left(\bm{x}^{\prime},\bm{x}^{\prime},\epsilon_{1}\right)$;

$\overline{\mathcal{A}}\left(\bm{x},\bm{x}^{\prime},\epsilon_{1}\right)=\bigcup_{h=0}^{L}\overline{\mathcal{A}}^{h}\left(\epsilon_{1}\right)$.

$\mathcal{B}^{h}\left(\bm{x},\bm{x}^{\prime},\epsilon_{2}\right)=\left\{\left|\left\langle\mathbf{b}^{(h)}(\bm{x}),\mathbf{b}^{(h)}(\bm{x}^{\prime})\right\rangle-\prod_{h=h}^{L}\dot{\Sigma}^{(h)}(\bm{x},\bm{x}^{\prime})\right|<\epsilon_{2}\right\}$;

$\overline{\mathcal{B}}^{h}\left(\bm{x},\bm{x}^{\prime},\epsilon_{2}\right)=\mathcal{B}^{h}\left(\bm{x},\bm{x},\epsilon_{2}\right)\cap\mathcal{B}^{h}\left(\bm{x},\bm{x}^{\prime},\epsilon_{2}\right)\cap\mathcal{B}^{h}\left(\bm{x}^{\prime},\bm{x}^{\prime},\epsilon_{2}\right)$;

$\overline{\mathcal{B}}\left(\bm{x},\bm{x}^{\prime},\epsilon_{2}\right)=\bigcup_{h=1}^{L+1}\overline{\mathcal{B}}^{h}\left(\epsilon_{2}\right)$;

$\overline{\mathcal{C}}\left(\bm{x},\bm{x}^{\prime},\epsilon_{3}\right)=\left\{|f(\bm{\theta},\bm{x})|\leq\epsilon_{3},|f(\bm{\theta},\bm{x}^{\prime})|\leq\epsilon_{3}\right\}$;

$\overline{\mathcal{D}}^{h}\left(\bm{x},\bm{x}^{\prime},\epsilon_{4}\right)=\mathcal{D}^{h}\left(\bm{x},\bm{x},\epsilon_{4}\right)\cap\mathcal{D}^{h}\left(\bm{x},\bm{x}^{\prime},\epsilon_{4}\right)\cap\mathcal{D}^{h}\left(\bm{x}^{\prime},\bm{x}^{\prime},\epsilon_{1}\right)$;

$\overline{\mathcal{D}}\left(\bm{x},\bm{x}^{\prime},\epsilon_{4}\right)=\bigcup_{h=1}^{L+1}\overline{\mathcal{D}}^{h}\left(\epsilon_{4}\right)$.

For simplicity, we will omit $\bm{x},\bm{x}^{\prime}$ when there’s no ambiguity. For events $\mathcal{A},\mathcal{B}$, we define the event $\mathcal{A}\Rightarrow\mathcal{B}$ as $\neg\mathcal{A}\wedge\mathcal{B}$.

For matrix $\bm{A}$, define the projection matrix for the column space of $\bm{A}$, $\bm{\Pi}_{\bm{A}}:=\bm{A}\bm{A}^{\dagger}$ and the orthogonal projection matrix $\bm{\Pi}_{\bm{A}}^{\perp}=I-\bm{A}\bm{A}^{\dagger}$. For two random variables $X$ and $Y$, $X\operatorname{\overset{\textrm{d}}{=\joinrel=}}_{\mathcal{A}}Y$ means $X$ is equal to $Y$ in distribution conditioned on the $\sigma$-algebra generated by $\mathcal{A}$.

where $\widetilde{\bm{w}}\sim\mathcal{N}(0,\bm{I}_{d})$ is a fresh i.i.d. copy of $\bm{w}$.

This lemma is straightforward when $\Pi^{\perp}_{\bm{G}}$ is a diagonal matrix.

where we used the fact that if $\bm{w}\sim\mathcal{N}(0,\bm{I}_{d})$, then for any orthogonal $\bm{U}$, $\bm{U}\bm{w}\sim\mathcal{N}(0,\bm{I}_{d})$ twice. ∎

E.2 Proof Sketch

Recall that $\Theta^{(L)}(\bm{x},\bm{x}^{\prime})=\sum_{h=1}^{L+1}\left(\Sigma^{(h-1)}(\bm{x},\bm{x}^{\prime})\cdot\prod_{h^{\prime}=h}^{L+1}\dot{\Sigma}^{(h^{\prime})}(\bm{x},\bm{x}^{\prime})\right)$, thus it suffices to show that if $\min_{h\in[L]}d_{h}\geq\Omega(\frac{L^{6}}{\epsilon^{4}}\log({L}/{\delta}))$, then w.p. $1-\delta$, for every $0\leq h\leq L$, it holds that

which is a direct consequence of Theorem E.2 ∎

In other words, if $\min_{h\in[L]}d_{h}\geq c_{1}\frac{L^{2}\log(\frac{L}{\delta_{1}})}{\epsilon_{1}^{2}}$, $\epsilon_{1}\leq\min(c_{2},\frac{1}{L})$, then for fixed $\bm{x},\bm{x}^{\prime}$,

In other words, if $\min_{h\in[L]}d_{h}\geq c_{1}\frac{L^{2}\log(\frac{L}{\delta_{1}})}{\epsilon^{4}_{1}}$, $\epsilon_{1}\leq\frac{c_{2}}{L}$, then for fixed $\bm{x},\bm{x}^{\prime}$,

Note that for $c_{\sigma}=2$ for $\sigma\left(z\right)=\max(0,z)$, by definition of $\mathbf{b}^{(h)}$, we have

The proof is based on a careful control of the following events.

We will use induction on Lemma E.7 to prove Theorem E.2. In the statement of Theorem E.1, we set $\delta_{1}=\frac{\delta}{4}$, $\epsilon_{1}=\frac{\epsilon^{2}}{8}$, for some $c_{1},c_{2}$, we have

In the statement of Lemma E.6, we set $\delta_{4}=\frac{\delta_{2}}{4}$, and $\epsilon_{1}=\frac{\epsilon}{2}$. Note that for $c_{1}$ large enough $\sqrt{\frac{2\log\frac{24L}{\delta}}{\min_{h}d_{h}}}\leq\frac{\epsilon}{2}$ and thus we have

In the statement of Lemma E.4, we set $\delta_{3}=\frac{\delta}{4}$, and $\epsilon_{1}=\frac{\epsilon^{2}}{8}$, we have

Using union bound on Equation (21),(22),(23), we have

For $1\leq h\leq L$ in the statement of Lemma E.7, we set $\epsilon_{2}=3(L+1-h)\epsilon$, $\epsilon_{3}=3\sqrt{\log\frac{16}{\delta}}$, $\epsilon_{2}=\epsilon$, $\delta_{4}=\frac{\delta}{4L}$. Note that for $c_{1}$ large enough, $C\sqrt{\frac{\log\frac{1}{\delta_{2}}}{d_{h}}}+C^{\prime}\sqrt{\frac{\log\frac{L}{\delta_{2}}}{d_{h}}}<\epsilon$. Thus we have

Using union bound again on Equation (24) and Equation (25) for every $h$ in $\{1,2,\ldots,L\}$, we have

E.3 Proof of Lemma E.4

For fixed $\bm{g}^{(L)}(\bm{x})$, $f(\bm{\theta},\bm{x})=\bm{W}^{(L+1)}\bm{g}^{(L)}(\bm{x})\operatorname{\overset{\textrm{d}}{=\joinrel=}}N(0,\left\|\bm{g}^{(L)}(\bm{x})\right\|^{2}$. Thus by subgaussian concentration[cite], we know w.p. $\geq 1-\delta$ over the randomness of $\bm{W}^{(L+1)}$, $\left|f(\bm{\theta},\bm{x})\right|\leq\sqrt{2\log\frac{2}{\delta}}\left\|\bm{g}^{(L)}(\bm{x})\right\|$.

For $\epsilon_{1}\leq 1$, we have $\epsilon_{1}^{2}/2<1$, which implies $\left\|\bm{g}^{(L)}(\bm{x})\right\|^{2}\leq 1+\frac{\epsilon_{1}^{2}}{2}\leq 2$, and thus taking union bound over $\bm{x},\bm{x}^{\prime}$, we have w.p. $\geq 1-\delta$, $|f(\bm{\theta},\bm{x})|\leq 2\sqrt{\log\frac{2}{\delta}}$,$|f(\bm{\theta},\bm{x}^{\prime})|\leq 2\sqrt{\log\frac{2}{\delta}}$.

E.4 Proof of Lemma E.5

Define $\bm{G}^{(h)}(\bm{x},\bm{x}^{\prime})=\begin{bmatrix}\bm{g}^{(h)}(\bm{x})^{\top}\bm{g}^{(h)}(\bm{x})&\bm{g}^{(h)}(\bm{x})^{\top}\bm{g}^{(h)}(\bm{x}^{\prime})\\ \bm{g}^{(h)}(\bm{x}^{\prime})^{\top}\bm{g}^{(h)}(\bm{x})&\bm{g}^{(h)}(\bm{x}^{\prime})^{\top}\bm{g}^{(h)}(\bm{x}^{\prime})\end{bmatrix}$, we have for every $1\leq h\leq L$,

For simplicity, we denote $\bm{G}^{(h)}(\bm{x},\bm{x}^{\prime})$, $\bm{\Lambda}^{(h)}(\bm{x},\bm{x}^{\prime})$ by $\bm{G}$,$\bm{\Lambda}$ respectively.

Since $\dot{\sigma}(z)=\bm{1}[z\geq 0]$ is 0-homogeneous, we have

It is easy to verify that $|\sqrt{G_{11}G_{22}}-1|\leq\epsilon^{2}/2$, and thus

Note that $\dot{\Sigma}^{(h)}(\bm{x},\bm{x}^{\prime})=t_{\sigma^{\prime}}\left(\left.\bm{\Sigma}^{(h)}\right|_{\bm{\bm{x}},\bm{\bm{x}^{\prime}}}\right)=\hat{t}_{\sigma^{\prime}}\left(\bm{\Lambda}^{(h)}(\bm{x},\bm{x}^{\prime})\right)$.

Combining Lemma E.8 and Lemma E.9, we have for any $(\bm{x},\bm{x}^{\prime})$,

Taking union bound over $(\bm{x},\bm{x}),(\bm{x},\bm{x}^{\prime}),(\bm{x}^{\prime},\bm{x}^{\prime})$ for the choice of $(\bm{x},\bm{x}^{\prime})$, we have

E.5 Proof of Lemma E.7

The proof of Lemma E.7 is based on the following 3 claims, Claim E.1, E.2 and E.3.

If $\overline{\mathcal{A}}^{L}\left(\epsilon_{1}^{2}/2\right)\bigwedge\overline{\mathcal{B}}^{h+1}\left(\epsilon_{2}\right)\bigwedge\overline{\mathcal{C}}\left(\epsilon_{3}\right)\bigwedge\overline{\mathcal{D}}^{h}\left(\epsilon_{4}\right)$, then we have

For any fixed $h$, let $\bm{G}=[\bm{g}^{(h)}(\bm{x})\ \bm{g}^{(h)}(\bm{x}^{\prime})]$,

w.p. $\geq 1-\frac{\delta_{2}}{2}$, if $\overline{\mathcal{A}}^{L}\left(\epsilon_{1}^{2}/2\right)\bigwedge\overline{\mathcal{B}}^{h+1}\left(\epsilon_{2}\right)\bigwedge\overline{\mathcal{C}}\left(\epsilon_{3}\right)\bigwedge\overline{\mathcal{D}}^{h}\left(\epsilon_{4}\right)$, then we have for any $(\bm{x}^{(1)},\bm{x}^{(2)})\in\{(\bm{x},\bm{x}),(\bm{x},\bm{x}^{\prime}),(\bm{x}^{\prime},\bm{x}^{\prime})\}$,

As a by-product, for any $\bm{x}^{(1)}\in\{\bm{x},\bm{x}^{\prime}\}$, we have

It suffices to prove this claim conditioned on every possible realization of

Recall that $\bm{G}=\left[\bm{g}^{(h)}(\bm{x}^{(1)})\ \bm{g}^{(h)}(\bm{x}^{(2)})\right]$, we further define $\bm{F}=\left[\bm{f}^{(h)}(\bm{x}^{(1)})\ \bm{f}^{(h)}(\bm{x}^{(2)})\right]$. Applying Lemma E.3 on each row of $\bm{W}^{h+1}$, we have

where $\widetilde{\bm{W}}$ is an iid copy of $\bm{W}^{(h+1)}$.

where $\bm{\xi}\sim N(\bm{0},\bm{I}_{2d_{h}})$.

Thus conditioned on $\{\mathbf{b}^{(h+1)}(\bm{x}^{(1)}),\mathbf{b}^{(h+1)}(\bm{x}^{(2)}),\bm{g}^{(h)}(\bm{x}^{(1)}),\bm{g}^{(h)}(\bm{x}^{(2)})\}$, we have

Now we are ready to prove Claim E.1 by applying Lemma E.10. Let $\bm{A}=\frac{1}{2}\bm{M}^{\top}\begin{bmatrix}\bm{0}&\Pi^{\perp}_{\bm{G}}\bm{D}\Pi^{\perp}_{\bm{G}}\\ \Pi^{\perp}_{\bm{G}}\bm{D}\Pi^{\perp}_{\bm{G}}&\bm{0}\end{bmatrix}\bm{M}$, we have

Thus by Lemma E.10 with $t=\log\frac{6}{\delta_{2}}$ we have w.p. $1-\frac{\delta_{2}}{6}$,

The main part of the claim is completed by taking union bound over $(\bm{x},\bm{x}),(\bm{x},\bm{x}^{\prime}),(\bm{x}^{\prime},\bm{x}^{\prime})$. For the by-product, let $\bm{x}^{(2)}=\bm{x}^{(1)}$, and we have

w.p. $\geq 1-\frac{\delta_{2}}{2}$, if $\overline{\mathcal{A}}^{L}\left(\epsilon_{1}^{2}/2\right)\bigwedge\overline{\mathcal{B}}^{h+1}\left(\epsilon_{2}\right)\bigwedge\overline{\mathcal{C}}\left(\epsilon_{3}\right)\bigwedge\overline{\mathcal{D}}^{h}\left(\epsilon_{4}\right)$, then

It suffices to prove the claim for $\bm{x}$. We will denote $\bm{x}$ by $\bm{x}$, $\bm{g}^{(h)}(\bm{x})$ by $\bm{g}^{(h)}$ and $\mathbf{b}^{(h+1)}(\bm{x})$ by $\mathbf{b}^{(h+1)}$. We also define $\Pi_{\bm{g}}$ as $\bm{g}\bm{g}^{\top}$, and $\Pi_{\bm{G}/\bm{g}}=\Pi_{\bm{G}}-\Pi_{\bm{g}}$. Clearly, $\Pi_{\bm{G}/\bm{g}}$ is still a projection matrix of rank or $1$.

Since $\left\|\bm{\Pi}_{\bm{G}}\left(\bm{W}^{(h+1)}\right)^{\top}\mathbf{b}^{(h+1)}(\bm{x})\right\|\leq\left\|\bm{\Pi}_{\bm{g}}\left(\bm{W}^{(h+1)}\right)^{\top}\mathbf{b}^{(h+1)}\right\|+\left\|\bm{\Pi}_{\bm{G}/\bm{g}}\left(\bm{W}^{(h+1)}\right)^{\top}\mathbf{b}^{(h+1)}\right\|$, it suffices to bound these two terms separately.

Recall $\mathbf{b}^{(h+1)}$ is defined as the gradient of $f(\bm{\theta},\bm{x})$ with respect to the pre-activation of layer $h+1$, $\bm{f}^{h+1}$, thus if we view $g$ as a function $\bm{g}^{(h)},\bm{W}^{(h+1)},\ldots,\bm{W}^{(L+1)}$, by the rule of back propagation, we have

Note that $\bm{g}^{(h)}(x^{(0)})^{\top}\bm{g}^{(h)}(\bm{x})\geq 1-\epsilon_{1}^{2}/2\geq\frac{1}{2}$, we have

where $t\sim N(0,\left\|\mathbf{b}^{(h+1)}\right\|)$. Hence w.p. $\geq 1-\delta_{2}/4$ over the randomness of $\bm{W}^{(L)}$, $\left\|\bm{\Pi}_{\bm{G}/\bm{g}}\left(\bm{W}^{(h+1)}\right)^{\top}\bm{b}^{(h+1)}\right\|\leq\sqrt{2\log\frac{8}{\delta_{2}}}\left\|\mathbf{b}^{(h+1)}\right\|\leq\sqrt{2\log\frac{8}{\delta_{2}}}\leq 2\sqrt{\log\frac{8}{\delta_{2}}}\ (\epsilon_{2}<1)$.

Thus by assumption $\log\frac{8}{\delta_{2}}\leq d_{h}$, we have $2\sqrt{\log\frac{8}{\delta_{2}}}+\sqrt{2}\epsilon_{3}\leq 2\sqrt{d_{h}}+\sqrt{2}\leq 3\sqrt{2d_{h}}$. ∎

Wrapping things up, by combining Claim E.2 and Claim E.3, we have w.p. $\geq 1-\delta_{2}$, for any pair of $(\bm{x}^{(1)},\bm{x}^{(2)})\in\{(\bm{x},\bm{x}),(\bm{x},\bm{x}^{\prime}),(\bm{x}^{\prime},\bm{x}^{\prime}))\}$,

Using Equation (28) together with Claim E.1 and Claim E.2, we’ve finished the proof for Lemma E.7.

Appendix F Proof of Theorem 3.2

i.e., the kernel induced from the gradient of the prediction with respect to the parameters of the neural network at time $t$.

We prove a lemma to reduce the prediction perturbation bound to the kernel perturbation bound.

Fix $\epsilon_{\bm{H}}\leq\frac{1}{2}\lambda_{0}$. Suppose $\left|f_{nn}(\bm{\theta}(0),\bm{x}_{i})\right|\leq\epsilon_{init}$ for $i=1,\ldots,n$ and $\left|f_{nn}\left(\bm{\theta}(0),\bm{x}_{te}\right)\right|\leq\epsilon_{init}$ and $\left\|\bm{u}_{nn}(0)-\bm{y}\right\|_{2}=O\left(\sqrt{n}\right)$. Furthermore, if for all $t\geq 0$ $\left\|\ker_{ntk}(\bm{x}_{te},\bm{X})-\ker_{t}(\bm{x}_{te},\bm{X})\right\|_{2}\leq\epsilon_{test}$ and $\left\|\bm{H}^{*}-\bm{H}(t)\right\|_{2}\leq\epsilon_{\bm{H}}$, then we have

Our proof relies a careful analysis on the trajectories induced by gradient flows for optimizing the neural network and the NTK predictor.

The solution to this program can written as applying gradient flow on

with initialization $\bm{\bm{\beta}}(0)=\bm{0}$. We use $\bm{\bm{\beta}}(t)$ to denote this parameter at time $t$ trained by gradient flow and $f_{ntk}\left(\bm{x}_{te},\bm{\bm{\beta}}(t)\right)$ be the predictor for $\bm{x}_{te}$ at time $t$. With these notations, we rewrite

where we have used the fact that the initial prediction is . Now we take a closer look at the time derivative:

We thus we analyze the difference between the NN predictor and NTK predictor via this integral form

For the second term, recall $\left\|\bm{H}^{*}-\bm{H}(t)\right\|_{2}\leq\epsilon_{\bm{H}}$ by our assumption so $\lambda_{\min}\left(\bm{H}(t)\right)\geq\frac{1}{2}\lambda_{0}$. Using this fact we know $\left\|\bm{u}_{nn}(t)-\bm{y}\right\|_{2}\leq\exp(-\frac{\kappa^{2}}{2}\lambda_{0}t)\left\|\bm{u}_{nn}(0)-\bm{y}\right\|_{2}$. Therefore, we can bound

To bound $\int_{t=0}^{\infty}\left\|\bm{u}_{nn}(t)-\bm{u}_{ntk}(t)\right\|_{2}$, we observe that $\bm{u}_{nn}(t)\rightarrow\bm{y}$ and $\bm{u}_{ntk}(t)\rightarrow\bm{y}$ with linear convergence rate. Therefore, we can choose some $t_{0}=\frac{C}{\lambda_{0}\kappa^{2}}\log\left(\frac{n}{\epsilon_{\bm{H}}\lambda_{0}\kappa}\right)$ so that

By Lemma F.1, the problem now reduces to (i) choose $\kappa$ small enough to make $\epsilon_{init}=O(\epsilon)$ and (ii) show when the width is large enough then $\epsilon_{\bm{H}}$ and $\epsilon_{test}$ are both $O(\epsilon)$. For (i), based on Theorem E.1 and the union bound, we can just choose $\kappa=O\left(\frac{\epsilon}{\log(n/\delta)}\right)$ to make $\epsilon_{init}=O(\epsilon)$ with probability $1-\delta$. For (ii), we will use Theorem 3.1 and Lemma F.2 below, and then apply the union bound. ∎

In this subsection we prove the following lemma.

If $\left\|\frac{\partial f(\bm{\theta}(t),\bm{x})}{\partial\bm{\theta}}-\frac{\partial f(\bm{\theta}(0),\bm{x})}{\partial\bm{\theta}}\right\|\leq\epsilon$ and $\left\|\frac{\partial f(\bm{\theta}(t),\bm{x}^{\prime})}{\partial\bm{\theta}}-\frac{\partial f(\bm{\theta}(0),\bm{x}^{\prime})}{\partial\bm{\theta}}\right\|\leq\epsilon$, we have

By the proof of Theorem 3.1, we know $\left\|\frac{\partial f(\bm{\theta}(0),\bm{x})}{\partial\bm{\theta}}\right\|_{2}=O\left(1\right)$. Then we can just use triangle inequality. ∎

For $h=1,\ldots,L$, $i=1,\ldots,m$, we denote

Note $\bm{z}^{(h)}=\sqrt{\frac{2}{m}}\bm{f}^{(h)}$. Here we use $\bm{z}^{(h)}$ instead of $\bm{f}^{(h)}$ for the ease of presentation.

Recall the gradient to $\bm{W}^{(h)}$ is:

Then with probability at least $1-\delta$ over random initialization, if $\left\|\triangle\bm{W}^{(h)}\right\|_{2}\leq\sqrt{m}\omega$ for all $h=1,\ldots,L$, we have $\left\|\triangle\bm{g}^{(h)}\right\|_{2}=O(\omega L^{5/2}\sqrt{\log m})$ for all $h=1,\ldots,L$.

While Allen-Zhu et al. [2018b] did not consider the perturbation on $\bm{W}^{(1)}$, by scrutinizing their proof, it is easy to see that the perturbation bounds still hold even if there is a small perturbation on $\bm{W}^{(1)}$.

The next lemma bounds the backward vector, adapted from

While Allen-Zhu et al. [2018b] did not consider the perturbation on $\bm{W}^{(L+1)}$, by scrutinizing their proof, it is easy to see that the perturbation bounds still hold even if there is a small perturbation on $\bm{W}^{(L+1)}$.

Combing these two lemmas and the result for the initialization (Theorem 3.1), we have the following “gradient-Lipschitz" lemma.

We use the triangle inequality to bound the perturbation

The following lemma shows for given weight matrix, if we have linear convergence and other weight matrices are only perturbed by a little, then the given matrix is only perturbed by a little as well.

We let $C,C_{0},C_{1},C_{2},C_{3}>0$ be some absolute constants.

The last step we used $\left\|\frac{\partial f_{nn}(\bm{\theta}(0),\bm{x}_{i})}{\partial\bm{W}^{(h)}}\right\|_{F}=O(1)$. Suppose there exists $t$ such that $\left\|\bm{W}^{(h)}(t)-\bm{W}^{(h)}(0)\right\|_{F}>\omega\sqrt{m}.$ Denote

For any $t<t_{0}$, we know for all $h^{\prime}\in[L+1]$, $\left\|\bm{W}^{(h^{\prime})}(t)-\bm{W}^{(h^{\prime})}(0)\right\|_{2}\leq\omega\sqrt{m}.$ Therefore, by Lemma F.6, we know

Therefore, using the fact that $\omega$ is sufficiently small we can bound

Since we also know $m$ is sufficiently large to make $\omega\sqrt{m}>\frac{C_{3}\sqrt{n}}{\lambda_{0}}$, we have a contradiction. ∎

The next lemma shows if all weight matrices only have small perturbation, then we still have linear convergence.

Under this assumption and the result of initialization, we know for all $t\geq 0$, $\lambda_{\min}\left(\bm{H}(t)\right)\geq\frac{1}{2}\lambda_{0}$. This in turn directly imply the linear convergence result we want. ∎

Lastly, with these lemmas at hand, using an argument similar to [Du et al., 2019], we can show during training, weight matrices do not move by much.

We analyze case by case. Suppose at time $t_{0}$, $\left\|\bm{W}^{(h)}(t_{0})-\bm{W}^{(h)}(0)\right\|_{F}>\omega\sqrt{m}$. By Lemma F.7, we know there exists some $0\leq t_{1}<t_{0}$ such that either there exists $h^{\prime}\neq h$ such that

However, this violates the minimality of $t_{0}$. For the other case, if

By Lemma F.8, we know there exists $t_{1}<t_{0}$ such that there exists $h\in[L+1]$,

However, again this violates the minimality of $t_{0}$. ∎

Now we can finish the proof of Lemma F.2.

By Lemma F.9, we know for $t\rightarrow\infty$, $\left\|\bm{W}^{(h)}(t)-\bm{W}^{(h)}(0)\right\|_{F}\leq O\left(\omega\sqrt{m}\right)$ for if $\omega$ is sufficiently. Applying Lemma F.6, we know we only have a small perturbation on the gradient. Applying Lemma F.3, we know we only have small perturbation on kernel values. ∎

Appendix G CNTK Derivation

i.e., this operator extracts the $(i,j)$-th patch. By this definition, we can rewrite the CNN definition:

For $h=1,\ldots,H$, $\beta=1,\ldots,{C}^{(h)}$, the intermediate outputs are defined as

We expand $\Theta^{(L)}(\bm{x},\bm{x}^{\prime})$ to show we can write it as the sum of $(L+1)$ terms with each term representing the inner product between the gradients with respect to the weight matrix of one layer. We first define an linear operator

With this operator, we can expand CNTK as (for simplicity we drop on $\bm{x}$ and $\bm{x}^{\prime}$)

which could be rewritten as the following by Property 31,

G.2 Derivation

We first compute the derivative of the prediction with respect to one single filter.

With this expression, we proceed to we compute the inner product between gradients with respect to the $h$-th layer matrix

Similar to our derivation to NTK, we can use the following approximation

Thus it remains to show that $\forall(i,j,i^{\prime},j^{\prime})\in[{P}]\times[{Q}]\times[{P}]\times[{Q}]$,

Now, we tackle the term $\left(\frac{\partial f(\bm{\theta},\bm{x})}{\partial[\bm{x}_{(\beta)}^{(h)}]_{ij}}\cdot\frac{\partial f(\bm{\theta},\bm{x}^{\prime})}{\partial[\bm{x}^{\prime(h)}_{(\beta)}]_{i^{\prime}j^{\prime}}}\right)$. Notice that

where the first approximation is due to our initialization of $\bm{W}^{(h+1)}$. In other words, we’ve shown

Since we use a fully-connected weight matrix as the last layer, we have $\bm{G}^{(L)}(\bm{x},\bm{x^{\prime}},\bm{\theta})\approx\bm{I}$.

Thus by induction with Equation 36, we have derived Equation 34, which completes the derivation of CNTK.

Appendix H Formula of CNTK with Global Average Pooling

In this section we define CNN with global average pooling considered in this paper and its corresponding CNTK formula.

For $h=1,\ldots,L$, $\beta=1,\ldots,{C}^{(h)}$, the intermediate outputs are defined as

Besides using global average pooling, another modification is that we do not train the first and the layer. This is inspired by Du et al. [2018a] in which authors showed that if one applies gradient flow, then at any training time $t$, the difference between the squared Frobenius norm of the weight matrix at time $t$ and that at initialization is same for all layers. However, note that $\bm{W}^{(1)}$ and $\bm{W}^{(L+1)}$ are special because they are smaller matrices compared with other intermediate weight matrices, so relatively, these two weight matrices change more than the intermediate matrices during the training process, and this may dramatically change the kernel. Therefore, we choose to fix $\bm{W}^{(1)}$ and $\bm{W}^{(L+1)}$ to the make over-parameterization theory closer to practice.

CNTK formula.

We let $\bm{x},\bm{x}^{\prime}$ be two input images. Note because CNN with global average pooling and vanilla CNN shares the same architecture except the last layer, $\bm{\Sigma}^{(h)}(\bm{x},\bm{x}^{\prime})$, $\dot{\bm{\Sigma}}^{(h)}(\bm{x},\bm{x}^{\prime})$ and $\bm{K}^{(h)}(\bm{x},\bm{x}^{\prime})$ are the same for these two architectures. the only difference is in calculating the final kernel value. To compute the final kernel value, we use the following procedure.

First, we define $\bm{\Theta}^{(0)}(\bm{x},\bm{x}^{\prime})=\bm{0}$. Note this is different from CNTK for vanilla CNN which uses $\bm{\Sigma}^{(0)}$ as the initial value because we do not train the first layer.

For $h=1,\ldots,L-1$ and $(i,j,i^{\prime},j^{\prime})\in[{P}]\times[{Q}]\times[{P}]\times[{Q}]$, we define

For $h=L$, we define $\bm{\Theta}^{(L)}(\bm{x},\bm{x}^{\prime})=\dot{\bm{K}}^{(L)}(\bm{x},\bm{x}^{\prime})\odot\bm{\Theta}^{(L-1)}(\bm{x},\bm{x}^{\prime})$.

Lastly, the final kernel value is defined as

Note that we ignore $\bm{K}^{(L)}$ comparing with the CNTK of CNN. This is because we do not train the last layer. The other difference is we calculate the mean over all entries, instead of calculating the summation over the diagonal ones. This is because we use global average pooling so the cross-variances between every two patches will contribute to the kernel.

Appendix I Fast Computation for ReLU-Activated CNTK

In this section we present our approach to compute CNTK exactly. Notably, most computation required by our new approach can be described as entry-wise operations over matrices and tensors, which allows efficient implementations on GPUs.

Following the formulas in Sections 4 and H, the trickiest part is computing the expectation of the post-activation output, i.e., Equations (11) and (12). These two expectations depend on (the same) $2\times 2$ matrices $\left[\bm{\Lambda}^{(h)}(\bm{x},\bm{x}^{\prime})\right]_{ij,i^{\prime}j^{\prime}}$. To obtain faster implementations, our key observation is that if the diagonal entries of $\left[\bm{\Lambda}^{(h)}(\bm{x},\bm{x}^{\prime})\right]_{ij,i^{\prime}j^{\prime}}$ are all ones and the activation function is ReLU, there are closed-form formulas for the the corresponding expectations. To see this, let us suppose for now that $\bm{\Lambda}=\begin{pmatrix}1&\lambda\\ \lambda&1\end{pmatrix}$ for some $\left|\lambda\right|\leq 1$. When the activation function $\sigma\left(\cdot\right)$ is ReLU, one can show that

Here, we interpret $\bm{\Sigma}^{(h-1)}(\bm{x},\bm{x})$, $\bm{\Sigma}^{(h-1)}(\bm{x},\bm{x}^{\prime})$, $\bm{\Sigma}^{(h-1)}\left(\bm{x}^{\prime},\bm{x}\right)$ and $\bm{\Sigma}^{(h-1)}\left(\bm{x}^{\prime},\bm{x}^{\prime}\right)$ as matrices of size ${P}{Q}\times{P}{Q}$. If the diagonal entries of $A^{(h)}$ are all ones, then the diagonal entries of $\left[\bm{\Lambda}^{(h)}(\bm{x},\bm{x}^{\prime})\right]_{ij,i^{\prime}j^{\prime}}$ are all ones for all possible $(i,j,i^{\prime},j^{\prime})\in[{P}]\times[{Q}]\times[{P}]\times[{Q}]$, in which case we can calculate $\bm{K}^{(h)}(\bm{x},\bm{x}^{\prime})$ and $\dot{\bm{K}}^{(h)}(\bm{x},\bm{x}^{\prime})$ by simply applying the closed-form formulas described in (37) and (38) on $A^{(h)}$.

However, in general, the diagonal entries of $\bm{A}^{(h)}$ are not always all ones, in which case we resort to the homogeneity of the ReLU activation function. Suppose $\bm{\Lambda}=\begin{pmatrix}1&\lambda\\ \lambda&1\end{pmatrix}$ for some $\left|\lambda\right|\leq 1$, and $\bm{D}=\begin{pmatrix}c_{1}&0\\ 0&c_{2}\end{pmatrix}$ for some $c_{1},c_{2}\geq 0$, then one can show that

Inspired by this, our final approach is described as follows.

Let $\bm{D}=\begin{pmatrix}\bm{D_{x}}&\bm{0}\\ \bm{0}&\bm{D_{x^{\prime}}}\end{pmatrix}$, where $\bm{D_{x}}$ and $\bm{D_{x^{\prime}}}$ are diagonal matrices whose diagonal entries are square roots of the diagonal entries of $\bm{\Sigma}^{(h-1)}(\bm{x},\bm{x})$ and $\bm{\Sigma}^{(h-1)}(\bm{x^{\prime}},\bm{x^{\prime}})$, respectively.

Applying Equations (39) and (40) on $\bm{A}^{(h)}=\bm{D}\bm{\Lambda}^{(h)}\bm{D}$, where the diagonal entries of $\bm{\Lambda}^{(h)}$ are all ones.

Notice that the implementation above requires us to store the whole $\bm{A}^{(h)}$ matrix, which has size $2{P}{Q}\times 2{P}{Q}$. To further optimize the efficiency, we notice that to implement the approach described above, we only need to store the diagonal entries of $\bm{\Sigma}^{(h-1)}(\bm{x},\bm{x})$ and $\bm{\Sigma}^{(h-1)}(\bm{x^{\prime}},\bm{x^{\prime}})$, together with the matrix $\bm{\Sigma}^{(h-1)}(\bm{x},\bm{x}^{\prime})$, which has size ${P}{Q}\times{P}{Q}$.