C++を使ったEMアルゴリズムの実装（＋Pythonによるプロット）

複数のガウス分布を重ねあわせてできる混合ガウス分布はある程度複雑な分布を近似するときに使われ、Fisher Vectorや声質変換、手書き漢字変換などの応用がある重要な分布です。今回はEMアルゴリズムを用いて混合ガウス分布による近似を求めるプログラムを実装してみました。
EMアルゴリズムは混合ガウス分布以外にも使えるアルゴリズムですが、ここでは混合ガウス分布に話を限定しています。

EMアルゴリズム

式はPRML第9章に従います。いつものようにベクトルの太字は省略しています。

混合ガウス分布は
$\sum_{k=1}^{K} \pi_k \mathcal{N}(x|\mu_k, \Sigma_k)$
という形で表される分布です。
$\pi_k$ はk番目のガウス分布の割合を表し、 $0 \geq \pi_k \geq 1$ と $\sum_{k=1}^{K} \pi_k =1$ を満たします。

D次元ガウス分布は
$\mathcal{N}(x|\mu, \Sigma) = \frac{1}{(2\pi)^{D/2}}\frac{1}{|\Sigma|^{1/2}}\exp\{-\frac{1}{2}(x-\mu)^T\Sigma^{-1}(x-\mu)\}$
という式で表されます。

式変形を楽にするために、潜在変数zを導入します。
K次元ベクトル $z$ を、どれか一つの要素だけが1で他の全ての要素が0であるようなベクトルとします。
zの確率分布は、 $0 \geq \pi_k \geq 1$ と $\sum_{k=1}^{K} \pi_k =1$ を満たす $\pi_k$ を用いて
$p(z_k = 1) = \pi_k$ ( $z_k$ はzのk番目の要素）
と表せるとします。

また、zが与えられたもとでのxの条件付き分布は
$p(x|z_k=1) = \mathcal{N}(x|\mu_k, \Sigma_k)$
と表せるとします。

このとき、xの周辺分布は
$p(x) = \sum_z p(z)p(x|z) = \sum_{k=1}^{K} \pi_k \mathcal{N}(x|\mu_k, \Sigma_k)$
となるので、混合ガウス分布になることが分かります。

また、xが与えられたもとでのzの条件付き確率 $\gamma(z_k)$ を、
$\gamma_k := p(z_k=1|x)$

$= \frac{p(z_k=1)p(x|z_k=1)}{\sum_{j=1}^K p(z_j=1)p(x|z_j=1)}$

$= \frac{\pi_k \mathcal{N}(x|\mu_k, \Sigma_k)}{\sum_{j=1}^K \pi_j\mathcal{N}(x|\mu_j, \Sigma_j)}$

と定義します。これは、k番目のガウス分布がxを説明する度合い、負担率(responsibility)として解釈できます。

ここで、観測したデータ集合 ${x_1 \dots x_N}$ をK個のD次元ガウス分布の混合で説明する問題を考えます。
データ集合を第n行が $x_n^T$ であるNxD行列で、潜在変数を第n行が $z_n^T$ であるNxK行列で表すことにすると、混合ガウス分布の対数尤度関数は
$\ln p(X|\pi, \mu, \Sigma) = \sum_{n=1}^{N} \ln \{\sum_{k=1}^{K} \pi_k \mathcal{N}(x_n|\mu_n, \Sigma_k)\}$
となります。

対数尤度関数の $\mu_k$ に関する微分を0として整理すると、
$\mu_k = \frac{1}{N_k}\sum_{n=1}^{N} \gamma(z_{nk})x_n$
となります。同様に、 $\Sigma_k$ に関する微分を0として整理すると、
$\Sigma_k = \frac{1}{N_k}\sum_{n=1}^{N} \gamma(z_{nk})(x_n-\mu_k)(x_n-\mu_k)^T$
を得ます。ここで $N_k$ は
$N_k = \sum_{n=1}^{N}\gamma(z_{nk})$
とします。この値はk番目のガウス分布に割り当てられる点の実効的な数と解釈できます。

さらに、対数尤度関数を $\sum_{k=1}^{K} \pi_k =1$ という制約条件のもとで $\pi_k$ についてラグランジュの未定係数法で最大化すると
$\pi_k = \frac{N_k}{N}$
となります。

このとき、これらのパラメータは $\gamma(z_{nk})$ について最適ですが、 $\gamma(z_{nk})$ が尤度関数について最適とは限らないため $\gamma(z_{nk})$ を繰り返し計算する必要があります。こうして、次のようなEMアルゴリズムが与えられます。

1. $\mu_k, \Sigma_k, \pi_k$ を適当な値に初期化する

2. $\gamma(z_{nk})= \frac{\pi_k \mathcal{N}(x|\mu_k, \Sigma_k)}{\sum_{j=1}^K \pi_j\mathcal{N}(x|\mu_j, \Sigma_j)}$
を現在のパラメータに基づいて計算する（Eステップ）

3. $N_k = \sum_{n=1}^{N}\gamma(z_{nk})$ として、
$\mu_k^{new} = \frac{1}{N_k}\sum_{n=1}^{N} \gamma(z_{nk})x_n$
$\Sigma_k^{new} = \frac{1}{N_k}\sum_{n=1}^{N} \gamma(z_{nk})(x_n-\mu_k)(x_n-\mu_k)^T$
$\pi_k^{new} = \frac{N_k}{N}$
でパラメータを更新する（Mステップ）
4. 対数尤度が変化しなくなるか、一定の回数に達するまでEステップとMステップを交互に繰り返す

続いてこれを実装します。

Eigenのインストール

行列演算がたくさん出てくるのでEigenという行列演算ライブラリを使うことにしました。
公式Wikiから最新のコードをダウンロードして、解凍した中にあるEigenフォルダを/usr/local/includeなどのフォルダにコピーすればインストール終了です。特にコンパイルの必要はありません。使い方はGetting Startedなどを見てください。

EMアルゴリズムの実装

Eigenを使うと上のアルゴリズムの式をほぼそのまま書くだけで動きます。
今回はPRMLで使われていたOld Faithfulをデータセットに使いました。

#include <iostream>
#include <vector>
#include <cmath>
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

double Norm(VectorXd x, VectorXd mu, MatrixXd sigma){
    return exp(((x-mu).transpose() * sigma.inverse() * (x-mu)/(-2.0)).value()) 
        / sqrt(sigma.determinant())
        / pow(2.0*M_PI, x.size()/2.0);
}

int main() {
    const int D = 2; // dimension
    const int K = 2; // number of distribution
    int N;           // number of data
    vector<VectorXd> x;

    while(cin){
        double hoge;
        VectorXd tmp(D);
        for(int i=0; i<D; ++i){cin>>hoge; tmp(i)=hoge;}
        x.push_back(tmp);
    }

    // last data is null
    N = x.size()-1;
   
    VectorXd mu[K];     // mean vector
    MatrixXd sigma[K];  // covariant matrix
    double pi[K];       // mixture
    double eSize[K];     // effective cluster size (N_k in PRML)
    vector<vector<double> > gamma; // responsibility
    gamma.resize(N);
    for(int i=0; i<N; ++i) gamma[i].resize(K);
  
    // initialization
    for(int i=0; i<K; ++i) mu[i] = VectorXd::Zero(D);

    // initialize mu as mean of input
    VectorXd init_mu = VectorXd::Zero(D);
    for(int i=0; i<N; ++i) init_mu += x[i];
    init_mu /= N;

    mu[0] = 1.1 * init_mu;
    mu[1] = 0.9 * init_mu;

    for(int i=0; i<K; ++i) sigma[i] = MatrixXd::Identity(D, D);
    for(int i=0; i<K; ++i) pi[i] = 1.0/K;
    
    for(int loop = 0; loop < 20; ++loop){
        // E step
        for(int n=0; n<N; ++n){
            for(int k=0; k<K; ++k){
                double tmp = 0;
                for(int j=0; j<K; ++j){
                    tmp += pi[j] * Norm(x[n], mu[j], sigma[j]);
                }
                gamma[n][k] = pi[k] * Norm(x[n], mu[k], sigma[k]) / tmp;
            }
        }

        // M step
        for(int k=0; k<K; ++k){
           eSize[k] = 0;
           for(int n=0; n<N; ++n){
               eSize[k] += gamma[n][k];
            }
        }

        for(int k=0; k<K; ++k){
            VectorXd tmp_mu = VectorXd::Zero(D);
            for(int n=0; n<N; ++n){
                tmp_mu += gamma[n][k]*x[n];
            }
            mu[k] = tmp_mu / eSize[k];

            MatrixXd tmp_sigma = MatrixXd::Zero(D, D);
            for(int n=0; n<N; ++n) tmp_sigma += gamma[n][k]*(x[n]-mu[k])*(x[n]-mu[k]).transpose();
            sigma[k] = tmp_sigma / eSize[k];

            pi[k] = eSize[k] / N;
        }
        // show results
        cout << "start loop " << loop << ":" << endl;
        for(int k=0; k<K; ++k){
            cout << "pi" << k << ": " << pi[k] << endl;
            cout << "mu" << k << ": " << endl << mu[k] << endl;
            cout << "sigma" << k << ": " << endl << sigma[k] << endl;
        }
        cout << endl;
    }
}

コンパイルしたら

./a.out < faithful.txt

のようにして実行すると、20ループの結果が表示されます。20ループ目の結果は、

pi0: 0.644127
mu0: 
4.28966
79.9681
sigma0: 
0.169968 0.940609
0.940609  36.0462
pi1: 0.355873
mu1: 
2.03639
54.4785
sigma1: 
0.0691677  0.435168
 0.435168   33.6973

となりました。これが正しいのか可視化してみようと思います。

Pythonによるプロット

matplotlibのドキュメントにあった例を改造しました。
numpyをちゃんと使えばいいのですが、まだPythonからの好感度が上がりきっていないので適当な実装です。

#!/usr/bin/env python
# -*- coding: utf-8 -*- 

import matplotlib.pyplot as plt
import numpy as np
import math

def gaussian_2d(x, y, x0, y0, a, b, c, d):
    return np.exp(-0.5*((x-x0)*((x-x0)*d-(y-y0)*b)+(y-y0)*(-(x-x0)*b+(y-y0)*a))/(a*d-b*c))/6.28/math.sqrt(a*d-b*c)
delta = 0.025
x = np.arange(1.0, 6.0, delta)
y = np.arange(30.0, 100.0, delta)
X, Y = np.meshgrid(x, y)
# ここの数字は20回目のループの結果を入力したもの
Z1 = gaussian_2d(X, Y, 4.29, 79.97, 0.171, 0.941, 0.941, 36.05)
Z2 = gaussian_2d(X, Y, 2.04, 54.48, 0.069, 0.435, 0.435, 33.70)
Z = 0.644*Z1 + 0.356*Z2

plt.clf()
CS = plt.contour(X, Y, Z)
plt.clabel(CS, inline=1, fontsize=10)


fp = open("faithful.txt")
data_x = []
data_y = []
for row in fp:
    data_x.append(float((row.split()[0])))
    data_y.append(float((row.split()[1])))
fp.close()
plt.plot(data_x, data_y, "ro")

plt.show()