numpyとscipyを用いて混合ポアソン分布の事後分布を計算する

問題設定

簡単のためデータ$x\in\mathbb{Z}_{\geq0}$は$1$次元データとする。

• 統計モデル：混合ポアソン分布
  • $$p(x|\pi,\lambda)=\sum_{c=1}^C\pi_c\mathrm{Poisson}(x|\lambda_c)=\sum_{c=1}^C\pi_c\frac{\lambda_c^x}{x!}e^{-\lambda_c}$$
  • $\pi=(\pi_c)_{c=1}^C\in \mathbb{S}^{C-1}$は混合比率、$\lambda=(\lambda)_{c=1}^C\in\mathbb{R}_{\geq0}^C$はポアソン分布の強度パラメータ。（$\mathbb{S}^{C-1}=\{s\in\mathbb{R}^C|s_c\geq0,\ \sum_{c=1}^Cs_c=1\}\subset\mathbb{R}^C$は単体。）
• 事前分布：
  混合比率と強度でそれぞれ共役で独立な事前分布を用いる。（混合ポアソン分布全体として共役にはならない。）
  • 混合比率：ディリクレ分布
    • $$\varphi(\pi|\alpha)=\mathrm{Dirichlet}(\pi|\alpha)\propto\prod_{c=1}^C\pi_c^{\alpha_c-1}$$
    • $\alpha=(\alpha_c)_{c=1}^C\in\mathbb{R}_{>0}^C$は集中度と呼ばれるハイパーパラメータ。
  • 強度：ガンマ分布
    • $$\varphi(\lambda_c|\beta_c, \gamma_c)=\mathrm{Gamma}(\lambda_c|\beta_c, \gamma_c)\propto\lambda^{\beta_c-1}e^{-\gamma_c\lambda_c}$$
    • $\beta_c>0$は形状母数、$1/\gamma_c>0$は尺度母数と呼ばれるハイパーパラメータ。

ギブスサンプリング

サメーション記号と潜在変数を用いることで解析的に事後分布を求めることは可能だが、項数がデータ個数の応じて指数関数的に増加するので厳密に計算するのは現実的ではない。従ってMCMCを用いて近似的に事後分布を導出する。本記事では潜在変数を導入したギブスサンプリングを実装する。用いる潜在変数は各データ$x_i (i=1,\ldots,n)$ごとに対応し、そのデータが属するクラスタとして解釈出来るものである。潜在変数はワンホットベクトル$y_i\in\mathbb{Z}_{\geq0}^C$で表す。

1. 潜在変数$(y_i)_{i=1}^n$を初期化する
2. 以下の分布を用いて混合比率を更新
   
   $$\pi \sim \mathrm{Dirichlet}(\alpha+\sum_{i=1}^ny_i)$$
3. 以下の分布を用いて強度を更新
   
   $$\lambda_c \sim \mathrm{Gamma}(\beta_c+\sum_{i=1}^nx_iy_{i,c},\ \gamma_c+\sum_{i=1}^ny_{i,c})$$
4. 以下の分布を用いて潜在変数を更新
   
   $$y_i \sim \mathrm{Categorical}(\frac{\pi_c\mathrm{Poisson}(x_i|\lambda_c)}{\sum_{c=1}^C\pi_c\mathrm{Poisson}(x_i|\lambda_c)})$$
5. 2.に戻る。適当な回数まで繰り返す。

高速化のための工夫

ギブスサンプリングはパラメータを互いに更新し合うため、サンプリングのイテレーションは原理上省くことは出来ない。従って、一回のイテレーションの速さを以下に短くするかが重要である。

• データと潜在変数を集約する
  扱うデータが離散であることとギブスサンプリングの更新式から、同じ値のデータを予め集約して計算することが可能である。データの次元が低次元であるか分散が小さければ著しい高速化が期待できる。高次元や分散が大きい場合は、データの種類がデータの個数とほぼ同じであるためそのまま実装する場合と速度はあまり変わらない。
• 複数の系列を同時に計算する
  計算能力に余裕がある場合は、複数のギブスサンプリングの系列を同時に配列に格納して計算することで、より短い時間で同じ数のギブスサンプルを取得することが出来る。
• 等価な分布により乱数生成する
  潜在変数の計算ではカテゴリカル分布を用いているが、np.random.multinomialを使用するとなるとデータの種類ごとに分けてイテレーションする必要があり計算が遅くなるが、np.random.binomialを使用してカテゴリに関するイテレーションで代用することが出来る。ただしカテゴリ数の方がデータの種類よりも十分に小さいという仮定を置いている。
• ufuncで可能な限り置き換える
  ギブスサンプリングの4.のCategorical分布のパラメータを計算する際に、Poisson質量関数を用いているが、scipy.statsの関数を使うよりはscipy.specialのufnucである階乗関数等を使って定義を直接書き下す方が高速に計算出来る。

実験

PyMCと本実装で計算速度を結果を比較した一例が以下である。

• 実験設定
  • データの個数：$N = 3000$
  • 一系列あたりのギブスサンプル数：$S = 500$
  • ポアソン分布の平均：$\lambda_1=3,\ \lambda_2=20$
  • 混合比率：$\pi_1=0.1,\ \pi_2=0.9$
  • バーンイン：$5\times N = 15000$
  • 系列数：$L=4$
• 計算結果

  PyMC	ours
  20.69[s]	1.19[s]

• 結果の詳細はこちら