gp2Scale: A Class of Compactly Supported Non-Stationary Kernels and Distributed Computing for Exact Gaussian Processes on 10 Million Data Points
Abstract
Despite a large corpus of recent work on scaling up Gaussian processes, a stubborn trade-off between computational speed, prediction, and uncertainty quantification accuracy, and customizability persists. This is because the vast majority of existing methodologies exploit various levels of approximation that lower accuracy and limit the flexibility of kernel and noise-model designs --- an unacceptable drawback at a time when expressive non-stationary kernels are on the rise in many fields. Here, we propose a methodology we term gp2Scale that scales exact Gaussian processes to more than 10 million data points without relying on inducing points, kernel interpolation, or neighborhood-based approximations, and instead leverages the existing capabilities of a GP: its kernel design. Highly flexible, compactly supported, and non-stationary kernels lead to the identification of naturally occurring sparse structure in the covariance matrix, which is then exploited to compute the linear system solution and the log-determinant for training. We demonstrate our method's functionality on several real-world datasets and compare it with state-of-the-art approximation algorithms. Although we show superior approximation performance in many cases, the method's real power lies in its agnosticism toward arbitrary GP customizations --- core kernel design, noise, and mean functions --- and the type of input space, making it optimally suited for modern Gaussian process applications.
Lay Summary
Scientists routinely need to make predictions from data and, crucially, need to know how confident those predictions are --- a requirement that is acute in manufacturing, weather predictions, materials, and drug discovery, and experiments at particle accelerators. A statistical tool called a Gaussian process does this exceptionally well, but it hits a hard computational wall: the calculations required grow so fast with the dataset size that, until now, it could not be applied to more than roughly ten thousand data points. The bottleneck turns out not to be a fundamental mathematical limit, but a design choice: standard Gaussian processes assume, by default, that every pair of data points has a non-zero influence on every other pair, filling the underlying matrix with unnecessary numbers. We developed gp2Scale, which introduces newly designed kernel functions --- the mathematical rules that describe how similar two measurements are --- that can return exactly zero whenever two data points are truly independent of each other. This creates a naturally sparse internal structure that we exploit using hundreds of GPUs working in parallel, allowing an exact, uncompromised Gaussian process to run on more than ten million data points --- a thousand-fold increase over the previous practical limit. Because gp2Scale never approximates, it preserves the rigorous, calibrated uncertainty estimates that make Gaussian processes scientifically trustworthy. We demonstrate the method on datasets ranging from U.S.\ topography to handwritten digit images, outperforming state-of-the-art approximation methods in most cases. The software is freely available as part of the open-source \href{https://gpcam.lbl.gov/}{gpCAM} package, and the method is fully customizable --- so scientists can tailor it to problems in climate science, materials characterization, or autonomous laboratory instruments, all without sacrificing scale or accuracy.