EMアルゴリズム
変分ベイズ法の前に、ここではEMアルゴリズムについて解説します。
EMアルゴリズムは隠れ変数を持つモデルの最尤推定値を求めるアルゴリズムですが、隠れ変数については分布を推定していて、変分ベイズ法の特別な場合と解釈できます。
混合正規分布
まず隠れ変数を持つ代表的なモデルの具体例として、混合正規分布を紹介します。
混合正規分布は
\begin{align}
p(\boldsymbol{x}|\boldsymbol{\theta})
&\equiv \sum_{k=1}^K \pi_k \mathcal{N}(\boldsymbol{x};\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})
\tag{1} \\
\mathcal{N}(\boldsymbol{x};\boldsymbol{\mu},\boldsymbol{\Lambda}^{-1})
&= (2\pi)^{-\frac{D}{2}} |\boldsymbol{\Lambda}|^{\frac{1}{2}}
\exp\left(
-\frac{1}{2} (\boldsymbol{x}-\boldsymbol{\mu})^T \boldsymbol{\Lambda} (\boldsymbol{x}-\boldsymbol{\mu})
\right)
\end{align}
というモデルで、複数の正規分布を混合比$\pi_k$($\sum \pi_k = 1$)で重み付けて和を取った分布です。パラメータは$\boldsymbol{\theta} = \{\pi_k, \boldsymbol{\mu}_k, \boldsymbol{\Lambda}_k\}_{k=1}^K$です。$\boldsymbol{\Lambda}_k$は精度行列で、共分散行列$\boldsymbol{\Sigma}_k$の逆行列:$\boldsymbol{\Lambda}_k=\boldsymbol{\Sigma}_k^{-1}$です。以前1次元正規分布のベイズ推定の導出の際に、分散ではなく逆数の精度$\lambda=(\sigma^2)^{-1}$を用いたのと同じで、一方が求まれば他方が求まるのと、この後の計算の利便性からこちらを用います。
隠れ変数
混合正規分布で、もし「個々のサンプル$\boldsymbol{x}_n$がどの要素分布から生成されたのか」がわかっていれば、「各要素分布のパラメータを別々に求めるだけ」なのですが、実際にはサンプル$\boldsymbol{x}_n$がどの要素分布から生成されたかはわかりません(観測されません)。
そこで、サンプル$\boldsymbol{x}_n$がどの分布から生成されたのかを指す『観測されていない変数』、すなわち隠れ変数$\boldsymbol{z}_n$を導入します。$\boldsymbol{z}_n=(z_{n1},\cdots,z_{nK})$は$\boldsymbol{x}_n$を生成した要素分布を指すインデックス$k$について$z_{nk}=1$、それ以外のインデックス$l \neq k$について$z_{nl}=0$とします。もし$\boldsymbol{z}$がわかっている、すなわち$\boldsymbol{x}$を生成する要素分布が特定されているとき、$\boldsymbol{x}$を生成する分布は$\boldsymbol{z}$でスイッチングされるので
\begin{align}
p(\boldsymbol{x}|\boldsymbol{z},\boldsymbol{\theta}) &= \prod_{k=1}^K \mathcal{N}(\boldsymbol{x};\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})^{z_k}
\tag{2}
\end{align}
と書けます。そして、$z_k=1$となる確率は混合比$\pi_k$なので、$\boldsymbol{z}$の事前分布はカテゴリカル分布:
\begin{align}
p(\boldsymbol{z}|\boldsymbol{\theta}) &= {\rm Cat}(\boldsymbol{z};\boldsymbol{\pi}) = \prod_{k=1}^K \pi_k^{z_k}
\tag{3}
\end{align}
となります。$\boldsymbol{z}$について周辺化すると
\begin{align}
p(\boldsymbol{x}|\boldsymbol{\theta})
&= \sum_\boldsymbol{z} p(\boldsymbol{x}|\boldsymbol{z},\boldsymbol{\theta}) p(\boldsymbol{z}|\boldsymbol{\theta}) \\
&= \sum_{k=1}^K \pi_k \mathcal{N}(\boldsymbol{x};\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})
\end{align}
となり、(1)の混合正規分布を再構成できます。
生成モデルで書くと
\begin{align}
\boldsymbol{z} &\sim p(\boldsymbol{z}) = {\rm Cat}(\boldsymbol{z};\boldsymbol{\pi}) \\
\boldsymbol{x} &\sim p(\boldsymbol{x}|\boldsymbol{z},\boldsymbol{\theta}) = \prod_{k=1}^K \mathcal{N}(\boldsymbol{x};\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})^{z_k}
\end{align}
となり、「まず確率変数$\boldsymbol{z}$が生成されて」「$\boldsymbol{z}$の条件付きで$\boldsymbol{x}$が生成される」という連鎖で表現されます。
天下りで$\boldsymbol{z}$を導入してきましたが、$\boldsymbol{z}$を導入したことでどのように$\boldsymbol{\theta}$を推定するのでしょうか。
隠れ変数を用いない場合
隠れ変数を用いない場合、混合正規分布の対数尤度関数は
\begin{align}
\ln p(\boldsymbol{X}|\boldsymbol{\theta})
&= \ln \prod_{n=1}^N \sum_{k=1}^K \pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1}) \\
&= \sum_{n=1}^N \ln \sum_{k=1}^K \pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})
\end{align}
となります。$\boldsymbol{X}=\{\boldsymbol{x}_n\}_{n=1}^N$です。偏微分を求めていくのですが、
\begin{align}
(\ln f)^\prime = \frac{1}{f} f^\prime
\end{align}
より$f^\prime = f (\ln f)^\prime$を使って、正規分布の微分について
\begin{align}
\mathcal{N}^\prime
&= \mathcal{N} (\ln \mathcal{N})^\prime
\end{align}
を使っていきます。
\begin{align}
&\frac{\partial}{\partial \pi_k} \ln p(\boldsymbol{X}|\boldsymbol{\theta}) \\
&= \sum_{n=1}^N
\frac{1}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\frac{\partial}{\partial \pi_k} \left\{
\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})
\right\} \\
&= \sum_{n=1}^N
\frac{\mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})} \\
&\frac{\partial}{\partial \boldsymbol{\mu}_k} \ln p(\boldsymbol{X}|\boldsymbol{\theta}) \\
&= \sum_{n=1}^N
\frac{1}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\frac{\partial}{\partial \boldsymbol{\mu}_k} \left\{
\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})
\right\} \\
&= \sum_{n=1}^N
\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\frac{\partial}{\partial \boldsymbol{\mu}_k} \ln \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1}) \\
&= \sum_{n=1}^N
\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\frac{\partial}{\partial \boldsymbol{\mu}_k} \ln \left\{
(2\pi_k)^{-\frac{D}{2}} |\boldsymbol{\Lambda}_k|^{\frac{1}{2}}
\exp\left(
-\frac{1}{2} (\boldsymbol{x}-\boldsymbol{\mu}_k)^T \boldsymbol{\Lambda}_k (\boldsymbol{x}-\boldsymbol{\mu}_k)
\right)
\right\} \\
&= \sum_{n=1}^N
\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\frac{\partial}{\partial \boldsymbol{\mu}_k} \left(
-\frac{1}{2} (\boldsymbol{x}-\boldsymbol{\mu}_k)^T \boldsymbol{\Lambda}_k (\boldsymbol{x}-\boldsymbol{\mu}_k)
\right) \\
&= \sum_{n=1}^N \frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})} \boldsymbol{\Lambda}_k (\boldsymbol{x}_n - \boldsymbol{\mu}_k) \\
&\frac{\partial}{\partial \boldsymbol{\Lambda}_k} \ln p(\boldsymbol{X}|\boldsymbol{\theta}) \\
&= \sum_{n=1}^N
\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\frac{\partial}{\partial \boldsymbol{\Lambda}_k} \ln \left\{
(2\pi_k)^{-\frac{D}{2}} |\boldsymbol{\Lambda}_k|^{\frac{1}{2}}
\exp\left(
-\frac{1}{2} (\boldsymbol{x}-\boldsymbol{\mu}_k)^T \boldsymbol{\Lambda}_k (\boldsymbol{x}-\boldsymbol{\mu}_k)
\right)
\right\} \\
&= \sum_{n=1}^N
\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\frac{\partial}{\partial \boldsymbol{\Lambda}_k} \left\{
\frac{1}{2} \ln |\boldsymbol{\Lambda}_k|
-\frac{1}{2} (\boldsymbol{x}-\boldsymbol{\mu}_k)^T \boldsymbol{\Lambda}_k (\boldsymbol{x}-\boldsymbol{\mu}_k)
\right\} \\
&= \sum_{n=1}^N
\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\frac{\partial}{\partial \boldsymbol{\Lambda}_k} \left\{
\frac{1}{2} \ln |\boldsymbol{\Lambda}_k|
-\frac{1}{2} {\rm tr}\Big( \boldsymbol{\Lambda}_k (\boldsymbol{x}-\boldsymbol{\mu}_k)(\boldsymbol{x}-\boldsymbol{\mu}_k)^T \boldsymbol{B}ig)
\right\} \\
&= \sum_{n=1}^N
\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\left\{
\frac{1}{2} \boldsymbol{\Lambda}_k^{-1}
-\frac{1}{2} (\boldsymbol{x}_n - \boldsymbol{\mu}_k) (\boldsymbol{x}_n - \boldsymbol{\mu}_k)^T
\right\}
\end{align}
です。$\boldsymbol{\Lambda}$による微分では
\begin{align}
\frac{\partial}{\partial \boldsymbol{A}} \ln |\boldsymbol{A}| &= \boldsymbol{A}^{-1} \\
\frac{\partial}{\partial \boldsymbol{A}} {\rm tr}(\boldsymbol{A} \boldsymbol{B}) &= \boldsymbol{B}^T
\end{align}
を用いました。勾配が求まったので、勾配法($\pi_k$について$\sum_k \pi_k = 1$の制約があるので、更新のたびに$\sum_k \pi_k = 1$となるように正規化する射影勾配法)で学習することはできます。
ですが、ちょっと大変そうですよね。。
EMアルゴリズム
EMアルゴリズムは最尤推定なので尤度$p(\boldsymbol{X}|\theta)$を最大化するのですが、
各サンプル$\boldsymbol{x}_n$を生成した要素分布を指す隠れ変数$\boldsymbol{z}_n$を導入して
\begin{align}
p(\boldsymbol{X}|\theta) = \int p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta}) d\boldsymbol{Z}
\end{align}
を考えます(混合正規分布のように離散値の場合は和なのですが、連続値か離散値か気にせず積分で表記しています)。$\boldsymbol{Z}=\{\boldsymbol{z}_n\}_{n=1}^N$です。
対数尤度は
\begin{align}
\ln p(\boldsymbol{X}|\boldsymbol{\theta})
&= \ln \int p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta}) d\boldsymbol{Z} \\
&= \ln \int q(\boldsymbol{Z}) \frac{p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})}{q(\boldsymbol{Z})} d\boldsymbol{Z} \tag{4}\\
&\geq \int q(\boldsymbol{Z}) \ln \frac{p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})}{q(\boldsymbol{Z})} d\boldsymbol{Z} \tag{5}\\
&= - {\rm KL}(q(\boldsymbol{Z}) \parallel p(\boldsymbol{Z}|\boldsymbol{X},\boldsymbol{\theta})) \tag{6} \\
&= \mathcal{F}(q(\boldsymbol{Z}),\boldsymbol{\theta})
\end{align}
と展開できます。(4)で隠れ変数$\boldsymbol{Z}$の分布$q(\boldsymbol{Z})$を導入しています。(5)はJensenの不等式で、上に凸な関数$f$の「$x$の平均値での関数値$f(\mathbb{E}[x])$」は「関数値$f(x)$の平均値$\mathbb{E}[f(x)]$」以上になる、というものです。EMアルゴリズムでは(5)の右辺である「対数尤度の下界」$\mathcal{F}(q(\boldsymbol{Z}),\boldsymbol{\theta})$を最大化します。
最適化する変数は$\boldsymbol{\theta}$と$q(\boldsymbol{Z})$です。
(6)の${\rm KL}()$は、分布がどれだけ異なっているかの尺度であるKLダイバージェンス:
\begin{align}
{\rm KL}(q \parallel p) = \int q(\boldsymbol{x}) \frac{q(\boldsymbol{x})}{p(\boldsymbol{x})} d\boldsymbol{x}
\end{align}
で、(5)の分子$p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})=p(\boldsymbol{Z}|\boldsymbol{X},\boldsymbol{\theta}) p(\boldsymbol{X}|\boldsymbol{\theta})$を$\boldsymbol{Z}$について規格化して$p(\boldsymbol{Z}|\boldsymbol{X},\boldsymbol{\theta})$が出てきています。これより$\boldsymbol{\theta}$の下での$q(\boldsymbol{Z})$の最適解は$p(\boldsymbol{Z}|\boldsymbol{X},\boldsymbol{\theta})$ということがわかります。
\begin{align}
{\rm argmax}_{q(\boldsymbol{Z})} \mathcal{F}(q(\boldsymbol{Z}),\boldsymbol{\theta}) = p(\boldsymbol{Z}|\boldsymbol{X},\boldsymbol{\theta}) = \prod_{n=1}^N p(\boldsymbol{z}_n|\boldsymbol{x}_n,\boldsymbol{\theta})
\end{align}
で、モデルのパラメータ$\boldsymbol{\theta}$の下でサンプル$\boldsymbol{x}_n$を見たときの$\boldsymbol{z}_n$の事後分布$p(\boldsymbol{z}_n|\boldsymbol{x}_n,\boldsymbol{\theta})$の積です。
EMアルゴリズムは、現在の推定値$\hat{\boldsymbol{\theta}}$ので$\hat{q}(\boldsymbol{Z})=p(\boldsymbol{Z}|\boldsymbol{X},\hat{\boldsymbol{\theta}})$を求め、$p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})$の$\hat{q}(\boldsymbol{Z})$による期待値((5)の右辺)を求めるExpectation Stepと、これを$\boldsymbol{\theta}$について最大化するMaximization Stepを反復することで最尤推定値$\hat{\boldsymbol{\theta}}$を求めます。(ExpectationとMaximizationの言葉通りに取るとこのようになると思うのですが、アルゴリズムの記載上は、単に$\hat{q}(\boldsymbol{Z})$を求めることを(まだExpectationしていないけど)Expectation Step、$\hat{q}(\boldsymbol{Z})$による期待値を$\boldsymbol{\theta}$について最大化するステップをMaximization Stepと書かれていることもあるように思います。)
また、(5)の右辺分母は$\boldsymbol{\theta}$に無関係なので、Maximization Stepは$\mathbb{E}_{\hat{q}(\boldsymbol{Z})}[\ln p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})]$を最大化します。
\begin{align*}
\begin{array}{ll}
{\rm {\bf Expectation Step:}} &
\hat{q}(\boldsymbol{Z}) \leftarrow {\rm argmax}_{q(\boldsymbol{Z})} \mathcal{F}(q(\boldsymbol{Z}),\hat{\boldsymbol{\theta}}) = p(\boldsymbol{Z}|\boldsymbol{X},\hat{\boldsymbol{\theta}}) \\
{\rm {\bf Maximization Step:}} &
\hat{\boldsymbol{\theta}} \leftarrow {\rm argmax}_{\boldsymbol{\theta}} \mathcal{F}(\hat{q}(\boldsymbol{Z}),\boldsymbol{\theta}) = {\rm argmax} \mathbb{E}_{\hat{q}(\boldsymbol{Z})}[\ln p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})]
\end{array}
\end{align*}
混合正規分布のEMアルゴリズム
Expectation Step
モデルのパラメータの現在の推定値$\hat{\boldsymbol{\theta}}$の下での$q(\boldsymbol{z}_n)$の最適解は
\begin{align}
\hat{q}(\boldsymbol{z}_n) &= p(\boldsymbol{z}_n|\boldsymbol{x}_n,\hat{\boldsymbol{\theta}}) = \frac{p(\boldsymbol{x}_n|\boldsymbol{z}_n,\hat{\boldsymbol{\theta}}) p(\boldsymbol{z}_n|\hat{\boldsymbol{\theta}})}{p(\boldsymbol{x}_n|\hat{\boldsymbol{\theta}})}
\end{align}
です。定数の分母の正規化項を無視して、分子に(2)と(3)を代入すると
\begin{align}
\hat{q}(\boldsymbol{z}_n)
&\propto p(\boldsymbol{x}_n|\boldsymbol{z}_n,\hat{\boldsymbol{\theta}}) p(\boldsymbol{z}_n|\hat{\boldsymbol{\theta}}) \\
&= \left( \prod_{k=1}^K \mathcal{N}(\boldsymbol{x}_n;\hat{\boldsymbol{\mu}}_k,\hat{\boldsymbol{\Lambda}}_k^{-1})^{z_{nk}} \right)
\left( \prod_{k=1}^K \hat{\pi}_k^{z_{nk}} \right) \\
&= \prod_{k=1}^K \left(\hat{\pi}_k \mathcal{N}(\boldsymbol{x}_n;\hat{\boldsymbol{\mu}}_k,\hat{\boldsymbol{\Lambda}}_k^{-1}) \right)^{z_{nk}}
\end{align}
$z_{nk}$についての関数形から、$z_{nk}$が$1$になる確率が$\hat{\pi}_k \mathcal{N}(\boldsymbol{x}_n;\hat{\boldsymbol{\mu}}_k,\hat{\boldsymbol{\Lambda}}_k^{-1})$に比例するカテゴリカル分布になっていることがわかります。これを正規化して
\begin{align}
\hat{q}(\boldsymbol{z}_n) &= {\rm Cat}(\boldsymbol{z}_n;\hat{\boldsymbol{r}}_n) \\
\hat{r}_{nk} &= \frac{\hat{\pi}_k \mathcal{N}(\boldsymbol{x}_n;\hat{\boldsymbol{\mu}}_k,\hat{\boldsymbol{\Lambda}}_k^{-1})}{\sum_{l=1}^K \hat{\pi}_l \mathcal{N}(\boldsymbol{x}_n;\hat{\boldsymbol{\mu}}_l,\hat{\boldsymbol{\Lambda}}_l^{-1})}
\end{align}
を得ます。$\hat{r}_{nk}$はカテゴリカル分布のパラメータで、言葉で言えば「サンプル$\boldsymbol{x}_n$が要素分布$k$から生成された確率」になっています。
Maximization Step
Maximization Stepでは$\mathbb{E}_{\hat{q}(\boldsymbol{Z})}[\ln p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})]$を最大化します。
\begin{align}
&\mathbb{E}_{\hat{q}(\boldsymbol{Z})}[\ln p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})] \\
&= \mathbb{E}_{\hat{q}(\boldsymbol{Z})}\left[\ln \prod_{n=1}^N p(\boldsymbol{x}_n,\boldsymbol{z}_n,\boldsymbol{\theta}) \right] \\
&= \mathbb{E}_{\hat{q}(\boldsymbol{Z})}\left[
\sum_{n=1}^N \ln p(\boldsymbol{x}_n|\boldsymbol{z}_n,\boldsymbol{\theta}) p(\boldsymbol{z}_n|\boldsymbol{\theta})
\right] \\
&= \mathbb{E}_{\hat{q}(\boldsymbol{Z})}\left[
\sum_{n=1}^N \ln \left( \prod_{k=1}^K \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})^{z_{nk}} \prod_{k=1}^K \pi_k^{z_{nk}} \right)
\right] \\
&= \mathbb{E}_{\hat{q}(\boldsymbol{Z})}\left[
\sum_{n=1}^N \sum_{k=1}^K z_{nk} \big( \ln \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1}) + \ln \pi_k \big)
\right] \\
&= \sum_{n=1}^N \sum_{k=1}^K \mathbb{E}_{\hat{q}(z_{nk})}[z_{nk}] \big( \ln \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1}) + \ln \pi_k \big) \\
&= \sum_{k=1}^K \left\{ \sum_{n=1}^N \hat{r}_{nk} \big( \ln \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1}) + \ln \pi_k \big) \right\}
\end{align}
最後の式で最初の和$\sum_{k=1}^K$は、その後の中括弧の「要素分布ごとに計算される($\hat{r}_{nk}$による重み付き)対数尤度」についての和になっています。なので、要素分布$k$のパラメータ$\{\pi_k, \boldsymbol{\mu}_k, \boldsymbol{\Lambda}_k\}$で偏微分をすると、他の要素分布$l \neq k$のパラメータに関する項は消えてしまいます。
さきほど**「個々のサンプル$\boldsymbol{x}_n$がどの要素分布から生成されたのか」がわかっていれば、「各要素分布のパラメータを別々に求めるだけ」と言っていたのですが、その「個々のサンプル$\boldsymbol{x}_n$がどの要素分布から生成されたのか」を指す隠れ変数$\boldsymbol{z}_n$を導入したことで、「隠れ変数$\boldsymbol{z}_n$の期待値$\hat{\boldsymbol{r}}_n$による重み付きで各要素分布のパラメータを別々に求めれば良い」**ことになっています。
$\boldsymbol{\mu}_k$についての偏微分は
\begin{align}
\frac{\partial}{\partial \boldsymbol{\mu}_k} \mathbb{E}_{\hat{q}(\boldsymbol{Z})}[\ln p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})]
&= \frac{\partial}{\partial \boldsymbol{\mu}_k} \sum_{k=1}^K
\left\{ \sum_{n=1}^N \hat{r}_{nk} \big( \ln \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1}) + \ln \pi_k \big) \right\} \\
&= \frac{\partial}{\partial \boldsymbol{\mu}_k} \sum_{n=1}^N \hat{r}_{nk} \ln \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1}) \\
&= \frac{\partial}{\partial \boldsymbol{\mu}_k} \sum_{n=1}^N \hat{r}_{nk} \ln \left\{
\frac{1}{\sqrt{2\pi}} |\boldsymbol{\Lambda}_k|^{\frac{1}{2}}
\exp\left( -\frac{1}{2} (\boldsymbol{x}_n-\boldsymbol{\mu}_k)^T \boldsymbol{\Lambda}_k (\boldsymbol{x}_n-\boldsymbol{\mu}_k) \right)
\right\} \\
&= \sum_{n=1}^N \hat{r}_{nk} \boldsymbol{\Lambda}_k (\boldsymbol{x}_n-\boldsymbol{\mu}_k)
\end{align}
より、$\boldsymbol{\mu}_k$の$\hat{q}(\boldsymbol{Z})$の下での最適解は
\begin{align}
\hat{\boldsymbol{\mu}}_k
&= {\rm argmax}_{\boldsymbol{\mu}_k} \mathcal{F}(\hat{q}(\boldsymbol{Z}),\boldsymbol{\theta})
= \frac{1}{N_k} \sum_{n=1}^N \hat{r}_{nk} \boldsymbol{x}_n \\
N_k
&\equiv \sum_{n=1}^N \hat{r}_{nk}
\end{align}
と、$\hat{r}_{nk}$の重み付きサンプル平均が求まります。同様に$\boldsymbol{\Lambda}_k$についての偏微分は
\begin{align}
\frac{\partial}{\partial \boldsymbol{\Lambda}_k} \mathbb{E}_{\hat{q}(\boldsymbol{Z})}[\ln p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})]
&= \frac{\partial}{\partial \boldsymbol{\Lambda}_k} \sum_{n=1}^N \hat{r}_{nk} \ln \left\{
\frac{1}{\sqrt{2\pi}} |\boldsymbol{\Lambda}_k|^{\frac{1}{2}}
\exp\left( -\frac{1}{2} (\boldsymbol{x}_n-\boldsymbol{\mu}_k)^T \boldsymbol{\Lambda}_k (\boldsymbol{x}_n-\boldsymbol{\mu}_k) \right)
\right\} \\
&= \frac{\partial}{\partial \boldsymbol{\Lambda}_k} \sum_{n=1}^N \hat{r}_{nk} \ln \left\{
\frac{1}{2} \ln |\boldsymbol{\Lambda}_k|
-\frac{1}{2} (\boldsymbol{x}_n-\boldsymbol{\mu}_k)^T \boldsymbol{\Lambda}_k (\boldsymbol{x}_n-\boldsymbol{\mu}_k)
\right\} \\
&= \sum_{n=1}^N \hat{r}_{nk}
\left(\frac{1}{2} \boldsymbol{\Lambda}_k^{-1} -\frac{1}{2} (\boldsymbol{x}_n-\boldsymbol{\mu}_k) (\boldsymbol{x}_n-\boldsymbol{\mu}_k)^T \right)
\end{align}
より、共分散行列$\boldsymbol{\Sigma}_k=\boldsymbol{\Lambda}_k^{-1}$の$\hat{q}(\boldsymbol{Z})$の下での最適解は
\begin{align}
\hat{\boldsymbol{\Sigma}}_k
&= \hat{\boldsymbol{\Lambda}}_k^{-1}
= {\rm argmax}_{\boldsymbol{\Lambda}_k^{-1}} \mathcal{F}(\hat{q}(\boldsymbol{Z}),\boldsymbol{\theta})
= \frac{1}{N_k} \sum_{n=1}^N \hat{r}_{nk} (\boldsymbol{x}_n-\boldsymbol{\mu}_k)(\boldsymbol{x}_n-\boldsymbol{\mu}_k)^T
\end{align}
と、$\hat{r}_{nk}$の重み付きサンプル共分散行列が求まります。
$\pi_k$については$\sum_{k=1}^K \pi_k = 1$の条件があるので、ラグランジュ乗数$\lambda$を導入して
\begin{align}
L(\boldsymbol{\pi},\lambda) = \mathbb{E}_{\hat{q}(\boldsymbol{Z})}[\ln p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})] + \lambda \left(\sum_{k=1}^K \pi_k - 1\right)
\end{align}
について
\begin{align}
\hat{\boldsymbol{\pi}} &= {\rm argmax}_{\boldsymbol{\pi}} \min_\lambda L(\boldsymbol{\pi},\lambda)
\end{align}
を解きます。$\min_\lambda L(\boldsymbol{\pi},\lambda)$は$\lambda$について最小化済みの関数を意味しているので引数は$\boldsymbol{\pi}$のみになります。$\sum_{k=1}^K \pi_k \neq 1$だと$\lambda$を調整することで$L(\boldsymbol{\pi},\lambda)$はいくらでも小さくできるので$\min_\lambda L(\boldsymbol{\pi},\lambda) = -\infty$になってしまいます。これを$\boldsymbol{\pi}$について最大化するので、制約$\sum_{k=1}^K \pi_k = 1$を満たす中で$\mathbb{E}_{\hat{q}(\boldsymbol{Z})}[\ln p(\boldsymbol{X},\boldsymbol{Z}|\boldsymbol{\theta})]$を最大化する$\boldsymbol{\pi}$が解になります。
偏微分は
\begin{align}
\frac{\partial}{\partial \pi_k} L(\boldsymbol{\pi},\lambda)
&= \frac{\partial}{\partial \pi_k} \left[
\sum_{k=1}^K \left\{
\sum_{n=1}^N \hat{r}_{nk} \big( \ln \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1}) + \ln \pi_k \big)
\right\}
+ \lambda \left( \sum_{k=1}^K \pi_k - 1 \right)
\right] \\
&= \sum_{n=1}^N \hat{r}_{nk} \frac{1}{\pi_k} + \lambda \tag{8} \\
\frac{\partial}{\partial \lambda} L(\boldsymbol{\pi},\lambda)
&= \sum_{k=1}^K \pi_k - 1 \tag{9}
\end{align}
であり、(8)$=0$の解:
\begin{align}
\hat{\pi}_k &= - \frac{\sum_{n=1}^N \hat{r}_{nk}}{\lambda} = - \frac{N_k}{\lambda}
\end{align}
を(9)に代入して$=0$を解いて
\begin{align}
- \sum_{k=1}^K \frac{N_k}{\lambda} - 1 = 0
\Leftrightarrow \hat{\lambda} = -\sum_{k=1}^K N_k
\end{align}
を得ます。ここで$\hat{r}_{nk}$はカテゴリカル分布のパラメータだったことを思い出すと、$\sum_{k=1}^K \hat{r}_{nk} = 1$なので$\sum_{k=1}^K N_k = \sum_{k=1}^K \sum_{n=1}^N \hat{r}_{nk} = \sum_{n=1}^N \sum_{k=1}^K \hat{r}_{nk} = N$より
\begin{align}
\hat{\lambda} = -N
\end{align}
であり、
\begin{align}
\hat{\pi}_k &= \frac{N_k}{N}
\end{align}
を得ます。
以上をまとめると混合正規分布のEMアルゴリズムは
\begin{align*}
\begin{array}{ll}
{\rm {\bf Expectation Step:}} &
\hat{r}_{nk} = \frac{\hat{\pi}_k \mathcal{N}(\boldsymbol{x}_n;\hat{\boldsymbol{\mu}}_k,\hat{\boldsymbol{\Lambda}}_k^{-1})}{\sum_{l=1}^K \hat{\pi}_l \mathcal{N}(\boldsymbol{x}_n;\hat{\boldsymbol{\mu}}_l,\hat{\boldsymbol{\Lambda}}_l^{-1})} \\
{\rm {\bf Maximization Step:}} &
\hat{\boldsymbol{\mu}}_k
= \frac{1}{N_k} \sum_{n=1}^N \hat{r}_{nk} \boldsymbol{x}_n \\
&
\hat{\boldsymbol{\Sigma}}_k
= \frac{1}{N_k} \sum_{n=1}^N \hat{r}_{nk} (\boldsymbol{x}_n-\boldsymbol{\mu}_k)(\boldsymbol{x}_n-\boldsymbol{\mu}_k)^T \\
&
\hat{\pi}_k
= \frac{N_k}{N}
\end{array}
\end{align*}
となります。
隠れ変数を用いない場合との比較
隠れ変数を用いない場合、勾配は次のようになるのでした。
\begin{align}
\frac{\partial}{\partial \boldsymbol{\mu}_k} \ln p(\boldsymbol{X}|\boldsymbol{\theta})
&= \sum_{n=1}^N \frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})} \boldsymbol{\Lambda}_k (\boldsymbol{x}_n - \boldsymbol{\mu}_k)
\tag{10}\\
\frac{\partial}{\partial \boldsymbol{\Lambda}_k} \ln p(\boldsymbol{X}|\boldsymbol{\theta})
&= \sum_{n=1}^N
\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
\left\{
\frac{1}{2} \boldsymbol{\Lambda}_k^{-1}
-\frac{1}{2} (\boldsymbol{x}_n - \boldsymbol{\mu}_k) (\boldsymbol{x}_n - \boldsymbol{\mu}_k)^T
\right\}
\tag{11}
\end{align}
(10)と(11)において、重み$\frac{\pi_k \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}$を$\hat{r}_{nk}$と置いて$\boldsymbol{\mu}_k$、$\boldsymbol{\Sigma}_k=\boldsymbol{\Lambda}_k^{-1}$について解くとEMアルゴリズムのMaximization Stepが得られます。また、$\pi_k$について$\sum_{k=1}^K \pi_k = 1$の制約を置いてラグランジュの未定乗数法を用いると
\begin{align}
\frac{\partial}{\partial \pi_k} \left\{
\ln p(\boldsymbol{X}|\boldsymbol{\theta}) + \lambda \left( \sum_{l=1}^K \pi_l - 1 \right)
\right\}
&= \sum_{n=1}^N
\frac{\mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_k,\boldsymbol{\Lambda}_k^{-1})}
{\sum_{l=1}^K \pi_l \mathcal{N}(\boldsymbol{x}_n;\boldsymbol{\mu}_l,\boldsymbol{\Lambda}_l^{-1})}
+ \lambda \\
&= \sum_{n=1}^N \frac{\hat{r}_{nk}}{\pi_k} + \lambda \\
\frac{\partial}{\partial \lambda} \left\{
\ln p(\boldsymbol{X}|\boldsymbol{\theta}) + \lambda \left( \sum_{l=1}^K \pi_l - 1 \right)
\right\}
&= \sum_{k=1}^K \pi_k - 1
\end{align}
で、これらを$0$と置いて解くとやはりEMアルゴリズムのMaximization Stepが得られます。