PaperlayerView exact revision on arXiv

Source-linked AI summary

GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration

Jacob R. Gardner, Geoff Pleiss, David Bindel, Kilian Q. Weinberger, Andrew Gordon Wilson

arXiv:1809.11165v6cs.LGstat.ML

TL;DR

GP inference tools underuse modern hardware and often require model-specific procedures. The paper introduces BBMM, which uses batched conjugate gradients and blackbox matrix-matrix products to compute inference quantities, with pivoted-Cholesky preconditioning. BBMM reduces exact-inference complexity and accelerates exact and scalable GP models, while GPyTorch provides the implementation platform.

  • Problem

    GP inference tools do not effectively use modern hardware, and tightly coupling models with inference procedures makes complex GP models require custom implementations.

  • Method

    BBMM uses modified batched conjugate gradients and requires only matrix-matrix multiplication routines for the kernel matrix and its derivative.

  • Results

    BBMM reduces exact GP inference complexity from O(n3) to O(n2) and achieves up to 20× exact-GP speedups, with up to 15× and 4× speedups for SKI and SGPR.

  • Takeaways & Limitations

    GPyTorch implements BBMM and several state-of-the-art GP models, enabling scalable GP approximations through efficient matrix-matrix multiplication routines.

  • Takeaways & Limitations

    The paper primarily focuses on regression, while exact details of ELBO derivations for non-Gaussian likelihoods are left to other work.

Abstract

from arXiv · show

Despite advances in scalable models, the inference tools used for Gaussian processes (GPs) have yet to fully capitalize on developments in computing hardware. We present an efficient and general approach to GP inference based on Blackbox Matrix-Matrix multiplication (BBMM). BBMM inference uses a modified batched version of the conjugate gradients algorithm to derive all terms for training and inference in a single call. BBMM reduces the asymptotic complexity of exact GP inference from $O(n^3)$ to $O(n^2)$. Adapting this algorithm to scalable approximations and complex GP models simply requires a routine for efficient matrix-matrix multiplication with the kernel and its derivative. In addition, BBMM uses a specialized preconditioner to substantially speed up convergence. In experiments we show that BBMM effectively uses GPU hardware to dramatically accelerate both exact GP inference and scalable approximations. Additionally, we provide GPyTorch, a software platform for scalable GP inference via BBMM, built on PyTorch.

1 Introduction

GP inference tools have not kept pace with modern hardware and often couple model specifications to custom inference procedures. BBMM addresses this gap with matrix-matrix computation, a modified conjugate-gradient method, preconditioning, and GPyTorch.

  • Motivation: GP inference tools underuse modern hardware and require substantial implementation effort for new models.Model and inference-engine entanglement makes complex models such as multi-output GPs and scalable approximations require custom procedures.
  • BBMM framework: BBMM requires only a blackbox routine for matrix-matrix multiplication with the kernel matrix and its derivative.This replaces the need to provide routines computing the full GP marginal log likelihood for complex models.
  • Algorithm: BBMM uses modified batched conjugate gradients to compute all quantities needed for the marginal likelihood and its derivatives.The method performs large matrix-matrix multiplications rather than relying on separate inference computations.
  • Complexity: BBMM reduces exact GP inference time complexity from O(n3) to O(n2).The approach also addresses space-complexity and numerical-stability issues in existing inference methods.
  • Preconditioning: A pivoted-Cholesky preconditioner significantly accelerates inference in practice and theory.The required preconditioner operations are efficient and require negligible time in practice.
  • Empirical results: Exact GPs achieve up to 20× speedups, while BBMM versions of SKI and SGPR achieve up to 15× and 4× speedups, respectively.The reported datasets contain up to 3,000 points for exact GPs and 500,000 points for scalable approximations.

2 Related Work

Related work applies iterative numerical-linear-algebra methods and preconditioning to GP inference, exploiting structure but retaining limitations in generality, parallelism, complexity, or stability.

  • Iterative methods: Conjugate gradients and Lanczos methods solve linear systems and eigenvalue problems without explicitly computing matrices.They are Krylov-subspace methods that access matrices through matrix-vector multiplications.
  • MVM-based GP inference: MVM-based GP methods exploit algebraic structure and can achieve improved asymptotic efficiency and space usage.SKI uses structured kernel matrices with fast MVMs, while Lanczos methods estimate log determinants and derivatives stochastically.
  • Preconditioning: Existing preconditioners do not provide a general-purpose solution for scalable GP inference.Jacobi preconditioning has no effect for stationary kernels, and many alternatives have Ω(n2) complexity.
  • Pivoted Cholesky: Pivoted Cholesky has been studied as a low-rank approximation, whereas this paper uses it primarily as a preconditioner.That use avoids accuracy loss from low-rank approximation and the complexity of computing derivatives.

3 Background

The background defines Gaussian processes through kernels and introduces the notation, predictive posterior, and negative log marginal likelihood used for inference and training.

  • Notation: X denotes n training examples in d dimensions, y denotes training labels, and KXX contains all pairwise kernel evaluations.The notation also defines kXx∗ for training-to-test kernel values and adds diagonal noise through bKXX = KXX + σ2I.
  • Gaussian processes: A Gaussian process defines a full distribution over the modeled function using a mean function and kernel.The passage names RBF and Matérn kernels as popular examples.
  • Prediction: GP predictions use the predictive posterior distribution conditioned on training inputs and labels.The predictive mean and covariance are defined for test inputs and pairs of test inputs.
  • Training: GP hyperparameters are learned by minimizing or sampling from the negative log marginal likelihood and its derivative.Examples include likelihood noise, kernel lengthscale, inducing-point locations, and neural-network parameters for deep kernel learning.

4 Gaussian process inference through blackbox matrix multiplication

BBMM unifies GP inference around batched matrix-matrix operations: modified conjugate gradients compute solves, trace terms, and log-determinant estimates in one call, while pivoted-Cholesky preconditioning improves convergence.

  • Framework: BBMM reduces GP inference to blackbox matrix-matrix multiplication routines for the kernel matrix and its derivative.This design targets efficient use of modern hardware while avoiding model-specific inference rules.
  • Inference operations: GP inference is dominated by linear solves, log determinants, and trace terms, traditionally computed through an expensive Cholesky decomposition.Cholesky-based computation requires O(n3) operations and does not effectively use parallel hardware.
  • Prior methods: MVM-based methods improve asymptotic and space efficiency but compute quantities through separate, sequential routines and can suffer Lanczos stability issues.Lanczos also requires O(np) space for p iterations and may lose orthogonality.
  • Modified CG: mBCG simultaneously solves multiple right-hand sides and returns partial Lanczos tridiagonalizations for those inputs.It takes a matrix containing y and random probe vectors as input.
  • Inference estimates: A single mBCG call supplies the GP solve, stochastic trace-estimation terms, and log-determinant estimate.The method derives these quantities from returned solves, derivative matrix products, and tridiagonalization terms.
  • Lanczos terms: mBCG avoids separate Lanczos iterations by recovering tridiagonalization terms from conjugate-gradient coefficients.This avoids extra computation, storage, and numerical instability associated with Lanczos iterations.
  • Runtime and space: Each mBCG iteration uses one kernel matrix-matrix multiply, with O(nt) space and O(p Ξ(bKXX)) time.For a standard matrix, the multiplication costs O(n2t), lower than the O(n3) complexity of Cholesky inference.
  • Preconditioning: Pivoted-Cholesky preconditioning is designed to provide efficient solves, log determinants, and Gaussian probe-vector sampling.Its operations cost O(nk2) for solves and log determinants and O(nk) for sampling, while empirically accelerating CG convergence.

5 Programmability with BBMM

BBMM makes complex and structured GP models programmable through blackbox matrix-matrix multiplication, while GPyTorch provides implementations built on this framework.

  • BBMM adapts readily to complex GP models and structured GP approximations.Its blackbox design requires only matrix multiplications with kernel matrices and their derivatives.
  • GPyTorch implements BBMM and several scalable GP models in a publicly available software library.The implementations include SGPR and KISS-GP-style models.
  • Bayesian linear regression fits BBMM with O(ptnd) time and exact computation in O(tnd^2) time.The kernel multiply costs O(tnd), matching existing efficient algorithms for Bayesian linear regression.
  • SGPR matrix-matrix multiplication costs O(tnm+tm^3), asymptotically faster than Cholesky inference at O(nm^2+m^3).This follows from distributing vector multiplications and grouping terms in the subset-of-regressors approximation.
  • SKI is a natural BBMM candidate because its structured kernel matrices provide fast matrix-vector multiplies and benefit from hardware acceleration.SKI specifies K_XU ≈ W K_UU using a sparse interpolation matrix W.
  • Kernel compositions can be handled automatically by composing the corresponding blackbox multiplication routines.For example, (K1K2 + K3)M becomes K1(K2M) + K3M.

6 Results

Experiments compare BBMM with Cholesky and other inference procedures across exact and scalable GP models, finding substantial speed and convergence benefits without worse final test error.

  • Baselines and datasets: The evaluation covers exact GPs, SGPR inducing-point models, and SKI models with Toeplitz K_UU and deep kernels.Exact and SGPR models are compared with GPFlow Cholesky engines, while SKI is compared with Dong et al.'s procedure.
  • Baselines and datasets: The datasets span exact models with up to 3500 training examples, SGPR with up to 50000, and SKI with up to 515000.Experiments use five datasets for each model family, with larger datasets assigned to scalable approximations.
  • Error comparison: BBMM is at least as accurate as Cholesky across all reported datasets and improves final test error on some exact-GP datasets.The reported improvements include Gas, Airfoil, and Wine with Exact GPs.
  • Error comparison: BBMM and Dong et al.'s SKI inference return identical outputs even though BBMM is faster.SKI models are therefore excluded from the Figure 3 error comparison.
  • Preconditioning: Increasing pivoted-Cholesky preconditioner rank substantially reduces the CG iterations required for convergence.The comparison uses no preconditioner and ranks 2, 5, and 9 on Protein and KEGG with deep RBF and Matern-5/2 kernels.
  • Preconditioning: Rank-5 preconditioning yields more accurate solves with virtually no impact on each CG iteration's runtime.The authors recommend using the pivoted-Cholesky preconditioner with BBMM because it rapidly accelerates convergence with negligible wall-clock overhead.

7 Discussion

The paper presents BBMM and GPyTorch as a hardware-efficient framework for GP inference, while noting compatibility with variational methods and the benefits of avoiding full Cholesky decomposition.

  • BBMM is a blackbox matrix-matrix framework implemented with several state-of-the-art GP models in GPyTorch.
  • BBMM is compatible with variational techniques for non-Gaussian likelihoods, although the paper focuses primarily on regression.A single mBCG call can compute the KL-divergence term that is the most computationally intensive part of the ELBO.
  • Avoiding full Cholesky decomposition reduces exact GP inference from O(n^3) to O(n^2) through conjugate gradients.The paper also reports that CG may provide better linear solves while early pivoted-Cholesky termination avoids bottlenecks and numerical instabilities.
  • The framework is intended to reduce the implementation complexity of new GP models while retaining efficient inference.

Supplementary Information for: GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference

The supplementary information identifies the paper's authors and their Cornell University affiliations and contact addresses.

  • The paper is authored by Jacob R. Gardner, Geoff Pleiss, David Bindel, Kilian Q. Weinberger, and Andrew Gordon Wilson.
  • The authors are affiliated with Cornell University.

A Analysis of the modified CG algorithm, mBCG

mBCG extends conjugate gradients from single-vector solves to batched matrix-matrix solves while also recovering Lanczos tridiagonalization matrices. This supplies the parallel computations needed for GP inference terms in one procedure.

  • Standard CG foundations: CG accesses A only through matrix-vector products, with p iterations costing O(p ξ(A)) and convergence governed more by conditioning than matrix size.The approximation can reach near-machine precision in p ≪ n iterations in favorable settings.
  • Batched conjugate gradients: mBCG simultaneously solves A^-1B for multiple right-hand sides using matrix-matrix multiplication.Its matrix operations are batched adaptations of standard conjugate-gradient updates.
  • Lanczos connection: mBCG returns Lanczos tridiagonalization matrices alongside the batched solves for estimating log determinants.The tridiagonal matrices are formed from the α_j and β_j coefficients generated during the iterations.
  • Batched conjugate gradients: Each solve has its own coefficient values, represented by coefficient vectors α_j and β_j.The algorithm updates and scales solution, residual, and search-direction matrices using these per-solve coefficients.
  • Lanczos connection: The Lanczos connection lets CG coefficients recover the same tridiagonal matrices produced by p Lanczos iterations.For systems with probe vectors b_1,...,b_t, mBCG recovers one corresponding tridiagonal matrix per probe.

B Runtime analysis of computing inference terms with mBCG

mBCG concentrates inference-term computation around repeated matrix-matrix multiplications, with low-cost postprocessing for solves, trace estimates, and log determinants. The remaining runtime and storage are dominated by the batched iterative computation.

  • mBCG runtime: mBCG performs one kernel matrix-matrix multiply before its loop and one during each iteration, giving runtime at least O(p Ξ(K̂_XX)).Ξ(K̂_XX) denotes the cost of multiplying the kernel matrix by an n × t matrix.
  • mBCG runtime: Operations on the n × t iterate matrices cost O(nt) per iteration and are lower-order than the matrix multiplications.This includes elementwise products and diagonal scaling of the matrix columns.
  • Inference-term recovery: The solve term requires no additional work beyond mBCG because the required matrix product is already its first output.Random-vector products and associated tridiagonal matrices support the remaining inference estimates.
  • Space complexity: The additional space for inner products is negligible compared with the dominant batched computation.The extra storage consists of 2t length-n vectors.
  • Log-determinant computation: The post-mBCG log-determinant computation costs O(tp^2) for eigendecomposing t tridiagonal matrices and uses O(tp^2) space.The eigendecomposition is performed separately for each p × p tridiagonal matrix.

C.1 Running time of the pivoted Cholesky decomposition.

The pivoted Cholesky preconditioner has efficient construction and application costs across exact and structured kernel settings. Its dominant operations scale with the retained rank rather than the full matrix factorization.

  • Exact-kernel setting: A rank-k pivoted Cholesky decomposition can be computed in O(ρ(K_XX)k^2) time.ρ(K_XX) is the time required to retrieve one row of K_XX.
  • Preconditioner operations: Solves with the preconditioner L_kL_k^T + σ^2I require O(nk^2) time.This is one of the stated properties of the pivoted Cholesky preconditioner.
  • Preconditioner operations: The preconditioner log determinant can also be computed in O(nk^2) time.The analysis uses the matrix determinant lemma for this computation.
  • Exact-kernel setting: For a standard matrix, constructing the decomposition costs O(nk^2), negligible relative to an O(n^2) matrix-vector multiplication.The construction reads the diagonal and k rows of K_XX.
  • Scalable approximations: Under SoR, rank-k construction costs O(nmk^2), roughly matching one MVM when k^2 ≤ m or k^2 ≈ m.Under SKI, accessing a row of K_XX costs O(n) using the interpolation structure.

D Convergence Analysis of Pivoted Cholesky Preconditioned CG

The convergence analysis links pivoted Cholesky quality to kernel eigenvalue decay and then to preconditioned CG conditioning. For univariate RBF kernels, superexponential eigenvalue decay supports rapid low-rank approximation and convergence.

  • Preconditioned CG: The CG error after k iterations is bounded using the condition number of the pivoted-Cholesky-preconditioned system.The exact solution is denoted u* = K̂_XX^-1y.
  • Eigenvalue decay: Pivoted Cholesky rapidly approximates a matrix when its eigenvalues decay exponentially.The analysis applies this principle to univariate RBF kernel matrices.
  • Eigenvalue decay: Univariate RBF kernel matrices have superexponentially decaying eigenvalues.The bound is established using modified Bessel functions.
  • Preconditioned CG: The resulting pivoted Cholesky preconditioner is analyzed through the condition number of P̂_k^-1 K̂_XX.Bounding this condition number connects the low-rank approximation error to CG convergence.
  • Preconditioned CG: The eigenvalue result yields a convergence bound containing 1 + O(exp(kb)/n).This bound is obtained after combining the preconditioner condition-number result with the standard CG convergence bound.
Loading 1809.11165v6…