Source-linked AI summary
Exact Gaussian Processes on a Million Data Points
Ke Alexander Wang, Geoff Pleiss, Jacob R. Gardner, Stephen Tyree, Kilian Q. Weinberger, Andrew Gordon Wilson
TL;DR
Exact GP inference has been limited by the computation and memory demands of standard methods, leaving comparisons with scalable approximations unresolved at very large data sizes. The paper uses preconditioned conjugate gradients and distributed kernel matrix multiplications across GPUs to scale exact inference beyond one million points. Exact GPs outperform popular approximations on nearly all studied large-scale datasets, while retaining fast prediction and broader data usage.
Problem
Exact GP inference is computationally intractable at large n under standard procedures, and its comparison with approximations has remained open in large-data regimes.
Method
The method uses preconditioned conjugate gradients with partitioned kernel matrix-vector multiplications distributed across multiple GPUs, without explicitly forming the kernel matrix.
Results
Exact GPs outperform popular approximate GP methods on nearly all benchmark datasets, including large datasets with n > 10^6.
Takeaways & Limitations
Exact GPs can use all available training data for predictions at large n while achieving fast prediction and linear training speedups from additional GPUs.
Takeaways & Limitations
Approximate methods may remain preferable with limited computational resources, in low-dimensional settings, or for non-Gaussian likelihoods requiring approximate inference.
Abstract
from arXiv · showhide
Gaussian processes (GPs) are flexible non-parametric models, with a capacity that grows with the available data. However, computational constraints with standard inference procedures have limited exact GPs to problems with fewer than about ten thousand training points, necessitating approximations for larger datasets. In this paper, we develop a scalable approach for exact GPs that leverages multi-GPU parallelization and methods like linear conjugate gradients, accessing the kernel matrix only through matrix multiplication. By partitioning and distributing kernel matrix multiplies, we demonstrate that an exact GP can be trained on over a million points, a task previously thought to be impossible with current computing hardware, in less than 2 hours. Moreover, our approach is generally applicable, without constraints to grid data or specific kernel classes. Enabled by this scalability, we perform the first-ever comparison of exact GPs against scalable GP approximations on datasets with $10^4 \!-\! 10^6$ data points, showing dramatic performance improvements.
1 Introduction
Exact GP inference has historically been difficult to scale because standard methods require quadratic storage and cubic computation. This paper combines matrix-multiplication-based inference with multi-GPU partitioning to train exact GPs beyond one million points and compare them directly with scalable approximations.
- O(n^3) computation and O(n^2) storage make exact GP inference intractable for large datasets, motivating scalable approximations.
- The paper scales exact GP training beyond previous limits using BBMM, conjugate gradients, and kernel matrix multiplications without approximations.BBMM reduces inference to iterations of matrix multiplication and avoids standard Cholesky-based implementations.
- Partitioning and distributing kernel matrix multiplications across GPUs reduces per-GPU memory to O(n), enabling training from seconds at n ≈ 10^4 to a few days at n ≈ 10^6.The approach also uses practical heuristics and improved hyperparameter initialization to accelerate training.
- Exact GP predictions remain available in less than 1 second on one GPU through a simple caching strategy.
- On UCI regression datasets, exact GPs often exceed a two-fold reduction in root-mean-squared error compared with scalable approximations.The comparison covers datasets in the large-data regime and supports continued benefits from adding training points.
2 Background
GPs define distributions over functions through mean and kernel functions, while exact inference commonly relies on costly linear algebra. The background introduces conjugate gradients as an iterative alternative requiring only kernel matrix-vector products.
- Gaussian processes: A GP places a distribution over functions using a prior mean, a kernel, and observed training data.The mean and kernel encode prior information about the data; RBF and Matérn kernels are common choices.
- Notation: The kernel matrix K_XX contains covariance terms for all training-point pairs, while bK_XX adds Gaussian observational noise.
- Training: GP hyperparameters are typically learned by maximizing the log marginal likelihood with gradient descent, and GPs usually have few such parameters.
- Predictions: With cached training quantities, predictive calculations can be reduced to O(n) matrix-vector multiplication and further accelerated through caching.The predictive posterior for a test point is Gaussian under a Gaussian likelihood, with mean and variance given by the GP equations.
- Cholesky decomposition: Cholesky factorization requires O(n^3) time and O(n^2) memory, and its recursive structure limits GPU parallelization.Distributed Cholesky approaches also incur quadratic communication and memory costs.
- Conjugate gradients: Conjugate gradients solves K_XXv = y iteratively, with each step requiring a kernel matrix-vector multiplication; preconditioning can accelerate convergence.A modified preconditioned method computes the terms needed for GP training and prediction using only kernel matrix-vector products.
3 Method
The method scales exact GP inference by replacing explicit Cholesky-based computation with partitioned kernel matrix-vector multiplications solved by preconditioned conjugate gradients and distributed across GPUs. This reduces per-device memory to O(n), supports large training sets, and preserves efficient exact predictions.
- Motivation: Exact GP inference with Cholesky becomes difficult beyond n >10^4 because of decomposition time, communication overhead, and quadratic storage.At n = 500,000, the decomposition requires a full terabyte of memory.
- O(n) memory MVM-based inference: O(n) memory replaces the O(n2) storage of explicit kernel matrices by computing matrix-vector products in separate pieces.Only constant-sized kernel-matrix portions are stored at a time, alongside the four PCG vectors requiring 4n space.
- Partitioned kernel MVMs: Partitioning the kernel matrix into row blocks lets each block’s matrix-vector product be computed separately, then concatenated while discarding the block.Each partition contains roughly 1/p of the full matrix entries and requires the full training data and vector v in memory.
- Multi-GPU parallelization: Distributed kernel matrix multiplications require O(n) communication, compared with O(n2) communication for distributed Cholesky decomposition.Devices exchange right-hand-side vectors and return partitioned output vectors rather than communicating full matrix factors.
- Predictions: Exact GP predictions remain efficient because the predictive mean uses a stored training-data solve, while predictive variance uses a training-data-dependent cache.Both predictive mean and variance can be computed in less than a second on a single GPU, even when training required days.
- PCG Convergence Criteria: Conjugate gradients is exact when run to a sufficiently strict tolerance, while hyperparameter training can tolerate looser convergence criteria than prediction solves.Predictive performance requires ϵ ≤0.01, whereas ϵ = 1 has little impact on final model performance during hyperparameter training.
4 Related Work
Prior scalable GP work includes structured-grid and inducing-point approaches, while recent exact methods use conjugate gradients on smaller or structured problems. The paper’s method extends MVM-based exact inference beyond those settings without requiring structured kernels or grid data.
- MVM-based GP inference: Earlier conjugate-gradient methods often relied on structured kernels, such as regularly spaced grids that enable O(n log n) matrix-vector multiplications.These methods reduce asymptotic complexity when the data lie on a regularly spaced grid.
- MVM-based GP inference: Cutajar et al. trained exact GPs with conjugate gradients on datasets containing up to 50,000 points.Their work investigated off-the-shelf preconditioners and inducing-point-based preconditioners.
- Approximate GP methods: Approximate GP methods scale through inducing points, structured kernel interpolation, or stochastic variational parameters rather than exact full-data inference.Inducing-point methods use O(nm2) time and O(nm) space, while SKI and SVGP provide alternative scalable constructions.
5 Results
The experiments compare exact GPs with scalable approximations across large UCI datasets, examining accuracy, training and prediction time, GPU scaling, and data or inducing-point effects. Exact GPs generally achieve lower error, while matrix-multiplication inference supports scalable training and fast prediction.
- Training time: Better hyperparameter initialization produced similar RMSEs with drastically less training time than 100 Adam iterations on large datasets.The initialization comparison evaluates exact GPs trained with the paper’s procedure against exact GPs trained using Adam.
- Accuracy: Exact GPs outperform popular approximate methods on nearly all benchmarking datasets, including datasets with over 1M points.The comparison covers SGPR and SVGP on UCI regression datasets.
- Training time: Exact GPs achieved near-linear training speedups from additional GPUs because inference is based on matrix multiplication.Kernel partitioning distributes the computation across devices for large datasets.
- Prediction time: Less than a second was required to compute 1,000 predictive means and variances for exact GPs across all evaluated dataset sizes.Prediction timing was measured on one NVIDIA RTX 2080 Ti GPU after the relevant prediction work was precomputed.
- Data scaling: An exact GP trained on less than a quarter of the training data outperformed approximate GPs trained on the entire dataset in all three examined cases.Test error for exact GPs continued to decrease as additional training data was included.
- Inducing-point ablation: Both inducing-point methods’ test RMSE saturated above the exact GP’s error as the number of inducing points increased.The inducing-point methods introduce O(nm^2 + m^3) time complexity, making much larger m difficult to train on one GPU.
6 Discussion
The paper extends exact GP inference to datasets exceeding one million training examples and finds exact GPs substantially outperform scalable approximations, while acknowledging regimes where approximations remain preferable.
- Over one million training examples are now within the demonstrated scope of exact GP inference through MVM-based methods and modern parallel hardware.The approach uses easily parallelizable routines and avoids the historical reliance on scalable approximations for large datasets.
- Exact GPs perform significantly better than approximate methods on large datasets while requiring fewer design choices.This comparison supports the paper’s claim that exact GPs are more broadly applicable than previously thought.
- Approximate GP methods may remain preferable when computational resources are limited, in low-dimensional settings, or for non-Gaussian likelihoods such as classification.Some low-dimensional approximations can achieve high accuracy faster, while classification requires an approximate inference strategy.
- Exact GPs are presented as an appealing option for substantially larger datasets, combining broad scalability with accuracy and limited expert intervention.The conclusion characterizes exact GPs as powerful yet simple and anticipates further scalability from hardware advances.
7 Appendix
Appendix experiments show that exact GPs remain generally more accurate than SGPR and SVGP with independent lengthscales, while also examining Adam-based training comparisons and optimization effort.
- Exact GPs with independent lengthscales: Exact GPs are generally more accurate than SGPR and SVGP when using independent lengthscales across input dimensions.Table 3 reports this setting over medium and large regression datasets, averaging experiments over three splits with pretraining.
- Exact GPs with independent lengthscales: Three-split experiments compare exact GPs with approximate methods on medium and large regression datasets using one independent lengthscale per dimension.The experiments use the same setup as the main paper with pretraining.
- Exact GPs with 100 steps of Adam: Exact GPs were also trained with 100 Adam steps to provide a fair comparison with SGPR and SVGP trained using Adam.The appendix reports this as a separate comparison from the pretraining and finetuning used in the main experiments.
- Exact GPs with 100 steps of Adam: Figure 5 compares exact GPs trained with 100 Adam steps against exact GPs trained partially with Adam.The comparison addresses whether the full 100 optimization steps are necessary.