Faster Kernel Ridge Regression Using Sketching and Preconditioning

Introduction

where $c_{1},\dots,c_{n}$ can be found by solving the linear equation

Compared to Kernel Support Vector Machines (KSVM), the computations involved in KRR are conceptually much simpler: solving a single linear system as opposed to solving a convex quadratic optimization problem. However, KRR has been observed experimentally to often perform just as well as KSVM . In this paper, we exploit the conceptual simplicity of KRR, and using advanced techniques in numerical linear algebra design an efficient method for solving (1).

One striking feature of the aforementioned papers is the use of a very large number of random features. Random feature maps typically provide only crude approximations to the kernel function, so to approach the full capacity of exact kernel learning (which is required to obtain state-of-the-art results) many random features are necessary. Nevertheless, even with a very large number of random features, we sometimes pay a price in terms of generalization performance. Indeed, in section 6.1 we show that in some cases, driving $s$ to be as large as $n$ is not sufficient to achieve the same test error rate as that of the full kernel method. Ultimately, methods that use approximations compromise in terms of performance in order to make the computation tractable.

We propose to use random feature maps as a means of forming a preconditioner for the kernel matrix. This preconditioner can be used to solve (1) to high accuracy using an iterative method. Thus, while training time still benefits from the use of high-quality random feature maps, there is no compromise in terms of modeling capabilities, modulo the decision to use kernel ridge regression and not some other learning method.

We provide a theoretical analysis that shows that at least for one kernel selection, the polynomial kernel, selecting the number of random features to be proportional to the so-called statistical dimension of the problem (this classical quantity is also frequently referred to as the effective degrees-of-freedom) yields a high-quality preconditioner in the sense that the relevant condition number is bounded by a constant. These can be viewed as a generalization of recent results on sharper bounds for linear regression and low-rank approximation with regularization . In addition, we discuss a method for testing whether the preconditioner computed by our algorithm is indeed such that the relevant condition number is bounded. While our analytical results are mostly of theoretical value (e.g. they are limited to only the polynomial kernel and are likely very pessimistic), they do expose an important connection between the statistical dimension and preconditioner size.

Finally, we report experimental results with a distributed-memory parallel implementation of our algorithm. We demonstrate the effectiveness of our algorithm for training high-quality models using both the Gaussian and polynomial kernel on datasets as large as one million training examples without compromising in terms of statistical capacity. For example, on one dataset with one million examples our code is able to solve (1) to relatively high accuracy in about an hour on resources readily available to researchers and practitioners (a cluster of EC2 instances).

An open-source implementation of the algorithm is available through the libSkylark library (http://xdata-skylark.github.io/libskylark/).

2 Related Work

Devising scalable methods for kernel methods has long been an active research topic. In terms of approximations, one dominant line of work constructs low-rank approximations of the Gram matrix. Two popular variants of this approach, as mentioned earlier, are the randomized feature maps and the Nyström methods . There are many variants of these approximation schemes, and it is outside the scope of this paper to mention all of them. Recent work has also focused on devising scalable methods that are capable of utilizing rather high rank approximations . In contrast, our goal is to develop a method which is capable of using lower rank approximation without paying a price in terms of model quality.

Another approach is to approximate the kernel matrix so that it will be more amenable to matrix-vector products or linear system solution. One idea is to use the Fast Gauss Transform to accelerate the matrix-vector products of the kernel matrix by an arbitrary vector . Another approach is to use a tree code to efficiently perform matrix-vector products . Related is also the hierarchical matrix approach in which an hierarchical matrix approximation to the kernel matrix is built . This representation is amenable to efficient implementation of wide range of operations on the approximate kernel matrix, including matrix-vector product and linear system solution.

The preconditioning approach has also been explored in the literature. Srinivasan et al. propose to use a regularized kernel matrix as a preconditioner in a flexible Krylov method . The regularized kernel matrix, which has a lower condition number, is solved using an inner conjugate gradient iteration. In parallel work to ours, Cutajar et al. recently discussed various preconditioning techniques for kernel matrices . One of the methods they propose is using random features to form a preconditioner. However, unlike our work, they do not include any theoretical analysis of this preconditioning approach. Furthermore, we propose additional algorithmic enhancements (multiple level preconditioning, testing preconditioners). Finally, it is worth mentioning that Cutajar et al. only experiment with small-scale low-dimensional datasets, while we present experimental results with large-scale high-dimensional datasets.

For a broad discussion of scalable methods for kernel learning, including of ideas not mentioned here, see Bottou et al. .

Preliminaries

2 Random Feature Maps

Throughout the paper, we use $s$ to denote the number of random features. This quantity is a parameter of the various algorithms.

where $p$ is some appropriate distribution that depends on the kernel function (e.g. Gaussian distribution for the Gaussian kernel). The existence of such a $p$ for every shift-invariant kernel function $k$ is a consequence of Bochner’s Theorem; see Rahimi and Recht for details. The feature map is then a Monte-Carlo sample: $\varphi({\mathbf{x}})=s^{-1/2}[\cos({\mathbf{w}}^{\textsc{T}}_{1}{\mathbf{x}}+b_{1})\dots\cos({\mathbf{w}}^{\textsc{T}}_{s}{\mathbf{x}}+b_{n})]^{\textsc{T}}$ where ${\mathbf{w}}_{1},\dots,{\mathbf{w}}_{s}$ are sampled from $p$ and $b_{1},\dots,b_{n}$ are sampled from a uniform distribution on $[0,2\pi]$.

A crucial observation that makes this transformation useful is that via a clever application of the Fast Fourier Transform, $\varphi({\mathbf{x}})$ can be computed in $O(q({\bf nnz}\left({\mathbf{x}}\right)+s\log{s}))$ (see Pagh for details), which allows for a fast application of the transform.

3 Fast Numerical Linear Algebra Using Sketching

Sketching has recently emerged as a powerful dimensionality reduction technique for accelerating numerical linear algebra primitives typically encountered in statistical learning such as linear regression, low rank approximation, and principal component analysis. The following description is only a brief semi-formal description of this emerging area. We refer the interested reader to a recent surveys for more information.

The underlying idea is to construct an embedding of a high-dimensional space into a lower-dimensional one, and use this to accelerate the computation. For example, consider the classical linear regression problem

The “sketch-and-solve” approach just described allows only crude approximations: the $\epsilon$-dependence for OSEs is $\epsilon^{-2}$. There is also an alternative “sketch-to-precondition” approach, which enjoy a much better $\log(1/\epsilon)$ dependence for $\epsilon$ and so supports very high accuracy approximations.

4 Random Feature Maps as Sketching

Thus, random feature maps can be viewed as a sketch that embeds

Viewed this way, the random features method is a “sketch-and-solve” approach. It is therefore not surprising that it produces suboptimal models. In this paper, we take the “sketch-to-precondition” approach to utilizing sketching.

Random Features Preconditioning

2 Analysis

We now analyze the algorithm when applied to kernel ridge regression with the polynomial kernel and TensorSketch as the feature map. In particular, in the following theorem we show that if $s$ is large enough then PCG will converge in $O(1)$ iterations (for a fixed convergence threshold).

with probability of at least $1-\delta$, after

In the above, ${\mathbf{c}}$ is the exact solution of the linear equation at hand (Equation 1).

The result is stated for the homogeneous polynomial kernel $k({\mathbf{x}},{\mathbf{z}})=({\mathbf{x}}^{\textsc{T}}{\mathbf{z}})^{q}$, but it can be easily generalized to the non-homogeneous case $k({\mathbf{x}},{\mathbf{z}})=({\mathbf{x}}^{\textsc{T}}{\mathbf{z}}+c)^{q}$ by adding a constant feature to each training point.

While the theorem gives an explicit formula for the number of iterations, we do not recommend to actually use this formula, and recommend instead the use of standard stopping criteria to declare convergence (these usually involve determining that the residual norm has dropped below some tolerance). The reason is that the iteration bound holds only with high probability. On the other hand, a higher probability bound bound holds when we consider a higher bound on the number of iterations. Thus, using standard stopping criteria renders the algorithm more robust. Furthermore, the bound holds only under exact arithmetic, while in practice PCG is used with inexact arithmetic.

Estimating the statistical dimension is a non-trivial task that is outside the scope of this paper (and can sometimes be avoided: see Section 5). Nevertheless, the bound does establish that the number of random features required for a constant number of iterations depends on the statistical dimension, which is always smaller than the number of training points. Since the kernel matrices often display quick decay in eigenvalues, it can be substantially smaller. In particular, the number of random features can be $o(n)$ when the statistical dimension is $o(n^{1/2})$. A discussion on how the statistical dimension behaves with regard to the training size is outside the scope of this paper. We refer the reader to a recent discussion by Bach on the subject .

Before proving the theorem we state some auxiliary definitions and lemmas.

In addition, we need the following lemma.

We prove that with probability of at least $1-\delta$

Thus, with probability of $1-\delta$ the relevant condition number is bounded by $3$. For PCG, if the condition number is bounded by $\kappa$, we are guaranteed to reduce the error (measured in the matrix norm of the linear equation) to an $\epsilon$ fraction of the initial guess after $\lceil\sqrt{\kappa}\ln(2/\epsilon)/2\rceil$ iterations . This immediately leads to the bound in the theorem statement.

A sufficient condition for (6) to hold is that

3 Other Kernels and Feature Maps

Close inspection of the proof reveals that the crucial ingredient is the matrix multiplication lemma (Lemma 4). In the following, we generalize Theorem 1 to feature maps which have similar structural properties. The proof, which is mostly analogous to the proof of Theorem 1, is included in the Appendix.

A linear feature map has an approximate multiplication property with $f(\nu,\delta)$ if $\varphi$ with at least $f(\nu,\xi,\delta)$ random features has that for all finite ordered sets ${\cal U},{\cal V}\subset{\cal H}_{k}$ the following holds with probability of at least $1-\delta$:

Currently, there is no proof that the approximate multiplication property holds for any feature map except for TensorSketch.

Multiple Level Sketching

We now show that the dependence on the statistical dimension can be improved by composing multiple sketching transforms. The crucial observation is that after the initial random feature transform, the training set is embedded in a Euclidean space. This suggests the composition of well-known transforms such as the Subsampled Randomized Hadamard Transform (SRHT) and the Johnson-Lindenstrauss transform with the initial random feature map. A similar idea, referred to as compact random features, was explored in the context of the random features method by Hamid et al. .

First, we consider the use of the SRHT after the initial TensorSketch for the polynomial kernel. To that end we recall the definition of the SRHT. Let $m$ be a power of 2. The $m\times m$ matrix of the Walsh-Hadamard Transform (WHT) is defined recursively as,

Let $m$ be some integer which is a power of 2, and $s$ an integer. A Subsampled Randomized Walsh-Hadamard Transform (SRHT) is an $s\times m$ matrix of the form

In the above, ${\mathbf{c}}$ is the exact solution of the linear equation at hand (Equation 1).

it suffices to prove that each of the two terms is bounded by $\frac{1}{4}$ with probability $1-\delta/2$. Due to the lower bound on $s_{1}$ this holds for the right term as explained in the proof of Theorem 1.

From a computational complexity point of view, in many cases one can set $s_{1}$ to be very large without increasing the asymptotic cost of the algorithm. For example, for random Fourier features, if $\sum^{n}_{i=1}{\bf nnz}\left({\mathbf{x}}_{i}\right)=O(n\log{s_{2}})$ then it is possible to set $s_{1}=\Theta(n)$ without increasing the asymptotic cost of the algorithm.

Adaptively Setting the Sketch Size

We show that there exists constants $m$ and $M$ such that

This hold if there is a constant $c<1$ such that for all ${\mathbf{y}}$

Without loss of generality we restrict ourselves to $\|{\mathbf{y}}\|_{2}=1$. We have,

Experiments

In this section we report experimental results with an implementation of our proposed one-level sketching algorithm. To allow the code to scale to datasets of size one million examples and beyond, we designed our code to use distributed memory parallelism using MPI. For most distributed matrix operations we rely on the Elemental library . We experiment on clusters composed on Amazon Web Services EC2 instances.

In this section we compare our method to the random features method, as defined in Section 2.2. Our goal is to demonstrate that our method, which solves the non-approximate kernel problem to relatively high accuracy, is able to fully leverage the data and deliver better generalization results than is possible using the the random features method.

We conducted experiments on the MNIST dataset. We use both the Gaussian kernel (with $\sigma=8.5$) and the polynomial kernel $k({\mathbf{x}},{\mathbf{z}})=(0.01{\mathbf{x}}^{\textsc{T}}{\mathbf{z}}+1)^{3}$. The regularization parameter was set to $\lambda=0.01$. Running time were measured on a single c4.8xlarge EC2 instance.

Results are shown in Figure 1. Inspecting the error rates (left plots), we see that while the random features method is able to deliver close to optimal error rates, there is a clear gap between the errors obtained using our method and the ones obtained by the random features method even for a very large number of random features. The gap persists even if we set the number of random features to be as large as the training size!

In terms of the running time (right plots), our method is, as expected, generally more expensive than the sketch-and-solve approach of the random features method, at least when the number of random features is the same. However, in order to get close to the performance of our method, the random features method requires so many features that its running time eventually surpasses that of our method with the optimal number of features (recall that for our method the number of random features only affects the running time, not the generalization performance.)

It is worth noting that the optimal training time for our method on MNIST was actually quite small: less than 2 minutes for the Gaussian kernel, and less than 4 minutes for the polynomial kernel, and this without sacrificing in terms of generalization performance by using an approximation.

In Figure 2 we further explore the complex interaction between algorithmic complexity, quantity of data and predicative performance. In this set of experiments we use subsamples of the COVTYPE dataset (to a maximum sample of 80% of the data; the rest is used for testing and validation) and subsamples from the extended MNIST-8M dataset (to a maximum of 450K data points). We compare the performance of our method as the number of samples increases to the performance of the random features method with three different profiles for setting $s$: $s=20000$ (an $O(n)$ training algorithm), $s=2.26\sqrt{nd}$ (an $O(n^{2}d)$ algorithm) and $s=0.2n$ (an $O(n^{3})$ time, $O(n^{2})$ memory algorithm). For both datasets we plot both error as a function of the datasize and error as a function of training time. The graphs clearly demonstrate the superiority of our method.

2 Resources and Running Time on a Cloud Service

Our implementation is designed to leverage distributed processing using clusters of machines. The wide availability of cloud-based computing services provides easy access to such platforms on an on-demand basis.

To give an idea of the running time and resources (and as a consequence the cost) of training models using our algorithm on public cloud services, we applied our method to various popular datasets on EC2. The results are summarized in Table 2. We can observe that using our method it is possible to train high-quality models on datasets as large as one million data points in a few hours. We remark that while we did tune $\sigma$ and $\lambda$ somewhat, we did not attempt to tune them to the best possible values.

3 Additional Experimental Results

In Table 4 we compare our method to an high-performance distributed block ADMM-based solver using the Random Features Method . We use the same resource configuration as in Table 2 (we remark that the ADMM solver is rather memory efficient and can function with less resources). The ADMM solver is more versatile in the choice of objective function, so we use hinge-loss (SVM). We use the same bandwidth ($\sigma$) and reguarization parameter ($\lambda$) and in Table 2. In general we set the number of random features ($s$) to be equal 25% of the dataset size. We clearly see that our method achieves better error rates, usually with better running times.

In Table 4 we examine whether the preconditioner indeed improves convergence and running time. We set the maximum iterations to 1000, and declare failure if failed to converge to $10^{-3}$ tolerance for classification and $10^{-5}$ for regression. We remark the following on the items labeled FAIL:

For MNIST-200k, without preconditioning CG failed to convergence but the error rate of the final model was just as good as our method.

For MNIST-300K the error rate of the final model deteriorated to 1.33% (compare to 0.92%).

For YEARMSD the final residual was $6.63\times 10^{-4}$ and the error deteriorated to $5.25\times 10^{-3}$ (compare to $4.58\times 10^{-3}$).

Almost always (with two exception, one of them a tiny dataset) our method was faster than the non-preconditioned method. More importantly our method is much more robust: the non-preconditioned algorithm failed in some cases.

In Table 5 we examine how classification generalization quality is affected by choosing a more relaxed tolerance criteria. In general, setting tolerance to $10^{-3}$ is just barely better than $10^{-2}$ in terms of test error. Only in one case (MNIST) the difference is bigger than 0.01%. Running time for $10^{-3}$ by is worse by a small factor and results are almost the same. We set the tolerance to $10^{-3}$ to be consistent with the notion of exploiting data to the fullest, although in practice $10^{-2}$ seems to be sufficient.

Conclusions

Kernel ridge regression is a powerful non-parametric technique whose solution has a closed form that involves the solution of a linear system, and thus is amenable to applying advanced numerical linear algebra techniques. A naive method for solving this system is too expensive to be realistic beyond “small data”. In this paper we propose an algorithm that solves this linear system to high accuracy using a combination of sketching and preconditioning. Under certain conditions, the running time of our algorithm is somewhere between $O(n^{2})$ and $O(n^{3})$, depending on properties of the data, kernel, and feature map. Empirically, it often behaves like $O(n^{2})$. As we show experimentally, our algorithm is highly effective on datasets with as many as one million training examples.

Obviously the main limitation of our algorithm is the $\Theta(n^{2})$ memory requirement for storing the kernel matrices. There are a few ways in which our algorithm can be leveraged to allow learning on much larger datasets. One non-algorithmic software-based idea is to use an out-of-core algorithm, i.e. use SSD storage, or even magnetic drive, to hold the kernel matrix. From an algorithmic perspective there are quite a few options. One idea is to use boosting to design a model that is an ensemble of a several smaller models based on non-uniform sampling of the data. Huang et al. recently showed that this can be highly effective in the context of kernel ridge regression . Another idea is to use our solver as the block solver in the block coordinate descent algorithm suggested by Tu et al. . We leave the exploration of these techniques to future work.

More importantly, one should note that even in the era of Big Data, it is not always the case that for a particular problem we have access to very big training set. In such cases it is even more important to fully leverage the data, and produce the best possible model. Our algorithm and implementation provide an effective way to do so.

Acknowledgments

The authors acknowledge the support from the XDATA program of the Defense Advanced Research Projects Agency (DARPA), administered through Air Force Research Laboratory contract FA8750-12-C-0323.

References

Appendix A Appendix - Proof of Theorem 5

Under the conditions of the previous proposition,

We prove that with probability of at least $1-\delta$

Thus, with probability of $1-\delta$ the relevant condition number is bounded by $3$. For PCG, if the condition number is bounded by $\kappa$, we are guaranteed to reduce the error (measured in the matrix norm of the linear equation) to an $\epsilon$ fraction of the initial guess after $\lceil\sqrt{\kappa}\ln(2/\epsilon)/2\rceil$ iterations . This immediately leads to the bound in the theorem statement.

A sufficient condition for (11) to hold is that