Scaling Neural Tangent Kernels via Sketching and Random Features
Authors: Amir Zandieh
Insu Han
Haim Avron
Neta Shoham
Chaewon Kim
Jinwoo Shin
Venue: NeurIPS 2021 Poster (2021)
Keywords: Neural Tangent Kernel
NTK
Convolutional Neural Tangent Kernel
CNTK
Sketching
Random Features
Kernel Approximation
Kernel Methods

Abstract: The Neural Tangent Kernel (NTK) characterizes the behavior of infinitely-wide neural networks trained under least squares loss by gradient descent. Recent works also report that NTK regression can outperform finitely-wide neural networks trained on small-scale datasets. However, the computational complexity of kernel methods has limited its use in large-scale learning tasks. To accelerate learning with NTK, we design a near input-sparsity time approximation algorithm for NTK, by sketching the polynomial expansions of arc-cosine kernels: our sketch for the convolutional counterpart of NTK (CNTK) can transform any image using a linear runtime in the number of pixels. Furthermore, we prove a spectral approximation guarantee for the NTK matrix, by combining random features (based on leverage score sampling) of the arc-cosine kernels with a sketching algorithm. We benchmark our methods on various large-scale regression and classification tasks and show that a linear regressor trained on our CNTK features matches the accuracy of exact CNTK on CIFAR-10 dataset while achieving 150x speedup.

Scaling Neural Tangent Kernels via Sketching and Random Features
Amir Zandieh Le-Plllnnk-nsstttu für Informatik azandieh@mpi- wpg.ppg.... de
Insu Han" Yale University insu. han@yale edu
Haim Avron
Tel Aviv University
haimavStauex tau. ac.
Neta Shoham
Chaewon Kim KAIST chaewonk@kaist.ac.kk
Jinwoo Shin KAIST
Tel Aviv University
shohamne@gmiil.cmmm com
j jinwoos@kaist ac kr
Abstract
The Neural Tangent Kernel (NTK) characterizes the behavior infinitely-wide neural networks trained under least squares loss by gradient descent. Recent works also report that NTK regression can outperform finitely-wide neural networks trained small-scale datasets. However, the computational complexity of kernel methods has limited its use in large-scale learning tasks. To accelerate learning with NTK, design near input-sparsity time approximation algorithm for NTK, by sketching the polynomial expansions arc-cosine kernels: our sketch for the convolutional counterpart of NTK (CNTK) transform any image using linear runtime in the number of pixels. Furthermore, we prove spectral approximation guarantee for the NTK matrix, by combining random features (based leverage score sampling) of the arc-cosine kernels with a sketching algorithm. We bench- mark our methods on various large-scale regression and classification tasks and show that linear regressor trained on our CNTK features matches the accuracy of exact CNTK on CIFAR-10 dataset while achieving 150x speedup.
1 Introduction
Recent results have shown that over-parameterized Deep Neural Networks (DNNs), generalize surprisingly well. In effort to understand this phenomena, researchers have studied ultra-wide DNNs and shown that in the infinite width limit, a fully connected DNN trained by gradient descent under least-squares loss is equivalent to kernel regression with respect the Neural Tangent Kernel (NTK) [5, 11, 22, 28]. This connection has shed light on DNNs' ability to generalize [10, 34] and optimize (train) their parameters efficiently [3, 4, 16). More recently, Arora al. [5] proved analo- gous equivalence between convolutional DNNs with infinite number of channels and Convolutional NTK (CNTK). Beyond the aforementioned theoretical purposes, several papers have explored the algorithmic use of this kernel. Arora et al. [6] and Geifman et al. [19] showed that NTK based kernel models can outperform trained DNNs (of finite width). Additionally, CNTK kernel regression sets an impressive performance record on CIFAR-10 10 for kernel methods without trainable kernels [5]. The NTK has also been used experimental design [39] and predicting training time [43].
However, the NaKeb approaches encounter the computational bottlenecks of kernel learning. In particular, for dataset of n images , Fr Rddd only writing down the CNTK kernel matrix requires Ω (d' n²) operations [5]. Running regression or PCA on the resulting kernel matrix takes additional cubic time in n, which is infeasible in large-scale setups.
"qqual contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

There is rich literature on kernel approximations for large-scale learning. One of the most popular approaches is the random features method which works by randomly sampling the feature space of the kernel function, originally due to the seminal work of Rahimi and Recht [37]. Another popular
approach which developed in linear sketching literature [41], works by designing sketches that
can be efficiently applied to the feature space of kernel without needing to explicitly form the high dimensional feature space. This approach has been successful at designing efficient subspace embeddings for the polynomial kernel [7. 1J. In this paper, we propose solutions for scaling the NTK and CNTK by building on both of these kernel approximations techniques and designing simply
efficient feature maps that approximate the NTK/CNTK evaluation. Consequently, transform the input dataset to these feature spaces, and then apply fast linear learning methods to approximate the answer of the corresponding nonlinear kernel method efficiently. The performance of such approximate methods is similar or sometimes better than the exact kernel methods due to implicit regularization effects of the approximation algorithms [37, 38, 23}.
1.1 Overview of Our Contributions
One of our results an efficient random features construction for the NTK. Our starting point is the explicit NTK feature map suggested by Bietti and Mairal [9] based on tensor product of the feature maps of arc-cosine kernels. We obtain our random features, by sampling the feature space of arc-cosine kernels [12]. However, the naive construction of the features would incur an exponential cost in the depth of the NTK, due to the tensor product of features generated in consecutive layers. We remedy this issue, by utilizing efficient sketching algorithm for tensor products known as TENsorSREH [1] which can effectively approximate the tensor products of vectors while preserving their inner products. We provide rigorous error analysis of the proposed scheme in Theorem 2.
Our next results are sketching methods for both NTK and CNTK using runtime that is linearly proportional to the sparsity of the input dataset (or number of pixels of images). Our methods rely on the arc-cosine kernels' feature space defined by their Taylor expansion. By careful truncation of their Taylor series, we approximate the are-cosine kernels with bounded-degree polynomial kernels. Because the feature space of polynomial kernel is the tensor product of its input space, its dimensionality is exponential in the degree of the kernel. Fortunately, Ahle et al. [1] have developed linear sketch known as POLYSKETCH that can reduce the dimensionality of high-degree tensor products very efficiently, therefore, we can sketch the resulting polynomial kernels using this technique. We then combine the transformed features from consecutive layers by further sketching their tensor products. In case of CNTK, we have an extra operation which sketches the direct sum of the features of neighbouring pixels each layer that precisely corresponds to the convolution operation in CNNs. We carefully analyze the errors introduced by polynomial approximations and various sketching steps in our algorithms and also bound their runtimes in Theorems and 4.
Furthermore, we improve the arc-cosine random features to spectrally approximate the entire kernel matrix, which is advocated in recent literature for ensuring high approximation quality in downstream tasks [8, 32]. Our construction is based on leverage score sampling, which entertains better convergence bounds [8, 28, 29]. However, computing this distribution is as expensive as solving the kernel methods exactly. We propose a simple distribution that tightly upper bounds the leverage scores of arc-cosine kernels and for further efficiency, use Gibbs sampling to generate random features from our proposed distribution. We provide our spectral approximation guarantee in Theorem 3.
Finally, we empirically benchmark our proposed methods on various classificationreeressson tasks and demonstrate that our methods perform similar to better than exact kernel method with NTK and CNTK while running extremely faster. In particular, we classify CIFAR-10 dataset 150x faster than exact CNTK and the same time achieve higher test accuracy.
1.2 Related Works
There has been a long line of work on the correspondence between DNN and kernel machines [26, 30, 35, 18, 42]. Furthermore, there has been many efforts in understanding variety of NTK properties including optimization [27, 3, 16, 44], generalization [10], loss surface [31], etc.
Novak et al. [35] tried accelerating CNTK computations via Monte Carlo methods by taking the gradient of a randomly initialized CNN with respect to its weights. Although they do not theoretically bound the number of required features, the fully-connected version of this method is analyzed in [5].
2

---

Particularly. for the gradient features to approximate the NTK up to E, the network width needs to be a 4 log 3 ), thus, transforming a single vector IR requires S2( L 26
L log : ga log & nnz(())
operations, which is slower than our Theorem 1 by factor of L / Furthermore, [5] shows that the performance of these random gradients is worse than exact CNTK by large margin, in practice. More recently. [28] proposed leverage score sampling for the NTK, however, their work is primarily theoretical and suggests no practical way of sampling the features. Another line of work on NTK approximation is an explicit feature map construction via tensor product proposed by Bictti and Mairal [9]. These explicit features can have infinite dimension in general and even if one uses a menniolieenio approximation to the features, the computational gain of random features will be lost due to expensive tensor product operations.
A popular line of work on kernel approximation problem is based on the Fourier features method [37], which works well for shift-invariant kernels and with some modifications can embed the Gaussian kernel near optimally [8]. Other random feature constructions have been suggested for variety of kernels, e.g., arc-cosine kernels [12], polynomial kernels [36]. In linear sketching literature, Avron al. [7] proposed a subspace embedding for the polynomial kernel which was recently extended to general dot product kernels [20]. The runtime of this method, while nearly linear in sparsity of the input dataset, scales exponentially in kernel's degree. Recently, Ahle al. [1] improved this exponential dependence to polynomial which enabled them to sketch high-degree polynomial kernels and led to near-optimal embeddings for Gaussian kernel. In fact, this sketching technology constitutes one of the main ingredients of our proposed methods. Additionally, combining sketching with leverage score sampling can improve the runtime of the polynomial kernel embeddings [40].
1.3 Preliminaries: POLYSKETCH and TENSORSRHT Transforms
Notations. We use [n] (1.....n} n). We denote the tensor (a.k.a. Kronecker) product by & and the element-wise (a.k.a. Hadamard) product of two vectors or matrices by O. Although tensor products are multidimensional objects, we often associate x @ y with single dimensional vector (119), .2.h.... Fm y1 2132,... ... 2 x. the pold self-tensoring of Another related operation that we use is the direct sum
For shorthand, we use the notation Sp to denote
terms
of vectors: By:= =- |xT , yy _.. We need notation for sub-tensors of a tensor. For instance, for 3-dimensional tensor Y E IRmxxxd and every E [d], we denote by Y(..)) the X n matrix that is defined as [Y....)). := Yi,,, for i [m].j E [m]. For square matrices A and B, we write A < B if B - A is positive semi semi-definite We also denote ReI.U(x:) = max(z; 0) and consider this element-wise operation when the input is a matrix. We use nnz(x) to denote the number of nonzero entries in z. Given a positive semidefinite matrix K and À > 0, the statistical dimension of K with is defined as sx (K) == tr K [K + -1)- For two functions f and g we denote their twofold composition by f g. defined as f o g(a)= == [g((a)..
The TENSORSRHT is a norn-preserving dimensionality reduction that can be applied to the tensor product of two vectors very quickly [1]. This transformation is generalization of the Subsampled Randomized Hadamard Transform (SRHT) [2] and can be computed in near linear time using the FFT algorithm. The POLY SOOLYSEER SKETCH extends the idea behind TENSORSRHT to high-degree tensor products by recursively sketching pairs of vectors in binary tree structure. This sketch preserves the norm of vectors in d with high probability and can be applied to tensor product vectors very quickly. The following Lemma, summarizes Theorems 1.2 and 1.3 of [1] and is proved in Appendix B.
Lemma (PoLYSKETCH). For every integers d 2 1 and every E, > O, there exists a distribution on random matrices Q" E Rm called degree p POLYSKETCH such that (1) for some m
O ( log" 1) and any y E d Pr [I[QPyll E +1 2 1 ö; (2) for any x Rd ifer E Rd is the standard basis vector along the first coordinate, the total time to compute (x & 2/22 log 1/2 nuv(x), pd log dl ) (3) for any collection of vectors UI,..., Up E Rd, the time to compute Q" (v1 @ ... 8 Up) is bounded by O pm log m + 3/2 dogg )): (4) for any > and any matrix E IR where the statis-
$i for all j = 0, 0..,...,p p is O (pm log" m + min
3

---

1.0
1: (a)
0.8
Degree-s polynomial
(a)
Kel (a)(o)//
16
32
0.4
0.2
1.0
0..5
0.0
1.0
-0.5
a
Figure 1: (Left) Normalized ReLU-NTK function K Kyauu (.) for L = (2, 4, 16, 32} and (Right)
degree-8 polynomial approximation of ReLU-NTK with =
tical dimension of A A is Sx, there exists some m = O Piss
log
n E8
such that,
[r[(1-e) ([) A A+ AT) M A A+ 21 (I)
(1+
2 ReLU Neural Tangent Kernel
Arora et al. [5] showed how to exactly compute the NTK of a L-layer fully-connected network, denoted by θ (y, 2). for any pair of vectors y, z E E using dynamic program in O(d + L) time. However, hard to gain insight into the structure of this kernel using that the dynamic program expression which involves recursive applications of nontrivial expectations. Fortunately, for the important case of ReLU activation this kernel takes an extremely nice and highly structured form.
The NTK in this case can be fully characterized by univariate function K(L) [-1, 1] - IR that we refer to as ReLU-NTK, which is the composition the arc-cosine kernels [12] and was recently derived in [9]. Exploiting this special structure is the key to designing efficient sketching methods and random features for this kernel.
Definition 1 (ReLU-NTK function). For every integer L the L-layer ReLU-NTK function K(L) [-1,1] - IR is defined via following procedure, for every a E [-1,1]:
Let 1o (a) and (a) be 6th and 1st order arc-cosine kernels [12] defined as follows,
1 (* arecos (a:), and ki(a) =
1 (v1 - a² + ( arccos (a)))
ko(a) =
(2)
Let (0) ((an(() (a) = a and for = 1,2, ... L, define {t)) I (a) and 5:(0) ((ala(a)
(a) as follows,
Eeauu(a) (a) = K1 o Ki K (a), and (:()) (a = No DDeen (a))
(3)
composition
Let (a) = 2: (a) = a and for 1, 1, .................... L. define (0 (a) recursively as follows, k (a) = ((6-1((a) (a) 11 (a:) + (1 (a).
The connection between ReLU-NTK function K-(L) and the NTK kernel (L) is formalized bellow,
@ (y,, z) = ||v|||||=||2 K(LL) relu
( for any y, z Rdd
(U.=)
This shows that the NTK is normalized dot-product kernel which be fully characterized by K(LL) [-1,1] R. plotted in Fig. 1. As shown in Fig. 1, this function is smooth and can be tightly approximated with a low-degree polynomial. It is evident that for larger values of L,, K (L) (.) converges knee shape, i.e., has a nearly constant value of roughly 0.3(L. + 1) on the interval [-1, ((]-))] and the interval [1 O(L-1), its value sharply increases to L + at a

---

Algorithm NTKSKETCH for fully-connected ReLU networks
input: vector a R2 network depth L, error and failure parameters E,
2: Choose integers = Ö
1. 2
Ö ( ( Ö
( 1.
0
1,5
and 8*
12
0 (5z log }) appropriately
For p 2L / and p' [9L2 /=²1
polynomials Prela (-) and pland (-) are defined as,
2p+2
(2i) 2i+2 1) 2)
9j a² p(1) (a) = Σ relu
a
2
2 (i1))
==0
T=0
2p'+1
(2i)!
Σ
bj a
a 2i+1
2 (2i+
4: (0)()) ||>|| Q' where Q' R is degree-1 POLYSKETCH per Lemma 1
5: (b(00( (x) V ()(0)))) where V Rexr is instance of SRHT [2]
6: for 1= I. do
7: Let @ E pm be degree-2p POLYSKETCH. Also, let T (2x(2)+)) be an
instance of SRHT. For every = 0, 1,.,,..,2 2p + 2, compute:
2p+2
81 e ( (lole-)(0))
(x) Q2
(x)) - T
v vazz (x) (x) (7)
8: Let Q20'+1 E R xp20'+s be a deg POLYSKETCH. Also, let W R
x2p (2p +2)
be an instance of SRHT. For every 1 = 1, 2p' 1, compute:
2p'+1
(x)+ + Q20'+1 (1o(E-)(())
Ø1 @ e1
( (x) W
(br (8)
9: Let Q² E Raxa be degree-2 POLTSKE Also, let R E ( be an SRHT. Compute:
(x) + R. (Q² (vlte) @ (x) + の (z))
(9)
10: Let G E I8* xs be matrix of i.i.d entries with distribution ((0, ) Compute:
V (2) (x) + || G. (6(l()))
(10)
II: return T (1)
Sketching and Random Features for NTK
The main results of this section are efficient oblivious sketching as well random features for the fully-connected NTK. As shown in Definition and Eq. (5), the NTK e(() is constructed by recursive composition of arc-cosine kernels x1 (.).) and ko(.). So, to design efficient sketches for the NTK we crucially need efficient methods for approximating these functions. Generally, there are two main approaches to approximating these functions; one is random features sampling and the other is truncated Taylor series expansion coupled with fast sketching. We design algorithms by exploiting both of these techniques.
3.1
NTK Sketch
Our main tool approximating the arc-cosine kernels with low-degree polynomials, and then applying POLYSK to the resulting polynomial kernels. The features for multi-layer NTK are the recursive tensor product of arc-cosine sketches at consecutive layers, which turn can be sketched efficiently using SKETCH. We present our oblivious sketch in Algorithm 1.
5

---

Now we present our main theorem on NTKSKETCH algorithm as follows.
Theorem 1. For every integers d 2 1 and 2 2, and any E, > O, let Out R x R² + R be the L-layer NTK with ReLU activation as per Definition I and Eq. (5). Then there exists randomized map v Vlll Rdd - TRd for some 8" = =0 ( log }) such that the following invariants hold,
1. For any vectors y, E ER: Pr [[ (V (v), v wild (2)) - - (v, 2)
Se (y, 2) 2 1 - 8.
2. For every vecor RR, the time to compute v Vull (s) O
log 1/2
log to2 nnz(())
For proof, see Appendix C. One can observe that the runtime of our NTKSKETCH is faster than the gradient features of an ultra-wide random DNN, studied by Arora et al. [5], by a factor of L L²/E".
3.2 NTK Random Features
The main difference between our random features construction and NTKSKETCH is the use of random features for approximating arc-cosine kernels Ki and X₁ in Eq. (2). For any x E R2, we denote
(x)) Step [101 [Iii..... Wmmll IT s). (x)
ReLU ([u(,..., wmmn] a),
(11)
where W1,.... Wing: w/: ........ W/ my E Rdd are i...d. samples from NO Ia). Cho and Saul [12] showed Fill and ( (v), 2 (+))) = || - Ki
that (V)( (o(=))] = KO Milla
(v.e)
TTlalll The
(M-2)
feature map for multi-layer NTK can be obtained by recursive tensoring of random feature maps for arc-cosine kernels at each layer of the network. However, one major drawback of such explicit tensoring is that the number of features, and thus the runtime, will be exponential in depth L. In order to make the feature map more compact, we utilize a degree-2 POLYSKETCH Q² to reduce the dimension of the tensor products at each layer and get rid of exponential dependence L. We present the performance guarantee of our random features, defined in Algorithm 2, in Theorem 2.
Theorem 2. Given y, = E RRd. and L 2 let @lLi) the L-layer fully-connected ReLU NTK. For (+t log 4) such that,
E, 8 > 0, there exist mo = O 1 log b) mm = O [ log 4 m O
Pr (1
1 (L)
rf (v), rt (z) - (L)
(L
nek (y, 2) S e o (v. 2) 2 1 0
(12)
where (LLL rf (v), (L() rf (2) E mmm. +ms are the outputs of Algorithm 2, using the same randomness.
The proof of Theorem 2 is provided in Appendix D. Arora al. [5] proved that the gradient of randomly initialized ReLU network with finite width can approximate the NTK, but their feature di-
mension should Ω क5
log? 1/2 + _o+ llog & -d) which is larger than ours by a factor of y log 4. In Section 5, also empirically show that Algorithm 2 requires far fewer features than random gradients.
3.3 Spectral Approximation for NTK via Leverage Scores Sampling
Although the above NTK approximations estimate the kernel function itself, it is still questionable how it affects the performance of downstream tasks. Several works on kernel approximation adopt
spectral approximation bound with regularization À > 0 and approximation factor 0, that is,
(1-E)(K'))
(L) ntk + XI) = (1) T y (L) + XI = (1 +
K
(L) ntk + AI),
(13)
where vr(4) : [v)()
(xr.)) and [K atk k.j = (z., The spectral bound can provide rigorous guarantees for downstream applications including kernel ridge regression [8], clustering and PCA [32]. We first provide spectral bounds for arc-cosine kernels, then we present our spectral approximation bound for two-layer ReL.U networks, which is the first in the literature.
1Ö(+) suppresses poly 4 factors.

---

Algorithm 2 Random Features for ReLU NTK via POLYSKETCH
1: imput: vector Rd, network depth I,, feature dimensions mg, mi, and ms
2: oir (x) f x 2/||=||2.
Pee (x)
t x/ |
3: for 8= 1 to
L do
4:
4o
Pet
(t-1) (3)) where pg is defines per Eq. (11) with my features
5: Ozt (z)
11
Pet
(t-1) (z) where $1 defines as per Eq. (11) with my features
6: Draw a degree-2 POLY Qa that maps R", and compute:
(0 (z) Ozt (2) Lt
(z) + Q²
Pot (z) Vet (x))
(L) 7: return V (x) + ||1|| V (x)
To guarantee that the arc-cosine random features in Eq. (11) provide spectral approximation, we will use the leverage score sampling framework of [8, 28]. We reduce the variance of random features by performing importance sampling. The challenge is to find proper modified distribution that certainly reduces the variance. It turns out that the original 0th order arc-cosine random features has a small enough variance. More precisely, let Kú be the 0th order arc-cosine kernel matrix, i.e., (m),m;)) [T m2 and Po == [fo(21),.... Po(zn)), where bo(x) defined in Eq. (11).
[Kola.- = Ko
If the number of features mo 2 3 X log 1682 ). then
[r[(1 ( XI v D Po + XI 3 (1++) + XD) 2 1 -
(14)
Next, we consider spectral approximation of the 1st order are-cosine kernel. Unlike the previous case, modifications of the sampling distribution are required. Specifically, for any x E R let
(1
क 1 -
2d m
ReLU
([ww)
W1
Wmi
x
(15)
l1
where wi....; Wmng Rd are i.i.d. samples from p(w) = (2ल]य/2a 1| exp (-2 |v lm For this modified features, let X e Rdxa be the dataset, K₁ be the 1st order arc-cosine kernel matrix. and o = (Z1). (xw)) If the number of
i... [Klk, = || 11 | ayll
(z..Zz)
features my 2 8 3 min (rank(22 (rank(X))
| 11: log (168a) ), then
Pr (11 XI) + + + T 3 +) + XI)] 2 -
(16)
The details provided in Appendix 1 and Appendix F.2. We now ready state spectral approximation bound for our modified random features.
Theorem 3. Given dataset X E IRdx with (Xaas,llaa S 1 for every i E [n], let Kat Ko. Ki
be kernel matrices for two-layer ReLU NTK and are-cosine kernels of tth and 1st order, respec- tively. For any > O, suppose 88 is the statisticul dimension of Kay- Modify Algorithm by
replacing (2(+) in line 5 with (-)- defined in Eq. (15). For any E, > 0, let y (L) E +m.m++m))x
be the output matrix of this algorithm with L = 1. There exist mo = O x log )), ,m =
O
1X12 2 min (rank(X)) A
log +)) mm = 0 1+x log n such that,
Pr (1 (1-5) (Kack + AΙ)
v of
+ 1 (1+ Keek + AA) 2 21-6. 8. of
(L)
(17)
For proof see Appendix E.3. To generalize the current proof technique to deeper networks, one needs monotone property of arc-cosine kernels, i.e., x1 (X) I (Y) for X Y. However, this property does not hold in general and we leave the extension deeper networks to future work.
7

---

0130
0.31
16:00
16:03
lest BIcuURScy
0.50
0.33
GRastty CNTKSuurez Eassc CNTX
0.75
NTKEE
0.5
NTKRF
GRANR?
Easd NIK
Ecoct NTK
CNTRSSRTTH
121
Seature dimarsion
wal clock time (sod)
teature dimansion
running time (sec)
(a) MNIST with Figure 2: Test accuracy of: (a) approximate NTK methods (GRADRF, NTKSKETCH and NTKRF) on MNIST and (b) approximate CNTK methods (GRADRF and CNTKSKETCH) on CIFAR-10.
(b) CIFAR-I0 with CNTK
4 Sketching Convolutional Neural Tangent Kernel
In this section, we design and analyze efficient sketching method for the Convolutional Neural Tangent Kernel (CNTK). We focus mainly on CNTK with Global Average Pooling (GAP), which exhibits superior empirical performance compared to vanilla CNTK with no pooling [5], however, our techniques can be applied to the vanilla version, as well. Using the DP of Arora et al. [5], the number of operations needed for exact computation of the depth-L CNTK value θ (L) (v. 2) for images y. 2 E pdxd is S2 (d* L), which is extremely slow particularly due to its quadratic dependence on the number of pixels of input images d² Fortunately, we are able to show that the CNTK for the important case of ReLU activation is a highly structured object that can be fully characterized in terms of tensoring and composition of arc-cosine kernels, and exploiting this special structure is key to designing efficient sketching methods for the CNTK. Unlike the fully-connected NTK, CNTK is not simple dot-product kernel function like Eq. (5). The key reason being that CNTK works by partitioning its input images into patches and locally transforming the patches at each layer, as opposed to the NTK which operates on the entire input vectors. We present our derivation of the ReLU CNTK function and its main propertics in Appendix
Similar to NTKSKETCH our method relies on approximating the arc-cosine kernels with low-degree polynomials via Taylor expansion, and then applying POLYSKETCH the resulting polynomial kernels. Our sketch computes the features for each pixel of the input image, by tensor product of arc- cosine sketches at consecutive layers, which in turn can be sketched efficiently using POLYSKETCH Additionally, the features of pixels that lie in the same patch get locally combined at each layer via direct sum operation. This precisely corresponds to the convolution operation in neural networks. We present our CNTKSKETCH algorithm in Appendix G and give its performance guarantee in the following theorem.
Theorem 4. For every positive integers di, d2, and L 2, and every E, 6 >> if we let O catz Rd Xd2xc x Rdd xdzxe + R be the L-layer CNTK with ReLU activation and GAP given in [5], then RRdixx2xx R for some s* = = { log })) such that:
(L))
there exist randomized map map vi cntx
For any images y, z E &Rdixx2xxe
Pr (WLL (v), T (())(())
(il.aaaa. (y, (v,=)) -
s (.o (u. (v) 2 d.
2. For every image a E Bdd xd time to compute v w([) (x) is O 20.7
Lt
(d1d2)
log dideb
The proof is in Appendix G. Runtime of our CNTKSKETCH is only linear in the number of image pixels di d2, which in stark contrast to quadratic scaling of the exact CNTK computation [5].
Experiments
In this section, we empirically show that rumning least squares regression on the features generated by our methods extremely fast and effective for learning with NTK and CNTK kernel machines. We run experiments on a system with an Intel E5-2630 CPU with 256 GB RAM and a single GeForce RTX 2080 GPUs with 12 GB RAM. Codes are available at https: //gi thub. com/ insuhan/ ntk- sketch-rf.
8

---

Table 1: Test accuracy and runtime to solve CNTK regression and its approximations on CIFAR-10. (*) means that the result is copied from Arora et [5].
CNTKSKETCH
(ours)
GRADRE
Exact CNTK
CNN
Feature dimension
4,096
8,192
16,384
9,328
17,040
42,816
Test accuracy (%)
67.58
70.46
72.06
62.49
62.57
65.21
70.47'
63.81
Time (8)
780
1,870
5,160
300
360
580
1.000.000
Table 2: MSF and runtime on lcage-scale UCI datasets. We measure the entire time to solve kernel ridge regression.
(-) means Out-of-Memory error.
Songs
WorkLoads
CT
Protein
# data points (n)
467,315
179,585
53,500
39,617
MSE
Time (s)
MSE
Time (8)
MSE
Time (s)
MSE
Time (8)
RBF Kernel
35.37
59.23
18.96
46.45
RFF
109.50
231
4.034 x 104
53.0
48.20
15.2
19.72
12.1
NTK
30.52
72.10
20.24
76.93
NTKRF (ours)
94.27
95
3.554 104
35.7
46.91
2.12
20.51
4.3
NTKSKETCH (ours)
92.83
36
3.538 x 104
27.5
46.52
18.8
21.19
14.91
5.1 NTK Classification on MNIST
We first benchmark our proposed NTK approximation algorithms on MNIST [25] dataset and compare against gradient-based NTK random features [5] (GRADRF) as a baseline method. To apply our methods and GRADRF into classification task, encode class labels into one-hot vectors with zero-mean and solve the ridge regression problem. We search the ridge parameter with a random subset of training set and choose the one that achieves the best validation accuracy. We use the ReLU network with depth L = 1. In Fig. 2a, we observe that our random features (NTKRF) achieves the best test accuracy. The NTKSKETCH narrowly follows the performance NTKRF and the Grad-RF is the worst method which confirms the observations of Arora et al. [5], i.e., gradient of finite width network degrades practical performances.
Remark (Optimizing NTKSKETCH for Deeper Nets). As shown in Eq. (5), the NTK normal- ized dot-product kernel characterized by the function (o)). This function can be easily computed using (I.) operations any desired a E [-,,1], therefore, we can efficiently fit polynomial to this function using numerical methods (for instance, it shown in Fig. that a degree-8 polynomial
tightly approximate the depth-3 ReLU-NTK function K{a) Then, resulting polynomial kernel using POLYSKETCH as was previously done for Gaussian and general dot-pne dot-product kernels [1, 40]. Therefore, can accelerate our NTKSKETCH for deeper networks (T. > 2), using this heuristic.
efficiently sketch the
5.2
CNTK Classification on CIFAR-10
Next test our CNTKSKETCH on CIFAR-10 dataset [24]. We choose convolutional network
of depth = 3 and compare CNTKSKETCH and GRADRF for various feature dimensions. We borrow results of both CNTK and CNN from Arora et al. [5]. The results are provided in Fig. 2b and Table 1. Somewhat surprisingly, CNTKSKETCH even performs better than the exact CNTK regression by achieving 72. 72.06% when feature dimension is set 16,384. The likely explanation that CNTKŠKETCH takes advantages of implicit regularization effects of approximate feature map and powerful expressiveness of the CNTK. Moreover, computing the CNTK matrix takes over 250 hours (12 days) under our setting which is at least 150x slower than our CNTKSKETCH.
5.3 Regression on Large-scale UCI Datasets
We also demonstrate the computational efficiency of our NTKSKETCH and NTKRF using 4 large- scale UCI regression datasets [17] by comparing against exact NTK, RBF as well as Random Fourier Features (RFF). For our methods and RFF, we fix the output dimension to m 8,192 for all datasets. In Table 2, we report the runtime compute feature map or kernel matrix and evaluate the averaged mean squared errors (MSE) on the test set via 4-fold cross validation. The exact kernel methods face

---

Out-of Out-of-Memory eтror on larger datasets. The proposed NTK approximations are significantly faster than the exact NTK, e.g., NTKRF shows up 30x speedup under CT dataset. We also verify that, except for Protein dataset, our methods outperform RFF.
6 Discussion and Conclusion
In this work, we propose efficient low-rank feature maps for the NTK and CNTK kernel matrices based on both sketching and random features. Computing NTK have been raised severe computational problems when they apply practical applications. Our methods runs remarkably faster than the NTK with performance improvement.
Potential negative societal impact. This is technical work proposing provable algorithms which stand alone independently of data, c.g., do not learn any private information of input data. We think there is particular potential negative societal impact due
work.
Limitations. This paper only considers fully-comnected and convolutional neural networks, and our ideas are not directly applicable to scale up NTK of other deep networks, e.g., transformers [21].
Acknowledgments and Disclosure of Funding
Amir Zandieh was partially supported by the Swiss NSF grant No. SDELP2_1951444 Haim Avron and Neta Shoham were partially supported by BSF grant 201 7698 and ISF grant 1272/17. Jinwoo Shin was partially supported by the Engineering Research Center Program through the National Research Foundation of Korea (NRF) funded by the Korean Government MSIT (NRF ASRRAA5A159922 and Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) 000_200100_000000 Artificial Intelligence Graduate School Program (KAIST)).
References
[1] Thomas D Ahle, Michael Kapralov, Jakob BT Knudsen, Rasmus Pagh, Ameya Velingker, David P Woodruff, and Amir Zandieh. Oblivious sketching high-degree polynomial kerels.. In Symposium on Discrete Algorithms (SODA), 2020.
[2] Nir Ailon and Bernard Chazelle. The fast Johnson Lindenstrauss transform and approximate nearest neighbors. SIAM Journal on computing. 2009.
[3] Zeyuan AZu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-paramettrrizaion.. In International Conference on Machine Learning (ICML), 2019.
[4] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference Machine Learning (ICML), 2019.
[5] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with infinitely wide neural net. In Neural Information Processing Systems (NeurIPS), 2019.
[6] 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. In International Conference on Learning Representations (ICLR), 2019.
[7] Haim Avron, Huy Nguyen, and David Woodruff. Subspace embeddings for the polynomial kernel. In Neural Information Processing Systems (NeurIPS), 2014.
[8] Haim Avron, Michael Kapralov, Cameron Musco, Christopher Musco, Ameya Velingker, and Amir Zandich. Random Fourier features for kernel ridge regression: Approximation bounds and statistical guarantees. In International Conference on Machine Learning (ICML), 2017.
[9] Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kernels. In Neural Information Processing Systems (NeurIPS), 2019.
10

---

[10] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Neural Information Processing Systems (NeurIPS), 2019.
[11] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. In Neural Information Processing Systems (NeurIPS), 2019.
[12] Youngmin Cho and Lawrence Saul. Kernel methods for deep learning. In Neural Information Processing Systems (NeurlPS), 2009.
[13] Michael B Cohen, Jelani Nelson, and David Woodruff. Optimal Approximate Matrix Prod- uct in Terms of Stable Rank. In International Colloquium Automata, Languages, and Programming (JCALP),
[14] Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and dual view on expressivity. In Neural Information Processing Systems (NeurIPS), 2016.
[15] Sanjoy Dasgupta and Anupam Gupta. An elementary proof of a theorem of Johnson and Lindenstrauss. Random Structures & Algorithms.
[16] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations (ICLR), 2019.
[17] Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URLhttp::/arcchiv... ics. uci. edu/ml.
[18] Adrià Marriga-Alonso, Carl Edward Rasmussen, and Laurence Aitchison. Deep Convolu- tional Networks as shallow Gaussian Processes. In International Conference Learning Representations (ICLR), 2018.
[19] Amnon Geifman, Abhay Yadav, Yoni Kasten, Meirav Galun, David Jacobs, and Ronen Basri. On the similarity between the laplace and neural tangent kernels. In Neural Information Processing Systems (NeurIPS), 2020.
[20] Insu Han, Haim Avron, and Jinwoo Shin. Polynomial Tensor Sketch for Element-wise Function of Low-Rank Matrix. In International Conference on Machine Learning (ICML), 2020.
[21] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: NNGP and NTK for deep attention networks. In International Conference on Machine Learning (ICML), 2020.
[22] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Neural Information Processing Systems (NeurIPS), 2018.
[23] Arthur Jacot, Berfin Simsek, Francesco Spadaro, Clément Hongler, and Franck Gabriel. Implicit regularization of random feature models. In International Conference Machine Learning (ICML), 2020.
[24] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
[25] Yann LeCun, Corinna Cortes, and CJ Burges. MNIST handwritten digit database. 2010.
[26] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep Neural Networks as Gaussian Processes. In International Conference on Learning Representations (ICLR), 2018
[27] 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 Neural Information Processing Systems (NeurIPS), 2019.
[28] Jason Lee, Ruoqi Shen, Zhang Song, Mendi Wang, and Zheng Yu. Generalized Leverage Score Sampling for Neural Networks. In Neural Information Processing Systems (NeurIPS), 2020.
11

---

[29] Zhu Li, Con-FFancois Ton, Dino Oglic, and Dino Sejdinovic. Towards a Unified Analysis of Random Fourier Features. Journal of Machine Learning Research (JMLR), 2021.
[30] Alexander de Matthews, Jiri Hron, Mark Rowland, Richard Turner, and Zoubin Ghahra- mani. Gaussian Process Behaviour in Wide Deep Neural Networks. In International Conference on Learning Representations (ICLR), 2018.
[31] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 2018.
[32] Cameron Musco and Christopher Musco. Recursive Sampling for the Nyström Method. In Neural Information Processing Systems (NeurIPS), 2017.
[33] Jelani Nelson and Huy L Nguyên. OSNAP: Faster numerical linear algebra algorithms via sparser subspace embeddings. In Foundations of Computer Science (FOCS), 2013.
[34] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias:
On the role of implicit regularization in deep learning. In International Conference Learning Representations (ICLR), 2019.
[35] Roman Novak, Lechao Xiao, Yasaman Bahri, Jaehoon Lee, Greg Yang, Jiri Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-dickstein. Bayesian Deep Convolutional Net- works with Many Channels are Gaussian Processes. In International Conference on Learning Representations (ICLR), 2018.
[36] Jeffrey Pennington, Felix Xinnan X Yu, and Sanjiv Kumar. Spherical random features for polynomial kemels. In Neural Information Processing Systems (NeurIPS), 2015.
[37] Ali Rahimi and Benjamin Recht. Random Features for Large-Scale Kernel Machines. In Neural Information Processing Systems (NeurIPS), 2007.
[38] Alessandro Rudi and Lorenzo Rosasco. Generalization properties learning with random
features. In Neural Information Processing Systems (NeurIPS), 2017.
[39] Neta Shoham and Haim Avron. Experimental Design for Overparameterized Learning with Application to Single Shot Deep Active Learning. In arXiv preprint arXiv:2009. /2820, 2020.
[40] David P Woodruff and Amir Zandieh. Near Input Sparsity Time Kernel Embeddings via Adaptive Sampling. International Conference on Machine Learning (ICML), 2020.
[41] David P Woodruff al. Sketching as Tool for Numerical Linear Algebra. Foundations and Trendy@ in Theoretical Computer Science, 2014.
[42] Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian pro- cess behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint 19X 2019.
[43] Luca Zancato, Alessandro Achille, Avinash Ravichandran, Rahul Bhotika, and Stefano Soatto. Predicting training time without training. In Neural Information Processing Systems (NeurIPS), 2020.
[44] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks. arXiv preprint arXiv:1811.08888, 2018.
12


Reviews:
[{"review_id": "oHiCnuVQ7W4", "reviewer": "Reviewer_M4ry", "summary": "This paper proposes a near input-sparsity time approximation method to scale up the NTK based on sketching and random feature. Besides, the proposed method can be extended for efficient CNTK computation with linear running time. The theoretical analysis and experimental results show the effectiveness of the proposed method. ", "questions": "", "limitations": "", "rating": 7, "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": "Strength:\n\n1. This paper proposed an interesting and useful method to accelerate learning with NTK. It also provides nontrivial theoretical analysis of the proposed method. \n2. The proposed method achieves efficient NTK and CNTK computation based on random features. Their experiment results show that the proposed feature mapping can greatly reduce the running time for NTK and CNTK with performance improvement.\n\nWeakness:\n\n- The writing of this paper is not clear enough.  For example, in Algorithm 1,  $P_{relu}^{(p)}(\\cdot)$ appears with few explanations within the main content.  Although I can find more information in the supplementary material. \n- It seems there is a typo in Theorem 2.  In **Theorem 2**, the probability might be $\\leq \\delta$.\n\nQuestions:\n\n- In Table 2, the running time for NTKRF is much smaller than NTKSKETCH in the latter two datasets while the opposite results can be observed in the first two datasets. Are there some insights into this? Are there any criteria for choosing between these two methods?\n\n- According to the results in (Arora et al.), the result of exact CNTK and CNN is based on setting the depth to 3. If the depth becomes 4, exact CNTK and CNN achieve the accuracy of 75.93 and 80.94, which greatly outperform the proposed method regarding accuracy. What are the results f the proposed method uses the larger feature dimension?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "eSOTt2QQALa", "reviewer": "Reviewer_98MV", "summary": "The aim of this work is to design approximation algorithms for Neural Tangent Kernel (NTK) learning to reduce the computational cost. To approximate the ReLU NTK, one approach is to approximate the arc-cosine kernel with low-degree polynomials and further apply the norm-preserving dimensionality reduction on the resulting polynomial kernels. The other approach is based on efficiently reducing the tensor product dimension for random features. To approximate the Convolutional NTK (CNTK), a similar sketching-based approach is proposed with additional techniques to do local combinations of same-patch features. The proposed approaches are measured empirically on MNIST, CIFAR-10, and large-scale UCI datasets.", "questions": "", "limitations": "", "rating": 7, "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": "[pros]\nThis work is significant given the growing interest in NTK-based models and the need to run NTK on large-scale datasets. The proposed random-feature-based approximation and sketching-based approximation for NTK, as well as the sketching approach for CNTK, all come with theoretical guarantees on approximation and thorough empirical evaluations on real-world datasets.\n\n[cons]\n- while both the sketching technique and the random feature technique for NTK have similar approximation guarantees, and also very similar empirical performance in Fig. 2 and Table 2, I would appreciate more qualitative analysis and discussions on when one technique would be more preferable than the others.\n- I wonder why using a shallow ReLU network with depth one in experiments presented in Fig. 2 since the theoretical analysis (Thm.1 and 2) are based on assumption L >= 2. It would be more interesting to see the empirical evaluations on cases when there are theoretical guarantees.\n- a natural question arises while reading this work: is there a random-feature based approximation for CNTK as a counterpart of what presented in Sec. 3.2? Since it's not mentioned in this work, I wonder if there are some difference between NTK and CNTK that prevents the random-feature based approach to generalize.\n- Some suggestions on improving clarity:\n\t- For the Eq. 12 in Thm. 2, should the \\geq be \\leq? I believe it's a typo.\n\t- to be more self-contained: to include the formal definition of NTK and CNTK.\n\t- in the first sentence of Sec. 2, the notation \\Theta and L are used before definition.\n\t- the second invariant in Thm. 1 use the notation nnz without definition.\n\t- The \\notin notation in the inequality in Thm. 4 is confusing since the term (1 +- epsilon)*\\Theta is clearly not a set.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "There is no ethical issue with this paper.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "eRBlk_5dYQ5", "reviewer": "Reviewer_hkr8", "summary": "Neural Tangent Kernel (NTK) is a recently developed theoretical tools to analyze the optimization of deep neural network under grandient descent. In particular, the infinitely wide neural network has been shown to be equivalent to kernel regression with NTK. However, a key issue of the NTK in practice is the scalability. This paper aims to tackle this problem by approximating the NTK with sketching and random features. Specifically, the authors propose the following two different approximation algorithms. \n\nThe first algorithm proposes to approximate the arc-cosine kernels with low degree plynomials combined with polynomial sketching in order to further reduce the dimensionality of the high-degree tensor products. This algorithm is later applied to CNTK with some modification to cope with the GAP. The authors also provide a detailed theoretical analysis on the approximation error of the NTK.\n\nThe second algorithm utilizes the random Fourier feature method to approximate the arc-cosine kernels. Since the number of parameters scale exponentially with the number of layers, the authors further employ the polysketch to reduce the feature dimensionality. Theoretical analysis of the approximation error is also provided.\n\nFInally, the authors demonstrate the efficacy of the proposed algorithms using extensive experiments on MNIST, CIFAR-10 and some UCI dataset. Experimental results shows that the proposed algorithm yields significant time savings with close to NTK prediction accuracy. ", "questions": "", "limitations": "", "rating": 7, "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": "Clarity: The paper is clearly structured and well-written, though the authors are encouraged to carefully proof-read the manuscript to avoid some typos.\n\nOriginality: The nolvety of the paper is to combine the polynomial sketching with either Taylor expansion approximation or random Fourier feature approximation of the arc-cosine kernels. Both sketching and random Fourier feature approximation are well studied in literature in terms of their theoretical properties. This allows the authors from the current paper to provide detailed theoretical analysis of their proposed algorithms.\n\nQuality \\& Significance: This paper addresses an important issue of NTK learning by combining well studied methods in literature. In my opinion, the porposed algorithms are an important contribution to the literature.\n\nMinor Comments:\nExcept the spectral approximation via leverage scores sampling, all the theoretical analyses focus on the kernel approximation error, which does not directly translate to the prediction accuracy. I am wondering if some analysis of the learning risk is possible.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "15"}, "extra_text": {}}, {"review_id": "W_ey0xYOoi-", "reviewer": "Reviewer_Rtas", "summary": "In this paper, a sketching algorithm for the Neural Tangent Kernel (NTK) is proposed to accelerate its computation. Besides, a proof of the spectral approximation guarantee for the NTK matrix is provided, which combines the random features of the arc-cosine kernels with a sketching algorithm. Empirical studies on the MNIST and CIFAR-10 image classification task show that the linear regressor trained the sketch CNTK features matches the accuracy of the exact CNTK while being 150x faster.", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "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": "The recent progress on the Neural Tangent Kernel (NTK) builds a bridge between the neural network and kernel learning, making it worthwhile to investigate efforts to study the property of NTK. The NTK and the traditional kernel methods suffers the same bottleneck, i.e., the computational complexity, which naturally inspires this paper to adopt the sketching ideas to accelerate the kernel computation. \n\nAlthough I did not go through all the equations of the paper, the success in the sketching idea in the traditional kernel learning makes me believe it can transfer to the NTK setting, and the experimental section confirms this point. In fact, the speedup of the proposed methods shown in Table 1 and 2 is quite cheerful, making the training of CNTK more practical than before.  \n\nOne question: is it possible to apply the idea in the following paper to make it even faster?\n\n[1] Fastfood — Approximating Kernel Expansions in Loglinear Time. Quoc Le, Tamas Sarlos and Alex Smola. ICML 2013.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "They indeed discussed the limitation of the methods. For example, they cannot be applied to transformers right now. \n\nThere would not be negative societal impact of the work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "CNIewbZG9xf", "reviewer": "Reviewer_Mq45", "summary": "The submission proposed an low-rank feature maps method to accelerate learning with the Neural Tangent Kernel (NTK). They utlized the sketching and random features to realize the approximation. A theoretical analysis from the spectral pespertive has been applied. Finally, experiments on large-scale regression and classification tasks show the time-efficent of the proposed method.", "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": "Originality: Approximating the NTK for the large-scale task is somewhat novel to me. The related works are adequately cited and it is clear how this work differs from the previous contribution according to the presentation by authors.\n\nQuality: The submission is technically sound and all claims including theoretical analysis and experimental results are well supported. However, I have some concerns regarding the weaknesses of this work:\n\n(1) This work only considers arc-cosine kernel (corresponding to ReLU activation), can you extend to more kernels (activations)?\n(2) This work only considers NTK for fully connected networks and CNTK, which you have mentioned in the main text.\n(3) The classification and regression tasked used in this work is somewhat fundamental, I would expect more specific but significant application.\n\nClarity: The submission is clearly written and well organized.\n\nSignificance: Since the (C)NTK regression usually performs worse than its counterpart neural network on computer vision tasks [1,2,3]. Thus the significance of accelerating NTK learning for computer vision tasks is somewhat weak to me and relevant to researchers in subareas only. \n\n[1] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In NeurIPS. arXiv preprint arXiv:1904.11955.\n\n[2] Ghorbani, Behrooz, et al. \"Limitations of lazy training of two-layers neural networks.\" arXiv preprint arXiv:1906.08899 (2019).\n\n[3] Lee, Jaehoon, et al. \"Finite versus infinite neural networks: an empirical study.\" arXiv preprint arXiv:2007.15801 (2020).\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": " The authors adequately addressed the limitations and potential negative societal impact of their work", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper seems to give a substantially more practical method for using neural tangent kernels in practice, along with theoretical guarantees. Some issues of clarity and various small improvements to make to the paper came up in the discussion; make sure to incorporate those into the final version of the paper. I'd also like to encourage the authors to work towards getting their algorithm implemented in the https://github.com/google/neural-tangents package (if they haven't already), which would make practical uptake by users far more likely.
Decision: Accept (Poster)