[论文解读] 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 Inference参考文献 33被引用 536

一句话总结

介绍 BBMM 推理，一种面向高斯过程的黑箱矩阵-矩阵方法，使用批量修正共轭梯度和 GPU 加速，在基于 PyTorch 的框架（GPyTorch）下实现更快的精确与近似 GP 推理。

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.

研究动机与目标

• 推动对硬件高效的 GP 推理工具的需求，这些工具将模型设定与推理过程解耦。
• 开发一个黑箱框架 BBMM，它依赖核矩阵-矩阵乘法，而非 Cholesky 分解。
• 提供一个可扩展、面向 GPU 的 GP 推理引擎，支持精确 GP 与流行的近似方法（SGPR、SKI）。
• 提供一个基于 PyTorch 构建的软件平台 GPyTorch，以简化复杂 GP 模型的原型设计。

提出的方法

• 使用修改后的批量共轭梯度（mBCG）在一次调用中计算所有推断项（K^{-1}y、log|K|，以及 Tr(K^{-1} dK/dθ)）。
• 通过使用探针向量 z_i 的随机迹估计来估计对数行列式和迹。
• 采用置换的 Cholesky 预条件 P = L_k L_k^T + σ^2 I，加速 CG 收敛并实现高效的对数行列式修正。
• 证明 BBMM 将精确 GP 推理复杂度从 O(n^3) 降至 O(n^2)，并可与 SGPR 和 SKI 框架集成。
• 展示 BBMM 如何利用 GPU 硬件，在多个 GP 模型和数据集上实现相对于 Cholesky 基方法的显著加速。

实验结果

研究问题

• RQ1 BBMM 推理能否在精确和近似 GP 模型中达到甚至超过基于 Cholesky 的 GP 推理的准确性？
• RQ2 具有置换 Cholesky 预条件的修改后的批量 CG（mBCG）在 GPU 硬件上的收敛速度和可扩展性表现如何？
• RQ3 BBMM 在多大程度上可以作为一个黑箱框架应用于各种 GP 模型和可扩展近似（SGPR、SKI、Toeplitz/KISS-GP），且推导工作最小化？

主要发现

• 与 CPU 的 Cholesky 基方法相比，BBMM 在 GPU 上对精确 GP 推理和可扩展近似的加速显著。
• 使用 BBMM 的精确 GP 在数据集大约 3000 点时可比 Cholesky 基方法快多达 20 倍（受限于 GPU 内存）。
• 使用 BBMM 的 SGPR 和 SKI 分别在数据集达到 50 万点时实现最多 15 倍和 4 倍的加速。
• 采用置换 Cholesky 的预条件显著加速 BBMM 框架中 CG 的收敛。
• BBMM 方法通过依赖高效的核矩阵乘法，使得许多 GP 模型的实现代码行数更少（通常少于 50 行）。