A Fast, Well-Founded Approximation to the Empirical Neural Tangent Kernel
Authors: Mohamad Amin Mohamadi
Danica J. Sutherland
Venue: Submitted to ICLR 2023 (2023)
Keywords: neural tangent kernels
deep learning theory

Abstract: Empirical neural tangent kernels (eNTKs) can provide a good understanding of a given network's representation: they are often far less expensive to compute and applicable more broadly than infinite-width NTKs. For networks with $O$ output units (e.g. an $O$-class classifier), however, the eNTK on $N$ inputs is of size $NO \times NO$, taking $\mathcal{O}\big( (N O)^2\big)$ memory and up to $\mathcal{O}\big( (N O)^3 \big)$ computation. Most existing applications have therefore used one of a handful of approximations yielding $N \times N$ kernel matrices, saving orders of magnitude of computation, but with limited to no justification. We prove that one such approximation, which we call ``sum of logits,'' converges to the true eNTK at initialization. Our experiments demonstrate the quality of this approximation for various uses across a range of settings.

Under review as a conference paper at ICLR 2023
A FAST, WELL-FOUNDER APPROXIMATION TO THE EMPIRICAL NEURAL TANGENT KERNEL
Anonymous authors Paper under double-blind review
ABSTRACT
Empirical neural tangent kernels (eNTKs) can provide good understanding of a given network's 's representation: they are often far less expensive to compute and applicable more broadly than infinite-width NTKs. For networks with O output units (e.g. an O-class classifier), however, the eNTK on N inputs is of size NO x NO, taking O ((NO)?) memory and up to O ((NO)s computation. Most existing applications have therefore used one of handful of approximations yielding N x N kernel matrices, saving orders of magnitude of computation, but with limited to no justification. We prove that one such approximation, which we call "sum of logits," converges to the true eNTK at initialization. Our experiments demonstrate the quality of this approximation for various uses across a range of settings.
INTRODUCTION
The pursuit of theoretical foundation for deep learning has lead researches to uncover interesting connections between neural networks (NNs) and kernel methods. has long been known that randomly initialized NNs in the infinite width limit Gaussian processes with what is termed the Neural Network Gaussian Process (NNGP) kernel, and training the last layer with gradient flow under squared loss corresponds to the posterior mean (Neal, 1996; Williams, 1996; Hazan & Jaakkola, 2015; Lee et al., 2017; Matthews et al., 2018; Novak et al., 2018; Yang, 2019). More recently. Jacot et al. (2018) (building off line of closely related prior work) showed that the same is true if we train all the parameters of the network, but using a different kernel called the Neural Tangent Kernel (NTK). Yang (2020); Yang & Littwin (2021) later showed that this connection is architecturally universal, extending the domain from fully-connected NNs to most of the currently-used networks in practice, such as ResNets and Transformers. Lee et al. (2019) also showed that the dynamics of training wide but finite-width NNs with gradient descent can be approximated by a linear model obtained from the first-order Taylor expansion of that network around its initialization. Furthermore, they experimentally showed that this approximation approximation excellently holds even for networks that are not so wide.
In addition theoretical insights from the results themselves, NTKs have had significant impact in diverse practical settings. Arora et al. (2019b) show very strong performance of NT -based models on variety of low-data classification and regression tasks. The condition number of an NN's NTK has been shown correlation directly with the trainability and generalization capabilities of the NN (Xiao al., 2018; 2020); thus, Park et al. (2020); Chen et al. (2021) have used this to develop practical algorithms for neural architecture search. Wei et al. (2022); Bachmann et al. (2022) estimate the generalization ability of specific network, randomly initialized or pre-trained on a different dataset, with efficient cross-validation. Zhou et al. (2021) use NTK regression for efficient meta-learning, and Wang et al. (2021); Holzmüller al. (2022); Mohamadi et al. (2022) use NTKs for active learning.
There has also been significant theoretical insight gained from empirical studies of networks' NTKs. Here are a few examples: Fort et al. (2020) use NTKs to study how the loss geometry the NN evolves under gradient descent. Franceschi et al. (2021) employ NTKs to analyze the behaviour of Generative Adverserial Networks (GANs). Nguyen et al. (2020; 2021) used NTKs for dataset distillation. He al. (2020); Adlam et al. (2020) used NTKs to predict and analyze the uncertainty of a NN's predictions. Tancik et al. (2020) use NTKs to analyze the behaviour of MLPs in learning high frequency functions, leading to new insights into our understanding of neural radiance fields.

---

Under review as a conference paper at ICLR 2023
Tlay ImageNot (200)
1011
10"
1014
10
10
10-15
pNTC
PNTC
JNTC
BeaNee deytl
BeaNet deytis
18:
BeaNeet deyti
Figure 1: Comparison of wall-clock time of evaluating the eNTK and pNTK of a pair of input datapoints over various datasets and ResNet depths.
We thus believe NTKs will continue to be used in both theoretical and empirical deep learning,
Unfortunately, however, computing the NTK for practical networks is extremely challenging, and most of the time not even computationally feasible. The NTK of NN is defined as the outer product of the Jacobians of the output of the NN with respect to its parameters:
eNTK == (0 x2) = [Jo (fo [Jo (fo(xz2)
(1)
where Je (fe (x)) denotes the Jacobian of the function f at point æ with respect to the flattened
vector of all its parameters, 0 E IR. Assuming f R2 + R where D is the input dimension and O the number of outputs, we have Jo (fo(z)) E ROX and @s(21, x22 E ROx0 Thus, computing
the NTK between a set of N1 data points and set of N₂ data points yields N2 matrices each of shape O x O, which we usually reshape into an N,O x N₂O matrix.
When computing an eNTK on tasks involving large datasets and with multiple output neurons, e.g. in a classification model with O classes, the eNTK quickly becomes impractical regardless of how fast each entry is computed due to its NO NO size. For instance, the full eNTK of classification model even on the relatively mild CIFAR-IO dataset (Krizhevsky, 2009), stored in double precision, takes over 1.8 terabytes in memory. For practical usage, we need to do something better.
In this work, we present a simple trick for a strong approximation of the eNTK that removes the O² from the size of the kernel matrix, resulting in a factor of O2 improvement in the memory and up to O in computation. Since for typical classification datasets O is at least 10 (e.g. CIFAR-10) and potentially 1 (100 or more (e.g. ImageNet, Deng al., 2009), this provides multiple orders of magnitude savings over the original eNTK (I). We prove that under appropriate initialization of the NN this approximation converges to the original eNTK at a rate of Onn ((n-1/3 122) for a network of depth L and width in each layer, and the predictions of kernel regression with the approximate kernel do the same. Finally, we present diverse experimental investigations supporting our theoretical results across range of different architectures and settings. We hope this approximation further enables researches to employ NTKs towards theoretical and empirical advances in wide networks.
In the infinite-width limit of properly initialized NNs, O₂ converges almost surely at initialization to particular kerne,, and remains constant over the course of training. Algorithms are available to compute this expectation exactly, but they tend to be substantially more expensive than computing (1) directly for all but extremely wide networks. The convergence to this infinite- wiith regime also slow in practice, and moreover it eliminates some of the interest of the framework: neural architecture search, predicting generalization of a pre-trained representation, and meta-learning are all considerably less interesting when we only consider infinite-wide networks that do essentially no feature learning, Thus, in this paper, we focus only the "empirical" eNTK (1).
Infinite NTKs
2 RELATED WORK
Among the numerous recent works that have used eNTKs either to gain insights about various phc- nomenons in deep learning propose new algorithms, not many have publicized the computational costs and implementation details of computing eNTKs. Nevertheless, all are in agreement about the expense of such computations (Park et al., 2020; Holzmüller et al., 2022; Fort et al., 2020).
Several recent works have, mostly "quietly," employed various techniques to avoid dealing with the full eNTK matrix; however, to the best of our knowledge, none provide any rigorous justifications.
2

---

Under review as a conference paper at ICLR 2023
Wei et al. (2022, Section 2.3) point out that if the final layer of NN is randomly initialized, the expected eNTK can be written as Ku 8 To for some kernel Ku. where Io is the O x O identity matrix and @ is the Kronecker product. Thus, they use the approximation in which they only compute the eNTK with respect to one of the logits of the NN. Although their approach to approximating the eNTK is similar to ours, they don't 't provide any rigorous bounds or empirical study of how closely the actual eNTK is approximated by its expectation in this regard. Wang et al. (2021) employs the same "single-logit" strategy, though they only mention the unfinit-wiidt limit as motivation supporting their trick. Despite of these claims, we will see in our experiments that the eNTK is generally nor diagonal. We will, however, prove upper bounds on distance of our approximation provide experimental support that this approximation captures the behaviour of the eNTK even when the NN's weights are not at initialization. Park al. (2020); Chen et al. (2021) also seem to use a form of "single-logit" approximation to eNTK, without explicitly mentioning it. Lee et al. (2019), by contrast, do use the full eNTK, and hence never compute the kernel on more than 256 datapoints.
the eNTK, and
More recently, Novak et al. (2021) performed an in-depth analysis of computational and memory complexity required for computing the eNTK, and proposed two new approaches to reduce the time complexity of computing the eNTK in special cases (depending on the architecture of the NN) over explictly implementing (1). We remark that as our approaches are complementary, as we propose an approximation the eNTK and Novak al. (2021) proposes an algorithm for computation of the eNTK, which also applies when to computing our approximation - in fact, we use their "structured derivatives" method in our experiments. Moreover, their approach does not reduce the memory complexity of computing the eNTK, typically the largest burden in practical applications.
3 PSEUDO-NTK
We define the pseudo-NTK (pNTK) of the network fo as
PNTK- = 00(01,
Vo
(x1)
Vo
1
(
(x22)
(2)
VO
V
i=1
1XP
PX1
where 16() (a:) denotes the ith output of fe on the input æ. Unlike eNTK, which is matrix-valued kernel for each pair of inputs, pNTK is a traditional scalar-valued kernel.
Some recent work (Arora et al., 2019a; Yang, 2020; Wei et al., 2022; Wang et al., 2021) has pointed out that in the infinite width limit the NTK lim ( x22 becomes a diagonal matrix. Thus, one can avoid computing the off-diagonal entries of the infinite-width NTK of each pair through using O0 22) @x(11, .22) Io. resulting in drastic ( time and memory complexity decrease.
Practitioners have accordingly used the same approach in computing the corresponding eNTK of a finite width network, but with little no further justification. We see in our experiments that for finite width networks, the NTK is not diagonal. In fact, we show that for most practical networks, it is very far from being diagonal, casting doubts on the validity of arguments justifying the approximation with asymptotic diagonality. We justify this category of approximation with thcoretical bounds on the difference of the true NTK from the approximation (2), which we also call "sum of logits."
Before turning to the formal results and experimental evaluation, we give some intuition. First, suppose that f(()) (x: <0x) Vi, so that U; PR- is the ith row of linear read-out layer; then I Ee (x) d(z). 1 Eo uz] If the vectors U; (0, 0 are independent, then vo
vo
Vo Σの ~ NO, 0² I1 ,) has the same distribution as any individual entry, say V1. Thus, at initialization, our of logits approximation agrees in distribution with the first-logit approximation. Our proof uses the sum-of of-loiss form, however, and we believe it may be more sensible for networks that are not at random initialization (like a pretrained network with randomized read-out layer).
Calling this vector (whether the first logit or sum of logits) v, we can think of (2) as the NTK of model with single scalar output as function of o, whose last layer has weights v. When linearize a network with that kernel for an O-class classification problem, getting the formula (6) discussed in Section 4.3, we end up effectively using one-vs-rest classifier scheme. Thus, we can
3

---

Under review as a conference paper at ICLR 2023
think of the pseudo-NTK, approximating the process of training O one-vs-rest classifiers, rather than a single O-way classifier.
4 APPROXIMATION QUALITY OF PSEUDO-NTK
We will now study various aspects of the approximation of (2) to (1), both in theory and empirically, To experimentally evaluate our claims, we present various experiments that compare the targeted characteristic of pNTK vs eNTK both at initialization and throughout training. We present our experiments on four widely-used architectures, namely, fully-connected networks (FCN) of depth 3 (as in Lee al., 2019; 2020), fully-convolutional networks (ConvNet) of depth (as in Arora et al., 2019a;b; Lee et al., 2020), ResNet18 (He et al., 2016) and WideResNet- 16-k (Zagoruyko & Komodakis, 2016). We evaluate each architecture different widths, as mentioned in the plot legends: we show exact widths for FCN, while for others we show widening factor (details in Appendix A). For consistency with most other recent papers studying NTKs and properties of NNs general, we chose CIFAR-10 10 (Krizhevsky, 2009) the dataset to evaluate our experiments on. Each experiment is repeated using three seeds, and the corresponding error bars are plotted. However, in some figures the error bars are not plotted. In those cases, it is because they interfered with clear interpretation of the plots, as the error bars were too large and not informative. We used Stochastic Gradient Descent (SGD) the optimizer of choice for training our networks. Details on the configurations used in optimization of each network can be found in the Appendix.
Each network each experiment is trained for 200 epochs using mini-batch stochastic gradient descent (SGD), run on NVIDIA V100 GPUs with 32GB of memory. Details on the batch sizes and learning rates used for each NN can be found in Appendix A. The measured statistic for each experiment has been reported after 0, 50, 100, 150, and 200 epochs.
4.1 PSEUDO-NTK CONVERGES TO ENTK AS THE NETWORK'S 'S WIDTH GROWS
The first crucial thing to verify is whether the PNTK kemel matrix approximates the true eNTK as whole. We study this first in terms of Frobenius norm.
Theorem 1 (Informal). Let fo RD RO be a fully-commected network with layers of width n whose parameters initialized according to He et al. (2015) initialization, and ReLU-type activations. Let @0(21, 222) be the corresponding pNTK of fo as in (2) and Oo(21, x22) the corresponding eNTK as in ff) a fixed pair of inputs Iı, 22. With high probability over the initialization,
|00(z1, x2) ) & lo 02(21, O( #))
O (21, FF
All of the results in the paper can be straightforwardly extended networks with different widths, as long as the consecutive layers' widths satisfy n4+1 (m1)). Moreover, the results can be made architecturally universal with the techniques of Yang (2020); Yang Littwin (2021).
Remark 2.
Theorem provides the first upper bound on the convergence rate of pNTK towards eNTK. A formal statement for Theorem and its proof are in Appendix B.2.
Remark 3. Based on the provided proof. it's 's straightforward that the ratio of information between off-diagonal and on- on-diagonal elements of the eNTK matrix converges to zero with rate of O(n 4) with high probability over random initialization, as depicted in Figure 2.
Experimental results in Figure 2 support the fact that as width grows, the ratio between off-diagonal and diagonal elements of O₂ (æ1 2222 converges to zero. Furthermore, Figure 13 provides experimental support that width grows, Ôp @ lo converges to Oş in terms of relative Frobenius norm. Note that the result of Theorem only applies to epoch zero of the depicted figures, as the theorem applies only to the networks whose weights are initialization (in the NTK parameterization).
However, as can be seen in the figures, these results don't 't necessarily apply to the NNs whose readout layers are not at initialization (i.e., after few epochs of training). This naturally gives rise to the question: Can the PNTK be used to analyze and represent NNs whose parameters are far from initialization? We will now take various experimental approaches towards studying this question.
4

---

Under review as a conference paper at ICLR 2023
4.2 PSEUDO-NTK' SPECTRUM CONVERGES TO ENTK'S 'S AS WIDTH GROWS
As discussed before, the conditioning of a network's s eNTK has been shown to be closely related to generalization properties of the network, such trainability and generalization risk (Xiao et al., 2018; 2020; Wei et al., 2022). Thus, we would like to know how well the pNTK's ' eigenspectrum approximates that of the eNTK. The following theorem gives a bound on the rate of convergence between the maximum cigenvalues of the two kernels.
Theorem 4 (Informal). Let Se RDD + BoO be a fully-connected network with layers of width n whose parameters initialized according to He et al. (2015) initialization, and ReLU-type activations. Let O0(x1, 22) be the corresponding pNTK of fe as in (2) and 00(21, 22) the corresponding eNTK as in for a fixed pair of imputs I1, 2/2. With high probability over the initialization,
Amaa (Go(s1, F2) & Io)
Amar (09 (x)) E Onnn" #).
Amar (Or (11, x2))
Theorem 4 provides an upper bound on the difference between the maximum eigenvalue of PNTK ver- sus the maximum eigenvalue of eNTK based on the NN's width. A formal statement for Theorem and its proof are given in Appendix B.3, Experimental results in Figure 4 support the fact that as width grows, the max eigenvalue of θg converge to the max eigenvalue of θg. Figure 5 provides similar experimental support regarding the minimum eigenvalues of Ô, and Op in terms of relative difference. Together, these two results intuitively imply that the condition number k(K) = Amax A ///min (K) should become similar as width grows; this is also supported by results in Figure 6. Again, it shall be noted that Theorem 4 only applies to the NNs whose widths are at initialization and as in the previous subsection, these results don't 't necessarily apply to the NNs not at initialization (i.e., after few epochs of training).
Interestingly, the rate of increase/decreass in the difference between maximum and minimum eigen- values and the condition numbers between PNTK and eNTK do not have necessarily have a monotonic behaviour as the training goes on. Observing the exact values of Amin Amax- and for different architectures, widths at initialization and throughout training reveals that in ConvNet, WideResNet and ResNet18 architectures, Amin is close to zero at initialization, but grows during training: the inverse phenomenon is observed with FCNs. Further investigations of these statistics might reveal interesting insights about the behaviour of NNs trained with SGD and the connections between eNTK and trainibility of the architecture.
4.3 KERNEL REGRESSION USING PNTK VS KERNEL REGRESSION USING ENTK
Lee et al. (2019) proved that as a finite NN's width grows, its training dynamics can be well approxi- mated using the first-order Taylor expansion of that NN around its initialization (a linearized neural network). Informally, they showed that when f is wide enough, its predictions after being trained using gradient descent with suitably small learning rate on the training set D can be approximated by those of the linearized network flin::
Min (z) = fo(x) + Oo(z, D) Oo(D,D) OpD - fo(D))
(5)
Ox1
Ox1
OxNO
NOXNO
NOX1
FCN
ConyNet
BecNe11
WideieeNNet
0.15
0.7
3.65
3.00
0.5
3.55
64
3.56
014
0.45
100
150
270
305
104
1000
126
35
Figure 2: Comparing the magnitude of sum of on-diagonal and off-diagonal elements of O₂ at initialization and throughout training, based on 1000 points from CIFAR- 10. The reported numbers the average of 1000 x 1000 kernels each having shape of 10 x 10. The same subset has then been used to train the NN using SGD. As the NN's width grows, the cNTK converges to being diagonal at initialization among all different architectures.
5

---

Under review as a conference paper at ICLR 2023
ConvWoe
Reshet11
WodeResNet
1.2
3.55
0.6
3.55
3.7
0.2
04
3.4
Epards
Rprct
Speck
Spoch
Figure 3: Evaluating the relative difference of Frobenius norm of ea(D, and @n(D, @ To initialization and throughout training, based on 1000 points from CIFAR-10. As the NN's width grows, the relative difference in ||@all and ||0b @ LollF decreases initialization.
where V/D is the matrix of one-hot labels for the training points D, and Oo is the eNTK of f at initialization fo- This is simply kernel regression the training data DD using the kernel O, and prior mean fo- Wei et al. (2022) use the same kernel in a generalized cross-validation estimator (Craven & Wahba, 1978) predict the generalization risk of the NN. As discussed before, using the eNTK in these applications is practically infcasible, due to huge time and memory complexity of the kernel, but we show the pNTK approximates fläm (x) with much improved time and memory complexity.
Theorem 5 (Informal). Let So RDD RO be fully-connected network with layers of width whose parameters initialized according to He et al. (2015) initialization, and ReLU-type activations. the corresponding eNTK
Let 0g(x1,42) x22 be the corresponding pNTK of fe as in (2) and Og (2), x2) as in (1) for fixed pair of imputs x1, Z2. Define
(x)(x)
fo(x) + Ōo(x, D) Oo(D, O0D
(O ('D-Jo(D))
(6)
NxO With proper reshaping, with high probability over random initialization,
1x0
1xO
IxN
NxN
[x) - -f Onn
(7)
Theorem 5 provides an upper bound on the norm difference of flin from (6) and flin from Equa- tion (5) based on the NN's s last layer's 's width. A formal statement is given and proved in Appendix B.4.
Experimental results in Figure show that width grows, the predictions of kernel regression using Ô converge to the prediction of those obtained from using Og while requiring orders of magnitude of less memory and time to compute. Figure 8 shows similar results for the difference in prediction accuracies achieved using kernel regression through θ and O kernels. Appendix B.4 also shows further analysis of how well the linearized network predicts the final accuracy of the trained model for each architecture and width pair. Although [[f flin (æ)]læ decreases with width of the network in Figure 7 at initialization, this does not necessarily translate monotonic behaviour in prediction accuracies, non-smooth function of the vector of predictions; we do see that the expected pattern or less holds, however.
surprising outcome depicted in Figures 7 and is that while training the model's parameters, predictions of flin and plin converge very quickly. This is particularly intriguing as it's in contrast
PCN
ConvMet
Bes
WideRosNet
1.0
3.5
0.0
0.6
04
35
60
220
5
320
153
0.7
130
200
150
Figure 4: Evaluating the relative difference of Amax of (a(D.D) and (, D) at initialization and throughout training, based on kernels on subset ([2) = 1000) of points from CIFAR-10. As the NN's width grows, the relative difference in Xuase ( (D, D)) and (m, D)) decreases.

---

Under review as a conference paper at ICLR 2023
CouxNest
ResNNt19
WeseReeNNet
3.5
1.3
3.3
1051
1.3
Spands
Eprct
Epech
Epech
Figure 5: Evaluating the relative difference of Amin of ( (D, D) and (D, D) at initialization and throughout training. The kernels have been evaluated using subset (iD) 1000) of points from CIFAR-10. As the NN's width grows, the relative difference in ( (Q,(D, and Amin (+o D)) decreases. Note the extremely large values reported for ConvNet. As observed previously in Lee et al. (2020) and Xiao et al. (2020), CNN-GAP ill-conditioned and Amin (On( ((D,D)), D)) + O, while Amii (D, D)) > 0.001, causing the huge discrepancy. More details in Appendix B.3.
PCN
ConvNet
BesNet11
WidoResiNet
14663 A066) (eobul
-4.5
0.5
02:28
-0.2
8-5
00-404
92-3192
-1.2
250
50
130
150
200
100
115
Spoch
Spock
Epock
Figure 6: Evaluating the relative difference of condition number of eNTK and pNTK at ini- tialization and throughout training. Condition number is defined as x(K) = Amax ( / (K). The kernels have been evaluated using subset ((2) 1000) of points from CIFAR-10. As the NN's width grows, the relative difference in Amin (O,(D, D)) and Amin (D, D)) decreases. The bizarre values found in the ConvNet plot can be addressed as discussed Figure 5; more details in Appendix B.3.
CoayNet
WideResNet
3.5
KF1
0.4
0.5
35
Epedds
Spacs
Speck
Figure 7: Evaluating the relative norm difference of kernel regression outputs using eNTK and pNTK as in Equation (5) and Equation (6) at initialization and throughout training. The kernel regression has been done on 'D = 1000 training points and | = 500 test points randomly selected from CIFAR-10's train and test sets. As the NN's width grows, the relative difference in flin(2) and plin (X) decreases at initialization. Surprisingly, the difference between these two continues to quickly vanish throughout Sthe training process using SGD.
with the experimental results depicted in Figures 2, 4, and 13. In other words, although the kernels Og and Ô @ Lo seem to be diverging in Frobenius norm, cigenspectrum, and so on, kernel regression using those two kernels converges quickly, so that after 50 epochs the difference in predictions almost totally vanishes. We believe further investigation of why this phenomenon observed could lead to new interesting insights about the training dynamics of NNs.
KERNEL REGRESSION USING PNTK ON FULL CIFAR- 10 DATASET
Thanks to the reduction in time and memory complexity of pNTK over eNTK, motivated by Theorem and experimental findings in Figure 8, we finally evaluate the corresponding pNTKs of the four architectures that we have used in experimental evaluations in different widths using full CIFAR-
7

---

Under review as a conference paper at ICLR 2023
FCN
CoovNet
ResNet18
WideR
N0.12 V0 2048
AaalAA 180
3.0
0700
12
Spech
Epech
Figure 8: Evaluating the accuracy difference of kernel regression outputs using eNTK and pNTK as in Equation (5) and Equation (6) at initialization and throughout training. Kernel regres- sion using pNTK always achieves a higher test accuracy in comparison to kernel regression using eNTK. The kernel regression has been done on [D] 1000 training points and | 500 test points randomly selected from CIFAR-10's train and test sets. As the NN's width grows, the prediction accuracy of in fiin and flin decreases initialization. Again, the difference between these two continues to vanish throughout the training process using SGD.
ConyNNt
ResNatt8
WinteBuNNat
70
73
61
150
Spech
Figure 9: Evaluating the test accuracy of kernel regression predictions using pNTK as in Equa- tion (6) on the full CIFAR-10 10 dataset at initialization and throughout training. As the NN's width grows, the test accuracy of flin is also improved, but eventually saturates with the growing width. Using trained weights in computation of PNTK results in improved test accuracy of flin.
PCN
ConviNtt
ResNett
Wide Raeeee
30
0.7
33
30
33
100
330
122
11.6
114
2.0
3:06
11.6
200
140
1.0
Spoch
Figure 10: Evaluating the test accuracy of model f throughout SGD training on the full CIFAR-
10 dataset. In contrast to flim, the test accuracy of f does not significantly improve with growing width.
BesNett8
Widee
vin
1014
40
512
300
33
27
12002
Spoch
Spech
Spech
Figure 11: Evaluating the difference in test accuracy of kernel regression using pNTK as in (6) vs the current model f throughout SGD training on the full CIFAR-10 10 dataset: how much does a linearized predictor with the current representation improve prediction accuracy over SGD?
10 data, at initialization and throughout training the models under SGD. As mentioned previously, running kernel regression with eNTK on all of CIFAR-10 would require evaluating 25 x 10" Jacobian-vectus products and more than ss 1.8 terabytes of memory; using pNTK, this can be done with far more reasonable 25 10 Jacobian-vector products and 2 18 gigabytes of memory. This is
8

---

Under review as a conference paper at ICLR 2023
still heavy load compared to, say, direct SGD training, but is within the reach of standard compute nodes.
Figure 9 shows the test accuracy of flin on the full train and test sets of CIFAR-10. In the infinite- width limit, the test accuracy of flin at initialization (and later, because the kernel stays constant in this regime) should match the final test accuracy of f: that is, the epoch 0 points in Figure 9 would agree with roughly the epoch 200 points in Figure 10. This comparison is plotted directly in Appendix C. Furthermore, the test accuracies of predictions of kernel regression using the pNTK are lower than those achieved by the NTKs of infinite-widd width counterparts for fully-connected and colly-convolution networks. This is consistent with results on eNTK by Arora et al. (2019a); Lee et al. (2020), although Arora et al. (2019a) studied only a "CIFAR-2" version.
It is worth noting from Figures 9 and 11 that, in contrast to the findings of Fort et al. (2020), we observe that the corresponding PNTK of the NN f continues change even after epoch 50 of SGD training. Although for fully-connected networks and some versions of ResNet18 this change is not significant, fully-convolutional networks and WideResNets the pNTK continues to exhibit changes until epoch 150, where the training error has vanished. We remark that Fort et al. (2020) analyzed eNTKs based on only 500 random samples from CIFAR- 10, while the pNTK approximation has enabled us to
analysis on the 100-times larger full dataset.
Lastly, to help the community better analyze the properties of NNs and their training dynamics and avoid wasting computation by redoing this work, we plan to share computed pNTKs for all the mentioned architectures and widths, as well as pNTKs of ResNets with 34, 50, 101 and 152 layers on CIFAR-10 and CIFAR-100 (Xiao al., 2017) datasets, both at initialization and using pretrained ImageNet (Deng et al., 2009) weights. We hope that our contribution will enable further analyses and breakthroughs towards stronger theoretical understanding of the training dynamics of deep neural networks.
6 DISCUSSION
Our pseudo-NTK approach to approximating the empirical Neural Tangent Kernel has provable bounds, good empirical performance, and multiple orders of magnitude improvement in runtime speed and memory requirements over the direct empirical NTK. We evaluate our claims and the quality of the approximation under diverse settings, giving new insights into the behaviour of the empirical NTK with trained representations. We help justify the correctness of recent approximation schemes, and hope that our rigorous results and thorough experimentation will help researchers develop a deeper understanding of the training dynamics of finite networks, and develop new practical applications of the NTK theory.
One major remaining question is to theoretically analyze what happens to the pNTK or eNTK during SGD training of the network. In particular, the fast convergence of flin and flin when training the network, as seen in Figures 7 and 8, runs counter to our expectations based on the approximation worsening in Frobenius norm (Figure 7), maximum eigenvalue (Figure 4), and condition number (Figure 6). This seems likely to be very important to practical use of the PNTK.
Perhaps relatedly, it is also unclear why pNTK consistently results in higher prediction accuracies than when kernel regression is done using eNTK, given that our motivation for PNTK is entirely in terms of approximating the eNTK (Figure 8). Intuitively, this may be related to regularization-tyy effect: the pNTK corresponds to a particularly limited choice of a "separable"
pverator-valued kernel (e.g.
Álvarez et al., 2012). Separable kernels are a common choice in that literature for both computational and regularization reasons; by enforcing this particularly simple form, we remove many degrees of freedom relating to the interaction between "tasks" (different classes) that may be unnecessary or hard to estimate accurately with the eNTK. This might, in some sense, correspond to one-vs-one rather than re-stt framework in the intuitive sense discussed in Section 3. Understanding this question in more detail might require a more detailed understanding of the structure of the eNTK at finite width, and/or a much more detailed understanding of the interaction between classes in the dataset with learning in the NTK regime.

---

Under review as a conference paper at ICLR 2023
Finally, even the pNTK is still rather expensive compared running SGD neural networks. It might make for a better starting point than the full eNTK for other speedup methods, however, like kernel approximation schemes or advanced linear algebra solvers (e.g. Rudi et al., 2017).
REFERENCES
Ben Adlam, Jachoon Lee, Lechao Xiao, Jeffrey Pennington, and Jasper Snock. Exploring the uncertainty properties of neural networks' implicit priors in the infinite-width limit. 2020.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In NewrIPS, 2019a.
Sanjeev Arora, Simon S Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang. and Dingli Yu. Harnessing the power of infinitely wide deep nets on small-data tasks. 2019b.
Devansh Arpit and Yoshua Bengio. The benefits of rereparramettiizationn at initialization in deep relu networks, 2019.
Gregor Bachmann, Thomas Hofmann, and Aurélien Lucchi. Generalization through the lens of leave-one-out error. In ICLR, 2022.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. 2021.
Peter Craven and Grace Wahba. Smoothing noisy data with spline functions. Numerische mathematik, 31(4):377-403, 1978.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. Advances in Neural Information Processing Systems, 33:5850-5861, 2020.
Jean-Y Yves Franceschi, Emmanuel de Bézenac, Ibrahim Ayed, Mickaël Chen, Sylvain Lamprier, and Patrick Gallinari. A neural tangent kernel perspective of gans. 2021.
Tamir Hazan and Tommi Jaakkola. Steps toward deep kernel methods from infinite neural networks. 2015.
Bobby He, Balaji Lakshninarayanan, and Yee Whye Teh. Bayesian deep ensembles via neural tangent kernel. Advances in Neural Information Processing Systems, 33:1010-1022, 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, 2015. URL httph:t////// /ttpa:///axxvvvvoaaeeeao abbs/ 1502. 01001115 852.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
David Holzmüller, Viktor Zaverkin, Johannes Kästner, and Ingo Steinwart. A framework and benchmark for deep batch active learning for regression. 2022.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In NeurlPS, 2018.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz. Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks gaussian processes. 2017.
10

---

Under review as a conference paper at ICLR 2023
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In NeurIPS, 2019.
Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study, Advances in Neural Information Processing Systems, 33:15156-1511 2020.
Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour wide deep neural networks. 2018.
Mohamad Amin Mohamadi, Wonho Bae, and Danica J. Sutherland. Making look-ahead active learning strategies feasible with neural tangent kernels, 2022.
Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pp. 29-53. Springer, 1996.
Timothy Nguyen, Zhourong Chen, and Jachoon Lec. Dataset meta-learning from kernel ridge- regression. 2020.
Timothy Nguyen, Roman Novak, Lechao Xiao, and Jachoon Lee. Dataset distillation with infinitely wide convolutional networks. Advances in Neural Information Processing Systems, 34, 2021.
Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are gaussian processes. 2018.
Roman Novak, Jascha Sohl-Dickstein, and Samuel Stern Schoenholz. Fast finite width neural tangent kernel. In Fourth Symposium Advances in Approximate Bayesian Inference, 2021.
Daniel S Park, Jachoon Lee, Daiyi Peng, Yuan Cao, and Jascha Sohl-Dickstein. Towards nngp- mugp-guided neural architecture search. 2020.
Alessandro Rudi, Luigi Carratino, and Lorenzo Rosasco. FALKON: An optimal large scale kernel method. In NeurlPS, 2017.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara -ridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in Neural Information Processing Systems, 33:7537-7547, 2020.
Martin Wainwright. High-DDienssioal Statistics: A Non-Asymptotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2019. doi:
10.1017/9781 108627771.
Haonan Wang, Wei Huang, Andrew Margenot, Hanghang Tong, and Jingrui Hc. Deep active learning by leveraging training dynamics. 2021.
Alexander Wei, Wei Hu, and Jacob Steinhardt. More than toy: Random matrix models predict how real-world neural representations generalize. 2022.
Christopher Williams. Computing with infinite networks. Advances in neural information processing systems, 9, 1996.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07744. 2017.
Lechao Xiao, Yasaman Bahri, Jascha Doil--iccstten, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and mean field theory of cnns: How to train 10,000-layer vanilla convolu- tional neural networks. In International Conference on Machine Learning, pp. 5393-5402. PMLR, 2018.
Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generaliza- tion in deep neural networks. In International Conference on Machine Learning, pp. 10462-10472. PMLR, 2020.
11

---

Under review as a conference paper at ICLR 2023
Greg Yang. Wide feedforward recurrent neural networks of any architecture gaussian processes. Advances in Neural Information Processing Systems, 32, 2019.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture, 2020.
Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics, 2021.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. BMVC, 2016.
Yufan Zhou, Zhenyi Wang, Jiayi Xian, Changyou Chen, and Jinhui Xu. Meta-learning with neural tangent kernels. 2021.
Mauricio A. Álvarez, Lorenzo Rosasco, and Neil D. Lawrence. Kernels for vector-valued functions: A review. Foundations and Trends@ in Machine Learning, 4(3):195-28) 2012. doi: 10.1561/ 2200000036.
12

---

Under review as a conference paper at ICLR 2023
A DETAILS OF EXPERIMENTAL SETUP
In this section, we present the details on the experimental setup used for the plots depicted in the main body of the paper. As mentioned, the exact width for FCNs have been reported. For WeetResNet-16-A we use two block layers, and the initial convolution in the network has a width of 16WF where WF is the reported WF. For instance, WF 16 means that the first block layer has width of 256 and the second block layer has a width of 512. For ResNet18, we also used the same approach, multiplying WF by 16. Thus, when WF 4, the constructed network will have the exact architecture as the classical ResNet18 8 architecture reported. A WF of 16 means ResNet18 with each layer being 4 times wider than the original width.
When training the neural networks using SGD, a constant batch size of 128 was used across all different networks and different dataset sizes used for training. The learning rate for all networks was also fixed to 0.1. However, not all networks were trainable with this fixed learning rate, as the gradients would sometimes blow up and give NaN training loss, typically for the largest width of each mentioned architecture. In those cases, we decreased the learning rate to 0.01 to train the networks.
Note that to be consistent with the literature on NTKs, techniques like data augmentation have been turned off, but weight decay of 0.0001 along with momentum of 0.9 for SGD is used. Data augmentation here plays an important role in the attained test accuracies of the fully trained networks.
B FURTHER RESULTS ON THE APPROXIMATION QUALITY
In this section we'll Ill lay out the proof's of the theorems provided the main text, mainly Theorems 1, and 5. Towards this, we first define some notation and show a simple recursive formula for computing the tangent kernel that we take advantage of to prove theorems. Consider a NN f RR. - R We assume the final read-out layer of the NN is dense layer with width w. Assuming the NN has L layers, we define 0₃ to be the corresponding parameters layer E {1, 2,....
: L}. Furthermore, let's I as the output of the immediate last layer of the NN f, such that f(xx) = (L9(2)
define g : IR_ for some 0 E R0x=
As shown by Lee et al. (2019); Yang (2020), the NTK can be reformulated as the layer wise sum of gradients (when the parameters of each layer Or are assumed to be vectorized) of the output with respect to 0g. Accordingly, denote eNTK of NN
0,(01, = Σ Vo.l(t))oaf(ss)
(8)
Now, noting that the final layer of f is dense layer, 0r. Thus, we can rewrite (8) as
the chain rule to write fz) as
af 8g(x) where Øf(x) 0g(ar) 00:
(g(a)
(2()) =
Σ
Or Vo.9(Z1) Ve.((Z)) To VeL
1=1
L
(9)
Σ Vo,g(z)) Vog(zz) 0r T +Tz))(z) g(22)
0r.
= Oz (x oz + g(x1) g(22)
Applying Equation (9), we can already see that the pNTK of network simply calculates weighted summation of all elements of eNTK into a scalar, since it can be seen as adding new final dense layer to the network / with the fixed weight vector 1 1o where 1o is the O-dimensional vector consisting of all 1s. Note that, however, this 10 vector is not trainable in this context and is a fixed vector.
vo
Before moving on with the approximation proofs, we would like to mention that the proofs in this section rely heavily concentration inequalities of sub-expor mib-exxonential random variables. Thus, start by providing some background about aab-exponential random variables and the related concentration inequalities that will use later on.
13

---

Under review as a conference paper at ICLR 2023
B.1 BACKGROUND ON SUB- EIB-EXPNENNIAL RANDOM VARIABLES
A real-valued random variable X with mean j is called sub-exponential (Wainwright, 2019) if there non-negative parameters (v, o:) such that
Ble w Se
22
for all ||| a
We use X ~ SE a:) to denote that X is a sub-exponential random variable with parameters (v, a). but note that this is not particular distribution.
A famous sub-exponential random variable is the product of two standard normal distributions, 24 (O, 1), such that the factors are independent (X = |z| ||| SE ap)) with mean
2/x) or the same (X₂ 2² SE(2, 4) with mean 1). We now present few lemmas regarding sub-exponential random variables that will come in handy in the later subsections the appendix.
Lemma 6. If random variable X is sub-exponential with parameters (v, a)), then the random variable sX where E Rt is also sub-exponential with parameters (sv, (s),
Proof. Consider X SF a)) and X' sX with E[X* = s E then according to the definition of sub-exponential random variable
2²12
E exp (A(X- - -x)) exp(
for all |\]
2
a
exp (ax - )46) s10)]
2²82
A
1
5 exp(
for all
1
(10)
2
80X
2²222x²
x=
1
E (exp -x'x1 exp(
for all |I|
80
Defining a' = and v' sv we recover that X' ~ (F(ss,80))
Proposition 7. If the random variables X, for i E [1 for N E N are all sub-exponential with purameters (Vi, as)) and independent, then EE X, SB(V E u7 max; az), and 1 max, as
1 N E
X SE
1 E
Proof. This a simplification of the discussion prior to equation 2.18 Wainwright (2019). Proposition 8. For random variable X SE(V, a), the following concentration inequality holds:
122 1
Pr([X n| 2 t) S 2exp
(-min
2y2
2a
Proof. The proof directly follows from applying scalar multiplication to the result derived in Equation 2.18 in Wainwright (2019).
Corollary For random variable SElu, a), ) the the following inequality holds with probability at least - 8:
2 XX 11 max p 2 log z: 2a log 2 g (2/2
B.2 PSEUDO-NTK RELATIVELY CONVERGES TO ENTK As WIDTH GROWS
Let's denote a neural network with L dense hidden layers whose width is n as:
(x))) æ
(x) = (W((++) f'(x)) f(x) = f'((x) = p(
(1I)
such that à is differentiable coordinate-wise activation function.
14

---

Under review as a conference paper at ICLR 2023
Setting A (ReLU-MLP). We assume the following assumptions hold in our setting:
We assume W [R for E 1,.... L is initialized according to the He et al. (2015) initialization, meaning that each scalar parameter is distributed according to N/0, 1/m--))-
We assume the width of all hidden layers are identical (and equal to n). The proof extends naturally to the case of non-equal widths as long m441/m + C₁ E (0, oc) for each consecutive pair of layers.
We assume e is the ReLU activation. This can be generalized to I-Lipschitz, ReLU-like functions such as GeLU, PReLU, and so on, as discussed in Appendix B.5.
We assume the training data X' is finite and contained in compact set and there are no overlapping datapoints.
Note On Parameterization Although we assume a Gaussian distribution for each scalar variable, the proof's in this section apply to any other distribution used for scalar parameters long as:
The variance of the parameters set according to He et al. (2015), and the mean zero.
Each scalar parameter is initialized independently of all other ones.
The distribution used is sub- uub-Gaussian.
This applies to all bounded initialization methods, like truncated normal or uniform on an interval. In what follows, we generally assume each w/+1 E SG(1/m) with mean zero. In general, the product of two sub-Gaussian distributions has sub-exponential distribution. For the product of two independent weights, (++ w//
J+1 with i and/or j b, denote the parameters 1. SE ap)- ni
For (wl) )², we use the parameters 1 S.VV,,,, (tg), whose mean is A4 + 0.
Note that we can recursively define the eNTK of pl+1 using the eNTK of ft
(mi,
(x++(x1)
(x)))(2))
(xt) (x2)
(1++) (11,
OW()
OW
awi+1
aWL+1
(o(W(tt))
(0(W(+)))
+ k (x1,
awG
OWG
36(W((1)
(m))
(x2)
(6(W(+)) (12))
+ kht (I1,
(f'(1))
OW
awG
0f((22)
de (((L)))))nn)
())) (12)
(
0f
aW
OWO
0f
06(WV-1)
@ x2)
+ +KD kdtt (21,02)
0f((z1)
0f((22)
(12)
where
26(W(+))((E)) w(+1) o ((ww(n(())) 1xn (x)
(13)
and 121, x2) T o (x2) following the same equations with ø the identity function, so that of (z) = 1. Furthermore, note that using the same approach we show that PNTK of the layer I can be derived as
is diagonal matrix. We think of the last layer as
( (21. =
1. ((+1) (2+1)(2110))) T22) n
T
(14)
Vn
where 1, is the vector of is with size
15

---

Under review as a conference paper at ICLR 2023
We now sketch the proof idea first and then move onto rigorously proving each part of the sketch. First, note that using Equation (12) we can recursively calculate the eNTK of a general MLP. We take advantage of this recursive definition and derive bounds for the magnitude of the elements of the eNTK on a layer-to-layer basis recursively. To do so, we first show that the eNTK of the first layer of the NN, @ (x1, I2).. ). is in general a diagonal matrix. Then, we present a series of Lemmas that bound the elements of the eNTK of layer I + based on the magnitude (bounds) of the eNTK of layer 1. Finally, based on the derived bounds on the magnitude of elements of the eNTK of a NN with I layers and Equation (14), we prove that the Frobenius norm of the pNTK relatively converges the Frobenius norm of the corresponding eNTK with high probability over random initialization.
Before moving on, it's useful to first show simple inequality on the elements of tangent kernel based on the Lipschitz-ness of the activation function; this will help us further in deriving the aforementioned bounds. Define (xx) = W o [O(W ( of (1++1) x2)
(1=1(=))
We can write each entry 1xn
@ (21,
va+ (x1) in (1 + (a(i) [(a)
b=1
(0
(21,
ia (( 1,000 (2I, (11()=) ((z))
ap a (11, + (x1) f (x:)) I(i j) (15)
where I denotes the 0-1 indicator function, the first inequality follows from the activation function ø being I-Lipschitz and the second incquality follows from the product of two sub-gaussian distributions being distributed as a sub-exponential variable. Note that we use S EEajo ap)) for the sake of generality but in the case i = A a = it actually refers to S S (V,
Lemma 10 (Diagonality of the first layer's tangent kernel). For a NN under Setting A, the correspond- ing eNTK of the first layer (x), 222) is diagonal. Moreover, there is corresponding constant (O)) 0 such that for all diagonal elements (0)) 22),0 have that
(0 (21, CO).
Proof. Consider the one layer NN (x) = (())))) For this case, we have:
if
e (21,
(16)
0
if
and thus, since the activation function Φ is I-Lipschitz we can conclude that for all i,
|
xel
ifi i
e (2),
(17)
0
Thus, the tangent kernel of the first layer is a diagonal matrix whose entries are independent of the width of the first layer (n), and can be bounded by positive constant, given by C(I) max xxx |x] m2)-
Next, we present series of lemmas that will help derive the bounds on the elements of the tangent kernel of layer /+1 based on the bounds of the tangent kernel of layer l. The next lemma specifically bounds the values of the tangent kernel of the second layer, based on the diagonality of the first layer's 's tangent kernel and (15).
Lemma 11. Consider a NN under Setting A with depth 2 ++ 1. Assume there is a constant ((1) > 0
such that ( (010 (x1, x2)aa < CD) for all i E [mu] and every non non-diagone element of (()) (z), x2) is
zero. Then for any small 8 0 there are corresponding constants c((+1) polylog(n/ (), C cll+1)
16

---

Under review as a conference paper at ICLR 2023
polylog(n/6)) n't >> such that for any >> ntt,, it holds with probability least 5 that, simultaneuusly for all i.
C ( i
10(+11) (21,
S
(18)
c //n i.
Proof. Recall that
((++))
(11,
02
(r1) v(+ (1) + ((a))
I(z)) 1),
a=1 b=1
where
(11+ -((+1) ta
Σ w w (l()) (
The weight W(++1) in is tab-Gaussian, and the å term has absolute value most one, their product is also sub-Gause GabSisssiinn with the same variance proxy parameter (as can be casily seen e.g. from the moment-based characterization of aub-Gausiianity). As noted before, when the indices match we replace (z1) (1+1) (x2) )ia by Ess (Va, ax)) (which has mean μ4 * 0). In the case when the indices are not the same, the two are independent but do not necessarily have mean zero. However, mean of each decays with 1 at least. In what follows, we show this for Gaussian initialization (extension to other Lub-Gaussian distributions is straightforward) To show that the (1+1) (x) decays with 1 we introduce an event
ny
mean of each t/
E: W W (z)
w f'(x)d) A
w(l+t) <0
c=1,caa
and compute the conditional mean on happening or not, as
E (((())
y (1(+1)(x))
(x) E Pr[p] + E V ((tt (z) -B Pr[-E]
E
where E ( (z) -E and Pr[E] Here, E corresponds the event where the
weight scalar w w
((+1) is correlated with the argument in à and d(.) = 1. Without loss of generality, (t++)())) (x) E on (x) 0 and 1,cpt (1+1) fl(x), 0 (as it indeed follows a zero mean normal distribution, hence introducing a factor of which we omit for simplicity) and calculate
can condition E
E (x)
E
E
(((+1)
(=(+1)
(()
Σ
w
in
(
a=1ocpa
wo(++)) > b)
(19)
E
MmN-(0,00
((b)
E ni 6~~
(0,02)
((b))
where x² [x))) @ is the standard normal PDF and the standard normal CDE Note f(x))n
that d(b) (b))
|b| + Hence, assuming x² = ( E [((+1)))) E decays with a rate
and so does E [x)a+ Accordingly, for the rest of the proof, we use the term "e to represent
the mean of sub-exponential variable (E ap) which stems from product of (1++) (2)+)))()))))))) ia and
((+1) (z22)s where the indices ia and jb don't match.
17

---

Under review as a conference paper at ICLR 2023
Based on Equation (15) expand the elements of ( (21. 222)
W w((+))
(21, #2) +ff((a)) (x2)) if
(21,
W(++1) W
jb 1 (31,
if
in
SEj (v., 0g a, + (x)) ((22) ifi=j
Ekaja a2 (00
if
(20)
We've ve assumed an upper bound on the diagonal elements of (21, 12) C(), WW have then that
1 SE(U,, a.) (21, I Σ SEia(us; ~ SE a=1
C(I)
V.C V,C(
a,CC
TI
vn
since each item in the last sum is independent of all others. Noting this term has mean C Ms. we get brggh-pbbbbblity upper bound via Corollary Since there are n independent such terms, one for each i E [ri], we'll want upper bound that holds over all of them: with probability at least 1 òi we have
(Eia (0)) az C00
ps max
20(()
2
log
2n
8C
2n log
71
n1
81
2n
< (())
M4 + 21
log
51
for n large enough that 2v2 > 8// 1 log
22 b, which will be the case for n = Ω (polyiog 2t)
Likewise, Lemma 17 shows high-probability upper bound on (f') f'(222) for all n > with probability least 5g. Combining the two with a union bound, we have with probability least 1 61 8g that
2n log
( S ((.)
(1s + V
+GOn.
81
This inequality is dominated by the G term. Thus can find upper bound (G()) +em, where for 2 max (Rm: (1polylog(1/66)) we have that O ( 3 can be chosen to be independent of n.
2m log
O (v log 1 bg
81
For the off-diagonal terms, have
C (cn
S (vp. ap)l-
n
Again applying Corollary 9 for each j and taking a union bound gives that with probability at least 1 8z.
max 0 (++ (x1, S ((l)
Mp
Up log nnn - 1)
+
n²
82
as long is large enough that u, ap log
m 62
again true long as Ω (polylog +
Hence, based on this bound and also the bound provided for the diagonal entries, we can achieve the desired result by setting n > n/+ with probability at least 1 (ő1 + 62 + dg)-
18

---

Under review as a conference paper at ICLR 2023
Lemma 12. Consider NN under Setting A with depth 2 ++ 1. Assume there constants c = polylog(n/6p) Clu polylog(n/6p) such that we have
< c.'n and
(21, C2 vn with probability at least 1- 1- and any width > up simultaneousty for all i.j. Then for any arbitrary small 5 >> O there are constants C/ ((++) polylog(n/8), C
polylog(n/0)) nl41 > 0 such that for any n > +e and i
(oltt2n chtt1)
(2(+))(Z1,22)) 22)gg
n
ifi
(21)
C2
Vn
i
with probability at least 1 - 6. In other words, the magnitude of elements of the tangent kernel in the recursive definition will not grow.
Proof. Using the same expansion that we utilized in the proof for the previous Lemma, for the diagonal elements of ((1+)))x1 we have:
(0(41)) (x1, x2
w(+1)
ia w((1)) w0 (I1, (00 + G
SEialVs a.), (1))
Σ
SE ap) ap)
+G
b=1,bşa
@((+t1) (01,
@ 2(11)) (01,22a
(22)
where
(1)
a'yn
10m-11
1) (I x2) <
Se (Ve, ~ 11) VrSE
(23)
n
with probability at least (1 - 5pp) and
c((()
n
0 () (I1, x2)a I
Σ
ap)
nVn
a=1 b= 1,bba
(24)
Cl
n²
( Vn
SE SE 0,) = SE Upp )
ny
with probability at least (1 - 6p) As shown before, both of these terms are dominated by the G00n term in the inequality for diagonal elements and thus, similar to the previous Lemma, we can (where accordingly nmm the minimum width coming from Lemma 17) such that for all > n +1; [@(++)(( 22) al (G(++))n with probability at least 1 - (51 + 5p + 5g) where 81 comes from bounding the sub-exponential variables above and can be replaced with positive constant (which can decay with O(n +)). Again, just like the previous Lemma, the same high probability bound (albeit with very slight modifications) holds for all diagonal elements.
show that there is an +1 max(Am; Tlp, polylog(1///)
For the nin-diagonal elements of (!-1)) (x1, we have:
(0(11) (2 22)
W(+1) ((+1)) 10 (000
a=1
13
n Σ
EEnajo (Mps a. (11,
SE (up, (x1,
6(4+)) (21,22)
(25)
19

---

Under review as a conference paper at ICLR 2023
where
c n
|0| (0 (21,
S Praja (Up Ap) ~ 11 Vn|SE
apvn
Upp
(26)
n
n
and
c vn
(x1,
5
Biigh (yp ap]l
am 1 b=1,bşa
Op (27)
( Vn
n²
SE (up: pp)]= C2 VTISE (uz-
n
each with probability at least n np
8.n Thus, we can claim that according to Corollary 9 that for all
1e (2 (21,22)is
S C(C)
n² + vT 20, log log
Pe
2 82
(28)
with probability at least 282 conditioned the entries of previous layer's tangent kernel being bounded as mention in the Lemma assumption (Note that we can take advantage of this assumption here as we will include the failure case on the diagonal elements and then have a union bound on diagonal and non -inagoaanal elements). Here 52 comes from applying the Bernstein inequality on the two sub-exponentia variables. Applying union bound on the bounds derived for each entry of the tangent kernel we can see that
(1(+))) (xI, Oll
Hp +
2up
4 log 82
+ 2llog
(29)
n2Vn
112
for all i # with probability at least 1
Thus, the Lemma's 's claim holds with probability least -61 + 8/2 + 8p + 0g) and n as desired.
An alert reader already can notice that connecting the previous three Lemmas would result in an upper bound for the diagonal and non-diagonal elements the eNTK the NN f at initialization. Lemma 13. Consider NN f under Setting A. For every arbitrary small 5 >> O, there are constants G = polylog(n/O) C2 = (olylog(n/o), no > such that with probability at least å over the random initialization, it holds sinultaneously for all i, that if n > no. the corresponding eNTK of on the arbitrary datapoints Iı and 22 satisfies
Cin
ifi=
10(1) -
(30)
C2vn if
Proof. Starting with Lemma 10, we have a bound for the entries of the first layer's tangent kernel. Plugging this bound into Lemma 11 we get bound on the elements of the second layer. Next, we can recursively apply Lemma 12. Note that, however, this recursion will induce a new factor in our bound that depends on the depth of our NN, and thus, would enforce the minimum width to also depend on depth. Assume we have applied the first three lemmas and we want to consider the next L 3 recursions of application of Lemma 12 for NN with L > layers. As our assumption, for any arbitrary small 8 > there are constants ((3) (ppolyoogn//8), cfa (polylog(n/88) (3)
such that for any n > (3) and i,j
(3) n
ifi=j
10 (21, 5
(31)
c
(3) Vn
if
20

---

Under review as a conference paper at ICLR 2023
Now, through recursively applying Lemma 12 that for the layer L for any and n nLL) there are constants C( O(polylog(n(L- 3)/8)), c O(polylog(n(L. 3)/6)), (L2) > such that
c i ((1) Vn i
O (I1, S
(32)
with probability at least 1- The change (1) is also trackable as presented Lemma 12 through max(n (3), H(polyyog((E. 33)/))) for each layer T.
Lemma 14. Consider a NN f under Setting A. For every arbitrary small 5 >> 0 and the arbitrary datapoints xn and V2, it holds that
10(x), 0(11, S ((y/llogn)
(33)
with probability least 8 8 over random initialization for any nl where ny = pply)oo(L//8) In other words, the Frobenius norm of the difference between eNTK and pNTK evaluated on two datapoints are bounded by O (v)
Proof. We note by D(11, ((1) - ( in Equation (14) we can write the
-60 (2) 22) @ lo. Using the expansion provided
0 0
SEins (ups ap (0))))1111)))))))))))))))))))) I2)
Eeads (Mps ap
0
12(2) S
b=1
c=1
d=
a=
S (vajo a (21,
(34)
where first option is for diagonal elements i and second for non-diagonal ones.
Applying Lemma 13 we can assume there are constants C(L-1) = polylog(n/6p). C chb-1) polylog(n/65,) such that (2) with probability at least 1 - 6p simultaneously for all a, b. Thus we can write the diagonal elements of (2))
c(T-1) n and (e(f-1) (x,, x2)abl cfi-1) Vn
D( S c(l-1)) n
Σ
SEia at.)
0 Σ
1
Sen,
C
0 (L-1)
S Ecana (Ups dp))
d=1,defc
Da(m
Cll (L-1) 1)
C (L-1) n
n
0
SEaat (pps ap)l
b=1.bb
d=1
Da(#i
21

---

Under review as a conference paper at ICLR 2023
We would like to bound each of (2(21 22) for k (1,2.3,4) and then find bound for the diagonal elements using the combination them. Starting with (21 x2)
"
(211 S 1
Σ
(us (c.)
S Eua (Ks)
0
c(1-11
SEiolUs: (ks) - Siia (Us: (s5)
1
Σ
S
(Us)
0
c=l,cii
-((
(1-1)
12
SE
a0 1) 11
SE
1 1)
022
0
0
0²
0
=((-11
10
SE
v2(1
1)2
+
(2 1)
max
a2
1)
02
02
0²
0
0
VSEE
V/
1
1
a
1 (1- ))
0 vn
(35)
Thus, using Corollary we can claim
(211) < c(4-1) vn
1/
(1
1
8 log σ Ms na/1
=
(36)
with probability at least 8/4..
For the other terms, we can simply the analysis through noting that they are all form of weighted summation of independent sub-exponential random variables of the same distribution. For such a summation with a weight of and b summation terms, we have
b
X
a
Σ
a Σ
SE ap)l ) = =aSE (12Vb,a)),
(37)
Thus, applying Corollary 9, we get that
8 2a max ypp b blog 8 ap log 8)
[x|I]]
(38)
with probability at least
8/4. Accordingly we can claim
(2)(1))22all <2c((-1)
max (vap/ (1- 02
1
log ap log
8
8 z
nス
up
(39)
vn max /mm/ (1 (Da(1) < 2c((t-11)
8 log , ap log 3
(40)
n?
n²
nn
(D((1))222l < 2c((-11) v max /m// (n/(1-1
1
8 ap
8 log
Mp
n² +
log
O2
O
ni /n
(41)
all independently and with probability at least 1 - 8/4. Moreover, one can easily apply the same technique and that D(x1,12)1 for i + j follows similar bound to the one of Equation (40).
22

---

Under review as a conference paper at ICLR 2023
Thus, loosening the off off-diagonal terms for simplicity, applying union bound on the previous three inequalities yields
8 8 8 ))
D(11, x2),,g (o + C/2 1 Vn max Up log 8r ap log
(42)
with probability at least - (5 + 6p)-
Finally, as [ x2)ll VEis 12) if each entry's 's absolute value less than >
then the Frobenius norm is less than to. Thus we can combine bound on each of the O entries to see that
Pr [I((01 = 80 (c + c( + Ms) Vnnax 2//log
802 log {0) ))
80 16(6++p
(43)
desired.
Lemma 15. Consider NN f under Setting A. For every arbitrary small 8 >> 0 and the arbitrary datapoints and T2, it holds that
(2)11 nogl
(44)
with probability at least 1 - å over random initialization. In other words, the Frobenius norm of the eNTK evaluated datapoints is lower bounded by O(n (og n)).
Proof. Considering that the dot product of post-activations appear in the diagonal elements of the eNTK in conjunction with Lemma 17, this is straightforward. Note that this bound also applies to the maximum eigenvalue of the eNTK matrix since the maximum eigenvalue is bigger (or equal to) than the sum of elements the matrix divided by the number of columns.
We are finally ready to present the proof of Theorem 1. Theorem 16. Consider a NN I under Setting A. For every arbitrary small 8 >> O and the arbitrary datapoints V and F2, there exists ng such that
(2)) x2) 61r) (2), IollF O
1 Vn
(45)
(11,
with probability at least 1 for >
Proof. The proof straightforward from applying Lemma 14 and Lemma 15.
Consider NN under Setting A with L 2 2 and ReLU activation function. The dot product of two post-activations (x,)) (() (x22) grows linearly with the width of the network with high probability over random initialization.
Lemma 17.
Proof. We begin by showing that the dot product of the post-activations of the first layer of the NN under setting Setting A grow linearly using a simple Hoeffding bound. Next, we apply Thorem I from Arpit & Bengio (2019) to show that the magnitude of this dot product is preserved in the next layers. First, note that as we assume the data lies in a compact set and the cost-activations are all positive, one can easily see that for each Ii, 22 E X' and for all | [1 L] we have that:
min ZEX
())) < (x)) (x2) & max (1)) (x))
(46)
ZEX
To simplify the proofs this Lemma, we use this fact and instead work with the norm the post- activations and note that the final result the norms can be accordingly applied dot products
23

---

Under review as a conference paper at ICLR 2023
of post-activations of different inputs. For the first layer, we have that (() o(W())) where w ~ (N0,, ) and æ E _ Hence, each f()()); is i..d and distributed as (0, Iall
"j
where ARR is the Rectified Normal Distribution, Hence, using the properties of the Rectified Normal distribution we get that:
-||=| =
B|| (=)]+
(47)
Next, as the Rectified Normal is sub-gaussian distribution, we can apply the Hoeffding bound to see that
(x) - -nuil < ov
2 log
(48)
with probability least 1 - 8 where Jr no random initialization of the weights of the first layer. Hence, combining this and Equation (46)
|=ll and is the standard deviation of (z) 12 over
come up with constants G and G such that for any > 0, 22, G n S (z1) (x2) <
Glin with probability least 1 8. Next, we can adapt Theorem from Arpit Bengio (2019) to see that for post activations of layer E [2 L.]
Pr (1- (1)-=))))1 < (x) s (1 (1+==)=11))) (1+ 2 21-8
(49)
where 8' Σt 2N exp (-n (5 + log (+ N the size of our dataset and any positive small constant. Combining this with the result from the first layer's post-activations see that
(1 e1-
22) s (x)) (1 E1 E1) (mp + 82)
(50)
with probability at least (1 (1 8¹). Hence, for any n T 8 > 0 and (21, X' x X' one can come up with constants G = polyloggn/6)). G = polylog(n/0) for post activations of layer I such that
a n S (z) (x2) G.-²n n
(51)
with probability at least 1-6 8, where 7mm depends on / and 8.
B.3 PSEUDO-NTK's 'S MAXIMUM EIGENVALUE CONVERGES TO ENTK'S MAXIMUM
EIGENVALUE AS WIDTH GROWS
In this subsection, we present a formal proof for Theorem 4.
Proof. Note that, as both PNTK and eNTK are symmetric PSD matrices, their maximum eigenvalues are equal to their spectral norm. Furthermore, the spectral norm of matrix is upper-bounded Frobenius norm. Now, note that according to the triangle inequality, we have
bounded by its
(2)1))2)) = [[((z,, .2) 3 lo + (2) - Ó X 10)
(52)
< 0g) 8 Toll + [(0()), (211)22 Toll
Thus
(2(11,0)] - (0( @ loll < [[((z), x2) - Õ(11, @ loll.
(53)
which according to (43) together with the fact that for any matrix A, Amax (A 8 I) = Amax(A) implies that with probability at least - 8,
Xmase (0(x1, Xuna (0(21,
80 (c(l-1) 1) + c((--1) Vn max (viog
802
802 .22log
(54)
8
24

---

Under review as a conference paper at ICLR 2023
Moreover, mentioned in the proof of Lemma 15, combining the previous inequality with the fact that Amax (0(z1, 22)) 2 (2(n) with high probability shows that there exists 8' and no such that
Amax (0(11, x2)) - Amax (6(21,
((1/V)
(55)
Amax (9(11, x2))
with probability 8' over random initialization for no as desired.
B.4 KERNEL REGRESSION USING PNTK VS KERNEL REGRESSION USING ENTK
In this subsection we provide a formal proof for Theorem 5.
Proof. We start by proving simpler version of theorem, and then show correspondence that expands the result of the simpler proof to the original Theorem. Assuming | = | N (training data), define
((x) - X10(X,) and h(x) = (6(z), 10) 1o)
-1 y. (56)
Note that as the result of kernel regression (without any regularization) does not change with scaling the kernel with fixed scalar, can use weighted version of the kernels mentioned in the previous equation without loss of generality. Accordingly, we define
(Ho(x)x)) y and & = (0)) a =
y. y
(57)
11
Using the fact that M- 1-11--1 --1= -M- M)M and (A@1) A-1 2 we can show that
à a-a=-6(xxx") -1
8 To - ( @ to 00(x ((11,22)))
(58)
Assume
min (Amin (0(X X)), Amin (0(X) Then
||á - all j n
0 lo 0xx,
(59)
n
Plugging into the formula for kernel regression, we get that
h(x)) - h(z) =
1 2 1o à e N)a
(xx
12
(60)
0(xx 8 1o - 1 ,)(,
(), - a)
à n
n
Thus
[1(x)) - -(x)) < - Ô(x, x) 8 lo 1 , (2, x3) |||||||| + -1-=((z,))) x) all
n
1 Ô X) To O(x !|
(61)
x n
Lo(z, X | 0 x) @ lo 1 x)xxx))
25

---

Under review as a conference paper at ICLR 2023
Cole. Actiralioz
Tusb Accovddie
Pall Actiralls
Truacated Nernal
$5.20
Figure 12: Comparing the magnitude of sum of on-diagonal and off-diagonal elements of Og initialization and throughout training, based on 1000 points from CIFAR-10 10. The reported numbers the average of 1000 x 1000 kernels each having shape of 10 x 10. The same subset has then been used to train the NN using SGD.
Artivation
Tath Activadion
Polli Acttratioz
Thuncared Nermal
24
03
20
Figure 13: Evaluating the relative difference of Frobenius norm of OD D) and Ox(D,D) @ at initialization and throughout training, based on 1000 points from CIFAR-10.
Now, note that as for block matrix A of Aig blocks we have that || S EigllAs) it follows that for any matrix valued kernel K
(K(X,)) S Σ
[K(x1,
(62)
Z1.72EX
Using this fact, we can rewrite the bound as
N
1
[lA(x) - h(x))(( s
x
x1)
2 10
O(x, xill
x
n
(63)
N
0
x
- 0 mg)) @ lo - 0 x3)) IP
x2
for some particular 2i, 12, X. Using (43),
with probability at least - 8 that
[M(x)) -(x))| S
8NOa A/nn
max (Niog
80 2og
80"
|il|
1
N - 0 x1)1, (64)
To show the correspondence between h(xx) and flin (x). as in (6), note that
h(x) (Ó(I, 8 10) (6(X, @ 10)
(0(z, @ 10)
(65)
= vec (10),,((,
where ),, vec vy) the result of inverse of the vectorization operation, converting the NO x vector to O N matrix. Thus, h(x)) (x, X') x) "'y where y' is the N O matrix derived from reshaping the NO x 1 vector y. The proof is complete.
B.5 EXTENDING THE PROOFS TO OTHER ARCHITECTURES
In this subsection we elaborate on how one can extend the current proofs to different architectures. We start by providing sketch how the dense weight vectors can be replaced by other layers
26

---

Under review as a conference paper at ICLR 2023
Cau: Actontiin
Tusb Actinadie
Palli Actirallez
Truacated Nernal
400
130
Figure 14: Evaluating the relative difference of Amax of (n(D, and @D D) at initialization and throughout training, based on kernels subset (iD) 1000) of points from CIFAR-1O.
Gkle Actheliee
"ash Accovudion
Ceta Actiraaiiss
Stuacoted Ntemal
0.1
08
Figure 15: Evaluating the relative norm difference of kernel regression outputs using eNTK and pNTK as in Equation (5) and Equation (6) at initialization and throughout training. The kernel regression has been done on 'DI = 1000 training points and | = 500 test points randomly selected from CIFAR-10's train and test sets.
of choice like convolutions. First, note how the linear weights are used in Equation (12). As mentioned in Section 6 of Yang (2020), we can accordingly write the same expansion for other forward computational graphs and derive the corresponding canonincal decomposition for them. In subsection 6.2.1, Yang (2020) provides a concrete example on how one can derive this expansion for a general RNN-like architecture. As the proofs provided in this section depend on the MLP structure only by means of the canonical decomposition, one can extend them to general architecture by deriving the corresponding canonical decomposition of that architecture.
Non-Gaussian Weights: According to the strategy used in the proofs, we need the individual weights be distributed such that the product of two independent scalar weights (as in Equation (12)) remain -ub-exxonential. Hence, any sub-gan uub-gaussian initialization method, such as any bounded initialization (e.g. truncated normal or uniform on an interval) can be used, and the same proof structure would support the same convergence rate, albeit with different constants in convergence (independent of n).
Non-ReLU activations: general, the proofs rely on the ReLU activation through Lemma 17, which gives concentration bound on the absolute value of the dot product of post-activations of each layer of the NN. To use other nonlinearities, we would only need an analagous result for that nonlinearity; the other proofs follow without requiring any other significant change.
Experimental Evaluation: To provide further experimental support for this argument, we have conducted an ablation study the FCN architecture with differrnt nonlincarities and with truncated Gaussian initialization (Figures 12, 13, 14 and 15). As seen in the provided figures, the impact of nonlinearity and initialization method long as they follow the provided setting in Setting A, is marginal.
C MORE DETAILS ON KERNEL. REGRESSION USING PNTK ON FULL CIFAR- 10 DATASET
In this section we provide another figure comparing the accuracy of flin (x) with parameters derived at epoch E E {0, 50, 100, 150, 200} of training the NN with SGD. On the y-axis, the reported number is glin (z) - fx) where f" denotes the final model obtained after training f for 200 epochs. As seen in Figure 16 the architecture of the model has a significant impact on how good the linearization predicts the final accuracy of the fully-trained model. However, as proven in Theorem in conjunction with the linearization approximations provided in Lee et al. (2019), width grows,
27

---

Under review as a conference paper at ICLR 2023
FCM
CoaviNet
BesNot18
WideBesNett
-20
10 1024
40 130
-20
90-4
106
Epach
Epads
Epends
Rpech
Figure 16: Evaluating the difference in test accuracy of kernel regression using pNTK in (6) vs the final model f throughout SGD training on the full CIFAR- 10 dataset. How much worse would be to "give up" on SGD at this point and train flin with the current representation?
this approximation becomes more accurate. One unexplored fact regarding this experiment that fact that lineraization with trained parameters significantly outperforms linearization at initialization, which is intuitive but not rigorously investigated yet.
28


Reviews:
[{"review_id": "lXsS1AMghN", "reviewer": "Reviewer_wjkj", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes an approximation of the neural tangent kernel (NTK) for models with a large output dimension. Specifically, considering a model with output dimension $O$, for a specific pair of data examples, the standard definition of NTK gives an $O\\times O$ matrix. This paper tries to approximate the standard NTK by the identity matrix multiplied with a scalar which is defined as the average of the original NTK. It is proved that as the width goes to infinity, the approximated NTK converges to the standard NTK in Frobenius norm and spectral norm, and the kernel regression solutions are also close. The theoretical results are further supported by experiments.", "summary_of_the_review": "This paper proposes an efficient way for NTK approximation and provides interesting results, but more justification of novelty is needed.", "strength_and_weaknesses": "I think this is an interesting work for NTK approximation, and indeed the proposed method can save the memory usage by a factor of $O^2$. However, I am not very sure about novelty, as detailed below. \n1. As mentioned on top of page 4 of this paper, some recent work already points out that the standard NTK converges to a diagonal matrix as width goes to infinity. This paper further provides formal theorems (Theorems 1, 4 and 5), but the technical challenge and innovation are not discussed. For example, for convergence in Frobenius norm, I imagine it just follows from randomness of the last layer, and we do not need to deal with some independence issue since other matrices and vectors involved in the gradient calculation do not depend on the last layer, which can simplify the proof a lot. Please correct me if I am wrong, but either way I think it is necessary to discuss the technical difficulties to justify innovation.\n2. An interesting result is shown in Figure 7, where the kernel regression outputs of the standard NTK and approximated NTK actually become closer during training. I believe this is a novel discovery, but there is not enough discussion.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "bLxBm8UBKw", "reviewer": "Reviewer_u749", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a fast approximation of the empirical neural tangent kernel. The authors theoretically prove the approximation accuracy of the proposed method, and use experiments to demonstrate its performance.", "summary_of_the_review": "The proposed method is reasonable and the theory seems correct. However, the significance and novelty of this paper need further explanation.", "strength_and_weaknesses": "Strengths:\n\nAlthough I did not thoroughly check all the proofs, the proposed methods and the theory look reasonable. \n\nThe experiment results are well presented and look convincing. \n\nWeaknesses:\n\nThe authors mentioned that the motivation for studying the eNTKs is that it provides a good understanding of a given network’s representation. However, by definition, it seems that whenever one can afford to train a neural network, the cost to calculate the eNTK should also be affordable. Therefore the significance of the proposed method is questionable. \n\nAs the authors mentioned, the fact that a eNTK with $NO \\times NO$ entries can be approximated by an $N \\times N$ kernel matrix is already a well-known result. Therefore the proposed approximation method is not surprising. \n\nMinor comments:\n\nAt the top of page 4, the statement \"in the infinite width limit the NTK becomes a diagonal matrix\" is inaccurate and may be misleading. If I understand correctly, $K(x_i,x_j)$ for $i \\neq j$ is not equal to zero.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "EIT1VZC4Hf", "reviewer": "Reviewer_JUEa", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In the paper, the author provided some theoretical and experimental results on the so-called \"sum of logits\" approximation of empirical neural tangent kernel (eNTK) having multiple output units.\nThe proposed approach consists in approximating the eNTK by a block diagonal matrix as the Kronecker product between some pseudo-NTK approach (pNTK) and $I_O$, the identity matrix of size $O$, if the network has $O$ output units. This saves an order $O^2$ of memory and (up to) order $O^3$ of computation in the evaluation of NTK.\nThe authors showed, in Theorem 1 and 4, respectively, that the relative difference in Frobenius norm and maximum eigenvalue between the proposed approximation and the true eNTK is of order $O(n^{-1/2})$, for deep networks of width $n$ and with Gaussian random weights. This provides theoretical justifications for the proposed approximation.\nFurther numerical experiments were provided to validate the proposed approach at and beyond initialization.", "summary_of_the_review": "I find the contributions of the paper marginally significant or novel. And some clarifications and revisions are needed to help better understand and possibly (re)evaluate the contribution of this work. See below.\n\nDetailed comments:\n* P1: I imagine NNGP stands for neural network Gaussian process?\n* i'th at the bottom of P3 versus ith in P4\n* P4: I get confused when reading \"Before turning to the formal results and experimental evaluation, we give some intuition. First, suppose that ...\": I suppose that here v_i should be the ith row or column, NOT entry, of a linear read-out layer? Since in the sentence that follows \"if the vectors v_i \\sim \\mathcal N(0, \\sigma^2 I_O)\"\n* Theorem 1: can something be said beyond the current setting? Say, beyond ReLU activation? Different initialization scheme, say, beyond Gaussian? Should we understand \\in O(n^{-1/2}) as of order O(n^{-1/2}) as in classical computer science or statistic literature?\n* Remark 3 and Figure 2: the referred \"ratio of information\" here should be explicitly defined and (at least briefly) discussed\n* I believe it is necessary to mention in the statement of Theorem 1 that the result holds ONLY for networks having random weights or networks at initialization: in the current version of the theorem this fact is somehow very implicit and this makes a huge difference.\n* Theorem 4 and its proof: in fact, the result (and proof) in Theorem 4 can be improved to obtain some control on any corresponding pair of eigenvalues of eNTK and its approximation: it in fact follows from  Weyl's inequality (in linear algebra) that max_{i} |\\lambda_i (A) - \\lambda_i (B)| \\leq \\| A - B \\|, with \\lambda_i(A) eigenvalues of A and \\| A \\| the spectral/operator norm of A. Some conclusions on the condition number can be drawn in a similar manner.", "strength_and_weaknesses": "**Strengths**: the paper is in general well written and easy to follow (a few clarifications are needed, though, see my detailed comments below in \"Summary Of The Review.\" The obtained theoretical results are of interest, believed to be true in the literature, but, to the best of my knowledge, rigorously established for the first time. I did not check the proof in great detail, but the line of arguments looks compelling.\n\n**Weaknesses**: the theoretical results obtained in the paper are not very strong and are limited in many aspects (layer width, Gaussian initialization, nonlinearity, etc). Some improvements can be obtained without much effort, see again my detailed comments below.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9Yfz8SUbpf4", "reviewer": "Reviewer_bBN3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the relationship between the \"pseudo-NTK\" and the \"empirical NTK\". They are different from each other only in the multi-class (multi-output) case: let $f^L_1,f^L_2, \\ldots,f^L_O$ denote the $O$ outputs of the network at the last layer $L$. \nIf $x_1$ and $x_2$ are two samples,  $eNTK(x_1,x_2)_{i,j}$   \n\nis the inner product between the gradients of $f^L_i$ at $x_1$ and that of $f^L_j$ at $x_2$. On the other hand, the $pNTK(x_1,x_2)$ is a single number equal to the inner product between the gradients of $1/\\sqrt{O}\\sum_{i=1}^O f^L_i$ at $x_1$ and $x_2$.  It is shown here that the pNTK approximated the diagonal elements of the eNTK, and the off diagonal elements of the eNTK are an order of $1/\\sqrt{n}$ smaller (than the diagonal elements), where $n$ is the width of the network. This allows one to construct a faster approximation to the eNTK in the multi-output case. It is also shown that the approximation stays good even if one solves the kernel regression problem associated to the relevant kernels, and many data experiments confirm the findings (though obviously not in terms of exact convergence rates. Some of the experiments show that the approximation holds even after training, though this is not proved from a theoretical standpoint. \n\nThe proof of the main result works by induction (similarly to other NTK papers), iteratively bounding terms in key equation (15), except that one is not taking expectations but instead leveraging the sub-gaussianity of the relevant variables to bound the activations' absolute values with high probability.  ", "summary_of_the_review": "This is an interesting paper in a very hot topic. I cannot 100 percent vouch for originality but if the results are indeed original, they are very worthy of consideration. There are correctness issues. Nothing fatal, but the paper needs polishing. \n\n\n\n\n=======================minor comments, typos=============\n\n\nMain paper: \n\npage 4: \"casting doubt on the correctness of previous results\". This type of provocative statement should be used parsimoniously. I know the authors end up saying that the main contribution of this work is to remove the doubt but still, I would advise reformulating that statement. \n\n\nPage 4 \" we prefer the sum-of-logits for our networks...\" The use of the word \"prefer\" makes it difficult to see your main point there. The one versus rest analogy probably needs more explanation as well. \n\nIn the experiments (cf. explanation in section 4), the authors seem to say that they provide error bars based on only three samples. Surely this can't be right!\n\nIn remark 3, \"on diagonal\" and \"off diagonal\" should be swapped. \n\n\nAppendix:\nAt the beginning of section B:   $g:\\mathbb{R}^O\\rightarrow \\mathbb{R}^w$ should be $g:\\mathbb{R}^D\\rightarrow \\mathbb{R}^w$\n\nAt the bottom of the same page (13), it would be nice to explain that the vector $1_O$ is not trainable (its gradient doesn't show in the kernel). \n\nIn the second point of \"setting\" on page 15, there is the condition that the layers should be within a constant ratio of each other. This is not formulated correctly: mathematically, the condition as formulated there is satisfied for any network as long as the widths are integers. \n\n\nAt the bottom of page 17, there is  $W^{(2)}_{ia}$  \n\nwhich should be  \n\n\n$W^{(l+1)}_{ia}$\n\n\n(line spacing due to formatting difficulties with tex in openreview)\n\n\nIn the proof of Theorem 17 (which is a reformulation of the main theorem of the paper), the proof consists in two cross references, one of which is \"??\".\n\n\n\n", "strength_and_weaknesses": "Strengths: \n   -This is a very cool topic. I am a little out of touch with the most recent NTK literature but it seems hard to find non-asymptotic results expressed as a function of the width. \n   -There is a good mix of theory and practice and assuming everything is correct, this paper can have an impact both in statistical learning theory and in practical applications. \n\n\nWeaknesses: \n\n  - It is not clearly stated whether or not the main result (Theorem 1, equation (3)) holds only because of the results in this paper or simply because both quantities inside the norm in the numerator converge to the actual (non empirical) NTK at the same rate. I went through the original NTK paper (Jacot et al 18) and it seems like the results are purely asymptotic, but it is still surprising to me that I cannot find a convergence rate (in terms of width) for the eNTK to the NTK. Can the authors provide the relevant references and explain whether their result follows directly from such a non-asymptotic result? \n\n- The proofs are at best not reader friendly and at worse, wrong: there are three main issues as far as I can see (see more details in \"clarity, ...\")\n\n1. The management of the high probability events (the \"deltas\") is very vague. I am quite confident that this can be fixed (possibly by also changing the wording of the main theorems slightly), but since the proofs are not extremely impressive in the first place, it doesn't seem out of place to request a picture perfect first submission for a conference such as ICLR. **Honestly, I think the results are (at least slightly) wrong.** There are a few union bounds arguments missing, which should make a factor of log(L) appear. \n\n2. The results assume a constant width in the proofs, quickly saying that using other works' would allow one to make the results architecturally robust. It would be nice to see those techniques in action. This is especially relevant in the sense that I think there is a hidden assumption here: I cannot see how to fix the proof of the main theorem without assuming that $O\\leq n$ (the number of outputs is less than the width of the network). This definitely restricts the setting and removes some of the typical Xtreme multiclass cases. \n\n3. The explanations in the proofs are rushed and it can be hard to read some of them. This is despite the fact that the main strategy of the proof is actually pretty straightforward. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject