Turbocharging Gaussian Process Inference with Approximate Sketch-and-Project

TL;DR

This paper introduces ADASAP, an accelerated sketch-and-project algorithm that significantly improves the scalability of Gaussian process inference, enabling efficient handling of datasets with over 300 million samples.

Contribution

The paper presents ADASAP, a novel distributed algorithm for Gaussian process inference that is faster, scalable, and theoretically grounded, outperforming existing methods.

Findings

ADASAP outperforms state-of-the-art solvers on benchmark datasets.

It scales to datasets with over 300 million samples.

Theoretical analysis shows rapid convergence of the posterior mean.

Abstract

Gaussian processes (GPs) play an essential role in biostatistics, scientific machine learning, and Bayesian optimization for their ability to provide probabilistic predictions and model uncertainty. However, GP inference struggles to scale to large datasets (which are common in modern applications), since it requires the solution of a linear system whose size scales quadratically with the number of samples in the dataset. We propose an approximate, distributed, accelerated sketch-and-project algorithm ($\texttt{ADASAP}$) for solving these linear systems, which improves scalability. We use the theory of determinantal point processes to show that the posterior mean induced by sketch-and-project rapidly converges to the true posterior mean. In particular, this yields the first efficient, condition number-free algorithm for estimating the posterior mean along the top spectral basis functions, showing that our approach is principled for GP inference. $\texttt{ADASAP}$ outperforms state-of-the-art solvers based on conjugate gradient and coordinate descent across several benchmark datasets and a large-scale Bayesian optimization task. Moreover, $\texttt{ADASAP}$ scales to a dataset with $> 3 \cdot 10^8$ samples, a feat which has not been accomplished in the literature.