甲斐性なしのブログ

うどんくらいしか食べる気がしない

EMアルゴリズムと変分ベイズEMアルゴリズムのメモ

表題の通り。この辺りややこしくていつもわからなくなるので、身につけるためにも自分なりの理解をまとめたメモを書く。

対数尤度と下限とKLダイバージェンス

まず、既知の変数をX、未知の変数(ここでは潜在変数か未知パラメータかは区別していない)をZとすると以下が成り立つ。

\displaystyle \ln p(X) = KL\left[q(Z)\parallel p(Z| X) \right] + {\mathcal L}\left[q(Z)\right]  \tag{1}

ただし、KL\left[q(Z) \parallel p(Z | X) \right] はKLダイバージェンスで、

\displaystyle KL\left[q(Z) \parallel  p(Z | X) \right] =\int q(Z) \ln \frac{q(Z)}{p(Z | X)} dZ = - \int q(Z) \ln \frac{p(Z | X)}{q(Z)} dZ \tag{2}

{\mathcal L}\left[q(Z)\right] は、対数尤度の下限で

\displaystyle {\mathcal L}\left[q(Z)\right] = \int q(Z) \ln \frac{p(X, Z)}{q(Z)} dZ \tag{3}

であり、常に

\displaystyle \ln p(X) \geq {\mathcal L}\left[q(Z)\right] \tag{4}

が成り立つ。これは関数\lnが上に凸でありイェンセン不等式より\int p(x)\ln f(x)dx \leq \ln \int p(x)f(x)dxであるため、\int q(Z) \ln \frac{p(X, Z)}{q(Z)} dZ \leq \ln \int q(Z)\frac{p(X, Z)}{q(Z)} dZ =\ln \int p(X, Z) dZ =  \ln p(X) として得られる。

また、式(1)は、式(2)と式(3)を実際に足してみると得られる。

EMアルゴリズム

考え方

  • 得られているものは既知の変数のX、求めたいのは未知パラメータ\Theta
  • 目的関数は対数尤度\ln p(X|\Theta)であり、これを最大化する\Thetaを求めたい。
    • ところが、これを直接的に最適化するのは一般に困難。
  • ここで、潜在変数Zを導入し、{\mathcal L}\left[q(Z), \Theta \right] = \int q(Z) \ln \frac{p(X, Z|\Theta)}{q(Z)} dZという量を考える。
    • 式(4)より、\ln p(X|\Theta) \geq {\mathcal L}\left[q(Z), \Theta \right]が成り立つ。
  • そのため、{\mathcal L}\left[q(Z), \Theta \right]を最大化する関数qとパラメータ\Thetaを求めることを考える。
    • ただし、この両者を一気に最適化する解析解は得られないので、座標降下法の考えで一方を固定し他方を最適化するという操作を繰り返して求める。

EMアルゴリズムの更新式

ということで、更新アルゴリズムは以下の通り。

  1. まずは、パラメータを既知の値 {\hat \Theta}と固定した場合に、{\mathcal L} \left[q(Z), {\hat \Theta} \right]を最大にする関数qを求める。
    • 式(1)より、 \displaystyle {\mathcal L}\left[q(Z), {\hat \Theta} \right]  = \ln p(X| {\hat \Theta}) - KL\left[q(Z) | p(Z| X, {\hat \Theta}) \right]
    • KL\left[q(Z) \parallel p(Z| X, {\hat \Theta})\right] \geq 0 なので、結局q(Z) = p(Z| X, {\hat \Theta})とするのが最適。
    • 色々分け方があるようだが、これをEステップと呼ぶことが多いらしい。
  2. 次に、関数qを上で求めたq(Z) = p(Z| X, {\hat \Theta})に固定し、{\mathcal L} \left[p(Z| X, {\hat \Theta}),  \Theta \right]を最大化する\Thetaを求める。
    • {\mathcal L} \left[p(Z| X, {\hat \Theta}),  \Theta \right] =  \int p(Z| X, {\hat \Theta})\ln \frac{p(X, Z|\Theta)}{p(Z| X, {\hat \Theta})} dZ =  \int p(Z| X, {\hat \Theta})\ln p(X, Z|\Theta) dZ + {\rm const}なので、 \int p(Z| X, {\hat \Theta})\ln p(X, Z|\Theta) dZ = \langle \ln p(X, Z|\Theta)\rangle _{p(Z| X, {\hat \Theta})}を最大化する\Thetaを求めればよい。
    • 期待値の最適化なので\ln p(X, Z|\Theta)さえ最適化できるなら適用可能。
    • こちらをMステップと呼ぶ。

変分ベイズEMアルゴリズム

変分近似の考え方

  • 得られているものは既知の変数のX、未知変数は潜在変数もパラメータもひっくるめてZとする。
  • モデルp(X, Z)が与えられていて、未知変数の事後分布p(Z | X)を求めたい。
    • これは、p(Z | X) = \frac{p(X, Z)}{p(X)}であり、厳密に求めるのは困難(p(X)が難しい)。
  • そこで、p(Z | X)を別の分布q(Z)で近似することを考える。
    • 当然q(Z)p(Z | X)に近い方がいいので、式(2)のKLダイバージェンスを最小化する関数qを求めることになる。
  • その際、q(Z)に何の制約も課さないと、q(Z) = p(Z | X) という無意味な解が得られるので、q(Z)に何らの制約を課して表現能力を限定する。
    • 典型的なのは、q(Z) = \prod_{i=1}^{D} q(z_i | X)のように未知変数間の独立性を制約にする近似(平均場近似)。
  • 式(1)より、KLダイバージェンスを最小化は{\mathcal L}\left[q(Z)\right] を最大化することと等価。
    • \ln p(X)は定数なので。
  • ということで、定めた制約のもと{\mathcal L}\left[q(Z)\right] を最大化するqを求めればよい。

ここまでは変分近似の考え方。おそらく、変分ベイズEMアルゴリズムと言うと、以下のようなもう少し限定されたアルゴリズムのことを言うと思う。

変分EMアルゴリズムの考え方

  • 今、未知変数がパラメータ\Thetaと潜在変数Zだとする。
  • これらの事後分布p(\Theta, Z | X)を求めたいが厳密に求めるのは困難なので、q(\Theta, Z)という分布で近似することを考える。
    • qの制約としては、q(\Theta, Z) = q(\Theta)q(Z)という独立性を制約とする。
    • モデルに応じてどういう独立性を仮定するのが適切なのかが変わってくる(\Thetaをいくつかの要素が独立だと仮定したり)。
  • 式(3)より、\displaystyle {\mathcal L}\left[q(\Theta)q(Z)\right] = \int \int q(\Theta)q(Z) \ln \frac{p(X|\Theta, Z)p(Z | \Theta)p(\Theta)}{q(\Theta)q(Z)} dZ d\Theta
    • この式はベイズの定理よりp(X, \Theta, Z)=p(X|\Theta, Z)p(Z | \Theta)p(\Theta)と展開している。
      • これはベイズの定理から導かれる展開で、モデル次第ではもっと色々展開できる可能性がある(p(Z|\Theta)=p(Z)とできたりとか)。
  • 後はこれを座標降下法の要領で、q(\Theta)を固定してq(Z)に関する最適化とq(Z)を固定してq(\Theta)に関する最適化を交互に繰り返せばよい。

変分EMアルゴリズムの更新式の導出

ここからはもう少し、この式を展開する。


\begin{align}
{\mathcal L} \left[q(\Theta)q(Z) \right] &= \int \int q(\Theta)q(Z) \ln \frac{p(X|\Theta, Z)p(Z | \Theta)p(\Theta)}{q(\Theta)q(Z)} dZ d\Theta \\
&= \int \int q(\Theta)q(Z) \ln p(X|\Theta, Z)  dZ d\Theta  + \int \int q(\Theta)q(Z)\ln p(Z | \Theta) dZ d\Theta \\
&\quad +\int \int q(\Theta) \ln \frac{p(\Theta)}{q(\Theta)}d\Theta +  \int q(Z)\ln\frac{1}{ q(Z)} dZ 
\end{align}

この式において、まずq(\Theta) = {\hat q(\Theta)}という既知の関数として、q(Z)について最大化することを考える。


\displaystyle
\begin{align}
{\mathcal L} \left[{\hat q}(\Theta) q(Z) \right] &=\int \int {\hat q}(\Theta)q(Z) \ln p(X|\Theta, Z)  dZ d\Theta  + \int \int {\hat q}(\Theta)q(Z)\ln p(Z | \Theta) dZ d\Theta \\
&\quad + \int q(Z)\ln\frac{1}{ q(Z)} dZ + {\rm const} \\
&=\int  q(Z) \langle \ln p(X|\Theta, Z) \rangle_{{\hat q}(\Theta)}  dZ   + \int q(Z)\langle \ln p(Z | \Theta) \rangle_{ {\hat q}(\Theta)} dZ \\ 
&\quad + \int q(Z)\ln\frac{1}{ q(Z)} dZ + {\rm const} \\
&=\int  q(Z) \ln \exp(\langle \ln p(X|\Theta, Z) \rangle_{{\hat q}(\Theta)})  dZ   + \int q(Z)\ln \exp(\langle \ln p(Z | \Theta) \rangle_{ {\hat q}(\Theta)}) dZ \\
&\quad + \int q(Z)\ln\frac{1}{ q(Z)} dZ + {\rm const} \\
&= \int q(Z)\ln\frac{\exp \left(\langle \ln p(X|\Theta, Z) \rangle_{{\hat q}(\Theta)} + \langle \ln p(Z | \Theta) \rangle_{ {\hat q}(\Theta)} \right)}{ q(Z)} dZ + {\rm const}
\end{align}

この最後の式をよく見ると、分子をZ積分した時1になれば(つまり確率分布であれば)、これはKLダイバージェンス(に-1をかけたもの)である。そのため、

\displaystyle r(Z) \propto \exp \left(\langle \ln p(X|\Theta, Z) \rangle_{{\hat q}(\Theta)} + \langle \ln p(Z | \Theta) \rangle_{ {\hat q}(\Theta)} \right)
\displaystyle \int r(Z) dZ = 1

と定義すれば、

\displaystyle {\mathcal L} \left[{\hat q}(\Theta) q(Z) \right] = -KL\left[ q(Z) \parallel r(Z) \right] + {\rm const}

になり、下限を最大化するためにはKLダイバージェンスを0にするq(Z)を求めればよい。すなわち、

\displaystyle q(Z) = r(Z) \propto exp \left(\langle \ln p(X|\Theta, Z) \rangle_{{\hat q}(\Theta)} + \langle \ln p(Z | \Theta) \rangle_{ {\hat q}(\Theta)} \right)

とすればよい。この潜在変数側の更新を変分Eステップという。ちなみに両辺logを取ると、

\displaystyle \ln q(Z) =  \langle \ln p(X|\Theta, Z) \rangle_{{\hat q}(\Theta)} + \langle \ln p(Z | \Theta) \rangle_{ {\hat q}(\Theta)}  + {\rm const}

と書ける。q(Z)がどんな分布かを求める際、このlogを取った形式で計算をすることが多い。

同様の式展開をq(Z) = {\hat q}(Z)という既知の関数としてq(\Theta)について最大化する場合に適用すると、


\displaystyle
\begin{align}
{\mathcal L} \left[q(\Theta) {\hat q}(Z) \right] =
\int q(\Theta)\ln\frac{\exp \left(\langle \ln p(X|\Theta, Z) \rangle_{{\hat q}(Z)} + \langle \ln p(Z | \Theta) \rangle_{ {\hat q}(Z)} \right) p(\Theta)}{ q(\Theta)} d\Theta + {\rm const}
\end{align}

なので、q(\Theta)に関する更新式は、

\displaystyle q(\Theta) \propto \exp \left( \langle \ln p(X | \Theta, Z) \rangle_{{\hat q}(Z)}  +  \langle  \ln p(Z | \Theta) \rangle_{{\hat q}(Z)}  \right) p(\Theta)

となる。こちらは変分Mステップと呼ばれる。こちらも両辺logを取った式を書いておくと、

\displaystyle \ln q(\Theta) =   \langle \ln p(X | \Theta, Z) \rangle_{{\hat q}(Z)}  +  \langle  \ln p(Z | \Theta) \rangle_{{\hat q}(Z)} + \ln  p(\Theta) + {\rm const}