動く変分混合ガウス分布（導出編）- 動く PRML シリーズ（２）

やりたいこと

動く PRML シリーズ、第２回は変分混合ガウス分布 (variational Bayesian Gaussian mixture model, VB-GMM) です。
はじめに、前回の繰り返しになりますが、反復繰り返し型の機械学習アルゴリズムを理解するためには、大きく分けて二つのステップがあることを再確認します。一つ目はもちろん、

更新式を導出すること。

反復アルゴリズムの理論的性質はすべて更新式の形に反映されています。従って、この更新式を自力で導出することはとても勉強になります。

そして、二つ目は、

各イテレーションの内容をグラフにプロットし、実際の挙動を体感すること。

更新式には確かに必要なこと全てが記されていますし、熟練の研究者であれば、更新式の形から、その実際の挙動をある程度予測することが可能です。

しかし、それは楽譜だけを見てオーケストラを聴き取るようなものではないでしょうか。そのような力を手に入れる最良の訓練は、もちろん、楽譜と実際の音楽を何度も繰り返し聴き比べることです。

そこで本記事では、変分混合ガウス分布 (VB-GMM) の導出から始め、繰り返しごとに結果をプロットするプログラムを python で実装し、初期値応答、局所解頑健性、収束性能などを体感することを目標とします。

生成モデル

VB-GMM では、GMM と同じく、K 個のクラスから N 個のデータ点が生成されていると考えます。つまり、この分布からデータを生成したければ、K 個の面があるサイコロを N 回振って、k の目が出た回数だけ k 番目の正規分布からデータ点をサンプルすればよいのです。簡単ですね。
尤度関数は GMM と同じく、

　　 $p(X|Z,\;\;\mu,\;\;\lambda)\;\;=\;\;\prod_{nk}\mathcal{N}(x_n|\mu_k,\;\;\lambda_k)^{z_{nk}}$
　　 $p(Z|\pi)\;\;=\;\;\prod_{nk}\pi_k^{z_{nk}}$
となります。
さらに、変分ベイズ法はベイズ推定アルゴリズムなので、パラメータの事前分布 $p(\pi,\;\;\mu,\;\;\lambda)$ を導入する必要があります。

ベイズ推定の位置付け

変分ベイズ法は EM アルゴリズムに基づく繰り返し最適化法のひとつですが、EM アルゴリズムが適用できる問題には大きく分けて最尤推定、MAP 推定、ベイズ推定の三種類があり、それぞれ推定の目標が異なります。
観測変数を X、モデルパラメータを $\theta$ とすると、最尤推定は

　　 $\theta^*=\mathrm{argmax}_\theta p(X|\theta)$

を、MAP 推定は

　　 $\theta^*=\mathrm{argmax}_\theta p(\theta|X)$

を、ベイズ推定は

　　 $p(\theta|X)=\frac{p(X,\;\;\theta)}{p(X)}$

を推定することに相当しています。MAP 推定やベイズ推定では、生成分布 $p(X|\theta)$ の他に事前分布 $p(\theta)$ が必要なことが、ベイズの定理から分かります。

事前分布

先に述べた通り、変分ベイズ法では事前分布が必要なため、これを導入します。事前分布として、閉形式で更新式が導出できることが保証されている共役事前分布

　　 $p(\pi)\;=\;\mathrm{Dir}(\pi|\alpha_0)\;=\;\frac{\Gamma(K\alpha_0)}{\Gamma(\alpha_0)^K}\prod_k \pi_k^{\alpha_0-1}$
　　 $p(\mu_k,\;\;\lambda_k)\;=\;\mathcal{N}(\mu_k|m_0,\;\;(\beta_0 \lambda_k)^{-1})\mathcal{W}(\lambda_k|w_0,\;\;\nu_0)$

を使用します。

変分事後分布

さてベイズ推定では、観測変数で条件付けられた潜在変数の事後分布 $p(\theta|X)$ を推定するわけですが、様々な理由から、これをこのまま閉形式で導出することは出来ません。
そこで、 $p(\theta|X)$ をいくつかの変分事後分布 $q(\theta_1),\;\;q(\theta_2),\;\;\cdots$ の積で近似し、変分事後分布の積 $q(\theta_1)q(\theta_2)\;\cdots$ と真の事後分布 $p(\theta|X)$ の KL ダイバージェンスが最小となるように各分布を更新していきます。変分法の一般的な導出は PRML に譲りますが、大事なことは、

共役事前分布を導入すれば、変分事後分布は事前分布と同じ形の指数型分布を使って書ける。例えば、正規分布なら正規分布、ディリクレ分布ならディリクレ分布というように。
さらに、 $q(\theta_1)$ の変分事後分布は、全ての確率変数の同時分布 $p(X,\;\;\theta_1,\;\;\theta_2,\;\;\cdots)$ の対数尤度の $\theta_2,\;\;\theta_3\;\;\cdots$ に関する期待値を用いて計算できる。

という二点です。変分混合ガウス分布では、様々な理由*1から、 $p(Z,\;\;\pi,\;\;\mu,\;\;\lambda|X)\;\approx\;q(Z)q(\pi,\;\;\mu,\;\;\lambda)$ の形に分解します。

変分 E ステップ (VB-E Step)

変分推論の更新式はひたすら機械的に導出することができます。
まず、Z は多項分布にしたがうので、Z の変分事後分布も多項分布の形で書くことができます。ここでは、多項分布のパラメータを $\gamma_{nk}$ とおきます。

　　 $q(Z)\;=\;\prod_{nk} \gamma_{nk}^{z_{nk}}$

Z の変分事後分布の対数形は、Z 以外の潜在変数に関する完全同時分布の期待値として書けるので、

　　 $\ln q(Z)\;=\; \mathbb{E}_{\pi,\;\;\mu,\;\;\lambda}[\ln p(X,\;\;Z,\;\;\pi,\;\;\mu,\;\;\lambda)]\;+\;\mathrm{const.}$
　　 $=\;\mathbb{E}[\ln p(X | Z,\;\;\mu,\;\;\lambda)]\;+\;\mathbb{E}[\ln p(Z|\pi)]\;+\;\mathrm{const.}$
　　 $=\;\sum_{nk} z_{nk} \{\mathbb{E}[\ln \mathcal{N}(x_n|\mu_k,\;\;\lambda_k)]\; +\;\mathbb{E}[\ln \pi_k]\}\;+\;\mathrm{const.}$

となります。従って
　　 $\gamma_{nk}\;\propto\;\exp(\mathbb{E}[\ln \mathcal{N}(x_n|\mu_k,\;\;\lambda_k)]\;+\;\mathbb{E}[\ln \pi_k])$
となります。

ちなみに具体的な値はというと、例えば
　　 $\mathbb{E}[\ln \pi_k]\;=\;\psi(\alpha_k)\;-\;\psi(\sum_k \alpha_k)$
となります。ここで、 $\alpha_k$ は事後ハイパーパラメータ（後述）であり、 $\psi$ はディガンマ関数です。 $\pi_k$ が常に 1 より小さいことを考えると、この期待値は通常 -3 とか -5 とかいう値になります。このことを覚えておくとデバッグの際に重宝するでしょう。

$\mathbb{E}[\ln \mathcal{N}(x_n|\mu_k,\;\;\lambda_k)]$ については書く気すら起こりません…。というか、覚えていません。PRML を見て関数として実装して、その中身については忘れてしまうのが吉でしょう。

変分 M ステップ (VB-M Step)

自然な分解

E-Step では明示しませんでしたが、変分事後分布を導出する際には、更新しようとしている変数のみを残し、不要な項をどんどん定数項に押し込んでいきます。次に述べる変分 M ステップでは $\pi,\;\;\mu,\;\;\lambda$ の変分事後分布を計算するわけですが、これをちゃんと計算してやると、なんと

　　 $q(\pi,\;\;\mu,\;\;\lambda) \;=\;q(\pi) \prod_k q(\mu_k,\;\;\lambda_k)$

の形で書けることが分かります。そこで、文献によっては始めからこの形の分解を与えているものもあります。この導出はかなり文章を食うので省略します。

クラス混合比の更新

モクモクと計算していきます。変分事後分布を事後ハイパーパラメータ (posterior hyperparameter) を用いて $q(\pi) \;=\; \mathrm{Dir}(\pi|\alpha_1,\;\;\cdots\;\;\alpha_K)$ と書き、混合比の対数事前確率

　　 $\ln p(\pi) \;=\; \sum_k (\alpha_0 - 1) \ln \pi_k\;+\;\mathrm{const.}$
を用いると、

　　 $\ln q(\pi)\;=\; \mathbb{E}_Z[\ln p(X,\;\;Z,\;\;\pi,\;\;\mu,\;\;\lambda)]\;+\;\mathrm{const.}$
　　 $=\;\mathbb{E}[\ln p(Z|\pi)]\;+\;\ln(\pi)\;+\;\mathrm{const.}$
　　 $=\;\sum_k (\sum_n\mathbb{E}[z_{nk}] \;+\;\alpha_0\;-\;1) \ln \pi_k \;+\;\mathrm{const.}$
　　 $=\;\sum_k (N_k\;+\;\alpha_0\;-\;1) \ln \pi_k \;+\;\mathrm{const.}$
となります。ここで、 $N_k\;=\;\sum_n \gamma_{nk}$ です。

さらに、 $q(\pi) \;=\; \mathrm{Dir}(\pi|\alpha_1,\;\;\cdots\;\;\alpha_K)$ と比較することで、クラス混合比に対する更新式

　　 $\alpha_k\;=\;\alpha_0\;+\;N_k$

を得ます。some cheat なんて言われることもありますが、変分事後分布の計算は基本的に係数合わせだけで出来るので、微分してゼロを解く必要はありません。らくちんですね。

クラス平均と精度の更新

簡単なんですが、手間はかかります。手で導出する時は正規分布のハイパーパラメータを先に計算して、後でウィシャート分布のパラメータに集中するのが吉だと思います。導出は、例によって一次元で行います。

一次元版の正規ウィシャート分布は、変数変換により正規ガンマ分布で書くことができるため、事前分布と事後分布を
　　 $p(\mu_k,\;\;\lambda_k)\;=\;\mathcal{N}(\mu_k|m_0,\;\;(\beta_0 \lambda_k)^{-1})\mathcal{G}(\lambda_k|a_0,\;\;b_0)$
　　 $q(\mu_k,\;\;\lambda_k)\;=\;\mathcal{N}(\mu_k|m_k,\;\;(\beta_k \lambda_k)^{-1})\mathcal{G}(\lambda_k|a_k,\;\;b_k)$
と書きます。事前分布の対数確率は

　　 $\ln p(\mu_k,\;\;\lambda_k)\;=\; \frac{1}{2} \ln \frac{\beta_0 \lambda_k}{2 \pi}\;-\;\frac{\beta_0 \lambda_k}{2}(\mu_k\;-\;m_0)^2 \;+\;(a_0 - 1)\ln \lambda_k\;-\;b_0\lambda_k\;+\;\mathrm{const.}$

と書けるので、
　　 $\ln q(\mu_k,\;\;\lambda_k)\;=\; \mathbb{E}_Z[\ln(X,\;\;Z,\;\;\pi,\;\;\mu,\;\;\lambda)]\;+\;\mathrm{const.}$
　　 $=\;\mathbb{E}[\ln(X|Z,\;\;\mu,\;\;\lambda)]\;+\ln p(\mu_k,\;\;\lambda_k)\;+\;\mathrm{const.}$
　　 $=\;\sum_n \gamma_{nk}\ln \mathcal{N}(x_n|\mu_k,\;\;\lambda_k)\;+\;\ln p(\mu_k,\;\;\lambda_k)\;+\;\mathrm{const.}$
　　 $=\;\sum_n \gamma_{nk} \left\{ \frac{1}{2} \ln \frac{\lambda_k}{2 \pi}\;-\;\frac{\lambda_k}{2}(x_n\;-\;\mu_k)^2 \right\} \;+\;\ln p(\mu_k,\;\;\lambda_k)\;+\;\mathrm{const.}$
となります。

係数を比べることにより、
　　 $m_k\;=\;\frac{m_0 \beta_0\;+\;\sum_n \gamma_{nk} x_n}{\beta_0\;+\;N_k}$
　　 $\beta_k\;=\;\beta_0\;+\;N_k$
　　 $a_k\;=\;a_0\;+\;\frac{N_k}{2}$
　　 $b_k\;=\;b_0\;+\;\frac{1}{2} \sum_n \gamma_{nk}(x_n\;-\;m_k)^2$
を得ます。

更に、ガンマ分布とウィシャート分布は
　　 $\mathcal{W}(\lambda|w,\;\;\nu)=\;\mathcal{G}(\lambda|\frac{\nu}{2},\;\;\frac{1}{2w})$
と変数変換できるため、
　　 $w_k^{-1}\;=\;w_0^{-1}\;+\;\sum_n \gamma_{nk}(x_n\;-\;m_k)^2$
　　 $\nu_k\;=\;\nu_0\;+\;N_k$
を得ます。多次元ウィシャート分布の場合は、
　　 $W_k^{-1}\;=\;W_0^{-1}\;+\;\sum_n\gamma_{nk}(x_n\;-\bar{x}_k)(x_n\;-\bar{x}_k)^{\mathrm{T}}\;+\;\frac{\beta_0N_k}{\beta_0\;+N_k}(\bar{x}_k\;-\;m_0)(\bar{x}_k\;-\;m_0)^{\mathrm{T}}$
となることが知られています。(PRML を参照のこと。)