混合ポアソン分布を題材に、変分ベイズ法を理解する③

明けましておめでとうございます。吉田弁二郎です。

前回の記事yoshidabenjiro.hatenablog.comで、混合ポアソン分布の変分ベイズ更新式を導出しました。今回は、これらを実装してアルゴリズムの動作を確認していきます。

コードについては後述するとして、まず、サンプルデータを生成してみます。
f:id:yoshidabenjiro:20170101023929p:plain
クラス数 $K=4$ , データ次元 $D=2$ , サンプル数 $N=500$ で、潜在変数 $Z$ の事前ディリクレ分布を定めるパラメータは $\alpha=[1, 1, 1, 1]$ , データを生成するポアソン分布のパラメータ $\lambda$ の事前分布、つまりガンマ分布を定めるパラメータは $r=\frac{300}{12000}$ , $s=\frac{300}{r}$ とし、10セット分のサンプルを用意しています。300 と 12,000 はそれぞれ、ガンマ分布の平均と分散の値です。

推論がうまくいくデータ、いかないデータ

上の図を見てみると、データセットによっては結構クラス間のかぶりがあったりして推論が難しそうな気がしてきます。そんな中でも trial=7 のものはクラスがきれいに分離していますね。まずはこちらをサンプルとして変分ベイズ推論を行ってみたものを下図に示します。
f:id:yoshidabenjiro:20170101023938p:plain
左がオリジナルのデータ、右が推定されたデータです。クラスが明確に４つに分かれていますね。今回は変分推論のセットアップとして $q(Z, \lambda, \pi)=q(Z)q(\lambda, \pi)$ という近似を行いましたが、結果としては正しく推論できているようです。同一のクラスに対して左と右とで異なった番号が割り当てられた部分がありますが、これはモデルにパラメータを入れ替える自由度があるせいで、推論の正確さとは違った問題として認識しておくべきだと思います。パラメータ入れ替えの自由度があるモデルは格別に興味深いので、そのうち改めて何か書くつもりです。

クラスが一部重なり合った trial=6 ではどうでしょうか。
f:id:yoshidabenjiro:20170101025832p:plain
推論をさせてみると、全体としてはうまくいっているように見えるものの、やはり境界領域では誤判断が起こっていることが分かります。データ数を増やすなどしないといけなさそうです。

変分ベイズ推論の実装

今回実装すべきパラメータ更新式を復習しておきましょう：
\begin{eqnarray}
\begin{cases}
\pi_{nk}^{new} \propto \exp\Bigl(\sum_{d}\bigl(x_{nd}\bigl(\psi(s^{old}_{kd})-\ln r^{old}_{k}\bigr) - \frac{s^{old}_{kd}}{r^{old}_{k}}\bigr) + \psi(\alpha_{k}^{old}) - \psi(\alpha_{0}^{old})\Bigr),& \\
s^{new}_{kd} = \sum_{n}\pi_{nk}^{old}x_{nd}+s,& \\
r^{new}_{k} = \sum_{n}\pi_{nk}^{old} + r,& \\
\alpha_{k}^{new} = \sum_{n}\pi_{nk}^{old} + \alpha_{k}.
\end{cases}
\end{eqnarray}記号の定義などは前回の記事にあるので適宜ご参照ください。

実装は Python3 で行いました。blog/pmm.py at master · dlnp2/blog · GitHub にコードを置いてあるので、ここでは要所だけ確認します*1。

""" Bayesian Poisson mixture model """
import numpy as np
from numpy import dot
from scipy.special import digamma as psi
from scipy.misc import logsumexp
from pandas import DataFrame, concat


def fit(self, X, init_pi, init_r, init_s, init_alpha, n_iter):
    self.pi_inf = np.array([np.array(init_pi)])
    self.r_inf = np.array([np.array(init_r)])
    self.s_inf = np.array([np.array(init_s)])
    self.alpha_inf = np.array([np.array(init_alpha)])
    for _ in range(n_iter):
        r_old = self.r_inf[-1]
        s_old = self.s_inf[-1]
        alpha_old = self.alpha_inf[-1]
        ln_pi_new = [dot(X, psi(s_old[k]) - log(r_old[k])) -
                     (s_old[k] / r_old[k]).sum() +
                     psi(alpha_old[k]) - psi(alpha_old.sum())
                     for k in range(self.K)]
        pi_new = np.exp(np.vstack(ln_pi_new) -
                        logsumexp(ln_pi_new, axis=0)).T
        s_new = dot(pi_new.T, X) + init_s
        n_new = pi_new.sum(axis=0)
        r_new = n_new + init_r
        alpha_new = n_new + init_alpha
        self.pi_inf = np.append(self.pi_inf, [pi_new], axis=0)
        self.r_inf = np.append(self.r_inf, [r_new], axis=0)
        self.s_inf = np.append(self.s_inf, [s_new], axis=0)
        self.alpha_inf = np.append(self.alpha_inf, [alpha_new], axis=0)

def predict(self, trial, init_pi, init_r, init_s, init_alpha, n_iter):
    X_original = self.pick_sample(trial)
    X_inferred = X_original.copy()
    X_original['trial'] = 'original'
    X_inferred['trial'] = 'inferred'
    self.fit(X_inferred[['x1', 'x2']].values, init_pi, init_r, init_s,
             init_alpha, n_iter)
    X_inferred['z'] = np.argmax(self.pi_inf[-1], axis=1)
    self.show_samples(data=concat([X_original, X_inferred]))

とあるクラスを作って作業したので self 引数が入っています。テクニカルに重要だったのは

from scipy.misc import logsumexp

の部分で、これは文字通り指数関数の和の対数を計算するものです。今回の場合もそうでしたが、指数関数の引数が大きな値になって overflow が生じるのを回避するのに必要です。この事情のため、 $\pi$ の計算に関しては

ln_pi_new = [...]
pi_new = np.exp(... ln_pi_new ...).T

として一度対数を取って計算した上で指数関数に入れ戻しています。また、１回のイテレーションの中で $r, s, \alpha$ が $\pi^{new}$ によって更新されるのがミソです。その他の部分については、更新式を素直に実装しているという感じ。NumPy はベクトル演算をしてくれるので本当に助かりますね。

このコードを使って遊ぶには、

K = 4
D = 2
N = 500
ALPHA = [1, 1, 1, 1]
mean_gamma = 300
var_gamma = 12000
R = mean_gamma / var_gamma
S = mean_gamma * R
pmm = PMM(K, D, N, ALPHA, S, R)
pmm.add_samples(10)
pmm.show_samples()

としてサンプルデータを作成した後、適当な trial 番号（下では 6 を選択）を選んで

init_pi = np.array([0.1, 0.4, 0.1, 0.2] * N).reshape((N, K))
init_r = [0.01, 0.01, 0.01, 0.01]
init_s = np.array([5, 5] * K).reshape((K, D))
init_alpha = [0.5, 0.7, 0.4, 0.8]
n_iter = 100
pmm.predict(6, init_pi, init_r, init_s, init_alpha, n_iter)

とすればOK。イテレーション回数が増えるにつれクラスがきれいに分かれていく様子を見るとほっこりします。

なお、動作を確認した環境は、

Mac OSX 10.11.16
Python3 (Anaconda3-4.0.0)
NumPy 1.11.0
SciPy 0.17.0
Pandas 0.19.2
Jupyter Notebook

です。

*1:ここに掲載したコードだけコピペしても動かないのでご注意を！

luggage baggage

Machine learning, data analysis, web technologies and things around me.

混合ポアソン分布を題材に、変分ベイズ法を理解する③

推論がうまくいくデータ、いかないデータ

変分ベイズ推論の実装