[论文解读] Fast Matrix Square Roots with Applications to Gaussian Processes and Bayesian Optimization

Geoff Pleiss, Martin Jankowiak+3 more|arXiv (Cornell University)|Jun 19, 2020

Gaussian Processes and Bayesian Inference参考文献 79被引用 11

一句话总结

本文提出了一种快速、无需矩阵存储的算法，通过矩阵-向量乘法（MVM）计算矩阵平方根及其逆矩阵，从而在高维高斯过程和贝叶斯优化中实现高效的采样与白化。通过结合有理逼近与预处理多移位MINRES求解器，该方法在100次以内的MVM内实现了4–5位小数的精度——在GPU加速下可扩展至50,000×50,000矩阵，支持可扩展的变分推断与吉布斯采样。

ABSTRACT

Matrix square roots and their inverses arise frequently in machine learning, e.g., when sampling from high-dimensional Gaussians $\mathcal{N}(\mathbf 0, \mathbf K)$ or whitening a vector $\mathbf b$ against covariance matrix $\mathbf K$. While existing methods typically require $O(N^3)$ computation, we introduce a highly-efficient quadratic-time algorithm for computing $\mathbf K^{1/2} \mathbf b$, $\mathbf K^{-1/2} \mathbf b$, and their derivatives through matrix-vector multiplication (MVMs). Our method combines Krylov subspace methods with a rational approximation and typically achieves $4$ decimal places of accuracy with fewer than $100$ MVMs. Moreover, the backward pass requires little additional computation. We demonstrate our method's applicability on matrices as large as $50,\!000 imes 50,\!000$ - well beyond traditional methods - with little approximation error. Applying this increased scalability to variational Gaussian processes, Bayesian optimization, and Gibbs sampling results in more powerful models with higher accuracy.

研究动机与目标

• 解决高维高斯过程与贝叶斯优化中计算K^{±1/2}b的计算瓶颈。
• 实现大规模高斯过程模型（最多10,000个诱导点）的可扩展推断。
• 在高维问题（如25,600维）中，利用矩阵平方根实现高效的吉布斯采样。
• 为自动微分在学习与优化流程中的反向传播开发可扩展的实现。
• 通过避免显式Cholesky分解并利用GPU加速下的MVM，降低内存与计算成本。

提出的方法

• 基于Hale等人（2020）的方法，使用矩阵平方根的有理逼近表示为一系列移位矩阵逆的和。
• 采用改进的多移位MINRES（msMINRES）算法，在单次迭代中求解多个移位系统(t_q I + K)^{-1}b，共享MVM计算。
• 应用单一预处理矩阵，尽管移位参数不同，仍能加速msMINRES在多个移位上的收敛。
• 使用分块MVM实现，保持O(N)内存使用，避免完整矩阵存储。
• 通过伴随敏感性方法推导出∂(K^{±1/2}b)/∂K与∂(K^{±1/2}b)/∂b的可扩展反向传播。
• 结合基于求积的有理逼近与迭代Krylov子空间方法，在精度与效率之间取得平衡。

实验结果

研究问题

• RQ1 能否在无需显式Cholesky分解的情况下，高效计算高维设置下的矩阵平方根？
• RQ2 能否使用单一预处理矩阵加速多个移位上的多移位Krylov求解器？
• RQ3 能否在少于100次矩阵-向量乘法内实现高精度（如4位以上小数）？
• RQ4 能否在保持近似误差极低的前提下，将该方法扩展至50,000×50,000的矩阵规模？
• RQ5 能否将该方法集成到端到端学习流程中，包括反向传播，用于变分高斯过程与贝叶斯优化？

主要发现

• 该方法在少于100次矩阵-向量乘法内实现了矩阵平方根计算的4–5位小数精度。
• 该算法在50,000×50,000矩阵规模下保持极低的近似误差，远超传统基于Cholesky的方法。
• K^{±1/2}b的反向传播计算开销极低，可高效集成于可微机器学习流程中。
• 通过基于O(M²) MVM的自然梯度更新，实现了最多10,000个诱导点的变分高斯过程推断。
• 在25,600维图像重建问题中，成功使用所提出的K^{-1/2}b模块执行了吉布斯采样。
• 理论分析表明，收敛速率依赖于K的条件数，误差界随迭代次数呈指数衰减。