PPL比較：Stan、NumPyro、PyMC

はじめに

最近はPPLのStan、NumPyro、PyMCを触っていたのですが、それぞれ一長一短だったため個人的な感覚での評価値ですが、各ライブラリの比較表を残しておきます。
各ライブラリの細かい使い方は記事一覧から確認してください。

比較表

階層モデル〜DL系までのカテゴリは書きやすさで評価しています。

	Stan	NumPyro	PyMC	MEMO
階層モデル	△	◯	◯	NumPyroとPyMCでは階層のところを行列演算で自然に書くことができる
ガウス過程	△	△	◯	PyMCではカーネル関数やスパースGPなどデフォルトで用意されている
時系列モデル（AR系）	◯	◯	◯	各言語で同じ程度の学習コストで書くことができる
状態空間モデル	◯	△	△	NumPyroやPyMCではscan関数に慣れる必要あり。PyMCではexperimentalとしてガウス・線形の場合は実装された
欠損値	◯	◯	◎	PyMCでは欠損値を自動で対応してくれる
離散潜在変数	◯	◎	◯	NumPyroでは周辺化を内部で実施してくれる
打ち切り/切断	◯	◎	◎	打ち切り/切断用のクラスが用意されている
DL系	×	◯	×	NumPyroはDLのライブラリと連携可能かつGPU対応
コンパイル速度	×	◎	◯	Stanは遅い、NumPyroは体感ほとんど感じない
実行速度	◯	◯	△	PyMCはバックエンドにNumPyroが設定できるモデルであれば速いが、現状一部のモデルで動作しない
フォーラムやドキュメントの充実度	◎	△	◯	歴史あるStanは強い。pymcは活発だが、活発が故にバージョン違いのコードや説明が混在
カスタマイズ性	◯	×	×	NumPyroとPyMCは既存の用意されていないことをやろうとするとソースコードを見に行かないと分からないことも多い

個人的な結論

現状、NumPyroやPyMCは絶賛改善中で機能やドキュメント不足、それらに伴いカスタマイズが困難という課題があります。そのため、ドキュメント内のExampleや単純な階層モデルなどは書きやすいのですが、一歩進んだ改良がしたくなった時にStanよりかなり大変でした。また、カスタマイズが難しい点は内部のロジックに由来するところも大きいので、根本的な解決は難しいのではないかと思っています。
また、Stanはカスタマイズしやすく、幅広いモデルに対応可能ですがPyMCのガウス過程のモジュールのようなまとまった部品が提供されていません。

将来的に、多様で適度な粒度の部品を持つPytorchのようなライブラリが出てくるといい（個人的には全てPythonで書きたいのでpymcの機能がどんどん拡充されていくことを期待してます）のですが、現状は幅広いモデルに対応したいのであればStanをベースにして自身で部品を揃えていくのが一番いい気がしています。一方で、階層モデル、ガウス過程やDL、因果推論などトピックに絞って使用するのであればそれぞれ得意なNumPyroやPyMCを選ぶことも良い選択肢かと思いました。

最後に、最初に書いた通り個人的な感覚で書いたので、雰囲気だけ参考にして細かいところは各ドキュメント等を参考にしてください。