PRML第13章 隠れマルコフモデルの最尤推定 Python実装

第１３章では系列データを扱うモデルとして隠れマルコフモデルが紹介されています。今回はその隠れマルコフモデル(HMM:Hidden Markov Model)での最尤推定を実装します。パラメータを固定してフォーワード-バックワードアルゴリズムを使って隠れ状態を推定するということは何度かしたことがあるのですが、パラメータまで推定するということはしたことなかったのでいい機会なのでやってみることにしました。ちなみに、HMMはカルマンフィルタや粒子フィルタなどの基礎になっています。

# 隠れマルコフモデル
隠れマルコフモデルでは系列データ（例えば、コインを百回投げて出た裏表）を考えます。そして、それらの観測された系列データの背後には潜在変数（例えば、裏（表）が出やすいコインが使われた）があるとしています。
隠れマルコフモデルのグラフィカルモデルは下の図(PRML図13.5より)のようになっています。
![graphical_model.png](https://qiita-image-store.s3.amazonaws.com/0/148017/9f198a21-ce27-db45-b601-95b872dbc61b.png)
$\\{x_i\\}$を観測変数（コインの裏表）、$\\{z_i\\}$を潜在変数（コインの偏り）としています。潜在変数はひとつ前の潜在変数にのみ依存し（例：表が出やすいコインが使われた後は、また同じ表に偏ったコインが使われる）、観測変数はそのときの潜在変数にのみ依存しています（例：表に偏ったコインが使われたら、表が出やすい）。このように潜在変数がひとつ前の潜在変数にのみ依存している場合、一次マルコフ連鎖といいます。

数式を用いてもう少し細かく見ていきます。
潜在変数$z_n$が$z_{n-1}$にのみ依存していて、その関係を遷移確率行列$A$で表します。行列$A$のそれぞれの成分は潜在変数の遷移確率を表します。成分$A_{ij}$は$z_{n-1}$が状態$i$（$z_{n-1,i}=1$）のときに$z_n$の状態が$j$（$z_{n,j}=1$）になる確率となっています。ただし、潜在変数には一対K符号化法が用いられているとします。これらを数式にして書き下すと、

```math
p(z_n|z_{n-1},A) = \prod_{i=1}^K\prod_{j=1}^KA_{ij}^{z_{n-1,i},z_{nj}}
```
となります。
ただし、一番最初の潜在変数$z_1$にはその前の潜在変数が存在しないので、K次元の確率ベクトル$\pi$を用いて

```math
p(z_1|\pi) = \prod_{i=1}^K\pi_i^{z_{1i}}
```
とします。

観測変数$x$がD個の離散的な状態をとる離散変数（例：コインの裏表なら２つの状態がある離散変数）とした場合、観測変数の条件付き分布は

```math
p(x_n|z_n,\mu) = \prod_{i=1}^D\prod_{k=1}^K\mu_{ik}^{x_{ni}z_{nk}}
```
となります。ただし、$x$にも一対K符号化法が用いられています。


# HMMの最尤推定
上で紹介した隠れマルコフモデルでは３つのパラメータ$\pi,A,\mu$が登場してきました。ここでは、観測された系列データ$X = \\{x_n\\}_{n=1}^N$からそれら３つのパラメータを最尤推定します。これら３つのパラメータをひとまとめにして$\theta$と表すと、最大化すべき尤度関数は

```math
p(X|\theta) = \sum_Zp(X,Z|\theta)
```
となります。ただし、$Z=\\{z_n\\}_{n=1}^N$。Zの場合の数は、$z$の状態数KのN乗となって（例えば：コインの偏りが２通りあってコインを１００回投げた場合、$2^{100}\approx10^{30}$通り）、直接計算できません。そこで、EMアルゴリズムの枠組みで一次の隠れマルコフモデルであることを活かして計算量を減らす方法を使います。

EMアルゴリズムでは、以下の式の計算と$\theta$についての最大化を繰り返します。

```math
Q(\theta,\theta^{old}) = \sum_Zp(Z|X,\theta^{old})\ln p(X,Z|\theta)
```

## Mステップ
$\gamma(z_n)$を潜在変数$z_n$の周辺事後分布、$\xi(z_{n-1},z_n)$を連続した潜在変数の同時事後分布とします。

```math
\begin{align}
\gamma(z_n) &= p(z_n|X,\theta^{old})\\
\xi(z_{n-1},z_n) &= p(z_{n-1},z_n|\theta^{old})
\end{align}
```
これらを用いて、最大化すべき式を書き直すと、

```math
Q(\theta,\theta^{old}) = \sum_{k=1}^K\gamma(z_{1k})\ln\pi_k + \sum_{n=2}^N\sum_{j=1}^K\sum_{k=1}^K\xi(z_{n-1,j},z_{n,k})\ln A_{jk} + \sum_{n=1}^N\sum_{k=1}^K \gamma(z_{nk})\ln p(x_n|z_n,\mu)
```
これをパラメータ$\pi,A,\mu$について最大化すると、そのときのパラメータはラグランジュ乗数法を用いて

```math
\begin{align}
\pi_k &= {\gamma(z_{1k})\over\sum_{k=1}^K\gamma(z_{1j})}\\
A_{jk} &= {\sum_{n=2}^N\xi(z_{n-1,j},z_{nk})\over\sum_{l=1}^K\sum_{n=2}^N\xi(z_{n-1,j},z_{nl})}\\
\mu_{ik} &= {\sum_{n=1}^N\gamma(z_{nk})x_{ni}\over\sum_{n=1}^N\gamma(z_{nk})}
\end{align}
```
と求まります。

## Eステップ
このステップでは、Mステップで用いた潜在変数についての２つの事後確率分布$\gamma(z_n),\xi(z_{n-1},z_n)$を計算します。その方法はフォーワード-バックワードアルゴリズムと言われていて、名前の通り系列データに関して前からと後ろからの２つの計算をすることで効率的に事後確率分布を計算できます。

### フォーワード
$\alpha(z_1)=p(x_1,z_1)=p(x_1|z_1)p(z_1)$として、$\alpha(z_n)=p(x_1,\dots,x_n,z_n)$を再帰的に計算します。

```math
\begin{align}
\alpha(z_n) &= p(x_1,\dots,x_n,z_n)\\
&= \sum_{z_{n-1}}p(x_1,\dots,x_n,z_n,z_{n-1})\\
&= \sum_{z_{n-1}}p(x_1,\dots,x_n|z_n,z_{n-1})p(z_n|z_{n-1})p(z_{n-1})\\
&= \sum_{z_{n-1}}p(x_1,\dots,x_{n-1}|z_{n-1})p(x_n|z_n)p(z_n|z_{n-1})p(z_{n-1})\\
&= p(x_n|z_n)\sum_{z_{n-1}}\alpha(z_{n-1})p(z_n|z_{n-1})
\end{align}
```

### バックワード
$\beta(z_N)=p(x_N|z_N)$として、$\beta(z_n)=p(x_{n+1},\dots,x_N|z_n)$を再帰的に計算します。

```math
\begin{align}
\beta(z_n)&=p(x_{n+1},\dots,x_N|z_n)\\
 &= \sum_{z_{n+1}}p(z_{n+1},x_{n+1},\dots,x_N|z_n)\\
&= \sum_{z_{n+1}}p(x_{n+1},\dots,x_N|z_n,z_{n+1})p(z_{n+1}|z_n)\\
&= \sum_{z_{n+1}}p(x_{n+2},\dots,x_N|z_{n+1})p(x_{n+1}|z_{n+1})p(z_{n+1}|z_n)\\
&= \sum_{z_{n+1}}\beta(z_{n+1})p(x_{n+1}|z_{n+1})p(z_{n+1}|z_n)
\end{align}
```

### 事後確率分布
上のフォーワードとバックワードより

```math
\begin{align}
\gamma(z_n) &= {\alpha(z_n)\beta(z_n)\over p(X)}\\
\xi(z_{n-1},z_n) &= {\alpha(z_{n-1})p(x_n|z_n)p(z_n|z_{n-1})\beta(z_n)\over p(X)}
\end{align}
```
のように事後確率分布が求まります。

# コード

## import
いつもどおりnumpyとmatplotlibを使います。

```python
import matplotlib.pyplot as plt
import numpy as np
```

## 隠れマルコフモデル

```python
class HiddenMarkovModel(object):

    def __init__(self, n_states_hidden, n_states_observe):
        # 潜在変数の状態数
        self.n_states_hidden = n_states_hidden

        # 観測変数の状態数
        self.n_states_observe = n_states_observe

        # 初期潜在変数の分布のパラメータpiの初期化
        self.initial = np.ones(n_states_hidden) / n_states_hidden

        # 遷移確率行列Aの初期化
        self.transition = np.ones((n_states_hidden, n_states_hidden)) / (2 * n_states_hidden)
        self.transition += np.eye(n_states_hidden) * 0.5

        # 観測変数の分布のパラメータmuの初期化
        self.observation = np.random.rand(n_states_observe, n_states_hidden)
        self.observation /= np.sum(self.observation, axis=0, keepdims=True)

    # pi,A,muの最尤推定
    def fit(self, sequence, iter_max=100):
        # EMステップを繰り返す
        for i in xrange(iter_max):
            params = np.hstack((self.transition.ravel(), self.observation.ravel()))

            # Eステップ
            p_hidden, p_transition = self.expectation(sequence)

            # Mステップ
            self.maximization(sequence, p_hidden, p_transition)

            # 収束しているか確認
            if np.allclose(params, np.hstack((self.transition.ravel(), self.observation.ravel()))):
                break

    # Eステップ
    def expectation(self, sequence):
        N = len(sequence)
        forward = np.zeros(shape=(N, self.n_states_hidden))
        # alpha(z_1)=p(x_1,z_1)を計算
        forward[0] = self.initial * self.observation[sequence[0]]
        backward = np.zeros_like(forward)
        # beta(z_N)=p(x_N|z_N)を計算
        backward[-1] = self.observation[sequence[-1]]

        # フォーワード
        for i in xrange(1, len(sequence)):
            # PRML式(13.36)
            forward[i] = self.transition.dot(forward[i - 1]) * self.observation[sequence[i]]

        # バックワード
        for j in xrange(N - 2, -1, -1):
            # PRML式(13.38)
            backward[j] = (self.observation[sequence[j + 1]] * backward[j + 1]).dot(self.transition)

        # PRML式(13.33) 潜在変数z_nの事後確率分布gamma(z_n)を計算
        p_hidden = forward * backward
        p_hidden /= np.sum(p_hidden, axis=-1, keepdims=True)

        # PRML式(13.43) 連続した潜在変数の同時事後確率分布xi(z_{n-1},z_n)を計算
        p_transition = self.transition * (self.observation[sequence[1:]] * backward[1:])[:, :, None] * forward[:-1, None, :]
        p_transition /= np.sum(p_transition, axis=(1, 2), keepdims=True)

        return p_hidden, p_transition

    # Mステップ
    def maximization(self, sequence, p_hidden, p_transition):
        # PRML式(13.18) 初期潜在変数の分布のパラメータの更新
        self.initial = p_hidden[0] / np.sum(p_hidden[0])

        # PRML式(13.19) 遷移確率行列の更新
        self.transition = np.sum(p_transition, axis=0) / np.sum(p_transition, axis=(0, 2))
        self.transition /= np.sum(self.transition, axis=0, keepdims=True)

        # PRML式(13.23) 観測モデルのパラメータ更新
        x = p_hidden[:, None, :] * (np.eye(self.n_states_observe)[sequence])[:, :, None]
        self.observation = np.sum(x, axis=0) / np.sum(p_hidden, axis=0)
```

## データ

```python
# コイン投げのデータをつくる
def create_toy_data(sample_size=100):

    # 偏りに応じてコインを投げる
    def throw_coin(bias):
        if bias == 1:
            # 表が0.8の確率で出るコインを投げる
            return np.random.choice(range(2), p=[0.2, 0.8])
        else:
            # 裏が0.8の確率で出るコインを投げる
            return np.random.choice(range(2), p=[0.8, 0.2])

    # 始めのコインの偏り
    bias = np.random.uniform() > 0.5
    coin = []
    cheats = []
    for i in xrange(sample_size):
        coin.append(throw_coin(bias))
        cheats.append(bias)

        # 0.01の確率でコインの偏りをかえる
        bias = bias + np.random.choice(range(2), p=[0.99, 0.01])
        bias = bias % 2
    coin = np.asarray(coin)

    return coin, cheats
```

## 全体のコード

```python
import matplotlib.pyplot as plt
import numpy as np


class HiddenMarkovModel(object):

    def __init__(self, n_states_hidden, n_states_observe):
        self.n_states_hidden = n_states_hidden
        self.n_states_observe = n_states_observe
        self.initial = np.ones(n_states_hidden) / n_states_hidden
        self.transition = np.ones((n_states_hidden, n_states_hidden)) / (2 * n_states_hidden)
        self.transition += np.eye(n_states_hidden) * 0.5
        self.observation = np.random.rand(n_states_observe, n_states_hidden)
        self.observation /= np.sum(self.observation, axis=0, keepdims=True)

    def fit(self, sequence, iter_max=100):
        for i in xrange(iter_max):
            params = np.hstack((self.transition.ravel(), self.observation.ravel()))
            p_hidden, p_transition = self.expectation(sequence)
            self.maximization(sequence, p_hidden, p_transition)
            if np.allclose(params, np.hstack((self.transition.ravel(), self.observation.ravel()))):
                break

    def expectation(self, sequence):
        N = len(sequence)
        forward = np.zeros(shape=(N, self.n_states_hidden))
        forward[0] = self.initial * self.observation[sequence[0]]
        backward = np.zeros_like(forward)
        backward[-1] = self.observation[sequence[-1]]
        for i in xrange(1, len(sequence)):
            forward[i] = self.transition.dot(forward[i - 1]) * self.observation[sequence[i]]
        for j in xrange(N - 2, -1, -1):
            backward[j] = (self.observation[sequence[j + 1]] * backward[j + 1]).dot(self.transition)
        p_hidden = forward * backward
        p_hidden /= np.sum(p_hidden, axis=-1, keepdims=True)
        p_transition = self.transition * (self.observation[sequence[1:]] * backward[1:])[:, :, None] * forward[:-1, None, :]
        p_transition /= np.sum(p_transition, axis=(1, 2), keepdims=True)

        return p_hidden, p_transition

    def maximization(self, sequence, p_hidden, p_transition):
        self.initial = p_hidden[0] / np.sum(p_hidden[0])
        self.transition = np.sum(p_transition, axis=0) / np.sum(p_transition, axis=(0, 2))
        self.transition /= np.sum(self.transition, axis=0, keepdims=True)
        x = p_hidden[:, None, :] * (np.eye(self.n_states_observe)[sequence])[:, :, None]
        self.observation = np.sum(x, axis=0) / np.sum(p_hidden, axis=0)


def create_toy_data(sample_size=100):

    def throw_coin(bias):
        if bias == 1:
            return np.random.choice(range(2), p=[0.2, 0.8])
        else:
            return np.random.choice(range(2), p=[0.8, 0.2])

    bias = np.random.uniform() > 0.5
    coin = []
    cheats = []
    for i in xrange(sample_size):
        coin.append(throw_coin(bias))
        cheats.append(bias)
        bias = bias + np.random.choice(range(2), p=[0.99, 0.01])
        bias = bias % 2
    coin = np.asarray(coin)

    return coin, cheats


def main():
    coin, cheats = create_toy_data(200)

    hmm = HiddenMarkovModel(2, 2)
    hmm.fit(coin, 100)
    p_hidden, _ = hmm.expectation(coin)

    plt.plot(cheats)
    plt.plot(p_hidden[:, 1])
    for i in xrange(0, len(coin), 2):
        plt.annotate(str(coin[i]), (i - .75, coin[i] / 2. + 0.2))
    plt.ylim(-0.1, 1.1)
    plt.show()


if __name__ == '__main__':
    main()
```

# 結果
上のコードを走らせると下の図のような結果が得られることもあると思います。
![result.png](https://qiita-image-store.s3.amazonaws.com/0/148017/243b75f4-7c44-cfcd-e7d4-fd2e5a35ef62.png)
横軸はコイン投げが何回目なのか、縦軸は確率を表しています。青色の線は実際に用いられたコインの偏り(表が出やすい、裏が出やすいの２種類のコイン)を表していて、0のときは裏に偏るコイン、1のときは表に偏るコインが用いられています。緑の線は、そのときに投げられたコインが状態1（どっちのコインが状態１なのかはわかりません)である確率を表しています。グラフ中の0,1は実際に観測されたコインの裏表を表しています。ただし、すべて表示すると読みにくいので二回に一回だけ表示しています。

EMアルゴリズムは初期値によっては局所最適解にはまってしまうことがあるので、毎回上のような結果が得られるとは限りません。

# 終わりに
パラメータも変数にしてしまうと結構な頻度で悪い局所最適解にはまるので、最尤推定ではなく事前知識を導入して最大事後確率推定なんかができるといいのかもしれません。