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.aac
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 31 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 Rmxdd" 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,, =) = ||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
9j a² p(1) (a) = Σ relu
a
2
2 (i1))
1) 2)
==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] s),
(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 (+))) = || - Ks
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:
(2) (0 (z) Ozt (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
GRastt CNTKSuurez Eossc 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 lcrge-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": "zgj66AqESXj", "reviewer": "Reviewer_tQEn", "summary": "In this manuscript the authors propose and analyze a sketching scheme that allows for more efficient computations involving so-called neural tangent kernels and their convolutional counterparts. By leveraging recent work on sketching tensors and polynomial kernels from [1] the authors provide algorithms for depth L networks with provable guarantees on the quality of the feature embedding (provided the sketch size is large enough). To complement these theoretical developments the authors provide spectral approximation bounds for a variant of their method in a simplified setting (L = 1). Lastly, some numerical results are provided to illustrate the potential efficacy of the method in practice.", "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": "Overall, this manuscript has nice results, though as written it feels like the big step/insight was in ref. [1] where the sketches are constructed. The main contribution here is then showing/observing the requisite structure is present in the (C)NTK. As the structure of the NTK is from prior work (e.g., refs. [9] and [12]) it does feel like some of the more substantial contributions relate to observations about the CNTK, (limited) spectral approximation bounds, and showing a level of numerical feasibility. I do think there are some improvements that could be made to the manuscript to better explain aspects of these contributions\n\nIn section 4 the authors note that there are specific structural observations they make about the CNTK that allow it to be efficiently sketched. Do those observations alone yield a faster algorithm than what is cited from [5]? Additionally, while the dependence on $d_1$ and $d_2$ is better than [5] the cost is worse in L and an epsilon is introduced. Thinking of the image sized as fixed (which seems reasonable in certain settings), this means there is a network depth (or desired accuracy) at which point the exact computation is seemingly preferable; is that the case? and how should that trade-off be evaluated?\n\nIn the numerical results the authors have to pick parameters (beyond just the final feature size), how are these chosen? Moreover, even the final feature size does not necessary seem to be directly guided by the theoretical results (i.e., the theory suggests larger sketch sizes may be warranted) the paper shows the sketch sizes chosen yield decent results. Is there some \"practical guidance\" that could be applied here? How were other internal parameters (e.g., in Alg. 1) chosen?\n\nIn some of the algorithms there are two distinct steps: (1) approximation by a polynomial kernel and (2) a specially designed sketch for that kernel. While the polynomial kernel order is chosen adaptability, is that \"convergence\" with epsilon balanced with the sketch convergence? Roughly speaking, as \\epsilon -> 0 would the method, at some point, be reasonably thought of as using an \"exact\" polynomial kernel? and, if so, what are the implications of that? Some of the properties of the NTK cited by the authors (e.g., generalization) depend on kernel properties such as minimal eigenvalues. The behavior of some of those properties could look quite different for a smooth polynomial kernel vs the less smooth zonal kernel induced by ReLU. This presumably relates to the authors informal discussion of regularization effects.\n\nWhile I think it is reasonable that this paper sets itself in the \"NTK regime\" and is presented as a computational improvement for that setting, it may be reasonable for the paper to briefly note that the NTK regime is not fully descriptive. See, e.g., When Do Neural Networks Outperform Kernel Methods? by Ghorbani et al. NeurIPS 2020 and references therein.\n\nAssorted comments/suggestions/questions:\n\n- In places the authors rely a bit too much on \"inheriting\" information from ref. [5]. For example, Theorem 4 explicitly calls out a model defined in [5] as do some of the CNTK numerical experiments. These details seem important enough that they warrant inclusion somewhere in this work (e.g., the supplement) so that the reader can actually see what the models are.\n\n- Early in the intro $x_i$ is introduced as a $d \\times d$ image and then later is generically a d dimensional vector. I think it would be better to keep x as a d-dimensional vector and reserve some other letters for matrix/tensor representations of images to avoid potential confusion (similarly, maybe it is better if images are always $d_1 \\times d_2$ to avoid overloading d).\n\n- Should eq. (12) have the inequality flipped or a $1-\\delta$ on the right hand side?\n\n- On line 136 the abbreviation DP used later for dynamic program should be properly \"defined.\"\n\n- The manuscript could be a bit more clear about what all operations are explicitly included in the timing reported. There is a brief note on this in 5.3, but what about for the CNTK?\n\n- The manuscript may benefit from having a bit more notation introduced earlier on. This would make it easier to talk concretely about kernel learning, features maps, etc. \n", "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": "4.5"}, "extra_text": {}}, {"review_id": "MqiNWKcuTaT", "reviewer": "Reviewer_agp1", "summary": "\n\nIn this paper, the authors propose a method to approximate NTK by sketching the polynomial expansions of arc-cosine kernels. The POLYSKETCH  sketch technique is used to reduce the dimensionality of high-degree tensor products.  Moreover,  the approximation error is theoretically analyzed.", "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": "\n\nPros.\n1. Approximate ReLU-NTK by truncation of Taylor series of expansion of arc-cosine kernels is interesting.  The corresponding approximation error bound is new. The complexity of the feature dimension is smaller than the gradient-based approximation scheme (Arora et al, 2019). \n2. The paper is well organized and clearly written.\n\nCons. \n1. Although I appreciate the theoretical contribution, the motivation of acceleration ReLU-NTK is unconvinced for me.  Theoretical analysis gradient-based approximation is for better understand deep neural network training. In contrast, this paper focus on the acceleration of learning with ReLU-NTK.  However, even exact ReLU-NTK works not well compared with deep neural networks.  The proposed acceleration may not be easy to find practical applications. \n\n2. The paper focuses on ReLU-NTK approximation and acceleration. However, the empirical test focuses on the downstream classification and regression.  It is better to show the empirical study of the approximation error compared with different baselines. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations are clearly stated. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 "}, "extra_text": {}}, {"review_id": "5ASDAGCzOds", "reviewer": "Reviewer_kQ2k", "summary": "This work combines dot product kernel approximation via Taylor series and arc-cosine kernel approximation for NTK approximation for deep ReLU neural networks. In particular, to avoid the exponential increase in $L$, POLYSKETCH is applied to guarantee the complexity to be linear in the depth $L$. Besides, the leverage sampling based strategy is also developed for arc-cosine kernels and further for NTK. ", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "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": "I happen to review this paper again. This submission merged two papers and addressed most of my previous concerns. Therefore, I have no further comments on this paper except the following minor issues:\n\n\n1.\tIn the experiments, the MSE metric should be changed to $|| K - \\tilde{K} ||/|| K ||$ in some norms for intuitive/fair comparison. Please follow experimental settings of previous works on random features/Nystrom approximation.\n\n2.\tThe obtained approximation bound in Theorem 1 might be loose due to two independent constants $\\epsilon$ and $\\delta$. It would be possible to obtain a tight bound? E.g., an exponential order decay following the fashion of concentration inequalities. I don’t check the proof carefully but I think the JL transform used in the proof of Theorem 1 could be improved?\n\n3.\tOne paper on leverage score based sampling for random features is missing:\n\n*Li, Z., Ton, J. F., Oglic, D., & Sejdinovic, D. Towards a unified analysis of random Fourier features. ICML2019.*\n\nThe used technical tools for spectral approximation is also based on the matrix concentration inequality, see Lemma 4 for details. This paper and other random features literature [37] require the random feature map to be continuous, which in fact excludes the zero-order arc-cosine kernel (corresponds to the step function). How to address this when applied to NTK?\n\n4.\tThe expression “low-dimensional feature maps” is inaccurate. The dimension of the obtained feature map is larger than d but is smaller than n. So it would be better to polish this.\n\n5.\tThe Gibbs sampling used in Algorithm 3 seeks for an approximation sampling as the modified distribution q(w) is intractable to be sampled.\n\n\n6.\tSome reference information should be checked, e.g., [36] published on NeurIPS2007 and [37] published on NeurIPS2017. \n", "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": "2"}, "extra_text": {}}, {"review_id": "-yH7bAhfNPp", "reviewer": "Reviewer_CT76", "summary": "This paper gives an approximation method for calculating neural tangent kernel with random features and sketching. The algorithm significantly speeds up the original NTK computation.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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 method described in the paper address the vital running time problem of the original NTK. The experiments verify the speed up by a large factor. The time complexity is well studied in the paper.\n\nI have some concerns in the writing, as it seldomly introduces the original NTK and then directly in the proposed algorithm itself, which is unfriendly to some readers. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "\nThere are many complexity notions in the paper that are hard to digest, e.g. in the feature dimension sample complexity in 192, is there a dominating factor within the \\Omega notion, as it later needs to compare with previous works?", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper considers sketching and random feature methods to accelerate learning with the Neural Tangent Kernel (NTK). NTK based methods are promising linearized approximations of non-convex neural network models, however, they suffer from the high computational complexity of kernel matrix operations. Scaling NTK methods via randomized approximations is a promising direction towards understanding deep neural networks and matching their performance using simpler architectures such as linear kernel machines. 

The reviews all agreed that the paper contains interesting theoretical and experimental results, and expressed minor concerns and suggestions.  One reviewer pointed out that a discussion of the regime in which NTK is not fully descriptive would be useful. Another reviewer commented that the complexity notions are hard to digest. Overall, the paper received generally positive reviews, while some of them pointed out certain minor problems and gave above borderline scores, they all recommended accepting the paper. Please take into account the updated reviews when preparing the final version to accommodate the requested changes.
Decision: Accept (Poster)