# Scaling Neural Tangent Kernels
via Sketching and Random Features

Amir Zandieh111Equal contribution. &Insu Han11footnotemark: 1 &Haim Avron Neta
Shoham &Chaewon Kim &Jinwoo Shin Max-Planck Institute for Informatics. Email:
<EMAIL_ADDRESS>Advanced Institute of Science and Technology.
Email<EMAIL_ADDRESS>Aviv University. Email:
<EMAIL_ADDRESS>Aviv University. Email<EMAIL_ADDRESS>Advanced Institute of Science and Technology. Email:
<EMAIL_ADDRESS>Advanced Institute of Science and Technology. Email:
<EMAIL_ADDRESS>
###### Abstract

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

##  1 Introduction

Recent results have shown that over-parameterized Deep Neural Networks (DNNs),
generalize surprisingly well. In an effort to understand this phenomena,
researchers have studied ultra-wide DNNs and shown that in the infinite width
limit, a fully connected DNN trained by gradient descent under least-squares
loss is equivalent to kernel regression with respect to the Neural Tangent
Kernel (NTK) [5, 11, 22, 28]. This connection has shed light on DNNs’ ability
to generalize [10, 33] and optimize (train) their parameters efficiently [3,
4, 16]. More recently, Arora et al. [5] proved an analogous equivalence
between convolutional DNNs with infinite number of channels and Convolutional
NTK (CNTK). Beyond the aforementioned theoretical purposes, several papers
have explored the algorithmic use of this kernel. Arora et al. [6] and Geifman
et al. [19] showed that NTK based kernel models can outperform trained DNNs
(of finite width). Additionally, CNTK kernel regression sets an impressive
performance record on CIFAR-10 for kernel methods without trainable kernels
[5]. The NTK has also been used in experimental design [38] and predicting
training time [42].

However, the NTK-based approaches encounter the computational bottlenecks of
kernel learning. In particular, for a dataset of $n$ images
$x_{1},x_{2},\ldots x_{n}\in\mathbb{R}^{d\times d}$, only writing down the
CNTK kernel matrix requires $\Omega\left(d^{4}\cdot n^{2}\right)$ operations
[5]. Running regression or PCA on the resulting kernel matrix takes additional
cubic time in $n$, which is infeasible in large-scale setups.

There is a rich literature on kernel approximations for large-scale learning.
One of the most popular approaches is the random features method which works
by randomly sampling the feature space of the kernel function, originally due
to the seminal work of Rahimi and Recht [36]. Another popular approach which
is developed in linear sketching literature [40], works by designing sketches
that can be efficiently applied to the feature space of a kernel without
needing to explicitly form the high dimensional feature space. This approach
has been successful at designing efficient subspace embeddings for the
polynomial kernel [7, 1]. In this paper, we propose solutions for scaling the
NTK and CNTK by building on both of these kernel approximations techniques and
designing low-dimensional feature maps that approximate the NTK/CNTK
evaluation. Consequently, we can simply transform the input dataset to these
low-dimensional feature spaces, and then apply fast linear learning methods to
approximate the answer of the corresponding nonlinear kernel method
efficiently. The performance of such approximate methods is similar or
sometimes better than the exact kernel methods due to implicit regularization
effects of the approximation algorithms [36, 37, 23].

###  1.1 Overview of Our Contributions

  * •

One of our results is an efficient random features construction for the NTK.
Our starting point is the explicit NTK feature map suggested by Bietti and
Mairal [9] based on tensor product of the feature maps of arc-cosine kernels.
We obtain our random features, by sampling the feature space of arc-cosine
kernels [12]. However, the naïve construction of the features would incur an
exponential cost in the depth of the NTK, due to the tensor product of
features generated in consecutive layers. We remedy this issue, by utilizing
an efficient sketching algorithm for tensor products known as TensorSRHT [1]
which can effectively approximate the tensor products of vectors while
preserving their inner products. We provide a rigorous error analysis of the
proposed scheme in Theorem 2.

  * •

Our next results are sketching methods for both NTK and CNTK using a runtime
that is linearly proportional to the sparsity of the input dataset (or number
of pixels of images). Our methods rely on the arc-cosine kernels’ feature
space defined by their Taylor expansion. By careful truncation of their Taylor
series, we approximate the arc-cosine kernels with bounded-degree polynomial
kernels. Because the feature space of a polynomial kernel is the tensor
product of its input space, its dimensionality is exponential in the degree of
the kernel. Fortunately, Ahle et al. [1] have developed a linear sketch known
as PolySketch that can reduce the dimensionality of high-degree tensor
products very efficiently, therefore, we can sketch the resulting polynomial
kernels using this technique. We then combine the transformed features from
consecutive layers by further sketching their tensor products. In case of
CNTK, we have an extra operation which sketches the direct sum of the features
of neighbouring pixels at each layer that precisely corresponds to the
convolution operation in CNNs. We carefully analyze the errors introduced by
polynomial approximations and various sketching steps in our algorithms and
also bound their runtimes in Theorems 1 and 4.

  * •

Furthermore, we improve the arc-cosine random features to spectrally
approximate the entire kernel matrix, which is advocated in recent literature
for ensuring high approximation quality in downstream tasks [8, 31]. Our
construction is based on leverage score sampling, which entertains better
convergence bounds [8, 28]. However, computing this distribution is as
expensive as solving the kernel methods exactly. We propose a simple
distribution that tightly upper bounds the leverage scores of arc-cosine
kernels and for further efficiency, use Gibbs sampling to generate random
features from our proposed distribution. We provide our spectral approximation
guarantee in Theorem 3.

  * •

Finally, we empirically benchmark our proposed methods on various
classification/regression tasks and demonstrate that our methods perform
similar to or better than exact kernel method with NTK and CNTK while running
extremely faster. In particular, we classify CIFAR-10 dataset 150$\times$
faster than exact CNTK and at the same time achieve higher test accuracy.

###  1.2 Related Works

There has been a long line of work on the correspondence between DNN and
kernel machines [26, 29, 34, 18, 41]. Furthermore, there has been many efforts
in understanding a variety of NTK properties including optimization [27, 3,
16, 43], generalization [10], loss surface [30], etc.

Novak et al. [34] tried accelerating CNTK computations via Monte Carlo methods
by taking the gradient of a randomly initialized CNN with respect to its
weights. Although they do not theoretically bound the number of required
features, the fully-connected version of this method is analyzed in [5].
Particularly, for the gradient features to approximate the NTK up to
$\varepsilon$, the network width needs to be
$\Omega(\frac{L^{6}}{\varepsilon^{4}}\log\frac{L}{\delta})$, thus,
transforming a single vector $x\in\mathbb{R}^{d}$ requires
$\Omega(\frac{L^{13}}{\varepsilon^{8}}\log^{2}\frac{L}{\delta}+\frac{L^{6}}{\varepsilon^{4}}\log\frac{L}{\delta}\cdot{\rm
nnz}(x))$ operations, which is slower than our Theorem 1 by a factor of
$L^{3}/\varepsilon^{2}$. Furthermore, [5] shows that the performance of these
random gradients is worse than exact CNTK by a large margin, in practice. More
recently, [28] proposed leverage score sampling for the NTK, however, their
work is primarily theoretical and suggests no practical way of sampling the
features. Another line of work on NTK approximation is an explicit feature map
construction via tensor product proposed by Bietti and Mairal [9]. These
explicit features can have infinite dimension in general and even if one uses
a finite-dimensional approximation to the features, the computational gain of
random features will be lost due to expensive tensor product operations.

A popular line of work on kernel approximation problem is based on the Fourier
features method [36], which works well for shift-invariant kernels and with
some modifications can embed the Gaussian kernel near optimally [8]. Other
random feature constructions have been suggested for a variety of kernels,
e.g., arc-cosine kernels [12], polynomial kernels [35]. In linear sketching
literature, Avron et al. [7] proposed a subspace embedding for the polynomial
kernel which was recently extended to general dot product kernels [20]. The
runtime of this method, while nearly linear in sparsity of the input dataset,
scales exponentially in kernel’s degree. Recently, Ahle et al. [1] improved
this exponential dependence to polynomial which enabled them to sketch high-
degree polynomial kernels and led to near-optimal embeddings for Gaussian
kernel. In fact, this sketching technology constitutes one of the main
ingredients of our proposed methods. Additionally, combining sketching with
leverage score sampling can improve the runtime of the polynomial kernel
embeddings [39].

###  1.3 Preliminaries: PolySketch and TensorSRHT Transforms

#### Notations.

We use $[n]:=\\{1,\dots,n\\}$. We denote the tensor (a.k.a. Kronecker) product
by $\otimes$ and the element-wise (a.k.a. Hadamard) product of two vectors or
matrices by $\odot$. Although tensor products are multidimensional objects, we
often associate $x\otimes y$ with a single dimensional vector
$(x_{1}y_{1},x_{2}y_{1},\ldots x_{m}y_{1},x_{1}y_{2},\ldots x_{m}y_{2},\ldots
x_{m}y_{n})$. For shorthand, we use the notation $x^{\otimes p}$ to denote
$\underbrace{x\otimes\ldots\otimes x}_{p\text{ terms}}$, the $p$-fold self-
tensoring of $x$. Another related operation that we use is the _direct sum_ of
vectors: $x\oplus y:=\left[x^{\top},y^{\top}\right]^{\top}$. We need notation
for _sub-tensors_ of a tensor. For instance, for a $3$-dimensional tensor
${\bm{Y}}\in\mathbb{R}^{m\times n\times d}$ and every $l\in[d]$, we denote by
${\bm{Y}}_{(:,:,l)}$ the $m\times n$ matrix that is defined as
$\left[{\bm{Y}}_{(:,:,l)}\right]_{i,j}:={\bm{Y}}_{i,j,l}$ for
$i\in[m],j\in[n]$. 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
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})$.
For two functions $f$ and $g$ we denote their twofold composition by $f\circ
g$, defined as $f\circ g(\alpha):=f(g(\alpha))$.

The TensorSRHT is a norm-preserving dimensionality reduction that can be
applied to the tensor product of two vectors very quickly [1]. This
transformation is a generalization of the Subsampled Randomized Hadamard
Transform (SRHT) [2] and can be computed in near linear time using the FFT
algorithm. The PolySketch extends the idea behind TensorSRHT to high-degree
tensor products by recursively sketching pairs of vectors in a binary tree
structure. This sketch preserves the norm of vectors in $\mathbb{R}^{d^{p}}$
with high probability and can be applied to tensor product vectors very
quickly. The following Lemma, summarizes Theorems 1.2 and 1.3 of [1] and is
proved in Appendix B.

######  Lemma 1 (PolySketch).

For every integers $p,d\geq 1$ and every $\varepsilon,\delta>0$, there exists
a distribution on random matrices ${\bm{Q}}^{p}\in\mathbb{R}^{m\times d^{p}}$,
called degree $p$ PolySketch such that (1) for some
$m=\mathcal{O}\left(\frac{p}{\varepsilon^{2}}\log^{3}\frac{1}{\varepsilon\delta}\right)$
and any $y\in\mathbb{R}^{d^{p}}$,
$\Pr\left[\|{\bm{Q}}^{p}y\|_{2}^{2}\in(1\pm\varepsilon)\|y\|_{2}^{2}\right]\geq
1-\delta$; (2) for any $x\in\mathbb{R}^{d}$, if $e_{1}\in\mathbb{R}^{d}$ is
the standard basis vector along the first coordinate, the total time to
compute ${\bm{Q}}^{p}(x^{\otimes(p-j)}\otimes{e}_{1}^{\otimes j})$ for all
$j=0,1,\dots,p$ is
$\mathcal{O}\left(pm\log^{2}m+\min\left\\{\frac{p^{3/2}}{\varepsilon}\log\frac{1}{\delta}~{}{\rm
nnz}(x),pd\log d\right\\}\right)$; (3) for any collection of vectors
$v_{1},\dots,v_{p}\in\mathbb{R}^{d}$, the time to compute
${\bm{Q}}^{p}\left(v_{1}\otimes\dots\otimes v_{p}\right)$ is bounded by
$\mathcal{O}\left(pm\log
m+\frac{p^{3/2}}{\varepsilon}d\log\frac{1}{\delta}\right)$; (4) for any
$\lambda>0$ and any matrix ${\bm{A}}\in\mathbb{R}^{d^{p}\times n}$, where the
statistical dimension of ${\bm{A}}^{\top}{\bm{A}}$ is $s_{\lambda}$, there
exists some
$m=\mathcal{O}\left(\frac{p^{4}s_{\lambda}}{\varepsilon^{2}}\log^{3}\frac{n}{\varepsilon\delta}\right)$
such that,

$\displaystyle\Pr\left[(1-\varepsilon)\left({\bm{A}}^{\top}{\bm{A}}+\lambda{\bm{I}}\right)\preceq({\bm{Q}}^{p}{\bm{A}})^{\top}({\bm{Q}}^{p}{\bm{A}})+\lambda{\bm{I}}\preceq(1+\varepsilon)\left({\bm{A}}^{\top}{\bm{A}}+\lambda{\bm{I}}\right)\right]\geq
1-\delta.$ (1)

##  2 ReLU Neural Tangent Kernel

Arora et al. [5] showed how to exactly compute the NTK of a fully-connected
network $\Theta_{\tt ntk}^{(L)}(y,z)$ for any pair of vectors
$y,z\in\mathbb{R}^{d}$ using a dynamic program in $\mathcal{O}(d+L)$ time.
However, it is hard to gain insight into the structure of this kernel using
that DP expression which involves recursive applications of nontrivial
expectations. Fortunately, for the important case of ReLU activation this
kernel takes an extremely nice and highly structured form. The NTK in this
case can be fully characterized by a univariate function $K_{\tt
relu}^{(L)}:[-1,1]\to\mathbb{R}$ that we refer to as _ReLU-NTK_ , which is the
composition of the arc-cosine kernels [12] and was recently derived in [9].
Exploiting this special structure is the key to designing efficient sketching
methods and random features for this kernel.

######  Definition 1 (ReLU-NTK function).

For every integer $L>0$, the $L$-layered ReLU-NTK function $K_{\tt
relu}^{(L)}:[-1,1]\to\mathbb{R}$ is defined via following procedure, for every
$\alpha\in[-1,1]$:

  1. 1.

Let $\kappa_{0}(\alpha)$ and $\kappa_{1}(\alpha)$ be $0^{th}$ and $1^{st}$
order arc-cosine kernels [12] defined as follows,

$\displaystyle\kappa_{0}(\alpha):=\frac{1}{\pi}\left(\pi-\arccos\left(\alpha\right)\right),\text{~{}~{}~{}and~{}~{}~{}}\kappa_{1}(\alpha):=\frac{1}{\pi}\left(\sqrt{1-\alpha^{2}}+\alpha\cdot\left(\pi-\arccos\left(\alpha\right)\right)\right).$
(2)

  2. 2.

Let $\Sigma_{\tt relu}^{(0)}(\alpha):=\alpha$ and for $\ell=1,2,\ldots L$,
define $\Sigma_{\tt relu}^{(\ell)}(\alpha)$ and $\dot{\Sigma}_{\tt
relu}^{(\ell)}(\alpha)$ as follows,

$\displaystyle\Sigma_{\tt
relu}^{(\ell)}(\alpha):=\underbrace{\kappa_{1}\circ\kappa_{1}\circ\dots\circ\kappa_{1}}_{\ell\text{-fold
self composition}}(\alpha),\text{~{}~{}~{}and~{}~{}~{}}\dot{\Sigma}_{\tt
relu}^{(\ell)}(\alpha):=\kappa_{0}\left(\Sigma_{\tt
relu}^{(\ell-1)}(\alpha)\right).$ (3)

  3. 3.

Let $K_{\tt relu}^{(0)}(\alpha):=\Sigma_{\tt relu}^{(0)}(\alpha)=\alpha$ and
for $\ell=1,2,\ldots L$, define $K_{\tt relu}^{(\ell)}(\alpha)$ recursively as
follows,

$\displaystyle K_{\tt relu}^{(\ell)}(\alpha):=K_{\tt
relu}^{(\ell-1)}(\alpha)\cdot\dot{\Sigma}_{\tt
relu}^{(\ell)}(\alpha)+\Sigma_{\tt relu}^{(\ell)}(\alpha).$ (4)

The connection between ReLU-NTK function $K_{\tt relu}^{(L)}$ and the NTK
kernel $\Theta_{\tt ntk}^{(L)}$ is formalized bellow,

$\displaystyle\Theta_{\tt ntk}^{(L)}(y,z)\equiv\|y\|_{2}\|z\|_{2}\cdot K_{\tt
relu}^{(L)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right),\text{~{}~{}~{}for any
}y,z\in\mathbb{R}^{d}.$ (5)

This shows that the NTK is a _normalized dot-product kernel_ which can be
fully characterized by $K_{\tt relu}^{(L)}:[-1,1]\to\mathbb{R}$, plotted in
??. As shown in ??, this function is smooth and can be tightly approximated
with a low-degree polynomial. It is evident that for larger values of $L$,
$K_{\tt relu}^{(L)}(\cdot)$ converges to a _knee shape_ , i.e., it has a
nearly constant value of roughly $0.3(L+1)$ on the interval
$[-1,1-\mathcal{O}(L^{-1})]$, and on the interval $[1-\mathcal{O}(L^{-1}),1]$
its value sharply increases to $L+1$ at $\alpha=1$.

Figure 1: (Left) Normalized ReLU-NTK function $K_{\tt relu}^{(L)}(\cdot)$ for
$L=\\{2,4,8,16,32\\}$ and (Right) a degree-$8$ polynomial approximation of
ReLU-NTK with $L=3$.

##  3 Sketching and Random Features for NTK

The main results of this section are efficient oblivious sketching as well as
random features for the fully-connected NTK. As shown in ?? and ??, the NTK
$\Theta_{\tt ntk}^{(L)}$, is constructed by recursive composition of arc-
cosine kernels $\kappa_{1}(\cdot)$ and $\kappa_{0}(\cdot)$. So, to design
efficient sketches for the NTK we crucially need efficient methods for
approximating these functions. Generally, there are two main approaches to
approximating these functions; one is random features sampling and the other
is truncated Taylor series expansion coupled with fast sketching. We design
algorithms by exploiting both of these techniques.

###  3.1 NTK Sketch

Our main tool is approximating the arc-cosine kernels with low-degree
polynomials, and then applying PolySketch to the resulting polynomial kernels.
The features for multi-layered NTK are the recursive tensor product of arc-
cosine sketches at consecutive layers, which in turn can be sketched
efficiently using PolySketch. We present our oblivious sketch in ??.

Algorithm 1 NTKSketch for fully-connected ReLU networks

1: input: vector $x\in\mathbb{R}^{d}$, network depth $L$, error and failure
parameters $\varepsilon,\delta>0$

2: Choose integers
$s=\widetilde{\mathcal{O}}\left(\frac{L^{2}}{\varepsilon^{2}}\right)$,
$n_{1}=\widetilde{\mathcal{O}}\left(\frac{L^{4}}{\varepsilon^{4}}\right)$,
$r=\widetilde{\mathcal{O}}\left(\frac{L^{6}}{\varepsilon^{4}}\right)$,
$m=\widetilde{\mathcal{O}}\left(\frac{L^{8}}{\varepsilon^{\frac{16}{3}}}\right)$,
and
$s^{*}=\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\log\frac{1}{\delta}\right)$
appropriately222$\widetilde{\mathcal{O}}(\cdot)$ suppresses
$\text{poly}\log\frac{L}{\varepsilon\delta}$ factors.

3: For $p=\left\lceil 2L^{2}/\varepsilon^{\frac{4}{3}}\right\rceil$ and
$p^{\prime}=\left\lceil 9L^{2}/\varepsilon^{2}\right\rceil$, polynomials
$P_{\tt relu}^{(p)}(\cdot)$ and $\dot{P}_{\tt relu}^{(p^{\prime})}(\cdot)$ are
defined as,  $\begin{split}P_{\tt
relu}^{(p)}(\alpha)\equiv\sum_{j=0}^{2p+2}c_{j}\cdot\alpha^{j}&:=\frac{1}{\pi}+\frac{\alpha}{2}+\frac{1}{\pi}\sum_{i=0}^{p}\frac{(2i)!\cdot\alpha^{2i+2}}{2^{2i}(i!)^{2}(2i+1)(2i+2)},\\\
\dot{P}_{\tt
relu}^{(p^{\prime})}(\alpha)\equiv\sum_{j=0}^{2p^{\prime}+1}b_{j}\cdot\alpha^{j}&:=\frac{1}{2}+\frac{1}{\pi}\sum_{i=0}^{p^{\prime}}\frac{(2i)!}{2^{2i}(i!)^{2}(2i+1)}\cdot\alpha^{2i+1}.\end{split}$
(6)

4: $\phi^{(0)}(x)\leftarrow\|x\|_{2}^{-1}\cdot{\bm{Q}}^{1}\cdot x$, where
${\bm{Q}}^{1}\in\mathbb{R}^{r\times d}$ is a degree-$1$ PolySketch as per ??

5: $\psi^{(0)}(x)\leftarrow{\bm{V}}\cdot\phi^{(0)}(x)$, where
${\bm{V}}\in\mathbb{R}^{s\times r}$ is an instance of SRHT [2]

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

7: Let ${\bm{Q}}^{2p+2}\in\mathbb{R}^{m\times r^{2p+2}}$ be a degree-$2p+2$
PolySketch. Also, let ${\bm{T}}\in\mathbb{R}^{r\times(2p+3)\cdot m}$ be an
instance of SRHT. For every $l=0,1,\ldots,2p+2$, compute:
$Z^{(\ell)}_{l}(x)\leftarrow{\bm{Q}}^{2p+2}\left(\left[\phi^{(\ell-1)}(x)\right]^{\otimes
l}\otimes{e}_{1}^{\otimes
2p+2-l}\right),~{}\phi^{(\ell)}(x)\leftarrow{\bm{T}}\cdot\bigoplus_{l=0}^{2p+2}\sqrt{c_{l}}Z^{(\ell)}_{l}(x)$
(7)

8: Let ${\bm{Q}}^{2p^{\prime}+1}\in\mathbb{R}^{n_{1}\times r^{2p^{\prime}+1}}$
be a degree-$2p^{\prime}+1$ PolySketch. Also, let
${\bm{W}}\in\mathbb{R}^{s\times(2p^{\prime}+2)\cdot r}$ be an instance of
SRHT. For every $l=0,1,\ldots,2p^{\prime}+1$, compute:
$Y^{(\ell)}_{l}(x)\leftarrow{\bm{Q}}^{2p^{\prime}+1}\left(\left[\phi^{(\ell-1)}(x)\right]^{\otimes
l}\otimes{e}_{1}^{\otimes
2p^{\prime}+1-l}\right),~{}\dot{\phi}^{(\ell)}(x)\leftarrow{\bm{W}}\cdot\bigoplus_{l=0}^{2p^{\prime}+1}\sqrt{b_{l}}Y^{(\ell)}_{l}(x)$
(8)

9: Let ${\bm{Q}}^{2}\in\mathbb{R}^{s\times s^{2}}$ be a degree-$2$ PolySketch.
Also, let ${\bm{R}}\in\mathbb{R}^{s\times(s+r)}$ be an SRHT. Compute:
$\psi^{(\ell)}(x)\leftarrow{\bm{R}}\cdot\left({\bm{Q}}^{2}\left(\psi^{(\ell-1)}(x)\otimes\dot{\phi}^{(\ell)}(x)\right)\oplus\phi^{(\ell)}(x)\right).$
(9)

10: Let ${\bm{G}}\in\mathbb{R}^{s^{*}\times s}$ be a matrix of i.i.d. entries
with distribution $\mathcal{N}(0,\frac{1}{s^{*}})$. Compute:  $\Psi_{\tt
ntk}^{(L)}(x)\leftarrow{\|x\|_{2}}\cdot{\bm{G}}\cdot\psi^{(L)}(x).$ (10)

11: return ${\Psi}_{\tt ntk}^{(L)}(x)$

Now we present our main theorem on NTKSketch algorithm as follows.

######  Theorem 1.

For every integers $d\geq 1$ and $L\geq 2$, and any $\varepsilon,\delta>0$,
let $\Theta_{\tt ntk}^{(L)}:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$
be the $L$-layered NTK with ReLU activation as per ?? and ??. Then there
exists a randomized map $\Psi_{\tt
ntk}^{(L)}:\mathbb{R}^{d}\to\mathbb{R}^{s^{*}}$ for some
$s^{*}=\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\log\frac{1}{\delta}\right)$
such that the following invariants hold,

  1. 1.

For any vectors $y,z\in\mathbb{R}^{d}$: $\Pr\left[\left|\left<\Psi_{\tt
ntk}^{(L)}(y),\Psi_{\tt ntk}^{(L)}(z)\right>-\Theta_{\tt
ntk}^{(L)}(y,z)\right|>\varepsilon\cdot\Theta_{\tt
ntk}^{(L)}(y,z)\right]\leq\delta$.

  2. 2.

For every vecor $x\in\mathbb{R}^{d}$, the time to compute $\Psi_{\tt
ntk}^{(L)}(x)$ is
$\mathcal{O}\left(\frac{L^{11}}{\varepsilon^{6.7}}\log^{3}\frac{L}{\varepsilon\delta}+\frac{L^{3}}{\varepsilon^{2}}\log\frac{L}{\varepsilon\delta}\cdot{\rm
nnz}(x)\right)$.

For a proof, see ??. One can observe that the runtime of our NTKSketch is
faster than the gradient features of an ultra-wide random DNN, studied by
Arora et al. [5], by a factor of $L^{3}/\varepsilon^{2}$.

###  3.2 NTK Random Features

The main difference between our random features construction and NTKSketch is
the use of random features for approximating arc-cosine kernels $\kappa_{0}$
and $\kappa_{1}$ in ??. For any $x\in\mathbb{R}^{d}$, we denote

$\displaystyle\Phi_{0}(x):=\sqrt{\frac{2}{m_{0}}}\
\mathrm{Step}\left(\left[w_{1},\ldots,w_{m_{0}}\right]^{\top}x\right),~{}~{}~{}\Phi_{1}(x):=\sqrt{\frac{2}{m_{1}}}\
\mathrm{ReLU}\left(\left[w^{\prime}_{1},\ldots,w_{m_{1}}^{\prime}\right]^{\top}x\right),$
(11)

where
$w_{1},\dots,w_{m_{0}},w^{\prime}_{1},\dots,w^{\prime}_{m_{1}}\in{\mathbb{R}}^{d}$
are i.i.d. samples from $\mathcal{N}({0},{\bm{I}}_{d})$. Cho and Saul [12]
showed that
${\mathbb{E}}[\left\langle{\Phi_{0}(y),\Phi_{0}(z)}\right\rangle]=\kappa_{0}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$ and
${\mathbb{E}}[\left\langle{\Phi_{1}(y),\Phi_{1}(z)}\right\rangle]=\|y\|_{2}\|z\|_{2}\cdot\kappa_{1}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$. The feature map for multi-layered NTK
can be obtained by recursive tensoring of random feature maps for arc-cosine
kernels at each layer of the network. However, one major drawback of such
explicit tensoring is that the number of features, and thus the runtime, will
be exponential in depth $L$. In order to make the feature map more compact, we
utilize a degree-2 PolySketch ${\bm{Q}}^{2}$ to reduce the dimension of the
tensor products at each layer and get rid of exponential dependence on $L$. We
present the performance guarantee of our random features, defined in ??, in
??.

######  Theorem 2.

Given $y,z\in\mathbb{R}^{d}$ and $L\geq 2$, let $\Theta_{\tt ntk}^{(L)}$ the
$L$-layered fully-connected ReLU NTK. For $\varepsilon,\delta>0$, there exist
$m_{0}=\mathcal{O}\left(\frac{L^{2}}{\varepsilon^{2}}\log\frac{L}{\delta}\right),m_{1}=\mathcal{O}\left(\frac{L^{6}}{\varepsilon^{4}}\log\frac{L}{\delta}\right),m_{s}=\mathcal{O}\left(\frac{L^{2}}{\varepsilon^{2}}\log^{3}\frac{L}{\varepsilon\delta}\right)$,
such that,

$\displaystyle\Pr\left[\left|\left\langle{{\Psi}_{\tt rf}^{(L)}(y),\Psi_{\tt
rf}^{(L)}(z)}\right\rangle-\Theta_{\tt
ntk}^{(L)}(y,z)\right|>\varepsilon\cdot\Theta_{\tt
ntk}^{(L)}(y,z)\right]\geq\delta,$ (12)

where $\Psi_{\tt rf}^{(L)}(y),\Psi_{\tt
rf}^{(L)}(z)\in{\mathbb{R}}^{m_{1}+m_{\mathtt{s}}}$ are the outputs of ??,
using the same randomness.

The proof of ?? is provided in ??. Arora et al. [5] proved 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}\frac{L}{\delta}+\frac{L^{6}}{\varepsilon^{4}}\cdot\log\frac{L}{\delta}\cdot
d\right)$ which is larger than ours by a factor of
$\frac{L^{7}}{\varepsilon^{4}}\log\frac{L}{\delta}$. In ??, we also
empirically show that ?? requires far fewer features than random gradients.

Algorithm 2 Random Features for ReLU NTK via PolySketch

1: input: vector $x\in\mathbb{R}^{d}$, network depth $L$, feature dimensions
$m_{0}$, $m_{1}$, and $m_{s}$

2: ${\psi}_{\tt rf}^{(0)}(x)\leftarrow{x}/{\|x\|_{2}},\phi_{\tt
rf}^{(0)}(x)\leftarrow{x}/{\|x\|_{2}}$

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

4: $\dot{\phi}_{\tt rf}^{(\ell)}(x)\leftarrow\Phi_{0}\left(\phi_{\tt
rf}^{(\ell-1)}(x)\right)$, where $\Phi_{0}$ is defines as per ?? with $m_{0}$
features

5: $\phi_{\tt rf}^{(\ell)}(x)\leftarrow\Phi_{1}\left(\phi_{\tt
rf}^{(\ell-1)}(x)\right)$, where $\Phi_{1}$ is defines as per ?? with $m_{1}$
features

6: Draw a degree-$2$ PolySketch ${\bm{Q}}^{2}$ that maps to
$\mathbb{R}^{m_{s}}$ and compute:  ${\psi}_{\tt
rf}^{(\ell)}(x)\leftarrow\phi_{\tt
rf}^{(\ell)}(x)\oplus{\bm{Q}}^{2}\cdot\left(\dot{\phi}_{\tt
rf}^{(\ell)}(x)\otimes{\psi}_{\tt rf}^{(\ell-1)}(x)\right)$

7: return ${\Psi}_{\tt rf}^{(L)}(x)\leftarrow\|x\|_{2}\cdot{\psi}_{\tt
rf}^{(L)}(x)$

###  3.3 Spectral Approximation for NTK via Leverage Scores Sampling

Although the above NTK approximations can estimate the kernel function itself,
it is still questionable how it affects the performance of downstream tasks.
Several works on kernel approximation adopt spectral approximation bound with
regularization $\lambda>0$ and approximation factor $\varepsilon>0$, that is,

$\displaystyle(1-\varepsilon)({\bm{K}}^{(L)}_{\tt
ntk}+\lambda{\bm{I}})\preceq(\bm{\Psi}^{(L)})^{\top}\bm{\Psi}^{(L)}+\lambda{\bm{I}}\preceq(1+\varepsilon)({\bm{K}}^{(L)}_{\tt
ntk}+\lambda{\bm{I}}),$ (13)

where
$\bm{\Psi}^{(L)}:=\left[\Psi^{(L)}(x_{1}),\dots,\Psi^{(L)}(x_{n})\right]$ and
$[{\bm{K}}^{(L)}_{\tt ntk}]_{i,j}=\Theta_{\tt ntk}^{(L)}(x_{i},x_{j})$. The
spectral bound can provide rigorous guarantees for downstream applications
including kernel ridge regression [8], clustering and PCA [31]. We first
provide spectral bounds for arc-cosine kernels, then we present our spectral
approximation bound for two-layered ReLU networks, which is the first in the
literature.

To guarantee that the arc-cosine random features in ?? provide spectral
approximation, we will use the leverage score sampling framework of [8, 28].
We reduce the variance of random features by performing importance sampling.
The challenge is to find a proper modified distribution that certainly reduces
the variance. It turns out that the original $0^{th}$ order arc-cosine random
features has a small enough variance. More precisely, let ${\bm{K}}_{0}$ be
the $0^{th}$ order arc-cosine kernel matrix, i.e.,
$[{\bm{K}}_{0}]_{i,j}=\kappa_{0}\left(\frac{\langle
x_{i},x_{j}\rangle}{\|x_{i}\|_{2}\|x_{j}\|_{2}}\right)$, and
$\bm{\Phi}_{0}:=\left[\Phi_{0}(x_{1}),\ldots,\Phi_{0}(x_{n})\right]$, where
$\Phi_{0}(x)$ is defined in ??. If the number of features
$m_{0}\geq\frac{8}{3}\frac{n}{\lambda\varepsilon^{2}}\log\left(\frac{16s_{\lambda}}{\delta}\right)$,
then

$\displaystyle\Pr\left[(1-\varepsilon)({\bm{K}}_{0}+\lambda{\bm{I}})\preceq\bm{\Phi}_{0}^{\top}\bm{\Phi}_{0}+\lambda{\bm{I}}\preceq(1+\varepsilon)({\bm{K}}_{0}+\lambda{\bm{I}})\right]\geq
1-\delta.$ (14)

Next, we consider spectral approximation of the $1^{st}$ order arc-cosine
kernel. Unlike the previous case, modifications of the sampling distribution
are required. Specifically, for any $x\in\mathbb{R}^{d}$, let

$\widetilde{\Phi}_{1}(x)=\sqrt{\frac{2d}{m_{1}}}\mathrm{ReLU}\left(\left[\frac{w_{1}}{\|w_{1}\|_{2}},\ldots,\frac{w_{m_{1}}}{\|w_{m_{1}}\|_{2}}\right]^{\top}x\right),$
(15)

where $w_{1},\dots,w_{m_{1}}\in{\mathbb{R}}^{d}$ are i.i.d. samples from
$p(w):=\frac{1}{(2\pi)^{d/2}d}\left\|w\right\|_{2}^{2}\exp\left(-\frac{1}{2}\left\|w\right\|_{2}^{2}\right)$.
For this modified features, let ${\bm{X}}\in\mathbb{R}^{d\times n}$ be the
dataset, ${\bm{K}}_{1}$ be the $1^{st}$ order arc-cosine kernel matrix, i.e.,
$[{\bm{K}}_{1}]_{i,j}=\|x_{i}\|_{2}\|x_{j}\|_{2}\cdot\kappa_{1}\left(\frac{\langle
x_{i},x_{j}\rangle}{\|x_{i}\|_{2}\|x_{j}\|_{2}}\right)$, and
$\bm{\Phi}_{1}:=\left[\widetilde{\Phi}_{1}(x_{1}),\ldots,\widetilde{\Phi}_{1}(x_{n})\right]$.
If the number of features
$m_{1}\geq\frac{8}{3}\frac{d}{\varepsilon^{2}}\cdot\min\left\\{\mathrm{rank}({\bm{X}})^{2},\frac{\|{\bm{X}}\|_{2}^{2}}{\lambda}\right\\}\log\left(\frac{16s_{\lambda}}{\delta}\right)$,
then

$\displaystyle\Pr\left[(1-\varepsilon)({\bm{K}}_{1}+\lambda{\bm{I}})\preceq\bm{\Phi}_{1}^{\top}\bm{\Phi}_{1}+\lambda{\bm{I}}\preceq(1+\varepsilon)({\bm{K}}_{1}+\lambda{\bm{I}})\right]\geq
1-\delta.$ (16)

The details are provided in ?? and ??. We are now ready to state our spectral
approximation bound for our modified random features.

######  Theorem 3.

Given a dataset ${\bm{X}}\in{\mathbb{R}}^{d\times n}$ with
$\|{\bm{X}}_{(:,i)}\|_{2}\leq 1$ for every $i\in[n]$, let ${\bm{K}}_{\tt
ntk},{\bm{K}}_{0},{\bm{K}}_{1}$ be kernel matrices for two-layered ReLU NTK
and arc-cosine kernels of $0^{th}$ and $1^{st}$ order, respectively. For any
$\lambda>0$, suppose $s_{\lambda}$ is the statistical dimension of
${\bm{K}}_{\tt ntk}$. Modify ?? by replacing $\Phi_{1}(\cdot)$ in line 5 with
$\widetilde{\Phi}_{1}(\cdot)$ defined in ??. For any $\varepsilon,\delta>0$,
let $\bm{\Psi}_{\tt rf}^{(L)}\in{\mathbb{R}}^{(m_{1}+m_{s})\times n}$ be the
output matrix of this algorithm with $L=1$. There exist
$m_{0}=\mathcal{O}\left(\frac{n}{\varepsilon^{2}\lambda}\log\frac{s_{\lambda}}{\delta}\right),m_{1}=\mathcal{O}\left(\frac{d}{\varepsilon^{2}}\cdot\min\left\\{\mathrm{rank}({\bm{X}})^{2},\frac{\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda}\right\\}\log\frac{s_{\lambda}}{\delta}\right),m_{s}=\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\cdot\frac{n}{1+\lambda}\log^{3}\frac{n}{\varepsilon\delta}\right)$
such that,

$\displaystyle\Pr\left[(1-\varepsilon)\left({\bm{K}}_{\tt
ntk}+\lambda{\bm{I}}\right)\preceq\left(\bm{\Psi}_{\tt
rf}^{(L)}\right)^{\top}\bm{\Psi}_{\tt
rf}^{(L)}+\lambda{\bm{I}}\preceq(1+\varepsilon)\left({\bm{K}}_{\tt
ntk}+\lambda{\bm{I}}\right)\right]\geq 1-\delta.$ (17)

For a proof see ??. To generalize the current proof technique to deeper
networks, one needs a monotone property of arc-cosine kernels, i.e.,
$\kappa_{1}({\bm{X}})\preceq\kappa_{1}({\bm{Y}})$ for
${\bm{X}}\preceq{\bm{Y}}$. However, this property does not hold in general and
we leave the extension to deeper networks to future work.

##  4 Sketching Convolutional Neural Tangent Kernel

In this section, we design and analyze an efficient sketching method for the
Convolutional Neural Tangent Kernel (CNTK). We focus mainly on CNTK with
Global Average Pooling (GAP), which exhibits superior empirical performance
compared to vanilla CNTK with no pooling [5], however, our techniques can be
applied to the vanilla version, as well. Using the DP of Arora et al. [5], the
number of operations needed for exact computation of the depth-$L$ CNTK value
$\Theta_{\tt cntk}^{(L)}(y,z)$ for images $y,z\in\mathbb{R}^{d\times d}$ is
$\Omega\left(d^{4}\cdot L\right)$, which is extremely slow particularly due to
its quadratic dependence on the number of pixels of input images $d^{2}$.
Fortunately, we are able to show that the CNTK for the important case of ReLU
activation is a highly structured object that can be fully characterized in
terms of tensoring and composition of arc-cosine kernels, and exploiting this
special structure is key to designing efficient sketching methods for the
CNTK. Unlike the fully-connected NTK, CNTK is not a simple dot-product kernel
function like ??. The key reason being that CNTK works by partitioning its
input images into patches and locally transforming the patches at each layer,
as opposed to the NTK which operates on the entire input vectors. We present
our derivation of the ReLU CNTK function and its main properties in ??.

Similar to NTKSketch our method relies on approximating the arc-cosine kernels
with low-degree polynomials via Taylor expansion, and then applying PolySketch
to the resulting polynomial kernels. Our sketch computes the features for each
pixel of the input image, by tensor product of arc-cosine sketches at
consecutive layers, which in turn can be sketched efficiently using PolySketch
. Additionally, the features of pixels that lie in the same patch get _locally
combined_ at each layer via direct sum operation. This precisely corresponds
to the convolution operation in neural networks. We present our CNTKSketch
algorithm in ?? and give its performance guarantee in the following theorem.

######  Theorem 4.

For every positive integers $d_{1},d_{2},c$ and $L\geq 2$, and every
$\varepsilon,\delta>0$, if we let $\Theta_{\tt
cntk}^{(L)}:\mathbb{R}^{d_{1}\times d_{2}\times
c}\times\mathbb{R}^{d_{1}\times d_{2}\times c}\to\mathbb{R}$ be the
$L$-layered CNTK with ReLU activation and GAP given in [5], then there exist a
randomized map $\Psi_{\tt cntk}^{(L)}:\mathbb{R}^{d_{1}\times d_{2}\times
c}\to\mathbb{R}^{s^{*}}$ for some
$s^{*}=\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\log\frac{1}{\delta}\right)$
such that:

  1. 1.

For any images $y,z\in\mathbb{R}^{d_{1}\times d_{2}\times c}$:
$\Pr\left[\left<\Psi_{\tt cntk}^{(L)}(y),\Psi_{\tt
cntk}^{(L)}(z)\right>\notin(1\pm\varepsilon)\cdot\Theta_{\tt
cntk}^{(L)}(y,z)\right]\leq\delta$.

  2. 2.

For every image $x\in\mathbb{R}^{d_{1}\times d_{2}\times c}$, time to compute
$\Psi_{\tt cntk}^{(L)}(x)$ is
$\mathcal{O}\left(\frac{L^{11}}{\varepsilon^{6.7}}\cdot(d_{1}d_{2})\cdot\log^{3}\frac{d_{1}d_{2}L}{\varepsilon\delta}\right)$.

The proof is in ??. Runtime of our CNTKSketch is only linear in the number of
image pixels $d_{1}d_{2}$, which is in stark contrast to quadratic scaling of
the exact CNTK computation [5].

##  5 Experiments

(a) MNIST with NTK

(b) CIFAR-10 with CNTK

Figure 2: Test accuracy of: (a) approximate NTK methods (GradRF, NTKSketch and
NTKRF) on MNIST and (b) approximate CNTK methods (GradRF and CNTKSketch) on
CIFAR-10.

In this section, we empirically show that running least squares regression on
the features generated by our methods is extremely fast and effective for
learning with NTK and CNTK kernel machines. We run experiments on a system
with an Intel E5-2630 CPU with 256 GB RAM and a single GeForce RTX 2080 GPUs
with 12 GB RAM.

###  5.1 NTK Classification on MNIST

We first benchmark our proposed NTK approximation algorithms on MNIST [25]
dataset and compare against gradient-based NTK random features [5] (GradRF) as
a baseline method. To apply our methods and GradRF into classification task,
we encode class labels into one-hot vectors with zero-mean and solve the ridge
regression problem. We search the ridge parameter with a random subset of
training set and choose the one that achieves the best validation accuracy. We
use the ReLU network with depth $L=1$. In ??, we observe that our random
features (NTKRF ) achieves the best test accuracy. The NTKSketch narrowly
follows the performance of NTKRF and the Grad-RF is the worst method which
confirms the observations of Arora et al. [5], i.e., gradient of a finite
width network degrades practical performances.

######  Remark 1 (Optimizing NTKSketch for Deeper Nets).

As shown in ??, the NTK is a normalized dot-product kernel characterized by
the function $K^{(L)}_{\tt relu}(\alpha)$. This function can be easily
computed using $\mathcal{O}(L)$ operations at any desired $\alpha\in[-1,1]$,
therefore, we can efficiently fit a polynomial to this function using
numerical methods (for instance, it is shown in ?? that a degree-$8$
polynomial can tightly approximate the depth-$3$ ReLU-NTK function $K_{\tt
relu}^{(3)}$). Then, we can efficiently sketch the resulting polynomial kernel
using PolySketch , as was previously done for Gaussian and general dot-product
kernels [1, 39]. Therefore, we can accelerate our CNTKSketch for deeper
networks ($L>2$), using this heuristic.

###  5.2 CNTK Classification on CIFAR-10

Next we test our CNTKSketch on CIFAR-10 dataset [24]. We choose a
convolutional network of depth $L=3$ and compare CNTKSketch and GradRF for
various feature dimensions. We borrow results of both CNTK and CNN from Arora
et al. [5]. The results are provided in ?? and ??. Somewhat surprisingly,
CNTKSketch even performs better than the exact CNTK regression by achieving
$72.06\%$ when feature dimension is set to $16{,}384$. The likely explanation
is that CNTKSketch takes advantages of implicit regularization effects of
approximate feature map and powerful expressiveness of the CNTK. Moreover,
computing the CNTK matrix takes over 250 hours (12 days) under our setting
which is at least $150\times$ slower than our CNTKSketch.

Table 1: Test accuracy and runtime to solve CNTK regression and its
approximations on CIFAR-10. (*) means that the result is copied from Arora et
al. [5].

|  CNTKSketch (ours) | GradRF | Exact CNTK | CNN
---|---|---|---|---
Feature dimension | $4{,}096$ | $8{,}192$ | $16{,}384$ | $9{,}328$ | $17{,}040$ | $42{,}816$ |  |
Test accuracy (%) | 67.58 | 70.46 | 72.06 | 62.49 | 62.57 | 65.21 | 70.47∗ | 63.81∗
Time (s) | 780 | 1,870 | 5,160 | 300 | 360 | 580 | > 1,000,000 |

Table 2: MSE and runtime on large-scale UCI datasets. We measure the entire
time to solve kernel ridge regression. $(-)$ means Out-of-Memory error.

| $\mathrm{MillionSongs}$ | $\mathrm{WorkLoads}$ | $\mathrm{CT}$ | $\mathrm{Protein}$
---|---|---|---|---
# of data points ($n$) | 467,315 | 179,585 | 53,500 | 39,617
| MSE | Time (s) | MSE | Time (s) | MSE | Time (s) | MSE | Time (s)
RBF Kernel | $-$ | $-$ | $-$ | $-$ | 35.37 | 59.23 | 18.96 | 46.45
RFF | 109.50 | 231 | $4.034\times 10^{4}$ | 53.0 | 48.20 | 15.2 | 19.72 | 12.1
NTK | $-$ | $-$ | $-$ | $-$ | 30.52 | 72.10 | 20.24 | 76.93
NTKRF (ours) | 94.27 | 95 | $3.554\times 10^{4}$ | 35.7 | 46.91 | 2.12 | 20.51 | 4.3
NTKSketch (ours) | 92.83 | 36 | $3.538\times 10^{4}$ | 27.5 | 46.52 | 18.8 | 21.19 | 14.91

###  5.3 Regression on Large-scale UCI Datasets

We also demonstrate the computational efficiency of our NTKSketch and NTKRF
using $4$ large-scale UCI regression datasets [17] by comparing against exact
NTK, RBF as well as Random Fourier Features (RFF). For our methods and RFF, we
fix the output dimension to $m=8{,}192$ for all datasets. In ??, we report the
runtime to compute feature map or kernel matrix and evaluate the averaged mean
squared errors (MSE) on the test set via $4$-fold cross validation. The exact
kernel methods face Out-of-Memory error on larger datasets. The proposed NTK
approximations are significantly faster than the exact NTK, e.g., NTKRF shows
up to 30$\times$ speedup under $\mathrm{CT}$ dataset. We also verify that,
except for $\mathrm{Protein}$ dataset, our methods outperform RFF.

##  6 Conclusion

In this work, we propose efficient low-rank feature maps for the NTK and CNTK
kernel matrices based on both sketching and random features. Computing NTK
have been raised severe computational problems when they apply to practical
applications. Our methods runs remarkably faster than the NTK with performance
improvement. We expect applications for our methods in scaling up other
variants of NTK, e.g., transformers [21].

##  7 Acknowledgements

Amir Zandieh was supported by the Swiss NSF grant No. P2ELP2_195140. Haim
Avron and Neta Shoham were partially supported by BSF grant 2017698 and ISF
grant 1272/17.

## References

  * Ahle et al. [2020] Thomas D Ahle, Michael Kapralov, Jakob BT Knudsen, Rasmus Pagh, Ameya Velingker, David P Woodruff, and Amir Zandieh.  Oblivious sketching of high-degree polynomial kernels.  In _Symposium on Discrete Algorithms (SODA)_ , 2020.
  * Ailon and Chazelle [2009] Nir Ailon and Bernard Chazelle.  The fast Johnson–Lindenstrauss transform and approximate nearest neighbors.  _SIAM Journal on computing_ , 2009.
  * Allen-Zhu et al. [2019] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  In _International Conference on Machine Learning (ICML)_ , 2019.
  * Arora et al. [2019a] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.  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] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  In _Neural Information Processing Systems (NeurIPS)_ , 2019b.
  * Arora et al. [2019c] Sanjeev Arora, Simon S Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu.  Harnessing the Power of Infinitely Wide Deep Nets on Small-data Tasks.  In _International Conference on Learning Representations (ICLR)_ , 2019c.
  * Avron et al. [2014] Haim Avron, Huy Nguyen, and David Woodruff.  Subspace embeddings for the polynomial kernel.  In _Neural Information Processing Systems (NeurIPS)_ , 2014.
  * Avron et al. [2017] Haim Avron, Michael Kapralov, Cameron Musco, Christopher Musco, Ameya Velingker, and Amir Zandieh.  Random Fourier features for kernel ridge regression: Approximation bounds and statistical guarantees.  In _International Conference on Machine Learning (ICML)_ , 2017.
  * Bietti and Mairal [2019] Alberto Bietti and Julien Mairal.  On the inductive bias of neural tangent kernels.  In _Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Cao and Gu [2019] Yuan Cao and Quanquan Gu.  Generalization bounds of stochastic gradient descent for wide and deep neural networks.  In _Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Chizat et al. [2019] Lenaic Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  In _Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Cho and Saul [2009] Youngmin Cho and Lawrence Saul.  Kernel methods for deep learning.  In _Neural Information Processing Systems (NeurIPS)_ , 2009.
  * Cohen et al. [2016] Michael B Cohen, Jelani Nelson, and David P Woodruff.  Optimal Approximate Matrix Product in Terms of Stable Rank.  In _International Colloquium on Automata, Languages, and Programming (ICALP)_ , 2016.
  * Daniely et al. [2016] Amit Daniely, Roy Frostig, and Yoram Singer.  Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity.  In _Neural Information Processing Systems (NeurIPS)_ , 2016.
  * [15] Sanjoy Dasgupta and Anupam Gupta.  An elementary proof of a theorem of Johnson and Lindenstrauss.  _Random Structures & Algorithms_.
  * Du et al. [2019] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  In _International Conference on Learning Representations (ICLR)_ , 2019.
  * Dua and Graff [2017] Dheeru Dua and Casey Graff.  UCI machine learning repository, 2017.  URL http://archive.ics.uci.edu/ml.
  * Garriga-Alonso et al. [2018] Adrià Garriga-Alonso, Carl Edward Rasmussen, and Laurence Aitchison.  Deep Convolutional Networks as shallow Gaussian Processes.  In _International Conference on Learning Representations (ICLR)_ , 2018.
  * Geifman et al. [2020] Amnon Geifman, Abhay Yadav, Yoni Kasten, Meirav Galun, David Jacobs, and Ronen Basri.  On the similarity between the laplace and neural tangent kernels.  In _Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Han et al. [2020] Insu Han, Haim Avron, and Jinwoo Shin.  Polynomial Tensor Sketch for Element-wise Function of Low-Rank Matrix.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Hron et al. [2020] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak.  Infinite attention: NNGP and NTK for deep attention networks.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Jacot et al. [2018] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  In _Neural Information Processing Systems (NeurIPS)_ , 2018.
  * Jacot et al. [2020] Arthur Jacot, Berfin Simsek, Francesco Spadaro, Clément Hongler, and Franck Gabriel.  Implicit regularization of random feature models.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Krizhevsky [2009] Alex Krizhevsky.  Learning multiple layers of features from tiny images.  Technical report, 2009.
  * LeCun et al. [2010] Yann LeCun, Corinna Cortes, and CJ Burges.  MNIST handwritten digit database.  2010\.
  * Lee et al. [2018] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein.  Deep Neural Networks as Gaussian Processes.  In _International Conference on Learning Representations (ICLR)_ , 2018.
  * Lee et al. [2019] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington.  Wide neural networks of any depth evolve as linear models under gradient descent.  In _Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Lee et al. [2020] Jason Lee, Ruoqi Shen, Zhang Song, Mendi Wang, and Zheng Yu.  Generalized Leverage Score Sampling for Neural Networks.  In _Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Matthews et al. [2018] Alexander G de G Matthews, Jiri Hron, Mark Rowland, Richard E Turner, and Zoubin Ghahramani.  Gaussian Process Behaviour in Wide Deep Neural Networks.  In _International Conference on Learning Representations (ICLR)_ , 2018.
  * Mei et al. [2018] Song Mei, Andrea Montanari, and Phan-Minh Nguyen.  A mean field view of the landscape of two-layer neural networks.  _Proceedings of the National Academy of Sciences_ , 2018.
  * Musco and Musco [2017] Cameron Musco and Christopher Musco.  Recursive Sampling for the Nyström Method.  In _Neural Information Processing Systems (NeurIPS)_ , 2017.
  * Nelson and Nguyên [2013] Jelani Nelson and Huy L Nguyên.  OSNAP: Faster numerical linear algebra algorithms via sparser subspace embeddings.  In _Foundations of Computer Science (FOCS)_ , 2013.
  * Neyshabur et al. [2019] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro.  In search of the real inductive bias: On the role of implicit regularization in deep learning.  In _International Conference on Learning Representations (ICLR)_ , 2019.
  * Novak et al. [2018] Roman Novak, Lechao Xiao, Yasaman Bahri, Jaehoon Lee, Greg Yang, Jiri Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-dickstein.  Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes.  In _International Conference on Learning Representations (ICLR)_ , 2018.
  * Pennington et al. [2015] Jeffrey Pennington, Felix Xinnan X Yu, and Sanjiv Kumar.  Spherical random features for polynomial kernels.  In _Neural Information Processing Systems (NeurIPS)_ , 2015.
  * Rahimi and Recht [2009] Ali Rahimi and Benjamin Recht.  Random Features for Large-Scale Kernel Machines.  In _Neural Information Processing Systems (NeurIPS)_ , 2009.
  * Rudi and Rosasco [2016] Alessandro Rudi and Lorenzo Rosasco.  Generalization properties of learning with random features.  In _Neural Information Processing Systems (NeurIPS)_ , 2016.
  * Shoham and Avron [2020] Neta Shoham and Haim Avron.  Experimental Design for Overparameterized Learning with Application to Single Shot Deep Active Learning.  In _arXiv preprint arXiv:2009.12820_ , 2020.
  * Woodruff and Zandieh [2020] David P Woodruff and Amir Zandieh.  Near Input Sparsity Time Kernel Embeddings via Adaptive Sampling.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Woodruff et al. [2014] David P Woodruff et al.  Sketching as a Tool for Numerical Linear Algebra.  _Foundations and Trends® in Theoretical Computer Science_ , 2014.
  * Yang [2019] Greg Yang.  Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation.  _arXiv preprint arXiv:1902.04760_ , 2019.
  * Zancato et al. [2020] Luca Zancato, Alessandro Achille, Avinash Ravichandran, Rahul Bhotika, and Stefano Soatto.  Predicting training time without training.  In _Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Zou et al. [2018] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu.  Stochastic gradient descent optimizes over-parameterized deep relu networks.  _arXiv preprint arXiv:1811.08888_ , 2018.

##  Appendix A ReLU-NTK Expression

Arora et al. [5] showed how to exactly compute the $L$-layered NTK with
activation $\sigma:\mathbb{R}\to\mathbb{R}$ using the following dynamic
program (DP):

  1. 1.

For every $y,z\in\mathbb{R}^{d}$, let $\Sigma^{(0)}(y,z):=\langle y,z\rangle$
and for every layer $h=1,2,\ldots L$, recursively define the covariance
$\Sigma^{(h)}:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$ as:

$\begin{split}&\Lambda^{(h)}(y,z):=\begin{pmatrix}\Sigma^{(h-1)}(y,y)&\Sigma^{(h-1)}(y,z)\\\
\Sigma^{(h-1)}(z,y)&\Sigma^{(h-1)}(z,z)\end{pmatrix},\\\
&\Sigma^{(h)}(y,z):=\frac{\mathbb{E}_{(u,v)\sim\mathcal{N}\left(0,\Lambda^{(h)}(y,z)\right)}\left[\sigma(u)\cdot\sigma(v)\right]}{\mathbb{E}_{x\sim\mathcal{N}(0,1)}\left[|\sigma(x)|^{2}\right]}.\end{split}$
(18)

  2. 2.

For $h=1,2,\ldots L$, define the derivative covariance as,

$\dot{\Sigma}^{(h)}(y,z):=\frac{\mathbb{E}_{(u,v)\sim\mathcal{N}\left(0,\Lambda^{(h)}(y,z)\right)}\left[\dot{\sigma}(u)\cdot\dot{\sigma}(v)\right]}{\mathbb{E}_{x\sim\mathcal{N}(0,1)}\left[|\sigma(x)|^{2}\right]}.$
(19)

  3. 3.

Let $\Theta_{\tt ntk}^{(0)}(y,z):=\Sigma^{(0)}(y,z)$ and for every integer
$L\geq 1$, the depth-$L$ NTK expression is defined recursively as:

$\Theta_{\tt ntk}^{(L)}(y,z):=\Theta_{\tt
ntk}^{(L-1)}(y,z)\cdot\dot{\Sigma}^{(L)}(y,z)+\Sigma^{(L)}(y,z).$ (20)

While using this DP, one can compute the kernel value $\Theta_{\tt
ntk}^{(L)}(y,z)$ for any pair of vectors $y,z\in\mathbb{R}^{d}$ using
$\mathcal{O}(d+L)$ operations, it is hard to gain insight into the structure
of this kernel using the expression above. In particular, the NTK expression
involves recursive computation of nontrivial expectations which makes it is
unclear whether there exist efficient sketching solutions for this kernel in
its current form. However, we show that for the important case of ReLU
activation, this kernel takes an extremely nice and highly structured form. In
fact, the NTK can be fully characterized by a _univariate function_ $K_{\tt
relu}^{(L)}:[-1,1]\to\mathbb{R}$, and exploiting this special structure is the
key to designing efficient sketching methods for this kernel.

Now we proceed to prove ??, that is,

$\displaystyle\Theta_{\tt ntk}^{(L)}(y,z)\equiv\|y\|_{2}\|z\|_{2}\cdot K_{\tt
relu}^{(L)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right),\text{~{}~{}~{}for any
}y,z\in\mathbb{R}^{d}.$ (5)

First note that the main component of the DP given in ??, ??, and ?? is the
_Activation Covariances_ :

$\mathbb{E}_{w\sim\mathcal{N}(0,{\bm{I}}_{d})}\left[\sigma(w^{\top}y)\cdot\sigma(w^{\top}z)\right],~{}\text{
and
}\mathbb{E}_{w\sim\mathcal{N}(0,{\bm{I}}_{d})}\left[\dot{\sigma}(w^{\top}y)\cdot\dot{\sigma}(w^{\top}z)\right]\text{~{}~{}
for every }y,z\in\mathbb{R}^{d}.$

It is worth noting that the above activation covariance functions are positive
definite and hence they both define valid kernel functions in
$\mathbb{R}^{d}\times\mathbb{R}^{d}$. The connection between the ReLU
activation covariance functions and arc-cosine kernel functions defined in ??
of ?? is proved in Cho and Saul [12]. We restate this result as follows,

$\begin{split}\mathbb{E}_{w\sim\mathcal{N}(0,{\bm{I}}_{d})}\left[\mathrm{ReLU}(w^{\top}y)\cdot\mathrm{ReLU}(w^{\top}z)\right]&=\frac{\|y\|_{2}\|z\|_{2}}{2}\cdot\kappa_{1}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\\\
\mathbb{E}_{w\sim\mathcal{N}(0,{\bm{I}}_{d})}\left[\mathrm{Step}(w^{\top}y)\cdot\mathrm{Step}(w^{\top}z)\right]&=\frac{1}{2}\cdot\kappa_{0}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right).\end{split}$ (21)

Proof of ??: Consider the NTK expression given in ??, ??, and ??. We first
prove by induction on $h=0,1,2,\ldots L$ that the covariance function
$\Sigma^{(h)}(y,z)$ defined in ?? satisfies:

$\Sigma^{(h)}(y,z)=\|y\|_{2}\|z\|_{2}\cdot\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right).$

The base of induction is trivial for $h=0$ due to $\Sigma^{(0)}(y,z)=\langle
y,z\rangle=\|y\|_{2}\|z\|_{2}\cdot\Sigma_{\tt relu}^{(0)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$.

To prove the inductive step, suppose that the inductive hypothesis holds for
$h-1$, i.e.,

$\Sigma^{(h-1)}(y,z)=\|y\|_{2}\|z\|_{2}\cdot\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$

The $2\times 2$ covariance matrix $\Lambda^{(h)}(y,z)$, defined in ??, can be
decomposed as $\Lambda^{(h)}(y,z)=\begin{pmatrix}f^{\top}\\\
g^{\top}\end{pmatrix}\cdot\begin{pmatrix}f&g\end{pmatrix}$, where
$f,g\in\mathbb{R}^{2}$. Now note that $\|f\|_{2}^{2}=\Sigma^{(h-1)}(y,y)$,
hence, by inductive hypothesis, we have,
$\|f\|_{2}^{2}=\|y\|_{2}^{2}\cdot\Sigma_{\tt relu}^{(h-1)}\left(\frac{\langle
y,y\rangle}{\|y\|_{2}^{2}}\right)=\|y\|_{2}^{2}\cdot\Sigma_{\tt
relu}^{(h-1)}(1)=\|y\|_{2}^{2}.$

Using a similar argument we have $\|g\|_{2}^{2}=\|z\|_{2}^{2}$. Therefore, by
??, we can write

$\displaystyle\Sigma^{(h)}(y,z)$
$\displaystyle=\frac{1}{\mathbb{E}_{x\sim\mathcal{N}(0,1)}\left[|\sigma(x)|^{2}\right]}\cdot\mathbb{E}_{w\sim\mathcal{N}\left(0,I_{2}\right)}\left[\sigma(w^{\top}f)\cdot\sigma(w^{\top}g)\right]$
$\displaystyle=\frac{2}{\kappa_{1}(1)}\cdot\mathbb{E}_{w\sim\mathcal{N}\left(0,I_{2}\right)}\left[\sigma(w^{\top}f)\cdot\sigma(w^{\top}g)\right]$
$\displaystyle=\|f\|_{2}\|g\|_{2}\cdot\kappa_{1}\left(\frac{\langle
f,g\rangle}{\|f\|_{2}\|g\|_{2}}\right)=\|y\|_{2}\|z\|_{2}\cdot\kappa_{1}\left(\frac{\langle
f,g\rangle}{\|y\|_{2}\|z\|_{2}}\right).$

Since we assumed that $\Lambda^{(h)}(y,z)=\begin{pmatrix}f^{\top}\\\
g^{\top}\end{pmatrix}\cdot\begin{pmatrix}f&g\end{pmatrix}$, we have $\langle
f,g\rangle=\Sigma^{(h-1)}(y,z)$. By inductive hypothesis along with ??, we
find that

$\Sigma^{(h)}(y,z)=\|y\|_{2}\|z\|_{2}\cdot\kappa_{1}\left(\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)=\|y\|_{2}\|z\|_{2}\cdot\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right),$

which completes the induction and proves that for every $h=0,1,\dots,L$,

$\Sigma^{(h)}(y,z)=\|y\|_{2}\|z\|_{2}\cdot\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right).$ (22)

For obtaining the final NTK expression we also need to simplify the derivative
covariance function defined in ??. Recall that we proved before, the
covariance matrix $\Lambda^{(h)}(y,z)$ can be decomposed as
$\Lambda^{(h)}(y,z)=\begin{pmatrix}f^{\top}\\\
g^{\top}\end{pmatrix}\cdot\begin{pmatrix}f&g\end{pmatrix}$, where
$f,g\in\mathbb{R}^{2}$ with $\|f\|_{2}=\|y\|_{2}$ and $\|g\|_{2}=\|z\|_{2}$.
Therefore, by ??, we can write

$\displaystyle\dot{\Sigma}^{(h)}(y,z)$
$\displaystyle=\frac{1}{\mathbb{E}_{x\sim\mathcal{N}(0,1)}\left[|\sigma(x)|^{2}\right]}\cdot\mathbb{E}_{w\sim\mathcal{N}\left(0,I_{2}\right)}\left[\dot{\sigma}(w^{\top}f)\cdot\dot{\sigma}(w^{\top}g)\right]$
$\displaystyle=\frac{2}{\kappa_{1}(1)}\cdot\mathbb{E}_{w\sim\mathcal{N}\left(0,I_{2}\right)}\left[\dot{\sigma}(w^{\top}f)\cdot\dot{\sigma}(w^{\top}g)\right]$
$\displaystyle=\kappa_{0}\left(\frac{\langle
f,g\rangle}{\|y\|_{2}\|z\|_{2}}\right).$

Since we assumed that $\Lambda^{(h)}(y,z)=\begin{pmatrix}f^{\top}\\\
g^{\top}\end{pmatrix}\cdot\begin{pmatrix}f&g\end{pmatrix}$, $\langle
f,g\rangle=\Sigma^{(h-1)}(y,z)=\|y\|_{2}\|z\|_{2}\cdot\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$.
Therefore, by ??, for every $h\in\\{1,2,\dots L\\}$,

$\dot{\Sigma}^{(h)}(y,z)=\kappa_{0}\left(\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)=\dot{\Sigma}_{\tt
relu}^{(h)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right).$ (23)

Now we prove by induction on integer $L$ that the NTK with $L$ layers and ReLU
activation given in ?? satisfies

$\Theta_{\tt ntk}^{(L)}(y,z)=\|y\|_{2}\|z\|_{2}\cdot K_{\tt
relu}^{(L)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right).$

The base of induction, trivially holds because, by ??:

$\Theta_{\tt ntk}^{(0)}(y,z)=\Sigma^{(0)}(y,z)=\|y\|_{2}\|z\|_{2}\cdot
K^{(0)}_{\tt relu}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right).$

To prove the inductive step, suppose that the inducive hypothesis holds for
$L-1$, that is $\Theta_{\tt ntk}^{(L-1)}(y,z)=\|y\|_{2}\|z\|_{2}\cdot K_{\tt
relu}^{(L-1)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$. Now
using the recursive definition of $\Theta_{\tt ntk}^{(L)}(y,z)$ given in ??
along with ?? and ?? we can write,

$\displaystyle\Theta_{\tt ntk}^{(L)}(y,z)$ $\displaystyle=\Theta_{\tt
ntk}^{(L-1)}(y,z)\cdot\dot{\Sigma}^{(L)}(y,z)+\Sigma^{(L)}(y,z)$
$\displaystyle=\|y\|_{2}\|z\|_{2}\cdot K_{\tt relu}^{(L-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\cdot\dot{\Sigma}_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)+\|y\|_{2}\|z\|_{2}\cdot\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$
$\displaystyle\equiv\|y\|_{2}\|z\|_{2}\cdot K_{\tt
relu}^{(L)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right).$

This completes the proof of ??. ∎

##  Appendix B Sketching Preliminaries: PolySketch and SRHT

Our sketching algorithms use the Subsampled Randomized Hadamard Transform
(SRHT) [2] to reduce the dimensionality of the intermediate vectors that arise
in our computations. Next lemma gives the performance of SRHT sketches which
is proved, for instance, in Theorem 9 of [13],

######  Lemma 2 (SRHT Sketch).

For every positive integer $d$ and every $\varepsilon,\delta>0$, there exists
a distribution on random matrices ${\bm{S}}\in\mathbb{R}^{m\times d}$ with
$m=\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\cdot\log^{2}\frac{1}{\varepsilon\delta}\right)$,
called SRHT sketch, such that for any vector $x\in\mathbb{R}^{d}$,
$\Pr\left[\|{\bm{S}}x\|_{2}^{2}\in(1\pm\varepsilon)\|x\|_{2}^{2}\right]\geq
1-\delta$. Moreover, ${\bm{S}}x$ can be computed in time
$\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\cdot\log^{2}\frac{1}{\varepsilon\delta}+d\log
d\right)$.

Now we restate the ?? and present the proof,

See 1

Proof of ??: The fourth statement of the lemma immediately follows from
Theorem 1.3 of Ahle et al. [1]. Moreover, by invoking Theorem 1.2 of Ahle et
al. [1], we find that there exists a random sketch
${\bm{Q}}^{p}\in\mathbb{R}^{m\times d^{p}}$ such that
$m=C\cdot\frac{p}{\varepsilon^{2}}\log^{3}\frac{1}{\varepsilon\delta}$, for
some absolute constant $C$, and for any $y\in\mathbb{R}^{d^{p}}$,

$\displaystyle\Pr\left[\|{\bm{Q}}^{p}y\|_{2}^{2}\in(1\pm\varepsilon)\|y\|_{2}^{2}\right]\geq
1-\delta.$

This immediately proves the first statement of the lemma.

${\bf S_{\text{base}}}$${\bf S_{\text{base}}}$${\bf T_{\text{base}}}$${\bf
T_{\text{base}}}$${\bf S_{\text{base}}}$${\bf T_{\text{base}}}$${\bf
T_{\text{base}}}$internal nodes: TensorSRHTleaves: OSNAP

Figure 3: The structure of sketch ${\bm{Q}}^{p}$ proposed in Theorem 1.2 of
[1]: the sketch matrices in nodes of the tree labeled with
${\bm{S}}_{\text{base}}$ and ${\bm{T}}_{\text{base}}$ are independent
instances of TensorSRHT and OSNAP, respectively.

As shown in [1], the sketch ${\bm{Q}}^{p}$ can be applied to tensor product
vectors of the form $v_{1}\otimes v_{2}\otimes\ldots v_{p}$ by recursive
application of $\mathcal{O}(p)$ independent instances of OSNAP transform [32]
and a novel variant of the SRHT , proposed in [1] called TensorSRHT, on
vectors $v_{i}$ and their sketched versions. The sketch ${\bm{Q}}^{p}$, as
shown in ??, can be represented by a binary tree with $p$ leaves where the
leaves are OSNAP sketches and the internal nodes are the TensorSRHT. The use
of OSNAP in the leaves of this sketch structure ensures excellent runtime for
sketching sparse input vectors. However, note that if the input vectors are
not sparse, i.e., ${\rm nnz}(v_{i})=\widetilde{\Omega}(d)$ for input vectors
$v_{i}$, then we can simply remove the OSNAP transforms from the leaves of
this structure and achieve improved runtime, without hurting the approximation
guarantee. Therefore, the sketch ${\bm{Q}}^{p}$ that satisfies the statement
of the lemma is exactly the one introduced in [1] for sparse input vectors and
for non-sparse inputs is obtained by removing the OSNAP transforms from the
leaves of the sketch structure given in ??.

Runtime analysis: By Theorem 1.2 of [1], for any vector $x\in\mathbb{R}^{d}$,
${\bm{Q}}^{p}x^{\otimes p}$ can be computed in time $\mathcal{O}\left(pm\log
m+\frac{p^{3/2}}{\varepsilon}\log\frac{1}{\delta}\cdot{\rm nnz}(x)\right)$.
From the binary tree structure of the sketch, shown in ??, it follows that
once we compute ${\bm{Q}}^{p}x^{\otimes p}$, then
${\bm{Q}}^{p}\left(x^{\otimes p-1}\otimes{e}_{1}\right)$ can be computed by
updating the path from one of the leaves to the root of the binary tree which
amounts to applying an instance of OSNAP transform on the ${e}_{1}$ vector and
then applying $\mathcal{O}(\log p)$ instances of TensorSRHT on the
intermediate nodes of the tree. This can be computed in a total additional
runtime of $\mathcal{O}(m\log m\log p)$. By this argument, it follows that
${\bm{Q}}^{p}\left(x^{\otimes p-j}\otimes{e}_{1}^{j}\right)$ can be computed
sequentially for all $j=0,1,2,\cdots p$ in total time $\mathcal{O}\left(pm\log
p\log m+\frac{p^{3/2}}{\varepsilon}\log\frac{1}{\delta}\cdot{\rm
nnz}(x)\right)$. By plugging in the value
$m=\mathcal{O}\left(\frac{p}{\varepsilon^{2}}\log^{3}\frac{1}{\varepsilon\delta}\right)$,
this runtime will be
$\mathcal{O}\left(\frac{p^{2}\log^{2}\frac{p}{\varepsilon}}{\varepsilon^{2}}\log^{3}\frac{1}{\varepsilon\delta}+\frac{p^{3/2}}{\varepsilon}\log\frac{1}{\delta}\cdot{\rm
nnz}(x)\right)$, which gives the second statement of the lemma for sparse
input vectors $x$. If $x$ is non-sparse, as we discussed in the above
paragraph, we just need to omit the OSNAP transforms from our sketch
construction which translates into a runtime of
$\mathcal{O}\left(\frac{p^{2}\log^{2}\frac{p}{\varepsilon}}{\varepsilon^{2}}\log^{3}\frac{1}{\varepsilon\delta}+pd\log
d\right)$. Therefore, the final runtime bound is
$\mathcal{O}\left(\frac{p^{2}\log^{2}\frac{p}{\varepsilon}}{\varepsilon^{2}}\log^{3}\frac{1}{\varepsilon\delta}+\min\left\\{\frac{p^{3/2}}{\varepsilon}\log\frac{1}{\delta}\cdot{\rm
nnz}(x),pd\log d\right\\}\right)$, which proves the second statement of the
lemma.

Furthermore, the sketch ${\bm{Q}}^{p}$ can be applied to tensor product of any
collection of $p$ vectors. The time to apply ${\bm{Q}}^{p}$ to the tensor
product $v_{1}\otimes v_{2}\otimes\ldots v_{p}$ consists of time of applying
OSNAP to each of the vectors $v_{1},v_{2}\ldots v_{p}$ and time of applying
$\mathcal{O}(p)$ instances of TensorSRHT to intermediate vectors which are of
size $m$. This runtime can be upper bounded by
$\mathcal{O}\left(\frac{p^{2}\log\frac{p}{\varepsilon}}{\varepsilon^{2}}\log^{3}\frac{1}{\varepsilon\delta}+\frac{p^{3/2}}{\varepsilon}d\cdot\log\frac{1}{\delta}\right)$,
which proves the third statement of the ??. ∎

##  Appendix C NTK Sketch: Claims and Invariants

We start by proving that the polynomials $P_{\tt relu}^{(p)}(\cdot)$ and
$\dot{P}_{\tt relu}^{(p^{\prime})}(\cdot)$ defined in ?? of ?? closely
approximate the arc-cosine functions $\kappa_{1}(\cdot)$ and
$\kappa_{0}(\cdot)$ on the interval $[-1,1]$.

Remark. Observe that
$\kappa_{0}(\alpha)=\frac{d}{d\alpha}\left(\kappa_{1}(\alpha)\right)$.

######  Lemma 3 (Polynomial Approximations to $\kappa_{1}$ and $\kappa_{0}$).

If we let $\kappa_{1}(\cdot)$ and $\kappa_{0}(\cdot)$ be defined as in ?? of
??, then for any integer $p\geq\frac{1}{9\varepsilon^{2/3}}$, the polynomial
$P_{\tt relu}^{(p)}(\cdot)$ defined in ?? of ?? satisfies,

$\max_{\alpha\in[-1,1]}\left|P_{\tt
relu}^{(p)}(\alpha)-\kappa_{1}(\alpha)\right|\leq\varepsilon.$

Moreover, for any integer $p^{\prime}\geq\frac{1}{26\varepsilon^{2}}$, the
polynomial $\dot{P}_{\tt relu}^{(p^{\prime})}(\cdot)$ defined as in ?? of ??,
satisfies,

$\max_{\alpha\in[-1,1]}\left|\dot{P}_{\tt
relu}^{(p^{\prime})}(\alpha)-\kappa_{0}(\alpha)\right|\leq\varepsilon.$

Proof of ??: We start by Taylor series expansion of $\kappa_{0}(\cdot)$ around
$\alpha=0$,
$\kappa_{0}(\alpha)=\frac{1}{2}+\frac{1}{\pi}\cdot\sum_{i=0}^{\infty}\frac{(2i)!}{2^{2i}\cdot(i!)^{2}\cdot(2i+1)}\cdot\alpha^{2i+1}$.
Therefore, we have

$\displaystyle\max_{\alpha\in[-1,1]}\left|\dot{P}_{\tt
relu}^{(p^{\prime})}(\alpha)-\kappa_{0}(\alpha)\right|$
$\displaystyle=\frac{1}{\pi}\cdot\sum_{i=p^{\prime}+1}^{\infty}\frac{(2i)!}{2^{2i}\cdot(i!)^{2}\cdot(2i+1)}$
$\displaystyle\leq\frac{1}{\pi}\cdot\sum_{i=p^{\prime}+1}^{\infty}\frac{e\cdot
e^{-2i}\cdot(2i)^{2i+1/2}}{2\pi\cdot 2^{2i}\cdot e^{-2i}\cdot
n^{2i+1}\cdot(2i+1)}$
$\displaystyle=\frac{e}{\sqrt{2}\pi^{2}}\cdot\sum_{i=p^{\prime}+1}^{\infty}\frac{1}{\sqrt{i}\cdot(2i+1)}$
$\displaystyle\leq\frac{e}{\sqrt{2}\pi^{2}}\cdot\int_{p^{\prime}}^{\infty}\frac{1}{\sqrt{x}\cdot(2x+1)}dx$
$\displaystyle\leq\frac{e}{\sqrt{2}\pi^{2}}\cdot\frac{1}{\sqrt{p^{\prime}}}\leq\varepsilon.$

To prove the second part of the lemma, we consider the Taylor expansion of
$\kappa_{1}(\cdot)$ at $\alpha=0$. Since
$\kappa_{0}(\alpha)=\frac{d}{d\alpha}\left(\kappa_{1}(\alpha)\right)$, the
Taylor series of $\kappa_{1}(\alpha)$ can be obtained from the Taylor series
of $\kappa_{0}(\alpha)$ as follows,

$\kappa_{1}(\alpha)=\frac{1}{\pi}+\frac{\alpha}{2}+\frac{1}{\pi}\cdot\sum_{i=0}^{\infty}\frac{(2i)!}{2^{2i}\cdot(i!)^{2}\cdot(2i+1)\cdot(2i+2)}\cdot\alpha^{2i+2}.$

Hence, we have

$\displaystyle\max_{\alpha\in[-1,1]}\left|P_{\tt
relu}^{(p)}(\alpha)-\kappa_{1}(\alpha)\right|$
$\displaystyle=\frac{1}{\pi}\cdot\sum_{i=p+1}^{\infty}\frac{(2i)!}{2^{2i}\cdot(i!)^{2}\cdot(2i+1)\cdot(2i+2)}$
$\displaystyle\leq\frac{1}{\pi}\cdot\sum_{i=p+1}^{\infty}\frac{e\cdot
e^{-2i}\cdot(2i)^{2i+1/2}}{2\pi\cdot 2^{2i}\cdot e^{-2i}\cdot
n^{2i+1}\cdot(2i+1)\cdot(2i+2)}$
$\displaystyle=\frac{e}{\sqrt{2}\pi^{2}}\cdot\sum_{i=p+1}^{\infty}\frac{1}{\sqrt{i}\cdot(2i+1)\cdot(2i+2)}$
$\displaystyle\leq\frac{e}{\sqrt{2}\pi^{2}}\cdot\int_{p}^{\infty}\frac{1}{\sqrt{x}\cdot(2x+1)\cdot(2x+2)}dx$
$\displaystyle\leq\frac{e}{\sqrt{2}\pi^{2}}\cdot\frac{1}{6\cdot
p^{3/2}}\leq\varepsilon.$

This completes the proof of ??. ∎.

Therefore, it is possible to approximate the function $\kappa_{0}(\cdot)$ up
to error $\varepsilon$ using a polynomial of degree
$\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\right)$. Also if we want to
approximate $\kappa_{1}(\cdot)$ using a polynomial up to error $\varepsilon$
on the interval $[-1,1]$, it suffices to use a polynomial of degree
$\mathcal{O}\left(\frac{1}{\varepsilon^{2/3}}\right)$. One can see that since
the Taylor expansions of $\kappa_{1}$ and $\kappa_{0}$ contain non-negative
coefficients only, both of these functions are positive definite.
Additionally, the polynomial approximations $P_{\tt relu}^{(p)}$ and
$\dot{P}_{\tt relu}^{(p^{\prime})}$ given in ?? of ?? are positive definite
functions.

In order to prove ??, we also need the following lemma on the error
sensitivity of polynomials $P_{\tt relu}^{(p)}$ and $\dot{P}_{\tt
relu}^{(p^{\prime})}$,

######  Lemma 4 (Sensitivity of $P_{\tt relu}^{(p)}$ and $\dot{P}_{\tt
relu}^{(p)}$).

For any integer $p\geq 3$, any $\alpha\in[-1,1]$, and any $\alpha^{\prime}$
such that $|\alpha-\alpha^{\prime}|\leq\frac{1}{6p}$, if we let the
polynomials $P_{\tt relu}^{(p)}(\alpha)$ and $\dot{P}_{\tt
relu}^{(p)}(\alpha)$ be defined as in ?? of ??, then

$\left|P_{\tt relu}^{(p)}(\alpha)-P_{\tt
relu}^{(p)}(\alpha^{\prime})\right|\leq|\alpha-\alpha^{\prime}|,$

and

$\left|\dot{P}_{\tt relu}^{(p)}(\alpha)-\dot{P}_{\tt
relu}^{(p)}(\alpha^{\prime})\right|\leq\sqrt{p}\cdot|\alpha-\alpha^{\prime}|.$

Proof of ??: Note that an $\alpha^{\prime}$ that satisfies the preconditions
of the lemm, is in the range $\left[-1-\frac{1}{6p},1+\frac{1}{6p}\right]$.
Now we bound the derivative of the polynomial $\dot{P}_{\tt relu}^{(p)}$ on
the interval $\left[-1-\frac{1}{6p},1+\frac{1}{6p}\right]$,

$\displaystyle\max_{\alpha\in\left[-1-\frac{1}{6p},1+\frac{1}{6p}\right]}\left|\frac{d}{d\alpha}\left(\dot{P}_{\tt
relu}^{(p)}(\alpha)\right)\right|$
$\displaystyle=\frac{1}{\pi}\cdot\sum_{i=0}^{p}\frac{(2i)!}{2^{2i}\cdot(i!)^{2}}\cdot\left(1+\frac{1}{6p}\right)^{2i}$
$\displaystyle\leq\frac{1}{\pi}+\frac{e^{4/3}}{\sqrt{2}\pi^{2}}\cdot\sum_{i=1}^{p}\frac{1}{\sqrt{i}}$
$\displaystyle\leq\frac{1}{\pi}+\frac{e^{4/3}}{\sqrt{2}\pi^{2}}\cdot\int_{0}^{p}\frac{1}{\sqrt{x}}dx$
$\displaystyle\leq\sqrt{p},$

therefore, the second statement of lemma holds.

To prove the first statement of lemma, we bound the derivative of the
polynomial $P_{\tt relu}^{(p)}$ on the interval
$\left[-1-\frac{1}{6p},1+\frac{1}{6p}\right]$ as follows,

$\displaystyle\max_{\alpha\in\left[-1-\frac{1}{6p},1+\frac{1}{6p}\right]}\left|\frac{d}{d\alpha}\left(P_{\tt
relu}^{(p)}(\alpha)\right)\right|$
$\displaystyle=\frac{1}{\pi}\cdot\sum_{i=0}^{p}\frac{(2i)!}{2^{2i}\cdot(i!)^{2}\cdot(2i+1)}\cdot\left(1+\frac{1}{6p}\right)^{2i+1}$
$\displaystyle\leq\frac{19}{18\pi}+\frac{e^{25/18}}{\sqrt{2}\pi^{2}}\cdot\sum_{i=1}^{p}\frac{1}{\sqrt{i}\cdot(2i+1)}$
$\displaystyle\leq\frac{19}{18\pi}+\frac{e^{25/18}}{\sqrt{2}\pi^{2}}\cdot\int_{0}^{p}\frac{1}{\sqrt{x}\cdot(2x+1)}dx$
$\displaystyle\leq 1,$

therefore, the second statement of the lemma follows. This completes the proof
of ??. ∎

For the rest of this section, we need two basic properties of tensor products
and direct sums:

$\displaystyle\langle x\otimes y,z\otimes w\rangle=\langle
x,z\rangle\cdot\langle y,w\rangle,\quad\langle x\oplus y,z\oplus
w\rangle=\langle x,z\rangle+\langle y,w\rangle$ (24)

for vectors $x,y,z,w$ with conforming sizes.

Now we are in a position to analyze the invariants that are maintained
throughout the execution of NTKSketch (??):

######  Lemma 5 (Invariants of the NTKSketch algorithm).

For every positive integers $d$ and $L$, every $\varepsilon,\delta>0$, every
vectors $y,z\in\mathbb{R}^{d}$, if we let $\Sigma_{\tt
relu}^{(h)}:[-1,1]\to\mathbb{R}$ and $K_{\tt relu}^{(h)}:[-1,1]\to\mathbb{R}$
be the functions defined in ?? and ?? of ??, then with probability at least
$1-\delta$ the following invariants hold for every $h=0,1,2,\ldots L$:

  1. 1.

The mapping $\phi^{(h)}(\cdot)$ computed by NTKSketch in line 4 and ?? of ??
satisfy

$\left|\left<\phi^{(h)}(y),\phi^{(h)}(z)\right>-\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq({h+1})\cdot\frac{\varepsilon^{2}}{60L^{3}}.$

  2. 2.

The mapping $\psi^{(h)}(\cdot)$ computed by NTKSketch in line 5 and ?? of ??
satisfy

$\left|\left<\psi^{(h)}(y),\psi^{(h)}(z)\right>-K_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq\varepsilon\cdot\frac{h^{2}+1}{10L}.$

Proof of ??: The proof is by induction on the value of $h=0,1,2,\ldots L$.
More formally, consider the following statements for every $h=0,1,2,\ldots L$:

  1. ${\bf P_{1}(h):}$

$\begin{split}&\left|\left<\phi^{(h)}(y),\phi^{(h)}(z)\right>-\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq({h+1})\cdot\frac{\varepsilon^{2}}{60L^{3}},\\\
&\left|\left\|\phi^{(h)}(y)\right\|_{2}^{2}-1\right|\leq({h+1})\cdot\frac{\varepsilon^{2}}{60L^{3}},\text{
and
}\left|\left\|\phi^{(h)}(z)\right\|_{2}^{2}-1\right|\leq({h+1})\cdot\frac{\varepsilon^{2}}{60L^{3}}.\end{split}$

  2. ${\bf P_{2}(h):}$

$\begin{split}&\left|\left<\psi^{(h)}(y),\psi^{(h)}(z)\right>-K_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq\varepsilon\cdot\frac{h^{2}+1}{10L},\\\
&\left|\left\|\psi^{(h)}(y)\right\|_{2}^{2}-K_{\tt
relu}^{(h)}(1)\right|\leq\varepsilon\cdot\frac{h^{2}+1}{10L},\text{ and
}\left|\left\|\psi^{(h)}(z)\right\|_{2}^{2}-K_{\tt
relu}^{(h)}(1)\right|\leq\varepsilon\cdot\frac{h^{2}+1}{10L}.\end{split}$

We prove that the following holds,

$\displaystyle\Pr[P_{1}(0)]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right),~{}~{}\text{ and
}~{}~{}\Pr[P_{2}(0)|P_{1}(0)]\geq 1-\mathcal{O}\left(\frac{\delta}{L}\right).$
(25)

Additionally, by induction, we prove that for every $h=1,2,\ldots L$,

$\displaystyle\Pr\left[P_{1}(h)|P_{1}(h-1)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right),~{}~{}\text{ and }$
$\displaystyle\Pr\left[P_{2}(h)|P_{2}(h-1),P_{1}(h),P_{1}(h-1)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (26)

(1) Base of induction ($h=0$): By line 4 of ??,
$\phi^{(0)}(y)=\frac{1}{\|y\|_{2}}\cdot{\bm{S}}\cdot{\bm{Q}}^{1}\cdot y$ and
$\phi^{(0)}(z)=\frac{1}{\|z\|_{2}}\cdot{\bm{S}}\cdot{\bm{Q}}^{1}\cdot z$,
thus, ?? implies the following

$\displaystyle\Pr\left[\left|\left<\phi^{(0)}(y),\phi^{(0)}(z)\right>-\frac{\left<{\bm{Q}}^{1}y,{\bm{Q}}^{1}z\right>}{\|y\|_{2}\|z\|_{2}}\right|\leq\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\cdot\frac{\left\|{\bm{Q}}^{1}y\right\|_{2}\|{\bm{Q}}^{1}z\|_{2}}{\|y\|_{2}\|z\|_{2}}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (27)

By using the above together with ?? and union bound as well as triangle
inequality, we have

$\displaystyle\Pr\left[\left|\left<\phi^{(0)}(y),\phi^{(0)}(z)\right>-\frac{\left<y,z\right>}{\|y\|_{2}\|z\|_{2}}\right|\leq\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (28)

Similarly, we can prove that

$\displaystyle\Pr\left[\left|\left\|\phi^{(0)}(y)\right\|_{2}^{2}-1\right|\leq\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right),\text{ and }$ (29)
$\displaystyle\Pr\left[\left|\left\|\phi^{(0)}(z)\right\|_{2}^{2}-1\right|\leq\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (30)

Using union bound, we prove the base of induction for statement $P_{1}(0)$,
i.e.,

$\displaystyle\Pr[P_{1}(0)]\geq 1-\mathcal{O}\left(\frac{\delta}{L}\right).$
(31)

Moreover, by line 5 of ??, $\psi^{(0)}(y)={\bm{V}}\cdot\phi^{(0)}(y)$ and
$\psi^{(0)}(z)={\bm{V}}\cdot\phi^{(0)}(z)$, thus, ?? implies that,

$\displaystyle\Pr\left[\left|\left<\psi^{(0)}(y),\psi^{(0)}(z)\right>-\left<\phi^{(0)}(y),\phi^{(0)}(z)\right>\right|\leq\mathcal{O}\left(\frac{\varepsilon}{L}\right)\cdot\left\|\phi^{(0)}(y)\right\|_{2}\left\|\phi^{(0)}(z)\right\|_{2}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (32)

By conditioning on $P_{1}(0)$ and using the above together with triangle
inequality it follows that,

$\displaystyle\Pr\left[\left|\left<\psi^{(0)}(y),\psi^{(0)}(z)\right>-K_{\tt
relu}^{(0)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq\frac{\varepsilon}{10L}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (33)

Similarly we can prove that with probability at least
$1-\mathcal{O}\left(\frac{\delta}{L}\right)$ we have
$\left|\left\|\psi^{(0)}(y)\right\|_{2}^{2}-K_{\tt
relu}^{(0)}(1)\right|\leq\frac{\varepsilon}{10L}$ and
$\left|\left\|\psi^{(0)}(z)\right\|_{2}^{2}-K_{\tt
relu}^{(0)}(1)\right|\leq\frac{\varepsilon}{10L}$, which proves the base of
induction for the second statement, i.e., $\Pr[P_{2}(0)|P_{1}(0)]\geq
1-\mathcal{O}\left({\delta}/{L}\right)$. This completes the base of induction.

(2) Inductive step: Assume that the inductive hypothesis holds for $h-1$.
First, note that by ?? and using ?? we have the following,

$\Pr\left[\left|\left<\phi^{(h)}(y),\phi^{(h)}(z)\right>-\sum_{j=0}^{2p+2}c_{j}\left<Z^{(h)}_{j}(y),Z^{(h)}_{j}(z)\right>\right|\leq\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\cdot
A\right]\geq 1-\mathcal{O}\left(\frac{\delta}{L}\right),$ (34)

where
$A:=\sqrt{\sum_{j=0}^{2p+2}c_{j}\|Z^{(h)}_{j}(y)\|_{2}^{2}}\cdot\sqrt{\sum_{j=0}^{2p+2}c_{j}\|Z^{(h)}_{j}(z)\|_{2}^{2}}$
and the collection of vectors $\left\\{Z^{(h)}_{j}(y)\right\\}_{j=0}^{2p+2}$
and $\left\\{Z^{(h)}_{j}(z)\right\\}_{j=0}^{2p+2}$ and coefficients
$\\{c_{j}\\}_{j=0}^{2p+2}$ are defined as per ?? and ??, respectively.

By ?? together with union bound, the following inequalities simultaneously
hold for all $j=0,\dots,2p+2$, with probability at least
$1-\mathcal{O}\left(\delta/L\right)$:

$\displaystyle\left|\left<Z^{(h)}_{j}(y),Z^{(h)}_{j}(z)\right>-\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>^{j}\right|\leq\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\cdot\left\|\phi^{(h-1)}(y)\right\|_{2}^{j}\left\|\phi^{(h-1)}(z)\right\|_{2}^{j}$
$\displaystyle\left\|Z^{(h)}_{j}(y)\right\|_{2}^{2}\leq\frac{11}{10}\cdot\left\|\phi^{(h-1)}(y)\right\|_{2}^{2j}$
(35)
$\displaystyle\left\|Z^{(h)}_{j}(z)\right\|_{2}^{2}\leq\frac{11}{10}\cdot\left\|\phi^{(h-1)}(z)\right\|_{2}^{2j}$

Therefore, by plugging ?? to ?? and using union bound, triangle inequality and
Cauchy–Schwarz inequality we find that,

$\Pr\left[\left|\left<\phi^{(h)}(y),\phi^{(h)}(z)\right>-P^{(p)}_{\tt
relu}\left(\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>\right)\right|\leq\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\cdot
B\right]\geq 1-\mathcal{O}\left(\frac{\delta}{L}\right),$ (36)

where $B:=\sqrt{P^{(p)}_{\tt
relu}\left(\|\phi^{(h-1)}(y)\|_{2}^{2}\right)\cdot P^{(p)}_{\tt
relu}\left(\|\phi^{(h-1)}(z)\|_{2}^{2}\right)}$ and $P^{(p)}_{\tt
relu}(\alpha)=\sum_{j=0}^{2p+2}c_{j}\cdot\alpha^{j}$ is the polynomial defined
in ??. Using the inductive hypothesis $P_{1}(h-1)$, we have that

$\displaystyle\left|\left\|\phi^{(h-1)}(y)\right\|_{2}^{2}-1\right|\leq
h\cdot\frac{\varepsilon^{2}}{60L^{3}},~{}~{}\text{ and
}~{}~{}\left|\left\|\phi^{(h-1)}(z)\right\|_{2}^{2}-1\right|\leq
h\cdot\frac{\varepsilon^{2}}{60L^{3}}.$ (37)

Therefore, by ?? we have $\left|P_{\tt
relu}^{(p)}\left(\|\phi^{(h-1)}(y)\|_{2}^{2}\right)-P_{\tt
relu}^{(p)}(1)\right|\leq h\cdot\frac{\varepsilon^{2}}{60L^{3}}$ and
$\left|P_{\tt relu}^{(p)}\left(\|\phi^{(h-1)}(z)\|_{2}^{2}\right)-P_{\tt
relu}^{(p)}(1)\right|\leq h\cdot\frac{\varepsilon^{2}}{60L^{3}}$.
Consequently, since $P_{\tt relu}^{(p)}(1)\leq P_{\tt relu}^{(+\infty)}(1)=1$,
we obtain that $B\leq\frac{11}{10}.$ By plugging this into ?? we have,

$\Pr\left[\left|\left<\phi^{(h)}(y),\phi^{(h)}(z)\right>-P^{(p)}_{\tt
relu}\left(\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>\right)\right|\leq\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\right]\geq
1-\mathcal{O}(\frac{\delta}{L}).$ (38)

Furthermore, the inductive hypothesis $P_{1}(h-1)$ implies that

$\displaystyle\left|\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>-\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq
h\cdot\frac{\varepsilon^{2}}{60L^{3}}.$ (39)

Hence, using ?? we find that,

$\displaystyle\left|P^{(p)}_{\tt
relu}\left(\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>\right)-P^{(p)}_{\tt
relu}\left(\Sigma_{\tt relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)\right|\leq
h\cdot\frac{\varepsilon^{2}}{60L^{3}}.$ (40)

By incorporating the above inequality into ?? using triangle inequality we
find that,

$\Pr\left[\left|\left<\phi^{(h)}(y),\phi^{(h)}(z)\right>-P^{(p)}_{\tt
relu}\left(\Sigma_{\tt relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)\right|\leq\frac{h\cdot\varepsilon^{2}}{60L^{3}}+\mathcal{O}\left(\frac{\varepsilon^{2}}{L^{3}}\right)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (41)

Now, by invoking ?? and using the fact that $p=\left\lceil
2L^{2}/{\varepsilon}^{4/3}\right\rceil$ we have,

$\displaystyle\left|P_{\tt relu}^{(p)}\left(\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)-\kappa_{1}\left(\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)\right|\leq\frac{\varepsilon^{2}}{76L^{3}}.$
(42)

By combining the above inequality with ?? using triangle inequality and using
the fact that $\kappa_{1}\left(\Sigma_{\tt relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)=\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$ (by
??), we get the following bound,

$\displaystyle\Pr\left[\left|\left<\phi^{(h)}(y),\phi^{(h)}(z)\right>-\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq(h+1)\cdot\frac{\varepsilon^{2}}{60L^{3}}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (43)

Similarly, we can prove that

$\displaystyle\Pr\left[\left|\left\|\phi^{(h)}(y)\right\|_{2}^{2}-1\right|>\frac{(h+1)\cdot\varepsilon^{2}}{60L^{3}}\right]\leq\mathcal{O}\left(\frac{\delta}{L}\right),\text{
and }$ (44)
$\displaystyle\Pr\left[\left|\left\|\phi^{(h)}(z)\right\|_{2}^{2}-1\right|>\frac{(h+1)\cdot\varepsilon^{2}}{60L^{3}}\right]\leq\mathcal{O}\left(\frac{\delta}{L}\right).$
(45)

This is sufficient to prove the inductive step by union bound, i.e.,
$\Pr[P_{1}(h)|P_{1}(h-1)]\geq 1-\mathcal{O}\left(\frac{\delta}{L}\right)$.

Now we prove the inductive step for statement $P_{2}(h)$, that is, we prove
that conditioned on $P_{2}(h-1),P_{1}(h),P_{1}(h-1)$, statement $P_{2}(h)$
holds with probability at least $1-\mathcal{O}\left(\frac{\delta}{L}\right)$.
First, note that by ?? and using ?? we have,

$\Pr\left[\left|\left<\dot{\phi}^{(h)}(y),\dot{\phi}^{(h)}(z)\right>-\sum_{j=0}^{2p^{\prime}+1}b_{j}\left<Y^{(h)}_{j}(y),Y^{(h)}_{j}(z)\right>\right|\leq
O\left(\frac{\varepsilon}{L}\right)\cdot\widehat{A}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right),$ (46)

where
$\widehat{A}:=\sqrt{\sum_{j=0}^{2p^{\prime}+1}b_{j}\|Y^{(h)}_{j}(y)\|_{2}^{2}}\cdot\sqrt{\sum_{j=0}^{2p^{\prime}+1}b_{j}\|Y^{(h)}_{j}(z)\|_{2}^{2}}$
and the collection of vectors
$\left\\{Y^{(h)}_{j}(y)\right\\}_{j=0}^{2p^{\prime}+1}$ and
$\left\\{Y^{(h)}_{j}(z)\right\\}_{j=0}^{2p^{\prime}+1}$ and coefficients
$\\{b_{j}\\}_{j=0}^{2p^{\prime}+1}$ are defined as per ?? and ??,
respectively. By invoking ?? along with union bound, with probability at least
$1-\mathcal{O}\left(\frac{\delta}{L}\right)$, the following inequalities hold
true simultaneously for all $j=0,1,\ldots 2p^{\prime}+1$

$\displaystyle\left|\left<Y^{(h)}_{j}(y),Y^{(h)}_{j}(z)\right>-\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>^{j}\right|\leq\mathcal{O}\left(\frac{\varepsilon}{L}\right)\cdot\left\|\phi^{(h-1)}(y)\right\|_{2}^{j}\left\|\phi^{(h-1)}(z)\right\|_{2}^{j}$
$\displaystyle\left\|Y^{(h)}_{j}(y)\right\|_{2}^{2}\leq\frac{11}{10}\cdot\left\|\phi^{(h-1)}(y)\right\|_{2}^{2j}$
(47)
$\displaystyle\left\|Y^{(h)}_{j}(z)\right\|_{2}^{2}\leq\frac{11}{10}\cdot\left\|\phi^{(h-1)}(z)\right\|_{2}^{2j}$

Therefore, by plugging ?? into ?? and using union bound, triangle inequality
and Cauchy–Schwarz inequality we find that,

$\Pr\left[\left|\left<\dot{\phi}^{(h)}(y),\dot{\phi}^{(h)}(z)\right>-\dot{P}^{(p^{\prime})}_{\tt
relu}\left(\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>\right)\right|\leq\mathcal{O}\left(\frac{\varepsilon}{L}\right)\cdot\widehat{B}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right),$ (48)

where $\widehat{B}:=\sqrt{\dot{P}^{(p^{\prime})}_{\tt
relu}\left(\|\phi^{(h-1)}(y)\|_{2}^{2}\right)\cdot\dot{P}^{(p^{\prime})}_{\tt
relu}\left(\|\phi^{(h-1)}(z)\|_{2}^{2}\right)}$ and
$\dot{P}^{(p^{\prime})}_{\tt
relu}(\alpha)=\sum_{j=0}^{2p^{\prime}+1}b_{j}\cdot\alpha^{j}$ is the
polynomial defined in ??. By inductive hypothesis $P_{1}(h-1)$ we have
$\left|\left\|\phi^{(h-1)}(y)\right\|_{2}^{2}-1\right|\leq
h\cdot\frac{\varepsilon^{2}}{60L^{3}}$ and
$\left|\left\|\phi^{(h-1)}(z)\right\|_{2}^{2}-1\right|\leq
h\cdot\frac{\varepsilon^{2}}{60L^{3}}$. Therefore, using the fact that
$p^{\prime}=\left\lceil 9L^{2}/\varepsilon^{2}\right\rceil$ and ??,
$\left|\dot{P}_{\tt
relu}^{(p^{\prime})}\left(\|\phi^{(h-1)}(y)\|_{2}^{2}\right)-\dot{P}_{\tt
relu}^{(p^{\prime})}(1)\right|\leq\frac{h\cdot\varepsilon}{20L^{2}}$ and
$\left|\dot{P}_{\tt
relu}^{(p^{\prime})}\left(\|\phi^{(h-1)}(z)\|_{2}^{2}\right)-\dot{P}_{\tt
relu}^{(p^{\prime})}(1)\right|\leq\frac{h\cdot\varepsilon}{20L^{2}}$.
Consequently, since $\dot{P}_{\tt relu}^{(p^{\prime})}(1)\leq\dot{P}_{\tt
relu}^{(+\infty)}(1)=1$, we find that $\widehat{B}\leq\frac{11}{10}.$ By
plugging this into ?? we have,

$\Pr\left[\left|\left<\dot{\phi}^{(h)}(y),\dot{\phi}^{(h)}(z)\right>-\dot{P}^{(p^{\prime})}_{\tt
relu}\left(\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>\right)\right|\leq\mathcal{O}\left(\frac{\varepsilon}{L}\right)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (49)

Furthermore, inductive hypothesis $P_{1}(h-1)$ implies
$\left|\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>-\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq
h\cdot\frac{\varepsilon^{2}}{60L^{3}}$, hence, by invoking ?? we find that,

$\displaystyle\left|\dot{P}^{(p^{\prime})}_{\tt
relu}\left(\left<\phi^{(h-1)}(y),\phi^{(h-1)}(z)\right>\right)-\dot{P}^{(p^{\prime})}_{\tt
relu}\left(\Sigma_{\tt relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)\right|\leq\frac{h\cdot\varepsilon}{20L^{2}}.$
(50)

By plugging the above inequality into ?? using triangle inequality, we find
that,

$\Pr\left[\left|\left<\dot{\phi}^{(h)}(y),\dot{\phi}^{(h)}(z)\right>-\dot{P}^{(p^{\prime})}_{\tt
relu}\left(\Sigma_{\tt relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)\right|\leq\frac{h\cdot\varepsilon}{20L^{2}}+\mathcal{O}\left(\frac{\varepsilon}{L}\right)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (51)

Now, by invoking ?? and using the fact that $p^{\prime}=\left\lceil
9L^{2}/{\varepsilon}^{2}\right\rceil$ we have,

$\displaystyle\left|\dot{P}_{\tt relu}^{(p^{\prime})}\left(\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)-\kappa_{0}\left(\Sigma_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)\right|\leq\frac{\varepsilon}{15L}.$
(52)

By combining the above inequality with ?? using triangle inequality and using
the fact that $\kappa_{0}\left(\Sigma_{\tt relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right)=\dot{\Sigma}_{\tt
relu}^{(h)}\left(\frac{\langle y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)$ (by
??), we get the following bound,

$\Pr\left[\left|\left<\dot{\phi}^{(h)}(y),\dot{\phi}^{(h)}(z)\right>-\dot{\Sigma}_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq\frac{\varepsilon}{8L}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (53)

Similarly we can show that,

$\Pr\left[\left|\left\|\dot{\phi}^{(h)}(y)\right\|-1\right|>\frac{\varepsilon}{8L}\right]\leq\mathcal{O}\left(\frac{\delta}{L}\right),~{}\text{
and}~{}\Pr\left[\left|\left\|\dot{\phi}^{(h)}(z)\right\|-1\right|>\frac{\varepsilon}{8L}\right]\leq\mathcal{O}\left(\frac{\delta}{L}\right).$
(54)

Let $f:=\psi^{(h-1)}(y)\otimes\dot{\phi}^{(h)}(y)$ and
$g:=\psi^{(h-1)}(z)\otimes\dot{\phi}^{(h)}(z)$. Then by ?? and using ?? we
have the following,

$\Pr\left[\left|\left<\psi^{(h)}(y),\psi^{(h)}(z)\right>-\left<{\bm{Q}}^{2}f\oplus\phi^{(h)}(y),{\bm{Q}}^{2}g\oplus\phi^{(h)}(z)\right>\right|\leq\mathcal{O}\left(\frac{\varepsilon}{L}\right)\cdot
D\right]\geq 1-\mathcal{O}\left(\frac{\delta}{L}\right),$ (55)

where
$D:=\left\|{\bm{Q}}^{2}f\oplus\phi^{(h)}(y)\right\|_{2}\left\|{\bm{Q}}^{2}g\oplus\phi^{(h)}(z)\right\|_{2}$.
By the fact that we conditioned on $P_{1}(h)$,

$D\leq\sqrt{\left\|{\bm{Q}}^{2}f\right\|_{2}^{2}+\frac{11}{10}}\cdot\sqrt{\left\|{\bm{Q}}^{2}g\right\|_{2}^{2}+\frac{11}{10}}.$

By ??, we can further obtain an upper bound:

$D\leq\frac{11}{10}\cdot\sqrt{\left\|f\right\|_{2}^{2}+1}\cdot\sqrt{\left\|g\right\|_{2}^{2}+1}.$

Now note that because we conditioned on $P_{2}(h-1)$ and using ??, with
probability at least $1-\mathcal{O}\left(\frac{\delta}{L}\right)$ the
following holds:

$\left\|f\right\|_{2}^{2}=\left\|\psi^{(h-1)}(y)\right\|_{2}^{2}\left\|\dot{\phi}^{(h)}(y)\right\|_{2}^{2}\leq\frac{11}{10}\cdot
K^{(h-1)}_{\tt relu}(1)=\frac{11}{10}h.$

Similarly, $\left\|g\right\|_{2}^{2}\leq\frac{11}{10}h$ with probability at
least $1-\mathcal{O}\left(\frac{\delta}{L}\right)$, thus, by union bound:

$\Pr[D\leq 2(h+1)|P_{2}(h-1),P_{1}(h),P_{1}(h-1)]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$

Therefore, by combining the above with ?? via union bound we find that,

$\Pr\left[\left|\left<\psi^{(h)}(y),\psi^{(h)}(z)\right>-\left<{\bm{Q}}^{2}f\oplus\phi^{(h)}(y),{\bm{Q}}^{2}g\oplus\phi^{(h)}(z)\right>\right|\leq\mathcal{O}\left(\frac{\varepsilon
h}{L}\right)\right]\geq 1-\mathcal{O}\left(\frac{\delta}{L}\right),$ (56)

Now note that
$\left<{\bm{Q}}^{2}f\oplus\phi^{(h)}(y),{\bm{Q}}^{2}g\oplus\phi^{(h)}(z)\right>=\left<{\bm{Q}}^{2}f,{\bm{Q}}^{2}g\right>+\left<\phi^{(h)}(y),\phi^{(h)}(z)\right>$.
We proceed by bounding the term
$\left|\left<{\bm{Q}}^{2}f,{\bm{Q}}^{2}g\right>-\left<f,g\right>\right|$ using
??, as follows,

$\displaystyle\Pr\left[\left|\left<{\bm{Q}}^{2}f,{\bm{Q}}^{2}g\right>-\left<f,g\right>\right|\leq\mathcal{O}\left(\frac{\varepsilon}{L}\right)\cdot\|f\|_{2}\|g\|_{2}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (57)

We proved that conditioned on $P_{2}(h-1)$ and $P_{1}(h-1)$,
$\left\|f\right\|_{2}^{2}\leq 11h/10$ and $\left\|g\right\|_{2}^{2}\leq
11h/10$ with probability at least
$1-\mathcal{O}\left(\frac{\delta}{L}\right)$. Hence, by union bound we find
that,

$\Pr\left[\left.\left|\left<{\bm{Q}}^{2}f,{\bm{Q}}^{2}g\right>-\left<f,g\right>\right|\leq\mathcal{O}\left(\frac{\varepsilon
h}{L}\right)\right|P_{2}(h-1),P_{1}(h-1)\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (58)

Note that
$\left<f,g\right>=\left<\psi^{(h-1)}(y),\psi^{(h-1)}(z)\right>\cdot\left<\dot{\phi}^{(h)}(y),\dot{\phi}^{(h)}(z)\right>$,
thus by conditioning on inductive hypothesis $P_{2}(h-1)$ and ?? we have,

$\begin{split}\left|\left<f,g\right>-K_{\tt relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\cdot\dot{\Sigma}_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|&\leq\frac{\varepsilon}{8L}\left(h+\varepsilon\cdot\frac{(h-1)^{2}+1}{10L}\right)+\varepsilon\cdot\frac{(h-1)^{2}+1}{10L}\end{split}$

By combining the above inequality with ??, $P_{1}(h)$, and ?? using triangle
inequality and union bound we get the following inequality,

$\Pr\left[\left|\left<\psi^{(h)}(y),\psi^{(h)}(z)\right>-K_{\tt
relu}^{(h-1)}\left(\frac{\langle
y,z\rangle}{\|y\|\|z\|}\right)\cdot\dot{\Sigma}_{\tt
relu}^{(h)}\left(\frac{\langle y,z\rangle}{\|y\|\|z\|}\right)-\Sigma_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|\|z\|}\right)\right|>\varepsilon\cdot\frac{h^{2}+1}{10L}\right]\leq
O\left(\frac{\delta}{L}\right).$

By noting that $K_{\tt relu}^{(h-1)}\left(\alpha\right)\cdot\dot{\Sigma}_{\tt
relu}^{(h)}\left(\alpha\right)+\Sigma_{\tt
relu}^{(h)}\left(\alpha\right)=K_{\tt relu}^{(h)}\left(\alpha\right)$ (see ??)
we have proved that

$\displaystyle\Pr\left[\left|\left<\psi^{(h)}(y),\psi^{(h)}(z)\right>-K_{\tt
relu}^{(h)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq\varepsilon\cdot\frac{h^{2}+1}{10L}\right]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$ (59)

Similarly we can prove the following inequalities hold with probability at
least $1-\mathcal{O}\left(\frac{\delta}{L}\right)$,

$\left|\left\|\psi^{(h)}(y)\right\|_{2}^{2}-K_{\tt
relu}^{(h)}(1)\right|\leq\varepsilon\cdot\frac{h^{2}+1}{10L},\text{ and
}\left|\left\|\psi^{(h)}(z)\right\|_{2}^{2}-K_{\tt
relu}^{(h)}(1)\right|\leq\varepsilon\cdot\frac{h^{2}+1}{10L}.$

This proves the inductive step for the statement $P_{2}(h)$ follows, i.e.,

$\Pr[P_{2}(h)|P_{2}(h-1),P_{1}(h),P_{1}(h-1)]\geq
1-\mathcal{O}\left(\frac{\delta}{L}\right).$

Therefore, by union bounding over all $h=0,1,2,\ldots L$, it follows that the
statements of the lemma hold simultaneously for all $h$ with probability at
least $1-\delta$. This completes the proof of ??. ∎

We now analyze the runtime of the NTKSketch algorithm:

######  Lemma 6 (Runtime of NTKSketch).

For every positive integers $d$ and $L$, every $\varepsilon,\delta>0$, every
vector $x\in\mathbb{R}^{d}$, the time to compute NTKSketch $\Psi_{\tt
ntk}^{(L)}(x)\in\mathbb{R}^{s^{*}}$, for
$s^{*}=\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\cdot\log\frac{1}{\delta}\right)$,
using the procedure given in ?? is bounded by,

$\mathcal{O}\left(\frac{L^{11}}{\varepsilon^{6.7}}\cdot\log^{3}\frac{L}{\varepsilon\delta}+\frac{L^{3}}{\varepsilon^{2}}\cdot\log\frac{L}{\varepsilon\delta}\cdot{\rm
nnz}(x)\right).$

Proof of ??: There are three main components to the runtime of this procedure
that we have to account for. The first is the time to apply the sketch
${\bm{Q}}^{1}$ to $x$ in line 4 of ??. By ??, the runtime of computing
${\bm{Q}}^{1}\cdot x$ is
$\mathcal{O}\left(\frac{L^{6}}{\varepsilon^{4}}\cdot\log^{3}\frac{L}{\varepsilon\delta}+\frac{L^{3}}{\varepsilon^{2}}\cdot\log\frac{L}{\varepsilon\delta}\cdot{\rm
nnz}(x)\right)$. The second heavy operation corresponds to computing vectors
$Z_{j}^{(h)}(x)={\bm{Q}}^{2p+2}\cdot\left(\left[\phi^{(h-1)}(x)\right]^{\otimes
j}\otimes{e}_{1}^{\otimes 2p+2-j}\right)$ for $j=0,1,2,\ldots 2p+2$ and
$h=1,2,\ldots L$ in ??. By ??, the time to compute $Z_{j}^{(h)}(x)$ for a
fixed $h$ and all $j=0,1,2,\ldots 2p+2$ is bounded by,

$\mathcal{O}\left(\frac{L^{10}}{\varepsilon^{20/3}}\cdot\log^{2}\frac{L}{\varepsilon}\cdot\log^{3}\frac{L}{\varepsilon\delta}+\frac{L^{8}}{\varepsilon^{16/3}}\cdot\log^{3}\frac{L}{\varepsilon\delta}\right)=\mathcal{O}\left(\frac{L^{10}}{\varepsilon^{6.7}}\cdot\log^{3}\frac{L}{\varepsilon\delta}\right).$

The total time to compute vectors $Z_{j}^{(h)}(x)$ for all $h=1,2,\ldots L$
and all $j=0,1,2,\ldots 2p+2$ is thus
$\mathcal{O}\left(\frac{L^{11}}{\varepsilon^{6.7}}\cdot\log^{3}\frac{L}{\varepsilon\delta}\right)$.
Finally, the last computationally expensive operation is computing vectors
$Y_{j}^{(h)}(x)={\bm{Q}}^{2p^{\prime}+1}\cdot\left(\left[\phi^{(h-1)}(x)\right]^{\otimes
j}\otimes{e}_{1}^{\otimes 2p^{\prime}+1-j}\right)$ for $j=0,1,2,\ldots
2p^{\prime}+1$ and $h=1,2,\ldots L$ in ??. By ??, the runtime of computing
$Y_{j}^{(h)}(x)$ for a fixed $h$ and all $j=0,1,2,\ldots 2p^{\prime}+1$ is
bounded by,

$\mathcal{O}\left(\frac{L^{6}}{\varepsilon^{6}}\cdot\log^{2}\frac{L}{\varepsilon}\cdot\log^{3}\frac{L}{\varepsilon\delta}+\frac{L^{8}}{\varepsilon^{6}}\cdot\log^{3}\frac{L}{\varepsilon\delta}\right)=\mathcal{O}\left(\frac{L^{8}}{\varepsilon^{6}}\cdot\log^{3}\frac{L}{\varepsilon\delta}\right).$

Hence, the total time to compute vectors $Y_{j}^{(h)}(x)$ for all
$h=1,2,\ldots L$ and all $j=0,1,2,\ldots 2p^{\prime}+1$ is
$\mathcal{O}\left(\frac{L^{9}}{\varepsilon^{6}}\cdot\log^{3}\frac{L}{\varepsilon\delta}\right)$.
The total runtime of the NTK Sketch is obtained by summing up these three
contributions. This completes the proof of ??. ∎

Now we are ready to prove the main theorem on NTKSketch. See 1

Proof of ??: Let $\psi^{(L)}:\mathbb{R}^{d}\to\mathbb{R}^{s}$ for
$s=\mathcal{O}\left(\frac{L^{2}}{\varepsilon^{2}}\cdot\log^{2}\frac{L}{\varepsilon\delta}\right)$
be the mapping defined in ?? of ??. By ??, the NTK Sketch $\Psi_{\tt
ntk}^{(L)}(x)$ is defined as

$\Psi_{\tt ntk}^{(L)}(x):={\|x\|_{2}}\cdot{\bm{G}}\cdot\psi^{(L)}(x).$

Because ${\bm{G}}$ is a matrix of i.i.d normal entries with
$s^{*}=C\cdot\frac{1}{\varepsilon^{2}}\cdot\log\frac{1}{\delta}$ rows, for a
large enough constant $C$, ${\bm{G}}$ is a JL transform [15] and hence
$\Psi_{\tt ntk}^{(L)}$ satisfies the following,

$\Pr\left[\left|\left<\Psi_{\tt ntk}^{(L)}(y),\Psi_{\tt
ntk}^{(L)}(z)\right>-{\|y\|_{2}}{\|z\|_{2}}\cdot\left<\psi^{(L)}(y),\psi^{(L)}(z)\right>\right|\leq\mathcal{O}(\varepsilon)\cdot
A\right]\geq 1-\mathcal{O}(\delta),$

where
$A:={\|y\|_{2}}{\|z\|_{2}}\left\|\psi^{(L)}({y})\right\|_{2}\left\|\psi^{(L)}({z})\right\|_{2}$.
By ?? and using the fact that $K_{\tt relu}^{(L)}(1)=L+1$, the following
bounds hold with probability at least $1-\mathcal{O}(\delta)$:

$\left\|\psi^{(L)}({y})\right\|_{2}^{2}\leq\frac{11}{10}\cdot(L+1),\text{ and
}\left\|\psi^{(L)}({z})\right\|_{2}^{2}\leq\frac{11}{10}\cdot(L+1).$

Therefore, by union bound we find that,

$\Pr\left[\left|\left<\Psi_{\tt ntk}^{(L)}(y),\Psi_{\tt
ntk}^{(L)}(z)\right>-{\|y\|_{2}}{\|z\|_{2}}\cdot\left<\psi^{(L)}({y}),\psi^{(L)}({z})\right>\right|\leq\mathcal{O}(\varepsilon
L)\cdot{\|y\|_{2}}{\|z\|_{2}}\right]\geq 1-\mathcal{O}(\delta).$

Additionally, by ??, the following holds with probability at least
$1-\mathcal{O}(\delta)$:

$\left|\left<\psi^{(L)}({y}),\psi^{(L)}({z})\right>-K_{\tt
relu}^{(L)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq\frac{\varepsilon(L+1)}{10}.$

Hence by union bound and triangle inequality we have,

$\Pr\left[\left|\left<\Psi_{\tt ntk}^{(L)}(y),\Psi_{\tt
ntk}^{(L)}(z)\right>-{\|y\|_{2}}{\|z\|_{2}}\cdot K_{\tt
relu}^{(L)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)\right|\leq\frac{\varepsilon(L+1)}{9}\cdot{\|y\|_{2}}{\|z\|_{2}}\right]\geq
1-\mathcal{O}(\delta).$

Now note that by ??, ${\|y\|_{2}}{\|z\|_{2}}\cdot K_{\tt
relu}^{(L)}\left(\frac{\langle
y,z\rangle}{\|y\|_{2}\|z\|_{2}}\right)=\Theta_{\tt ntk}^{(L)}(y,z)$, and also
note that for every $L\geq 2$ and any $\alpha\in[-1,1]$, $K_{\tt
relu}^{(L)}\left(\alpha\right)\geq(L+1)/9$, therefore,

$\Pr\left[\left|\left<\Psi_{\tt ntk}^{(L)}(y),\Psi_{\tt
ntk}^{(L)}(z)\right>-\Theta_{\tt
ntk}^{(L)}(y,z)\right|\leq\varepsilon\cdot\Theta_{\tt
ntk}^{(L)}(y,z)\right]\geq 1-\delta.$

Remark on the fact that $K_{\tt relu}^{(L)}\left(\alpha\right)\geq(L+1)/9$ for
every $L\geq 2$ and any $\alpha\in[-1,1]$. Note that from the definition of
$\Sigma_{\tt relu}^{(h)}$ in ??, we have that for any $\alpha\in[-1,1]$:
$\Sigma_{\tt relu}^{(0)}(\alpha)\geq-1$, $\Sigma_{\tt relu}^{(1)}(\alpha)\geq
0$, $\Sigma_{\tt relu}^{(2)}(\alpha)\geq\frac{1}{\pi}$, and $\Sigma_{\tt
relu}^{(h)}(\alpha)\geq\frac{1}{2}$ for every $h\geq 3$ because
$\kappa_{1}(\cdot)$ is a monotonically increasing function on the interval
$[-1,1]$. Moreover, using the definition of $\dot{\Sigma}_{\tt relu}^{(h)}$ in
??, we have that for any $\alpha\in[-1,1]$: $\dot{\Sigma}_{\tt
relu}^{(1)}(\alpha)\geq 0$, $\dot{\Sigma}_{\tt
relu}^{(2)}(\alpha)\geq\frac{1}{2}$, and $\dot{\Sigma}_{\tt
relu}^{(h)}(\alpha)\geq\frac{3}{5}$ for every $h\geq 3$ because
$\kappa_{0}(\cdot)$ is a monotonically increasing function on the interval
$[-1,1]$. By an inducive proof and using the definition of $K_{\tt
relu}^{(L)}$ in ??, we can show that $K_{\tt
relu}^{(L)}\left(\alpha\right)\geq(L+1)/9$ for every $L\geq 2$ and any
$\alpha\in[-1,1]$

Runtime analysis: By ??, runtime to compute the NTKSketch is

$\displaystyle\mathcal{O}\left(\frac{L^{11}}{\varepsilon^{6.7}}\log^{3}\frac{L}{\varepsilon\delta}+\frac{L^{3}}{\varepsilon^{2}}\log\frac{L}{\varepsilon\delta}\cdot{\rm
nnz}(x)\right).$ (60)

This completes the proof of ??.∎

##  Appendix D NTK Random Features: Claims and Proofs

###  D.1 Proof of ??

In this section we prove ??. We first restate the theorem: See 2

In the proof of this theorem we use the following results from the literature,

######  Lemma 7 (Corollary 16 in [14]).

Given integer $\ell>0$ and $x,x^{\prime}\in\mathbb{R}^{d}$ such that
$\left\|y\right\|_{2}=\left\|z\right\|_{2}=1$, let $\phi_{\tt
rf}^{(\ell)}(y),\phi_{\tt rf}^{(\ell)}(z)$ be defined as per line 5 of ??. For
$\delta_{1},\varepsilon_{1}\in(0,1)$, there exists a constant $C_{1}>0$ such
that for any $m_{1}\geq
C_{1}\frac{L^{2}}{\varepsilon_{1}^{2}}\log\left(\frac{L}{\delta_{1}}\right)$
the following holds:

$\displaystyle\Pr\left[\left|\left\langle{\phi_{\tt rf}^{(\ell)}(y),\phi_{\tt
rf}^{(\ell)}(z)}\right\rangle-\Sigma_{\tt relu}^{(\ell)}(\langle
y,z\rangle)\right|\leq\varepsilon_{1}\right]\geq 1-\delta_{1}.$ (61)

We also need the following lemma,

######  Lemma 8 (Lemma E.5 in [5]).

Given $x,x^{\prime}\in{\mathbb{R}}^{d}$ with
$\left\|x\right\|_{2}=\left\|x^{\prime}\right\|_{2}=1$, integer $\ell>0$, let
$\phi_{\tt rf}^{(\ell)}(x),\phi_{\tt rf}^{(\ell)}(x^{\prime})$ be defined as
per line 5 of ??. For any $\varepsilon_{2}\in(0,1)$ assume that

$\displaystyle\left|\left\langle{\phi_{\tt rf}^{(\ell)}(x),\phi_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle-\Sigma_{\tt relu}^{(\ell)}(\langle
x,x^{\prime}\rangle)\right|\leq\frac{\varepsilon_{2}^{2}}{2}.$ (62)

Then, for $\dot{\phi}_{\tt rf}^{(\ell)}(x),\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})$ defined as in line 4 of ??, and any $\delta_{2}>0$
the following holds:

$\displaystyle\Pr\left[\left|\left\langle{\dot{\phi}_{\tt
rf}^{(\ell)}(x),\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle-\dot{\Sigma}_{\tt
relu}^{(\ell)}(\langle
x,x^{\prime}\rangle)\right|\leq\varepsilon_{2}+\sqrt{\frac{2}{m_{0}}\log\left(\frac{6}{\delta_{2}}\right)}\right]\geq
1-\delta_{2}.$ (63)

Proof of ??: For fixed $y,z\in\mathbb{R}^{d}$ and $\ell=0,\dots,L$, we denote
the estimation error as

$\displaystyle\Delta_{\ell}:=\max_{(x,x^{\prime})\in\\{(y,z),(y,y),(z,z)\\}}\left|\left\langle{\psi_{\tt
rf}^{(\ell)}(x),\psi_{\tt rf}^{(\ell)}(x^{\prime})}\right\rangle-K_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|$

and note that $\Delta_{0}=0$. Recall that, by ?? and ??:

$\displaystyle\Theta_{\tt
ntk}^{(\ell)}(x,x^{\prime})=\|x\|_{2}\|x^{\prime}\|_{2}\cdot K_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right),$ (64)

where for every $\alpha\in[-1,1]$ and $\ell=1,\dots,L$:

$\displaystyle K_{\tt relu}^{(\ell)}(\alpha)$ $\displaystyle:=K_{\tt
relu}^{(\ell-1)}(\alpha)\cdot\dot{\Sigma}_{\tt
relu}^{(\ell)}(\alpha)+\Sigma_{\tt relu}^{(\ell)}(\alpha),$
$\displaystyle\dot{\Sigma}_{\tt relu}^{(\ell)}(\alpha)$
$\displaystyle:=\kappa_{0}\left(\Sigma_{\tt relu}^{(\ell-1)}(\alpha)\right),$
$\displaystyle\Sigma_{\tt relu}^{(\ell)}(\alpha)$
$\displaystyle:=\kappa_{1}\left(\Sigma_{\tt relu}^{(\ell-1)}(\alpha)\right).$

We use the recursive relation to approximate:

$\displaystyle{\left\langle{\psi_{\tt rf}^{(\ell)}(x),\psi_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle}$ $\displaystyle=\left\langle{\phi_{\tt
rf}^{(\ell)}(x),\phi_{\tt rf}^{(\ell)}(x^{\prime})}\right\rangle$
$\displaystyle\qquad+\left\langle{{\bm{Q}}^{2}\cdot\left(\dot{\phi}_{\tt
rf}^{(\ell)}(x)\otimes{\psi}_{\tt
rf}^{(\ell-1)}(x)\right),{\bm{Q}}^{2}\cdot\left(\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})\otimes{\psi}_{\tt
rf}^{(\ell-1)}(x^{\prime})\right)}\right\rangle$
$\displaystyle\approx\left\langle{\phi_{\tt rf}^{(\ell)}(x),\phi_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle+\left\langle{\dot{\phi}_{\tt
rf}^{(\ell)}(x)\otimes{\psi}_{\tt rf}^{(\ell-1)}(x),\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})\otimes{\psi}_{\tt
rf}^{(\ell-1)}(x^{\prime})}\right\rangle$
$\displaystyle\approx\left(\Sigma_{\tt relu}^{(\ell)}+\dot{\Sigma}_{\tt
relu}^{(\ell)}\cdot K_{\tt relu}^{(\ell-1)}\right)\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)=K_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right).$

For notational simplicity, we define the following events:

$\displaystyle\mathcal{E}_{\phi}^{(\ell)}(\varepsilon)$
$\displaystyle:=\left\\{\left|\left\langle{\phi_{\tt rf}^{(\ell)}(x),\phi_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle-\Sigma_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|\leq\varepsilon:\forall(x,x^{\prime})\in\\{(y,z),(y,y),(z,z)\\}\right\\},$
(65) $\displaystyle\dot{\mathcal{E}}_{\phi}^{(\ell)}(\varepsilon)$
$\displaystyle:=\left\\{\left|\left\langle{\dot{\phi}_{\tt
rf}^{(\ell)}(x),\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle-\dot{\Sigma}_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|\leq\varepsilon:\forall(x,x^{\prime})\in\\{(y,z),(y,y),(z,z)\\}\right\\}.$
(66)

Our proof is based on the following claims:

First we claim that, there exists a constant $C_{1}>0$ such that for any
$m_{1}\geq
C_{1}\frac{L^{6}}{\varepsilon^{4}}\log\left(\frac{L}{\delta}\right)$:

$\Pr\left[\mathcal{E}_{\phi}^{(\ell)}\left(\frac{\varepsilon^{2}}{100L^{2}}\right)\right]\geq
1-\frac{\delta}{3L}.$ (67)

which directly follows by invoking ?? with
$(x,x^{\prime})\in\left\\{\left(\frac{y}{\|y\|_{2}},\frac{z}{\|z\|_{2}}\right),\left(\frac{y}{\|y\|_{2}},\frac{y}{\|y\|_{2}}\right),\left(\frac{z}{\|z\|_{2}},\frac{z}{\|z\|_{2}}\right)\right\\}$
and setting $\varepsilon_{1}=\frac{\varepsilon^{2}}{100L^{2}}$ and
$\delta_{1}=\frac{\delta}{9L}$ and applying union bound over choices of
$(x,x^{\prime})$.

Our second claim is that 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

$\Pr\left[\left.\dot{\mathcal{E}}_{\phi}^{(\ell)}\left(\frac{\varepsilon}{8L}\right)\right|\mathcal{E}_{\phi}^{(\ell)}\left(\frac{\varepsilon^{2}}{100L^{2}}\right)\right]\geq
1-\frac{\delta}{3L}.$ (68)

The above statement is a direct consequence of invoking ?? with
$(x,x^{\prime})\in\left\\{\left(\frac{y}{\|y\|_{2}},\frac{z}{\|z\|_{2}}\right),\left(\frac{y}{\|y\|_{2}},\frac{y}{\|y\|_{2}}\right),\left(\frac{z}{\|z\|_{2}},\frac{z}{\|z\|_{2}}\right)\right\\}$
and union bounding over choices of $(x,x^{\prime})$. In ??, we choose
$\varepsilon_{2}=\frac{\varepsilon}{10L},\delta_{2}=\frac{\delta}{9L}$ and
$m_{0}\geq\frac{3200L^{2}}{\varepsilon^{2}}\log\left(\frac{54L}{\delta}\right)$
for $\varepsilon,\delta\in(0,1)$ to obtain ?? by union bound.

Our next claim is the following,

######  Claim 1.

Let $f(x):=\dot{\phi}_{\tt rf}^{(\ell)}(x)\otimes{\psi}_{\tt
rf}^{(\ell-1)}(x)$ for every $x\in\mathbb{R}^{d}$. There exists a constant
$C_{2}>0$ such that if $m_{s}\geq
C_{2}\frac{L^{2}}{\varepsilon^{2}}\log\frac{L}{\varepsilon\delta}$ then
conditioned on the events
$\dot{\mathcal{E}}_{\phi}^{(\ell)}\left(\frac{\varepsilon}{8L}\right)$ and
$\mathcal{E}_{\phi}^{(\ell)}\left(\frac{\varepsilon^{2}}{100L^{2}}\right)$ the
following holds with probability at least $1-\frac{\delta}{3L}$, for all
$(x,x^{\prime})\in\left\\{\left(\frac{y}{\|y\|_{2}},\frac{z}{\|z\|_{2}}\right),\left(\frac{y}{\|y\|_{2}},\frac{y}{\|y\|_{2}}\right),\left(\frac{z}{\|z\|_{2}},\frac{z}{\|z\|_{2}}\right)\right\\}$,

$\displaystyle\left|\left\langle{{\bm{Q}}^{2}\cdot f(x),{\bm{Q}}^{2}\cdot
f(x^{\prime})}\right\rangle-\left\langle{f(x),f(x^{\prime})}\right\rangle\right|\leq\frac{\varepsilon}{100L}(\ell+\Delta_{\ell-1}).$

The proof of the above claim is provided in ??. Now, by combining ??, ?? and
??, we can prove the following claim which provides a recursive relation for
bounding $\Delta_{\ell}$.

######  Claim 2.

If the events in ??, ?? and ?? hold, then

$\displaystyle\Delta_{\ell}\leq\left(1+\frac{\varepsilon}{7L}\right)\Delta_{\ell-1}+\frac{\varepsilon}{7L}\ell.$
(69)

The proof of ?? is provided in ??.

Note that by union bound, with probability $1-\frac{\delta}{L}$, all
preconditions of ?? hold and thus ?? holds with probability
$1-\frac{\delta}{L}$. Applying union bound on ?? for all $\ell\in[L]$ and
solving the recurrence, we obtain that with probability at least $1-\delta$,
the following bound holds

$\displaystyle\Delta_{\ell}\leq\frac{\varepsilon}{9L}\cdot\ell^{2}.$ (70)

When $L\geq 2$, we showed in the proof of ?? that $K_{\tt
relu}^{(L)}(\cdot)\geq\frac{L+1}{9}$, therefore,

$\displaystyle\left|\left\langle{\psi_{\tt rf}^{(L)}(x),\psi_{\tt
rf}^{(L)}(x^{\prime})}\right\rangle-K_{\tt relu}^{(L)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|\leq\varepsilon\cdot
K_{\tt relu}^{(L)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right).$ (71)

Since ${\Psi}_{\tt rf}^{(L)}(x)=\|x\|_{2}\cdot{\psi}^{(L)}(x)$ and
${\Psi}_{\tt
rf}^{(L)}(x^{\prime})=\|x^{\prime}\|_{2}\cdot{\psi}^{(L)}(x^{\prime})$, this
implies that,

$\Pr\left[\left|\left\langle{\Psi_{\tt rf}^{(L)}(x),\Psi_{\tt
rf}^{(L)}(x^{\prime})}\right\rangle-\Theta_{\tt
ntk}^{(L)}(x,x^{\prime})\right|\leq\varepsilon\cdot\Theta_{\tt
ntk}^{(L)}(x,x^{\prime})\right]\geq 1-\delta.$

This completes the proof of ??. ∎

###  D.2 Proof of Auxiliary Claims

See 1 Proof of ??: The proof is based on ?? that provides an upper bound on
variance of the PolySketch. By using the definition of $f(x),f(x^{\prime})$
and ??, with probability at least $1-\frac{\delta}{9L}$, we have

$\displaystyle\left|\left\langle{{\bm{Q}}^{2}\cdot f(x),{\bm{Q}}^{2}\cdot
f(x^{\prime})}\right\rangle-\left\langle{f(x),f(x^{\prime})}\right\rangle\right|$
$\displaystyle\qquad\leq\frac{\varepsilon}{200L}\left\|\dot{\phi}_{\tt
rf}^{(\ell)}(x)\right\|_{2}\left\|\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})\right\|_{2}\left\|\psi_{\tt
rf}^{(\ell-1)}(x)\right\|_{2}\left\|\psi_{\tt
rf}^{(\ell-1)}(x^{\prime})\right\|_{2}$
$\displaystyle\qquad\leq\frac{\varepsilon}{200L}\left(1+\frac{\varepsilon}{8L}\right)\left\|\psi_{\tt
rf}^{(\ell-1)}(x)\right\|_{2}\left\|\psi_{\tt
rf}^{(\ell-1)}(x^{\prime})\right\|_{2}$
$\displaystyle\qquad\leq\frac{\varepsilon}{100L}\left(\ell+\Delta_{\ell-1}\right)$

where the second inequality follows from the assumption that
$\dot{\mathcal{E}}_{\phi}^{(\ell)}\left(\frac{\varepsilon}{8L}\right)$ holds
and the third one follows from the fact that $K_{\tt
relu}^{(\ell-1)}(\alpha)\leq\ell$ for any $\alpha\in[-1,1]$ and

$\displaystyle\left\|\psi_{\tt rf}^{(\ell-1)}(x)\right\|_{2}^{2}\leq K_{\tt
relu}^{(\ell-1)}(x,x^{\prime})+\Delta_{\ell-1}\leq\ell+\Delta_{\ell-1}.$ (72)

Union bounding over the choices of $(x,x^{\prime})$ completes the proof of ??.
∎

See 2 Proof of ??: Recall that

$\displaystyle\Delta_{\ell}:=\max_{(x,x^{\prime})\in\\{(y,z),(y,y),(z,z)\\}}\left|\left\langle{\psi_{\tt
rf}^{(\ell)}(x),\psi_{\tt rf}^{(\ell)}(x^{\prime})}\right\rangle-K_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|.$

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

$\displaystyle\left|\left\langle{\psi_{\tt rf}^{(\ell)}(x),\psi_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle-K_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|\leq\left|\left\langle{\phi_{\tt
rf}^{(\ell)}(x),\phi_{\tt rf}^{(\ell)}(x^{\prime})}\right\rangle-\Sigma_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|$
$\displaystyle\qquad\qquad+\left|\left\langle{{\bm{Q}}^{2}\cdot
f(x),{\bm{Q}}^{2}\cdot
f(x^{\prime})}\right\rangle-\left\langle{f(x),f(x^{\prime})}\right\rangle\right|$
(73) $\displaystyle\qquad\qquad+\left|\left\langle{\dot{\phi}_{\tt
rf}^{(\ell)}(x)\otimes{\psi}_{\tt rf}^{(\ell-1)}(x),\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})\otimes{\psi}_{\tt
rf}^{(\ell-1)}(x^{\prime})}\right\rangle-\dot{\Sigma}_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)K_{\tt
relu}^{(\ell-1)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|.$

for $(x,x^{\prime})\in\\{(y,z),(y,y),(z,z)\\}$. By the assumption that ??
holds, the definition of
$\mathcal{E}_{\phi}^{(\ell)}\left(\frac{\varepsilon^{2}}{100L^{2}}\right)$
implies that,

$\displaystyle\left|\left\langle{\phi_{\tt rf}^{(\ell)}(x),\phi_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle-\Sigma_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|\leq\frac{\varepsilon^{2}}{100L^{2}}$
(74)

Furthermore, the assumption that ?? holds, implies that,

$\displaystyle\left|\left\langle{{\bm{Q}}^{2}\cdot f(x),{\bm{Q}}^{2}\cdot
f(x^{\prime})}\right\rangle-\left\langle{f(x),f(x^{\prime})}\right\rangle\right|\leq\frac{\varepsilon}{100L}(\ell+\Delta_{\ell-1}).$
(75)

For the third part in ??, we observe that

$\displaystyle\left|\left\langle{\dot{\phi}_{\tt
rf}^{(\ell)}(x)\otimes{\psi}_{\tt rf}^{(\ell-1)}(x),\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})\otimes{\psi}_{\tt
rf}^{(\ell-1)}(x^{\prime})}\right\rangle-\dot{\Sigma}_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)K_{\tt
relu}^{(\ell-1)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|$
$\displaystyle\leq\left|\left\langle{\dot{\phi}_{\tt
rf}^{(\ell)}(x),\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle\right|\cdot\left|\left\langle{{\psi}_{\tt
rf}^{(\ell-1)}(x),{\psi}_{\tt rf}^{(\ell-1)}(x^{\prime})}\right\rangle-K_{\tt
relu}^{(\ell-1)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|$
$\displaystyle\quad\qquad+K_{\tt relu}^{(\ell-1)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\cdot\left|\left\langle{\dot{\phi}_{\tt
rf}^{(\ell)}(x),\dot{\phi}_{\tt
rf}^{(\ell)}(x^{\prime})}\right\rangle-\dot{\Sigma}_{\tt
relu}^{(\ell)}\left(\frac{\langle
x,x^{\prime}\rangle}{\|x\|_{2}\|x^{\prime}\|_{2}}\right)\right|$
$\displaystyle\leq\left(1+\frac{\varepsilon}{8L}\right)\cdot\Delta_{\ell-1}+\ell\cdot\frac{\varepsilon}{8L},$
(76)

where the second inequality comes from that the assumption that ?? holds along
with $\left|\dot{\Sigma}_{\tt relu}^{(\ell)}(\cdot)\right|\leq 1$ and
$\left|K_{\tt relu}^{(\ell-1)}(\cdot)\right|\leq\ell$. Putting ??, ?? and ??
into ??, we have

$\displaystyle\Delta_{\ell}\leq\frac{\varepsilon^{2}}{100L^{2}}+\frac{\varepsilon}{100L}\left(\ell+\Delta_{\ell-1}\right)+\left(1+\frac{\varepsilon}{8L}\right)\Delta_{\ell-1}+\frac{\ell\cdot\varepsilon}{8L}=\left(1+\frac{\varepsilon}{7L}\right)\Delta_{\ell-1}+\frac{\varepsilon}{7L}\ell.$
(77)

This completes the proof of ??. ∎

##  Appendix E Spectral Approximation via Leverage Scores Sampling

###  E.1 Zeroth Order Arc-Cosine Kernels

The proofs here rely on Theorem 3.3 in [28] which states spectral
approximation bounds of random features for general kernels equipped with the
leverage score sampling. This result is a generalization of [8] on the Random
Fourier Features.

######  Theorem 5 (Theorem 3.3 in [28]).

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]^{\top}.$
(78)

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}}^{\top}\overline{\bm{\Phi}}+\lambda{\bm{I}}\preceq\left(1+\varepsilon\right)\left({\bm{K}}+\lambda{\bm{I}}\right)$
(79)

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

We now ready to provide spectral approximation guarantee for arc-cosine
kernels of order zero.

######  Theorem 6.

Given dataset ${\bm{X}}\in{\mathbb{R}}^{d\times n}$, let
${\bm{K}}_{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{W}}{\bm{X}})\in{\mathbb{R}}^{m\times
n}$ where each entry in ${\bm{W}}\in{\mathbb{R}}^{m\times d}$ is an i.i.d.
sample from $\mathcal{N}(0,1)$. For
$\lambda\in(0,\left\|{\bm{K}}_{0}\right\|_{2})$, let $s_{\lambda}$ be the
statistical dimension of ${\bm{K}}_{0}$. Given $\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 that

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

with probability at least $1-\delta$.

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

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

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

$\displaystyle\tau_{\lambda}(w)$
$\displaystyle:=p(w)\cdot\bm{\Phi}_{0}(w)^{\top}\left({\bm{K}}_{0}+\lambda{\bm{I}}\right)^{-1}\bm{\Phi}_{0}(w)$
(81) $\displaystyle\leq
p(w)\left\|({\bm{K}}_{0}+\lambda{\bm{I}})^{-1}\right\|_{2}\left\|\bm{\Phi}_{0}(w)\right\|_{2}^{2}$
(82) $\displaystyle\leq
p(w)\frac{\left\|\bm{\Phi}_{0}(w)\right\|_{2}^{2}}{\lambda}$ (83)
$\displaystyle\leq p(w)\frac{2n}{\lambda}$ (84)

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{K}}_{0}+\lambda{\bm{I}}$ is equal to or greater
than $\lambda$. The last inequality is from that
$\left\|\mathrm{Step}(x)\right\|_{2}^{2}\leq n$ for any
$x\in{\mathbb{R}}^{n}$. Note that
$\int_{{\mathbb{R}}^{d}}p(w)\frac{2n}{\lambda}dw=\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 ??. ∎

###  E.2 First Order Arc-Cosine Kernels

######  Theorem 7.

Given dataset ${\bm{X}}\in{\mathbb{R}}^{d\times n}$, let
${\bm{K}}_{1}\in{\mathbb{R}}^{n\times n}$ be the arc-cosine kernel matrix of
$1$-th order with ${\bm{X}}$ and $v_{1},\dots,v_{m}\in{\mathbb{R}}^{d}$ be
i.i.d. random vectors from probability distribution
$q(v)=\frac{1}{(2\pi)^{d/2}d}\left\|v\right\|_{2}^{2}\exp\left(-\frac{1}{2}\left\|v\right\|_{2}^{2}\right).$
Denote
$\bm{\Phi}_{1}:=\sqrt{\frac{2d}{m}}\left[\frac{\mathrm{ReLU}({\bm{X}}^{\top}v_{1})}{\left\|v_{1}\right\|_{2}},\
\dots\
,\frac{\mathrm{ReLU}({\bm{X}}^{\top}v_{m})}{\left\|v_{m}\right\|_{2}}\right]^{\top}$
and for $\lambda\in(0,\left\|{\bm{K}}_{1}\right\|_{2})$, let $s_{\lambda}$ be
the statistical dimension of ${\bm{K}}_{1}$. Given $\varepsilon\in(0,1/2)$ and
$\delta\in(0,1)$, if
$m\geq\frac{8}{3}\frac{d}{\varepsilon^{2}}\min\left\\{\mathrm{rank}({\bm{X}})^{2},\frac{\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda}\right\\}\log\left(\frac{16s_{\lambda}}{\delta}\right)$,
then it holds that

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

with probability at least $1-\delta$.

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

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

Then,

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

where the inequality in fourth line holds from the fact that
$\left\|\mathrm{ReLU}(x)\right\|_{2}^{2}\leq\left\|x\right\|_{2}^{2}$ for any
vector $x$.

On the other hand, if we write the ReLU function in terms of a matrix form,
i.e.,
$\bm{\Phi}_{1}(w)=\sqrt{2}~{}\mathrm{ReLU}({\bm{X}}^{\top}w)=\sqrt{2}~{}{\bm{S}}{\bm{X}}^{\top}w$
where ${\bm{S}}$ is a diagonal matrix such that ${\bm{S}}_{ii}=1$ if
$[{\bm{X}}^{\top}w]_{i}>0$ else ${\bm{S}}_{ii}=0$ for $i\in[n]$, then we can
obtain that

$\displaystyle\tau_{\lambda}(w)$
$\displaystyle:=p(w)\cdot\bm{\Phi}_{1}(w)^{\top}({\bm{K}}_{1}+\lambda{\bm{I}})^{-1}\bm{\Phi}_{1}(w)$
$\displaystyle=2~{}p(w)\cdot
w^{\top}{\bm{X}}{\bm{S}}({\bm{K}}_{1}+\lambda{\bm{I}})^{-1}{\bm{S}}{\bm{X}}^{\top}w.$
(86)

Now note that by definition of ${\bm{K}}_{1}$, we have
$[{\bm{K}}_{1}]_{i,j}=\|x_{i}\|_{2}\|x_{j}\|_{2}\cdot\kappa_{1}\left(\frac{\langle
x_{i},x_{j}\rangle}{\|x_{i}\|_{2}\|x_{j}\|_{2}}\right)$. Therefore, using the
Taylor expansion of the function
$\kappa_{1}(\alpha)=\frac{1}{\pi}+\frac{\alpha}{2}+\frac{1}{\pi}\cdot\sum_{i=0}^{\infty}\frac{(2i)!}{2^{2i}\cdot(i!)^{2}\cdot(2i+1)\cdot(2i+2)}\cdot\alpha^{2i+2}$
and the fact that its Taylor coefficients are all non-negative, we have,

$\frac{1}{2}{\bm{X}}^{\top}{\bm{X}}\preceq{\bm{K}}_{1}.$

Using the above inequality along with ??, we can write,

$\displaystyle\tau_{\lambda}(w)$ $\displaystyle\leq 2~{}p(w)\cdot
w^{\top}{\bm{X}}{\bm{S}}\left(\frac{1}{2}{\bm{X}}^{\top}{\bm{X}}+\lambda{\bm{I}}\right)^{-1}{\bm{S}}{\bm{X}}^{\top}w$
$\displaystyle\leq
2~{}p(w)\cdot\left\|w\right\|_{2}^{2}\cdot\left\|{\bm{X}}{\bm{S}}\left(\frac{1}{2}{\bm{X}}^{\top}{\bm{X}}+\lambda{\bm{I}}\right)^{-1}{\bm{S}}{\bm{X}}^{\top}\right\|_{2}$
(87)

To obtain an upper bound of the third term in ??, we consider the singular
value decomposition of ${\bm{X}}={\bm{V}}\bm{\Sigma}{\bm{U}}^{\top}$. And we
have

$\displaystyle\left\|{\bm{X}}{\bm{S}}\left(\frac{1}{2}{\bm{X}}^{\top}{\bm{X}}+\lambda{\bm{I}}\right)^{-1}{\bm{S}}{\bm{X}}^{\top}\right\|_{2}$
$\displaystyle=\left\|{\bm{V}}\bm{\Sigma}{\bm{U}}^{\top}{\bm{S}}\left(\frac{1}{2}{\bm{U}}\bm{\Sigma}^{2}{\bm{U}}^{\top}+\lambda{\bm{I}}\right)^{-1}{\bm{S}}{\bm{U}}\bm{\Sigma}{\bm{V}}^{\top}\right\|_{2}$
$\displaystyle=\left\|\bm{\Sigma}{\bm{U}}^{\top}{\bm{S}}\left(\frac{1}{2}{\bm{U}}\bm{\Sigma}^{2}{\bm{U}}^{\top}+\lambda{\bm{I}}\right)^{-1}{\bm{S}}{\bm{U}}\bm{\Sigma}\right\|_{2}$
$\displaystyle=\left\|\bm{\Sigma}{\bm{U}}^{\top}{\bm{S}}{\bm{U}}\left(\frac{1}{2}\bm{\Sigma}^{2}+\lambda{\bm{I}}\right)^{-1}{\bm{U}}^{\top}{\bm{S}}{\bm{U}}\bm{\Sigma}\right\|_{2}$
(88)

Now we observe that

$\displaystyle[{\bm{U}}^{\top}{\bm{S}}{\bm{U}}]_{ij}\leq\left\|u_{i}\right\|_{2}\left\|u_{j}\right\|_{2}$
(89)

which leads us to
$\left\|{\bm{U}}^{\top}{\bm{S}}{\bm{U}}\right\|_{F}\leq\mathrm{rank}({\bm{U}})=\mathrm{rank}({\bm{X}})$.
Since ${\bm{U}}^{\top}{\bm{S}}{\bm{U}}$ is a positive semi-definite matrix we
have,

$0\preceq{\bm{U}}^{\top}{\bm{S}}{\bm{U}}\preceq\mathrm{rank}({\bm{X}})\cdot{\bm{I}}.$

Therefore, plugging this into ?? and ?? gives,

$\tau_{\lambda}(w)\leq
2~{}p(w)\cdot\left\|w\right\|_{2}^{2}\cdot\mathrm{rank}({\bm{X}})^{2}.$

Denote
$\widetilde{\tau}(w):=2~{}p(w)\left\|w\right\|_{2}^{2}\min\left\\{\mathrm{rank}({\bm{X}})^{2},\frac{\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda}\right\\}$
and it holds that

$\displaystyle\int_{{\mathbb{R}}^{d}}\widetilde{\tau}(w)dw=2d\min\left\\{\mathrm{rank}({\bm{X}})^{2},\frac{\left\|{\bm{X}}\right\|_{2}^{2}}{\lambda}\right\\}$
(90)

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

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

and recall that the modified random features are defined as

$\displaystyle\bm{\Phi}_{1}$
$\displaystyle=\frac{1}{\sqrt{m}}\left[\sqrt{\frac{p(w_{1})}{q(w_{1})}}\
\bm{\Phi}_{1}(w_{1}),\ \dots\ ,\sqrt{\frac{p(w_{m})}{q(w_{m})}}\
\bm{\Phi}_{1}(w_{m})\right]^{\top}$ (92)
$\displaystyle=\sqrt{\frac{2d}{m}}\left[\frac{\mathrm{ReLU}({\bm{X}}^{\top}w_{1})}{\left\|w_{1}\right\|_{2}},\
\dots\
,\frac{\mathrm{ReLU}({\bm{X}}^{\top}w_{m})}{\left\|w_{m}\right\|_{2}}\right]^{\top}.$
(93)

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

Approximate sampling. It is not trivial to sample a vector
$w\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 $w$ and then iteratively
replaces $[w]_{i}$ with a sample

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

for $i\in[d]$ and repeats this process for $T$ iterations. Sampling from
$q([w]_{i}|[w]_{\setminus\\{i\\}})$ can be done via the inverse transformation
method.333It requires the CDF of $q([w]_{i}|[w]_{\setminus\\{i\\}})$ which is
equivalent to
$\frac{\mathrm{erf}\left({[w]_{i}}/{\sqrt{2}}\right)+1}{2}-\frac{[w]_{i}\exp\left(-[w]_{i}^{2}/2\right)}{\sqrt{2\pi}(1+\left\|w\right\|_{2}^{2}-[w]_{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(v)$. This is
negligible compared to the feature map construction with PolySketch for
$T=\mathcal{O}(1)$. The pseudo-code for the modified random features of
$A_{1}$ using Gibbs sampling is outlined in ??.

Algorithm 3 Gibbs Sampling for ?? via Inverse Transformation Method