[論文レビュー] BlackJAX: Composable Bayesian inference in JAX

Alberto Cabezas, Adrien Corenflos+20 more|arXiv (Cornell University)|Feb 16, 2024

Bayesian Methods and Mixture Models被引用数 6

ひとこと要約

BlackJAX は、MCMC、SMC、SG-MCMC、および VI のための JAX における機能的で組み合わせ可能なベイズ推論ビルディングブロックのセットを提供し、低レベルのコンポーネントから新しいアルゴリズムを組み立てることを可能にします。CPU、GPU、TPU を対象とし、正規化されていない log-density targets と統合します。

ABSTRACT

BlackJAX is a library implementing sampling and variational inference algorithms commonly used in Bayesian computation. It is designed for ease of use, speed, and modularity by taking a functional approach to the algorithms' implementation. BlackJAX is written in Python, using JAX to compile and run NumpPy-like samplers and variational methods on CPUs, GPUs, and TPUs. The library integrates well with probabilistic programming languages by working directly with the (un-normalized) target log density function. BlackJAX is intended as a collection of low-level, composable implementations of basic statistical 'atoms' that can be combined to perform well-defined Bayesian inference, but also provides high-level routines for ease of use. It is designed for users who need cutting-edge methods, researchers who want to create complex sampling methods, and people who want to learn how these work.

研究の動機と目的

• ブラックボックス・サンプラーと並んで、構造を意識したベイズ推論の必要性を動機づける。
• 基本的なコンポーネントから推論アルゴリズムを構築するための、組み合わせ可能で機能的な API を導入する。
• 標準的なパイプラインを超えて、推論手法をカスタマイズ・拡張するための低レベル API を提供する。
• 正規化されていない log-density ターゲットと最新のハードウェアとの統合を実証する。

提案手法

• 各反復が副作用なしに新しい状態を更新して返す、機能的で状態駆動の API を提案する。
• 基本的な推論コンポーネントのライブラリを提供する（例：MH acceptance/reject ステップ、HMC/Langevin ダイナミクス、SMC、VI ツール、温度調整、リサンプリング）。
• これらのコンポーネントの組み合わせを有効化して、新規または既存のアルゴリズムを形成し、HMC のウィンドウや適応スキームに例示される。
• 特定用途に合わせたカーネル構築と複数粒子やカーネルによる並列実験のための低レベル API を提供する。
• すべてのコンポーネントはターゲット log-density 関数で動作し、JAX を介して CPU/GPU/TPU 実行をサポートします。

実験結果

リサーチクエスチョン

• RQ1 ベイズ推論手法を再利用可能で組み合わせ可能なコンポーネントにどのように分解できるか？
• RQ2 副作用のない機能的設計は、MCMC/SMC/VI アルゴリズムの組み立てにおける並列化と柔軟性を向上させるか？
• RQ3 BlackJAX は正規化されていない log-density ターゲットと確率的プログラミングのワークフローとどの程度効果的に統合できるか？
• RQ4 高レベル API と低レベル API の両方を公開することが、アルゴリズムのカスタマイズと実験に与える影響は何か？

主な発見

• 状態を持つ組み合わせ可能な設計で、サンプリングと近似推論の統一的で純粋に機能的なインターフェースを提供します。
• コアコンポーネントの複数のバリアントを含む（例：MH acceptance/reject、さまざまな積分器を用いたHMC、NUTS、一般化HMC、非可逆スライスサンプリング）。
• 正規化されていない log-density ターゲット上で動作することにより、確率的プログラミング言語との相互運用性を示す。
• SMC や並列設定でカーネルをカスタマイズし、ベースカーネルを混合する低レベル API を提供します。
• JAX ベースのエコシステムとの広い互換性と寛容な Apache License v2.0 を備え、オープンソースのガバナンスと包括的なテストスイート（99% カバレッジ）を持つ。
• 論文、チュートリアル、オープンソースの講座への採用を通じて教育・研究への影響を記録する。