PyMC3 vs NumPyro

2021年11月11日2022年12月3日

最近、PyMC3 と NumPyro の両方を触る機会がありましたので、その感想をまとめたいと思います。

PyMC3 の特徴

PyMC3 は、Python 上で確率プログラミングを行える環境のひとつですが、一番の特徴は超絶簡単なそのインターフェースにあるのではないかと思っています。例えば、線形回帰の場合ですと、こんな風に非常にスッキリ書けます。

真ん中に書いてある 

mu = a * x_scaled + b

というのが線形回帰の式で、その前に書いてある

a = pm.Normal('a’, mu=0, sigma=10)
b = pm.Normal('b’, mu=0, sigma=10)

という式が、パラメータの事前分布を定義する式になります。

NumPyro の場合は、もうちょっとだけ複雑な感じになりますが、ほぼ一緒くらいの感じです。

NumPyro の場合は、モデリングの際にプレートを使えるのが大きな特徴のひとつになりますが、プレートを使わなくても表現できるモデルも多いので、人によりモデルの書き方が変わってしまう、という現象も起きてしまいます。

統計モデリングに慣れている人であれば、どっちで書いても迷いが少ないかと思うのですが、初心者がコードを見る場合は、いろいろな書き方があるとそこで迷ってしまうことがあるので、その点からも PyMC3 の方が若干有利なのかな…と個人的には思っています。

NumPyro の特徴

NumPyro の方の最大の特徴は、やはりそのサンプリングの速度にあるのではないかと思っています。PyMC3 と比較した感じだと、やはりかなり速いという印象を持ちました。

Hello Cybernetics さんのブログで紹介されている「近年の PPL の動向」というスライドにも、速度比較の結果が出ており、モデルによっては Stan より速いケースがある模様です。

また、別の方のブログでは、CPU と GPU で速度比較をしたケースについても書かれており、やはりモデルによっては GPU によりかなり高速化するケースもあるような印象を持ちました。

ただ、私の講座の方で紹介しているような簡単なモデルだと、その差があまりなかったり、かえって遅くなるケースもあったりして、その辺の挙動はいまだ謎です（使い方が悪いだけ？？）。

JAX

こうした NumPyro の速さの秘密は、実はそのバックエンドの計算エンジンである JAX にあります。

JAX は、基本的には高速な線形代数演算のためのパッケージなのですが、そこに自動微分や並列化・ベクトル化などの機能も加わり、機械学習で使われる多くの基本的な演算がこのパッケージひとつでカバーされてしまいます。

実際、DeepMind の公開している Haiku というニューラルネットワークのパッケージも実はこの JAX の上に構築されており、新しく登場する機械学習系のパッケージが段々と JAX を採用しつつあるな…という流れをなんとなく感じます。

線形代数演算だけでなく、自動微分も、並列化も、GPU / TPU による高速化も…と JAX ひとつで何でもできちゃうのね…というのは、長年 MathWorks に籍をおいて、MATLAB-MAN として生きてきた身としては何だかとても感慨深いものがありますが、そういう時代になってしまったということなのでしょう。

JAX 恐るべし…

JAX については、最近公開した新しい Udemy の講座の無料プレビューの部分でも紹介しているので、ぜひご覧になってみて下さい。

実際のところ、NumPyro を使うだけならば、JAX については使い方を知らなくても全く問題ないのですが、JAX に感銘を受けたこともあり、なんとなく 紹介したくなってしまいました ^^;

講座の無料プレビューの部分では、JIT と GPU / TPU による高速化の部分を簡単に紹介しています。

JAX に関しては、ブログの方でチュートリアルを紹介されている方もいらっしゃいますので、ぜひそちらの方も参考にしてみて下さい。私も参考にさせて頂きました。

一方で、NumPyro の弱点としては、日本語の情報があまり多くないことが挙げられます。いまのところ和書などもほとんどありません（2021年11月11日現在）。

また、バージョンも 0.8.0 と、まだバージョン１にも達していないので、その辺りは若干不安ですが、Ad-hoc な解析や勉強に使う分には十分に安定していると思うので、用途によっては積極的に使えるのではないかと思います。ちょっと青田刈りの感じもしますが、おそらく数年のうちにはバージョン１になるのではないかと思うので、青田刈りの嫌いでない方はぜひお試し頂けたらと思います。