Scaling Neural Tangent Kernels via Sketching and Random Features

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.


Reviews and author responses:

 Thanks for the authors' feadback.
I won't change my review score as it is good enough in terms of these minor issues.


 I would like to thank the authors for their detailed response to my review.

While the response does clarify some points, it does not do so in a manner that appreciably changes my perspective on the paper. For example, motivation for simply setting all the number of features (dimensions) equal and a discussion of scaling in different "regimes" (e.g., w.r.t. $L$) seem to remain limited. The necessity of further discussion on some of these points in the present work is debatable, but these limitations in conjunction with my thoughts on the overall advancements/novelties of the work does mean that I will keep my score as is.

 We thank the reviewer for the detailed and insightful comments. Regarding the comment:

- Feels like the big step/insight was in ref [1] where the sketches are constructed:
  * Our sketching method works by first approximating the arc-cosine kernels in each layer of the network by a polynomial and then applying dimensionality reduction using the PolySketches proposed in ref [1]. We then combine the resulting features of arc-cosine kernels in consecutive layers of the network using the PolySketch. Furthermore, our random features method works in the same fashion except that we approximate the arc-cosine kernel by random features sampling as opposed to polynomial expansion. If we focused on a single layer of the network then it might appear that most of the analysis relies on the guarantees of the PolySketch presented in ref [1]. However, we need to analyze a multi-layered network where the sketched output of one layer is fed to the next layer and bounding the effects of errors from previous layers in the output of each layer is non-trivial. While our algorithm heavily relies on the PolySketch, the analysis of our method is completely new and does not immediately follow from prior works.




- Structural observation of CNK: 
  * Our structural observations about CNTK do not imply a faster kernel computation than the DP given in [5] and only enable our low-rank kernel approximation.
Additionally, exact computation of CNTK requires $\mathcal{O}(L (d_1 d_2)^2)$ operations. This means that computing the kernel matrix for a dataset of $n$ images requires $\mathcal{O}(L (d_1 d_2)^2 \cdot n^2)$ operations. Our method does not need to compute the kernel matrix anymore and one can simply transform each of the data points to the (approximate) feature space using our method and then run linear learning methods such as kernel ridge regression on the transformed data points. Transforming $n$ images using our method takes $\mathcal{O}(L^{11}/\varepsilon^{6.7} \cdot (d_1 d_2) \log^3(n) \cdot n)$ and since our feature dimension is only $\mathcal{O}(\log(n) / \varepsilon^2)$, we can run approximate KRR using the same asymptotic runtime. This is faster than exact CNTK computation in two ways: our method scales linearly in n as opposed to quadratic; our method scales linear in the number of pixels $(d_1 d_2)$ as opposed to quadratic. We will emphasize this in the intro section.





- Parameter choice:
  * For our random features method, while the theoretical result involves 3 distinct parameters, which are the target dimension of the PolySketch and the number of random features sampled for arc-cosine kernels of order 0 and 1, to simplify things in practice we set all these quantities to be equal to the final feature size. This way there is no hyperparameter and we only need to choose our desired final feature size.
For our (C)NTK Sketches, there are various hyperparameters such as the target dimensions of the intermediate sketch matrices that we use. To simplify the algorithm in practice we set all target dimensions of all sketch matrices in our algorithm to be equal to the final target dimension. This gets rid of most hyperparameters. 
The only hyperparameter left is the degree of the polynomial we use to approximate the arc-cosine kernels, i.e., the degrees of $P_{relu}^{(p)}(\cdot)$ and $\dot{P}_{relu}^{(p)}(\cdot)$ in Eq (6). We choose the degree from the set {$2, 4, 8$} by running a hyperparameter search with cross-validation on the training data. However, our practical guidance is that choosing a polynomial degree of 4 works well, almost always. We will add these details of the practical guidance in our final draft.

- Does convergence with epsilon balance polynomial and sketch?
  * In the analysis of our (C)NTK Sketch, $\varepsilon$ controls the errors of both polynomial approximation of the arc-cosine functions and the skekching of the resulting polynomial kernels. More specifically, $\varepsilon$ controls the degree of the polynomials in Eq (6) of Algorithm1 (i.e., $p,p’$) as well as the target dimensions of various sketch matrices in Algorithm 1. Hence, $\varepsilon \rightarrow 0$ implies infinite degree polynomials and infinite sketching dimension, thus both of these approximations become exact. 




- Usage of “NTK regime”:
  * We agree that the NTK regime is not fully descriptive and do not intend to claim that our work can in any way replace the trained neural networks. We believe that learning with NTK kernels is one of few promising research directions that can increase the practical relevance of kernel methods. So, our paper tries to lift one of the main obstacles in kernel-based learning with NTKs which was their slow runtime. We will add a remark that clearly describes our setting in the intro.



- We will also update all typos and minor issues in our final manuscript.


 We thank the reviewer for the detailed and insightful comments. Regarding the comment:

- MSE metric:
  * In Table 2, we evaluate the MSE of the different (approximate) kernel regression methods on the test dataset, not the kernel approximation quality. For the final version of our paper, we will make sure to add a separate kernel approximation experiment with the suggested norm evaluation.



- Tighter bound with delta and epsilon:
  * Our NTK Sketch is constructed by recursively composing various independent sketch matrices and we analyzed it by various applications of the union bound and triangle inequality, which might be loose. In other words, in our analysis, we always assumed that the errors of different sketch matrices accumulate in the worst case while in reality, since the sketch matrices are independent, we should get a random walk on the real line. However, we leave a tight analysis of our sketches as an open question for future works. 



- Missing reference:
  * Thanks a lot for pointing this out. We will update our references.



- Continuity of feature map:
  * Our analysis of the spectral approximation relies on Theorem 3.3 in [Lee et al., 2020], which does not require continuity of the feature map. The proof is mainly based on some version of the matrix Bernstein inequality and its only requirement is the finiteness of the integral of the leverage score function. All feature maps in our paper satisfy this condition.

    [Lee et al., 2020] Lee, Jason D., et al. Generalized Leverage Score Sampling for Neural Networks. NeurIPS, 2020.



We will also revise all minor issues in our submission including the term “low-dimensional feature maps”, Gibbs sampling as well as correct reference information as suggested.


 We thank the reviewer for the detailed and insightful comments. Regarding the comment:

- Unconvinced motivation: 
  * We agree with the reviewer that the kernel methods in general have lost some of their practical relevance in recent years and deep learning has replaced them in many applications. We view NTK as one of the most promising directions for bringing the kernel methods back to some applications again. The main obstacle to this is the slow runtime of the kernel based learning with NTK and our paper tries to address this issue.
We remark that although the NTK does not currently perform as well as deep neural networks, it does not require training (possibly huge) parameters of a network and it would be exciting to have a purely untrained kernel based model that works reasonably well and fast.
Moreover, recent works study interesting practical applications of the NTK such as active learning [Shoham et al., 2020] and neural architect search [Chen et al., 2021]. However, these literature commonly point out the computational problem of the NTK. Our motivation on accelerating the NTK is thus clear and convincing to us.

    [Shoham et al., 2021] Shoham, Neta, and Haim Avron. Experimental Design for Overparameterized Learning with Application to Single Shot Deep Active Learning. arXiv:2009.12820. 2020.

    [Chen et al., 2021] Chen, Wuyang, Xinyu Gong, and Zhangyang Wang. Neural Architecture Search on Imagenet in Four GPU Hours: A Theoretically Inspired Perspective. ICLR. 2021.



- Approximation error: 
    * Thanks a lot for the suggestion. We will provide some experiments on the kernel approximation errors of our algorithms in the final draft of our work.



 We thank the reviewer for the detailed and insightful comments. We provide basic concepts of the NTK in the supplementary material (please see Appendix A) due to the space limitation. We will make sure to improve the presentation by adding references to Appendix A.
Regarding the feature dimension in line 192, as the reviewer suggested, unless the dimensionality of the input dataset is huge, the first term is enough to compare with previous work. We will update this as well as all notations in the final draft.


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.  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.

I 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.  
There 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?

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.  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

In 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?

In 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?

In 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.

While 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.

Assorted comments/suggestions/questions:

- 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.

- 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).

- Should eq. (12) have the inequality flipped or a $1-\delta$ on the right hand side?

- On line 136 the abbreviation DP used later for dynamic program should be properly "defined."

- 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?

- 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. 
 Yes



In 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.  

Pros.
1. 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). 
2. The paper is well organized and clearly written.

Cons. 
1. 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. 

2. 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.  The limitations are clearly stated. 

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.   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:


1.	In 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.

2.	The 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?

3.	One paper on leverage score based sampling for random features is missing:

*Li, Z., Ton, J. F., Oglic, D., & Sejdinovic, D. Towards a unified analysis of random Fourier features. ICML2019.*

The 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?

4.	The 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.

5.	The Gibbs sampling used in Algorithm 3 seeks for an approximation sampling as the modified distribution q(w) is intractable to be sampled.


6.	Some reference information should be checked, e.g., [36] published on NeurIPS2007 and [37] published on NeurIPS2017. 
 yes

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