機械学習でも平均場近似使うの知ってる?

EMアルゴリズムと関連しますが、確率分布の近似手法について解説したいと思います。
今回の記事は、平均場近似です。物理学科の学生は必ず習うやつですが、機械学習でも活躍します。データを弄る時に、変数が一つだけというのは本当に稀です。一変数と多変数では面倒くささと難しさが全然違います。
そんな時に、変数同士の繋がりを断ち切って一変数ずつ考えてみるのが平均場近似です。変数が大量で、複雑な関係がある場合も簡単な計算に帰着できたり、確定的な解が求まるのがメリットです。しかし、近似的な解である事には注意が必要です。
EMアルゴリズムを知っていると議論の流れが分かりやすいと思います。以下の記事を読んでみてください。

EMアルゴリズム

扱うデータが複数の分布を持っていそうな時、分布の混合割合と、パラメーターを推定する方法があります。それはEMアルゴリズムと呼ばれています。EMアルゴリズムの原理について解説します。

EMアルゴリズムのpythonによる実装と一般化

EMアルゴリズムのまとめをして、python 上で実装します。最後に、EMアルゴリズムの一般化について触れます。

EMアルゴリズムの一般化と計算例

EMアルゴリズムを色々な確率分布に対して使えるようにします。正規分布とベルヌーイ分布の場合で、EMアルゴリズムを実装するための計算をします。

PRMLが参考図書です。日本語だと下巻に書いてると思います。

EMアルゴリズムの復習
平均場近似
平均場近似の計算例
まとめ

EMアルゴリズムの復習

簡単にEMアルゴリズムの復習をしておきます。
例えば、食の好みのアンケートを取った時にラーメン焼肉ハンバーガーが上位に並んでいたら、なんとなく男性ばかりがアンケートに答えたのかなと思いますね。(私はそう思ってしまいます。)
もしかしたら、男性だけがアンケートに答えたのかもしれませんし、ラーメン屋の前でアンケートを取ったのかもしれません。こんな感じに、ぱっと見のデータには表れてこない情報を隠れ変数と呼びます。そのような状況の問題を考えます。
EMアルゴリズムで解く問題設定は以下のようにします。
未知の確率分布
$$\begin{eqnarray}
p(X| \theta )
\end{eqnarray}$$
の$ \theta $(が従う確率分布)を決めたいのですが、隠れ変数$Z $の存在と、その確率分布の形$q(Z) $、更に$ p(X,Z) $が分かっている事を仮定します。
この時、対数尤度関数が以下のように計算出来ます。¹
$$\begin{eqnarray}
\log p(X| \theta ) &=& \mathcal{L}(q) +KL(q \| p) 　\\
\mathcal{L}(q, \theta) &=& \ \int_{Z} q(Z) \log \left\{ \frac{p(X,Z| \theta )}{q(Z)} \right\} dZ \\
KL(q \| p) &=& – \int_{Z} q(Z) \log \left\{ \frac{p(Z|X,\theta)}{q(Z)} \right\} dZ
\end{eqnarray} $$
ここで、$KL (q\| p) $はKLダイバージェンスと呼ばれる量で、確率分布同士の距離を測る量です。一応距離みたいなもので、常に0以上です。また、 $KL (q\| p) =0 $となるのは$ p=q $となる時だけです。 ²
この事から、常に $ \log p(X| \theta ) \geq \mathcal{L} (q ) $ です。適当に$ \theta $を固定し、$ q=p $ と置いて $ \mathcal{L} (q ) $ を最大化する。更新されたパラメーターで、またしても $q=p $と置いて $ \mathcal{L} (q ) $ を最大化する。というのを繰り返すのがEMアルゴリズム ³ です。
大事なのは、対数尤度を最大にするパラメーター$ \theta $を求めるのが目標である事と、その為に対数尤度を$ q $ の関数だと思うという事です。

平均場近似

EMアルゴリズムでは、対数尤度を最大にする為に、$q $を色々動かして最適なパラメーター$ \theta $を得る事で$ p(X,\theta ) $を求めました。次に考える問題は、$ p(X|Z) $の形を仮定して、$ p(Z|X) $を求める問題です。
EMアルゴリズムのように、対数尤度を最大にする$q(Z) $を求めますが、ここで話を簡単にするために$q(Z) $に制限を付けます。ここで、平均場近似が出てきます。
平均場近似では、$ q(Z) = \prod q_i (Z_i ) $と仮定します。つまり、隠れ変数$Z_i $ たちが独立である事を仮定します。
大体は独立同分布( independent and identically distributed )⁴でもあることを仮定します。つまり、$q_i (Z_i) $ が正規分布に従っていたら他の$ q_j (Z_j )$も正規分布に従うという感じです。
この仮定の下で$q_j (Z_j ) $達がどう予測されるか調べましょう。$q(Z) $の形を指定してしまっているので、対数尤度をEMアルゴリズムのように最大化していくことは出来ません。
具体的には、いきなりKLダイバージェンスを0にすることが出来ません。一方で、$ \log p(X) \geq \mathcal{L} (q) $が成り立つので、$ \mathcal{L} (q) $を大きくすれば、対数尤度を大きくすることが出来ます。 $ \mathcal{L} (q) $ を最大化するための $q_i (Z_i) $ 達を求めます。
まずは、 $ \mathcal{L} (q) $ を、一つの $q_i (Z_i) $ についてまとめてみます。
$$\begin{eqnarray}
\mathcal{L}(q) &=& \ \int_{Z} \prod q_i (Z_i ) \log \left\{ \frac{p(X,Z )}{ \prod q_i (Z_i ) } \right\} dZ \\
&=& \int \left \{ \int \prod_{i \neq j} q_i (Z_i)\log p(X,Z) dZ_{i\neq j} \right \} dZ_j \\
&-& \int q_j \log q_j dZ_j + {\rm const }
\end{eqnarray}$$
ここで、$ \int dZ_{i \neq j } $ は$ i=j $以外の変数で積分することを示しています。また、 $ Z_j $ に関係ない項は$ { \rm const } $でまとめました。以下のように関数を定義して、式を見やすくしましょう。
$$\begin{eqnarray}
E_{i\neq j}[\log p(X, Z)] &=& \int \log p(X,Z) \prod_{i\neq j} q_i dZ_i \\
\log \tilde{p} (X, Z_j) &=& E_{i\neq j}[\log p(X, Z)] + {\rm const}
\end{eqnarray}$$
この式を使うと、$ \mathcal{L} (q) $は以下のように書き直せます。
$$\begin{eqnarray}
\mathcal{L}(q) &=& \int q_j \log \tilde{p} (X, Z_j) dZ_j – \int q_j \log q_j dZ_j \\
&=& – KL(q_j \| \tilde{p}(X,Z_j) )
\end{eqnarray}$$
$q_j $ だけで考えると、$ \mathcal{L}(q) $を最大化するには、$ q_j =\tilde{p} (X,Z_j ) $ と置けばよい事が分かります。⁵つまり、各$ q_j $に対して、以下のように確率分布を更新すれば良いです。
$$\begin{eqnarray}
q^{\ast}_j (Z_j) &=& \frac{\exp( E_{i\neq j} [\log p(X, Z)] )}{ \int \exp( E_{i\neq j} [\log p (X, Z)] ) dZ_j}
\end{eqnarray}$$
この手法の良い所は、確率分布が比較的簡単になることと、計算が一回でおしまいな所です。もちろん、何回も計算しても良いですが。
平均場近似の理論はこんなところにしておいて、簡単な例を計算してみます。

平均場近似の計算例

簡単な例として、多次元正規分布を平均場近似で求めてみます。
具体的には、以下のような状況を考えます。
$$\begin{eqnarray}
\mu &=& (\mu_1 , \mu _2 ) , \Lambda = \{ \Lambda _{ij} \}_{i,j =1,2} \\
p(Z) &=& \mathcal{N}(Z|\mu, \Lambda ^{-1} )
\end{eqnarray}$$
ただし、$ \Lambda_{12} =\Lambda_{21} $とします。この確率分布を平均場近似$ q(Z) = q_1 (z_1) q_2 (z_2) $で再現してみます。
$$\begin{eqnarray}
\log q_1 ^{\ast} (z_1) &\sim & E_{z_2} [\log p(z) ] \\
&=& E_{z_2}\left[ -\frac{1}{2} (z_1 – \mu _1 )^2 \Lambda_{11} – (z_1 -\mu _1)\Lambda _{12} (z_2 -\mu _2 )\right] + {\rm const} \\
&=& -\frac{1}{2} z_1 ^{2} \Lambda _{11} + z_1 \mu _1 \Lambda _{11} -z_1 \Lambda _{12} (E_{z_2} [z_2] -\mu _2 ) + {\rm const}
\end{eqnarray}$$
式変形に際して、$z_1 $に関係ない項は const にまとめています。また、$E_{z_2} $ は、$q_2(z_2 )$で期待値を取ることを表しており、$z_2 $を含まない項に対しては $ \int q_2(z_2) dz_2 =1 $を掛けるだけである事に注意しましょう。
また、$ \log $を取った結果が二次形式になっているので、 $ q_1 ^{\ast} (z_1) $は一次元の正規分布になること⁶
が分かります。
$$\begin{eqnarray}
q_1 ^{\ast} (z_1) &=& \mathcal{N}(z_1 | m_1, \Lambda _{11} ^{-1} )\\
m_1 &=& \mu _1 – \Lambda ^{-1} _{11} \Lambda _{12} (E_{z_2} [z_2] -\mu _2 )
\end{eqnarray}$$
$q_2 ^{\ast} $も同様に求める事が出来て、
$$\begin{eqnarray}
q_2 ^{\ast} (z_2) &=& \mathcal{N}(z_2 | m_2, \Lambda _{22} ^{-1} )\\
m_2 &=& \mu _2 – \Lambda ^{-1} _{22} \Lambda _{21} (E_{z_1} [z_1] -\mu _1 )
\end{eqnarray}$$
最後に、$ E_{1_1} [z_1] , E_{z_2} [z_2] $ についてですが、それぞれが正規分布に従うことが分かっているので、 $ E[z_i] =m_i $ となる必要があります。これは、$ E_{z_i} [z_i] = \mu _i $⁷と置くことで上手くいきます。
結果的に、平均場近似で求められる確率分布は以下のようになります。
$$\begin{eqnarray}
q_1 ^{\ast} (z_1) &=& \mathcal{N}(z_1 | \mu _1, \Lambda _{11} ^{-1} )\\
q_2 ^{\ast} (z_2) &=& \mathcal{N}(z_2 | \mu _2, \Lambda _{22} ^{-1} )\\
q(Z) = &=& q_1 ^{\ast} (z_1) q_2 ^{\ast} (z_2)
\end{eqnarray}$$
結果的に、多次元正規分布の共分散行列の対角成分以外は捨て去った確率分布が出てきました。共分散行列の非対角成分はそれぞれの変数の相互作用と思える訳ですが、それらの効果は平均して消してしまったと捉えられます。
相関が小さい場合には上手く近似出来て、そうでない場合は微妙な近似になります。平均値だけは合うんですけどね。
平均場近似は、変分推論とか、ベイズ統計学周りで良く使われるようです。別の記事で、変分推論について書きたいと思います。

まとめ

EMアルゴリズムを簡単に復習した
確率分布の近似手法の一つとして、平均場近似を紹介した。
多次元正規分布に対して、平均場近似を計算してみた。

ベイズの定理で対数尤度の中身を変形してます
KLダイバージェンスの解説記事も一応あります。https://masamunetogetoge.com/entropy
EMアルゴリズムの解説記事では実例でも抽象的にも書いてるので、分からなかったら上に初めに挙げた記事を読んでください。
i.i.d. ってやつです。
もちろん$ q $を $q_j $達で分解しているので、本当の最大値になる保証はありません。
正規分布の$ \exp $の中を展開して、変数に関係ある項だけを見ると、以下のようになる事に注意しましょう。
$$\begin{eqnarray}
-\frac{1}{2} ( x^2 -2 \mu x )/\sigma^2
\end{eqnarray}$$
$ m _i = \mu _i $