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

1 Introduction

• GP inference tools often underuse modern hardware and tightly couple model specification with inference, forcing complex models and scalable approximations to use custom procedures.The paper motivates matrix-matrix operations as a way to use GPU acceleration and ease implementation across a wide class of GP models.Source: §1 · pp. 1–2
• BBMM requires only a blackbox routine for matrix-matrix multiplication with the kernel and its derivative, easing implementation across complex GP models.This contrasts with approaches that require routines to compute the full GP marginal log likelihood; existing model-specific inference engines can impede rapid prototyping.Source: §1 · pp. 1–2
• Modified batched conjugate gradients computes the marginal likelihood and its derivatives, uses hardware-efficient matrix-matrix operations, and reduces exact GP inference time complexity from O(n^3) to O(n^2).The approach also addresses space-complexity and numerical-stability issues in existing methods.Source: §1 · p. 2
• A preconditioner based on pivoted Cholesky has efficient operations that require negligible time in practice and significantly accelerates inference.The paper reports empirical and theoretical evidence for the acceleration.Source: §1 · p. 2
• Exact BBMM GPs are up to 20× faster than Cholesky on datasets up to 3,000 points; BBMM SKI and SGPR yield up to 15× and 4× speedups on datasets as large as 500,000 points.The paper also reports that SKI, SGPR, and other scalable approximations can be implemented in fewer than 50 lines of code using an efficient matrix-matrix multiplication routine.Source: §1 · p. 2
• GPyTorch is a PyTorch-based software platform for scalable Gaussian processes that uses BBMM inference.Its implementations of SKI, SGPR, and other scalable approximations require only an efficient matrix-matrix multiplication routine.Source: §1 · p. 2

• Conjugate gradients and Lanczos solve linear systems or eigenproblems through matrix-vector multiplies, historically for large numerical problems, especially those involving sparse matrices [11, 12, 17, 33, 37, 38, 43].Source: §2 · p. 2
• Prior GP work uses matrix-vector methods for efficiency [13] [15] [40] [51]; this paper uses Dong et al.’s [13] log-determinant estimator while avoiding Lanczos storage and stability issues.SKI [51] uses structured kernel matrices with fast matrix-vector multiplies to achieve improved asymptotic complexity.Source: §2 · p. 2
• Preconditioning accelerates conjugate-gradient convergence, but Cutajar et al.’s [10] GP preconditioners are not general-purpose: Jacobi has no effect for stationary kernels, and many alternatives cost Ω(n^2).Source: §2 · p. 2
• Unlike prior low-rank uses of pivoted Cholesky [4], this paper uses it primarily as a preconditioner, avoiding approximation accuracy loss and derivative-computation complexity.Harbrecht et al. studied it mainly in scientific computing, and the paper uses theoretical results from to prove preconditioner convergence bounds; Bach [4] treated it as an approximate training method.Source: §2 · pp. 2–3

3 Background

• X and y denote training examples and labels; KXX contains pairwise kernel values, while adding σ^2I gives the noisy kernel matrix.The notation also defines xi as the ith training example and kXx∗ as kernel values between training examples and a test point.Source: §3 · p. 3
• A Gaussian process assigns a full function distribution, with popular kernels including RBF and Matérn.Source: §3 · p. 3
• GP predictions use the predictive posterior to obtain a test-point predictive mean and covariance between two test inputs.Source: §3 · p. 3
• GP hyperparameters, including likelihood noise and kernel lengthscale, are learned by minimizing or sampling the negative log marginal likelihood.Examples also include inducing point locations and neural network parameters for deep kernel learning [53].Source: §3 · p. 3

4 Gaussian process inference through blackbox matrix multiplication

• BBMM reduces inference to a user-supplied matrix-multiplication routine for the kernel and its derivative, without requiring additional inference rules for complex GP models.Its goal is a unified framework that uses modern hardware efficiently.Source: §4 · p. 3
• One modified batched conjugate-gradients (mBCG) call computes the predictive linear solve, log determinant, and trace term needed for GP inference.mBCG handles multiple right-hand sides simultaneously and returns partial Lanczos tridiagonalizations alongside the solves.Source: §4 · pp. 3–4
• mBCG derives the log-determinant estimate from its tridiagonalization terms, avoiding separate Lanczos iterations and their extra computation, storage, and numerical instability.The tridiagonal terms are obtained from conjugate-gradients coefficients with O(1) additional work per iteration.Source: §4 · p. 4
• For p iterations, BBMM uses O(nt) space and O(p Ξ(K̂XX)) time, or O(n²t) per matrix multiply, versus O(n³) for standard Cholesky inference.The authors note that structured data or sparse GP approximations can further reduce the time complexity.Source: §4 · p. 5

4.1 Preconditioning

• A BBMM preconditioner must support roughly linear-time solves and linear space, efficient log-determinant computation, and sampling probe vectors from N(0, P).These requirements let preconditioning avoid dominating the inference routine's running time while correcting the log-determinant estimate.Source: §4.1 · p. 5
• Pivoted-Cholesky preconditioning provides O(nk²) solves and log determinants and O(nk) probe sampling; increasing its rank improves CG convergence and makes the final log-determinant estimate more exact.For univariate RBF kernels, the paper states that CG convergence is expected to improve exponentially with rank; as k approaches n, the log determinant of the preconditioned system tends to zero and the final estimate approaches the exact value.Source: §4.1 · p. 6

• BBMM is black-box: implementing a model requires only kernel and derivative matrix-multiplication routines, and examples need at most 50 lines of Python in GPyTorch.The resulting implementations are available through the GPyTorch library.Source: §5 · p. 7
• Bayesian linear regression uses BBMM in O(ptnd) time and is exact in O(tnd^2), matching efficient existing algorithms; multitask GPs adapt similarly.The linear-regression implementation needs no additional derivation.Source: §5 · p. 7
• For SoR-based SGPR, matrix multiplication costs O(tnm+tm^3), asymptotically faster than Cholesky inference at O(nm^2+m^3); FITC diagonal correction is straightforward.The same passage identifies SGPR and other sparse GP techniques as using the subset-of-regressors approximation.Source: §5 · p. 7
• SKI matrix multiplies cost O(tn+tm^2) without K_UU structure, while Toeplitz KISS-GP reduces them to O(tn+tm log m), with Toeplitz MVMs costing O(m log m).KISS-GP can also use Kronecker structure in K_UU to accelerate matrix-vector multiplies.Source: §5 · p. 7
• Kernel compositions can be evaluated automatically from component BBMM routines; SGPR and KISS-GP use this approach, and product kernels follow Gardner et al. [15].With predefined compositional strategies, SGPR reduces multiplication to applying K_UU^-1, while KISS-GP uses multiplication by Toeplitz K_UU.Source: §5 · p. 7

Results

• BBMM provides substantial speed benefits over Cholesky and standard MVM-based CG inference, especially on GPUs, across Exact, SGPR, and SKI comparisons.Exact and SGPR are compared with GPFlow Cholesky engines; SKI is compared with Dong et al.’s inference procedure.Source: §6 · p. 7, Figure 1 · p. 8

Error comparison

• Across Exact and SGPR models and both tested kernels, BBMM matches or improves final test MAE versus Cholesky, with gains on Gas, Airfoil, and Wine for Exact GPs.The paper attributes possible improvements to CG regularizing effects; SKI models are excluded because BBMM and Dong et al.’s procedure return identical outputs.Source: §6 · p. 8, Figure 3 · p. 8
• For SKI, BBMM and Dong et al.’s procedure produce identical outputs, although BBMM is faster.SKI results are therefore excluded from the test-MAE figure.Source: §6 · p. 8
• Unlike Cholesky, BBMM avoids adding diagonal jitter for stability without requiring double precision, which can reduce numerical instability at increased computational cost.Cholesky may add jitter; the paper notes double precision as an alternative and points to Figure 2.Source: Figure 2 · p. 8, §6 · p. 9

Preconditioning

• Increasing pivoted-Cholesky preconditioner rank substantially reduces CG iterations to convergence on Protein and KEGG, with deep RBF and deep Matern-5/2 kernels.The evaluation compares no preconditioner with ranks 2, 5, and 9.Source: §6 · p. 9, Figure 4 · p. 9
• Rank-5 preconditioning yields more accurate solves and test MAE at virtually no per-iteration runtime cost, so the authors recommend using it with BBMM.Prediction MAE is compared against total wall-clock time while varying CG iterations.Source: §6 · p. 9, Figure 4 · p. 9

• The paper presents BBMM, a GP-inference framework based on blackbox matrix-matrix multiplication with kernel matrices, implemented alongside state-of-the-art GP models in the publicly available GPyTorch package.Source: §7 · p. 9

Non-Gaussian likelihoods

• For non-Gaussian likelihoods, BBMM supports variational techniques such as [54], and one mBCG call computes the ELBO’s most computationally intensive term: the KL divergence between multivariate Gaussians.These approaches compute the variational lower bound (ELBO) rather than the GP marginal log likelihood.Source: §7 · p. 9

Avoiding the Cholesky decomposition

• Cholesky decomposition takes O(n^3), while CG takes O(n^2); CG may even provide better linear solves, and early pivoted preconditioning avoids the bottleneck and many numerical instabilities.The preconditioner computes only the first five rows of the pivoted decomposition.Source: §7 · p. 10

• The authors hope the work will reduce the complexity of implementing new GP models while enabling efficient inference.Source: §7 · p. 10

A Analysis of the modified CG algorithm, mBCG

• CG accesses A only through matrix-vector multiplications, so p iterations cost O(p ξ(A)); in practice, accurate solves can require p ≪ n iterations, with convergence depending more on conditioning than matrix size.For positive definite A, the cited convergence observation bounds error in terms of the condition number.Source: A Analysis of the modified CG… · p. 14
• mBCG simultaneously solves A^-1B = [A^-1b1, …, A^-1bt] with matrix-matrix multiplication and returns a Lanczos tridiagonal matrix for each solve to estimate log determinants.It modifies standard CG by batching multiple solves and returning the associated tridiagonal matrices.Source: A Analysis of the modified CG… · p. 14

A.1 Adaptation to multiple right hand sides.

• Replacing each right-hand-side matrix-vector multiply with one matrix-matrix multiply lets mBCG perform multiple solves at once.Each column’s result is equivalent to applying the standard CG matrix-vector operation separately.Source: §A.1 · p. 15
• For each right-hand side, mBCG uses its own α and β coefficients, stored as vectors, while batched updates and scaling preserve the individual standard-CG recurrences.The coefficient vectors allow all solve operations to be performed in batch.Source: §A.1 · p. 15

A.2 Obtaining Lanczos tridiagonal matrices from mBCG.

• After p standard preconditioned CG iterations on A^-1z, the iteration coefficients α and β determine the same approximate Lanczos tridiagonal matrix as p Lanczos iterations.This established connection lets mBCG recover the matrix by reusing coefficients generated during its solves.Source: §A.2 · p. 16
• mBCG uses each solve’s α and β coefficient vectors to form its tridiagonal matrix, corresponding to a Lanczos run with that solve’s right-hand side as probe.For systems with right-hand sides b1, …, bt, the recovered matrices correspond to probes b1, …, bt.Source: §A.2 · p. 16

B Runtime analysis of computing inference terms with mBCG

• mBCG runs in O(pΞ(K)) time and uses O(nt) additional space beyond storing K, because its other operations are dominated by kernel matrix–matrix multiplies.Here Ξ(K) is the time to multiply K by an n × t matrix.Source: B Runtime analysis of computing… · p. 16
• Post-mBCG recovery of the solve requires no extra work, and trace recovery adds O(tn) time, leaving runtime dominated by mBCG.The trace step’s additional space for 2t length-n vectors is described as negligible.Source: B Runtime analysis of computing… · pp. 16–17
• Estimating log |K| after mBCG takes O(tp^2) time for tridiagonal eigendecompositions and O(tp^2) space for eigenvectors.Each tridiagonal eigendecomposition takes O(p^2) time; computing the remaining eigenvalue expression adds only O(p) work per matrix.Source: B Runtime analysis of computing… · p. 17

C The Pivoted Cholesky Decomposition

• Standard Cholesky can be viewed as building progressively higher-rank approximations to K by recursively factoring blocks and their Schur complements.The derivation obtains a leading block recursively, solves for the off-diagonal factor, and factors the Schur complement.Source: C The Pivoted Cholesky Decomposition · p. 17
• Pivoted Cholesky permutes rows and columns so the next pivot is the largest diagonal entry of the current Schur complement, aiming to reduce its norm.The first permutation moves the maximum-diagonal entry of K to the upper-left position.Source: C The Pivoted Cholesky Decomposition · p. 18

C.1 Running time of the pivoted Cholesky decomposition.

• O(ρ(KXX)k^2) time computes rank-k pivoted Cholesky; for a standard positive-definite matrix, this is O(nk^2), negligible beside an O(n^2) matrix-vector multiply.Computing the factor requires reading the matrix diagonal and k rows.Source: §C.1 · p. 18
• O(nk^2) time suffices for both linear solves with the preconditioner and its log determinant.The solve uses the Woodbury formula; the log determinant uses the matrix determinant lemma.Source: §C.1 · pp. 18–19
• O(nmk^2) time computes rank-k pivoted Cholesky for the SoR approximation in SGPR, costing roughly one MVM when k^2 ≤ m or k^2 ≈ m.For a general approximation, the cost is O(ρ(A)k^2), where ρ(A) is the time to access a row of A.Source: §C.1 · p. 19
• O(nk^2) time computes pivoted Cholesky for an SKI kernel matrix, comparable to its MVM complexity, which is linear in n.The row-access cost follows from the sparse interpolation matrix W and constant-time access to elements of KUU.Source: §C.1 · p. 19
• O(nk) time is required for the matrix-vector multiplication by Lk when computing the stated samples.Source: §C.1 · p. 19

D Convergence Analysis of Pivoted Cholesky Preconditioned CG

• For univariate RBF kernels, the section establishes a convergence bound for CG preconditioned by rank-k pivoted Cholesky.The theorem considers the kth CG solution and its error relative to the exact solution.Source: D Convergence Analysis of Pivoted… · pp. 19–20
• The section proves that univariate RBF kernel-matrix eigenvalues decay superexponentially, so the pivoted-Cholesky convergence result applies.It invokes Harbrecht et al.'s result relating eigenvalue decay to rapid convergence of the rank-k pivoted Cholesky decomposition.Source: D Convergence Analysis of Pivoted… · p. 20
• A condition-number bound for the preconditioned system connects the pivoted Cholesky approximation to the rate of CG convergence.The section relates this bound to the standard CG convergence result.Source: D Convergence Analysis of Pivoted… · p. 20

E.1 Proof of Lemma 1

• For a univariate RBF kernel matrix, the proof defines E as the residual after its rank-k pivoted Cholesky approximation.The proof then applies Cauchy–Schwarz and the triangle inequality in bounding the residual.Source: §E.1 · pp. 20–21
• Harbrecht et al. show that E is positive semidefinite, so its spectral norm is at most its trace.The proof recalls that Tr(E) is the trace of the difference between KXX and its pivoted Cholesky approximation.Source: §E.1 · p. 21

E.2 Proof of Lemma 3

• If a degree-d polynomial approximates exp(−γr^2) within ε on [−1, 1], the proof establishes λ_{d+1}(KXX) ≤ nε.This is the polynomial-approximation step used in the proof of Lemma 3.Source: §E.2 · p. 21
• A degree-2l truncated Chebyshev expansion approximates exp(−γr^2) with error below 2e^(−γ/4)I_{l+1}(γ/4).The proof bounds the expansion tail using |(−1)^jT_{2j}(x)| ≤ 1 on [−1, 1].Source: §E.2 · p. 21
• The auxiliary Bessel-function bound is obtained by expanding in power series, reordering absolutely convergent sums, and rewriting an inner sum as an exponential-series tail.The proof overview says this bound combines with the Chebyshev approximation to yield the theorem.Source: §E.2 · pp. 21–22
• The polynomial kernel matrix has rank at most d, while entrywise errors bounded by ε give a spectral-norm error at most nε and hence the eigenvalue bound.The proof uses the low rank of the sampled polynomial matrix and bounds its error matrix using Gershgorin’s circle theorem.Source: §E.2 · p. 21