Fast Sketching of Polynomial Kernels of Polynomial Degree

Introduction

Kernel methods are a powerful tool for solving non-parametric learning problems, such as kernel regression, support vector machines (SVM), principal component analysis (PCA), and many others. A typical burden for kernel methods is that they suffer from scalability, since computing a kernel matrix requires computing a quadratic (in the number of input points) number of entries in the matrix. A direction that has received less attention but still of particular interest is the regime where the dimension $d$ of the data points is large. Typically, applying the kernel function to each pair of data points takes $O(d)$ time. This is especially undesirable in applications for natural language processing [DL20] and computational biology [TPK02], where $d$ can be as large as $\operatorname{poly}(n)$, with $n$ being the number of data points. To compute the kernel matrix, the algorithm does have to read the $d\times n$ input matrix. Therefore, algorithms that have a nearly linear dependence on $nd$ are of particular interest.

To accelerate the computation of kernel matrices from the naïve $O(n^{2}d)$ time algorithm, a lot of work has focused on finding a good approximation to a kernel matrix efficiently [RR07, AM15, MM17, AKK+20, WZ20]. All of these methods make use of randomized algorithmic primitives such as sampling or sketching. Roughly speaking, the idea is to randomly generate a “sketching matrix" with a small number of rows, multiply the sketching matrix with the input matrix, and show that the resulting matrix approximately preserves the length of vectors in the row or column space of the original matrix.

Does there exist a sketch for polynomial kernels of degree $p$, such that the runtime is nearly linear in $nd$, and with an improved dependence on $p$?

Notice this is especially desirable for kernels such as the neural tangent kernel ($\mathsf{NTK}$) [JGH18] and the arc-cosine kernel [CS09], whose Taylor series have a much slower decay rate ($1/n^{c}$ for some $c$) compared to the Gaussian kernel (which is $1/n!$).

We develop an efficient algorithm that computes a sketch of the polynomial kernel of degree $p$ in time linear in $p^{2}$ and nearly linear in $nd$.

Our algorithm only uses two distinct sketches compared to the $O(p)$ independent sketches of [AKK+20]. This enables us to use repeated powering to compute our sketch very efficiently.

We characterize kernel matrices by considering their Taylor series, and provide different algorithmic schemes to solve them. Our characterization includes a family of interesting and popular kernels. We also use this sketch as a preconditioner for solving linear systems involving a kernel matrix, and we extend our sketch to solve kernel ridge regression, by composing it with another sketch that depends on the statistical dimension.

Sketching techniques for tensor-related problems

Sketching techniques have been used extensively in tensor-related problems, e.g., for linear-algebraic problems involving polynomial kernels [ANW14, AKK+20, WZ20], for tensor low-rank approximation [SWZ19], and for tensor regression [HLW17, DSSW18, DJS+19].

Subspace embeddings

An (oblivious) subspace embedding is a useful concept in randomized numerical linear algebra introduced by Sárlos [Sar06]. Many applications rely on subspace embeddings or their variants, such as linear regression, low-rank approximation [CW13, NN13, MM13, BW14, BWZ16, SWZ17, ALS+18], tensor decomposition [SWZ19], cutting plane methods [JLSW20], and linear programming [LSZ19, JSWZ21, SY21]

Roadmap

In Section 2, we introduce definitions, notations and some basic facts. In Section 3, we present a technical overview of our results. In Section 4, we propose an efficient algorithm to generate a sketch and apply it to a polynomial kernel of arbitrary positive integer degree $p$. In Section 5, we analyze our algorithm with a specific sketching matrix. In Section 6, we present applications to the Gaussian kernel and a more general class of kernels, which can be characterized through the coefficients of their Taylor expansion. We also discuss how to use our sketch as a preconditioner for solving kernel linear systems, and solve sketched kernel ridge regression.

Preliminaries

We define an oblivious subspace embedding ([Sar06]) as follows:

We also introduce tensor products of vectors and Kronecker products of matrices.

The Kronecker product of matrices is a natural extension of the tensor product of vectors:

An important property of the Kronecker product is the so-called mixed product property:

For matrices $A,B,C,D$ with appropriate sizes, the following holds:

We will extensively use the following notation:

2 Sketching Matrices

We recall the Subsampled Randomized Hadamard Transform ($\mathsf{SRHT}$), which is a Fast Johnson-Lindenstrauss transform [AC06].

Using the Fast Fourier Transform (FFT) [CT65], $S$ can be applied to a vector in time $O(d\log d)$.

We also introduce a sketching matrix for degree-$2$ tensors, which is a generalization of the $\mathsf{SRHT}$.

By leveraging the FFT algorithm in the sketch space, $S(x^{\otimes 2})$ can be computed in time $O(d\log d+m)$.

We will use the following properties of the $\mathsf{SRHT}$ and $\mathsf{TensorSRHT}$.

Let $T$ be an $\mathsf{SRHT}$ matrix defined in Definition 2.7. If $m=O(n\log(nd/\delta)\epsilon^{-2})$, then $T$ is an $(\epsilon,\delta,d,n)$-$\mathsf{OSE}$.

Let $S$ be a $\mathsf{TensorSRHT}$ matrix defined in Definition 2.9. If $m=O(n\log^{3}(nd/\epsilon\delta)\epsilon^{-2})$, then $S$ is an $(\epsilon,\delta,d,n)$-$\mathsf{OSE}$ for degree-$2$ tensors.

3 Kernels

We introduce several kernels that are widely-used in practice, e.g., see [GE08, CHC+10] for the polynomial kernel, and see [NJW01] for the Gaussian kernel.

Technical Overview

We first consider one way to compute the polynomial kernel $P$ via the identity $P=(X^{\otimes p})^{\top}X^{\otimes p}$. Our algorithm will try to compute $X^{\otimes p}$ quickly.

This allows for a much faster way to compute $x^{\otimes p}$: “square” a vector by computing the tensor product with itself, then apply a sketch, and repeat this process. By doing so, we reduce the dependence on $p$ in the first level of the tree from linear to logarithmic. However, this will incur a $p^{2}$ factor in the dimension of the sketch, and so we will pay more for $p$ in levels other than the first level. Fortunately, levels other than the first level apply sketches to lower dimensional vectors. By carefully balancing the complexity of applying $T$ and $S$, we achieve an improved running time, which is useful when the degree $p$ is large.

Fast Sketching Algorithm for the Polynomial Kernel

We introduce our algorithm that sketches a single vector $x^{\otimes p}$ and extend it to each column of a matrix. In Section 4.1 we give some definitions. In Section 4.2 we prove several technical tools related to tensors. In Section 4.3 we show our sketch preserves the column space of the polynomial kernel. In Section 4.4 we prove our main result for this section.

We define the sketching matrix formed by Algorithm 1 as follows:

where T^{q}=\underbrace{T\times T\times\ldots\times T}_{\text{qtimes}} and $Q^{q}=S^{1}\cdot S^{2}\cdot S^{4}\cdot\ldots\cdot S^{q/2}$, and S^{l}=\underbrace{S\times S\times\ldots\times S}_{\text{ltimes}}.

We will design an algorithm that achieves the following goal:

Case 1 If $p$ is a power of $2$, then it computes $\Pi^{p}X^{\otimes p}$ efficiently.

Case 2 If $p$ is not a power of $2$, then let $b$ be its binary representation and let

We will iterate through all indices in $E$ and continue tensoring two vectors where $b_{i}=1$, and apply $S$ to them.

2 Equivalence Results for Tensors

We provide two technical tools for handling tensors.

If $p$ is a power of $2$, then Algorithm 1 will output $z$ on line 8. We will exploit the fact that although Algorithm 1 only computes one vector at a time, we can view it as computing $p/2^{i}$ identical vectors in the $i^{\text{th}}$ iteration. On line 3, we can treat it as computing $p$ copies of $w_{0}$, and therefore, by Claim 2.5, we have

We can apply the same line of reasoning to line 4 of the algorithm. In the $i^{\text{th}}$ iteration, we can treat it as

a total of $p/2^{i}$ times. Again, using Claim 2.5, we have

Recursively applying this identity, we will end up with

Next, we wish to show that if $p$ is a power of $2$, then $\Pi^{p}$ preserves the subspace spanned by the columns of $X^{\otimes p}$ within a factor of $1\pm\frac{\epsilon}{2p}$. Notice this is weaker than $\Pi^{p}$ being an $\mathsf{OSE}$, but sufficient for our application to polynomial kernels. We will show that

The following lemma outlines the main technique in our proof of this property. It establishes a one-to-one mapping between a vector in the column span of $X^{\otimes p}$ and a matrix. We then use this equivalence to inductively prove that $\Pi^{p}$ preserves the target subspace.

On the other hand, we can write $X^{\otimes(p-1)}$ in its column form:

Recall that $Y$ is a diagonal matrix and therefore, the product $X^{\otimes(p-1)}Y$ can be expressed as

Using the outer product definition of matrix product, we have

This is a matrix of size $d^{p-1}\times d$. Therefore, we can write its Frobenius norm as

3 Preserving the Polynomial Kernel Subspace

We prove the sketch generated by Algorithm 1 preserves the column space of the polynomial kernel.

We proceed by induction on $p$. For $p=1$, by Definition 2.1, we have

For the inductive step, we will prove this for a general positive integer $p>1$. We break $p$ into $p-1$ and $1$. By Lemma 4.4, we have

Recall that $T$ is an $(\epsilon,\delta,d,n)$ $\mathsf{OSE}$ for $X$. This means right multiplying by $T^{\top}$ preserves the length of all columns of $(TX)^{\otimes(p-1)}YX^{\top}$, and therefore we have

Applying this to each column of $YX^{\top}$, we have

As we motivated in Section 3, if we view Algorithm 1 as a binary tree, Lemma 4.5 effectively proves that the bottom layer of the tree preserves the column space of $X^{\otimes p}$. We then pick $S$ to be an $\mathsf{OSE}$ for degree-$2$ tensors, and inductively establish our embedding.

We will prove this by induction on the number of iterations of Algorithm 1. Let $k=\log_{2}p$; we will induct on the parameter $l$ from $1$ to $k$. Let $\Upsilon^{2^{l}}$ denote the sketching matrix at level $l$: $\Upsilon^{2^{l}}=S^{p/2^{l}}\cdot S^{p/2^{l-1}}\cdot\ldots S^{p/2}\cdot T^{p},\forall l\in[k]$ and $\Upsilon^{0}=T^{p}$. We will prove the following statement: $\forall l\in\{0,\ldots,k\}$, we have

Note that when $l=k$, we have $\Upsilon^{2^{k}}=\Pi^{p}$, and therefore it gives us the desired result.

For the base case, note that Lemma 4.5 automatically gives our desired result. For the inductive step, we assume it holds for some $l-1$, so

Notice that $\Upsilon^{2^{l}}=S^{p/2^{l}}\cdot\Upsilon^{2^{l-1}}$. Let $Z$ be defined as the matrix

where we use $x_{j}^{i}$ to denote the $i^{\text{th}}$ column of $X$ after the $j^{\text{th}}$ iteration. From Algorithm 1, we have that $Z=\Upsilon^{2^{l}}X^{\otimes p}$, and so the product $S^{p/2^{l}}\cdot\Upsilon^{2^{l-1}}X^{\otimes p}$ can be written as $(SZ)^{\otimes p/2^{l}}$.

If $p/2^{l}>1$, then similar to Lemma 4.5,

The third step uses the same reasoning as Lemma 4.5, i.e., we can pull out $S$ by paying an extra $(1\pm\epsilon)^{p/2^{l}-1}$ factor. The last line uses the inductive hypothesis.

If $p/2^{l}=1$, then we will end up with $SZy$, and can simply use the fact that $S$ is an $\mathsf{OSE}$ to argue that $SZy$ preserves the length of $Zy$. We then use the inductive hypothesis on $Z$ to conclude the proof. ∎

Below, we state and prove a theorem that establishes the correctness of Algorithm 1 without instantiating the sketching matrix $T$ and $S$. This enables us to use different sketches with various trade-offs.

4 Main Result

We prove the main result of this section, which establishes the correctness of Algorithm 1.

Let $b$ be the binary representation of $p$, and let $E=\left\{i:b_{i}=1,i\in\{0,1,\ldots,\log_{2}p\}\right\}$. If $p$ is a power of $2$, by Lemma 4.7, we are done. So suppose $p$ is not a power of $2$. Let $q=2^{\lfloor\log_{2}p\rfloor}$. Algorithm 1 computes $\Pi^{q}X^{\otimes q}$ and combines intermediate results with indices in $E$ to form the final result. We will again prove this by induction on the indices in $E$, from smallest to largest. For the base case, let $i_{1}$ be an index in $E$ and let $q_{1}=2^{i_{1}}$. Since $q_{1}$ is a power of $2$, Lemma 4.7 establishes this case.

For the inductive step, suppose this holds for $i_{1},i_{2},\ldots,i_{j-1}\in E$, and let $q_{1}=2^{i_{1}},q_{2}=2^{i_{2}},\ldots,q_{j-1}=2^{i_{j-1}}$. We will prove this holds for $q_{j}=2^{i_{j}}$. Let $Z$ denote the matrix after the $(j-1)^{\text{th}}$ application of this recursive process. We will show that

We first use the fact $S$ is an $\mathsf{OSE}$ to obtain

Combining Eq. (2) and Eq. (4.4), we obtain Eq. (1), which is our desired result. ∎

In this section, we analyze the runtime of Algorithm 1 with $T$ being an $\mathsf{SRHT}$ sketch (Definition 2.7) and $S$ being a $\mathsf{TensorSRHT}$ sketch (Definition 2.9).

The goal of this section is to give a runtime analysis of Algorithm 1 using $\mathsf{SRHT}$ as $T$ and $\mathsf{TensorSRHT}$ as $S$.

Moreover, using Algorithm 1, $\Pi X^{\otimes p}={\cal Z}(S,T,X)$ can be computed in time $\widetilde{O}(nd+\epsilon^{-2}n^{2}p^{2})$.

We will use an $\mathsf{SRHT}$ for $T$ and a $\mathsf{TensorSRHT}$ for $S$. We pick both of these sketches to be $(\widehat{\epsilon},\delta,d,n)$-$\mathsf{OSE}$s where $\widehat{\epsilon}=\frac{\epsilon}{3p}$. Let $Z={\cal Z}(S,T,X)$ be the matrix generated by Algorithm 1 with these parameters. By Theorem 4.8, we have

By Taylor expanding $(1+x/n)^{n}$ around $x=0$, we have

Thus, by picking $\widehat{\epsilon}=\frac{\epsilon}{3p}$, we have

For both $\mathsf{SRHT}$ and $\mathsf{TensorSRHT}$ to be $(\epsilon/3p,\delta,d,d,n)$ $\mathsf{OSE}$s, we need $m=\widetilde{\Theta}\left(n/(\epsilon/3p)^{2}\right)=\widetilde{\Theta}(p^{2}n/\epsilon^{2})$.

We now analyze the runtime of Algorithm 1 under $\mathsf{SRHT}$ and $\mathsf{TensorSRHT}$. On line 2, we compute $TX$ in time $\widetilde{O}(nd)$ since $T$ is an $\mathsf{SRHT}$. We then enter a loop with $O(\log p)$ iterations, where in each iteration we apply $S$ to the tensor product of a column with itself resulting from the previous iteration. Since $S$ is a $\mathsf{TensorSRHT}$, this takes $O(m)=\widetilde{O}(p^{2}n/\epsilon^{2})$ time per column, and there are $n$ columns, so $\widetilde{O}(p^{2}n^{2}/\epsilon^{2})$ time in total for this step. We also compute each bit in the binary representation, which incurs an $O(\log p)$ factor in the final runtime. So it takes Algorithm 1 $\widetilde{O}(nd+p^{2}n^{2}/\epsilon^{2})$ time to compute ${\cal Z}(S,T,X)$. This completes the proof. ∎

2 Discussion

We compare our result with the results obtained in [AKK+20, WZ20]. The setting we are considering is 1) matrix $X$ is dense, i.e., $\operatorname{nnz}(X)\approx nd$, and 2) $d\gg n$. In such a scenario, [AKK+20] obtains a sketching dimension $m=\Omega(\epsilon^{-2}n^{2}p)$ and the runtime of applying the sketch to $X$ is $\widetilde{O}(pnd+\epsilon^{-2}n^{3}p^{2})$, so our result improves the dependence on the $nd$ term and pays only $n^{2}$ instead of $n^{3}$ on the second term. Another result from [WZ20] has $m=\widetilde{\Theta}(\epsilon^{-2}n)$ but the time to apply sketching is $\widetilde{O}(p^{2.5}nd+\operatorname{poly}(\epsilon^{-1},p)n^{3})$, which is much worse in the leading $nd$ term, compared to our result. However, we also point out the results obtained in these two works are more general than ours in the sense that their sketches have the $\mathsf{OSE}$ property, while our sketch only preserves the column space of $X^{\otimes p}$. Nevertheless, the latter suffices for our applications. We use Table 1 to summarize and compare the different results.

Also, the prior results mentioned are stated in terms of the statistical dimension, while we do not directly obtain bounds in terms of the statistical dimension, though our sketches can be composed with sketches that do. Therefore, we consider the case when there is no regularization $(\lambda=0)$ and $X^{\otimes p}$ is of full rank. In this case, the statistical dimension reduces to $n$.

Applications

In this section, we introduce various applications using our sketch. In Section 6.1, we study approximating the Gaussian kernel using our algorithm. In Section 6.2 we extend the analysis to a class of slow-decaying kernels. In Section 6.3 we illustrate an efficient algorithm to solve kernel linear systems. In Section 6.4 we show how to solve kernel ridge regression using our sketch.

We provide the fastest algorithm to preserve the column space of a Gaussian kernel when $d$ is large.

where $m=\widetilde{\Theta}(q^{3}n/\epsilon^{2})$ and $q=\Theta(r^{2}+\log(n/\epsilon))$.

We provide a sketch of the proof here, and further details can be found in the appendix. The Taylor expansion of the Gaussian kernel can be written as

where $D$ is a diagonal matrix with $D_{i,i}=\exp(-\|x_{i}\|_{2}^{2}/2)$. Let

If we set $q=\Omega(r^{2}+\log(n/\epsilon))$ and just use the first $q$ terms of $K$:

2 General p𝑝p-convergent Kernels

A key advantage of Algorithm 1 is its moderate dependence on the degree $p$, which gives it more leverage when $p$ is large. We introduce a characterization of kernels, based on the series of the coefficients in the Taylor expansion of the kernel. As we will later see in the proof of Theorem B.2, the decay rate of coefficients has a direct relation with the degree $p$ we need for approximating a kernel.

We say the kernel matrix $K$ for data matrix $X$ is $p$-convergent if its corresponding Taylor expansion series can be written as follows: $K=\sum_{l=0}^{\infty}C_{l}\cdot(X^{\otimes l})^{\top}X^{\otimes l}$, where the coefficients $C_{l}=(l+1)^{-\Theta(p)}$.

For the sake of illustration, suppose $r=1$. Then the first term in the running time becomes $\epsilon^{-2-\frac{3}{p}}n^{2+\frac{3}{p}}$. When $p$ is large, Theorem 6.3 gives a fast algorithm for approximating the kernel, but the runtime becomes much slower when $p\in(1,3)$. Therefore, we propose a novel sampling scheme to deal with small values of $p$. Roughly speaking, we exactly compute the first $s$ terms in the Taylor expansion, while for the remaining $q-s$ terms we sample only $s$ of them proportional to their coefficient. Using a matrix Bernstein bound (Theorem B.4), we obtain an even faster algorithm. We apply our result to the neural tangent kernel ($\mathsf{NTK}$), which is a $1.5$-convergent kernel.

We remark that our definition of $p$-convergent kernels captures a wide range of kernels that have slow decay rate in their coefficients in their Taylor expansion, such as $\mathsf{NTK}$ and arc-cosine kernels. Typically, the coefficients are of the form $1/n^{c}$ for some $c>1$. In contrast, Gaussian kernels enjoy a much faster decay rate, and therefore, designing algorithm for the Gaussian kernel is considerably simpler, since the number of terms we need to approximate it with in its Taylor expansion is small, and no sampling is necessary.

3 Kernel Linear Systems

Another interesting application of our sketching scheme is to constructing a preconditioner for solving PSD systems involving a kernel matrix [COCF16]. In order to apply algorithms such as Conjugate Gradient [She94], one has to obtain a good preconditioner for a potentially ill-conditioned kernel system.

in $\widetilde{O}\left(\epsilon^{-2}n^{2}\log(\kappa/\epsilon)+n^{\omega}+nd\right)$ time, where $\omega$ is the exponent of matrix multiplication (currently $\omega\approx 2.373$ [Wil12, LG14]).

In certain NLP [DL20] and biological tasks [TPK02] where $d=n^{c}$ for a positive integer $c$, Theorem 6.6 provides a fast algorithm for which the running time depends nearly linearly on $nd$. We also remark that the algorithm we use for Theorem 6.6 is inspired by the idea of [BPSW21] (their situation involves $c=4$). It is also interesting that in their applications, regularization is not needed since solving a kernel system is equivalent to training an over-parametrized ReLU network without regularization.

4 Kernel Ridge Regression

for $\lambda>0$. A relevant notion is the statistical dimension:

One drawback of our sketch is that we cannot obtain a dimension depending on $s_{\lambda}(K)$ instead of $n$, since it does not have the approximate matrix product property. To mitigate this effect, we propose the following composition of sketches.

Acknowledgments: D. Woodruff would like to thank partial support from NSF grant No. CCF-1815840, Office of Naval Research grant N00014-18-1-256, and a Simons Investigator Award.

References

Appendix

In Section A, we give an algorithm to compute a subspace embedding for the Gaussian kernel using Theorem 4.8. In Section B, we characterize a large class of kernels based on the coefficients in their Taylor expansion, and develop fast algorithms for different scenarios. In Section C, we apply our results in Section B to the Neural Tangent kernel. In Section D, we use our sketch in conjunction with another sketch to compute a good preconditioner for the Gaussian kernel. In Section E, we compose our sketch with our sketching matrices to solve Kernel Ridge Regression.

Notation

We use $\widetilde{O}(f)$ to denote $f\operatorname{poly}(\log f)$ and use $\widetilde{\Omega}(f)$ to denote $f/\operatorname{poly}(\log f)$.

Appendix A Gaussian Kernel

We remark that our construction of a sketch for the Gaussian kernel and its corresponding analysis is inspired by [AKK+20], and is thus similar to the proof of Theorem 5 in their paper. For completeness, we include a proof here.

where $m=\Omega(\epsilon^{-2}nq^{3}\log^{3}(nd/\epsilon\delta))$ and $q=\Theta(r^{2}+\log(n/\epsilon))$.

Let $q=C\cdot(r^{2}+\log(n/\epsilon))$ for a sufficiently large constant $C$, and let $Q=\sum_{l=0}^{q}\frac{(X^{\otimes l})^{\top}X^{\otimes l}}{l!}$ be the first $q$ terms of $K$. By the triangle inequality we have:

with probability at least $1-\frac{\delta}{q+1}$. Moreover, $Z_{l}$ can be computed in time

Our algorithm will simply compute $Z_{l}$ from $l=0$ to $q$, normalize each $Z_{l}$ by $\frac{1}{\sqrt{l!}}$, and then multiply by $D$. More precisely, the approximation $W_{g}(X)$ will be

By combining terms in (4) and using a union bound over all $0\leq l\leq q$, we obtain that with probability at least $1-\delta$, we have the following:

Also, by Theorem 5.1, the time to compute $W_{g}(X)$ is

Notice we will have to pay an additive $nd\log(nd/\epsilon\delta)$ due to line 2 of Algorithm 1, when applying the SRHT to $X$. However, we only need to perform this operation once for the term with the highest degree, or the terms with lower degree that can be formed by combining nodes computed with the highest degree. Thus, the final runtime is

Appendix B General p𝑝p-Convergent Sequences

We consider general $p$-convergent kernels defined below in Definition B.1. We apply our proposed Algorithm 1 to compute a subspace embedding with a fast running time.

In this section, we state a general theorem for $p>1$. The proof is similar to the proof for Theorem A.1. We start by restating the definition of a $p$-convergent kernel.

where the positive coefficients $C_{l}>0$ are a function of $l$, and $C_{l}$ satisfies

Similar to the Gaussian kernel, here we use the first $q$ terms to approximate the kernel matrix $K$.

Let $q=C\cdot(r^{2}+(n/\epsilon)^{1/p})$ for a sufficiently large constant $C$, and let $Q=\sum_{l=0}^{q}C_{l}(X^{\otimes l})^{\top}X^{\otimes l}$ be the first $q$ terms of $K$. By the triangle inequality, we have

The proof is identical to the proof of Theorem A.1, with the target dimension of $W_{g}$ being $m=m_{0}+m_{1}+\cdots+m_{q}=\Omega(\epsilon^{-2}nq^{3}\log(nd/\delta\epsilon)\log(n/\delta))$.

Similar to Theorem A.1, we have to pay an extra $nd\log(nd/\epsilon\delta)$ term to apply the $\mathsf{SRHT}$ to $X$, so the final running time is

Recall our setting is when $d=\operatorname{poly}(n)$, so if $p\geq 3$, Theorem B.2 gives a running time of $\widetilde{O}(n^{3}/\epsilon^{2}+nd)$, which is better than the classical result of $O(n^{2}d)$ as long as $d>n/\epsilon^{2}$. However, if $p\in(1,3)$, Theorem B.2 gives a worse dependence on $n$, which can be further optimized.

B.2 Sampling Scheme for 1<p<31𝑝31<p<3

We next describe a novel sampling scheme if $p\in(2,3)$, with a better dependence on $n$ compared to Theorem B.2. We first state some probability tools.

Let $S_{1},\ldots,S_{n}$ be independent, zero-mean random matrices with common size $d_{1}\times d_{2}$, and assume each one is uniformly bounded:

Let $q$ be the degree used in Theorem B.2 where $q=\Theta((n/\epsilon)^{1/p})$, and let $s$ be some positive integer smaller than $q$. We will consider the following scheme:

For the first $s$ terms in the Taylor expansion, we approximate each term directly using Theorem B.2.

For each of the next $q-s$ terms, we sample proportional to their coefficient $C_{l}$, taking only $s$ samples in total.

We now consider the operator norm of the expectation of $S_{i}^{\top}S_{i}$:

Let $Z=\sum_{i=1}^{m}S_{i}$. Since each sample is sampled independently, we have

Picking $m=\Theta(\epsilon^{-2}n^{2}s^{-2p}\log(n/\delta))$, and then averaging over $m$ samples, we get that

where we use the fact that the operator norm of $K$ is at least $1$, by our choice of $s$. We now compute the expected running time of this algorithm.

Runtime part 1: Computing the first s𝑠s terms

For the first $s$ terms, we can apply the same reasoning as in Theorem B.2 to get a running time of $O(\epsilon^{-2}n^{2}s^{3}\operatorname{poly}(\log(nd/\epsilon\delta))+nd\log(nd/\epsilon\delta))$.

Runtime part 2: Sampling the next s𝑠s terms

For the sampling part, we consider the expected degree $D$ of the sample we will be working with:

Now we are ready to compute the expected running time of the sampling phase:

Additionally, we need to apply the $\mathsf{SRHT}$ to $X$ at most twice, once for the initial phase, and once for the sampling phase, so the final running time is

∎ When $p\in(1,2]$, we use the largest degree $q$ as an upper bound for analyzing our running time.

In addition, using that $p\in(1,2]$, the first part of our running time can be upper bounded by

The proof is almost identical to the proof of Theorem B.5. The only difference is when considering the expected degree $D$, we use $q$ as an upper bound. The number of terms $s$ we approximate in the initial phase will be

Simplifying the Exponent For the exponent of $\epsilon$, we have

Appendix C Properties of the Neural Tangent Kernel

We discuss an application of our sampling algorithm for $p\in(1,2]$ (Corollary B.6) to the Neural Tangent Kernel ($\mathsf{NTK}$). We will first formally define the $\mathsf{NTK}$, then consider its Taylor expansion, and then use a $p$-convergent kernel to bound it.

In this section, we give the Taylor expansion of the $\mathsf{NTK}$, by first examining its corresponding function in a single variable, and then extend it to the matrix case.

Consider a simple two-layer (an alternative name is one-hidden-layer) ReLU network with input layer initialized to standard Gaussians, activation function ReLU, and output layer initialized to uniform and independent Rademacher ($\{-1,1\}$) random variables. Suppose we fix the output layer. Then the neural network can be characterized by a function

The above formulation is standard for convergence analysis of neural networks [LL18, DZPS19, AZLS19b, AZLS19a, SY19, BPSW21, HLSY21, MOSW21].

For the sake of simplicity, assume all $|a_{r}|=1$, $\forall r\in[m]$, and consider an individual summand, which gives rise to

The Taylor expansion of the $\mathsf{NTK}$ is

C.2 Approximating the 𝖭𝖳𝖪𝖭𝖳𝖪\mathsf{NTK}

In this section, we will use a $p$-convergent kernel to bound the $\mathsf{NTK}$, then apply Corollary B.6 to approximate it.

Let $C_{l}$ denote the coefficient of the $l^{th}$ term in the Taylor expansion of the $\mathsf{NTK}$:

The term $\binom{2l}{l}$ is the central binomial coefficient. We will use the following bound on it:

This gives upper and lower bounds on $C_{l}$:

Thus, $C_{l}=\Theta(\frac{1}{l^{1.5}})$, and we can use a $1.5$-convergent kernel for our approximation. Using Corollary B.6 with $p=1.5$, we obtain an $\epsilon$-approximation in time

Appendix D Preconditioning to Solve a Kernel Linear System

In our result, we follow a similar approach as in [BPSW21] and the proof is similar to the proof of Lemma 4.2 in their paper. The main novelty of our framework is that we use a spectral approximation to the kernel matrix and analyze the error and runtime under our approximation. For completeness, we include a proof in this setting.

Moreover, $\widehat{x}$ can be computed in time

where $\omega$ is the matrix multiplication exponent.

Before the proof, we define some notation and corresponding facts specifically about a PSD matrix.

Let $A,B$ be conforming square matrices. Then the following inequality holds:

where $\kappa(A)=\frac{\sigma_{\max}(A)}{\sigma_{\min}(A)}$ is the condition number of $A$.

We will make use of Lemma B.2 in [BPSW21].

Suppose $B$ is a PSD matrix for which $\frac{3}{4}\leq\|Bx\|_{2}\leq\frac{5}{4}$ holds for all $\|x\|_{2}=1$. Using gradient descent for $t$ iterations, we obtain

where $x_{0}$ is our initial guess, $x^{*}$ is the optimal solution, and $c\in(0,0.9]$.

Throughout the proof, we will set $\widehat{\epsilon}=\epsilon/4$. By Theorem A.1, we can compute an $\epsilon$-approximation to $Z$ and $W_{g}(X)$ in time

with solution $\widehat{x}$, then we have

Suppose $R$ is the $n\times n$ matrix computed via a QR decomposition, so that $SW_{g}(X)R$ has orthonormal columns. Then for any $\|x\|_{2}=1$, we have

Now, pick $\epsilon_{0}=0.1$ and solve the following regression problem:

Notice that Algorithm 2 implements gradient descent. Using Lemma D.3, after $t=\log(1/\widehat{\epsilon})$ iterations, we have

where $z^{*}=(R^{\top}W_{g}(X)^{\top}W_{g}(X)R)^{-1}R^{\top}y$ is the optimal solution to Equation (7). We will show the following for $x_{t}=Rz_{t}$:

Recalling that $z_{0}=0$, plugging into Eq. (8) we get

The second inequality uses that $R$ is a square matrix, the third inequality uses Fact D.2, and the second-to-last inequality uses that we have a $(1\pm\widehat{\epsilon})$-subspace embedding.

This means by setting the number of iterations to $t=\log(\kappa/\epsilon)$, we obtain

Computing $W_{g}(X)$, by Theorem A.1, takes time

Applying $S$ to $W_{g}(X)$, using the FFT algorithm, takes time

A QR decomposition algorithm, due to [DDH07], can be computed in time $n^{\omega}$.

The cost of each iteration is bounded by the cost of taking a matrix-vector product, which is at most $\widetilde{O}(n^{2}/\epsilon^{2})$, and there are $O(\log(\kappa/\epsilon))$ iterations in total. Thus, we obtain a final runtime of

Appendix E Kernel Ridge Regression

In this section, we show how to compose our sketch with other sketches whose dimensions depend on the statistical dimension of $K$ instead of $n$. Before proceeding, we introduce the notion of the statistical dimension.

Solving ridge regression with runtime depending on thhe statistical dimension is done in a number of works, for example [RR07, AM15, AKM+17, ACW17a, MM17].

We state and prove our main result in this section below.

Moreover, there exists a matrix $S$ with $m=\widetilde{O}(\epsilon^{-1}s_{\lambda}(K))$ rows such that if $x^{*}$ is the optimal solution to $\|S(Z^{\top}Zx-y)\|_{2}^{2}+\lambda\|Zx\|_{2}^{2}$, then

Finally, The time to solve above $\mathsf{KRR}$ is $\widetilde{O}(\epsilon^{-2}p^{2}n(n+m^{2})+n^{\omega})$.

Before starting the proof, we introduce a key lemma regarding using the $\mathsf{SRHT}$ to approximate the solution of $\mathsf{KRR}$.

Throughout the proof, we assume $K$ has full rank and set $S$ to be a $\mathsf{SRHT}$ matrix with $m=\widetilde{O}(\epsilon^{-1}s_{\lambda}(K))$ rows. We also use $A$ to denote $Z^{\top}Z$.

Part 1: Provide a construction of matrix $Z$;

Part 2: Provide a sketching matrix $S$ with the solution guarantee;

Part 3: Provide a runtime analysis for solving $\mathsf{KRR}$.

Note that part 1 can be solved using Theorem 5.1. As a side note, since $Z^{\top}Z$ is a $1\pm\epsilon$ approximation to $K$, with high probability it also has full rank. Consequently, $Z$ has full rank as well.

To show part 2, we will show the following:

The optimal solution to $\|S(Z^{\top}Zx-y)\|_{2}^{2}+\lambda\|Zx\|_{2}^{2}$ is a $(1\pm\epsilon)$ approximation to the optimum of $\|Z^{\top}Zx-y\|_{2}^{2}+\lambda\|Zx\|_{2}^{2}$;

The optimum of $\|Z^{\top}Zx-b\|_{2}^{2}+\lambda\|Zx\|_{2}^{2}$ is a $(1\pm\epsilon)$ approximation to the optimum of $\|Kx-y\|_{2}^{2}+\lambda\|X^{\otimes p}x\|_{2}^{2}$.

which can be solved using Lemma E.3. The only thing we need to justify is that the statistical dimension of $A$ gives a good approximation to the statistical dimension of $K$. Note that

Thus, the dimension $O(\epsilon^{-1}s_{\lambda}(K))=O(\epsilon^{-1}s_{\lambda}(A))$, which means we can invoke Lemma E.3.

For the final part, note that applying the sketch takes $\widetilde{O}(\epsilon^{-2}p^{2}n^{2})$ time. To solve the regression problem, we instead solve:

Since $Z$ has full rank, we know the argument $z$ realizing the minimum is the $x^{*}$ we are looking for. To output an $x$, we can simply solve the linear system $Zx=z$, which takes $\widetilde{O}(nt+n^{\omega})$ time. Finally, solving the above regression problem takes $\widetilde{O}(m^{2}t)$ time (see [SGV98]). This concludes our runtime analysis. ∎