EMアルゴリズムを可視化で理解

EMアルゴリズムは潜在変数を含む確率モデルのパラメータ推定に用いられる最適化手法。有名な手法だが、あまりちゃんと理解していなかったので、PRMLで勉強してみた。

ここでは~~数式を書くのは辛いので~~、可視化でイメージを掴んでもらえればと思う。潜在変数を含むモデルとして、混合正規分布を題材とする。

混合正規分布とは？

複数の正規分布の重み付き和が混合正規分布となる。ここで重み $\pi_k$ は混合係数と呼ばれ、 $\sum ^ {K} _ {k=1}\pi _ {k}=1$ である。

$\displaystyle{p(\boldsymbol x) =\sum ^ {K} _ {k=1}\pi _ {k} \mathcal{N}( \boldsymbol x | \boldsymbol \mu _ {k}, \boldsymbol \Sigma _{k})}$

下図は、３つの二変量正規分布からなる混合分布の例。それぞれの分布の平均を✖️で、共分散行列を95%信頼区間を示す楕円で表現している。

f:id:kibya:20190910214621p:plain — 混合正規分布

EMアルゴリズムはこの混合分布のパラメータ、すなわち平均、共分散行列、混合係数をデータから推定するための方法である。
つまり、下図のようなラベルのないデータが与えられた際に、上述のパラメータを求める際に使う。なお、混合した分布の数は既知とする。

f:id:kibya:20190910214833p:plain — 混合正規分布（ラベルなし）

パラメータの最尤推定を行う際は、対数尤度関数をパラメータで偏微分して0とおいた式を解くのが基本的なアプローチだが、潜在変数を含むモデルでは、解析的にこの式を解くことができない。そこでEMアルゴリズムの出番となる。

EMアルゴリズムの流れ

平均、共分散行列、混合係数をランダムに初期化。
Eステップ：負担率(responsibility)を算出。
Mステップ：潜在変数を固定した上で、最尤法により平均、共分散行列、混合係数を新たに算出。
対数尤度が収束するまで上記のE Step, M Stepを繰り返す。

順を追って説明する。

1. パラメータの初期化

ここでは平均は一様分布からサンプリング、共分散行列は $0.2I$ 、混合係数は $\alpha=1000$ のディリクレ分布からサンプリングした。

f:id:kibya:20190910220353p:plain — パラメータの初期化

2. E step

E Stepでは潜在変数の事後確率、すなわち負担率(responsibility) $\gamma(z _ {i,k})$ を算出する。ここで $\gamma(z _ {i,k})$ はデータ $i$ がクラスター $k$ から得られた確率を表す。

下図は $\gamma(z _ {i,0}), \gamma(z _ {i,1}), \gamma(z _ {i,2})$ をRGBとして各データ点に色をつけている。
つまり赤色の点は赤い✖️のクラスターに属する確率が高く、紫の点は赤・青のクラスターに属する確率が半々くらい、と言うことになる。

負担率の計算には初期化したパラメータを用いる。2回目以降のE StepではM Stepで更新したパラメータを用いる。

f:id:kibya:20190910220413p:plain — E Step

3. M step

M Stepでは最尤法により平均、共分散行列、混合係数の各パラメータを算出する。
詳細な説明は省くが、混合正規分布の対数尤度を各パラメータで微分して0とおくと、各パラメータの最尤推定値は負担率を用いて表せることが分かる。
つまりE Stepで得られた負担率を用いて、より良いパラメータを再計算できることになる。

その様子を表したのが下図で、平均、共分散行列が修正されたことが分かる。

f:id:kibya:20190910220435p:plain — M Step

4. E Step, M Stepの繰り返し

上記のE Step, M Stepを繰り返すと、対数尤度が常に増加することを示せる。
実用上は、対数尤度の増加があるしきい値より小さくなった時点で、アルゴリズムが収束したと判定する。

そもそもなぜE Step, M Stepに分ける必要があるかと言うと、E Stepでは負担率の計算にパラメータが必要、M Stepではパラメータの計算に負担率が必要だからである。
言い換えると、「パラメータを固定して負担率を計算 → 負担率を固定してパラメータを計算」を繰り返して最適化していく方がEMアルゴリズムだと言える。

EMアルゴリズムが収束する様子をアニメーションで示す。この場合は最初の数回でほぼ最適解が得られていることが分かる。

f:id:kibya:20190910230921g:plain — EMアルゴリズム　アニメーション

今回の場合、対数尤度の変化量のしきい値を0.001とすると64 iterationで収束した。
対数尤度は最初のステップで大幅に増加し、以降は小刻みに増え続けている。

結び

EMアルゴリズムは古典的で重要な手法だが、最尤推定には特異性の問題など重大な限界がある。
ベイズ的なアプローチである変分推論では最尤推定が抱える問題を解決できるとのことなので、次回はこちらを解説したい。

今回作成したコードはこちら。

import numpy as np
from numpy import random
from scipy import stats
import matplotlib.pyplot as plt
import seaborn as sns
from collections import Counter
import matplotlib.animation as ani
from matplotlib.patches import Ellipse
%matplotlib inline


# draw 95% confidence interval of multivariate normal distribution
def draw_ellipse(mu, cov, ax, **kwargs):
    var, U = np.linalg.eig(cov)
    if U[1, 0] == 0:
        angle = 0
    else:
        angle = 180. / np.pi * np.arctan(U[0, 0]/ U[1, 0])

    e = Ellipse(mu, 2 * np.sqrt(5.991 * var[0]), 2 * np.sqrt(5.991 * var[1]), 
                angle=angle, facecolor='none', **kwargs)
    ax.add_artist(e)
    return e


def make_plot(data, mu, cov, scatter_color, ax):
    artists = []
    im1 = ax.scatter(data[:,0], data[:,1], s=20, c=scatter_color, alpha=0.5)
    artists.append(im1)
    for i in range(n_cluster):
        im2 = ax.scatter(mu[i, 0], mu[i, 1], marker="x", s=50, c=colors[i])
        artists.append(im2)
        im3 = draw_ellipse(mu[i], cov[i], ax, edgecolor=colors[i])
        artists.append(im3)
    return artists


def calc_L(data, mu, cov, pi):
    L = np.zeros((len(data), len(pi)))
    for i in range(len(pi)):
        L[:,i] = pi[i] * stats.multivariate_normal(mean=mu[i], cov=cov[i]).pdf(data)
    return L


#＃# サンプルデータ作成
n_cluster = 3
n_sample = 500
z_list = np.array([0.2, 0.5, 0.3])
colors = ["r", "g", "b"]

random.seed(18)

mu_list = stats.uniform(loc=-3, scale=6).rvs(size=(3, 2))
cov_list = stats.invwishart(df=3, scale=np.eye(2)).rvs(3)

# 各サンプルをどの分布から発生させるか
lab = random.choice(n_cluster, n_sample, p = z_list)
class_count = Counter(lab)

data = []
for i in range(n_cluster):
    y = stats.multivariate_normal(mean=mu_list[i], cov=cov_list[i]).rvs(class_count[i])
    data.append(y)
    
D = np.vstack(data)  
    

#＃# EMアルゴリズム
threshold = 1e-4
fig, ax = plt.subplots()
ax.set_xlim(-3, 6)
ax.set_ylim(-5, 4)
ims = []

## mu, Sigma, piの初期化
(xmin, ymin), (xmax, ymax) = D.min(axis=0), D.max(axis=0)

random.seed(0)
mu_x = random.uniform(xmin, xmax, n_cluster)
mu_y = random.uniform(ymin, ymax, n_cluster)
mu = np.c_[mu_x, mu_y]

cov = np.tile(0.1 * np.eye(2), (n_cluster, 1, 1))
pi = random.dirichlet(1000 * np.ones(n_cluster))

## E Step 
L = calc_L(D, mu, cov, pi)

# log likelihood
loglike = np.log(L.sum(axis=1)).sum()


for step in range(50):
    loglike_prev = loglike
    
    ## E Step 
    # responsibility
    gamma = L / L.sum(axis=1, keepdims=True)

    ## plot
    artists = make_plot(D, mu, cov, gamma, ax)
    ims.append(artists) 
    
    ## M Step
    # 各クラスターの実質的なデータ数
    n_k = gamma.sum(axis=0)

    # 負担率による重み付き和
    mu = np.zeros((n_cluster, 2))
    for k in range(n_cluster):
        mu[k] = (gamma[:,k, np.newaxis] * D).sum(axis=0) / n_k[k]

    # 共分散行列
    cov = np.zeros((n_cluster, 2, 2))
    for k in range(n_cluster):
        diff = D - mu[k]
        sigma = np.dot(diff.T, gamma[:, k, np.newaxis] * diff) / n_k[k]
        cov[k] = sigma

    # 混合係数
    pi = n_k / n_sample
    
    ## calculate log likelihood
    L = calc_L(D, mu, cov, pi)
    loglike = np.log(L.sum(axis=1)).sum()

    loglike_diff = loglike - loglike_prev
    
    ## plot
    artists = make_plot(D, mu, cov, gamma, ax)
    ims.append(artists) 
    
    print("decrease of log likelihood: {0:.4f}".format(loglike_diff))
    
    if loglike_diff < threshold:
        print(f"log likelihood converged after {step} iteration")
        break
        
# アニメーションを作成する。
plt.tight_layout()
anim = ani.ArtistAnimation(fig, ims, interval=100, repeat_delay=1000)
# anim.save('animation.gif', writer='imagemagick')