Have ASkotch: Fast Methods for Large-scale, Memory-constrained Kernel Ridge Regression

Kernel ridge regression (KRR) is a fundamental computational tool, appearing in problems that range from computational chemistry to health analytics, with a particular interest due to its starring role in Gaussian process regression. However, it is challenging to scale KRR solvers to large datasets: with n training points, a direct solver (i.e., Cholesky decomposition) uses O(n²) storage and O(n³) flops. Iterative methods for KRR, such as preconditioned conjugate gradient (PCG), avoid the cubic scaling of direct solvers and often use low-rank preconditioners; a rank r preconditioner uses $O(rn)$ storage and each iteration requires O(n²) flops. To reduce the storage and iteration complexity of iterative solvers for KRR, we propose ASkotch (Accelerated scalable kernel optimization using block coordinate descent with Hessian preconditioning). For a given block size $|b| << n$, each iteration of ASkotch uses $O(r|b| + n)$ storage and $O(n|b|)$ flops, so ASkotch scales better than Cholesky decomposition and PCG. We prove that ASkotch obtains linear convergence to the optimum, with the convergence rate depending on the square roots of the preconditioned block condition numbers. Furthermore, we solve KRR problems that were considered to be impossibly large while using limited computational resources: we show that ASkotch outperforms PCG methods with respect to generalization error on large-scale KRR (up to n = 10⁸) and KRR classification tasks (up to n = 10⁷) while running each of our experiments on a single 12 GB Titan V GPU. Our work opens up the possibility of as-yet-unimagined applications of KRR across a wide range of disciplines.