Large Scale Kernel Learning using Block Coordinate Descent

Introduction

Kernel methods are a powerful tool in machine learning, allowing one to discover non-linear structure by mapping data into a higher dimensional, possibly infinite, feature space. However, a known issue is that kernel methods do not scale favorably with dataset size. For instance, a naïve implementation of a kernel least squares solver requires $O(n^{2})$ space and $O(n^{3})$ time to store and invert the full kernel matrix. The prevailing belief is that when $n$ reaches the millions, kernel methods are impractical.

This paper challenges the conventional wisdom by pushing kernel methods to the limit of what is practical on modern distributed compute platforms. We show that approximately solving a full kernel least squares problem with $n=2\times 10^{6}$ can be done in a matter of hours, and the resulting model achieves competitive performance in terms of classification errors. Mimicking the successes of the early 2000s, our algorithm is based on block coordinate descent and avoids full materialization of the kernel matrix [Joa99, FCL05].

Furthermore, in contrast to running multiple iterations in parallel and aggregating updates [AD11, NRRW11, ZWSL11, JST+14, LWR+15], we exploit distributed computation to parallelize individual iterations of block coordinate descent. We deliberately make this choice to alleviate communication overheads. Our resulting implementation inherits the linear convergence of block coordinate descent while efficiently scaling up to 1024 cores on 128 machines.

The capability to solve full kernel systems allows us to perform a direct head-to-head empirical comparison between popular kernel approximation techniques and the full kernel at an unprecedented scale. We conduct a thorough study of random features [RR07] and Nyström [WS01] approximations on three large datasets from speech, text, and image classification domains. Extending prior work comparing kernel approximations [YLM+12], our study is the first to work with multi-terabyte kernel matrices and to quantify computational versus statistical performance tradeoffs between the two methods at this scale. More specifically, we identify situations where the Nyström system requires significantly more iterations to converge than a random features system of the same size, but yields a better estimator when it does.

Finally, motivated by the empirical effectiveness of primal block coordinate descent methods in our own study and in related work that inspired our investigations [HAS+14], we derive a new rate of convergence for block coordinate descent on strongly convex smooth quadratic functions. Our analysis shows that block coordinate descent has a convergence rate that is no worse than gradient descent plus a small additive factor which is inversely proportional to the block size. Specializing this result to random features, Nyström, and kernel risk minimization problems corroborates our experimental findings regarding the iteration complexity of the three methods.

Background

This section concisely overviews the techniques used in this paper, and more importantly defines the specific optimization problems we solve. The theoretical underpinnings of kernel methods and their various approximations are well established in the literature; see e.g. [SS01] for a thorough treatment.

1 Kernel methods and approximations

Given a set of data points $\{(x_{i},y_{i})\}_{i=1}^{n}$ with $x_{i}\in\mathcal{X}$ and $y_{i}\in\{1,2,...,k\}$, we use the standard one-versus-all (OVA) approach [RK04] to turn a multiclass classification problem into $k$ binary classification problems of the form

Let $I\in 2^{[n]}$ denote an index set of size $p$, and let $K_{I}:=K([n],I)$, $K_{II}:=K(I,I)$. One common variant of the Nyström method [WS01, DM05, GM13, BJ05] is to use the matrix $\widehat{K}:=K_{I}K_{II}^{\dagger}K_{I}^{\mathsf{T}}$ as a low rank approximation to $K$ (2). An alternative approach is to first replace the minimization over $\mathcal{H}$ in (1) with $\mathcal{H}_{I}$, where $\mathcal{H}_{I}:=\operatorname*{span}\{\upkappa(x_{i},\cdot):i\in I\}$. We then arrive at the optimization problem

This extra regularization is justified statistically by [Bac13, EM15].

Random features.

2 Related work

An empirical comparison on Nyström versus random features was done by Yang et al. [YLM+12]. This study demonstrated that the Nyström method outperformed random features on every dataset in their experiments. Our experimental efforts differ from this seminal work in several ways. First, we quantify time versus statistical performance tradeoffs, instead of studying only the empirical risk minimizer. Second, we describe a scalable algorithm which allows us to compare performance with the full kernel. Finally, our datasets are significantly larger, and we also sweep across a much wider range of number of random features.

On the algorithms side, the inspiration for this work was by Huang et al. [HAS+14], who devised a similar block coordinate descent algorithm for solving random feature systems. In this work, we extend the block coordinate algorithm to both the full kernel and Nyström systems. This enables us to train the full kernel on the entire TIMIT dataset, achieving a lower test error than the random feature approximations.

Algorithms

The optimal solutions written in Section 2.1 require solving large linear systems where the data cannot be assumed to fit entirely in memory. This necessitates a different algorithm than the least squares solvers implemented in standard library routines. Fortunately, for the statistical problems we are interested in, obtaining a high accuracy solution is not as important. Hence, we propose to use block coordinate descent [BT89], which admits a natural distributed implementation, and, in our experience, converges to a reasonable accuracy after only a few passes through the data.

Coordinate methods in machine learning date back to the late 90s with SVMLight [Joa99] and SMO [Pla98]. More recently, many researchers [Yan13, RT13, JST+14, MSJ+15] have proposed using distributed computation to run multiple iterations of coordinate descent in parallel. As noted previously, we take a different approach and use distributed computing to accelerate within an iteration. This is similar to [HCL+08, YHCL10], both who describe block coordinate algorithms for solving SVMs. However, using the square loss instead of hinge loss simplifies our analysis and implementation.

1 Block coordinate descent

In the case where $f$ is least squares, the update (5) is equivalent to block Gauss-Seidel on the normal equations. For instance, for (4), the update (5) reduces to solving the $b\times b$ equation

where $Z_{I_{\tau}}:=Z([n],I_{\tau})$. The $w_{I_{\tau}}$ notation means we set only the coordinates in $I_{\tau}$ equal to the RHS, and the coordinates not in $I_{\tau}$ remain the same from the previous iteration.

We solve block coordinate descent in parallel by distributing the computation of $Z_{I_{\tau}}^{\mathsf{T}}Z_{I_{\tau}}$ and $Z_{I_{\tau}}^{\mathsf{T}}Y$. To do this, we partition the rows of $Z_{I_{\tau}}$, $Y$ across all the machines in a cluster and compute the sum of outer products from each machine. The result of this distributed operation is a $b\times b$ matrix and we pick $b$ such that the solve for $w^{\tau+1}_{I_{\tau}}$ can be computed quickly using existing lapack solvers on a single machine.

Choosing an appropriate value of $b$ is important as it affects both the statistical accuracy and run-time performance. Using a larger value for $b$ leads to improved convergence and is also helpful for using BLAS-3 primitives in single machine operations. However, a very large value for $b$ increases the serial execution time and the communication costs. In practice, we see that setting $b$ in the range 2,000 to 8,000 offers a good trade-off.

Block generation primitives.

As mentioned previously, our algorithms only require a procedure that materializes a column block at a time. We denote this primitive by $\textsc{KernelBlock}(X,I)$, where $X$ represents the data matrix and $I$ is a list of column indices. The output of KernelBlock is $K([n],I)$. After a column block is used in a block coordinate descent update of the model, it can be immediately discarded. We also use distributed computation to parallelize the generation of a block $K_{I_{\tau}}$ of the kernel matrix. We also define a similar primitive, $\textsc{RandomFeaturesBlock}(X,I)$, which returns $Z([n],I)$ for random feature systems.

2 Algorithm descriptions

Our full kernel solver is described in Algorithm 1. Algorithm 1 is actually Gauss-Seidel on the linear system $(K+n\lambda I_{n})\alpha=Y$, but as we will discuss in Section 4.3, this is equivalent to block coordinate descent on a modified objective function (which is strongly convex, even when $K$ is rank deficient). See [HNR15] for a similar discussion in the context of ridge regression.

Nyström block coordinate descent.

Unlike the full kernel case, our Nyström implementation operates directly on the normal equations. A notable point of our algorithm is that it does not require computation of the pseudo-inverse $K_{II}^{\dagger}$. When the number of Nyström features is large, calculating the pseudo-inverse $K_{II}^{\dagger}$ is expensive in terms of computation and communication. By making $K_{II}$ a part of the block coordinate descent update we are able to handle large number of Nyström features while only needing a block of features at a time.

We denote $\textsc{Selector}(n,I)$ as the function which returns an $\{0,1\}^{n\times|I|}$ matrix $S$ such that $S_{I(j)j}=1$ and zero otherwise, for $j=1,...,|I|$; this is simply the column selector matrix associated with the indices in $I$. Using the above notation, the Nyström algorithm is described in Algorithm 2.

Random features block coordinate descent.

Our random features solver is the same as Algorithm 2 from [HAS+14]. We include it in Algorithm 3 for completeness.

Computation and communication overheads.

Table 1 summarizes the computation and communication costs of the algorithms presented below. The computation costs in the full kernel are associated with computing the residual $R$ and solving a $b\times b$ linear system. For the Nyström method (and random features), the computation costs include computing $K_{I}^{\mathsf{T}}K_{I}$, $K_{I}^{\mathsf{T}}Y$ in parallel and a similar local solve. Computing the gram matrix however requires adding $M$ matrices of size $b\times b$. Using a tree-based aggregation, this results in $O(\log(M)b^{2})$ bytes being transferred. We study how these costs matter in practice in Section 5.

Computing the regularization path.

Algorithms 1, 2, and 3 are all described for a single input $\lambda$. In practice, for model selection, one often computes an estimator for multiple values of $\lambda$. The naïve way of doing this is to run the algorithm again for each value of $\lambda$. However, a faster approach, which we use in our experiments, is to maintain separate models $\alpha_{\lambda}$ and seperate residuals $R_{\lambda}$ for each value of $\lambda$, and reuse the computation of the block matrices $K_{b}$ for full kernel, $K_{b}^{\mathsf{T}}K_{b}$ for Nyström, and $Z_{b}^{\mathsf{T}}Z_{b}$ for random features. We can do this because the block matrices do not depend on the value of $\lambda$. As we show in Section 5, in each iteration of our algorithms, a large fraction of time is spent in computing these block matrices; thus this optimization allows us compute solutions for multiple $\lambda$ values for essentially the price of a single solution.

We would like to note that in the case of Nyström approximations, [RCR15] provides an algorithm for computing the regularization path along $p$, the number of Nyström samples, by using rank-one Cholesky updates. We leave it as future work to see if a similar technique can be applied to our Nyström block coordinate algorithm.

Optimization and statistical rates

In this section we present our theoretical results which characterize optimization error for kernel methods. All proofs are deferred to the appendix.

We start by stating the existing rates for block coordinate descent. To do this, we define a restricted Lipschitz constant as follows. For any $Q(\cdot)$ such that $Q(x)\succeq 0$ for all $x$, define

iterations, where $L_{\max,b}:=L_{\max,b}(\nabla^{2}f(\cdot))$.

1 Improved rate for quadratic functions

Let $w^{\tau}$ denote the $\tau$-th iterate of block coordinate descent with the index set $I_{\tau}$ consisting of $b\in\{1,...,d\}$ indices drawn uniformly at random without replacement from $\{1,...,d\}$. The iterate $w^{\tau}$ satisfies

Theorem 4.1 states that in order to reach an $\epsilon$-sub-optimal solution for $f$, the number of iterations required is at most

That is, block coordinate descent pays the rate of gradient descent plus $1/b$ times the rate of standard ($b=1$) coordinate descent (ignoring log factors). To see that this can be much better than the standard rate (7), suppose that $d=p^{2}$ for some $p\geq 1$, and consider any quadratic with Hessian

Equation (8) suggests setting $b$ such that the second term matches $L/m$ order wise. That is, as long as $b\gtrsim d\log{d}\lVert\operatorname*{diag}(\nabla^{2}f)\rVert_{\infty}/L$, we have that at most $\widetilde{O}(L/m)$ iterations are necessary We use the notation $x\gtrsim y$ to mean there exists an absolute constant $C>0$ such that $x\geq Cy$, and $\widetilde{O}(\cdot)$ to suppress dependence on poly-logarithmic terms. . In the sequel, we will assume this setting of $b$.

We highlight the main ideas behind the proof of Theorem 4.1. The proof proceeds in two steps. First, we establish a structural result which states that, given a large set $\mathcal{G}$ of indices where the restricted Lipschitz constant of the Hessian is well controlled, the overall dependence on Lipschitz constant is not much worse than the maximum Lipschitz constant restricted to $\mathcal{G}$. Second, we use a probabilistic argument to show that such a set $\mathcal{G}$ does indeed exist. The first result is based on a modification of the standard coordinate descent proof, whereas the second result is based on a matrix Chernoff argument.

2 Rates for kernel optimization

Table 2 quantifies the iteration complexity of solving the full kernel system versus the Nyström and random features approximation. Our worst case analysis shows that the Nyström system requires roughly $p$ times more iterations to solve than random features. This difference is due to the inability to reduce the Nyström normal equation from quadratic in $K$ to linear in $K$, as is done in the full kernel normal equation. Indeed, the Nyström method is less well conditioned in practice, and we observe similar phenomena in our experiments below. The derivation of the bounds in Table 2 is deferred to Appendix B.

3 Primal versus dual coordinate methods

Duality gives us a choice as to whether to solve the primal or dual problem; strong duality asserts that both solutions are equivalent. We can use this freedom to our advantage, picking the formulation which yields the most numerically stable system. For instance, in the full kernel solver we chose to work with the system $(K+n\lambda I_{n})\alpha=Y$ instead of $K(K+n\lambda I_{n})\alpha=KY$. The former is actually the dual system, and the latter is the primal. Here, the primal system has a condition number which is roughly the square of the dual.

On the other hand, for both Nyström and random features, our system works on the primal formulation. This is intuitively desirable since $p\ll n$ and hence the primal system is much smaller. However, some authors including [SSZ13] advocate for the dual formulation even when $p\ll n$. We claim that, at least in the case of random Fourier features, their argument does not apply.

To do this, we consider the random features program with $b=k=1$, which fits the framework of [SSZ13] the closest. By the primal-dual correspondence $w=\frac{1}{n\lambda}Z^{\mathsf{T}}\alpha$, the dual program is

Theorem 5 from [SSZ13] states that $O((n+L_{\max,1}(ZZ^{\mathsf{T}})/\lambda)\log{\epsilon^{-1}})$ iterations of dual coordinate ascent are sufficient to reach an $\epsilon$-sub-optimal primal solution. On the other hand, Equation (7) yields that at most $O((pL_{\max,1}(Z^{\mathsf{T}}Z)/n\lambda)\log{\epsilon^{-1}})$ iterations of primal coordinate descent are sufficient to reach the same accuracy.

For random Fourier features, both $L_{\max,1}(ZZ^{\mathsf{T}})$ and $L_{\max,1}(Z^{\mathsf{T}}Z)$ can be easily upper bounded, since $|\langle z(x_{i}),z(x_{i})\rangle|\leq\frac{2}{p}\sum_{k=1}^{p}|\cos(x_{i}^{\mathsf{T}}w_{k}+b_{k})|\leq 2$ and also $\lVert\operatorname*{diag}(Z^{\mathsf{T}}Z)\rVert_{\infty}=\max_{1\leq k\leq p}\frac{2}{p}\sum_{i=1}^{n}\cos^{2}(w_{k}^{\mathsf{T}}x_{i}+b_{k})\leq\frac{2n}{p}$. Therefore, the dual rate is $\widetilde{O}(n+1/\lambda)$ and the primal rate is $\widetilde{O}(1/\lambda)$. That is, for random Fourier features, the primal rate upper bound beats the dual rate upper bound.

Experiments

This section describes our experimental evaluation. We implement our algorithms in Scala on top of Apache Spark [ZCD+12]. Our experiments are run on Amazon EC2, with a cluster of 128 r3.2xlarge machines, each of which has 4 physical cores and 62 GB of RAM.

We measure classification accuracy for three large datasets spanning speech (TIMIT), text (Yelp), and vision (CIFAR-10). The size of these datasets are summarized in Table 3. For all our experiments, we set the block size to $b=6144$. We shuffle the raw data at the beginning of the algorithm, and select blocks in a random order for block coordinate descent. For the Nyström method, we uniformly sample $p$ columns without replacement from the full kernel matrix.

We evaluate a phone classification task on the TIMIT datasethttps://catalog.ldc.upenn.edu/LDC93S1, which consists of spoken audio from 462 speakers. We use the same preprocessing pipeline as [HAS+14], resulting in $2.25\times 10^{6}$ training examples and $10^{5}$ test examples. The preprocessing pipeline produces a dense vector with $440$ features and we use a shuffled version of this as the input to our kernel methods. We apply a Gaussian (RBF) kernel for the Nyström and exact methods and use random cosines [RR07] for the random feature method. Figure 1 shows the top-1 test error for each technique. From the figure, we can see that while the exact method takes the longest to complete a full epoch (around $2.5$ hours), it achieves the lowest top-1 test-error ($33.51\%$) among all methods after $3$ epochs. Furthermore, unlike the exact method, the data for the Nyström and random features with $p=200,000$ can be cached in memory; as a result, the approximate methods run much faster after the first epoch compared to the exact method.

We also compare Nyström and random features by varying $p$ in Figure 2 and find that for $p\geq 100,000$ both methods approach the test error of the full kernel within $1\%$.

2 Yelp Reviews

We next evaluate a text classification task where the goal is to predict a rating from one to five stars from the text of a review. The data comes from Yelp’s academic datasethttps://www.yelp.com/academic_dataset, which consists of $1.5\times 10^{6}$ customer reviews. We set aside 20% of the reviews for test, and train on the remaining 80%. For preprocessing, we use nltkhttp://www.nltk.org/ for tokenizing and stemming documents. We then remove English stop words and create $3$-grams, resulting in a sparse vector with dimension $6.52\times 10^{7}$. For the exact and Nyström experiments, we apply a linear kernel, which when combined with the $3$-grams can be viewed as an instance of a string kernel [SRR07]. For random features, we apply a hash kernel [WDL+09] using MurmurHash3 as our hash function. Since we are predicting ratings for a review, we measure accuracy by using the root mean square error (RMSE) of the predicted rating as compared to the actual rating. Figure 1 shows how various kernel methods perform with respect to wall clock time. From the figure, we can see that the string kernel performs much better than the hash-based random features for this classification task. We also see that the Nyström method achieves almost the same RMSE ($0.861$) as the full kernel ($0.854$) when using $122,000$ features. Finally, Figure 2 shows that the improved accuracy from using the string kernel over hashing holds as we vary the number of features ($p$) for the Nyström and random feature methods.

3 CIFAR-10

Our last task involves image classification for the CIFAR-10 dataset cs.toronto.edu/~kriz/cifar.html. We perform the same data augmentation as described in cuda-convnet2 github.com/akrizhevsky/cuda-convnet2, which results in 500,000 train images. For preprocessing, we use a pipeline similar to [CN12], replacing the $k$-means step with random image patches. Using 512 random image patches, we get $4096$ features per image and fitting a linear model with these features gives us $25.7\%$ test error. For our kernel methods, we start with these $4096$ features as the input and we use the RBF kernel for the exact and Nyström method and random cosines for the random features method.

From Figure 1, we see that on CIFAR-10 the full kernel takes around the same time as Nyström and random features. This is because we have fewer examples ($n=500,000$) and this leads to fewer blocks that need to be solved per-epoch. We are also able to cache the entire kernel matrix in memory ($\sim$ 2TB) in this case and this provides a speedup after the first epoch.

Furthermore, as shown in Figure 2, we see that applying a non-linear kernel to the output of convolutions using random patches can result in significant improvement in accuracy. With the non-linear kernel, we achieve a test error of $16.4\%$, which is $9.3\%$ lower than a linear model trained with the same features.

When comparing random features and Nyström after $5$ epochs for various values of $p$, we see that they perform similarly for smaller number of features but that random features performs better with larger number of features. We believe that this is due to the Nyström normal equations having a larger condition number for the CIFAR-10 augmented dataset which leads to a worse convergence rate. We verify this in Figures 4 and 4 by running by Nyström and random feature solvers for 50 epochs. In Figure 4, we fix the number of random features to $p\in\{60,000,200,000\}$, and we see that Nyström takes more epochs to converge but reaches a better test error. In Figure 4, we perform the same sweep as in Figure 2 except we stop at 50 epochs instead of 5. Indeed, when we do this, the difference between Nyström and random features matches the trends in Figures 2 and 2.

4 Performance

We next study the runtime performance characteristics of each method. Figure 3 shows a timing breakdown for running one block of block coordinate descent on the three datasets. From the figure, we see that the choice of the kernel approximation can significantly impact performance since different kernels take different amounts of time to generate. For example, the hash random feature used for the Yelp dataset is much cheaper to compute than the string kernel. However, computing a block of the RBF kernel is similar in cost to computing a block of random cosine features. This results in similar performance characteristics for the Nyström and random feature methods on TIMIT.

We also observe that the full kernel takes the least amount of time to solve one block. This is primarily because the full kernel does not compute a gram matrix $Z_{b}^{\mathsf{T}}Z_{b}$ and only extracts a block of the kernel matrix $K_{bb}$. Thus, when the number of blocks is small, as is the case for CIFAR-10 in Figure 1, the full kernel’s performance becomes comparable to the Nyström method.

5 Scalability of RBF kernel generation

Figure 3 also demonstrates that computing the gram matrix and generating the kernel block are the two most expensive steps in our algorithm. Computing the gram matrix uses distributed matrix multiplication, which is well studied [VDGW97]. To see how the cost of kernel generation changes as dataset size grows, we perform a weak scaling experiment where we increase the number of examples and the number of machines used while keeping the number of examples per machine constant ($n=16,384$). We run this experiment for $d=440$ and $d=4096$, which are the number of features in TIMIT and CIFAR-10 respectively. Figure 5 contains results from this experiment. In the weak scaling scenario, ideal scaling implies that the time to generate a block of the kernel matrix remains constant as we increase both the data and the number of machines. However, computing a block of the RBF kernel involves broadcasting a $b\times d$ matrix to all the machines in the cluster. This causes a slight decrease in performance as we go from $8$ to $128$ machines. As broadcast routines scale as $O(\log{M})$, we believe that our kernel block generation methods will continue to scale well for larger datasets.

Conclusion

This paper shows that scalable kernel machines are feasible with distributed computation. There are several theoretical and experimental continuations of this work.

On the theoretical side, a limitation of our current analysis of block coordinate descent is that we cannot hope to achieve rates better than gradient descent. We believe it is possible to leverage the direct solve in (5) to improve our rate, since when $b=d$ the algorithm reduces to Newton’s method. We are also interested in seeing if acceleration techniques can be applied to substantially reduce the number of iterations needed.

On the experimental side, we would like to extend our algorithm to handle other losses than the square loss; ADMM might be one approach for this. More broadly, since solving a least squares program is a core primitive for many optimization algorithms, we are interested to see if our techniques can be applied in other domains.

Acknowledgements

The authors thank Vikas Sindhwani and the IBM corporation for providing access to the derived TIMIT dataset used in our experiments. BR is generously supported by ONR awards N00014-14-1-0024, N00014-15-1-2620, and N00014-13-1-0129, and NSF awards CCF-1148243 and CCF-1217058. RR is supported by the U.S. Department of Energy under award numbers DE-SC0008700 and AC02-05CH11231. This research is supported in part by NSF CISE Expeditions Award CCF-1139158, LBNL Award 7076018, DARPA XData Award FA8750-12-2-0331, and gifts from Amazon Web Services, Google, SAP, The Thomas and Stacey Siebel Foundation, Adatao, Adobe, Apple, Inc., Blue Goji, Bosch, C3Energy, Cisco, Cray, Cloudera, EMC2, Ericsson, Facebook, Guavus, HP, Huawei, Informatica, Intel, Microsoft, NetApp, Pivotal, Samsung, Schlumberger, Splunk, Virdata and VMware.

References

Appendix A Proof of Theorem 4.1

Block Lipschitz constants.

We now define a restricted notion of Lipschitz continuity which works on blocks. For an index set $I$, define

Update rule.

where $I_{0},I_{1},...\in\Omega_{b}$ are chosen by some (random) strategy. A common choice is to choose $I_{k}$ uniformly at random from $\Omega_{b}$, and to make this choice independent of the history up to time $k$. This is the sampling strategy we will study. We now have enough notation to state and prove our basic inequality for coordinate descent. This is not new, but we record it for completeness, and because it is simple.

For every $k\geq 0$, we have that the $k+1$-th iterate satisfies the inequality

Now, by Taylor’s theorem, for some $t\in(0,1)$, setting $\alpha_{k}:=1/L_{I_{k}}$,

where (a) uses the fact that Euclidean projection is idempotent and also the definition of $L_{I_{k}}$. ∎

We now prove a structural result. The main idea is as follows. Suppose we have some subset $\mathcal{G}\subset\Omega_{b}$ where $\max_{I\in\mathcal{G}}L_{I}$ is much smaller compared to $L_{\max,b}$. If this subset is a significant portion of $\Omega_{b}$, then we expect to be able to improve the basic rate. The following result lays the groundwork for us to be able to make this kind of claim.

Let $\mathcal{G}\subset\Omega_{b}$ be such that $\frac{|\mathcal{G}^{c}|}{|\Omega_{b}|}=\alpha\frac{b}{d}$ for $\alpha\in$. Let each $I_{k}$ be independent and drawn uniformly from $\Omega_{b}$. Then, after $\tau$ iterations, the iterate $w^{\tau}$ satisfies

The basic proof structure is based on Theorem 1 of [Wri15]. The idea here is to compute the conditional expectation of $\frac{1}{L_{I_{k}}}\lVert P_{I_{k}}\nabla f(w^{k})\rVert^{2}$ w.r.t. $w^{k}$, taking advantage of the structure provided by $\mathcal{G}$. Put $t:=\max_{I\in\mathcal{G}}L_{I}$. Then,

Combining (11) and (12) with Proposition A.1 followed by iterating expectations, we conclude that

The rest of the proof proceeds identically to Theorem 1 of [Wri15], using $m$-strong convexity to control $\lVert\nabla f(w^{k})\rVert^{2}$ from below. ∎

The remainder of the proof involves showing the existence of a set $\mathcal{J}\subset\Omega_{b}$ that satisfies the hypothesis of Lemma A.2. To show this, we need some basic tools from random matrix theory. The following matrix Chernoff inequality is Theorem 2.2 from [Tro11].

The inequality of Theorem A.3 can be weakened to a more useful form, which we will use directly (see e.g. Section 5.1 of [Tro15]). The following bound holds for all $t\geq e$,

We now establish a result controlling the size of $\mathcal{J}_{t}(A)$ from below. We do this via a probabilistic argument, taking advantage of the matrix Chernoff inequality.

and plug into (14). The result follows by noting that $\max_{1\leq i\leq p}\lVert A_{:i}\rVert^{2}=\max_{1\leq i\leq p}e_{i}^{\mathsf{T}}A^{\mathsf{T}}Ae_{i}=\max_{1\leq i\leq p}(A^{\mathsf{T}}A)_{ii}$. ∎

We are now in a position to prove Theorem 4.1.

Setting $t$ as in (15) and invoking Lemma A.4, we have that every $I\in\mathcal{J}_{t}(Q)$ satisfies

The result follows immediately by an application of Lemma A.2 ∎

Appendix B Proofs for Section 4.2

As noted in Section 3, we actually run Gauss-Seidel on $(K+n\lambda I_{n})\alpha=Y$, which can be seen as coordinate descent on the program

Note that the objective is a strongly convex function with Hessian given as $D^{2}f(\alpha)[H,H]=\langle H,(K+n\lambda I_{n})H\rangle$. Theorem 4.1 tells us that setting $b\gtrsim n\log{n}\frac{\lVert\operatorname*{diag}(K+n\lambda I_{n})\rVert_{\infty}}{\lVert K+n\lambda I_{n}\rVert}$, $\widetilde{O}(1/\lambda)$ iterations are sufficient. Plugging values in, we get for $b\gtrsim\log^{2}{n}$ under (a) and $b\gtrsim n^{1/(2\beta+1)}\log{n}$ under (b), the number of iterations is bounded under (a) by $\widetilde{O}(n)$ and under (b) by $\widetilde{O}(n^{\frac{2\beta}{2\beta+1}})$.

Nyström approximation.

Theorem 4.1 tells us that we want to set $b\gtrsim p\log{p}\frac{\lVert\operatorname*{diag}(S^{\mathsf{T}}K(K+n\lambda I_{n})S+n\lambda\gamma I_{p})\rVert_{\infty}}{\lVert S^{\mathsf{T}}K(K+n\lambda I_{n})S+n\lambda\gamma I_{p}\rVert}$. Applying a matrix Chernoff argument (Lemma A.4) to control $\lVert S^{\mathsf{T}}K(K+n\lambda I_{n})S\rVert$ from above, we have w.h.p. that the number of iterations is $\widetilde{O}(p/\lambda\gamma)$. Under (a) this is $\widetilde{O}(np/\gamma)$ and under (b) this is $\widetilde{O}(pn^{\frac{2\beta}{2\beta+1}}/\gamma)$.

To control $b$, we apply a matrix Bernstein argument (Lemma B.2) to control $\lVert S^{\mathsf{T}}K(K+n\lambda I_{n})S\rVert$ from below w.h.p. This argument shows that when $\lambda\leq O(1)$ and $p\gtrsim\log{n}$, $\lVert S^{\mathsf{T}}K(K+n\lambda I_{n})S\rVert\gtrsim np$, from which we conclude that $b\gtrsim(1+\gamma)\log{n}$.

Random features approximation.

We now derive a rate for coordinate descent on (4). The Hessian of (4) is given by $nD^{2}f(\alpha)[H,H]=\langle H,(Z^{\mathsf{T}}Z+n\lambda I_{p})H\rangle$. Thus by Theorem 4.1, setting $b\gtrsim p\log{p}\frac{\lVert\operatorname*{diag}(Z^{\mathsf{T}}Z+n\lambda I_{p})\rVert_{\infty}}{\lVert Z^{\mathsf{T}}Z+n\lambda I_{p}\rVert}$ and applying a matrix Bernstein argument to control $\lVert Z^{\mathsf{T}}Z+n\lambda I_{p}\rVert$ from both directions (Lemma B.4), then as long as $p\gtrsim\log{n}$, we have w.h.p. that this is at most $\widetilde{O}(1/\lambda)$, which is the same rate as the full kernel. Furthermore, the block size is $b\gtrsim\log{n}$.

B.2 Supporting lemmas for Section B.1

For a fixed symmetric $Q$ and random $I$, we want to control $\lambda_{\max}(P_{I}QP_{I})$ from below. The matrix Chernoff arguments do not allow us to do this, so we rely on matrix Bernstein. The following result is Theorem 2 from [EM15].

This paves the way for the following lemma.

Put $p_{i}=1/m$ for $i=1,...,m$ and and $\beta=\frac{\lVert\Psi\rVert_{F}^{2}}{mB}$. By definition, $\beta\leq 1$. In this case, $\widetilde{S}=\sqrt{\frac{m}{p}}S$. Plugging these constants into Theorem B.1, we get that

Setting the RHS equal to $\delta$, we get that $t$ is the roots of the quadratic equation

Since solutions to $t^{2}-at-b=0$ satisfy $t\leq 2(a+\sqrt{b})$ when $a,b\geq 0$, from this we conclude

By the convexity of $\lambda_{\max}(\cdot)$,

Combining (16) and (17) yields the result. ∎

We now study random features. To do this, we need the following general variant of matrix Bernstein. The following is Corollary 6.2.1 of [Tro15].

where each $R_{k}$ is an independent copy of $R$. Then for all $t\geq 0$,

This variant allows us to easily establish the following lemma.