[論文レビュー] Pigeons.jl: Distributed Sampling From Intractable Distributions

Nikola Surjanovic, Miguel Biron-Lattes+4 more|arXiv (Cornell University)|Aug 18, 2023

Gaussian Processes and Bayesian Inference被引用数 6

ひとこと要約

Pigeons.jl は、複雑で扱いにくい確率分布（例えば、複数モードを持つ事後分布や高次元分布）からの分散型で再現可能なサンプリングを可能にする Julia パッケージであり、強力な並列処理不変性（strong parallelism invariance）という新しい手法を採用している。同じシードを用いたシングルスレッド、マルチスレッド、分散（MPI）実行において、出力が同一であることを保証しており、科学的再現性を確保するとともに、ランダム化された並列 MCMC アルゴリズムのデバッグを簡素化する。

ABSTRACT

We introduce a software package, Pigeons.jl, that provides a way to leverage distributed computation to obtain samples from complicated probability distributions, such as multimodal posteriors arising in Bayesian inference and high-dimensional distributions in statistical mechanics. Pigeons.jl provides simple APIs to perform such computations single-threaded, multi-threaded, and/or distributed over thousands of MPI-communicating machines. In addition, Pigeons.jl guarantees a property that we call strong parallelism invariance: the output for a given seed is identical irrespective of the number of threads and processes, which is crucial for scientific reproducibility and software validation. We describe the key features of Pigeons.jl and the approach taken to implement a distributed and randomized algorithm that satisfies strong parallelism invariance.

研究の動機と目的

• 分散および並列 MCMC サンプリングにおける再現性の課題に取り組むこと。通常、出力はスレッド数やプロセス数に依存して変動するためである。
• ベイズ推論や統計物理学で一般的な複雑で高次元、あるいはマルチモーダルな分布からのサンプリングに向け、スケーラブルで使いやすいインターフェースを提供すること。
• 既存のツールではしばしばアクセス不可能な、期待値 ∫f(x)π(x)dx と正規化定数 Z の両方の正確な推定を可能にすること。
• 連続的または離散的なパrameter空間にとどまらず、系統発生樹などの組み合わせオブジェクトを含む任意のデータ型をサポートすること。
• 強力な並列処理不変性（SPI）により正しく動作し、実行環境の違いに関わらずデバッグやテストを決定的に行えるように、正しさと検証のしやすさを確保すること。

提案手法

• パッケージは、複雑なターゲット分布の効率的探索を可能にする、分散型非可逆的平行温度法（parallel tempering）をコアなサンプリングアルゴリズムとして実装している。
• 単一のシードから独立で再現可能な乱数ストリームを生成するため、SplittableRandoms.jl パッケージを用いて、スレッド数やプロセス数に関わらず決定論的になるようにしている。
• Julia のマルチスレッドおよび MPI インターフェースを活用し、単一マシン、マルチコアシステム、数千ノードのクラスタまでをカバーする実行を可能にしている。
• ユーザーは非正規化密度 γ(x) を指定するだけで、Pigeons.jl は期待値の計算や正規化定数 Z の推定に必要なサンプルを返す。
• 強力な並列処理不変性は、すべてのスレッドおよびプロセスでマスターランダム生成器を決定的に分割することで実現されており、乱数生成における非決定論的振る舞いを排除している。
• 低レベルの分散コンピューティングの複雑さを抽象化したシンプルで高水準な API を提供するとともに、実行構成に対する完全な制御を維持している。

実験結果

リサーチクエスチョン

• RQ1 同じ乱数シードを用いた場合、分散型 MCMC サンプラーがスレッド数やプロセス数が異なる状況でも同一の出力を保証できるか？
• RQ2 スケーラビリティを損なわずに、高性能で分散型かつランダムなアルゴリズムにおいて、強力な並列処理不変性をどのように達成できるか？
• RQ3 並列および分散型の乱数生成における非決定論的要因は、科学的ソフトウェアの再現性を損なう主な要因として、どのようなものか？
• RQ4 汎用的なサンプリングパッケージとして、連続的および組み合わせ的パrameter空間の両方を効率的に処理できるか、正規化定数の推定も可能か？
• RQ5 スプライタブル乱数生成の統合は、分散型 MCMC アルゴリズムの正しさとテスト可能性をどのように向上させるか？

主な発見

• Pigeons.jl は強力な並列処理不変性を保証している：同じシードを用いた場合、スレッド数やプロセス数に関わらず出力が同一であるため、デバッグや検証を決定的に行える。
• パッケージは、数千台の MPI 通信を行うマシンを用いた分散実行を成功裏にサポートしており、シングルスレッド、マルチスレッド、分散構成のすべてで同一の結果を得ている。
• SplittableRandoms.jl を使用することで、すべての実行ユニットで独立で再現可能な乱数ストリームを実現しており、乱数生成における非決定論的振る舞いを排除している。
• フレームワークにより、事後分布の期待値と正規化定数 Z の両方の正確な推定が可能であり、既存のツールではしばしば省略または近似されている。
• 実装は、高次元やマルチモーダルな事後分布に対しても、高性能でスケーラブルかつ再現可能なサンプリングが実際可能であることを示している。
• パッケージは Julia 1.8 以降と互換性があり、複数の確率的プログラミング言語とも統合可能であり、科学的分野における広範な応用性を有している。