[Paper Review] GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration

Jacob R. Gardner, Geoff Pleiss+3 more|arXiv (Cornell University)|Sep 28, 2018

Gaussian Processes and Bayesian Inference33 references536 citations

TL;DR

Introduces BBMM inference, a blackbox matrix-matrix approach for Gaussian processes that uses batched modified conjugate gradients and GPU acceleration, enabling faster exact and approximate GP inference with a PyTorch-based framework (GPyTorch).

ABSTRACT

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.

Motivation & Objective

• Motivate the need for hardware-efficient GP inference tools that decouple model specification from inference procedures.
• Develop a blackbox framework (BBMM) that relies on kernel matrix-matrix multiplies rather than Cholesky factorizations.
• Provide a scalable, GPU-friendly GP inference engine that supports exact GPs and popular approximations (SGPR, SKI).
• Offer a software platform (GPyTorch) built on PyTorch to simplify prototyping of complex GP models.

Proposed method

• Use a modified batched conjugate gradients (mBCG) to compute all inference terms (K^{-1}y, log|K|, and Tr(K^{-1} dK/dθ)) in a single call.
• Estimate log-determinants and traces via stochastic trace estimation using probe vectors z_i.
• Employ a pivoted Cholesky preconditioner P = L_k L_k^T + σ^2 I to accelerate CG convergence and enable efficient log-determinant corrections.
• Show that BBMM reduces exact GP inference complexity from O(n^3) to O(n^2) and integrates with SGPR and SKI frameworks.
• Demonstrate that BBMM exploits GPU hardware to achieve substantial speedups over Cholesky-based methods across multiple GP models and datasets.

Experimental results

Research questions

• RQ1 Can BBMM inference match or exceed the accuracy of Cholesky-based GP inference across exact and approximate GP models?
• RQ2 How does modified batched CG (mBCG) with pivoted Cholesky preconditioning perform in terms of convergence speed and scalability on GPU hardware?
• RQ3 To what extent can BBMM be applied as a blackbox framework to a range of GP models and scalable approximations (SGPR, SKI, Toeplitz/KISS-GP) with minimal derivation effort?

Key findings

• BBMM on GPUs dramatically accelerates both exact GP inference and scalable approximations compared to CPU Cholesky-based methods.
• Exact GPs with BBMM can be up to 20× faster than Cholesky-based approaches on datasets up to about 3000 points (limited by GPU memory).
• SGPR and SKI with BBMM achieve up to 15× and 4× speedups respectively on datasets up to 500,000 points.
• Preconditioning with pivoted Cholesky significantly speeds up convergence of CG in the BBMM framework.
• The BBMM approach enables implementation of many GP models with fewer lines of code (often <50) by relying on efficient kernel matrix multiplies.