# Random Features for the Neural Tangent Kernel

Insu Han School of Electrical Engineering, Korea Advanced Institute of Science
and Technology. Email<EMAIL_ADDRESS>Haim Avron Department of Applied
Mathematics, Tel Aviv University. Email<EMAIL_ADDRESS>Neta Shoham
Department of Applied Mathematics, Tel Aviv University. Email:
<EMAIL_ADDRESS>Chaewon Kim Graduate School of AI, Korea Advanced Institute
of Science and Technology. Email<EMAIL_ADDRESS>Jinwoo Shin Graduate
School of AI, Korea Advanced Institute of Science and Technology. Email:
<EMAIL_ADDRESS>
###### Abstract

The Neural Tangent Kernel (NTK) has discovered connections between deep neural
networks and kernel methods with insights of optimization and generalization.
Motivated by this, recent works report that NTK can achieve better
performances compared to training neural networks on small-scale datasets.
However, results under large-scale settings are hardly studied due to the
computational limitation of kernel methods. In this work, we propose an
efficient feature map construction of the NTK of fully-connected ReLU network
which enables us to apply it to large-scale datasets. We combine random
features of the arc-cosine kernels with a sketching-based algorithm which can
run in linear with respect to both the number of data points and input
dimension. We show that dimension of the resulting features is much smaller
than other baseline feature map constructions to achieve comparable error
bounds both in theory and practice. We additionally utilize the leverage score
based sampling for improved bounds of arc-cosine random features and prove a
spectral approximation guarantee of the proposed feature map to the NTK matrix
of two-layer neural network. We benchmark a variety of machine learning tasks
to demonstrate the superiority of the proposed scheme. In particular, our
algorithm can run tens of magnitude faster than the exact kernel methods for
large-scale settings without performance loss.

##  1 Introduction

Recent literature have shown that trained overparameterized Deep Neural
Networks (DNNs), i.e., neural networks with substantially more parameters than
training data points, generalize surprisingly well. In an effort to understand
this phenomena, recently researchers have studied the infinite width limit of
DNNs (i.e., the number of neurons in each hidden layer goes to infinity) and
has shown that in that limit, deep learning is equivalent to kernel regression
where the kernel is the so-called Neural Tangent Kernel (NTK) of the network
Arora et al. (2019b); Chizat et al. (2019); Jacot et al. (2018); Lee et al.
(2020). This connection has been used to shed light on the ability of DNNs to
generalize Cao & Gu (2019); Neyshabur et al. (2019) and the ability to
optimize (train) their parameters efficiently Allen-Zhu et al. (2019); Arora
et al. (2019a); Du et al. (2019).

Beyond the aforementioned theoretical purpose, several papers have explored
the algorithmic use of the NTK. Arora et al. (2019c) and Geifman et al. (2020)
showed that NTK based kernel models can sometimes perform better than trained
neural networks. The NTK has also been used in experimental design for neural
networks Shoham & Avron (2020) and predicting training time Zancato et al.
(2020).

Although the NTK of a given network can sometimes be expressed as closed-form
formulas Arora et al. (2019b); Novak et al. (2020), actually using it to learn
kernel models encounters the computational bottlenecks of kernel learning,
e.g., $\mathcal{O}(n^{3})$ time and $\mathcal{O}(n^{2})$ space complexity for
kernel ridge regression with $n$ data points. With the NTK, the situation is
much worse since the cost required to compute the kernel matrix can be huge
Novak et al. (2020). This makes exact kernel learning with NTK infeasible
under large-scale setups.

There is a rich literature on using kernel approximations in order to enable
large-scale learning. One of the most popular approaches is the random
features approach, originally due to Rahimi & Recht (2009). Following the
seminal work, many random feature constructions have been suggested for a
variety of kernels, e.g., arc-cosine kernels Cho & Saul (2009), polynomial
kernels Pham & Pagh (2013); Pennington et al. (2015), general dot product
kernels Han et al. (2020), just to name a few. These low-dimensional features
enable us to apply them to fast linear methods, saving time and space
complexity drastically. Furthermore, their performances are similar or
sometimes better than the exact kernel methods due to implicit regularization
effect Rahimi & Recht (2009); Rudi & Rosasco (2016); Jacot et al. (2020).

In this paper, we propose an efficient random features construction for the
NTK of fully-connected neural networks with ReLU activations. Our starting
point is the explicit feature map for the NTK suggested by Bietti & Mairal
(2019). That feature map uses known explicit feature maps for the arc-cosine
kernel. By replacing the explicit feature map of the arc-cosine kernel with a
random feature map for the same kernel Cho & Saul (2009), we obtain a random
feature map for the NTK. However, the size of that feature map for that
construction can be even larger than the number of input data points $n$. The
underlying reason is from the tensor products between features generated in
consecutive layers. To avoid the issue, we utilize an efficient sketching
algorithm known as TensorSketch transform Pham & Pagh (2013); Ahle et al.
(2020); Woodruff & Zandieh (2020) which can effectively approximate the tensor
products of vectors while preserving their inner products. We provide a
rigorous error analysis of the proposed scheme. The resulting random features
have smaller dimension than the previous NTK feature map constructions.

Furthermore, in order to approximate the NTK with less features, we improve
the underlying existing random feature map of the arc-cosine kernel. Our
construction is based on a modified leverage score based sampling. Recent
literature has shown that random features that use leverage score sampling
entertain better convergence bounds Avron et al. (2017b); Lee et al. (2020).
However, computing the exact leverage scores requires the inversion of a
$n$-by-$n$ matrix which is equivalent to the cost for solving the kernel
method exactly. Luckily, Avron et al. (2017b); Lee et al. (2020) showed that
sampling from the upper bound of leverage score is enough to provide tight
error bounds. Motivated by these results, we propose simple and closed-form
upper bounds of leverage scores regarding to arc-cosine kernels. For further
efficiency, we make use of Gibbs sampling to generate random features from the
proposed modified distribution.

To theoretically justify our construction, we provide a spectral approximation
guarantee for the proposed random features for two-layer neural network.
Recent literature has advocated the use of such spectral bounds as a more
general metric to measure kernel approximation quality as it pertains to many
downstream tasks Avron et al. (2017b).

Finally, we empirically benchmark the proposed random feature methods under
machine learning tasks including classification/regression under UCI datasets,
and active learning under MNIST datasets. We demonstrate that our random
features method can perform similar to or better than the kernel method with
NTK. We further show that the random features approach can run up to 17 times
much faster with tested large-scale datasets, without loss on performance.

Related works.  Many literature have studied a variety of NTK properties
including optimization Allen-Zhu et al. (2019); Du et al. (2019),
generalization Cao & Gu (2019), loss surface Mei et al. (2018) and so on.
Recent works have more focused on NTK kernel itself. Geifman et al. (2020);
Chen & Xu (2021) have discovered that NTK are similar to the Laplace kernel in
term of spectral information when data points are in hypersphere. Fan & Wang
(2020) studied eigenvalue distributions of NNGP and NTK and showed that they
converge to deterministic distribution. However, to the best of our knowledge,
not many works focus on approximating the NTK. Arora et al. (2019b) studied
that gradient of randomly initialized network with finite widths can
approximate the NTK in theory. However, they report that practical
performances of random gradients are worse than that of the exact NTK by a
large margin. Another line of work on NTK approximation is an explicit feature
map construction via tensor product proposed by Bietti & Mairal (2019). These
explicit features can have infinite dimension in general. Hence, it is
impossible to use their features in practice. Even though one can use a
finite-dimensional feature map, the computational gain of random features can
be lost due to expensive tensor product operations.

##  2 Preliminaries

##### Notations.

We use $[n]:=\\{1,\dots,n\\}$. We denote $\otimes$ by the tensor (a.k.a.
Kronecker) product and $\odot$ by the element-wise (a.k.a. Hadamard) product
of two matrices. For square matrices ${\bm{A}}$ and ${\bm{B}}$, we write
${\bm{A}}\preceq{\bm{B}}$ if ${\bm{B}}-{\bm{A}}$ is positive semi-definite. We
write $[{\bm{A}}]_{i,j}$ as an entry of ${\bm{A}}$ in $i$-th row and $j$-th
column. Similarly, $[{\bm{v}}]_{i}$ is used for an $i$-th entry of vector
${\bm{v}}$. We also denote $\mathrm{ReLU}(x)=\max(x,0)$ and consider this
element-wise operation when the input is a matrix. Given a positive
semidefinite matrix ${\bm{K}}$ and $\lambda>0$, the statistical dimension of
${\bm{K}}$ with $\lambda$ is defined as
$s_{\lambda}({\bm{K}}):=\mathtt{tr}({\bm{K}}({\bm{K}}+\lambda{\bm{I}})^{-1})$.

###  2.1 NTK of Fully-connected Deep Neural Networks

Given an input ${\bm{x}}\in{\mathbb{R}}^{d}$, consider a fully-connected ReLU
network with input dimension $d$, hidden layer dimensions $d_{1},\dots,d_{L}$
as

$\displaystyle
f({\bm{x}},\bm{\theta})={\bm{h}}_{L}^{\top}{\bm{w}},\quad{\bm{h}}_{0}={\bm{x}},\quad{\bm{h}}_{\ell}=\sqrt{\frac{2}{d_{\ell}}}~{}\mathrm{ReLU}\left({\bm{h}}_{\ell-1}^{\top}{\bm{W}}_{\ell}\right),$
(1)

where $\bm{\theta}:=({\bm{W}}_{1},\dots,{\bm{W}}_{L},{\bm{w}})$ for
${\bm{W}}_{\ell}\in{\mathbb{R}}^{d_{\ell-1}\times d_{\ell}}$,
${\bm{w}}\in{\mathbb{R}}^{d_{L}}$, $\ell\in[L]$ represents the trainable
parameters and $d_{0}=d$. The Neural Tangent Kernel (NTK) is defined as

$\displaystyle
K_{\mathrm{NTK}}^{(L)}({\bm{x}},{\bm{x}}^{\prime})={\mathbb{E}}_{\bm{\theta}}\left[\left\langle{\frac{\partial
f({\bm{x}},\bm{\theta})}{\partial\bm{\theta}},\frac{\partial
f({\bm{x}}^{\prime},\bm{\theta})}{\partial\bm{\theta}}}\right\rangle\right]$
(2)

where $\bm{\theta}$ is from the standard Gaussian distribution. Given $n$ data
points
${\bm{X}}=[{\bm{x}}_{1},\dots,{\bm{x}}_{n}]^{\top}\in{\mathbb{R}}^{n\times
d}$, we will write
$f({\bm{X}},\bm{\theta}):=[f({\bm{x}}_{1},\bm{\theta}),\dots,f({\bm{x}}_{n},\bm{\theta})]^{\top}\in{\mathbb{R}}^{n}$
and the NTK matrix as ${\bm{K}}_{\mathrm{NTK}}^{(L)}\in{\mathbb{R}}^{n\times
n}$ whose $(i,j)$-th entry is
$K_{\mathrm{NTK}}^{(L)}({\bm{x}}_{i},{\bm{x}}_{j})$ for $i,j\in[n]$.

The motivation for the definition of the NTK is as follows. Consider learning
the network parameters $\bm{\theta}$ by minimizing the squared loss
$\frac{1}{2}\left\|f({\bm{X}},\bm{\theta})-{\bm{y}}\right\|_{2}^{2}$ for the
target ${\bm{y}}\in{\mathbb{R}}^{n}$ using gradient descent with
infinitesimally small learning rate. Regard the parameters as a time-evolving
continuous variable $\bm{\theta}_{t}$ for $t\geq 0$ that develops in the
course of the optimization. Then, Arora et al. (2019b) showed that

$\displaystyle\frac{d}{dt}f\left({\bm{X}},\theta_{t}\right)=-{\bm{K}}_{t}\cdot\left(f({\bm{X}},\theta_{t})-{\bm{y}}\right)$
(3)

where ${\bm{K}}_{t}:=\frac{\partial
f({\bm{X}},\bm{\theta}_{t})}{\partial\bm{\theta}}\left(\frac{\partial
f({\bm{X}},\bm{\theta}_{t})}{\partial\bm{\theta}}\right)^{\top}\in{\mathbb{R}}^{n\times
n}$. In the infinite width limit, i.e., $d_{1},\dots,d_{L}\rightarrow\infty$,
recent works analyzed that $\bm{\theta}_{t}$ remains constant during
optimization, i.e., equals to $\bm{\theta}_{0}$ Chizat et al. (2019); Allen-
Zhu et al. (2019); Du et al. (2019) and ${\bm{K}}_{t}={\bm{K}}_{0}$.
Furthermore, under a certain random initialization, in the same infinite width
limit, ${\bm{K}}_{0}$ converges in probability to
${\bm{K}}_{\mathrm{NTK}}^{(L)}$. This implies equivalence between the
prediction of neural network under the some initialization and kernel
regression with NTK (Arora et al., 2019b).

In addition, when the parameters of only last layer are updated, the network
prediction corresponds to the kernel so-called the neural network Gaussian
Process (NNGP):

$\displaystyle
K_{\mathrm{NNGP}}^{(L)}({\bm{x}},{\bm{x}}^{\prime})={\mathbb{E}}_{\bm{\theta}}\left[f({\bm{x}},\bm{\theta})\cdot
f({\bm{x}}^{\prime},\bm{\theta})\right]$ (4)

where the expectation is same as the NTK and we denote
$[{\bm{K}}_{\mathrm{NNGP}}^{(L)}]_{i,j}:=K_{\mathrm{NNGP}}^{(L)}({\bm{x}}_{i},{\bm{x}}_{j})$.

NTK computation. The NTK matrix of a fully-connected ReLU network can be
computed by the following recursive relation Jacot et al. (2018); Chizat et
al. (2019); Arora et al. (2019b):

$\displaystyle\begin{aligned}
&{\bm{K}}_{\mathrm{NTK}}^{(0)}={\bm{K}}_{\mathrm{NNGP}}^{(0)}={{\bm{X}}{\bm{X}}^{\top}},\\\
&{\bm{K}}_{\mathrm{NNGP}}^{(\ell)}=F_{1}\left({\bm{K}}_{\mathrm{NNGP}}^{(\ell-1)}\right),\\\
&{\bm{K}}_{\mathrm{NTK}}^{(\ell)}={\bm{K}}_{\mathrm{NNGP}}^{(\ell)}+{\bm{K}}_{\mathrm{NTK}}^{(\ell-1)}\odot
F_{0}\left({\bm{K}}_{\mathrm{NNGP}}^{(\ell-1)}\right),\end{aligned}$ (5)

where $F_{0},F_{1}:{\mathbb{R}}^{n\times n}\rightarrow{\mathbb{R}}^{n\times
n}$ are defined as

$\displaystyle[F_{0}({\bm{K}})]_{i,j}$
$\displaystyle:=1-\frac{1}{\pi}\cos^{-1}\left(\frac{[{\bm{K}}]_{i,j}}{\sqrt{[{\bm{K}}]_{i,i}[{\bm{K}}]_{j,j}}}\right),$
$\displaystyle[F_{1}({\bm{K}})]_{i,j}$
$\displaystyle:=\sqrt{[{\bm{K}}]_{i,i}[{\bm{K}}]_{j,j}}\cdot
f\left(\frac{[{\bm{K}}]_{i,j}}{\sqrt{[{\bm{K}}]_{i,i}[{\bm{K}}]_{j,j}}}\right)$

where $f(x)=\frac{1}{\pi}\left(\sqrt{1-x^{2}}+(\pi-\cos^{-1}(x))x\right)$ for
$x\in[-1,1]$ and ${\bm{K}}$ is an arbitrary positive semidefinite matrix. Note
that these matrix functions are derived from arc-cosine kernels Cho & Saul
(2009):

$\displaystyle\begin{aligned}
A_{0}({\bm{x}},{\bm{x}}^{\prime})&:=1-\frac{1}{\pi}\cos^{-1}\left(\frac{\left\langle{{\bm{x}},{\bm{x}}^{\prime}}\right\rangle}{\left\|{\bm{x}}\right\|_{2}\left\|{\bm{x}}^{\prime}\right\|_{2}}\right),\\\
A_{1}({\bm{x}},{\bm{x}}^{\prime})&:=\left\|{\bm{x}}\right\|_{2}\left\|{\bm{x}}^{\prime}\right\|_{2}f\left(\frac{\left\langle{{\bm{x}},{\bm{x}}^{\prime}}\right\rangle}{\left\|{\bm{x}}\right\|_{2}\left\|{\bm{x}}^{\prime}\right\|_{2}}\right).\end{aligned}$
(6)

Computing NTK of the network with $L$ layers takes $\mathcal{O}(n^{2}(d+L))$
time and $\mathcal{O}(n(n+d))$ space complexity.

###  2.2 Random Features and Spectral Approximation

Random features Rahimi & Recht (2009) is a methodology for scaling kernel
methods that saves both time and storage. In most general terms, the random
features model targets kernels
$K:{\mathbb{R}}^{d}\times{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$ that can be
written as
$K({\bm{x}},{\bm{x}}^{\prime})={\mathbb{E}}_{{\bm{v}}\sim{p}}\left[{\bm{\Phi}({\bm{x}},{\bm{v}})\cdot\bm{\Phi}({\bm{x}}^{\prime},{\bm{v}})}\right]$
for some distribution $p$ and a function
$\bm{\Phi}:{\mathbb{R}}^{d}\times{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$. The
random features approximations works as follows. First, we generate $m$
vectors ${\bm{v}}_{1},\dots,{\bm{v}}_{m}\in{\mathbb{R}}^{d}$ sampled from $p$.
We then define the feature map as

$\displaystyle\bm{\Phi}_{m}({\bm{x}}):=\frac{1}{\sqrt{m}}\left[\bm{\Phi}({\bm{x}},{\bm{v}}_{1}),\dots,\bm{\Phi}({\bm{x}},{\bm{v}}_{m})\right]^{\top}\in{\mathbb{R}}^{m}$

and the approximate kernel is
$K^{\prime}({\bm{x}},{\bm{x}}^{\prime})=\left\langle{\bm{\Phi}_{m}({\bm{x}}),\bm{\Phi}_{m}({\bm{x}}^{\prime})}\right\rangle$.

The main utility of the random feature map is due to the fact that the kernel
matrix ${\bm{K}}^{\prime}$ associated with $K^{\prime}$ is a low-rank matrix
with a known factorization. In particular, let
$\bm{\Phi}:=\left[\bm{\Phi}_{m}({\bm{x}}_{1}),\dots,\bm{\Phi}_{m}({\bm{x}}_{n})\right]^{\top}\in{\mathbb{R}}^{n\times
m}$, then ${\bm{K}}^{\prime}=\bm{\Phi}\bm{\Phi}^{\top}\approx{\bm{K}}$. The
rank of the approximate kernel matrix is $m$, which allows faster computation
and less storage. The parameter $m$ trades between computational complexity
and approximation quality. A small $m$ results in faster speedup but less
accurate kernel approximation.

Although the random features can approximate well 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 $\lambda>0$, that is,

$\displaystyle(1-\varepsilon)\left({\bm{K}}+\lambda{\bm{I}}\right)\preceq\bm{\Phi}\bm{\Phi}^{\top}+\lambda{\bm{I}}\preceq(1+\varepsilon)\left({\bm{K}}+\lambda{\bm{I}}\right)$

for $\varepsilon>0$ and show that it can provide rigorous guarantees of
downstream applications including kernel ridge regression Avron et al.
(2017b), clustering and PCA Musco & Musco (2017).

###  2.3 CountSketch and TensorSketch Transforms

The CountSketch transform is a norm-preserving dimensionality reduction
technique Charikar et al. (2002). Formally, let $h:[d]\rightarrow[m]$ be a
pairwise independent hash function whose bins are chosen uniformly at random
and $s:[d]\rightarrow\\{+1,-1\\}$ be a pairwise independent sign function
where signs are chosen uniformly at random. Given
${\bm{x}}\in{\mathbb{R}}^{d}$ and $m\in\mathbb{N}$, we define
$\mathcal{C}:{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}^{m}$ such that for
$i\in[m]$

$\displaystyle[\mathcal{C}({\bm{x}})]_{i}=\sum_{j:h(j)=i}s(j)[{\bm{x}}]_{j}$
(7)

and it is well-known that
${\mathbb{E}}[\left\langle{\mathcal{C}({\bm{x}}),\mathcal{C}({\bm{y}})}\right\rangle]=\left\langle{{\bm{x}},{\bm{y}}}\right\rangle$.
Observe that it requires a single pass over the input, hence the running time
becomes $\mathcal{O}(d)$.

Pham & Pagh (2013) proposed an efficient algorithm to apply CountSketch to the
tensor product of vectors and referred to this as TensorSketch. Let
$h_{1}:[d_{1}]\rightarrow[m]$, $h_{2}:[d_{2}]\rightarrow[m]$ be pairwise
independent random hash functions and $s_{1}:[d_{1}]\rightarrow\\{-1,1\\}$,
$s_{2}:[d_{2}]\rightarrow\\{-1,1\\}$ be pairwise independent sign functions.
We denote the corresponding CountSketches by $\mathcal{C}_{1}$ and
$\mathcal{C}_{2}$, respectively. Now consider a new transform
$\mathcal{C}:{\mathbb{R}}^{d_{1}d_{2}}\rightarrow{\mathbb{R}}^{m}$ whose hash
and sign functions are defined as $H(j_{1},j_{2})\equiv
h_{1}(j_{1})+h_{2}(j_{2})\pmod{m}$ and $S(j_{1},j_{2})=s_{1}(j_{1})\cdot
s_{2}(j_{2})$ for $j_{1}\in[d_{1}],j_{2}\in[d_{2}]$. Given
${\bm{x}}\in{\mathbb{R}}^{d_{1}},{\bm{y}}\in{\mathbb{R}}^{d_{2}}$, Pham & Pagh
(2013) showed that $\mathcal{C}\left({\bm{x}}\otimes{\bm{y}}\right)$ equals
the convolution between $\mathcal{C}_{1}({\bm{x}})$ and
$\mathcal{C}_{2}({\bm{y}})$ and its computation can be amortized as

$\displaystyle\mathcal{C}\left({\bm{x}}\otimes{\bm{y}}\right)=\mathrm{FFT}^{-1}\left(\mathrm{FFT}(\mathcal{C}_{1}({\bm{x}}))\odot\mathrm{FFT}(\mathcal{C}_{2}({\bm{y}}))\right)$
(8)

where $\mathrm{FFT}$ and $\mathrm{FFT}^{-1}$ are the Fast Fourier Transform
and its inverse. By the inner product preserving property, TensorSketch also
can be used as a low-rank approximation of element-wise product between two
Gramian matrices. More specifically, given ${\bm{X}}\in{\mathbb{R}}^{n\times
d_{1}},{\bm{Y}}\in{\mathbb{R}}^{n\times d_{2}}$, it holds

$\displaystyle\begin{aligned}
&\left({\bm{X}}{\bm{X}}^{\top}\right)\odot\left({\bm{Y}}{\bm{Y}}^{\top}\right)=\left({\bm{X}}\otimes{\bm{Y}}\right)\left({\bm{X}}\otimes{\bm{Y}}\right)^{\top}={\mathbb{E}}\left[\mathcal{C}\left({\bm{X}}\otimes{\bm{Y}}\right)\mathcal{C}\left({\bm{X}}\otimes{\bm{Y}}\right)^{\top}\right]\end{aligned}$
(9)

where $\otimes$ and $\mathcal{C}$ are performed in a row-wise manner. Note
that $\mathcal{C}\left({\bm{X}}\otimes{\bm{Y}}\right)\in{\mathbb{R}}^{n\times
m}$ can be computed in $\mathcal{O}(n(d_{1}+d_{2}+m\log m))$ time using ??.
This is much cheaper than that of computing $n$-by-$n$ dense matrix when $n\gg
d_{1},d_{2},m$. A larger $m$ guarantees better approximation quality but also
increases its running time. Avron et al. (2017a); Ahle et al. (2020); Woodruff
& Zandieh (2020) analyzed a spectral approximation guarantee of TensorSketch
transform.

For simplicity, we presented TensorSketch for vectors that are the tensor
product of only two vectors. This is enough for our needs, though we mention
that the TensorSketch transform can be defined for an arbitrary number of
tensor products.

##  3 Efficient NTK Random Features via Sketching Method

Our goal is to design efficient such random features for the NTK. Seemingly,
one can obtain such random features from definition of NTK (i.e., ??) by using
gradients of the randomly initialized networks as features. However, the NTK
is the infinite-width limit, while in practice we need to fix some finite
width, which will introduce a bias. Moreover, Arora et al. (2019b) showed that
the gradient features from a network with finite widths degrade the practical
performance by a huge gap. Instead, we focus on the closed-form expression of
NTK with ReLU activations.

Algorithm 1 Random Features for NTK of ReLU network via CountSketch

1: Input: ${\bm{x}}\in{\mathbb{R}}^{d}$, network depth $L$, feature dimensions
$m_{0}$, $m_{1}$ and $m_{\mathtt{cs}}$

2:
${\bm{\Phi}}^{(0)}({\bm{x}})\leftarrow{\bm{x}},\bm{\Psi}^{(0)}({\bm{x}})\leftarrow{\bm{x}}$,
and $m\leftarrow d$

3: for $\ell=1$ to $L$ do

4: Draw i.i.d. ${\bm{w}}_{i}\sim\mathcal{N}(\bm{0},{\bm{I}}_{m})$ for
$i\in[m_{0}]$ and
$\bm{\Lambda}^{(\ell)}({\bm{x}})\leftarrow\sqrt{\frac{2}{m_{0}}}\
\mathrm{Step}\left(\begin{bmatrix}{\bm{w}}_{1},\dots,{\bm{w}}_{m_{0}}\end{bmatrix}^{\top}\bm{\Psi}^{(\ell-1)}({\bm{x}})\right)$

5: Draw i.i.d. ${{\bm{w}}}^{\prime}_{j}\sim\mathcal{N}(\bm{0},{\bm{I}}_{m})$
for $j\in[m_{1}]$ and
$\bm{\Psi}^{(\ell)}({\bm{x}})\leftarrow\sqrt{\frac{2}{m_{1}}}\
\mathrm{ReLU}\left(\begin{bmatrix}{\bm{w}}^{\prime}_{1},\dots,{\bm{w}}^{\prime}_{m_{1}}\end{bmatrix}^{\top}\bm{\Psi}^{(\ell-1)}({\bm{x}})\right)$

6: Draw two independent CountSketch transforms $\mathcal{C}_{0}^{(\ell)}$ and
$\mathcal{C}_{1}^{(\ell)}$ that map to ${\mathbb{R}}^{m_{\mathtt{cs}}}$ and
$\bm{\Gamma}^{(\ell)}({\bm{x}})\leftarrow\mathrm{FFT}^{-1}\left(\mathrm{FFT}(\mathcal{C}^{(\ell)}_{0}(\bm{\Lambda}^{(\ell)}({\bm{x}})))\odot\mathrm{FFT}(\mathcal{C}^{(\ell)}_{1}({\bm{\Phi}}^{(\ell-1)}({\bm{x}})))\right)$

7:
${\bm{\Phi}}^{(\ell)}({\bm{x}})\leftarrow\begin{bmatrix}\bm{\Psi}^{(\ell)}({\bm{x}}),\
\bm{\Gamma}^{(\ell)}({\bm{x}})\end{bmatrix}$, $m\leftarrow m_{1}$

8: end for

9: return ${\bm{\Phi}}^{(L)}({\bm{x}})$

We begin by introducing random features of arc-cosine kernels $A_{0}$ and
$A_{1}$ originally due to Cho & Saul (2009):

$\displaystyle a_{0}({\bm{x}})$ $\displaystyle=\sqrt{\frac{2}{m_{0}}}\
\mathrm{Step}\left(\left[{\bm{w}}_{1},\dots,{\bm{w}}_{m_{0}}\right]^{\top}{\bm{x}}\right),$
(10) $\displaystyle a_{1}({\bm{x}})$ $\displaystyle=\sqrt{\frac{2}{m_{1}}}\
\mathrm{ReLU}\left(\left[{\bm{w}}^{\prime}_{1},\dots,{\bm{w}}_{m_{1}}^{\prime}\right]^{\top}{\bm{x}}\right)$
(11)

where
${\bm{w}}_{1},\dots,{\bm{w}}_{m_{0}},{\bm{w}}^{\prime}_{1},\dots,{\bm{w}}^{\prime}_{m_{1}}\in{\mathbb{R}}^{d}$
are sampled from $\mathcal{N}(\bm{0},{\bm{I}}_{d})$. It is known that
${\mathbb{E}}[\left\langle{a_{0}({\bm{x}}),a_{0}({\bm{x}}^{\prime})}\right\rangle]=A_{0}({\bm{x}},{\bm{x}}^{\prime})$
and
${\mathbb{E}}[\left\langle{a_{1}({\bm{x}}),a_{1}({\bm{x}}^{\prime})}\right\rangle]=A_{1}({\bm{x}},{\bm{x}}^{\prime})$
for ${\bm{x}},{\bm{x}}^{\prime}\in{\mathbb{R}}^{d}$.

Recently, Bietti & Mairal (2019) presented an explicit infinite-dimensional
feature map for the NTK of ReLU networks by using recursive tensoring of
explicit feature maps for the arc-cosine kernel. Replacing each explicit
feature map with a random feature map for the corresponding kernel we can
obtain a random feature map for the NTK. The resulting construction is:

$\displaystyle\begin{aligned}
\bm{\Psi}^{(\ell+1)}({\bm{x}})&=a_{1}\left(\bm{\Psi}^{(\ell)}({\bm{x}})\right),\quad\bm{\Phi}^{(0)}({\bm{x}})=\bm{\Psi}^{(0)}({\bm{x}})={\bm{x}},\\\
\bm{\Phi}^{(\ell+1)}({\bm{x}})&=\left[\bm{\Psi}^{(\ell+1)}({\bm{x}}),\quad
a_{0}\left(\bm{\Psi}^{(\ell)}({\bm{x}})\right)\otimes\bm{\Phi}^{(\ell)}({\bm{x}})\right],\end{aligned}$
(12)

for $\ell=0,\dots,L-1$. These features can be used for approximating both NTK
and NNGP as
$K_{\mathrm{NTK}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\approx\left\langle{\bm{\Phi}^{(\ell)}({\bm{x}}),{\bm{\Phi}^{(\ell)}}({\bm{x}}^{\prime})}\right\rangle$
and
$K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\approx\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}),{\bm{\Psi}^{(\ell)}}({\bm{x}}^{\prime})}\right\rangle$.

However, one major drawback of the last construction is that the number of
features is exponential in the depth. Indeed, the dimension of output features
$\bm{\Phi}^{(L)}({\bm{x}})$ is
$\left(\sum_{k=0}^{L-1}m_{0}^{k}\right)m_{1}+m_{0}^{L}d=\mathcal{O}(m_{0}^{L}(m_{1}+d))$.
This also leads to $\mathcal{O}(m_{0}^{L}(m_{1}+d)+m_{1}^{2})$ time
complexity. The exponential growth in depth $L$ is due to the tensor product
$\otimes$ in ??. For a large $L$, the number of features can easily be larger
than the number of data points $n$ and any computational saving is hopeless.

In order to make the feature map more compact, we utilize a TensorSketch to
reduce the dimension of
$a_{0}\left(\bm{\Psi}^{(\ell)}({\bm{x}})\right)\otimes\bm{\Phi}^{(\ell)}({\bm{x}})$.
We do so by replacing it with

$\displaystyle\bm{\Gamma}^{(\ell)}({\bm{x}}):=\mathrm{FFT}^{-1}\left(\mathrm{FFT}\left(\mathcal{C}_{1}(a_{0}(\bm{\Psi}^{(\ell)}({\bm{x}})))\right)\odot\mathrm{FFT}\left(\mathcal{C}_{2}(\bm{\Phi}^{(\ell)}({\bm{x}}))\right)\right)$

where $\mathcal{C}_{1}$ and $\mathcal{C}_{2}$ are independent111i.e., hash and
sign functions of $\mathcal{C}_{1}$ and $\mathcal{C}_{2}$ are independent.
CountSketch transforms that map to ${\mathbb{R}}^{m_{\mathtt{cs}}}$. Denote
$\widehat{\bm{\Phi}}^{(\ell)}({\bm{x}}):=\begin{bmatrix}\bm{\Psi}^{(\ell)}({\bm{x}}),\
\bm{\Gamma}^{(\ell)}({\bm{x}})\end{bmatrix}$ and one can expect that

$\displaystyle\left\langle{\bm{\Phi}^{(\ell)}({\bm{x}}),\bm{\Phi}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle\approx\left\langle{\widehat{\bm{\Phi}}^{(\ell)}({\bm{x}}),\widehat{\bm{\Phi}}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle$
(13)

from the property in ??. The process is repeated for every layer. A pseudo-
code for the proposed feature construction is described in ??.

We now provide that the approximation error bound of generated features from
??.

######  Theorem 1.

Given ${\bm{x}},{\bm{y}}\in\mathbb{R}^{d}$ such that
$\left\|{\bm{x}}\right\|_{2}=\left\|{\bm{x}}^{\prime}\right\|_{2}=1$ and
$L\geq 1$, let $K_{\mathrm{NTK}}^{(L)}$ the NTK of $L$-layer fully-connected
ReLU network. Given $\delta,\in(0,1)$, $\varepsilon\in(0,1/L)$, there exist
constants $C_{0},C_{1},C_{2}>0$ such that

$m_{0}\geq
C_{0}\frac{L^{2}}{\varepsilon^{2}}\log\left(\frac{L}{\delta}\right),\ \
m_{1}\geq
C_{1}\frac{L^{6}}{\varepsilon^{4}}\log\left(\frac{L}{\delta}\right),\ \
m_{\mathtt{cs}}\geq C_{2}\frac{L^{3}}{\varepsilon^{2}\delta}$

and

$\displaystyle\Pr\left(\left|\left\langle{{\bm{\Phi}}^{(L)}({\bm{x}}),\bm{\Phi}^{(L)}({\bm{x}}^{\prime})}\right\rangle-
K_{\mathrm{NTK}}^{(L)}({\bm{x}},{\bm{x}}^{\prime})\right|\leq
L\varepsilon\left(1+\frac{\varepsilon}{2}\right)^{2}+\varepsilon\right)\geq
1-\delta$

where
$\bm{\Phi}^{(L)}({\bm{x}}),\bm{\Phi}^{(L)}({\bm{x}}^{\prime})\in{\mathbb{R}}^{m_{1}+m_{\mathtt{cs}}}$
be the output of ?? of ${\bm{x}},{\bm{x}}^{\prime}$, respectively, using the
same CountSketch transforms.

The proof of ?? is provided in ??. We note that the restriction of inputs to
the hypersphere (i.e., $\left\|{\bm{x}}_{i}\right\|_{2}=1$) is a common
assumption used in the NTK analysis Bietti & Mairal (2019); Geifman et al.
(2020). This can be easily achieved by normalizing input data points. From ??,
the dimension of the proposed random features can be
$\mathcal{O}\left(\frac{L^{6}}{\varepsilon^{4}}\log\left(\frac{L}{\delta}\right)+\frac{L^{3}}{\varepsilon^{2}\delta}\right)$,
which gets rid of the exponential dependency on $L$, to guarantee the above
error bound. Furthermore, Arora et al. (2019b) studied that the gradient of
randomly initialized ReLU network with finite width can approximate the NTK,
but their feature dimension should be
$\Omega\left(\frac{L^{13}}{\varepsilon^{8}}\log^{2}\left(\frac{L}{\delta}\right)+\frac{L^{6}}{\varepsilon^{4}}\log\left(\frac{L}{\delta}\right)d\right)$
to guarantee an approximation error of $(L+1)\varepsilon$ with probability at
least $1-\delta$. This error bound is smaller than that in ?? by a factor of
$\left(1+\frac{\varepsilon}{2}\right)^{2}$, but their feature dimension is
much larger by a factor of
$\mathcal{O}\left(\frac{L^{7}}{\varepsilon^{4}}\log\left(\frac{L}{\delta}\right)\right)$.
We empirically observe that ?? requires much fewer dimension than both the
random gradient and the naïve feature map construction in ?? to achieve the
same error and provide these results in ??.

##  4 Spectral Approximation for the NTK

Our ultimate goal is to provide lower bounds on the parameter
$m_{0},m_{1},m_{\mathtt{cs}}$ to achieve tight error bound in terms of
spectral approximation of the NTK, i.e.,

$(1-\varepsilon)\left({\bm{K}}^{(L)}_{\mathrm{NTK}}+\lambda{\bm{I}}\right)\preceq\bm{\Phi}^{(L)}(\bm{\Phi}^{(L)})^{\top}+\lambda{\bm{I}}\preceq(1+\varepsilon)\left({\bm{K}}^{(L)}_{\mathrm{NTK}}+\lambda{\bm{I}}\right),$

where
$\bm{\Phi}^{(L)}:=\left[\bm{\Phi}^{(L)}({\bm{x}}_{1}),\dots,\bm{\Phi}^{(L)}({\bm{x}}_{n})\right]^{\top}$.
We first provide spectral bounds of the arc-cosine kernels in ??, which are
necessary prerequisites of our analysis on the NTK random features. Based on
these results, we present spectral bound of a two-layer ReLU network (i.e.,
$L=1$ in ??) and discuss hardness on generalizing this result to networks with
deeper layer. To the best of our knowledge, this has not been studied in
previous literature.

###  4.1 Spectral Approximation for Arc-cosine Kernels

Recently, Avron et al. (2017b); Lee et al. (2020) proposed that random
features with sampling from a modified distribution can give better
approximation guarantee. More precisely, suppose
$\bm{\Phi}:{\mathbb{R}}^{d}\times{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$ is a
function for random features of kernel $K$ with distribution $p$. Consider
random vectors ${\bm{z}}_{1},\dots,{\bm{z}}_{m}\in{\mathbb{R}}^{d}$ sampled
from some distribution $q$. Denote that

$\displaystyle\overline{\bm{\Phi}}_{m}({\bm{x}}):=\frac{1}{\sqrt{m}}\left[\sqrt{\frac{p({\bm{z}}_{1})}{q({\bm{z}}_{1})}}\bm{\Phi}({\bm{x}},{\bm{z}}_{1}),\dots\sqrt{\frac{p({\bm{z}}_{m})}{q({\bm{z}}_{m})}}\bm{\Phi}({\bm{x}},{\bm{z}}_{m})\right]^{\top}\in{\mathbb{R}}^{m}$

then one can verify that ${\mathbb{E}}_{{\bm{z}}\sim
q}\left[\left\langle{\overline{\bm{\Phi}}_{m}({\bm{x}}),\overline{\bm{\Phi}}_{m}({\bm{x}}^{\prime})}\right\rangle\right]=K({\bm{x}},{\bm{x}}^{\prime})$
for all ${\bm{x}},{\bm{x}}^{\prime}\in{\mathbb{R}}^{d}$. Now assume that the
distribution $q$ is defined by a measurable function
$q_{\lambda}:{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$ satisfies that
$q({\bm{v}})=q_{\lambda}({\bm{v}})/\int_{{\mathbb{R}}^{d}}q_{\lambda}({\bm{v}})d{\bm{v}}$
and

$\displaystyle q_{\lambda}({\bm{v}})\geq
p({\bm{v}})\cdot\bm{\Phi}({\bm{X}},{\bm{v}})^{\top}\left({\bm{K}}+\lambda{\bm{I}}\right)^{-1}\bm{\Phi}({\bm{X}},{\bm{v}})$
(14)

for $\lambda>0$ where
$\bm{\Phi}({\bm{X}},{\bm{v}}):=[\bm{\Phi}({\bm{x}}_{1},{\bm{v}}),\dots,\bm{\Phi}({\bm{x}}_{n},{\bm{v}})]^{\top}\in{\mathbb{R}}^{n}$.
Then, they proved that with high probability it holds

$\displaystyle(1-\varepsilon)\left({\bm{K}}+\lambda{\bm{I}}\right)\preceq\overline{\bm{\Phi}}_{m}\overline{\bm{\Phi}}_{m}^{\top}+\lambda{\bm{I}}\preceq(1+\varepsilon)\left({\bm{K}}+\lambda{\bm{I}}\right)$

where
$\overline{\bm{\Phi}}_{m}=\left[\overline{\bm{\Phi}}_{m}({\bm{x}}_{1}),\dots,\overline{\bm{\Phi}}_{m}({\bm{x}}_{n})\right]^{\top}\in{\mathbb{R}}^{n\times
m}$ and $\varepsilon\in(0,1)$ is a given parameter.

Observe that the lower bound in ?? requires $\mathcal{O}(n^{3})$ operations to
compute due to the matrix inverse. This can hurt the computational advantage
of random features. Hence, it is important to find such distribution $q$ that
is easy to sample while holding the ??.

In what follows, we provide that the original arc-cosine random features of
$0$-th order in ?? can indeed guarantee a spectral approximation bound.

######  Theorem 2.

Given ${\bm{X}}\in{\mathbb{R}}^{n\times d}$, let
${\bm{A}}_{0}\in{\mathbb{R}}^{n\times n}$ be the arc-cosine kernel matrix of
$0$-th order with ${\bm{X}}$ and denote
$\bm{\Phi}_{0}:=\sqrt{\frac{2}{m}}\mathrm{Step}({\bm{X}}{\bm{W}})\in{\mathbb{R}}^{n\times
m}$ where each entry in ${\bm{W}}\in{\mathbb{R}}^{d\times m}$ is an i.i.d.
sample from $\mathcal{N}(0,1)$. Let $s_{\lambda}$ be the statistical dimension
of ${\bm{A}}_{0}$. Given $\lambda\in(0,\left\|{\bm{A}}_{0}\right\|_{2})$,
$\varepsilon\in(0,1/2)$ and $\delta\in(0,1)$, if
$m\geq\frac{8}{3}\frac{n}{\lambda\varepsilon^{2}}\log\left(\frac{16s_{\lambda}}{\delta}\right)$,
then it holds

$\displaystyle(1-\varepsilon)({\bm{A}}_{0}+\lambda{\bm{I}})\preceq\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}+\lambda{\bm{I}}\preceq(1+\varepsilon)({\bm{A}}_{0}+\lambda{\bm{I}})$

with probability at least $1-\delta$.

The proof of ?? is provided in ??. The analysis is similar to that studied by
Avron et al. (2017a), i.e.,
$q_{\lambda}({\bm{v}})=({n}/{\lambda})p({\bm{v}})$, which implies that the
modified distribution is identical to the original one.

Next, we present our result on spectral approximation for arc-cosine random
features of $1$-st order. Unlike the previous case, sampling vectors from the
modified distribution in the form of the Gaussian scaled by squared
$\ell_{2}$-norm is required. The formal statement is provided in ??.

Algorithm 2 Modified Random Arc-cosine Features of $1$-st order ?? via Gibbs
Sampling

1: Input: ${\bm{X}}\in{\mathbb{R}}^{n\times d}$, feature dimension $m_{1}$,
Gibbs iterations $T$

2: Draw i.i.d. ${\bm{v}}_{i}\sim\mathcal{N}({\bf 0},{\bm{I}}_{d})$ for
$i\in[m_{1}]$

3: for $i=1$ to $m_{1}$ do

4: $q(x,z)\leftarrow$ inverse of
$\frac{\mathrm{erf}\left({x}/{\sqrt{2}}\right)+1}{2}-\frac{x\exp\left(-x^{2}/2\right)}{\sqrt{2\pi}(z+1)}$
(corresponds to the CDF of
$\Pr([{\bm{v}}_{i}]_{j}|[{\bm{v}}_{i}]_{\setminus\\{j\\}})$)

5: for $t=1$ to $T$ do

6: for $j=1$ to $d$ do

7: $u\leftarrow$ sample from $[0,1]$ at uniformly random

8: $[{\bm{v}}_{i}]_{j}\leftarrow
q\left(u,\sum_{k\in[d]\setminus\\{j\\}}[{\bm{v}}_{i}]_{k}^{2}\right)$

9: end for

10: end for

11: end for

12: return
$\sqrt{\frac{2d}{m}}\left[\frac{\mathrm{ReLU}({\bm{X}}{\bm{v}}_{1})}{\left\|{\bm{v}}_{1}\right\|_{2}},\
\dots\
,\frac{\mathrm{ReLU}({\bm{X}}{\bm{v}}_{m})}{\left\|{\bm{v}}_{m}\right\|_{2}}\right]$

######  Theorem 3.

Given ${\bm{X}}\in{\mathbb{R}}^{n\times d}$, let
${\bm{A}}_{1}\in{\mathbb{R}}^{n\times n}$ be the arc-cosine kernel matrix of
$1$-th order with ${\bm{X}}$ and
${\bm{v}}_{1},\dots,{\bm{v}}_{m}\in{\mathbb{R}}^{d}$ be i.i.d. random vectors
from probability distribution

$\displaystyle
q({\bm{v}})=\frac{1}{(2\pi)^{d/2}d}\left\|{\bm{v}}\right\|_{2}^{2}\exp\left(-\frac{1}{2}\left\|{\bm{v}}\right\|_{2}^{2}\right).$
(15)

Denote

$\displaystyle\bm{\Phi}_{1}:=\sqrt{\frac{2d}{m}}\left[\frac{\mathrm{ReLU}({\bm{X}}{\bm{v}}_{1})}{\left\|{\bm{v}}_{1}\right\|_{2}},\
\dots\
,\frac{\mathrm{ReLU}({\bm{X}}{\bm{v}}_{m})}{\left\|{\bm{v}}_{m}\right\|_{2}}\right]$
(16)

and let $s_{\lambda}$ be the statistical dimension of ${\bm{A}}_{1}$. Given
$\lambda\in(0,\left\|{\bm{A}}_{1}\right\|_{2})$, $\varepsilon\in(0,1/2)$ and
$\delta\in(0,1)$, if
$m\geq\frac{8}{3}\frac{d\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda\varepsilon^{2}}\log\left(\frac{16s_{\lambda}}{\delta}\right)$,
then it holds that

$\displaystyle(1-\varepsilon)({\bm{A}}_{1}+\lambda
I)\preceq\bm{\Phi}_{1}\bm{\Phi}_{1}^{\top}+\lambda{\bm{I}}\preceq(1+\varepsilon)({\bm{A}}_{1}+\lambda
I)$

with probability at least $1-\delta$.

The proof of ?? is provided in ??. We note that the modified distribution can
be expressed as a closed-form formula as in ??. Once random vectors are
sampled from this distribution, the modified random features in ?? can be
computed at the same cost of the original features in ??. In addition, the
lower bound on feature dimension depends on the square of the spectral norm of
input.

Approximate sampling. It is not trivial to sample a vector
${\bm{v}}\in{\mathbb{R}}^{d}$ from the distribution $q(\cdot)$ defined in ??.
Thus, we suggest to perform an approximate sampling via Gibbs sampling. The
algorithm starts with a random initialized vector ${\bm{v}}$ and then
iteratively replaces $[{\bm{v}}]_{i}$ with a sample from
$q([{\bm{v}}]_{i}|[{\bm{v}}]_{\setminus i})$ for $i\in[d]$ and repeat this
process for $T$ iterations. One can derive the conditional distribution

$\displaystyle
q([{\bm{v}}]_{i}|[{\bm{v}}]_{\setminus\\{i\\}})\propto\frac{\left\|{\bm{v}}\right\|_{2}^{2}}{1+\left\|{\bm{v}}\right\|_{2}^{2}-[{\bm{v}}]_{i}^{2}}\exp\left(-\frac{[{\bm{v}}]_{i}^{2}}{2}\right)$
(17)

and sampling a single random variable from ?? can be done via the inverse
transformation method.222It requires the CDF of
$q([{\bm{v}}]_{i}|[{\bm{v}}]_{\setminus\\{i\\}})$ which is equivalent to
$\frac{\mathrm{erf}\left({[{\bm{v}}]_{i}}/{\sqrt{2}}\right)+1}{2}-\frac{[{\bm{v}}]_{i}\exp\left(-[{\bm{v}}]_{i}^{2}/2\right)}{\sqrt{2\pi}(1+\left\|{\bm{v}}\right\|_{2}^{2}-[{\bm{v}}]_{i}^{2})}$.
We empirically verify that $T=1$ is enough for promising performances. The
running time of Gibbs sampling becomes $\mathcal{O}(m_{1}dT)$ where $m_{1}$
corresponds to the number of independent samples from $q({\bm{v}})$. This is
negligible compared to the feature map construction of CountSketch for
$T=\mathcal{O}(1)$. The pseudo-code for the modified random features of
$A_{1}$ using Gibbs sampling is outlined in ??.

(a) $\mathtt{abalone}$

(b) $\mathtt{ecoli}$

(c) $\mathtt{titanic}$

Figure 1: Generalized condition numbers of arc-cosine random features of Cho &
Saul (2009) and the proposed approach (??) under real-world UCI datasets.

We evaluate approximation quality of the proposed method (??) to that of the
random features Cho & Saul (2009) in ??. In particular, we compute the
condition number (i.e., ratio of the largest and smallest eigenvalues) of
$({\bm{A}}_{1}+\lambda{\bm{I}})^{-1/2}\left(\bm{\Phi}_{1}\bm{\Phi}_{1}^{\top}+\lambda{\bm{I}}\right)({\bm{A}}_{1}+\lambda{\bm{I}})^{-1/2}.$
If $(\bm{\Phi}_{1}\bm{\Phi}_{1}^{\top}+\lambda{\bm{I}})$ is spectrally close
to $({\bm{A}}_{1}+\lambda{\bm{I}})$, then the corresponding condition number
will be close to $1$. We evaluate the condition numbers of those random
features using 3 UCI datasets and set $\lambda=10^{-4}\cdot n$ when $n$ data
points are given. For each dataset, we increase $m_{1}$ from $100$ to
$1{,}000$. Observe that the proposed random features for arc-cosine features
have smaller condition numbers than the previous method for all datasets. We
provide more experimental results that the modified random features can
improve performance on downstream tasks in ??.

###  4.2 Spectral Approximation for the NTK of Two-layer ReLU Network

We are now ready to state a spectral approximation bound for our NTK random
features of a two-layer ReLU network, i.e., $L=1$.

######  Theorem 4.

Given
${\bm{X}}=[{\bm{x}}_{1},\dots,{\bm{x}}_{n}]^{\top}\in{\mathbb{R}}^{n\times
d}$, assume that $\left\|{\bm{x}}_{i}\right\|_{2}=1$ for $i\in[n]$. Let
${\bm{K}}_{\mathrm{NTK}}$ be the NTK of two-layer ReLU network, i.e., $L=1$ in
??, and ${\bm{A}}_{0},{\bm{A}}_{1}$ denote the arc-cosine kernels of $0$-th,
$1$-st order with ${\bm{X}}$, respectively, as in ??. For any
$\lambda\in(0,2\min(\left\|{\bm{A}}_{0}\right\|_{2},\left\|{\bm{A}}_{1}\right\|_{2})]$,
suppose $s_{\lambda}$ is an upper bound of statistical dimensions of both
${\bm{A}}_{0},{\bm{A}}_{1}$. Given $\varepsilon\in(0,1/2)$, $\delta\in(0,1)$,
let $\bm{\Phi}\in{\mathbb{R}}^{n\times(m_{1}+m_{\mathtt{cs}})}$ be the first
output of ?? with $L=1$ and

$m_{0}\geq\frac{48n}{\varepsilon^{2}\lambda}\log\left(\frac{48s_{\lambda}}{\delta}\right),\
\
m_{1}\geq\frac{16}{3}\frac{d\left\|{\bm{X}}\right\|_{2}^{2}}{\varepsilon^{2}\lambda}\log\left(\frac{48s_{\lambda}}{\delta}\right),\
\
m_{\mathtt{cs}}\geq\frac{297}{\varepsilon^{2}\delta}\left(\frac{n}{\lambda+1}\right)^{2}$

Then, with probability at least $1-\delta$, it holds that

$\displaystyle(1-\varepsilon)\left({\bm{K}}_{\mathrm{NTK}}+\lambda{\bm{I}}\right)\preceq\bm{\Phi}\bm{\Phi}^{\top}+\lambda{\bm{I}}\preceq(1+\varepsilon)\left({\bm{K}}_{\mathrm{NTK}}+\lambda{\bm{I}}\right).$
(18)

The proof of ?? is provided in ??. We note that the ridge regularization
parameter typically set to $\lambda=\lambda^{\prime}n$ where
$\lambda^{\prime}$ is a small constant, e.g., $10^{-4}$ Rudi & Rosasco (2016);
Avron et al. (2017b); Geifman et al. (2020). Combining this setting with the
fact that $\left\|{\bm{x}}_{i}\right\|_{2}=1$ yields that
${\bm{s}}_{\lambda}=\mathcal{O}(1)$. Hence, it is enough to choose
$m_{0}={\mathcal{O}}\left(\frac{1}{\varepsilon^{2}}\log\left(\frac{1}{\delta}\right)\right),m_{\mathtt{cs}}=\mathcal{O}\left(\frac{1}{\varepsilon^{2}\delta}\right)$
and
$m_{1}=\mathcal{O}\left(\frac{d}{\varepsilon^{2}}\log\left(\frac{1}{\delta}\right)\right)$
to achieve the spectral approximation in ?? since
$\left\|{\bm{X}}\right\|_{2}^{2}\leq\left\|{\bm{X}}\right\|_{F}^{2}=n$. This
leads us to
$\mathcal{O}\left(\frac{d}{\varepsilon^{2}}\log\left(\frac{1}{\delta}\right)+\frac{1}{\varepsilon^{2}\delta}\right)$
feature dimension which is nearly linear in the input dimension $d$.

The current proof technique cannot be used to generalize the result in ?? to
deeper networks (i.e., $L\geq 2$). For the proof to work, one needs a monotone
property of arc-cosine kernels, i.e., $F_{1}({\bm{X}})\preceq F_{1}({\bm{Y}})$
for ${\bm{X}}\preceq{\bm{Y}}$. However, this property does not hold in
general. Thus, we leave the extension to deeper networks to future work.

##  5 Experiments

(a) $L=1$

(b) $L=2$

(c) $L=3$

Figure 2: Mean squared error of entries in NTK matrix approximated by (1)
gradient of randomly initialized networks (Grad-RF), (2) the naive feature map
construction in Bietti & Mairal (2019) (Naïve-RF) and our algorithm based on a
sketching method under subsampled MNIST dataset.

In this section, we provide experimental results of our method on kernel
approximation and various kernel learning tasks including classification,
regression and active learning.

###  5.1 Kernel Approximation on MNIST Dataset

We first explore ?? for approximating the NTK matrices. We compare to
gradient-based NTK random features Arora et al. (2019b) (Grad-RF) and the
naïve random features without sketching Bietti & Mairal (2019) (Naïve-RF) as
baseline methods. To compute the exact NTK, we randomly choose $n=10{,}000$
data samples from MNIST dataset and evaluate the mean squared error (MSE) of
all approximate entries in NTK. We use the ReLU network with depths $L=1,2,4$.
For Grad-RF, we use an implementation proposed by Novak et al.
(2020).333https://github.com/google/neural-tangents In particular, it returns
an approximate NTK matrix rather than random features because the dimension of
features can be larger than $n$ which loses the computational gain of random
features. For example, gradient of a two-layer, 16-width ReLU network for
MNIST has $127{,}040$ dimension. We vary the network width in
$\\{2^{1}\dots,2^{7}\\}$. For fair dimension comparisons, we report the
expected dimension of Grad-RF. For Naïve-RF, we set $m_{0}=m_{1}$ and
$m_{0}\in\\{2^{1},\dots,2^{7}\\}$. For our method, we set
$m_{0}=m_{1}=m-m_{\mathtt{cs}}$ and
$m_{\mathtt{cs}}\in\\{\frac{m}{10},\dots,\frac{9m}{10}\\}$ for each
$m\in\\{10\cdot 2^{4},\dots,10\cdot 2^{16}\\}$ and report the average MSE of
$9$ different values of $m_{\mathtt{cs}}$. We omit to report the result when
memory overflow causes.

In ??, we observe that our random features achieves the lowest MSE for the
same dimension compared to other competitors. The Grad-RF is the worst method
and this corresponds to observations reported in Arora et al. (2019b), i.e.,
gradient features from a finite width network can degrade practical
performances. As the number of layers $L$ increases, the performance gaps
between Naïve-RF and other methods become large because its dimension grows
exponential in $L$.

###  5.2 Classification on Small-scale UCI Datasets

Next, we run our algorithm under 90 small-scale UCI classification datasets.
The number of data points $n$ varies from $10$ to $5{,}000$. We choose
hyperparameters using validation data and evaluate the test accuracy using
4-fold cross-validation provided in Fernández-Delgado et al.
(2014).444http://persoal.citius.usc.es/manuel.fernandez.delgado/papers/jmlr/data.tar.gz
We also consider the following additional metrics used in Arora et al.
(2019c); Geifman et al. (2020); P90 and P95 are the ratios of datasets where a
classifier reaches at least 90% and 95% of the maximum accuracy, and PMA
(percentage of the maximum accuracy) is the average ratio its accuracy to the
maximum among 90 datasets.

We run ?? with and without Gibbs sampling (GS) (i.e., ??) where the number of
Gibbs iteration is set to $T=1$ throughout all experiments. We also execute
various classifiers including AdaBoost, random forest, $k$-nearest neighbors
and support vector classifier (SVC). For methods running with SVC, we search
the cost value $C$ in $\\{2^{-19},\dots,2^{-3},10^{-7},\dots,10^{3}\\}$ and
choose the best one that achieves the best validation accuracy. We use the
support vector classifier (SVC) for random features methods (ours, RFF). For
methods using SVC, the cost value $C$ is chosen by searching in
$\\{2^{-19},\dots,2^{20}\\}$ that achieves the best validation accuracy. For
our algorithm and RFF, we consider the output dimension $m$ as a
hyperparameter. We search $m$ in $\\{10,20,\dots,100\\}$ for datasets with
$n\leq 600$ and explore $m$ in $\\{20,40,\dots,200\\}$ for datasets with
$n>600$ that achieves the best validation accuracy. For NTK, the network depth
$L$ changes from $1$ to $5$ which is the same setup in Arora et al. (2019c);
Shankar et al. (2020). We also compare test accuracy of fully-connected ReLU
network. We explore the network depth in $\\{1,2,3,4,5\\}$ and width in
$\\{2^{6},\dots,2^{11}\\}$. The ReLU network is trained by Adam optimizer for
$100$ epochs with an initial learning rate $0.1$ and cosine annealing is used
to schedule learning rate.

In ??, the average test accuracy with 95% confidence interval, P90/95 and PMA
are reported. Observe that the NTK achieves the best results while the NTK
Random Features with GS is the second best. The NTK Random Features performs
better than the Random Fourier Features because the NTK is more appropriate
choice compared to the RBF kernel. Finally, our method with GS shows higher
accuracy than that without GS.

Table 1: Results of average test accuracy, P90, P95 and PMA (percentage of the
maximum accuracy) on 90 UCI classification datasets. Bold entries indicate the
best results.

Method | Test Accuracy (%) | P90 | P95 | PMA
---|---|---|---|---
AdaBoost | 76.32 $\pm$ 3.56 | 66.67 | 37.78 | 89.44
Random Forest | 77.46 $\pm$ 3.75 | 79.31 | 57.47 | 90.84
$k$-Nearest Neighbors | 76.95 $\pm$ 3.42 | 72.22 | 43.33 | 90.31
Fully-connected ReLU Network | 81.10 $\pm$ 3.11 | 85.56 | 78.89 | 95.33
Polynomial Kernel | 79.54 $\pm$ 3.41 | 79.31 | 66.67 | 93.31
RBF Kernel | 81.79 $\pm$ 2.95 | 91.11 | 75.56 | 95.97
Random Fourier Features | 81.61 $\pm$ 2.98 | 88.89 | 71.11 | 95.74
NTK | 82.24 $\pm$ 2.94 | 92.22 | 80.00 | 96.53
NTK Random Features | 81.84 $\pm$ 2.89 | 92.22 | 75.56 | 96.16
NTK Random Features with GS | 81.85 $\pm$ 2.98 | 92.22 | 75.56 | 96.05

Table 2: Results of the mean squared errors (MSE) and wall-clock time (sec) on
large-scale UCI regression datasets. We measure the entire time for solving
the kernel ridge regression. Bold entries indicate the best MSE or time for
each dataset. (-) means the Out-of-Memory error.

| $\mathtt{MillionSongs}$ | $\mathtt{WorkLoads}$ | $\mathtt{Protein}$ | $\mathtt{SuperConduct}$
---|---|---|---|---
# of Training Data $n$ | 467,315 | 179,585 | 39,617 | 19,077
| MSE | Time (s) | MSE | Time (s) | MSE | Time (s) | MSE | Time (s)
RBF Kernel | (-) | (-) | (-) | (-) | 112.82 | 110.2 | 2239.83 | 19.5
Random Fourier Features | 108.50 | 159 | 7.05$\times 10^{4}$ | 63.7 | 81.98 | 14.6 | 1175.13 | 7.1
NTK | (-) | (-) | (-) | (-) | 90.03 | 243 | 513.25 | 51.9
NTK Random Features | 80.77 | 149.7 | 2.43$\times 10^{4}$ | 50.9 | 90.28 | 13.9 | 527.58 | 7.9
NTK Random Features with GS | 80.94 | 168.5 | $\mathbf{2.38\times 10^{4}}$ | 53.0 | 85.99 | 16.1 | 492.78 | 12.7

###  5.3 Regression on Large-scale UCI Datasets

We also demonstrate the computational efficiency of our method using $4$
large-scale UCI regression datasets. In particular, we consider kernel ridge
regression (KRR) problem. For a kernel function
$K:{\mathbb{R}}^{d}\times{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$, KRR problem
can be formulated as

$\displaystyle
y_{\mathrm{test}}=K({\bm{x}}_{\mathrm{test}},{\bm{X}})\left({\bm{K}}+\lambda
n{\bm{I}}\right)^{-1}{\bm{y}}.$ (19)

where
${\bm{X}}=[{\bm{x}}_{1},\dots,{\bm{x}}_{n}]^{\top}\in{\mathbb{R}}^{n\times d}$
is training data, ${\bm{y}}\in{\mathbb{R}}^{n}$ is training label,
${\bm{x}}_{\mathrm{test}}\in{\mathbb{R}}^{d}$ is a test data,
$y_{\mathrm{test}}$ is a predicted label,
$[{\bm{K}}]_{ij}=K({\bm{x}}_{i},{\bm{x}}_{j})$ and
$K({\bm{x}}_{\mathrm{test}},{\bm{X}})=[K({\bm{x}}_{\mathrm{test}},{\bm{x}}_{1}),\dots,K({\bm{x}}_{\mathrm{test}},{\bm{x}}_{n})]\in{\mathbb{R}}^{1\times
n}$. Note that solving the problem can require $\mathcal{O}(n^{3})$ time
complexity due to the matrix inversion in general. Consider a feature map
$\bm{\Phi}:{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}^{m}$ can approximate the
kernel $K$ such that
$K({\bm{x}},{\bm{x}})\approx\hat{K}({\bm{x}},{\bm{x}})=\left\langle{\bm{\Phi}({\bm{x}}),\bm{\Phi}{\bm{x}})}\right\rangle$.
Then, the computation can be amortized as
$y_{\mathrm{test}}=\bm{\Phi}({\bm{x}}_{\mathrm{test}})^{\top}\left(\bm{\Phi}({\bm{X}})^{\top}\bm{\Phi}({\bm{X}})+\lambda
n{\bm{I}}_{m}\right)^{-1}\bm{\Phi}({\bm{X}})^{\top}{\bm{y}}$ which requires
$\mathcal{O}(nm^{2}+m^{3})$ time to run. When $n\gg m$, this is much efficient
than solving the problem with the exact kernel.

We compare our methods to NTK, RBF and RFF. For ours and RFF, we choose the
output dimension to $m=10{,}000$ for all datasets, which is much smaller than
the number of data samples $n$. In ??, we report the wall-clock times and mean
squared errors (MSE) of test prediction. We face Out-of-Memory errors when
running kernel methods using $\mathtt{MillionSongs}$ and $\mathtt{WorkLoads}$
datasets. Observe that our random features are significantly faster than NTK,
e.g., up to $\times$17 speedup for $\mathtt{Protein}$ dataset, without
performance loss. We also verify that the NTK features achieve lower MSE than
RFF only for $\mathtt{Protein}$ but it outperforms with a huge gap for the
rest of the datasets.

###  5.4 Active Learning on MNIST Dataset

We finally apply the proposed method to active learning using MNIST dataset.
The goal is to select training data of fixed size $k$ that maximizes the
performance. Recently, Shoham & Avron (2020) suggested an active learning
strategy based on the NTK. They propose a novel criteria that can be an upper
bound of statistical risk for general kernel learning and present an algorithm
that greedily minimizes this criteria with the NTK. Their greedy process
begins with an empty set and iteratively appends singleton that minimizes the
proposed risk bound. It takes $\mathcal{O}(n^{2}k^{2})$ time to obtain $k$
data points which equals to the budget size to acquire labels which can be
prohibitive if $n$ is large.

Figure 3: Test accuracy (%) of random selection versus greedy minimization of
statistical bound using NTK random features.

Motivated by this, we apply the proposed NTK random features to their greedy
algorithm that can improve the running time. Recall that our random features
builds an approximation
${\bm{K}}^{\prime}=\bm{\Phi}\bm{\Phi}^{\top}\approx{\bm{K}}$ where
$\bm{\Phi}\in{\mathbb{R}}^{n\times m}$. Under certain parameter regimes, the
low rank structure of ${\bm{K}}^{\prime}$ can be used to implement a faster
version of the greedy algorithm. Specifically, after $\mathcal{O}(nm^{2})$
preprocessing, the cost per iteration of the greedy algorithm can be reduced
to $\mathcal{O}(nmj^{2})$, and the cost of finding the design of size $k$ to
$\mathcal{O}(nm(k^{3}+m))$. We provide more details in the supplementary
material.

?? illustrates performance of greedy minimization using our NTK random
features compared to randomly generated designs under MNIST dataset. We use a
$4$-layer fully-connected ReLU network with width $1{,}000$ and the dimension
of NTK random features is $m=10{,}000$. We clearly see that using the random
NTK features we can generate much better designs than randomly chosen data
points. It justifies that our random features plays a crucial role for active
learning tasks. We expect that the proposed method can be applied to various
machine learning applications with remarkable performance and computation
gains.

##  6 Conclusion

In this work, we propose an efficient algorithm for generating random features
of the Neural Tangent Kernel (NTK). We utilize TensorSketch transform combined
with the arc-cosine random features with an importance sampling. We also
provide spectral approximation bound to the NTK with layer $2$. Our
experiments validate the effectiveness of the proposed methods. We believe
that our method would be a broad interest both in theoretical and practical
domains.

## References

  * Ahle et al. (2020) Ahle, T. D., Kapralov, M., Knudsen, J. B., Pagh, R., Velingker, A., Woodruff, D. P., and Zandieh, A.  Oblivious sketching of high-degree polynomial kernels.  In _Symposium on Discrete Algorithms (SODA)_ , 2020.
  * Allen-Zhu et al. (2019) Allen-Zhu, Z., Li, Y., and Song, Z.  A convergence theory for deep learning via over-parameterization.  In _International Conference on Machine Learning (ICML)_ , 2019.
  * Arora et al. (2019a) Arora, S., Du, S., Hu, W., Li, Z., and Wang, R.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In _International Conference on Machine Learning (ICML)_ , 2019a.
  * Arora et al. (2019b) Arora, S., Du, S. S., Hu, W., Li, Z., Salakhutdinov, R. R., and Wang, R.  On exact computation with an infinitely wide neural net.  In _Neural Information Processing Systems (NeurIPS)_ , 2019b.
  * Arora et al. (2019c) Arora, S., Du, S. S., Li, Z., Salakhutdinov, R., Wang, R., and Yu, D.  Harnessing the Power of Infinitely Wide Deep Nets on Small-data Tasks.  In _International Conference on Learning Representations (ICLR)_ , 2019c.
  * Avron et al. (2014) Avron, H., Nguyen, H., and Woodruff, D.  Subspace embeddings for the polynomial kernel.  In _Neural Information Processing Systems (NeurIPS)_ , 2014.
  * Avron et al. (2017a) Avron, H., Clarkson, K. L., and Woodruff, D. P.  Faster Kernel Ridge Regression Using Sketching and Preconditioning.  In _SIAM Journal on Matrix Analysis and Applications (SIMAX)_ , 2017a.
  * Avron et al. (2017b) Avron, H., Kapralov, M., Musco, C., Musco, C., Velingker, A., and Zandieh, A.  Random Fourier features for kernel ridge regression: Approximation bounds and statistical guarantees.  In _International Conference on Machine Learning (ICML)_ , 2017b.
  * Bietti & Mairal (2019) Bietti, A. and Mairal, J.  On the inductive bias of neural tangent kernels.  In _Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Bossard et al. (2014) Bossard, L., Guillaumin, M., and Van Gool, L.  Food-101–mining discriminative components with random forests.  In _Proceedings of the European Conference on Computer Vision(ECCV)_ , 2014.
  * Cao & Gu (2019) Cao, Y. and Gu, Q.  Generalization bounds of stochastic gradient descent for wide and deep neural networks.  In _Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Charikar et al. (2002) Charikar, M., Chen, K., and Farach-Colton, M.  Finding frequent items in data streams.  In _International Colloquium on Automata, Languages, and Programming (ICAMP)_ , 2002.
  * Chen & Xu (2021) Chen, L. and Xu, S.  Deep neural tangent kernel and laplace kernel have the same RKHS.  In _International Conference on Learning Representations (ICLR)_ , 2021.
  * Chizat et al. (2019) Chizat, L., Oyallon, E., and Bach, F.  On lazy training in differentiable programming.  In _Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Cho & Saul (2009) Cho, Y. and Saul, L.  Kernel methods for deep learning.  In _Neural Information Processing Systems (NeurIPS)_ , 2009.
  * Daniely et al. (2016) Daniely, A., Frostig, R., and Singer, Y.  Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity.  In _Neural Information Processing Systems (NeurIPS)_ , 2016.
  * Du et al. (2019) Du, S. S., Zhai, X., Poczos, B., and Singh, A.  Gradient descent provably optimizes over-parameterized neural networks.  In _International Conference on Learning Representations (ICLR)_ , 2019.
  * Everingham et al. (2010) Everingham, M., Van Gool, L., Williams, C. K., Winn, J., and Zisserman, A.  The pascal visual object classes (voc) challenge.  _International Journal of Computer Vision_ , 2010.
  * Fan & Wang (2020) Fan, Z. and Wang, Z.  Spectra of the conjugate kernel and neural tangent kernel for linear-width neural networks.  In _Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Fei-Fei et al. (2004) Fei-Fei, L., Fergus, R., and Perona, P.  Learning generative visual models from few training examples: An incremental bayesian approach tested on 101 object categories.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshop_ , 2004.
  * Fernández-Delgado et al. (2014) Fernández-Delgado, M., Cernadas, E., Barro, S., and Amorim, D.  Do we need hundreds of classifiers to solve real world classification problems? In _Journal of Machine Learning Research (JMLR)_ , 2014.
  * Geifman et al. (2020) Geifman, A., Yadav, A., Kasten, Y., Galun, M., Jacobs, D., and Basri, R.  On the similarity between the laplace and neural tangent kernels.  In _Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Goyal et al. (2019) Goyal, P., Mahajan, D., Gupta, A., and Misra, I.  Scaling and benchmarking self-supervised visual representation learning.  In _International Conference on Computer Vision (ICCV)_ , 2019.
  * Han et al. (2020) Han, I., Avron, H., and Shin, J.  Polynomial Tensor Sketch for Element-wise Function of Low-Rank Matrix.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * He et al. (2016) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 2016.
  * Jacot et al. (2018) Jacot, A., Gabriel, F., and Hongler, C.  Neural tangent kernel: Convergence and generalization in neural networks.  In _Neural Information Processing Systems (NeurIPS)_ , 2018.
  * Jacot et al. (2020) Jacot, A., Simsek, B., Spadaro, F., Hongler, C., and Gabriel, F.  Implicit regularization of random feature models.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Khosla et al. (2011) Khosla, A., Jayadevaprakash, N., Yao, B., and Li, F.-F.  Novel dataset for fine-grained image categorization: Stanford dogs.  In _Proceedings of CVPR Workshop on Fine-Grained Visual Categorization (FGVC)_ , 2011.
  * Krizhevsky (2009) Krizhevsky, A.  Learning multiple layers of features from tiny images.  2009\.
  * Lee et al. (2020) Lee, J., Shen, R., Song, Z., Wang, M., and Yu, Z.  Generalized Leverage Score Sampling for Neural Networks.  In _Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Mei et al. (2018) Mei, S., Montanari, A., and Nguyen, P.-M.  A mean field view of the landscape of two-layer neural networks.  _Proceedings of the National Academy of Sciences_ , 2018.
  * Musco & Musco (2017) Musco, C. and Musco, C.  Recursive Sampling for the Nyström Method.  In _Neural Information Processing Systems (NeurIPS)_ , 2017.
  * Neyshabur et al. (2019) Neyshabur, B., Tomioka, R., and Srebro, N.  In search of the real inductive bias: On the role of implicit regularization in deep learning.  In _International Conference on Learning Representations (ICLR)_ , 2019.
  * Nilsback & Zisserman (2008) Nilsback, M.-E. and Zisserman, A.  Automated flower classification over a large number of classes.  In _2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing_, 2008.
  * Novak et al. (2020) Novak, R., Xiao, L., Hron, J., Lee, J., Alemi, A. A., Sohl-Dickstein, J., and Schoenholz, S. S.  Neural tangents: Fast and easy infinite neural networks in python.  In _International Conference on Learning Representations (ICLR)_ , 2020.
  * Pennington et al. (2015) Pennington, J., Yu, F. X. X., and Kumar, S.  Spherical random features for polynomial kernels.  In _Neural Information Processing Systems (NeurIPS)_ , 2015.
  * Pham & Pagh (2013) Pham, N. and Pagh, R.  Fast and scalable polynomial kernels via explicit feature maps.  In _Conference on Knowledge Discovery and Data Mining (KDD)_ , 2013\.
  * Rahimi & Recht (2009) Rahimi, A. and Recht, B.  Random Features for Large-Scale Kernel Machines.  In _Neural Information Processing Systems (NeurIPS)_ , 2009.
  * Rudi & Rosasco (2016) Rudi, A. and Rosasco, L.  Generalization properties of learning with random features.  In _Neural Information Processing Systems (NeurIPS)_ , 2016.
  * Shankar et al. (2020) Shankar, V., Fang, A., Guo, W., Fridovich-Keil, S., Ragan-Kelley, J., Schmidt, L., and Recht, B.  Neural kernels without tangents.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Shoham & Avron (2020) Shoham, N. and Avron, H.  Experimental Design for Overparameterized Learning with Application to Single Shot Deep Active Learning.  In _arXiv preprint arXiv:2009.12820_ , 2020.
  * Welinder et al. (2010) Welinder, P., Branson, S., Mita, T., Wah, C., Schroff, F., Belongie, S., and Perona, P.  Caltech-UCSD Birds 200.  Technical Report CNS-TR-2010-001, California Institute of Technology, 2010\.
  * Woodruff & Zandieh (2020) Woodruff, D. P. and Zandieh, A.  Near Input Sparsity Time Kernel Embeddings via Adaptive Sampling.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Zancato et al. (2020) Zancato, L., Achille, A., Ravichandran, A., Bhotika, R., and Soatto, S.  Predicting training time without training.  In _Neural Information Processing Systems (NeurIPS)_ , 2020.

##  Appendix A Proof of Theorems

###  A.1 Proof of ??

See 1

Proof of ??. For fixed ${\bm{x}},{\bm{x}}^{\prime}\in\mathbb{S}^{d-1}$ and
$\ell=0,\dots,L$, we denote the estimate error as

$\displaystyle\Delta_{\ell}:=\max_{({\bm{x}}_{1},{\bm{x}}_{2})\in\\{({\bm{x}},{\bm{x}}^{\prime}),({\bm{x}},{\bm{x}}),({\bm{x}}^{\prime},{\bm{x}}^{\prime})\\}}\left|\left\langle{\bm{\Phi}^{(\ell)}({\bm{x}}_{1}),\bm{\Phi}^{(\ell)}({\bm{x}}_{2})}\right\rangle-
K_{\mathrm{NTK}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|$

and note that $\Delta_{0}=0$. Recall that

$\displaystyle
K_{\mathrm{NTK}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})=K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})+\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\cdot{K}_{\mathrm{NTK}}^{(\ell-1)}({\bm{x}},{\bm{x}}^{\prime})$
(20)

where

$\displaystyle\dot{K}_{\mathrm{NTK}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})$
$\displaystyle:=1-\frac{1}{\pi}\cos^{-1}\left(\frac{K_{\mathrm{NNGP}}^{(\ell-1)}({\bm{x}},{\bm{x}}^{\prime})}{\sqrt{K_{\mathrm{NNGP}}^{(\ell-1)}({\bm{x}},{\bm{x}})K_{\mathrm{NNGP}}^{(\ell-1)}({\bm{x}}^{\prime},{\bm{x}}^{\prime})}}\right)$
(21)
$\displaystyle\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})$
$\displaystyle:=f\left(\frac{K_{\mathrm{NNGP}}^{(\ell-1)}({\bm{x}},{\bm{x}}^{\prime})}{\sqrt{K_{\mathrm{NNGP}}^{(\ell-1)}({\bm{x}},{\bm{x}}^{\prime})K_{\mathrm{NNGP}}^{(\ell-1)}({\bm{x}}^{\prime},{\bm{x}}^{\prime})}}\right)$
(22)

and $f(x)=\frac{1}{\pi}\left(\sqrt{1-x^{2}}+(\pi-\cos^{-1}(x))x\right)$ for
$x\in[-1,1]$.

We use the recursive relation to approximate:

$\displaystyle\left\langle{\bm{\Phi}^{(\ell)}({\bm{x}}),\bm{\Phi}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle$
$\displaystyle=\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}),\bm{\Psi}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle+\left\langle{\bm{\Gamma}^{(\ell)}({\bm{x}}),\bm{\Gamma}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle$
$\displaystyle\approx\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}),\bm{\Psi}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle+\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}),\bm{\Lambda}^{(\ell)}({\bm{x}}^{\prime})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}^{\prime})}\right\rangle$
$\displaystyle\approx
K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})+\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\cdot
K_{\mathrm{NTK}}^{(\ell-1)}({\bm{x}},{\bm{x}}^{\prime})=K^{(\ell)}_{\mathrm{NTK}}({\bm{x}},{\bm{x}}^{\prime}).$

For notational simplicity, we define the following events:

$\displaystyle\mathcal{E}_{\bm{\Psi}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime},\varepsilon)$
$\displaystyle:=\left\\{\left|\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}),\bm{\Psi}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle-
K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\right|\leq\varepsilon\right\\},$
(23)
$\displaystyle\mathcal{E}_{\bm{\Lambda}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime},\varepsilon)$
$\displaystyle:=\left\\{\left|\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}),\bm{\Lambda}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle-\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\right|\leq\varepsilon\right\\},$
(24)
$\displaystyle\mathcal{E}_{\bm{\Gamma}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime},\varepsilon)$
$\displaystyle:=\left\\{\left|\left\langle{\bm{\Gamma}^{(\ell)}({\bm{x}}),\bm{\Gamma}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle-\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}),\bm{\Lambda}^{(\ell)}({\bm{x}}^{\prime})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}^{\prime})}\right\rangle\right|\leq\varepsilon\right\\}$
(25)

and
$\mathcal{E}_{\bm{\Omega}}^{(\ell)}(\varepsilon):=\mathcal{E}_{\bm{\Omega}}^{(\ell)}({\bm{x}},{\bm{x}},\varepsilon)\bigcap\mathcal{E}_{\bm{\Omega}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime},\varepsilon)\bigcap\mathcal{E}_{\bm{\Omega}}^{(\ell)}({\bm{x}}^{\prime},{\bm{x}}^{\prime},\varepsilon)$
for ${\bm{\Omega}}=\\{\bm{\Psi},\bm{\Lambda},\bm{\Gamma}\\}$. Our proof is
based on the following claims:

######  Claim 1.

There exists a constant $C_{1}>0$ such that if $m_{1}\geq
C_{1}\frac{L^{6}}{\varepsilon^{4}}\log\left(\frac{L}{\delta}\right)$ then

$\displaystyle\Pr\left(\mathcal{E}_{\bm{\Psi}}^{(\ell)}\left(\frac{\varepsilon^{2}}{32L^{2}}\right)\right)\geq
1-\frac{\delta}{3L}.$

The proof of ?? is provided in ??.

######  Claim 2.

There exists a constant $C_{0}>0$ such that if $m_{0}\geq
C_{0}\frac{L^{2}}{\varepsilon^{2}}\log\left(\frac{L}{\delta}\right)$ then

$\displaystyle\Pr\left(\mathcal{E}_{\bm{\Lambda}}^{(\ell)}\left(\frac{3\varepsilon}{8L}\right)\bigg{|}\mathcal{E}_{\bm{\Psi}}^{(\ell)}\left(\frac{\varepsilon^{2}}{32L^{2}}\right)\right)\geq
1-\frac{\delta}{3L}.$

The proof of ?? is provided in ??.

######  Claim 3.

There exists a constant $C_{2}>0$ such that if $m_{\mathtt{cs}}\geq
C_{2}\frac{L^{3}}{\varepsilon^{2}\delta}$ then

$\displaystyle\Pr\left(\mathcal{E}_{\bm{\Gamma}}^{(\ell)}\left(\frac{\varepsilon}{8L}\left(\ell+\Delta_{\ell-1}\right)\right)\bigg{|}\mathcal{E}_{\bm{\Lambda}}^{(\ell)}\left(\frac{3\varepsilon}{8L}\right)\bigcap\mathcal{E}_{\bm{\Psi}}^{(\ell)}\left(\frac{\varepsilon^{2}}{32L^{2}}\right)\right)\geq
1-\frac{\delta}{3L}.$

The proof of ?? is provided in ??.

Combining ??, ?? and ??, we have

$\displaystyle\Pr\left(\mathcal{E}_{\bm{\Gamma}}^{(\ell)}\left(\frac{\varepsilon}{8L}\left(\ell+\Delta_{\ell}\right)\right)~{}\bigcap~{}\mathcal{E}_{\bm{\Lambda}}^{(\ell)}\left(\frac{3\varepsilon}{8L}\right)~{}\bigcap~{}\mathcal{E}_{\bm{\Psi}}^{(\ell)}\left(\frac{\varepsilon^{2}}{32L^{2}}\right)\right)\geq
1-\frac{\delta}{L}.$

Next, we claim that the above event can provide the recurrence relation of
$\Delta_{\ell}$ as described in below.

######  Claim 4.

For $\varepsilon\in(0,1)$, $L\geq 1$, if event

$\displaystyle\mathcal{E}_{\bm{\Gamma}}^{(\ell)}\left(\frac{\varepsilon}{8L}\left(\ell+\Delta_{\ell}\right)\right)~{}\bigcap~{}\mathcal{E}_{\bm{\Lambda}}^{(\ell)}\left(\frac{3\varepsilon}{8L}\right)~{}\bigcap~{}\mathcal{E}_{\bm{\Psi}}^{(\ell)}\left(\frac{\varepsilon^{2}}{32L^{2}}\right)$

holds for $\ell\in[L]$ then

$\displaystyle\Delta_{\ell}\leq\left(1+\frac{\varepsilon}{2L}\right)\Delta_{\ell-1}+\frac{\varepsilon}{2L}\ell+\frac{\varepsilon^{2}}{32L^{2}}.$
(26)

The proof of ?? is provided in ??.

Applying union bound on ?? for all $\ell\in[L]$ and solving the recurrence, we
obtain that with probabilty at least $1-\delta$

$\displaystyle\Delta_{L}\leq\left(1+\frac{\varepsilon}{2L}+\frac{\varepsilon^{2}}{32L^{2}}\right)\frac{\left(1+\frac{\varepsilon}{2L}\right)^{L}-1}{\frac{\varepsilon}{2L}}-L.$
(27)

When $L=1$, the statement in ?? holds since
$\Delta_{1}\leq\frac{\varepsilon}{2}+\frac{\varepsilon^{2}}{32}\leq\varepsilon(1+\frac{\varepsilon}{2})^{2}+\varepsilon$
for $\varepsilon\in(0,1)$. Assume that $L\geq 2$, we obtain

$\displaystyle\Delta_{L}$
$\displaystyle\leq\left(1+\frac{\varepsilon}{2L}+\frac{\varepsilon^{2}}{32L^{2}}\right)L\left(1+\frac{\varepsilon}{2L}\right)^{L-1}-L$
(28) $\displaystyle\leq
L\left(\left(1+\frac{\varepsilon}{L}\right)\left(1+\frac{\varepsilon}{2L}\right)^{L-1}-1\right)$
(29) $\displaystyle\leq
L\left(\left(1+\frac{\varepsilon}{L}\right)\left(\frac{\varepsilon}{2L}(L-1)\left(1+\frac{\varepsilon}{2L}\right)^{L-2}+1\right)-1\right)$
(30)
$\displaystyle=L\left(1+\frac{\varepsilon}{L}\right)\frac{\varepsilon}{2L}(L-1)\left(1+\frac{\varepsilon}{2L}\right)^{L-2}+\varepsilon$
(31) $\displaystyle\leq
L\left(1+\frac{\varepsilon}{2}\right)\frac{\varepsilon}{2}\left(1+\varepsilon\right)+\varepsilon$
(32) $\displaystyle\leq
L\varepsilon\left(1+\frac{\varepsilon}{2}\right)^{2}+\varepsilon$ (33)

where inequalities in the first and third line are from the fact that
$\frac{(1+x)^{k}-1}{x}\leq kx(1+x)^{k-1}$ for $x>0,k\geq 1$, the fifth line
follows from that

$\displaystyle\left(1+\frac{\varepsilon}{2L}\right)^{L-2}\leq\exp\left(\frac{\varepsilon}{2L}(L-2)\right)\leq\exp\left(\frac{\varepsilon}{2}\right)\leq
1+\varepsilon.$

Hence, we conclude that

$\displaystyle\Pr\left(\Delta_{L}\leq
L\varepsilon\left(1+\frac{\varepsilon}{2}\right)^{2}+\varepsilon\right)\geq
1-\delta.$ (34)

This completes the proof of ??. ∎

####  A.1.1 Proof of ??

See 1

Proof of ??. Recall that
$\mathcal{E}_{\bm{\Psi}}^{(\ell)}\left(\frac{\varepsilon^{2}}{32L^{2}}\right)$
is equivalent to

$\displaystyle\left|\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}_{1}),\bm{\Psi}^{(\ell)}({\bm{x}}_{2})}\right\rangle-
K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|\leq\frac{\varepsilon^{2}}{32L^{2}}$

for
$({\bm{x}}_{1},{\bm{x}}_{2})\in\\{({\bm{x}},{\bm{x}}^{\prime}),({\bm{x}},{\bm{x}}),({\bm{x}}^{\prime},{\bm{x}}^{\prime})\\}$.
The proof is directly followed by below lemma.

######  Lemma 1 (Corollary 16 in (Daniely et al., 2016)).

Given ${\bm{x}},{\bm{x}}^{\prime}\in\mathbb{R}^{d}$ such that
$\left\|{\bm{x}}\right\|=\left\|{\bm{x}}^{\prime}\right\|=1$, consider a ReLU
network with $L$ layers. For $\delta,\varepsilon\in(0,1)$, there exist
constants $C_{1},C_{2}>0$ such that $m_{1}\geq
C_{1}\frac{L^{2}}{\varepsilon_{1}^{2}}\log\left(\frac{L}{\delta}\right)$,
$\varepsilon_{1}\leq\min(c_{2},\frac{1}{L})$ and for
$({\bm{x}}_{1},{\bm{x}}_{2})\in\\{({\bm{x}},{\bm{x}}^{\prime}),({\bm{x}},{\bm{x}}),({\bm{x}}^{\prime},{\bm{x}}^{\prime})\\}$
and $\ell\in[L]$ it holds that with probability at least $1-\delta$

$\displaystyle\left|\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}_{1}),\bm{\Psi}^{(\ell)}({\bm{x}}_{2})}\right\rangle-
K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|\leq\varepsilon_{1}.$
(35)

In ??, setting $\varepsilon_{1}=\frac{\varepsilon^{2}}{32L^{2}}$ for
$\varepsilon\in(0,1)$ provides the result. This completes the proof of ??. ∎

####  A.1.2 Proof of ??

See 2

Proof of ??. Recall that

$\displaystyle\mathcal{E}_{\bm{\Lambda}}^{(\ell)}\left(\frac{3\varepsilon}{8L}\right)=\left\\{\left|\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}_{1}),\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})}\right\rangle-\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|\leq\frac{3\varepsilon}{8L}\right\\}$
$\displaystyle\mathcal{E}_{\bm{\Psi}}^{(\ell)}\left(\frac{\varepsilon^{2}}{32L^{2}}\right)=\left\\{\left|\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}_{1}),\bm{\Psi}^{(\ell)}({\bm{x}}_{2})}\right\rangle-
K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|\leq\frac{\varepsilon^{2}}{32L^{2}}\right\\}$

for
$({\bm{x}}_{1},{\bm{x}}_{2})\in\\{({\bm{x}},{\bm{x}}^{\prime}),({\bm{x}},{\bm{x}}),({\bm{x}}^{\prime},{\bm{x}}^{\prime})\\}$.
The proof is a direct consequence of the following lemma.

######  Lemma 2 (Lemma E.5 in Arora et al. (2019b)).

Given ${\bm{x}},{\bm{x}}^{\prime}\in{\mathbb{R}}^{d}$, $\ell\in[L]$ and
$\varepsilon_{2}\in(0,1)$, assume that

$\displaystyle\left|\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}),\bm{\Psi}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle-
K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\right|\leq\frac{\varepsilon_{2}^{2}}{2}.$
(36)

Then, it holds that with probability at least $1-\delta_{2}$

$\displaystyle\left|\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}),\bm{\Lambda}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle-\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\right|\leq\varepsilon_{2}+\sqrt{\frac{2}{m_{0}}\log\left(\frac{6}{\delta_{2}}\right)}.$
(37)

In ??, we choose
$\varepsilon_{2}=\frac{\varepsilon}{4L},\delta_{2}=\frac{\delta}{3L}$ and
$m_{0}\geq\frac{128L^{2}}{\varepsilon^{2}}\log\left(\frac{18L}{\delta}\right)$
for $\varepsilon,\delta\in(0,1)$ to obtain that

$\displaystyle\Pr\left(\left|\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}),\bm{\Lambda}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle-\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\right|\leq\frac{3\varepsilon}{8L}~{}\bigg{|}~{}\left|\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}),\bm{\Psi}^{(\ell)}({\bm{x}}^{\prime})}\right\rangle-
K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}},{\bm{x}}^{\prime})\right|\leq\frac{\varepsilon^{2}}{32L^{2}}\right)\geq
1-\frac{\delta}{3L}.$

This completes the proof of ??. ∎

####  A.1.3 Proof of ??

See 3 Proof of ??. Recall that
$\mathcal{E}_{\bm{\Gamma}}^{(\ell)}\left(\frac{\varepsilon}{8L}(\ell+\Delta_{\ell-1})\right)$
is equivalent to

$\displaystyle\left|\left\langle{\bm{\Gamma}^{(\ell)}({\bm{x}}_{1}),\bm{\Gamma}^{(\ell)}({\bm{x}}_{2})}\right\rangle-\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}_{1})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1}),\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})}\right\rangle\right|\leq\frac{\varepsilon}{8L}(\ell+\Delta_{\ell-1})$

for
$({\bm{x}}_{1},{\bm{x}}_{2})\in\\{({\bm{x}},{\bm{x}}^{\prime}),({\bm{x}},{\bm{x}}),({\bm{x}}^{\prime},{\bm{x}}^{\prime})\\}$.
The proof is based on the following lemma that provides an upper bound on
variance of the CountSketch transform.

######  Lemma 3.

Given ${\bm{x}},{\bm{x}}^{\prime}\in{\mathbb{R}}^{m}$ and
${\bm{y}},{\bm{y}}^{\prime}\in{\mathbb{R}}^{m^{\prime}}$, let
$\mathcal{C}_{1}:{\mathbb{R}}^{m}\rightarrow{\mathbb{R}}^{m_{\mathtt{cs}}},\mathcal{C}_{2}:{\mathbb{R}}^{m^{\prime}}\rightarrow{\mathbb{R}}^{m_{\mathtt{cs}}}$
be two independent CountSketch transforms for some $m_{\mathtt{cs}}>0$. Denote

$\displaystyle\bm{\Gamma}:=\mathrm{FFT}^{-1}\left(\mathrm{FFT}\left(\mathcal{C}_{1}({\bm{x}})\right)\odot\mathrm{FFT}\left(\mathcal{C}_{2}({\bm{y}})\right)\right),\quad\bm{\Gamma}^{\prime}:=\mathrm{FFT}^{-1}\left(\mathrm{FFT}\left(\mathcal{C}_{1}({\bm{x}}^{\prime})\right)\odot\mathrm{FFT}\left(\mathcal{C}_{2}({\bm{y}}^{\prime})\right)\right).$
(38)

Then, it holds that with probability at least $1-\delta_{3}$

$\displaystyle\left|\left\langle{\bm{\Gamma},\bm{\Gamma}^{\prime}}\right\rangle-\left\langle{{\bm{x}}\otimes{\bm{y}},{\bm{x}}^{\prime}\otimes{\bm{y}}^{\prime}}\right\rangle\right|\leq\sqrt{\frac{11}{\delta_{3}m_{\mathtt{cs}}}}\left\|{\bm{x}}\right\|_{2}\left\|{\bm{x}}^{\prime}\right\|_{2}\left\|{\bm{y}}\right\|_{2}\left\|{\bm{y}}^{\prime}\right\|_{2}.$
(39)

The proof of ?? is provided in ??. In ??, we choose
$\delta_{3}=\frac{\delta}{3L}$,
$m_{\mathtt{cs}}\geq\frac{33L(8L+3\varepsilon)^{2}}{\varepsilon^{2}\delta}$
for $\varepsilon,\delta\in(0,1)$ to satisfies that

$\displaystyle\sqrt{\frac{11}{\delta_{3}m_{\mathtt{cs}}}}\leq\frac{\frac{\varepsilon}{8L}}{1+\frac{3\varepsilon}{8L}}.$
(40)

Then, with probability at least $1-\frac{\delta}{3L}$ we have

$\displaystyle\left|\left\langle{\bm{\Gamma}^{(\ell)}({\bm{x}}_{1}),\bm{\Gamma}^{(\ell)}({\bm{x}}_{2})}\right\rangle-\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}_{1})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1}),\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})}\right\rangle\right|$
(41)
$\displaystyle\leq\frac{\frac{\varepsilon}{8L}}{1+\frac{3\varepsilon}{8L}}\left\|\bm{\Lambda}^{(\ell)}({\bm{x}}_{1})\right\|_{2}\left\|\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})\right\|_{2}\left\|\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1})\right\|_{2}\left\|\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})\right\|_{2}$
(42)
$\displaystyle\leq\frac{\frac{\varepsilon}{8L}}{1+\frac{3\varepsilon}{8L}}\left(1+\frac{3\varepsilon}{8L}\right)\left\|\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1})\right\|_{2}\left\|\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})\right\|_{2}$
(43)
$\displaystyle\leq\frac{\frac{\varepsilon}{8L}}{1+\frac{3\varepsilon}{8L}}\left(1+\frac{3\varepsilon}{8L}\right)\left(\ell+\Delta_{\ell-1}\right)=\frac{\varepsilon}{8L}\left(\ell+\Delta_{\ell-1}\right)$
(44)

where the second inequlity holds from the fact that
$\mathcal{E}_{\bm{\Lambda}}^{(\ell)}\left(\frac{3\varepsilon}{8L}\right)$
implies that for ${\bm{x}}^{\prime}\in\\{{\bm{x}}_{1},{\bm{x}}_{2}\\}$

$\displaystyle\left\|\bm{\Lambda}^{(\ell)}({\bm{x}}^{\prime})\right\|_{2}^{2}\leq\dot{K}_{\mathrm{NNGP}}^{(\ell-1)}({\bm{x}}^{\prime},{\bm{x}}^{\prime})+\frac{3\varepsilon}{8L}=1+\frac{3\varepsilon}{8L}$
(45)

since
$\dot{K}_{\mathrm{NNGP}}^{(\ell-1)}({\bm{x}}^{\prime},{\bm{x}}^{\prime})=1$
and the third one follows from that
$K_{\mathrm{NTK}}^{(\ell-1)}({\bm{x}},{\bm{y}})\leq\ell$ for
${\bm{x}},{\bm{y}}\in\mathbb{S}^{d-1},\ell\geq 1$ and

$\displaystyle\left\|\bm{\Phi}^{(\ell-1)}({\bm{x}}^{\prime})\right\|_{2}^{2}\leq
K_{\mathrm{NTK}}^{(\ell-1)}({\bm{x}}^{\prime},{\bm{x}}^{\prime})+\Delta_{\ell-1}\leq\ell+\Delta_{\ell-1}.$
(46)

This completes the proof of ??. ∎

####  A.1.4 Proof of ??

See 4 Proof of ??. Recall that

$\displaystyle\Delta_{\ell}:=\max_{({\bm{x}}_{1},{\bm{x}}_{2})\in\\{({\bm{x}},{\bm{x}}^{\prime}),({\bm{x}},{\bm{x}}),({\bm{x}}^{\prime},{\bm{x}}^{\prime})\\}}\left|\left\langle{\bm{\Phi}^{(\ell)}({\bm{x}}_{1}),\bm{\Phi}^{(\ell)}({\bm{x}}_{2})}\right\rangle-
K_{\mathrm{NTK}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|.$

Observe that the estimate error $\Delta_{\ell}$ can be decomposed into three
parts:

$\displaystyle\left|\left\langle{\bm{\Phi}^{(\ell)}({\bm{x}}_{1}),\bm{\Phi}^{(\ell)}({\bm{x}}_{2})}\right\rangle-
K_{\mathrm{NTK}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|\leq\left|\left\langle{\bm{\Psi}^{(\ell)}({\bm{x}}_{1}),\bm{\Psi}^{(\ell)}({\bm{x}}_{2})}\right\rangle-
K_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|$
$\displaystyle\qquad\qquad\qquad\qquad\qquad+\left|\left\langle{\bm{\Gamma}^{(\ell)}({\bm{x}}_{1}),\bm{\Gamma}^{(\ell)}({\bm{x}}_{2})}\right\rangle-\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}_{1})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1}),\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})}\right\rangle\right|$
(47)
$\displaystyle\qquad\qquad\qquad\qquad\qquad+\left|\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}_{1}),\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})}\right\rangle\left\langle{\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1}),\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})}\right\rangle-\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})K_{\mathrm{NTK}}^{(\ell-1)}({\bm{x}}_{1},{\bm{x}}_{2})\right|.$

for
$({\bm{x}}_{1},{\bm{x}}_{2})\in\\{({\bm{x}},{\bm{x}}^{\prime}),({\bm{x}},{\bm{x}}),({\bm{x}}^{\prime},{\bm{x}}^{\prime})\\}$.
By definition, the event
$\mathcal{E}_{\bm{\Psi}}^{(\ell)}\left(\frac{\varepsilon^{2}}{32L^{2}}\right)$
implies that

$\displaystyle\left|\left\langle{\bm{\Phi}^{(\ell)}({\bm{x}}_{1}),\bm{\Phi}^{(\ell)}({\bm{x}}_{1})}\right\rangle-\bm{\Theta}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|\leq\frac{\varepsilon^{2}}{32L^{2}}$
(48)

and the event
$\mathcal{E}_{\bm{\Gamma}}^{(\ell)}\left(\frac{\varepsilon}{8L}\left(\ell+\Delta_{\ell}\right)\right)$
implies that

$\displaystyle\left|\left\langle{\bm{\Gamma}^{(\ell)}({\bm{x}}_{1}),\bm{\Gamma}^{(\ell)}({\bm{x}}_{2})}\right\rangle-\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}_{1})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1}),\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})\otimes\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})}\right\rangle\right|\leq\frac{\varepsilon}{8L}\left(\ell+\Delta_{\ell}\right).$
(49)

For the third part in ??, we observe that

$\displaystyle\left|\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}_{1}),\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})}\right\rangle\left\langle{\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1}),\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})}\right\rangle-\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})K_{\mathrm{NTK}}^{(\ell-1)}({\bm{x}}_{1},{\bm{x}}_{2})\right|$
$\displaystyle\leq\left|\left\langle{\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1}),\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})}\right\rangle\right|\left|\left\langle{\bm{\Lambda}^{(\ell)}({\bm{x}}_{1}),\bm{\Lambda}^{(\ell)}({\bm{x}}_{2})}\right\rangle-\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|$
$\displaystyle\quad\qquad+\left|\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|\left|\left\langle{\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1}),\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})}\right\rangle-
K_{\mathrm{NTK}}^{(\ell-1)}({\bm{x}}_{1},{\bm{x}}_{2})\right|$
$\displaystyle\leq\left\|\bm{\Phi}^{(\ell-1)}({\bm{x}}_{1})\right\|_{2}\left\|\bm{\Phi}^{(\ell-1)}({\bm{x}}_{2})\right\|_{2}\frac{3\varepsilon}{8L}+\Delta_{\ell-1}$
$\displaystyle\leq\left(\ell+\Delta_{\ell-1}\right)\frac{3\varepsilon}{8L}+\Delta_{\ell-1}$
(50)

where the second inequality comes from that the event
$\mathcal{E}_{\bm{\Lambda}}^{(\ell)}\left(\frac{3\varepsilon}{8L}\right)$ and
$\left|\dot{K}_{\mathrm{NNGP}}^{(\ell)}({\bm{x}}_{1},{\bm{x}}_{2})\right|\leq
1$ and the last one follows from ??. Putting ??, ?? and ?? into ??, we have

$\displaystyle\Delta_{\ell}\leq\frac{\varepsilon^{2}}{32L^{2}}+\frac{\varepsilon}{8L}\left(\ell+\Delta_{\ell}\right)+\left(\ell+\Delta_{\ell-1}\right)\frac{3\varepsilon}{8L}+\Delta_{\ell-1}=\left(1+\frac{\varepsilon}{2L}\right)\Delta_{\ell-1}+\frac{\varepsilon}{2L}\ell+\frac{\varepsilon^{2}}{32L^{2}}.$
(51)

This completes the proof of ??. ∎

###  A.2 Proof of ??

The proofs here rely on Theorem 3.3 in Lee et al. (2020) which states spectral
approximation bounds of random features for general kernels equipped with the
leverage score sampling. This result is a generalization of Avron et al.
(2017b) working on the Random Fourier Features.

######  Theorem 5 (Theorem 3.3 in (Lee et al., 2020)).

Suppose ${\bm{K}}\in{\mathbb{R}}^{n\times n}$ is a kernel matrix with
statistical dimension $s_{\lambda}$ for some
$\lambda\in(0,\left\|{\bm{K}}\right\|_{2})$. Let
$\bm{\Phi}({\bm{w}})\in{\mathbb{R}}^{n}$ be a feature map with a random vector
${\bm{w}}\sim p({\bm{w}})$ satisfying that
${\bm{K}}={\mathbb{E}}_{{\bm{w}}}\left[\bm{\Phi}({\bm{w}})\bm{\Phi}({\bm{w}})^{\top}\right]$.
Define
$\tau_{\lambda}({\bm{w}}):=p({\bm{w}})\cdot\bm{\Phi}({\bm{w}})^{\top}({\bm{K}}+\lambda{\bm{I}})^{-1}\bm{\Phi}({\bm{w}})$.
Let $\widetilde{\tau}({\bm{w}})$ be any measurable function such that
$\widetilde{\tau}({\bm{w}})\geq\tau_{\lambda}({\bm{w}})$ for all ${\bm{w}}$.
Assume that $s_{\widetilde{\tau}}:=\int\widetilde{\tau}({\bm{w}})d{\bm{w}}$ is
finite. Consider random vectors ${\bm{w}}_{1},\dots,{\bm{w}}_{m}$ sampled from
$q({\bm{w}}):=\widetilde{\tau}({\bm{w}})/s_{\widetilde{\tau}}$ and define that

$\displaystyle\overline{\bm{\Phi}}:=\frac{1}{\sqrt{m}}\left[\sqrt{\frac{p({\bm{w}}_{1})}{q({\bm{w}}_{1})}}\bm{\Phi}({\bm{w}}_{1}),\
\dots\
,\sqrt{\frac{p({\bm{w}}_{m})}{q({\bm{w}}_{m})}}\bm{\Phi}({\bm{w}}_{m})\right].$
(52)

If
$m\geq\frac{8}{3}\varepsilon^{-2}s_{\widetilde{\tau}}\log\left(16s_{\lambda}/\delta\right)$
then

$\displaystyle\left(1-\varepsilon\right)\left({\bm{K}}+\lambda{\bm{I}}\right)\preceq\overline{\bm{\Phi}}\overline{\bm{\Phi}}^{\top}+\lambda{\bm{I}}\preceq\left(1+\varepsilon\right)\left({\bm{K}}+\lambda{\bm{I}}\right)$
(53)

holds with probability at least $1-\delta$.

We now ready to provide proofs of ??.

See 2

Proof of ??. Let $\bm{\Phi}_{0}({\bm{w}}):=\sqrt{2}\
\mathrm{Step}({\bm{X}}{\bm{w}})\in{\mathbb{R}}^{n}$ for
${\bm{w}}\in{\mathbb{R}}^{d}$ and $p({\bm{w}})$ be the probability density
function of the standard normal distribution. As studied in Cho & Saul (2009),
$\bm{\Phi}_{0}({\bm{w}})$ is a random feature of ${\bm{A}}_{0}$ such that

$\displaystyle{\bm{A}}_{0}={\mathbb{E}}_{{\bm{w}}\sim
p({\bm{w}})}\left[\bm{\Phi}_{0}({\bm{w}})\bm{\Phi}_{0}({\bm{w}})^{\top}\right].$
(54)

In order to utilize ??, we need an upper bound of $\tau_{\lambda}({\bm{w}})$
as below:

$\displaystyle\tau_{\lambda}({\bm{w}})$
$\displaystyle:=p({\bm{w}})\cdot\bm{\Phi}_{0}({\bm{w}})^{\top}\left({\bm{A}}_{0}+\lambda{\bm{I}}\right)^{-1}\bm{\Phi}_{0}({\bm{w}})$
(55) $\displaystyle\leq
p({\bm{w}})\left\|({\bm{A}}_{0}+\lambda{\bm{I}})^{-1}\right\|_{2}\left\|\bm{\Phi}_{0}({\bm{w}})\right\|_{2}^{2}$
(56) $\displaystyle\leq
p({\bm{w}})\frac{\left\|\bm{\Phi}_{0}({\bm{w}})\right\|_{2}^{2}}{\lambda}$
(57) $\displaystyle\leq p({\bm{w}})\frac{2n}{\lambda}$ (58)

where the inequality in second line holds from the definition of matrix
operator norm and the inequality in third line follows from the fact that
smallest eigenvalue of ${\bm{A}}_{0}+\lambda{\bm{I}}$ is equal to or greater
than $\lambda$. The last inequality is from that
$\left\|\mathrm{Step}({\bm{x}})\right\|_{2}^{2}\leq n$ for any
${\bm{x}}\in{\mathbb{R}}^{n}$. Note that
$\int_{{\mathbb{R}}^{d}}p({\bm{w}})\frac{2n}{\lambda}d{\bm{w}}=\frac{2n}{\lambda}$
and since it is constant the modified random features correspond to the
original ones. Putting all together into ??, we can obtain the result. This
completes the proof of ??. ∎

###  A.3 Proof of ??

See 3

Proof of ??. Let
$\bm{\Phi}_{1}({\bm{w}}):=\sqrt{2}~{}\mathrm{ReLU}({\bm{X}}{\bm{w}})\in{\mathbb{R}}^{n}$
for ${\bm{w}}\in{\mathbb{R}}^{d}$ and $p({\bm{w}})$ be the probability density
function of standard normal distribution. Cho & Saul (2009) also showed that
$\bm{\Phi}_{1}({\bm{w}})$ is a random feature of ${\bm{A}}_{1}$ such that

$\displaystyle{\bm{A}}_{1}={\mathbb{E}}_{{\bm{w}}\sim
p({\bm{w}})}\left[\bm{\Phi}_{1}({\bm{w}})\bm{\Phi}_{1}({\bm{w}})^{\top}\right].$
(59)

Again, we use the below upper bound as follow:

$\displaystyle\tau_{\lambda}({\bm{w}})$
$\displaystyle:=p({\bm{w}})\cdot\bm{\Phi}_{1}({\bm{w}})^{\top}({\bm{A}}_{1}+\lambda{\bm{I}})^{-1}\bm{\Phi}_{1}({\bm{w}})$
(60) $\displaystyle\leq
p({\bm{w}})\left\|({\bm{A}}_{1}+\lambda{\bm{I}})^{-1}\right\|_{2}\left\|\bm{\Phi}_{1}({\bm{w}})\right\|_{2}^{2}$
(61)
$\displaystyle=2~{}p({\bm{w}})\frac{\left\|\mathrm{ReLU}({\bm{X}}{\bm{w}})\right\|_{2}^{2}}{\lambda}$
(62) $\displaystyle\leq
2~{}p({\bm{w}})\frac{\left\|{\bm{X}}{\bm{w}}\right\|_{2}^{2}}{\lambda}$ (63)
$\displaystyle\leq
2~{}p({\bm{w}})\left\|{\bm{w}}\right\|_{2}^{2}\frac{\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda}$
(64)

where the inequality in fourth line holds from that
$\left\|\mathrm{ReLU}({\bm{x}})\right\|_{2}^{2}\leq\left\|{\bm{x}}\right\|_{2}^{2}$
for any vector ${\bm{x}}$. Denote
$\widetilde{\tau}({\bm{w}}):=2~{}p({\bm{w}})\left\|{\bm{w}}\right\|_{2}^{2}\frac{\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda}$
and it holds that

$\displaystyle\int_{{\mathbb{R}}^{d}}\widetilde{\tau}({\bm{w}})d{\bm{w}}=\int_{{\mathbb{R}}^{d}}2~{}p({\bm{w}})\left\|{\bm{w}}\right\|_{2}^{2}\frac{\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda}d{\bm{w}}=2d\frac{\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda}$
(65)

since
$\int_{{\mathbb{R}}^{d}}p({\bm{w}})\left\|{\bm{w}}\right\|_{2}^{2}=\mathtt{tr}({\bm{I}}_{d})=d$
for ${\bm{w}}\sim\mathcal{N}(\bm{0},{\bm{I}}_{d})$. We define the modified
distribution as

$\displaystyle
q({\bm{w}}):=\frac{\widetilde{\tau}({\bm{w}})}{\int_{{\mathbb{R}}^{d}}\widetilde{\tau}({\bm{w}})d{\bm{w}}}=p({\bm{w}})\frac{\left\|{\bm{w}}\right\|_{2}^{2}}{d}=\frac{1}{(2\pi)^{d/2}d}\left\|{\bm{w}}\right\|_{2}^{2}\exp\left(-\frac{1}{2}\left\|{\bm{w}}\right\|_{2}^{2}\right)$
(66)

and recall the modified random features as

$\displaystyle\bm{\Phi}_{1}$
$\displaystyle=\frac{1}{\sqrt{m}}\left[\sqrt{\frac{p({\bm{w}}_{1})}{q({\bm{w}}_{1})}}\
\bm{\Phi}_{1}({\bm{w}}_{1}),\ \dots\
,\sqrt{\frac{p({\bm{w}}_{m})}{q({\bm{w}}_{m})}}\
\bm{\Phi}_{1}({\bm{w}}_{m})\right]$ (67)
$\displaystyle=\sqrt{\frac{2d}{m}}\left[\frac{\mathrm{ReLU}({\bm{X}}{\bm{w}}_{1})}{\left\|{\bm{w}}_{1}\right\|_{2}},\
\dots\
,\frac{\mathrm{ReLU}({\bm{X}}{\bm{w}}_{m})}{\left\|{\bm{w}}_{m}\right\|_{2}}\right].$
(68)

Putting all together into ??, we derive the result. This completes the proof
of ??. ∎

###  A.4 Proof of ??

Before diving into detailed algorithmic analysis, we introduce spectral
approximation bounds of CountSketch when it applies to Hadamard product of two
PSD matrices. Recall that the CountSketch plays a key role for reducing the
feature map dimensionality and below theorem is used in the proof of ??.

######  Lemma 4.

Given ${\bm{X}}\in{\mathbb{R}}^{n\times d_{1}}$ and
${\bm{Y}}\in{\mathbb{R}}^{n\times d_{2}}$, let
$\mathcal{C}_{1},\mathcal{C}_{2}$ be the two independent CountSketch
transforms from ${\mathbb{R}}^{d_{1}},{\mathbb{R}}^{d_{2}}$ to
${\mathbb{R}}^{m}$, respectively. Denote that

$\displaystyle\bm{\Gamma}:=\mathrm{FFT}^{-1}(\mathrm{FFT}(\mathcal{C}_{1}({\bm{X}}))\odot\mathrm{FFT}(\mathcal{C}_{2}({\bm{Y}}))).$
(69)

Given $\varepsilon,\delta\in(0,1)$, $\lambda\geq 0$ and
$m\geq\frac{11}{\varepsilon^{2}\delta}\left(\frac{\mathtt{tr}({\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top})}{\mathtt{tr}({\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top})/n+\lambda}\right)^{2}$,
then it holds

$\displaystyle(1-\varepsilon)\left({\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top}+\lambda{\bm{I}}\right)\preceq\bm{\Gamma}\bm{\Gamma}^{\top}+\lambda{\bm{I}}\preceq(1+\varepsilon)\left({\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top}+\lambda{\bm{I}}\right)$
(70)

with probability at least $1-\delta$.

The proof of ?? is provided in ??.

See 4

Proof of ??. Note that the NTK of two-layer ReLU network can be formulated as

$\displaystyle{\bm{K}}_{\mathrm{NTK}}={\bm{A}}_{1}+{\bm{A}}_{0}\odot({\bm{X}}{\bm{X}}^{\top})$
(71)

where ${\bm{A}}_{0}$ and ${\bm{A}}_{1}$ are the arc-cosine kernel matrices of
order $0$ and $1$ with ${\bm{X}}$, respectively. Let $\bm{\Phi}_{0}$ and
$\bm{\Phi}_{1}$ be the random features of ${\bm{A}}_{0}$ and ${\bm{A}}_{1}$,
respectively, satisfying that
${\bm{A}}_{0}={\mathbb{E}}\left[\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\right]$ and
${\bm{A}}_{1}={\mathbb{E}}\left[\bm{\Phi}_{1}\bm{\Phi}_{1}^{\top}\right]$.
Based on the property of CountSketch , one can check that
$\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\odot{\bm{X}}{\bm{X}}^{\top}={\mathbb{E}}[\bm{\Gamma}\bm{\Gamma}^{\top}]$
where we recall that

$\displaystyle\bm{\Gamma}:=\mathrm{FFT}^{-1}(\mathrm{FFT}(\mathcal{C}_{1}(\bm{\Phi}_{0}))\odot\mathrm{FFT}(\mathcal{C}_{2}({\bm{X}}))).$
(72)

Our proof is a combination of spectral analysis of
$\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top},\bm{\Phi}_{1}\bm{\Phi}_{1}^{\top}$ and
$\bm{\Gamma}\bm{\Gamma}^{\top}$ which are stated in ??, ?? and ??,
respectively.

From ??, if
$m_{1}\geq\frac{16}{3}\frac{d\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda\varepsilon^{2}}\log\left(\frac{48s_{\lambda}}{\delta}\right)$
then with probability at least $1-\frac{\delta}{3}$ it holds

$\displaystyle(1-\varepsilon)\left({\bm{A}}_{1}+\frac{\lambda}{2}{\bm{I}}\right)\preceq\bm{\Phi}_{1}\bm{\Phi}_{1}^{\top}+\frac{\lambda}{2}{\bm{I}}\preceq(1+\varepsilon)\left({\bm{A}}_{1}+\frac{\lambda}{2}{\bm{I}}\right).$
(73)

From ??, if $m_{0}\geq
48\frac{n}{\lambda\varepsilon^{2}}\log\left(\frac{48s_{\lambda}}{\delta}\right)$
then with probability at least $1-\frac{\delta}{3}$ it holds

$\displaystyle\left(1-\frac{\varepsilon}{3}\right)\left({\bm{A}}_{0}+\frac{\lambda}{2}{\bm{I}}\right)\preceq\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}+\frac{\lambda}{2}{\bm{I}}\preceq\left(1+\frac{\varepsilon}{3}\right)\left({\bm{A}}_{0}+\frac{\lambda}{2}{\bm{I}}\right)$
(74)

Rearranging ??, we get

$\displaystyle\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\preceq\left(1+\frac{\varepsilon}{3}\right){\bm{A}}_{0}+\frac{\varepsilon}{6}\lambda{\bm{I}}.$
(75)

To guarantee spectral approximation of $\bm{\Gamma}\bm{\Gamma}^{\top}$, we
will use the result of ??. Before applying it, we provide an upper bound of
the trace of $\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\odot{\bm{X}}{\bm{X}}^{\top}$.
Consider
$\bm{\Phi}_{0}=\sqrt{\frac{2}{m_{0}}}\left[\mathrm{Step}({\bm{Z}}{\bm{w}}_{1}),\dots,\mathrm{Step}({\bm{Z}}{\bm{w}}_{m_{0}})\right]$
for some
${\bm{Z}}=\left[{\bm{z}}_{1},\dots,{\bm{z}}_{n}\right]\in{\mathbb{R}}^{n\times
d}$. Then, we have

$\displaystyle\mathtt{tr}\left(\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\odot{\bm{X}}{\bm{X}}^{\top}\right)$
$\displaystyle=\sum_{j=1}^{n}\left[\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\right]_{jj}\cdot\left[{\bm{X}}{\bm{X}}^{\top}\right]_{jj}$
(76)
$\displaystyle=\sum_{j=1}^{n}\left(\frac{2}{m_{0}}\sum_{i=1}^{m_{0}}\mathrm{Step}\left(\left\langle{{\bm{z}}_{j},{\bm{w}}_{i}}\right\rangle\right)^{2}\right)\cdot\left\|{\bm{x}}_{j}\right\|_{2}^{2}$
(77)
$\displaystyle\leq\sum_{j=1}^{n}1\cdot\left\|{\bm{x}}_{j}\right\|_{2}^{2}=n$
(78)

where the inequality in third line holds from that $\mathrm{Step}(x)\leq 1$
for any $x\in{\mathbb{R}}$ and the last equality follows from the assumption
that $\left\|{\bm{x}}_{j}\right\|_{2}=1$ for all $j\in[n]$.

Hence, using ?? with the fact that
$m_{\mathtt{cs}}\geq\frac{297}{\varepsilon^{2}\delta}\left(\frac{n}{1+\lambda}\right)^{2}\geq\frac{297}{\varepsilon^{2}\delta}\left(\frac{\mathtt{tr}(\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\odot{\bm{X}}{\bm{X}}^{\top})}{\mathtt{tr}(\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\odot{\bm{X}}{\bm{X}}^{\top})/n+\lambda}\right)^{2}$,
we have

$\displaystyle\left(1+\frac{\varepsilon}{3}\right)\left(\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\odot{\bm{X}}{\bm{X}}^{\top}+\frac{\lambda}{2}{\bm{I}}\right)\preceq\bm{\Gamma}\bm{\Gamma}^{\top}+\frac{\lambda}{2}{\bm{I}}\preceq\left(1+\frac{\varepsilon}{3}\right)\left(\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\odot{\bm{X}}{\bm{X}}^{\top}+\frac{\lambda}{2}{\bm{I}}\right)$
(79)

with probability at least $1-\frac{\delta}{3}$. Combining ?? with ??, with
probability at least $1-\frac{2}{3}\delta$, we get that

$\displaystyle\bm{\Gamma}\bm{\Gamma}^{\top}+\frac{\lambda}{2}{\bm{I}}$
$\displaystyle\preceq\left(1+\frac{\varepsilon}{3}\right)\left(\bm{\Phi}_{0}\bm{\Phi}_{0}^{\top}\odot{\bm{X}}{\bm{X}}^{\top}+\frac{\lambda}{2}I\right)$
(80)
$\displaystyle\preceq\left(1+\frac{\varepsilon}{3}\right)\left(\left[\left(1+\frac{\varepsilon}{3}\right){\bm{A}}_{0}+\frac{\varepsilon}{6}\lambda{\bm{I}}\right]\odot{\bm{X}}{\bm{X}}^{\top}+\frac{\lambda}{2}{\bm{I}}\right)$
(81)
$\displaystyle=\left(1+\frac{\varepsilon}{3}\right)\left(\left(1+\frac{\varepsilon}{3}\right)\left({\bm{A}}_{0}\odot{\bm{X}}{\bm{X}}^{\top}\right)+\frac{\varepsilon}{6}\lambda({\bm{I}}\odot{\bm{X}}{\bm{X}}^{\top})+\frac{\lambda}{2}{\bm{I}}\right)$
(82)
$\displaystyle=\left(1+\frac{\varepsilon}{3}\right)\left(1+\frac{\varepsilon}{3}\right)\left({\bm{A}}_{0}\odot{\bm{X}}{\bm{X}}^{\top}+\frac{\lambda}{2}{\bm{I}}\right)$
(83)
$\displaystyle\preceq\left(1+\varepsilon\right)\left({\bm{A}}_{0}\odot{\bm{X}}{\bm{X}}^{\top}+\frac{\lambda}{2}{\bm{I}}\right)$
(84)

where the equality in second line follows from the fact that
${\bm{A}}\odot{\bm{C}}\preceq{\bm{B}}\odot{\bm{C}}$ holds if
${\bm{A}}\preceq{\bm{B}}$ for positive semidefinite matrices
${\bm{A}},{\bm{B}}$ and ${\bm{C}}$ 555It is enough to show that
$({\bm{B}}-{\bm{A}})\odot{\bm{C}}\succeq 0$. Since
${\bm{B}}-{\bm{A}},{\bm{C}}\succeq 0$, this holds from the Schur product
theorem. and the fourth equality is from the assumption
$\left\|{\bm{x}}_{i}\right\|_{2}=1$ for all $i\in[n]$ which leads that
${\bm{I}}\odot({\bm{X}}{\bm{X}}^{\top})={\bm{I}}$. The last inequality holds
since $\varepsilon\in(0,1/2)$.

Similarly, we can obtain the following lower bound:

$\displaystyle\bm{\Gamma}\bm{\Gamma}^{\top}+\frac{\lambda}{2}{\bm{I}}\succeq\left(1-\varepsilon\right)\left({\bm{A}}_{0}\odot{\bm{X}}{\bm{X}}^{\top}+\frac{\lambda}{2}{\bm{I}}\right).$
(85)

Combining ??, ?? with ?? gives

$\displaystyle(1-\varepsilon)\left({\bm{A}}_{1}+{\bm{A}}_{0}\odot{\bm{X}}{\bm{X}}^{\top}+\lambda{\bm{I}}\right)\preceq\bm{\Phi}\bm{\Phi}^{\top}+\lambda{\bm{I}}\preceq(1+\varepsilon)\left({\bm{A}}_{1}+{\bm{A}}_{0}\odot{\bm{X}}{\bm{X}}^{\top}+\lambda{\bm{I}}\right).$
(86)

Furthermore, by taking a union bound over all events, ?? holds with
probability at least $1-\delta$. This completes the proof of ??. ∎

##  Appendix B Proof of Lemmas

The proofs of ?? and ?? are obtained from Lemma 2 in Avron et al. (2014) that
provides an upper bound on variance of TensorSketch transform of order $q\geq
2$.

######  Lemma 5 (Lemma 2 in Avron et al. (2014)).

For $q\geq 2$, consider $q$ of 3-wise independent hash functions $h_{1},\dots
h_{q}:[d]\rightarrow[m]$ and $q$ of 4-wise independent sign functions
$s_{1},\dots,s_{q}:[d]\rightarrow\\{+1,-1\\}$. Define the hash function
$H:[d^{q}]\rightarrow[m]$ and the sign function
$S:[d^{q}]\rightarrow\\{-1,+1\\}$ such that

$\displaystyle H(j)$ $\displaystyle\equiv
h_{1}(i_{1})+h_{2}(i_{2})+\dots+h_{q}(i_{q})\pmod{m},$ $\displaystyle S(j)$
$\displaystyle=s_{1}(i_{1})\cdot s_{2}(i_{2})\cdot\dots\cdot s_{q}(i_{q}).$

where $j\in[d^{q}]$ and $i_{1},\dots,i_{q}\in[d]$ such that
$j=i_{q}d^{q-1}+\cdots+i_{2}d+i_{1}$. Denote sketch matrix
${\bm{S}}\in\\{-1,0,+1\\}^{d^{q}\times m}$ satisfying that
${\bm{S}}_{j,H(j)}=S(j)$ for $j\in[d^{q}]$ and other entries are set to zero.
For any ${\bm{A}},{\bm{B}}\in{\mathbb{R}}^{n\times d^{q}}$, it holds

$\displaystyle{\mathbb{E}}_{{\bm{S}}}\left[\left\|{\bm{A}}{\bm{S}}{\bm{S}}^{\top}{\bm{B}}^{\top}-{\bm{A}}{\bm{B}}^{\top}\right\|_{F}^{2}\right]\leq\frac{(2+3^{q})}{m}\left\|{\bm{A}}\right\|_{F}^{2}\left\|{\bm{B}}\right\|_{F}^{2}.$
(87)

###  B.1 Proof of ??

See 3

Proof of ??. By the Markov’s inequality, we have

$\displaystyle\Pr\left(\left|\left\langle{\bm{\Gamma},\bm{\Gamma}^{\prime}}\right\rangle-\left\langle{{\bm{x}}\otimes{\bm{y}},{\bm{x}}^{\prime}\otimes{\bm{y}}^{\prime}}\right\rangle\right|\geq\varepsilon\right)\leq\frac{1}{\varepsilon^{2}}{\mathbb{E}}\left[\left|\left\langle{{\bm{x}}\otimes{\bm{y}},{\bm{x}}^{\prime}\otimes{\bm{y}}^{\prime}}\right\rangle-\left\langle{\bm{\Gamma},\bm{\Gamma}^{\prime}}\right\rangle\right|^{2}\right]\leq\frac{11\left\|{\bm{x}}\right\|_{2}^{2}\left\|{\bm{x}}^{\prime}\right\|_{2}^{2}\left\|{\bm{y}}\right\|_{2}^{2}\left\|{\bm{y}}^{\prime}\right\|_{2}^{2}}{\varepsilon^{2}m_{\mathtt{cs}}}$
(88)

where the last inequality follows from ?? with $q=2$. This completes the proof
of ??. ∎

###  B.2 Proof of ??

See 4

Proof of ??. Let $s_{1}:[d_{1}]\rightarrow\\{-1,+1\\}$ and
$h_{1}:[d_{1}]\rightarrow[m]$ be the random sign and hash function of
$\mathcal{C}_{1}$, respectively. Similarly, denote $s_{2}$ and $h_{2}$ by that
of $\mathcal{C}_{2}$, respectively. Then, $\bm{\Gamma}$ is the output of Count
Sketch $\mathcal{C}:{\mathbb{R}}^{d_{1}d_{2}}\rightarrow{\mathbb{R}}^{m}$
applying to ${\bm{X}}\otimes{\bm{Y}}$ whose sign and hash functions are
defined as

$\displaystyle s(i,j)$ $\displaystyle=s_{1}(i)\cdot s_{2}(j)$ (89)
$\displaystyle h(i,j)$ $\displaystyle\equiv h_{1}(i)+h_{2}(j)\pmod{m}$ (90)

for $i\in[d_{1}],j\in[d_{2}]$. Here, index $(i,j)$ can be considered as some
$k\in[d_{1}d_{2}]$ by transforming $i=\lfloor k/d_{2}\rfloor$ and $j\equiv
k\pmod{d_{2}}$.

Let ${\bm{S}}\in\\{-1,0,+1\\}^{d_{1}d_{2}\times m}$ be the sketch matrix of
$\mathcal{C}$ and we write ${\bm{Z}}:={\bm{X}}\otimes{\bm{Y}}$ for notational
simplicity. As shown in Avron et al. (2014), it is easy to check that
$\bm{\Gamma}={\bm{Z}}{\bm{S}}$ and we have

$\displaystyle{\mathbb{E}}[\bm{\Gamma}\bm{\Gamma}^{\top}]$
$\displaystyle={\mathbb{E}}[({\bm{X}}\otimes{\bm{Y}}){\bm{S}}{\bm{S}}^{\top}({\bm{X}}\otimes{\bm{Y}})^{\top}]$
(91)
$\displaystyle=({\bm{X}}\otimes{\bm{Y}}){\mathbb{E}}[{\bm{S}}{\bm{S}}^{\top}]({\bm{X}}\otimes{\bm{Y}})^{\top}$
(92) $\displaystyle=({\bm{X}}\otimes{\bm{Y}})({\bm{X}}\otimes{\bm{Y}})^{\top}$
(93)
$\displaystyle={\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top}={\bm{Z}}{\bm{Z}}^{\top}.$
(94)

Rearranging ??, we have

$\displaystyle-\varepsilon\left({\bm{Z}}{\bm{Z}}^{\top}+\lambda{\bm{I}}\right)\preceq{\bm{Z}}{\bm{S}}{\bm{S}}^{\top}{\bm{Z}}^{\top}-{\bm{Z}}{\bm{Z}}^{\top}\preceq\varepsilon\left({\bm{Z}}{\bm{Z}}^{\top}+\lambda{\bm{I}}\right).$
(95)

By multiplying $({\bm{Z}}{\bm{Z}}^{\top}+\lambda{\bm{I}})^{-{1}/{2}}$ by both
left and right sides in (95), it is enough to show that

$\displaystyle\left\|({\bm{Z}}{\bm{Z}}^{\top}+\lambda{\bm{I}})^{-{1}/{2}}{\bm{Z}}{\bm{S}}{\bm{S}}^{\top}{\bm{Z}}^{\top}({\bm{Z}}{\bm{Z}}^{\top}+\lambda{\bm{I}})^{-{1}/{2}}-({\bm{Z}}{\bm{Z}}^{\top}+\lambda{\bm{I}})^{-{1}/{2}}{\bm{Z}}{\bm{Z}}^{\top}({\bm{Z}}{\bm{Z}}^{\top}+\lambda{\bm{I}})^{-{1}/{2}}\right\|_{2}\leq\varepsilon$

By denoting
${\bm{A}}:=({\bm{Z}}{\bm{Z}}^{\top}+\lambda{\bm{I}})^{-{1}/{2}}{\bm{Z}}$, it
is equivalent to prove that

$\displaystyle\left\|{\bm{A}}{\bm{S}}{\bm{S}}^{\top}{\bm{A}}^{\top}-{\bm{A}}{\bm{A}}^{\top}\right\|_{2}\leq\varepsilon.$
(96)

By Markov’s inequality, we have

$\displaystyle\Pr\left[\left\|{\bm{A}}{\bm{S}}{\bm{S}}^{\top}{\bm{A}}^{\top}-{\bm{A}}{\bm{A}}^{\top}\right\|_{2}\geq\varepsilon\right]$
$\displaystyle\leq\Pr\left[\left\|{\bm{A}}{\bm{S}}{\bm{S}}^{\top}{\bm{A}}^{\top}-{\bm{A}}{\bm{A}}^{\top}\right\|_{F}\geq\varepsilon\right]$
$\displaystyle\leq\varepsilon^{-2}{\mathbb{E}}\left[\left\|{\bm{A}}{\bm{S}}{\bm{S}}^{\top}{\bm{A}}^{\top}-{\bm{A}}{\bm{A}}^{\top}\right\|_{F}^{2}\right]$
$\displaystyle\leq\frac{11}{\varepsilon^{2}m}\left\|{\bm{A}}\right\|_{F}^{4}$
$\displaystyle=\frac{11}{\varepsilon^{2}m}\left[\mathtt{tr}({\bm{A}}^{\top}{\bm{A}})\right]^{2}$
$\displaystyle=\frac{11}{\varepsilon^{2}m}\left[\mathtt{tr}\left(({\bm{Z}}{\bm{Z}}^{\top}+{\bm{I}})^{-1}{\bm{Z}}{\bm{Z}}^{\top}\right)\right]^{2}$
$\displaystyle\leq\frac{11}{\varepsilon^{2}m}\left(\frac{\mathtt{tr}({\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top})}{\mathtt{tr}({\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top})/n+\lambda}\right)^{2}$

where the inequality in third line holds from ?? with $q=2$ and the last
inequality follows from ??. Taking
$m\geq\frac{11}{\varepsilon^{2}\delta}\left(\frac{\mathtt{tr}({\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top})}{\mathtt{tr}({\bm{X}}{\bm{X}}^{\top}\odot{\bm{Y}}{\bm{Y}}^{\top})/n+\lambda}\right)^{2}$,
?? holds with probability at least $1-\delta$. This completes the proof of ??.
∎

######  Lemma 6.

Let ${\bm{A}},{\bm{B}}$ be positive semidefinite matrices and let
$s_{\lambda}$ be the statistical dimension of ${\bm{A}}\odot{\bm{B}}$ for any
$\lambda>0$, i.e.,
$s_{\lambda}:=\mathtt{tr}\left({({\bm{A}}\odot{\bm{B}})({\bm{A}}\odot{\bm{B}}+\lambda{\bm{I}})^{-1}}\right)$.
Then, it holds that
$s_{\lambda}\leq\frac{\mathtt{tr}({\bm{A}}\odot{\bm{B}})}{\mathtt{tr}({\bm{A}}\odot{\bm{B}})/n+\lambda}$.

Proof of ??. By the Schur product theorem, ${\bm{A}}\odot{\bm{B}}$ is a
positive semidefinite matrix. Let
$\lambda_{1}\geq\lambda_{2}\geq\dots\geq\lambda_{n}\geq 0$ be the eigenvalues
of ${\bm{A}}\odot{\bm{B}}$. By the definition of statistical dimension, it
holds that

$\displaystyle s_{\lambda}$
$\displaystyle=\mathtt{tr}\left(({\bm{A}}\odot{\bm{B}})\left({\bm{A}}\odot{\bm{B}}+\lambda{\bm{I}}\right)^{-1}\right)=\sum_{i=1}^{n}\frac{\lambda_{i}}{\lambda_{i}+\lambda}$
(97) $\displaystyle\leq
n\frac{\left(\sum_{i}\lambda_{i}\right)/n}{\left(\sum_{i}\lambda_{i}\right)/n+\lambda}=\frac{\sum_{i}\lambda_{i}}{(\sum_{i}\lambda_{i})/n+\lambda}=\frac{\mathtt{tr}({\bm{A}}\odot{\bm{B}})}{\mathtt{tr}({\bm{A}}\odot{\bm{B}})/n+\lambda}$
(98)

where the inequality holds from the Jensen’s inequality. This completes the
proof of ??. ∎

##  Appendix C Mathematical Details on Deep Active Learning

In this section, we give additional details on how NTK random features can be
used to accelerate the greedy selection algorithm in (Shoham & Avron, 2020).
In actuality, the improvement is not specific to NTK features, but works for
every kernel for which have a low-rank factorization. It applies to NTK random
features by virtue of the low-rank factorization the approximate kernel
induces.

For a matrix ${\bm{A}}$ and index sets ${\cal S}$ and ${\cal T}$, let
${\bm{A}}_{{\cal S},{\cal T}}$ denote the matrix obtained by restricting to
the rows whose index is in ${\cal S}$ and the columns whose index is in ${\cal
T}$. Using $:$ as the index set denotes the entire relevant index set.
Consider kernel ridge regression, and assume the kernel matrix is ${\bm{K}}$.
One important variant of the criteria developed in (Shoham & Avron, 2020) is
the minimization of

$J({\cal S})=\mathtt{tr}(-{\bm{K}}_{:,{\cal S}}({\bm{K}}_{{\cal S},{\cal
S}}+\lambda{\bm{I}}_{|{\cal S}|})^{-1}{\bm{K}}^{\top}_{:,{\cal S}})$

In order to perform greedy minimization of $J({\cal S})$ we need to be able to
evaluate $J({\cal S})$ quickly for a given ${\cal S}$.

Assume that ${\bm{K}}=\bm{\Phi}\bm{\Phi}^{\top}$ where $\bm{\Phi}$ has $m$
columns. We now show how after preprocessing $J({\cal S})$ can be computed in
$\mathcal{O}(|{\cal S}|^{2}m)$ time as long as $|{\cal S}|\leq m$. First,
notice that

$\displaystyle J({\cal S})$ $\displaystyle=$
$\displaystyle\mathtt{tr}(-{\bm{K}}_{:,{\cal S}}({\bm{K}}_{{\cal S},{\cal
S}}+\lambda{\bm{I}}_{|{\cal S}|})^{-1}{\bm{K}}^{\top}_{:,{\cal S}})$
$\displaystyle=$ $\displaystyle\mathtt{tr}(-({\bm{K}}_{{\cal S},{\cal
S}}+\lambda{\bm{I}}_{|{\cal S}|})^{-1}{\bm{K}}^{\top}_{:,{\cal
S}}{\bm{K}}_{:,{\cal S}})$ $\displaystyle=$
$\displaystyle\mathtt{tr}(-\underbrace{({\bm{K}}_{{\cal S},{\cal
S}}+\lambda{\bm{I}}_{|{\cal S}|})^{-1}\bm{\Phi}_{{\cal S},:}}_{{\bm{A}}({\cal
S})}\cdot\underbrace{\bm{\Phi}^{\top}\bm{\Phi}\bm{\Phi}^{\top}_{{\cal
S},:}}_{{\bm{B}}({\cal S})})$

Now, notice that ${\bm{A}}({\cal S})$ can be computed in $\mathcal{O}(|{\cal
S}|^{2}m)$ time if we assume that $|{\cal S}|\leq m$. As for ${\bm{B}}({\cal
S})$, this matrix consists exactly of the columns in ${\cal S}$ of
$\bm{\Phi}^{\top}\bm{\Phi}\bm{\Phi}^{\top}$. To take advantage of that we
precompute $\bm{\Phi}^{\top}\bm{\Phi}\bm{\Phi}^{\top}$ in
$\mathcal{O}(nm^{2})$. Finally, note that we are only interested in the trace
of ${\bm{A}}({\cal S})\cdot{\bm{B}}({\cal S})$. There is no need to compute
the entire product; we can compute only the diagonal elements. We see that
after the $\mathcal{O}(nm^{2})$ preprocessing, we can compute $J({\cal S})$ in
$\mathcal{O}(|{\cal S}|^{2}m)$ time.

In order to greedily minimize $J({\cal S})$, we start with ${\cal
S}=\emptyset$, and and add at each iteration the index that will minimize
$J({\cal S})$. To do so, we scan the entire index set, evaluating $J({\cal
T})$ for each candidate ${\cal T}$ that consists of the current ${\cal S}$ and
the addition index. Since there are $n$ data-points, the scan takes
$\mathcal{O}(n|{\cal S}|^{2}m)$ which is the cost per iteration. Including
preprocessing time, the cost of finding a design of size $k$ is
$\mathcal{O}(nm(k^{3}+m))$.

##  Appendix D Additional Experiments on Image Classification

###  D.1 Classification on Image Datasets

We additionally conduct experiments on image classification using $8$ fine-
grained datasets: $\mathtt{CIFAR}$-10/100 Krizhevsky (2009), $\mathtt{VOC07}$
Everingham et al. (2010), $\mathtt{Caltech}$-101 Fei-Fei et al. (2004),
$\mathtt{CUB}$-200 Welinder et al. (2010), $\mathtt{Dog}$-120 Khosla et al.
(2011), $\mathtt{Flower}$-102 Nilsback & Zisserman (2008), $\mathtt{Food}$-101
Bossard et al. (2014). In particular, we follow the transfer learning
mechanism Goyal et al. (2019) where we extract image features from the
penultimate layer of the pretrained ResNet18 He et al. (2016) with dimension
$d=512$. These features are then leveraged as inputs to be transformed to
random features. In particular, we follow low-shot setting used in Arora et
al. (2019c); we randomly choose 5 image data from each class of training set,
and use the whole test set for evaluation. We repeat sampling training images
$50$ times for $\mathtt{CIFAR}$-10 and $\mathtt{VOC07}$, $10$ times for other
datasets. This is because these datasets have relatively small classes, i.e.,
10 and 20, respectively.

We run ?? with and without Gibbs sampling (GS) (i.e., ??). The number of Gibbs
iteration is set to $T=1$. The output dimension $m$ is fixed to $4{,}000$ and
CountSketch dimension is considered $m_{\mathtt{cs}}$ as a hyperparameter. We
set $m_{0}=m_{1}=m-m_{\mathtt{cs}}$ and choose
$m_{\mathtt{cs}}\in\\{0,\frac{m}{10},\dots,\frac{9m}{10},m\\}$ for the best
validation accuracy. We normalize the output of ?? so that the corresponding
features lie in $\mathbb{S}^{d-1}$. We also benchmark the Random Fourier
Features (RFF) with the same dimension $m=4{,}000$. Once features are
generated, we train a linear classifier with SGD optimizer for $500$ epochs
where inputs can be the pretrained features or random features. We perform
grid search for finding the best learning rate in $\\{0.01,0.02,0.1,0.2\\}$
and momentum in $\\{0,0.9,0.99\\}$ and report the best test accuracy among 12
different setups. We also execute various classifiers including AdaBoost,
random forest, $k$-nearest neighbors and support vector classifier (SVC). For
methods running with SVC, we search the cost value $C$ in
$\\{2^{-19},\dots,2^{-3},10^{-7},\dots,10^{3}\\}$ and choose the best one that
achieves the globally maximum accuracy. Hence, a single cost value is globally
used for all iterations. The number of network depth for NTK is chosen by
$L=2$. For RBF kernel (i.e.,
$K({\bm{x}},{\bm{y}})=\exp(-\gamma\left\|{\bm{x}}-{\bm{y}}\right\|_{2}^{2})$),
we choose the best $\gamma$ in $\\{\frac{1}{d},\frac{1}{d\sigma}\\}$ that
achieves the globally maximum accuracy where $\sigma$ is the variance of
training data. For $k$-nearest neighbor classifier, we search the best $k$ in
$\\{1,2,\dots,\min(20,n/2)\\}$ where $n$ is the number of training data. This
allows the number of training instance per class to be roughly larger than
$2$. For AdaBoost and Random Forest, we search the number of ensembles in
$\\{50,100,200,500\\}$.

In ??, the average test accuracy with 95% confidence interval is reported.
Observe that NTK and the corresponding random features show better performance
than other competitors for all datasets. These observations match with the
previous result Arora et al. (2019c) that the NTK can outperform on small-
scale datasets. We additionally verify that NTK random features can perform
similar or even better than NTK for most datasets, especially with Gibbs
sampling. Such performance gaps are also observed between the Random Fourier
Features and RBF kernel. This is likely due to the fact that random features
has an implicit regularization effect which can lead to better generalization
ability.

Table 3: Results of average test accuracy for image classification using
features from the pretrained ResNet-18. We measure the 95% confidence interval
across $50$ iterations for $\mathtt{CIFAR}$-10 and $\mathtt{VOC07}$, and $10$
iterations for the rest. Bold entries indicate the best mean accuracy for each
dataset.

Method |  $\mathtt{CIFAR}$-10 |  $\mathtt{CIFAR}$-100 | $\mathtt{VOC07}$ |  $\mathtt{Caltech}$-101 |  $\mathtt{CUB}$-200 |  $\mathtt{Dog}$-120 |  $\mathtt{Flower}$-102 |  $\mathtt{Food}$-101
---|---|---|---|---|---|---|---|---
Linear Classifier | 59.97 $\pm$ 0.61 | 37.61 $\pm$ 0.42 | 59.91 $\pm$ 0.75 | 82.41 $\pm$ 0.39 | 37.92 $\pm$ 0.46 | 66.09 $\pm$ 0.51 | 73.56 $\pm$ 0.30 | 29.04 $\pm$ 0.33
$k$-Nearest Neighbors | 47.64 $\pm$ 0.77 | 26.09 $\pm$ 0.50 | 50.32 $\pm$ 0.89 | 73.37 $\pm$ 0.46 | 24.78 $\pm$ 0.48 | 54.33 $\pm$ 0.59 | 56.36 $\pm$ 0.49 | 18.59 $\pm$ 0.43
AdaBoost | 32.09 $\pm$ 0.90 | 9.56 $\pm$ 0.57 | 31.56 $\pm$ 0.99 | 26.41 $\pm$ 2.34 | 7.01 $\pm$ 0.87 | 22.72 $\pm$ 1.95 | 20.87 $\pm$ 1.81 | 6.20 $\pm$ 0.45
Random Forest | 56.79 $\pm$ 0.63 | 30.09 $\pm$ 0.43 | 56.16 $\pm$ 0.75 | 75.70 $\pm$ 0.74 | 29.71 $\pm$ 0.42 | 61.31 $\pm$ 0.45 | 64.49 $\pm$ 0.37 | 22.47 $\pm$ 0.39
Linear SVM | 58.59 $\pm$ 0.63 | 35.74 $\pm$ 0.48 | 60.63 $\pm$ 0.79 | 81.39 $\pm$ 0.55 | 36.68 $\pm$ 0.33 | 66.32 $\pm$ 0.49 | 70.94 $\pm$ 0.25 | 28.44 $\pm$ 0.37
RBF Kernel SVM | 59.16 $\pm$ 0.63 | 36.42 $\pm$ 0.54 | 60.79 $\pm$ 0.76 | 82.14 $\pm$ 0.45 | 36.67 $\pm$ 0.44 | 66.49 $\pm$ 0.57 | 71.41 $\pm$ 0.14 | 29.03 $\pm$ 0.34
Random Fourier Features | 59.63 $\pm$ 0.68 | 37.74 $\pm$ 0.50 | 60.95 $\pm$ 0.73 | 82.17 $\pm$ 0.51 | 37.50 $\pm$ 0.48 | 67.38 $\pm$ 0.43 | 72.83 $\pm$ 0.20 | 30.02 $\pm$ 0.36
NTK SVM | 60.48 $\pm$ 0.60 | 37.53 $\pm$ 0.55 | 61.19 $\pm$ 0.70 | 82.83 $\pm$ 0.41 | 37.95 $\pm$ 0.40 | 67.72 $\pm$ 0.49 | 72.32 $\pm$ 0.23 | 29.63 $\pm$ 0.37
NTK Random Features | 60.63 $\pm$ 0.62 | 38.53 $\pm$ 0.50 | 61.44 $\pm$ 0.67 | 82.65 $\pm$ 0.53 | 38.11 $\pm$ 0.56 | 68.06 $\pm$ 0.54 | 73.62 $\pm$ 0.27 | 30.49 $\pm$ 0.34
NTK Random Features with GS | 60.66 $\pm$ 0.60 | 38.49 $\pm$ 0.51 | 61.48 $\pm$ 0.68 | 82.67 $\pm$ 0.51 | 38.22 $\pm$ 0.51 | 68.07 $\pm$ 0.55 | 73.67 $\pm$ 0.28 | 30.50 $\pm$ 0.32

We also investigate the effect of feature dimension $m$ to image
classification performance. In ??, we plot the test accuracy of the proposed
NTK random features with Gibbs sampling (blue, triangle), Random Fourier
Features (red, cross) and features from the pretrained ResNet-18 (green,
circle) when feature dimension $m$ changes from $500$ to $8{,}000$. Note that
the pretrained features has a fixed dimension $d=512$. The hyperparameters are
chosen by the same approach as described above. We observe that a larger $m$
can lead to higher test accuracy for both ours and RFF. It suffices to set
$m=4{,}000$ for the NTK features to achieve higher accuracy than the
pretrained features. We also verify that the proposed NTK random features
shows better performance than RFF for the same dimension $m$. This justifies
that ours is more effective for fine-grained image classifications.

(a) $\mathtt{CIFAR}$-10

(b) $\mathtt{CIFAR}$-100

(c) $\mathtt{VOC07}$

(d) $\mathtt{Caltech}$-101

(e) $\mathtt{CUB}$-200

(f) $\mathtt{Dog}$-120

(g) $\mathtt{Flower}$-102

(h) $\mathtt{Food}$-101

Figure 4: Test accuracy of the proposed NTK random features with Gibbs
sampling and Random Fourier Features when the number of feature dimension $m$
changes from $500$ to $8{,}000$.