Batched Conjugate Gradients (mBCG) for Scalable GP Inference

• Batched Conjugate Gradients (mBCG) is a matrix-Krylov method that extends classical CG to solve multiple right-hand sides in parallel, crucial for large-scale Gaussian Process inference.
• The algorithm utilizes on-the-fly Lanczos tridiagonal accumulation to efficiently estimate log-determinants and trace terms, integrating seamlessly with stochastic Lanczos quadrature.
• A specialized low-rank pivoted Cholesky preconditioner is employed to accelerate convergence, reducing the number of iterations and computational complexity on modern GPU hardware.

Batched Conjugate Gradients (mBCG) is a matrix-Krylov method extending the classical Conjugate Gradients (CG) algorithm to support simultaneous solution of multiple right-hand sides and streamlined estimation of matrix functionals. It plays a central role in enabling scalable, highly parallelizable Gaussian Process (GP) inference—especially on modern GPU hardware—by collapsing what would otherwise be multiple serial CG or Lanczos passes into a single, hardware-efficient, batched operation. mBCG directly yields the quantities needed for both predictive posterior and log determinant computations in exact and approximate GP models, and supports stochastic trace and log-determinant estimation in a single sweep. The fundamental advances and practical implications of mBCG are presented below as realized in Blackbox Matrix-Matrix Multiplication (BBMM) inference and related approaches (Gardner et al., 2018).

1. Extension from Standard Conjugate Gradients

Classical CG solves a single linear system $K x = z$ for $x$ with $K \in \mathbb{R}^{n \times n}$ symmetric positive definite, generating iterates $x_0,\ldots,x_p$ that converge to $x^* = K^{-1}z$. mBCG generalizes this to simultaneously solve $t$ right-hand sides, i.e., $K X = Z$ with $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$. In doing so, mBCG further accumulates Lanczos tridiagonal coefficients during the same sweep, enabling on-the-fly estimation of stochastic Lanczos quadrature (SLQ) approximations for log-determinants and trace terms without separate Lanczos passes.

At each iteration $j$, mBCG maintains matrices $X_j, R_j, P_j \in \mathbb{R}^{n \times t}$ for the solution, residual, and search directions. The following matrix-matrix recurrences are computed:

• $x$0 (dense matrix-matrix multiplication, $x$1)
• Step-size vector: $x$2
• Solution update: $x$3
• Residual update: $x$4
• Conjugacy vector: $x$5
• Search direction: $x$6

During iteration, for each right-hand side $x$7, the scalars

• $x$8
• $x$9 populate the diagonals and off-diagonals, respectively, of the partial Lanczos tridiagonal $K \in \mathbb{R}^{n \times n}$0. After $K \in \mathbb{R}^{n \times n}$1 iterations, mBCG yields $K \in \mathbb{R}^{n \times n}$2 and partial tridiagonals for SLQ. This routine enables estimating:
• $K \in \mathbb{R}^{n \times n}$3 via $K \in \mathbb{R}^{n \times n}$4
• $K \in \mathbb{R}^{n \times n}$5 by Hutchinson’s method with random $K \in \mathbb{R}^{n \times n}$6
• $K \in \mathbb{R}^{n \times n}$7 by stochastic Lanczos quadrature on the tridiagonals

The batched nature allows efficient use of GPU Batched-BLAS through large fused matrix-matrix multiplies (Gardner et al., 2018).

2. Algorithmic Structure and Pseudocode

The essential mBCG algorithm proceeds as follows (unpreconditioned):

$X_j, R_j, P_j \in \mathbb{R}^{n \times t}$1

These outputs are then directly used for the primary quantities of GP inference:

• Posterior mean: $K \in \mathbb{R}^{n \times n}$8
• Derivatives: $K \in \mathbb{R}^{n \times n}$9
• Log-determinant: $x_0,\ldots,x_p$0

Preconditioning is introduced by transforming the residuals and inner products as discussed below (Gardner et al., 2018).

3. Low-Rank Pivoted Cholesky Preconditioning

mBCG achieves rapid convergence in practice by utilizing a specialized preconditioner based on the pivoted low-rank Cholesky decomposition. For an $x_0,\ldots,x_p$1 covariance matrix $x_0,\ldots,x_p$2, the rank-$x_0,\ldots,x_p$3 pivoted Cholesky factorization $x_0,\ldots,x_p$4 yields the preconditioning matrix $x_0,\ldots,x_p$5. The preconditioned system is

$x_0,\ldots,x_p$6

However, only $x_0,\ldots,x_p$7 is required within the mBCG recurrences: initial search directions as $x_0,\ldots,x_p$8, and inner products as $x_0,\ldots,x_p$9. Log-determinant computation proceeds via

$x^* = K^{-1}z$0

where the correction term is estimated via SLQ on the preconditioned tridiagonals.

Key computational properties:

• $x^* = K^{-1}z$1 can be computed in $x^* = K^{-1}z$2
• $x^* = K^{-1}z$3 solves in $x^* = K^{-1}z$4 (Woodbury identity)
• Sampling $x^* = K^{-1}z$5 is $x^* = K^{-1}z$6
• $x^* = K^{-1}z$7 in $x^* = K^{-1}z$8

For scalar RBF kernels, the preconditioned condition number satisfies $x^* = K^{-1}z$9, decaying exponentially in $t$0. The CG error accordingly falls as $t$1 (Theorem 4.2 in (Gardner et al., 2018)).

4. Computational Complexity and Parallelism

Each mBCG iteration requires one dense matrix-matrix multiply $t$2 at $t$3 cost for dense kernels, plus $t$4 operations for vector reductions and bookkeeping. For $t$5 iterations with $t$6 probe vectors, total work is $t$7. Cholesky factorization and classical CG require $t$8 and additional separate linear algebra passes for log-determinants and traces.

On modern GPU hardware, the matrix-matrix multiplies in mBCG enable full utilization of hardware (e.g., via cuBLAS GEMM), while runtime complexity for GP inference is reduced from $t$9 to $K X = Z$0. Further complexity reductions are achieved for structured kernel approximations (e.g., SKI, SGPR), where $K X = Z$1 decreases to $K X = Z$2 or $K X = Z$3. mBCG therefore unifies these techniques into a single, scalable, matrix-matrix Krylov routine (Gardner et al., 2018).

5. Practical GPU Implementation Strategies

The practical realization of mBCG for BBMM inference involves:

• Storing all block variables ($K X = Z$4) as $K X = Z$5 device-resident tensors
• Performing kernel matrix-matrix products $K X = Z$6 using user-supplied GPU routines (either recomputed on the fly or via fast operator representations)
• Allocating all workspace buffers ($K X = Z$7) a priori to prevent allocation overhead during iteration
• Compactly storing tridiagonal data as $K X = Z$8, where $K X = Z$9 and $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$0 require negligible memory
• Achieving $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$1 TFLOP/s performance per GEMM call by leveraging hardware-batched BLAS kernels

This approach ensures that each CG iteration corresponds to a single fused GEMM plus $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$2 sequence of small vector operations, maximizing hardware efficiency and enabling large-scale GP inference (Gardner et al., 2018).

6. Empirical Performance and Analytical Results

Empirical results on a Titan Xp GPU demonstrate that:

• For exact GPs ($Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$3), BBMM with mBCG is up to $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$4 faster than CPU-based Cholesky, and $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$5 faster than GPU-based Cholesky
• For SGPR with $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$6, BBMM achieves $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$7 speedup over CPU Cholesky
• For SKI+DKL with $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$8, BBMM is up to $Z = [z_1, ..., z_t] \in \mathbb{R}^{n \times t}$9 faster than state-of-the-art CPU-based MVM+Lanczos inference

Convergence curves indicate that unpreconditioned mBCG can require $j$0–100 iterations for $j$1 relative residuals, while a rank-5 Cholesky preconditioner reduces this to $j$2–20, and increasing $j$3 to $j$4 halves the required iterations further. Test error evaluations match or slightly improve upon Cholesky-based inference, with early-terminated CG occasionally regularizing tiny eigenvalues (Gardner et al., 2018).

7. Connections to Batched and Block Conjugate Gradient Variants

mBCG (as in BBMM) and the closely related cooperative CG (cCG) both generalize classical CG by introducing matrix-valued step sizes, updating multiple descent and conjugate directions in parallel. cCG, analyzed by Bhaya et al. (Bhaya et al., 2012), introduces block updates in $j$5 via $j$6 and enforces block (bi)-orthogonality and $j$7-conjugacy. The per-iteration cost per thread is $j$8 with $j$9 worst-case complexity given $X_j, R_j, P_j \in \mathbb{R}^{n \times t}$0. mBCG, as adopted in GP inference, prioritizes matrix-matrix operations to maximize GPU efficiency and tightly integrates with stochastic Lanczos quadrature for scalable kernel learning.

Both mBCG and cCG maintain the Krylov-optimality and finite step termination characteristics of classic CG, but are tailored for parallel and batched environments, with design tradeoffs for block orthogonalization and rank preservation (Bhaya et al., 2012, Gardner et al., 2018).