[論文レビュー] JAXNS: a high-performance nested sampling package based on JAX

Joshua G. Albert|arXiv (Cornell University)|Dec 30, 2020

Gaussian Processes and Bayesian Inference参考文献 11被引用数 4

ひとこと要約

JAXNS は JAX と XLA のプリミティブに基づいた高性能なネストド・サブセットリングパッケージであり、PolyChord や MultiNEST や dynesty よりも 2–3 時間の速度向上を達成しながらも、正確性を保持している。静的メモリ計算と JIT コンパイルを活用して、多様なベイズ推論問題における尤度評価を高速化し、機械学習およびデータサイエンス分野での実用的利用を可能にしている。

ABSTRACT

Since its debut by John Skilling in 2004, nested sampling has proven a valuable tool to the scientist, providing hypothesis evidence calculations and parameter inference for complicated posterior distributions, particularly in the field of astronomy. Due to its computational complexity and long-running nature, in the past, nested sampling has been reserved for offline-type Bayesian inference, leaving tools such as variational inference and MCMC for online-type, time-constrained, Bayesian computations. These tools do not easily handle complicated multi-modal posteriors, discrete random variables, and posteriors lacking gradients, nor do they enable practical calculations of the Bayesian evidence. An opening thus remains for a high-performance out-of-the-box nested sampling package that can close the gap in computational time, and let nested sampling become common place in the data science toolbox. We present JAX-based nested sampling (JAXNS), a high-performance nested sampling package written in XLA-primitives using JAX, and show that it is several orders of magnitude faster than the currently available nested sampling implementations of PolyChord, MultiNEST, and dynesty, while maintaining the same accuracy of evidence calculation. The JAXNS package is publically available at \url{https://github.com/joshuaalbert/jaxns}.

研究の動機と目的

• 時間制約や高スルーレートなデータサイエンス応用におけるネストド・サブセットリングの計算的ボトル neck を解消すること。
• 複雑でマルチモーダルかつ勾配ベースでない事後分布に対して、効率的かつ即時利用可能なベイズ証拠計算を可能にすること。
• JAX などの現代の機械学習フレームワークにネストド・サブセットリングを統合すること。ここには、かつてはネイティブ実装が存在しなかった。
• 既存のネストド・サブセットリングツールの性能制限を克服し、リアルタイムまたは大規模推論に適した速度を達成すること。
• 今後の研究で離散的事前分布およびスケーラブルでメモリ効率の良いサンプリング戦略への拡張を図ること。

提案手法

• JAX の XLA コンパイル済みプリミティブを用いて、ネストド・サブセットリングを静的メモリアルゴリズムとして再定式化し、アクセラレータ上で高性能な実行を実現すること。
• JAX の JIT コンパイルと自動微分を用いて尤度評価を実装し、高速かつベクトル化された計算を可能にすること。
• XLA の低レベル最適化を活用してメモリオーバーヘッドを最小限に抑え、計算スルーブプットを最大化すること。
• すべての中間点を保存せずに代表的なサンプルを維持できる静的メモリ版のリザーバー・サンプリングを設計すること。
• スキンリング（2004）で元々定義された連続的および離散的確率変数の両方と互換性を保つこと。
• Python スタイルのインターフェースを備えたモジュラーでオープンソースのパッケージ（JAXNS）を構築し、確率的プログラミングワークフローへのシームレスな統合を可能にすること。

実験結果

リサーチクエスチョン

• RQ1 ネストド・サブセットリングは、データサイエンス分野におけるリアルタイムおよび大規模ベイズ推論に実用的になるほど十分に高速化可能だろうか？
• RQ2 JAX の XLA コンパイルと静的メモリ設計は、従来の実装と比較してネストド・サブセットリングの性能をどの程度向上させられるだろうか？
• RQ3 ネストド・サブセットリングは、JAX などの現代の機械学習バックエンドに効率的に統合可能だろうか？これによりニューラルネットワークにおける証拠に基づくモデル選択が可能になるか？
• RQ4 JAXNS は、既存のツールよりも顕著な高速化を達成しながらも、どのようにしてベイズ証拠推定の正確性を維持しているのだろうか？
• RQ5 静的メモリネストド・サブセットリングの限界は何か？オンラインリサンプリング戦略はそれらをどのように緩和できるだろうか？

主な発見

• JAXNS は、PolyChord や MultiNEST や dynesty の順次評価モードと比較して、証拠計算が 2–3 時間の速度向上を達成している。
• この性能向上は、XLA コンパイル済みの静的メモリ計算に起因し、アクセラレータ向けにメモリアクセスと計算を最適化している。
• JAXNS は、既存のツールと同等の正確性を維持しており、信頼性のあるベイズモデル比較を保証している。
• パッケージは https://github.com/joshuaalbert/jaxns で公開されており、研究コミュニティにおける広範な採用を可能にしている。
• JAXNS は、強化学習における証拠に基づく訓練やマルチモーダル事後分布を有するベイズニューラルネットワークなど、かつては不可能とされていた応用を可能にしている。
• 今後の拡張では、離散的事前分布のサポートおよび静的メモリリザーバー・サンプリングのバリエーションを導入し、統計的忠実性を損なわずにメモリ効率を向上させることを計画している。