変分推論をpythonで実装する

マサムネのイメージ図 Python R

変分推論によるポアソン混合モデルの近似を、pythonで実装して遊んでみる記事です。
詳しい計算は、変分推論の記事を読んでみてください。

変分推論入門
変分推論でポアソン混合モデルを近似する式を導きます。最後に、EMアルゴリズムとの接点を紹介します。

記事で使っているソースコードはgithub に置いてあります。
https://github.com/msamunetogetoge

元ネタは、PRMLの10章 と緑本の4章です。

変分推論

初めに、使うモデルや、更新式をメモしておきます。モデルはポアソン混合モデルで、以下の式で定義されます。
$$\begin{eqnarray}
p(X, S , \lambda , \pi ) &=& p(X|S, \lambda )p(S|\pi)p(\pi ) p(\lambda ) \\
p(X| S , \lambda ) &=& \prod_{n} p(x_n | s_n , \lambda ) \\
p(x_n | s_n , \lambda ) &=& \prod_{k} {\rm Poi} (x_n |\lambda _k ) ^{s_{n,k} } \\
p(S |\pi ) &=& \prod p(s_n|\pi) \\
p(s_n|\pi) &=& {\rm Cat}(s_n , \pi ) = \prod \pi_k ^{s_{n,k}} \\
p(\lambda _k ) &=&{ \rm Gam } (\lambda _k | a,b ) \\
&=& C_{G}(a, b ) \lambda ^{a-1} \exp(-b\lambda _k )\\
p(\pi )= &=& { \rm Dir} (\pi |\alpha ) \\
&=&C_{D}(\alpha ) \prod \pi _{k} ^{\alpha _k -1}
\end{eqnarray}$$
1
このモデルは、\(S , \lambda \)を与える事で、データ\(x \)をサンプリングすることが出来ます。
\(S , \lambda \)も確率分布からサンプリングする必要がありますが、それには、\(a, b, \alpha \)を決定しなくてはなりません。
\( a, b, \alpha \)を決めるための更新式は以下の式になります。
$$\begin{eqnarray}
q^{\ast}(s_n ) &=& {\rm Cat }(s_n|\eta _n )\\
\eta _{n,k} &\propto & x_n E_{\lambda _k }[ \log \lambda _{k} ]- E_{\lambda _k }[ \lambda _k ]
+ E_{\pi _k } [\log \pi _k] \\
q^{\ast} (\lambda _k ) &=& { \rm Gam } (\lambda _k|\hat{a}_k , \hat{b}_k )\\
\hat{a}_k &=& \sum _{n} E_S [s_{n,k} ]x_n +a \\
\hat{b}_k &=& \sum_{n} E_S [s_{n,k} ] + b \\
q^{\ast} (\pi ) &=& {\rm Dir } (\pi |\hat{\alpha } )\\
\hat{\alpha _k } &=& \sum_{n} E_{S} [s_{n,k} ] +\alpha _{k}
\end{eqnarray}$$
期待値は以下の式から計算出来ます。
$$ \begin{eqnarray}
E_{\lambda _k }[ \lambda _{k} ] &=& \hat{a} /\hat{b}\\
E_{\lambda _k }[ \log \lambda _{k} ] &=& \psi (\hat{a} ) – \log \hat{b} \\
E_{\pi _k }[ \log \pi _{k} ] &=& \psi (\hat{\alpha _k }) -\psi (\sum \hat{ \alpha _k })\\
E_S[ s_{n,k} ] &=& \eta _{n,k}
\end{eqnarray}$$
事前分布と、パラメーター更新後の分布が一致しているので、パラメーターを更新する数を決めて計算することで、与えられたデータの、ポアソン混合モデルによる近似が得られます。

変分推論のpython による実装

適当なデータを生成しておいて、上記の変分推論でフィッティングしてみます。

変分推論を行うクラスをpythonで実装します。パラメーターの意味は、上でした計算と同じです。

一応クラスの中身の説明をします。
Estep でパラメーターを更新するのに必要な期待値を計算しています。
Mstepでパラメーターを更新しています。
itr_calc で指定した回数だけ計算し、\( \lambda \)を返しています。この時に、\( \pi \)の値が分かっていれば、その割合でポアソン分布を生成すればよいので、クラスで値を保持しています。
計算させてみましょう。

データは\(\lambda =3, 15 \)のポアソン分布を\( 7:3 \)の割合で生成していました。結構いい感じにフィッティング出来ています。重ねてヒストグラムにしてみましょう。

元データと予測したパラメーターを元に生成したデータ
元データと予測したパラメーターを元に生成したデータ

\( \lambda \)がどのように推移しているかを見てみます。

\( \lambda \)の更新の様子

グラフを見ると、500回計算したころには、収束しているような気がします。パラメーターの更新を、計算回数で制御するのでなくて、ELBOで制御する方法があります。

EM-like 変分推論

ELBO \( \mathcal{L}(q, \theta) \)は以下の式で定義されます。
$$\begin{eqnarray}
\log p(X| \theta ) &=& \mathcal{L}(q, \theta) +KL(q \| p) \\
\mathcal{L}(q, \theta) &=& E_{Z}\left[ \log \left\{ \frac{p(X,Z| \theta )}{q(Z)} \right\} \right]
\end{eqnarray} $$
ポアソン混合モデルの、場合は以下のように計算出来ます。
$$\begin{eqnarray}
\mathcal{L}(q, \theta) &=& \sum_n
\left\{ E_{\lambda} [ E_{s_n} [ \log p(x_n |s_n , \lambda ) ]]+
E_{s_n} [E_{\pi}[ \log p(s_n | \pi ) ]] – E_{s_n} [\log q_{ {\rm new } }(s_n )] \right\} \\
&-& KL[ q_{ {\rm new } }( \lambda ) ][ q_{ {\rm old } }( \lambda )] – KL[ q_{ {\rm new } }( \pi ) ][ q_{ {\rm old } }( \pi )]
\end{eqnarray} $$
期待値を計算します。
$$\begin{eqnarray}
E_{\lambda} [ E_{s_n} [ \log p(x_n |s_n , \lambda ) ]] &=& \sum_k \left( E_{s_n} [s_{n,k} ] ( x_n E_{\lambda _k }[ \log \lambda _{k} ]- E_{\lambda _k }[ \lambda _k ] ) \right) \\
E_{s_n} [E_{\pi}[ \log p(s_n | \pi ) ]] &=& \sum_k \left( E_{s_n} [s_{n,k} ] E_{\pi}[\log \pi _k ] \right)\\
E_{s_n} [\log q_{ {\rm new } }(s_n )] &=&\sum_k E_{s_n} [s_{n,k} ] \pi_{k} \\
KL[ q_{ {\rm new } }( \lambda ) ][ q_{ {\rm old } }( \lambda )] &=&
\sum_k \left( ( {a_{\rm new }}_k – {a_{\rm old }}_k )E_{\lambda _k } [\log \lambda _k ] –
({b_{\rm new }}_k – {b_{\rm old }}_k ) E_{\lambda _k } [ \lambda _k ] \right) \\
KL[ q_{ {\rm new } }( \pi ) ][ q_{ {\rm old } }( \pi )] &=&
\sum_k \left( {\alpha_{\rm new }}_k – {\alpha_{\rm old }}_k \right) E_{\pi _k} [\pi _k]
\end{eqnarray}$$
python でELBOを計算する関数を定義して、計算させてみましょう。

前に定義したクラスVI を継承して、Estep, Mstepをリサイクルしています。
itr_elbo を呼び出すと、更新後にELBOの差が0.01 になるまで計算を繰り返し、selg.Eに計算値が格納されます。遊ぶときは、max_itr とELBOの差を変数にした方が良いかもしれません。2
前に生成したデータをフィッティングして、グラフを描いてみます。

ELBOの計算値
ELBOの計算値

ELBOを監視することにより、計算回数を大幅に減らすことが出来ました。3
github のjupyter notebook には、混合ポアソンモデルでは上手くフィッティング出来ないデータも載せているので、よかったら見てください。

まとめ

  • 混合ポアソンモデルの変分推論をpythonで実装した。
  • 混合ポアソンモデルの場合に、elboを計算した。
  • elboを計算するクラスをpythonで実装した。
  1. numpy.random.Gamma()では、\( \theta =1/b \)がパラメーターで使われているので注意が必要です。
  2. ELBOが少し振動することがあるので、計算を打ち切るのが上手くいかなかったりします。
  3. 初期値によっては、上手くいかない時もあります。
タイトルとURLをコピーしました