今回は、EMアルゴリズムを使ってポアソン混合モデルによる学習を行うことにする。理論編です。

EMアルゴリズムは、E-StepとM-Stepから構成される最適化の手法である。今回の記事では、どこがE-StepだとかM-Stepだとかは最後に紹介するだけにとどめる。

混合分布を考える

ポアソン混合モデルでは、ポアソン混合分布、つまり、いくつかのポアソン分布を合わせた分布を扱う。

下記に、3つのポアソン分布を混ぜたポアソン混合分布を描いてみた。

import numpy as np
from scipy.special import factorial
import matplotlib.pyplot as plt

x = np.linspace(start=0, stop=50, num=1000)
# ポアソン分布の尤度関数だけ作っておいた
uni_poisson = lambda x, lambda_: (lambda_ ** x) / factorial(x) * np.exp(-lambda_)
poisson1 = uni_poisson(x, 3)
poisson2 = uni_poisson(x, 10)
poisson3 = uni_poisson(x, 40)

# 合計が1になるような3つの重み(2次元標準シンプレックス)を作って足し合わせる
mixed_poisson = poisson1 * 0.2 + poisson2 * 0.3 + poisson3 * 0.5

plt.plot(x,mixed_poisson)

前回までの復習

ポアソン混合モデルからデータが生成する過程は以下。
(前回までの記事と順序が異なる箇所があるが、問題ない。)

１ー１．クラスタの混合割合 $ \textbf{π} $ が、$ Dir(\textbf{π}|\textbf{α}) = C_D(\textbf{α}) \prod_{k=1}^{K} {\pi_k}^{\alpha_k - 1} $からサンプルされる。

１ー２．選択されたクラスタ $k$ のパラメータ $ λ_k $ が $ Gam(λ_k|a_k, b_k) $ からサンプルされる。

２．クラスタ番号を表すベクトル $ \textbf{s} $ が、$ Cat(\textbf{s}|\textbf{π}) = \prod_{k=1}^{K} {\pi_k}^{s_k} $ からサンプルされる。

３．$ Poi(x|λ_k) $ から $ x $ がサンプルされる。

今回、パラメータは事前分布からサンプルしません。

ポアソン混合モデルでは、パラメータの事後分布を解析的に求めやすくするために、パラメータに共役事前分布からの発生を仮定している。

この記事で採用するアルゴリズムでは、パラメータの事後分布は計算しないので、てきとうに初期化しちゃえ。

ということで、データの生成過程を以下のように簡略化する。

１．クラスタの混合割合 $ \textbf{π} $ と $ λ_k　(k = 1, 2, .. , K) $ をてきとうに初期化する

２．クラスタ番号を表すベクトル $ \textbf{s} $ が、$ Cat(\textbf{s}|\textbf{π}) = \prod_{k=1}^{K} {\pi_k}^{s_k} $ からサンプルされる。

３．$ Poi(x|λ_k) $ から $ x $ がサンプルされる。

上記の1～3の手順を数式で簡単に表せば

$$
p(x, \textbf{s} | \textbf{λ}, \textbf{π}) =
p(x|\textbf{s}, \textbf{λ})p(\textbf{s}|\textbf{π})
$$

となる。式を右から見て、

・$ \textbf{π} $ に基づいて $ \textbf{s} $ が生成される。

・$ \textbf{s} $ と $ \textbf{λ} $ に基づいて $ x $ が生成される。

という流れが見て取れるであろう。

ポアソン混合モデルにおける尤度関数は $ p(x, \textbf{s} | \textbf{λ}, \textbf{π})$ だという認識でいてくれればいいと思う。

一応 $ \textbf{s} $ について言及しておく。

$ \textbf{s} $ はパラメータじゃないの？尤度関数は $ p(x | \textbf{s} ,\textbf{λ}, \textbf{π})$ じゃないの？　と一瞬思ってしまった昔の私のために。

$ p(x | \textbf{s} ,\textbf{λ}, \textbf{π}) $ は $ x $ の所属クラスタとポアソン分布のパラメータ $ λ $ が所与の状態での $ x $ の尤度であり、$ p(x | \textbf{s} ,\textbf{λ}) $ である。

でも、そもそも $ \textbf{s} $ は観測できないわけで。観測できないが $ x $ と共に存在する隠れたデータとみなす、みたいな解釈で一旦はよさそう。

尤度関数を具体的に求めてみる

一般に、$N$個のデータ $ x_n $ について尤度関数を求めたい。順を追って、まずは一つのデータに対する尤度関数を求め、その次に$N$個のデータに尤度関数を求めよう。

データが一つの場合

\begin{eqnarray*}
p(x, \textbf{s} | \textbf{π}, \textbf{λ})
&=& p(x|\textbf{s}, \textbf{λ}) p(\textbf{s}|\textbf{π}) \\
&=& \prod_{k=1}^{K} {Poi(x|λ_k)}^{s_k} ・ Cat(\textbf{s}|\textbf{π}) \\
&=& \prod_{k=1}^{K} {Poi(x|λ_k)}^{s_k} ・ \prod_{k=1}^{K} {\pi_k}^{s_k} \\
&=& \prod_{k=1}^{K} (\pi_k Poi(x|λ_k))^{s_k} \\
\end{eqnarray*}

一旦ここまで。

データがN個の場合

$\textbf{X} = (x_1, x_2, .. , x_N)$ 、 $\textbf{S} = (\textbf{s}_1, \textbf{s}_2, .. , \textbf{s}_N)$ とする。

以下の式変形では、パラメータ所与の下で条件付き独立性を仮定していることに注意。

\begin{eqnarray*}
p(\textbf{X}, \textbf{S} | \textbf{π}, \textbf{λ})
&=& \prod_{n=1}^{N} p(x_n, \textbf{s}_n | \textbf{π}, \textbf{λ}) \\
&=& \prod_{n=1}^{N} \prod_{k=1}^{K} (\pi_k Poi(x_n|λ_k))^{s_{n,k}} \\
\end{eqnarray*}

以降では、$N$個のデータを扱うことにしよう。

ポアソン混合モデルにおける学習・推論

まずはパラメータ s の推定から始めよう

ポアソン混合モデルで、データから各ポアソン分布のパラメータ $ λ_k $ を推定するには、一つのデータがどのクラスタに属するかを推定する必要がある。

データ $ x_n $ 所与の下で $ \textbf{s}_n $ を推定してみよう。

直感的にいけるところは直感的にいこう。$ x_n $ 所与での $ s_{n,k} $ の事後確率を $ γ(s_{n,k}) $ とおく。$ γ(s_{n,k}) $ は、$ x_n $ における尤度(確率密度)の比でいい。ベイズの定理で丁寧に示すこともできるので興味ある方は調べてみてね。

\begin{eqnarray*}
γ(s_{n,k})
&=& p(s_{n,k} = 1|x_n) \\
&=& \frac{π_k Poi(x_n|λ_k)}{\sum_{k=1}^{K} π_k Poi(x_n|λ_k)} \\
\end{eqnarray*}

ここでは、数式での表現をわかりやすくするために、 $ \textbf{s}_n $ の第 $ k $ 成分 $ \textbf{s}_{n,k} $ についての事後確率を求めていることに注意してほしい。

もう少しだけ $ \textbf{s} $ の性質を。

\begin{eqnarray*}
\mathbb{E}[s_{n,k}|x_n]
&=& \sum_{s_{n,k} = \{0,1\}}^{} s_{n,k} p(s_{n,k} | x_n) \\
&=&  0 ・ p(s_{n,k} = 0 | x_n) + 1 ・ p(s_{n,k} = 1 | x_n) \\
&=&  p(s_{n,k} = 1 | x_n) \\
&=& γ(s_{n,k}) \\
\end{eqnarray*}

ということで、 $ s_{n,k} $ の事後分布と期待値が一致することがわかる。

\mathbb{E}[s_{n,k}|x_n] = γ(s_{n,k})

なんで唐突にパラメータ sの推定したのさ

さきほど求めた以下の尤度関数でパラメータ $ \textbf{π} $ 、 $ \textbf{λ} $ を最尤推定したい気持ちがある。

ただ、指数に登場する $ s_{n,k} $ は観測できない。そこで、データ $ x_n $ 所与の下での $ s_{n,k} $ の期待値を $ s_{n,k} $ の代わりに採用して話を進めたかった。

p(\mathbf{X}, \textbf{S} | \textbf{π}, \textbf{λ})
= \prod_{n=1}^{N} \prod_{k=1}^{K} (\pi_k Poi(x_n|λ_k))^{s_{n,k}}

対数尤度関数として扱った方が楽そうなのでそうしてみる。

\begin{eqnarray*}
\ln p(\mathbf{X}, \textbf{S} | \textbf{π}, \textbf{λ})
&=& \ln \prod_{n=1}^{N} \prod_{k=1}^{K} (\pi_k Poi(x_n|λ_k))^{s_{n,k}} \\
&=& \sum_{n=1}^{N} \sum_{k=1}^{K} \ln (\pi_k Poi(x_n|λ_k))^{s_{n,k}} \\
&=& \sum_{n=1}^{N} \sum_{k=1}^{K} s_{n,k} \ln (\pi_k Poi(x_n|λ_k)) \\
&=& \sum_{n=1}^{N} \sum_{k=1}^{K} s_{n,k} \ln \pi_k + \sum_{n=1}^{N} \sum_{k=1}^{K} s_{n,k} \ln Poi(x_n|λ_k) \\
\end{eqnarray*}

上述したとおり、$ s_{n,k} $ は観測できないため、これを $ s_{n,k} $ の期待値に置き換える。

関数としては別物なので、$ Q $関数と名付けておこう。

\begin{eqnarray*}
Q(\textbf{π}, \textbf{λ})
&=& \sum_{n=1}^{N} \sum_{k=1}^{K} \mathbb{E}[s_{n,k}] \ln \pi_k + \sum_{n=1}^{N} \sum_{k=1}^{K} \mathbb{E}[s_{n,k}] \ln Poi(x_n|λ_k) \\
&=& \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k + \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln Poi(x_n|λ_k) \\
\end{eqnarray*}

$$
a
$$

ここで、$ \textbf{s}_n $ は $ s_{n,k} \in {0, 1} $ である一方、その期待値 $ \mathbb{E}[s_{n,k}] $ の各成分は$ s_{n,k} \in (0, 1) $ であり、例えば $ \mathbb{E}[\textbf{s}_1] = (0.2, 0.1, 0.7) $ のように $ 0 $ と $ 1 $ の間の実数値をとることに注意。

$Q$関数は最尤推定できそう。最尤推定して求まるパラメータをそれぞれ $ \textbf{λ}^* $、$ \textbf{π}^* $ と表すことにしよう。

λ に関する最尤推定を行う。

\begin{eqnarray*}
\textbf{λ}^*
&=& \underset{λ}{argmax} \quad Q(\textbf{π}, \textbf{λ}) \\
&=& \underset{λ}{argmax} \quad \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k + \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln Poi(x_n|λ_k) \\
&=& \underset{λ}{argmax} \quad \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln Poi(x_n|λ_k) \\
\end{eqnarray*}

これを具体的に書き下してみよう。

\begin{eqnarray*}
\sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln Poi(x_n|λ_k)
&=& \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k})・\ln (\frac{{λ_k}^{x_n}}{x_n!} e^{-λ_k})
\end{eqnarray*}

$ λ_k $ で偏微分して「$=0$」とおき、極値を求めてみよう。

\begin{eqnarray*}
\frac{\partial}{\partial λ_k} \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k})・\ln (\frac{{λ_k}^{x_n}}{x_n!} e^{-λ_k})
&=& \frac{\partial}{\partial λ_k} \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k})・\ln (\frac{{λ_k}^{x_n}}{x_n!} e^{-λ_k}) \\
&=& \frac{\partial}{\partial λ_k} \sum_{n=1}^{N} (γ(s_{n,k})・\ln (\frac{{λ_k}^{x_n}}{x_n!} e^{-λ_k}) + \sum_{j\neq k}^{K} γ(s_{n,j})・\ln (\frac{{λ_j}^{x_n}}{x_n!} e^{-λ_j})) \\
&=& \frac{\partial}{\partial λ_k} \sum_{n=1}^{N} γ(s_{n,k})・\ln (\frac{{λ_k}^{x_n}}{x_n!} e^{-λ_k}) \\
&=& \sum_{n=1}^{N} γ(s_{n,k})・\frac{\partial}{\partial λ_k} \ln (\frac{{λ_k}^{x_n}}{x_n!} e^{-λ_k}) \\
&=& \sum_{n=1}^{N} γ(s_{n,k})・\frac{\partial}{\partial λ_k} (x_n \ln λ_k - \ln x_n ! - λ_k) \\
&=& \sum_{n=1}^{N} γ(s_{n,k})・(\frac{x_n}{λ_k} - 1) \\
\end{eqnarray*}

最後の式 $ = 0 $ として $λ_k$ を求めよう。

\begin{eqnarray*}
\sum_{n=1}^{N} γ(s_{n,k})・(\frac{x_n}{λ_k} - 1) &=& 0 \\
\sum_{n=1}^{N} γ(s_{n,k})・(x_n - λ_k) &=& 0 \\
\sum_{n=1}^{N} γ(s_{n,k}) x_n &=& \sum_{n=1}^{N} γ(s_{n,k}) λ_k \\
λ_k &=& \frac{\sum_{n=1}^{N} γ(s_{n,k}) x_n}{\sum_{n=1}^{N} γ(s_{n,k})} \\
\end{eqnarray*}

これを $ k = 1, 2, .. , K $ について計算すれば、$ \textbf{λ}^* $ が求まる！やった！

πに関する最尤推定を行う。

\begin{eqnarray*}
\textbf{π}^*
&=& \underset{π}{argmax} \quad Q(\textbf{π}, \textbf{λ}) \\
&=& \underset{π}{argmax} \quad \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k + \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln Poi(x_n|λ_k) \\
&=& \underset{π}{argmax} \quad \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k \\
\end{eqnarray*}

$ π_k $ で偏微分して「$=0$」とおき、極値を求めてみよう。

とはいかない。なぜなら、$ \textbf{π} $ には $ \sum_{k=1}^{K} π_k = 1 $ という制約条件があるからだ。

ここで採用する最適化の手法はラグランジュの未定乗数法だ。

ラグランジュ関数を以下のように定義する。ラグランジュ乗数としては一般に $ λ $ が用いられるが、ここではポアソン分布のパラメータ $ λ $ との混同を避けるために、 $ l $ とした。

f = \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k + l(\sum_{k=1}^{K} π_k - 1)

この先は $ λ_k $ を求めた作業とほとんど同じ流れ。

$ f $ にはパラメータとして $ π_k $ と $ l $ が存在するため、各々について偏微分を行い、「$=0$」とおき、極値を求める。

まずは $ l $ について偏微分し、「$=0$」としよう。

\begin{eqnarray*}
\frac{\partial}{\partial l} f
&=& \frac{\partial}{\partial l} (\sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k + l(\sum_{k=1}^{K} π_k - 1)) \\
&=& \frac{\partial}{\partial l} l(\sum_{k=1}^{K} π_k - 1) \\
&=& \sum_{k=1}^{K} π_k - 1 = 0 \\
\sum_{k=1}^{K} π_k &=& 1 \\
\end{eqnarray*}

と、まぁご存知の式が得られる。

次は $ π_k $ について偏微分し、「$=0$」としよう。

\begin{eqnarray*}
\frac{\partial}{\partial π_k} f
&=& \frac{\partial}{\partial π_k} (\sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k + l(\sum_{k=1}^{K} π_k - 1)) \\
&=& \frac{\partial}{\partial π_k} \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k + \frac{\partial}{\partial π_k} l(\sum_{k=1}^{K} π_k - 1) \\
&=& \frac{\partial}{\partial π_k} \sum_{n=1}^{N} γ(s_{n,k}) \ln \pi_k + \frac{\partial}{\partial π_k} l(π_k -1) \\
&=& \sum_{n=1}^{N} γ(s_{n,k}) \frac{\partial}{\partial π_k} \ln \pi_k + \frac{\partial}{\partial π_k} l(π_k -1) \\
&=& \sum_{n=1}^{N} γ(s_{n,k}) \frac{1}{π_k} + l = 0 \\
\sum_{n=1}^{N} γ(s_{n,k}) &=& - π_k l \\
\end{eqnarray*}

$ \sum_{n=1}^{N} γ(s_{n,k}) $ ってどうやって計算するんだ？？

いっそのこと $ \sum_{k=1}^{K} $ で両辺を $ k $ について足し合わせるのがセオリー。

\begin{eqnarray*}
\sum_{k=1}^{K} \sum_{n=1}^{N} γ(s_{n,k}) &=& \sum_{k=1}^{K} - π_k l \\
\sum_{n=1}^{N} (\sum_{k=1}^{K} γ(s_{n,k})) &=& - l (\sum_{k=1}^{K} π_k) \\
\sum_{n=1}^{N} 1 &=& - l \\
l = - N
\end{eqnarray*}

さっき放置した式に代入してしまえ。

\begin{eqnarray*}
\sum_{n=1}^{N} γ(s_{n,k}) &=& - π_k l \\
\sum_{n=1}^{N} γ(s_{n,k}) &=& π_k N \\
π_k  &=& \frac{\sum_{n=1}^{N} γ(s_{n,k})}{N} \\
\end{eqnarray*}

これを $ k = 1, 2, .. , K $ について計算すれば、$ \textbf{λ}^* $ が求まる！

ここまでのまとめ

１．尤度関数の対数をとり、最尤推定したいが、 $ s_{n,k} $ は観測できない。そこで、代用として期待値 $ \mathbb{E}[s_{n,k}|x_n] $ が欲しい。

ポアソン混合モデルでは、$ \mathbb{E}[s_{n,k}|x_n] $ は $ s_{n,k} $ の事後分布 $ γ(s_{n,k})
= p(s_{n,k} = 1|x_n) $ と一致する。

\begin{eqnarray*}
\mathbb{E}[s_{n,k}|x_n]
&=& γ(s_{n,k}) \\
&=& p(s_{n,k} = 1|x_n) \\
&=& \frac{π_k Poi(x_n|λ_k)}{\sum_{k=1}^{K} π_k Poi(x_n|λ_k)} \\
\end{eqnarray*}

２．対数尤度関数 $ \ln p(\mathbf{X}, \textbf{S} | \textbf{π}, \textbf{λ}) $ の式中の $ s_{n,k} $ をその期待値 $ γ(s_{n,k}) $ で置き換えた $ Q $ 関数を作る。

p(\textbf{X}, \textbf{S} | \textbf{π}, \textbf{λ})
= \prod_{n=1}^{N} \prod_{k=1}^{K} (\pi_k Poi(x_n|λ_k))^{s_{n,k}}

\ln p(\mathbf{X}, \textbf{S} | \textbf{π}, \textbf{λ})
= \sum_{n=1}^{N} \sum_{k=1}^{K} s_{n,k} \ln \pi_k + \sum_{n=1}^{N} \sum_{k=1}^{K} s_{n,k} \ln Poi(x_n|λ_k)

Q(\textbf{π}, \textbf{λ})
= \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln \pi_k + \sum_{n=1}^{N} \sum_{k=1}^{K} γ(s_{n,k}) \ln Poi(x_n|λ_k)

３．$ Q $ 関数を $ λ_k $ について偏微分し、最尤推定を行う。

\begin{eqnarray*}
\frac{\partial}{\partial λ_k} Q(\textbf{π}, \textbf{λ}) &=& 0 \\
λ_k &=& \frac{\sum_{n=1}^{N} γ(s_{n,k}) x_n}{\sum_{n=1}^{N} γ(s_{n,k})} \\
\end{eqnarray*}

これを $ k = 1, 2, .. , K $ について行い、$ λ^* $ を得る。

４。$ Q $ 関数の極値を、 $ \sum_{k=1}^{K} π_k = 1 $ の条件下で、ラグランジュの未定乗数法で解く。

\begin{eqnarray*}
f &=& Q(\textbf{π}, \textbf{λ}) + l(\sum_{k=1}^{K} π_k - 1) \\
\frac{\partial}{\partial π_k} f &=& 0 \\
\frac{\partial}{\partial l} f &=& 0 \\
π_k  &=& \frac{\sum_{n=1}^{N} γ(s_{n,k})}{N} \\
\end{eqnarray*}

これを $ k = 1, 2, .. , K $ について行い、$ π^* $ を得る。

実際には、この工程は1回では終わらず、何度も繰り返すことで(局所)最適解を得る。

今回は、あまり明記しなかったものの、EM(Expectation Maximization)アルゴリズムによってポアソン混合モデルの推論を行った。

EM(Expectation Maximization：期待値最大化)アルゴリズムは、E-StepとM-Stepを繰り返すことによって局所最適解を得る。

E(Expectation)-Stepは上記の手順1-2に相当し、M(Maximization)-Stepは上記の手順3-4に相当する。

EMアルゴリズムを実装する際に注意が必要な点がいくつかあり、ひとつには局所最適化を行う点である。

そのため、パラメータの初期化は複数パターンで行い、どの推定量が適切か、確認する必要がある。

次回はとうとうポアソン混合モデルの実装！

お楽しみに！