[論文レビュー] tfp.mcmc: Modern Markov Chain Monte Carlo Tools Built for Modern Hardware

Junpeng Lao, Christopher Suter+8 more|arXiv (Cornell University)|Feb 4, 2020

Markov Chains and Monte Carlo Methods参考文献 15被引用数 17

ひとこと要約

この論文では、TensorFlow Probability用に構築された現代的なマルコフ連鎖モンテカルロ（MCMC）ツールキットである tfp.mcmc を紹介する。このツールキットは、広範なデータ並列処理、ベクトル化計算、および合成可能でモジュラーなカーネルを活用することで、ハードウェア最適化された高効率なMCMC推論を実現する。SIMDおよびCPU、GPU、TPUにおけるマルチコア実行を活用して、巨大な並列性を達成しつつ、単純なターゲット対数尤度（TLP）インタフェースにより、あらゆる確率的モデルと互換性を保つ。

ABSTRACT

Markov chain Monte Carlo (MCMC) is widely regarded as one of the most important algorithms of the 20th century. Its guarantees of asymptotic convergence, stability, and estimator-variance bounds using only unnormalized probability functions make it indispensable to probabilistic programming. In this paper, we introduce the TensorFlow Probability MCMC toolkit, and discuss some of the considerations that motivated its design.

研究の動機と目的

• 現代のハードウェア機能（SIMD、マルチコアCPU、GPU、TPU）を最大限に活用し、スケーラブルな確率的推論を実現するMCMCフレームワークの設計。
• 複数のチェーンにおけるバッチ処理とベクトル化演算を活用することで、効率的かつ自明に並列化可能なMCMCの実現。
• ドメイン固有言語を必要とせず、特定のモデリングフレームワークに強く結合しない、モジュラーで合成可能なAPIを提供して、複雑なMCMC遷移カーネルの構築を可能にする。
• 正規化されていない対数尤度関数と自動微分を両方サポートすることで、多様な確率的モデルとの柔軟な統合を可能にする。
• 合成可能カーネルコンponentを用いて、適応的ステップサイズ、前処理、ハイパーパramータチューニングといった高度なMCMC技術を、一般用途のフレームワーク内でモジュラーかつ合成可能に実装する。

提案手法

• フレームワークは、TensorFlowのネイティブなバッチ処理と自動微分を活用し、複数のMCMCチェーンを並列にベクトル化計算可能にしている。
• ターゲット対数尤度（TLP）をPythonのコールバック関数として指定する関数型インタフェースを導入し、MCMCエンジンとモデル仕様を分離している。
• コアな抽象化として、MCMC遷移論理をカプセル化する TransitionKernel が存在し、Metropolis-Hastingsに未補正提案を組み込むような複雑なアルゴリズムのネスティングをサポートしている。
• sample_chainドライバは、TensorFlowの tf.while_loop とXLAコンパイルを用いて、バーンイン、サンプリング、トレースを調整し、ハードウェア加速と低レベル最適化を可能にしている。
• MetropolisHastings、SimpleStepSizeAdaptation、TransformedTransitionKernel などの合成可能カーネルにより、高度なMCMCワークフローをモジュラーに構築できる。
• Bijectorsによる再パrameterizationにより、状態空間を変換し、ハミルトニアンモンテカルロなどのアルゴリズムにおけるサンプリング効率を向上させている。

実験結果

リサーチクエスチョン

• RQ1 どのようにして、巨大な並列性を持つ現代のハードウェア上でMCMC推論を効率的にスケーリングできるか？
• RQ2 多様なハードウェアとモデルで動作する、高性能で合成可能なMCMCカーネルを実現するためのアーキテクチャパターンは何か？
• RQ3 モジュラリティや正しさを損なわずに、MCMCにベクトル化計算とバッチ処理をシームレスに統合する方法は何か？
• RQ4 SIMDおよびマルチコア実行を用いて数千の並列MCMCチェーンを実行した場合の、性能と収束への影響は何か？
• RQ5 一般用途のフレームワーク内で、適応的MCMC技術をモジュラーかつ合成可能に実装する方法は何か？

主な発見

• tfp.mcmcライブラリは、AVX512ベクトル化を活用した単一の32コアCPU上で最大1024本の並列MCMCチェーンを実行可能であり、従来のタスク並列アプローチに比べて桁違いの高速化を達成している。
• TensorFlowのバッチ処理のセマンティクスに準拠することで、チェーン間のベクトル化計算がネイティブにサポートされており、1つの tf.Tensor 入力で100本のチェーンを同時に実行可能である。
• 合成可能カーネル設計により、適応的ステップサイズ付きHMCやBijectorsによる前処理付きサンプリングといった複雑なMCMCアルゴリズムをモジュラーに構築できる。
• 正規化されていない対数尤度関数と自動微分の両方をサポートしており、TensorFlowで実装されたあらゆるモデルとシームレスに統合可能である。
• XLAコンパイルとTensorFlowの低レベルAPIの使用により、GPUおよびTPUでのハードウェア加速実行が可能となり、性能が顕著に向上している。
• ストリーミング期待値推定やマルチカーネルドライバといった、新たなワークフローを可能にし、高度な診断と適応的MCMC戦略を支援している。