[論文レビュー] Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro

Du Phan, Neeraj Pradhan+1 more|arXiv (Cornell University)|Dec 24, 2019

Parallel Computing and Optimization Techniques参考文献 15被引用数 226

ひとこと要約

本論文は NumPyro を紹介する。NumPyro は NumPy ベースの確率的プログラミングライブラリで、JAX 変換を用いた合成可能なエフェクトハンドラを利用し、エンドツーエンドの JIT コンパイルと大幅なスピードアップを実現する。これには、既存の実装を上回る反復的な NUTS サンプラーが含まれる。

ABSTRACT

NumPyro is a lightweight library that provides an alternate NumPy backend to the Pyro probabilistic programming language with the same modeling interface, language primitives and effect handling abstractions. Effect handlers allow Pyro's modeling API to be extended to NumPyro despite its being built atop a fundamentally different JAX-based functional backend. In this work, we demonstrate the power of composing Pyro's effect handlers with the program transformations that enable hardware acceleration, automatic differentiation, and vectorization in JAX. In particular, NumPyro provides an iterative formulation of the No-U-Turn Sampler (NUTS) that can be end-to-end JIT compiled, yielding an implementation that is much faster than existing alternatives in both the small and large dataset regimes.

研究の動機と目的

• Pyro に似たエフェクトハンドラを JAX ベースのバックエンドと統合して、NumPyro のモデリングと推論インターフェイスを拡張できることを示す。
• 可 composable transformations (jit, grad, vmap) が、Pyro 互換のモデリング言語を維持しつつ、推論サブルーチンを加速することを示す。
• CPU および GPU での高速化を XLA によって実現する、エンドツーエンドの JIT で反復的な NUTS 実装を提示する。
• vmap を用いたサブルーチンのベクトル化が、バッチ推論と予測を可能にする利点を示す。
• 小規模および大規模データセットにわたる性能向上を定量化するため、NumPyro を Stan および Pyro とベンチマークする。

提案手法

• NumPyro 内で Pyro に似たエフェクト処理抽象化（seed、trace、condition）を採用し、JAX バックエンド上で合成可能なプログラム変換を可能にする。
• JAX transformations（jit、grad、vmap）を活用して、関数型でトレース可能なコードを用いた推論サブルーチンを構築・最適化する。
• NUTS BuildTree サブルーチンを反復形に変換して、エンドツーエンドの JIT コンパイルと XLA 経由の演算子融合を改善する。
• prior sampling、posterior predictive sampling、log-likelihood 計算など、共通の推論タスクをバッチ処理するために vmap を使用する。
• エフェクトハンドラと JAX transform の統合を実演し、加速を得つつ Pyro 互換のモデリングインターフェイスを維持する。

実験結果

リサーチクエスチョン

• RQ1 Pyro スタイルのエフェクトハンドラを JAX 変換と効果的に組み合わせて、NumPyro で確率的プログラミングのワークロードを加速できるか。
• RQ2 NumPyro における NUTS などの推論サブルーチンのエンドツーエンド JIT コンパイルから得られる性能向上は、Pyro および Stan と比較してどうか。
• RQ3 vmap によるベクトル化が、エフェクトハンドラとどのように相互作用して、スケーラブルなバッチ推論と予測を可能にするか。
• RQ4 反復的で JIT でコンパイルされる NUTS 実装は、データセットのサイズが異なる場合にも大幅なスピードアップを提供しつつ正確性を維持できるか。
• RQ5 大規模モデルで NumPyro の合成可能な変換を使用した場合、CPU、GPU、TPU バックエンドでの実践的な利得は何か。

主な発見

• NumPyro の反復的 NUTS 実装は、テストされたモデルで Pyro および Stan よりも著しく高速であり、特定のタスクでは Pyro に約340倍、Stan に約6倍の速度向上が報告されている。
• End-to-end JIT コンパイル via XLA yields substantial speedups by enabling full fusion and optimized execution of inference subroutines.
• Vectorizing inference with vmap enables efficient batched sampling from the prior and posterior predictive distributions as well as batched log-likelihood computations.
• NumPyro’s framework maintains a Pyro-compatible modeling language while leveraging JAX transformations for acceleration across CPU, GPU, and TPU backends.
• The paper shows end-to-end acceleration benefits are observed in both small-scale CPU regimes and large-scale GPU regimes, motivating broader applicability of composable effects with JAX.