[論文レビュー] 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次元）における効率的なギブスサンプリングを支援すること。
• 学習および最適化パイプラインにおける自動微分のバックワードパスを構築すること。
• 明示的なコレスキー分解を回避し、GPUアクセラレーションを活用したMVMを活用することで、メモリおよび計算コストを低減すること。

提案手法

• Haleら（2020）に基づき、行列平方根の有理近似を、シフト付き行列逆行列の和として用いる。
• 複数のシフト系 (t_q I + K)^{-1}b を1回の反復で解くために、修正されたマルチシフトMINRES（msMINRES）アルゴリズムを採用し、MVMを共有する。
• 異なるシフトパラメータを持つ複数のシフトに対して、1つのプリコンディショナーを適用し、msMINRESの収束を加速する。
• パーティショニングされたMVMを用いることで、完全な行列格納を回避し、O(N)のメモリ使用量を維持する。
• 随伴感度法を用いて、∂(K^{±1/2}b)/∂K および ∂(K^{±1/2}b)/∂b のスケーラブルなバックワードパスを導出する。
• 四則演算に基づく有理近似と反復的Krylov部分空間法を組み合わせ、精度と効率のバランスを図る。

実験結果

リサーチクエスチョン

• RQ1 明示的なコレスキー分解を用いずに、高次元設定でも行列平方根を効率的に計算できるか？
• RQ2 複数のシフトに対して、1つのプリコンディショナーでマルチシフトKrylovソルバを加速できるか？
• RQ3 100回未満の行列-ベクトル乗算で、4桁以上（例：4+小数桁）の精度を達成できるか？
• RQ4 50,000×50,000の行列サイズに対しても、近似誤差が小さく保たれるか？
• RQ5 変分ガウス過程およびベイズ最適化のエンドツーエンド学習パイプライン（バックプロパゲーション含む）に統合可能か？

主な発見

• 100回未満の行列-ベクトル乗算で、行列平方根計算において4〜5桁の精度を達成。
• 50,000×50,000の行列サイズに対しても、近似誤差が無視できるほど小さく、従来のコレスキーに基づく手法をはるかに上回る性能を発揮。
• K^{±1/2}b のバックワードパスは追加計算が最小限であり、微分可能な機械学習パイプラインへの効率的統合を可能にする。
• O(M²)のMVMに基づく自然勾配更新を用いることで、最大10,000個の誘導点を有する変分ガウス過程推論が実現可能。
• 25,600次元の画像再構成問題において、提案された K^{-1/2}b ルーチンを用いてギブスサンプリングを成功裏に実行。
• 理論的解析により、収束速度がKの条件数に依存し、反復回数が増加するに従い誤差境界が指数関数的に減少することが示された。