---
title: Batched Conjugate Gradients (mBCG) for Scalable GP Inference
url: https://www.emergentmind.com/topics/batched-conjugate-gradients-mbcg
type: topic
---

# Batched Conjugate Gradients (mBCG) for Scalable GP Inference

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 [1809.11165].

## 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:

- $V_j = K P_j$ (dense matrix-matrix multiplication, $n \times t$)
- Step-size vector: $\alpha_j = \operatorname{diag}(R_j^\top R_j) \oslash \operatorname{diag}(P_j^\top V_j)$
- Solution update: $X_{j+1} = X_j + P_j\,\operatorname{Diag}(\alpha_j)$
- Residual update: $R_{j+1} = R_j - V_j\,\operatorname{Diag}(\alpha_j)$
- Conjugacy vector: $\beta_j = \operatorname{diag}(R_{j+1}^\top R_{j+1}) \oslash \operatorname{diag}(R_j^\top R_j)$
- Search direction: $P_{j+1} = R_{j+1} + P_j\,\operatorname{Diag}(\beta_j)$

During iteration, for each right-hand side $i$, the scalars
- $\gamma_j^{(i)} = 1/\alpha_j[i] + \beta_{j-1}[i]/\alpha_{j-1}[i]$
- $\delta_j^{(i)} = \sqrt{\beta_j[i]}/\alpha_j[i]$
populate the diagonals and off-diagonals, respectively, of the partial Lanczos tridiagonal $\tilde T_i \in \mathbb{R}^{p \times p}$. After $p$ iterations, mBCG yields $X_p \approx K^{-1}Z$ and partial tridiagonals for SLQ. This routine enables estimating:

- $K^{-1}y$ via $Z = [y]$
- $\operatorname{trace}(K^{-1} \partial K/\partial\theta)$ by Hutchinson’s method with random $Z$
- $\log |K|$ by stochastic Lanczos quadrature on the tridiagonals

The batched nature allows efficient use of GPU Batched-BLAS through large fused matrix-matrix multiplies [1809.11165].

## 2. Algorithmic Structure and Pseudocode

The essential mBCG algorithm proceeds as follows (unpreconditioned):

```pseudo
Input: K (mat-mat operator), Z ∈ ℝⁿ×ᵗ, p (CG iterations)
Output: X ≈ K⁻¹Z, tridiagonals {T_i}₁^t
X ← 0ₙ×ᵗ
R ← Z
P ← R
for j = 0,…,p−1:
    V ← K(P)
    α ← diag(Rᵀ R) ⧸ diag(Pᵀ V)
    X ← X + P Diag(α)
    R_new ← R – V Diag(α)
    β ← diag(R_newᵀ R_new) ⧸ diag(Rᵀ R)
    for i=1…t:
        γ_j^{(i)} ← 1/α[i] + (j>0 ? β_prev[i]/α_prev[i] : 0)
        δ_j^{(i)} ← sqrt(β[i])/α[i]
        Update T_i with γ, δ
    P ← R_new + P Diag(β)
    R ← R_new
    α_prev ← α, β_prev ← β
return X, {T_i}
```

These outputs are then directly used for the primary quantities of GP inference:
- Posterior mean: $K^{-1} y$
- Derivatives: $(1/t) \sum_i X[:,i]^\top (\partial K/\partial\theta) Z[:,i]$
- Log-determinant: $(1/t)\sum_i \mathbf{e}_1^\top \log(\tilde T_i) \mathbf{e}_1$

Preconditioning is introduced by transforming the residuals and inner products as discussed below [1809.11165].

## 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 $n \times n$ covariance matrix $K = K_\text{train} + \sigma^2 I$, the rank-$k$ pivoted Cholesky factorization $K_k = L_k L_k^\top$ yields the preconditioning matrix $M = K_k + \sigma^2 I$. The preconditioned system is

$$
M^{-1/2} K M^{-1/2} u = M^{-1/2} y, \qquad K^{-1} y = M^{-1/2} u
$$

However, only $M^{-1}$ is required within the mBCG recurrences: initial search directions as $P_0 = M^{-1} R_0$, and inner products as $R^\top M^{-1} R$. Log-determinant computation proceeds via
$$
\log |K| = \log |M| + \log |M^{-1/2} K M^{-1/2}|
$$
where the correction term is estimated via SLQ on the preconditioned tridiagonals.

Key computational properties:
- $L_k$ can be computed in $O(nk^2)$
- $M v = b$ solves in $O(nk^2)$ (Woodbury identity)
- Sampling $z \sim \mathcal{N}(0, M)$ is $O(nk)$
- $\log|M| = \sum_{i=1}^k \log(1+\lambda_i/\sigma^2) + n\log \sigma^2$ in $O(k)$

For scalar RBF kernels, the preconditioned condition number satisfies $\kappa(M^{-1} K) \le (1 + n e^{-b k})^2$, decaying exponentially in $k$. The CG error accordingly falls as $(1 + c e^{-b k}/n)^p$ (Theorem 4.2 in [1809.11165]).

## 4. Computational Complexity and Parallelism

Each mBCG iteration requires one dense matrix-matrix multiply $K \cdot P$ at $O(n^2 t)$ cost for dense kernels, plus $O(nt)$ operations for vector reductions and bookkeeping. For $p$ iterations with $t$ probe vectors, total work is $O(p n^2 t)$. Cholesky factorization and classical CG require $O(n^3)$ 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 $O(n^3)$ to $O(n^2)$. Further complexity reductions are achieved for structured kernel approximations (e.g., SKI, SGPR), where $T_\text{mult}$ decreases to $O(n m^2)$ or $O(n m \log m)$. mBCG therefore unifies these techniques into a single, scalable, matrix-matrix Krylov routine [1809.11165].

## 5. Practical GPU Implementation Strategies

The practical realization of mBCG for BBMM inference involves:

- Storing all block variables ($P, R, X$) as $n \times t$ device-resident tensors
- Performing kernel matrix-matrix products $K \cdot P$ using user-supplied GPU routines (either recomputed on the fly or via fast operator representations)
- Allocating all workspace buffers ($V, R, P, X$) a priori to prevent allocation overhead during iteration
- Compactly storing tridiagonal data as $p \times p \times t$, where $p \approx 20$ and $t \approx 10$ require negligible memory
- Achieving $>10$ 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 $O(1)$ sequence of small vector operations, maximizing hardware efficiency and enabling large-scale GP inference [1809.11165].

## 6. Empirical Performance and Analytical Results

Empirical results on a Titan Xp GPU demonstrate that:

- For exact GPs ($n \lesssim 3500$), BBMM with mBCG is up to $20\times$ faster than CPU-based Cholesky, and $4\times$ faster than GPU-based Cholesky
- For SGPR with $n \le 50,\!000$, BBMM achieves $10\times$ speedup over CPU Cholesky
- For SKI+DKL with $n \le 500,\!000$, BBMM is up to $15\times$ faster than state-of-the-art CPU-based MVM+Lanczos inference

Convergence curves indicate that unpreconditioned mBCG can require $p \approx 50$–100 iterations for $10^{-6}$ relative residuals, while a rank-5 Cholesky preconditioner reduces this to $10$–20, and increasing $k$ to $9$ halves the required iterations further. Test error evaluations match or slightly improve upon Cholesky-based inference, with early-terminated CG occasionally regularizing tiny eigenvalues [1809.11165].

## 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. [1204.0069], introduces block updates in $X_k \in \mathbb{R}^{n \times p}$ via $X_{k+1} = X_k + D_k \Alpha_k$ and enforces block (bi)-orthogonality and $A$-conjugacy. The per-iteration cost per thread is $O(n^2)$ with $O(n^{2+1/3})$ worst-case complexity given $p^* \sim n^{2/3}$. 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 [1204.0069, 1809.11165].

Source: https://www.emergentmind.com/topics/batched-conjugate-gradients-mbcg