今日はGMM（Gaussian Mixture Model）の学習をEMアルゴリズムを通して実装します。

モチベーション

クラスタリングをしたいのです。GMMはその名のとおり統計モデルであり、これから分類しようとしているデータ点が複数のガウス分布から生成されていると仮定をしています。EM法はGMMのパラメータ推定に用いられる最尤推定のアルゴリズムです。クラスタリングに最適な分布を知りたいということですね。

データ点が複数のガウス分布から生成されているという仮定は以下の式で表すことができます。$p(x)$は混合ガウス分布の確率密度関数であり、$K$はクラスタの数です。

p(\boldsymbol x) = \sum_{k=1}^K \pi_k \mathscr{N}(\boldsymbol x|\boldsymbol \mu_k, \boldsymbol \Sigma_k)\,,\quad \pi_k \geq 0\,,\quad \sum_{k=1}^K\pi_k = 1

$\pi_k$は混合係数と言って、各正規分布の重みを意味しており、各重みの合計は1にならなけれいけばいけません。

これまで扱ってきたK-meansはHardクラスタリングというクラスタリングで、それは各データ点が1つのクラスタに属するようなクラスタリングを意味します。しかし、GMM法では、データ点が複数のクラスタに属する確率を持ちます。この特徴からGMMは対照的にSoftクラスタリングとも呼ばれています。また、GMMはFuzzy K-meansとも言われており、K-meansと比較対象のアルゴリズムであることも分かります。

EM法の概要

EM法は、先ほどのガウス混合分布のパラメータ$\Theta = \lbrace \pi_k, \boldsymbol \mu_k, \boldsymbol \Sigma_k \rbrace_{k=1}^{K}$の最尤推定を目標とします。

以下の関数を目的関数とし、最小化を行います。（参考までに、マイナスを取って、最大化問題として解く場合もあります。）

E_{gmm}(\Theta):=-\sum_{n}^{N}\text{log}(\sum_{k}^{K}\pi_k \mathscr{N}(x|\boldsymbol \mu_k, \boldsymbol \Sigma_k))

Eステップ

Eステップでは、現在のパラメータ$\Theta$に対し、各データ点（$n=1, ..., N$）に対しての、確率密度関数$k=(1, ..., K)$の負担率を求めます。

ここがまた厄介な点で、僕は理解にとても苦労しました。負担率（Responsibility）ってなんや...となっていましたが、図を見れば一目瞭然です。@kenmatsu4(まつけん)さんの記事の画像を引用します。

例として、K=3と仮定してください。これは、これから分類しようとしているデータ点が3つのガウス分布から生成されているということを意味します。横軸がデータ、縦軸が確率密度です。

すると、各データ（横軸）に対して、それぞれの確率密度関数（赤青緑）の占める割合が違うことが分かりますよね。この割合こそが負担率（Responsibility）です。それを分かりやすく表したのが、以下の図です。（@kenmatsu4(まつけん)）

一般化すると負担率（Responsibility）は以下のように書くことができます。

$z_{kn}=p(k|\boldsymbol x_n;\Theta)=\frac{\pi_k\mathscr{N}(\boldsymbol x_n|\boldsymbol \mu_k, \boldsymbol \Sigma_k)}{\Sigma_{j=1}^{K}\pi_j\mathscr{N}(x_n;\boldsymbol \mu_j, \boldsymbol \Sigma_j)}$

こちらは、n（データ点）とk（どのガウス分布か）を入れてやれば、その点での注目しているガウス分布の占める割合が分かるという次第です。分母の$j$という文字は、$k$にかぶらないように導入しているだけです。

Mステップ

パラメータ調整をしていきます。すべてのパラメータは、負担率（Responsibility）を基に計算されます。

$\pi_k=\frac{1}{N}\Sigma_{n=1}^{N}z_{nk}$
$\boldsymbol \mu_k=\frac{\Sigma_{n=1}^{N}z_{nk}x_n}{\Sigma_{n=1}^{N}z_{nk}}$
$\boldsymbol \Sigma_{k}=\frac{\Sigma_{n=1}^{N}z_{nk}(x_n-\mu_k)(x_n-\mu_k)^T}{\Sigma_{n=1}^{N}z_{nk}}$

$\boldsymbol \mu_k, \boldsymbol \Sigma_k$はデータ点が一次元とも限らないので、ベクトルです。今回は二次元データを扱っています。

とにかく実装

アルゴリズムに関して大雑把に理解したら、実行してしまって、実行しながら間庭にを深めていきましょう。

環境

安定のGoogle Colab様

ソースコード解説

# imports
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
import scipy.linalg as la
import matplotlib.cm as cm
from matplotlib import rc
import time
import math
from scipy.stats import multivariate_normal
from IPython import display

%matplotlib inline

np.random.seed(42)

まずは必要なライブラリをimportします。

# Choose a GMM with 3 components

# means
m = np.zeros((3,2))
m[0] = np.array([1.2, 0.4])
m[1] = np.array([-4.4, 1.0])
m[2] = np.array([4.1, -0.3])

# covariances
S = np.zeros((3,2,2))
S[0] = np.array([[0.8, -0.4], [-0.4, 1.0]])
S[1] = np.array([[1.2, -0.8], [-0.8, 1.0]])
S[2] = np.array([[1.2, 0.6], [0.6, 3.0]])

# mixture weights
w = np.array([0.3, 0.2, 0.5])

データを生成するGMMを定義します。ここでは、3つのガウス分布からデータを生成します。

N = 600 # total number of data points
N_split = (N*w).astype(int) # number of data points per mixture component
x = []
y = []
for k in range(3):
    x_tmp, y_tmp = np.random.multivariate_normal(m[k], S[k], N_split[k]).T
    x = np.hstack([x, x_tmp])
    y = np.hstack([y, y_tmp])

data = np.vstack([x, y])
print(data.shape)

600のデータを生成します。

# plot the dataset
plt.figure()
plt.title("Mixture components")
plt.plot(x, y, 'ko', alpha=0.3)
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")

# plot the individual components of the GMM
plt.plot(m[:,0], m[:,1], 'or')

X, Y = np.meshgrid(np.linspace(-10,10,100), np.linspace(-10,10,100))
pos = np.dstack((X, Y))

for k in range(3):
    mvn = multivariate_normal(m[k,:].ravel(), S[k,:,:])
    xx = mvn.pdf(pos)
    plt.contour(X, Y, xx,  alpha = 1.0, zorder=10)

# plot the GMM
plt.figure()
plt.title("GMM")
plt.plot(x, y, 'ko', alpha=0.3)
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")

# build the GMM
gmm = 0
for k in range(3):
    mix_comp = multivariate_normal(m[k,:].ravel(), S[k,:,:])
    gmm += w[k]*mix_comp.pdf(pos)

plt.contour(X, Y, gmm,  alpha = 1.0, zorder=10);

データを可視化します。このデータを今からGMMを用いてクラスタリングしていきます。

K = 3 # number of clusters

means = np.zeros((K,2))
covs = np.zeros((K,2,2))
for k in range(K):
    means[k] = np.random.normal(size=(2,))
    covs[k] = np.eye(2)

weights = np.ones((K,1))/K
print("Initial mean vectors (one per row):\n" + str(means))
print("Initial covariance matrices:\n" + str(covs))

EM法のためにパラメータを初期化します。

#EDIT THIS FUNCTION
NLL = [] # negative log-likelihood of the GMM
gmm_nll = 0
likelihood_sum = 0
for k in range(K):
    likelihood_sum += multivariate_normal.pdf(mean=means[k], cov=covs[k], x=data.T) * weights[k]
    #<-- REPLACE THIS LINE, You can use multivariate_normal.pdf(mean=, cov=, x=) function
gmm_nll += -np.sum(np.log((likelihood_sum)))  # log(0) エラーを防ぐため 1e-9 を加える
NLL += [gmm_nll]#<-- REPLACE THIS LINE

print(NLL)

plt.figure()
plt.plot(x, y, 'ko', alpha=0.3)
plt.plot(means[:,0], means[:,1], 'oy', markersize=25)

for k in range(K):
    rv = multivariate_normal(means[k,:], covs[k,:,:])
    plt.contour(X, Y, rv.pdf(pos), alpha = 1.0, zorder=10)

plt.xlabel("$x_1$");
plt.ylabel("$x_2$");

配列NLLは、各イテレーションごとに更新されるNLL Lossが格納されます。

実行結果は以下のようになるはずです。

[7038.321287578101]

初期状態ではガウス分布の位置がかなり中央に偏っており、分布を表すガウス分布のパラメータが最適でないことが読み取れます。また、NLL Lossが7038という値になっていることが分かります。訓練によりこの値を最小化することが目標です。各イテレーションで、NLL Lossが小さくなっていくはずです。そのNLL Lossの情報をNLL配列にイテレーション毎に格納していきます。つまりNLL配列は各イテレーション毎にサイズが1増えていきます。最初コーディングをしていた時は、これを理解するのに多くの時間を要しました。

r = np.zeros((K,N)) # will store the responsibilities
for em_iter in range(100):
    print(em_iter)
    means_old = means.copy()

    # E-step: update responsibilities
    # 負担率の分子を計算している
    for k in range(K):
        r[k, :] = weights[k] * multivariate_normal.pdf(mean=means[k], cov=covs[k], x=data.T)
    print(r)
    # 確率にする
    r = r / np.sum(r, axis=0)
    
    # M-step
    N_k = np.sum(r, axis=1)

    for k in range(K):
        N_ave = 0
        D_ave = N_k[k]
        # update the means
        for n in range(N):
            N_ave += r[k][n] * data[:, n]
        means[k] = N_ave / D_ave 

        # update the covariances
        N_cov = 0
        D_cov = N_k[k]
        for n in range(N):
            N_cov += r[k][n] * np.outer(data[:, n] - means[k], data[:, n] - means[k])
        covs[k] = N_cov / D_cov

    # weights
    weights = np.zeros((K,))
    for k in range(K):
        weights[k] = N_k[k] / N

    # negative log-likelihood
    gmm_nll = 0
    likelihood_sum = 0
    for k in range(K):
        likelihood_sum += multivariate_normal.pdf(mean=means[k], cov=covs[k], x=data.T) * weights[k]
    gmm_nll += -np.sum(np.log((likelihood_sum)))
    # NLL配列にgmm_nllという要素を追加. gmm_nllは各イテレーション毎の目的関数の値です。
    NLL += [gmm_nll]

    print("NLL:")
    print(NLL)

    plt.figure()
    plt.plot(x, y, 'ko', alpha=0.3)
    plt.plot(means[:,0], means[:,1], 'oy', markersize=25)
    for k in range(K):
        rv = multivariate_normal(means[k,:], covs[k])
        plt.contour(X, Y, rv.pdf(pos), alpha = 1.0, zorder=10)

    plt.xlabel("$x_1$")
    plt.ylabel("$x_2$")
    plt.text(x=3.5, y=8, s="EM iteration "+str(em_iter+1))

    # NLL[em_iter+1]は新しい対数尤度、NLL[em_iter]は元の対数尤度. それらの差がある程度収束したら、そのgmm_nllを最適値とする. 
    if la.norm(NLL[em_iter+1]-NLL[em_iter]) < 1e-6:
        print("Converged after iteration ", em_iter+1)
        break

# plot final the mixture model
plt.figure()
gmm = 0
for k in range(3):
    mix_comp = multivariate_normal(means[k,:].ravel(), covs[k,:,:])
    gmm += weights[k]*mix_comp.pdf(pos)

plt.plot(x, y, 'ko', alpha=0.3)
plt.contour(X, Y, gmm,  alpha = 1.0, zorder=10)
plt.xlim([-8,8]);
plt.ylim([-6,6]);

こちらでは、実際にイテレーションを回していきます。ガウス分布が最適な位置に近づいていく様子も、プロットしています。

iteration2-60は省略

plt.figure()
plt.semilogy(np.linspace(1,len(NLL), len(NLL)), NLL)
plt.xlabel("EM iteration");
plt.ylabel("Negative log-likelihood");

idx = [0, 1, 9, em_iter+1]

for i in idx:
    plt.plot(i+1, NLL[i], 'or')

最後に各イテレーションとNLLの関係をグラフにします。

初期のイテレーションで、かなり損失が減少しており、イテレーションが進むにつれて小さな変化が続いていることが分かります。

最後に負担率rをプロットしてみます。

print(r)

結果は、
[[8.01547888e-04 2.20475080e-02 4.76704134e-03 ... 9.90050452e-01
9.99028612e-01 9.99991227e-01]
[9.99198452e-01 9.77952492e-01 9.95232959e-01 ... 9.94954813e-03
9.71387572e-04 8.77273282e-06]
[2.08381626e-11 2.96124653e-18 3.89865034e-11 ... 1.33661412e-12
1.28613589e-26 1.38198714e-33]]
となりました。ちなみに、rのサイズは(3, 600)です。600点の各データが確率的に（softに）各クラスタに所属していることが分かります。またそれらの和は各点で1となることも確認できます。

感想

実行時間はK-meansやk-mediansよもかなり短かったようです。アルゴリズムを理解するのが非常に難しかったですが、配列のサイズを表示してみたり、図示すると理解が深まるところもあります。

参考文献

EMアルゴリズム徹底解説