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

こんにちは。吉田弁二郎です。

前回の記事yoshidabenjiro.hatenablog.comでは、変分ベイズ法の定式化についてまとめました。今回は、この方法を具体的な事例に応用していきます。

混合ポアソンモデルの構成

ここでは混合ポアソン分布を例に取り、変分ベイズにおける更新式
\begin{equation}
\ln q^{*}(Z_{i}) = \ln \tilde{P}(X, Z_{i}) + const. \tag{1} \label{1}
\end{equation}を実際に使ってみます。ここで、
\begin{equation}
\ln \tilde{P}(X, Z_{i}) = \int \prod_{j\neq i}dZ_{j}q_{j}\ln P(X, Z).
\end{equation}です。

$K$ 個のクラスから $N$ 個のデータ $X = \{x_{i} \}$ が得られているとします。ただし、各データの次元は $D$ とし、各データおよびデータの各次元は独立にポアソン分布に従うとします。

f:id:yoshidabenjiro:20161229210337p:plain
こちらは手元で作成したサンプルデータで、 $K=4$ , $N=500$ , $D=2$ かつ各データ点ごとに、 $x_{1}$ , $x_{2}$ が独立にポアソン分布に従っています。データ点は、割り当てられた $K$ によって色分けされてます。こちらの生成コードは次回紹介したいと思います。

さて、同時分布は
\begin{align}
P(X, Z, \lambda, \pi) &= P(X|Z, \lambda)P(Z|\pi)P(\lambda)P(\pi)\\
&= \prod_{n,k,d}Pois(x_{nd}|\lambda_{kd})^{z_{nk}}\prod_{n,k}\pi_{k}^{z_{nk}}\prod_{k,d}Gam(\lambda_{kd}|r, s)Dir(\pi|\alpha)
\end{align}と書けます。ここで $Z=\{z_{nk}\}$ は各データが生成されたクラスを区別する潜在変数で、one-of-k 符号化されているとします*1。つまり、データ $x_{n}$ がクラス $k$ から生成された場合、 $z_{nk}=1$ かつ $z_{nj\neq k}=0$ という2値を取る変数です。各データ $x_{n}$ がポアソン分布に従い $z_{nk}$ がカテゴリカル（離散確率）分布に従うのに応じて、パラメータ $\lambda_{kd}$ と $\pi_{k}$ はそれぞれ共役分布となるガンマ分布とディリクレ分布に従うとしました。本来なら各 $\lambda_{kd}$ ごとに別々のパラメータで決まるガンマ分布を採用してもよいのですが、簡単のために共通のパラメータ（上式の $r, s$ ）で決まることにしています。

混合ポアソンモデルに対して変分ベイズ法を適用する

変分ベイズ法の一般論と混合ポアソンモデルの同時分布が定義できたので、具体的に分布の更新式を書き下していきます。

まず、変分ベイズの常套手段として、事後分布 $q(Z, \lambda, \pi)$ に対して次のような近似を課します：
\begin{equation}
q(Z, \lambda, \pi) = q(Z)q(\lambda, \pi).
\end{equation}潜在変数とそれ以外のパラメータを独立なものとして扱ってみましょう、ということです。

初めに $q(Z)$ について\eqref{1}式を適用すると、
\begin{align}
\ln q^{*}(Z) &= \int d\lambda d\pi q(\lambda, \pi)\ln P(X, Z, \lambda, \pi) + const.\\
&= \int d\lambda d\pi q(\lambda, \pi)\{\ln P(X|Z, \lambda) + \ln P(Z|\pi) + \ln P(\lambda) + \ln P(\pi)\} + const.\\
&= \langle \ln P(X|Z, \lambda)\rangle_{\lambda} + \langle\ln P(Z|\pi)\rangle_{\pi} + const.
\end{align}となります（ $\langle \rangle$ は分布 $q$ を用いた平均操作）。ポアソン分布の具体形
\begin{equation}
Pois(x_{nd}|\lambda_{kd}) = \frac{\lambda_{kd}^{x_{nd}}}{x_{nd}!}e^{-\lambda_{kd}}
\end{equation}を使うと
\begin{align}
\langle \ln P(X|Z, \lambda)\rangle_{\lambda} &= \sum_{n, k, d}z_{nk}(x_{nd}\langle \ln \lambda_{kd}\rangle_{\lambda_{kd}} - \langle \lambda_{kd}\rangle_{\lambda_{kd}}) + const., \\
\langle \ln P(Z|\pi)\rangle_{\pi} &= \sum_{n,k}z_{nk}\langle \ln \pi_{k}\rangle_{\pi_{k}}
\end{align}となり、結論、最適解 $q^{*}(Z)$ の関数形として
\begin{equation}
\ln q^{*}(Z) = \sum_{n,k}z_{nk}\Bigl\{ \sum_{d}\Bigl (x_{nd}\langle \ln \lambda_{kd}\rangle_{\lambda_{kd}} - \langle \lambda_{kd}\rangle_{\lambda_{kd}}\Bigr) + \langle\ln \pi_{k}\rangle_{\pi_{k}}\Bigr\} + const. \tag{2} \label{2}
\end{equation}が得られます。よく見てみると $n, k$ の和に対して $z_{nk}$ がくくり出された形をしているため、元の分布同様にカテゴリカル（離散確率）分布になっていることが分かると思います。ただ、各クラスのデータが生じる確率がデータを区別する添字 $n$ にも依存するようになっている点は元と異なっています。

次に $q(\lambda, \pi)$ についても同様に計算すると、
\begin{split}
\ln q^{*}(\lambda, \pi) &= \int dZq(Z)\ln P(X, Z, \lambda, \pi) + const.\\
&= \ln P(\lambda) + \ln P(\pi) + \langle \ln q(X|Z, \lambda)\rangle_{Z} + \langle \ln P(Z|\pi)\rangle_{Z} + const.
\end{split}平均値の計算については、
\begin{align}
\langle\ln P(X|Z, \lambda) \rangle_{Z} &= \sum_{n,k,d} \langle z_{nk}\rangle \ln Pois(x_{nd}|\lambda_{kd})\\
&= \sum_{n,k,d}\langle z_{nk}\rangle( x_{nd}\ln\lambda_{kd}-\lambda_{kd}) + const.\\
\langle\ln P(Z|\pi)\rangle &= \sum_{n,k}\langle z_{nk}\rangle\ln\pi_{k}
\end{align}より最適解 $q^{*}(\lambda, \pi)$ として
\begin{align}
\ln q^{*}(\lambda, \pi) &= \sum_{k}\Bigl\{\Bigl(\sum_{n}\langle z_{nk}\rangle + \alpha_{k} - 1\Bigr)\ln\pi_{k}\\
&\quad\quad + \sum_{d}\Bigl(\sum_{n}\langle z_{nk}\rangle x_{nd} + s -1\Bigr)\ln\lambda_{kd} \\
&\quad\quad - \sum_{d}(\sum_{n}\langle z_{nk}\rangle + r)\lambda_{kd} \Bigr\}\\
&\quad + const. \tag{3} \label{3}\\
&= \ln q^{*}(\pi) + \ln q^{*}(\lambda)
\end{align}を得ます。特徴として、1) $q^{*}(\lambda)$ と $q^{*}(\pi)$ が分離し $\lambda$ と $\pi$ が独立になることが自然に導かれる、2) $q^{*}(\lambda)$ はガンマ分布、 $q^{*}(\pi)$ はディリクレ分布となり元の分布と変わらない、という2点が挙げられます。

長くなってきましたが、あと一歩です。\eqref{2}と\eqref{3}で用いていた各パラメータについての平均値は、簡単な形で書くことができます。まず $\langle z_{nk}\rangle = \pi_{nk}$ は離散確率分布の関数形から明らかで、 $\ln\lambda$ については、
\begin{align}
\langle \ln\lambda \rangle &= \int_{0}^{\infty}d\lambda \frac{r^{s}}{\Gamma(s)}e^{-r\lambda}\lambda^{s-1}\ln\lambda\\
&= \int d\lambda \frac{r^{s}}{\Gamma(s)}e^{-r\lambda}\frac{d}{ds}\lambda^{s-1}\\
&= -\int d\lambda \frac{d}{ds}\Bigl(\frac{r^{s}}{\Gamma(s)}\Bigr)e^{-r\lambda}\lambda^{s-1}\\
&= -\int d\lambda \ln r\frac{r^{s}}{\Gamma(s)}\lambda^{s-1}e^{-r\lambda} + \int d\lambda r^{s}\frac{\Gamma'(s)}{\Gamma(s)^{2}}\lambda^{s-1}e^{-r\lambda}\\
&= -\ln r + \frac{\Gamma'(s)}{\Gamma(s)}\\
&= -\ln r + \psi (s).
\end{align}ここで $\psi(s)\equiv \frac{d}{ds}\ln\Gamma(s)$ はディガンマ関数です。また $\ln\pi_{k}$ についても、ディリクレ分布の規格化定数を $C(\alpha)\equiv \frac{\Gamma(\alpha_{0})}{\prod_{k}\Gamma(\alpha_k)}$ , $\alpha_{0}\equiv\sum_{k}\alpha_{k}$ として
\begin{align}
\langle \ln\pi_{k}\rangle &= C(\alpha)\int \prod_{j}d\pi_{j}\ln\pi_{k}\prod_{l}\pi_{l}^{\alpha_{l}-1}\\
&= C(\alpha)\int \prod_{j}d\pi_{j}\frac{\partial}{\partial\alpha_{k}}\prod_{l}\pi_{l}^{\alpha_{l}-1}\\
&= -\frac{\partial}{\partial \alpha_{k}}\ln C(\alpha)\\
&= \frac{d}{d\alpha_{k}}\ln\Gamma(\alpha_{k}) - \frac{\partial}{\partial \alpha_{k}}\ln\Gamma(\alpha_{0})\\
&= \psi(\alpha_{k}) - \psi(\alpha_{0}).
\end{align}これらを\eqref{2}と\eqref{3}に代入すれば、混合ポアソン分布に対する変分ベイズ法の更新式が得られます：
\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}
\tag{4} \label{4}
\end{eqnarray}ただし $\pi_{nk}$ は都度規格化しておく必要があります。これらに応じて $Z, \lambda, \pi$ の確率分布関数は
\begin{eqnarray}
\begin{cases}
q^{*}(Z, \lambda, \pi) = q^{*}(Z)q^{*}(\lambda)q^{*}(\pi), & \\
q^{*}(Z) = \prod_{n,k}\pi_{nk}^{z_{nk}},& \\
q^{*}(\lambda) = \prod_{k,d}\frac{r_{k}^{s_{kd}}\lambda_{kd}^{s_{kd}-1}e^{-r_{k}\lambda_{kd}}}{\Gamma(s_{kd})},& \\
q^{*}(\pi) = \frac{\Gamma(\alpha_{0})}{\prod_{j}\Gamma(\alpha_{j})}\prod_{k}\pi_{k}^{\alpha_{k}-1}
\end{cases}
\tag{5} \label{5}
\end{eqnarray}と決定されます。