[論文レビュー] Simple, Distributed, and Accelerated Probabilistic Programming

Dustin Tran, Matthew W. Hoffman+4 more|arXiv (Cornell University)|Jan 1, 2018

Machine Learning and Data Classification被引用数 30

ひとこと要約

この論文は、TensorFlowに埋め込まれた軽量で低レベルの確率的プログラミングフレームワークを紹介しており、1つの抽象化概念—確率的変数—に焦点を当て、効率的なモデル並列およびデータ並列推論を可能にしている。VAEおよび画像変換器モデルにおいて、256個のTPUv2チップで最適な線形スケーリングを達成し、マルチGPU NUTSサンプリングにおいてStanより100倍の高速化を実現した。

ABSTRACT

We describe a simple, low-level approach for embedding probabilistic programming in a deep learning ecosystem. In particular, we distill probabilistic programming down to a single abstraction—the random variable. Our lightweight implementation in TensorFlow enables numerous applications: a model-parallel variational auto-encoder (VAE) with 2nd-generation tensor processing units (TPUv2s); a data-parallel autoregressive model (Image Transformer) with TPUv2s; and multi-GPU No-U-Turn Sampler (NUTS). For both a state-of-the-art VAE on 64x64 ImageNet and Image Transformer on 256x256 CelebA-HQ, our approach achieves an optimal linear speedup from 1 to 256 TPUv2 chips. With NUTS, we see a 100x speedup on GPUs over Stan and 37x over PyMC3.

研究の動機と目的

• 確率的プログラミングを単一の抽象化概念、すなわち確率的変数に簡略化することで、簡素化すること。
• 最小限の工学的オーバーヘッドで、ディープラーニングエコシステムにおける効率的かつスケーラブルな推論を可能にすること。
• モデル並列およびデータ並列の確率的モデルにおいて、数千のアクセceleratorにわたる線形スケーリングを達成すること。
• マルチGPU No-U-Turn Sampler (NUTS) を用いたマルコフ連鎖モンテカルロ（MCMC）サンプリングを、既存のツールを凌駕する速度で高速化すること。

提案手法

• フレームワークは、複雑な分布や推論メカニズムを抽象化するため、確率的プログラムを唯一の核心抽象化概念、すなわち確率的変数でモデル化する。
• TensorFlowにネイティブに統合されており、自動微分および分散実行機能を活用して効率的な計算を実現する。
• モデル並列のVAEおよびデータ並列の自己回帰的モデル（例：Image Transformer）は、TPUv2ハードウェア上でこの抽象化を用いて実装されている。
• マルチGPU NUTS推論は、デバイス間でのパラメータシャーディングと勾配同期を通じて実現されている。
• 最小限のコード変更で変分推論とMCMCサンプリングの両方をサポートする。
• 1〜256個のTPUv2チップにわたる効率的なデータ並列およびモデル並列により、線形スケーリングが達成されている。

実験結果

リサーチクエスチョン

• RQ1 表現力や性能を損なわせることなく、確率的プログラミングを単一の抽象化概念に簡略化できるか？
• RQ2 最小限で低レベルのインターフェースは、モデル並列およびデータ並列の両設定において、数千のアクセceleratorにわたる最適なスケーリングを可能にするか？
• RQ3 このようなシステムは、既存の確率的プログラミングツールと比較して、MCMCサンプリングにおいて優れたパフォーマンスを達成できるか？
• RQ4 このフレームワークは、分散実行下でも、VAEやTransformerのような最先端のモデルをどの程度サポートできるか？

主な発見

• 64×64 ImageNetにおける最先端のVAEに対して、1〜256個のTPUv2チップで最適な線形スケーリングを達成した。
• 256×256 CelebA-HQにおける画像変換器モデルでも、データ並列設定において最適な線形スケーリングを示した。
• マルチGPU NUTSサンプリングは、Stanに対して100倍、PyMC3に対して37倍の高速化を達成した。
• 同一の低レベル抽象化を用いて、変分自己符号化器および自己回帰的モデルの両方で効率的かつスケーラブルな推論を可能にした。
• さまざまなモデルアーキテクチャおよびハードウェア構成（TPUやマルチGPU環境を含む）において、高いパフォーマンスを維持した。
• 確率的変数抽象化の単純さのおかげで、パフォーマンスやスケーラビリティを損なわずに広範な適用が可能になった。