「数式を numpy に落とし込むコツ」を HMM に当てはめてみる

数式をnumpyに落としこむコツ

View more presentations from Shuyo Nakatani

という発表を Tokyo.SciPy #2 でさせてもらったのだが、発表＆資料作成の時間の関係で、実際に数式を解釈する例を2つしか入れられなかったのが残念なところ。
今、社内 PRML 読書会で 13章の隠れマルコフをやっつけていて、その Baum-Welch の更新式がちょうどいい題材になっていることに気付いたので、ここで取り上げてみる。

$\alpha({\boldsymbol z}_n)=p({\boldsymbol x}_n|{\boldsymbol z}_n)\sum_{{\boldsymbol z}_{n-1}}\alpha({\boldsymbol z}_{n-1})p({\boldsymbol z}_n|{\boldsymbol z}_{n-1})$ 　 (PRML 式 13.36)

結構複雑な印象のある数式だが、こいつも資料の流れに従えば簡単に実装できてしまうことを見ていこう。

数式を読み解く
数式を書き換える
numpy に「逐語訳」する

というわけでまず「読み解き」だが、これが一番重要なステップ。
特に今回の式の場合は $\alpha({\boldsymbol z}_n)$ , $p({\boldsymbol x}_n|{\boldsymbol z}_n)$ , $p({\boldsymbol z}_n|{\boldsymbol z}_{n-1})$ の正体をちゃんと見極めておかないといけない。
「正体」とは、つまりスカラーかベクトルか行列か、ベクトルや行列ならその次数は何か、そしてその要素はどのように書けるか、という点だ。

まず $\alpha({\boldsymbol z}_n)$ は次のように定義されている。

$\alpha({\boldsymbol z}_n)=p({\boldsymbol x}_1,\cdots,{\boldsymbol x}_n,{\boldsymbol z}_n)$ 　 (PRML 式 13.34)

x_n たちは観測値で固定されていて、z_n は 1-of-K 表現で K 通りの値を取り得る。
というわけで $\alpha({\boldsymbol z}_n)$ は K 次のベクトルであることわかる。
z_n は本来 1-of-K 表現だが、z_n の k 番目だけが 1 であることを "z_n=k" と書くことにすると、

$\alpha({\boldsymbol z}_n)_k = p({\boldsymbol x}_1,\cdots,{\boldsymbol x}_n,{\boldsymbol z}_n=k) \hspace{4ex} (k=1,\cdots,K)$

となる。
alpha(z_n) なんて書き方では、とてもベクトルには見えないが実はベクトルだ、と言っても紛らわしいだけなので、今後 ${\boldsymbol \alpha}^{(n)}=\alpha({\boldsymbol z}_n)$ と書くことにする。その k 番目の要素は $\alpha_k^{(n)}$ と書けて、わかりやすい。

次に $p({\boldsymbol z}_n|{\boldsymbol z}_{n-1})$ の定義は

$p({\boldsymbol z}_n|{\boldsymbol z}_{n-1},{\boldsymbol A})=\prod_{k=1}^K\prod_{j=1}^K A_{jk}^{z_{n-1,j}z_{nk}}$ 　 (PRML 式 13.7)

と書かれている。
これは 1-of-K 表現のせいで多項分布がややこしく表現されてしまっているだけで、上の記法を許せば、以下のように書ける。

$p({\boldsymbol z}_n=k|{\boldsymbol z}_{n-1}=j,{\boldsymbol A})=A_{jk} \hspace{4ex} (j=1,\cdots,K,\;k=1,\cdots,K)$

つまり $p({\boldsymbol z}_n|{\boldsymbol z}_{n-1},{\boldsymbol A})$ は K×K 行列であったことがわかる。

最後に $p({\boldsymbol x}_n|{\boldsymbol z}_n)$ だが、今回は出力分布に多項分布を使うことにすると、その定義は (PRML 式 13.22) に書かれているが、それは上の p(z_n|z_{n-1}) と全く同じ形なので、同様に次の形になることがわかる。

$p({\boldsymbol x}_n=i|{\boldsymbol z}_n=k)=\mu_{ik} \hspace{4ex} (i=1,\cdots,D,\;k=1,\cdots,K)$

ここから $p({\boldsymbol x}_n|{\boldsymbol z}_n)$ は D×K 行列だったわけだ。
どうだろう、ここまで読み解いただけでも、最初の式がすでにだいぶ簡単に見えてきたのではないだろうか。

読み解きステップが終わったので、数式を書き換えよう。
資料から、書き換え方針について書いたページを引用してみる。

ベクトル・行列の積の数式は、このように要素の数式に書き換えて、上の3種類の形に当てはめてあげたらいい。
元の数式を要素の形に書き換えるには、まず左辺が K 次のベクトルだったことを思い出し、その k 番目の要素 $\alpha_k^{(n)}$ を求める式を立てることを考える。このとき $\alpha_k^{(n)} = p({\boldsymbol x}_1,\cdots,{\boldsymbol x}_n,{\boldsymbol z}_n=k)$ だったので、z_n=k を右辺に当てはめる。
次に z_{n-1} で和を取っているところがある。これは 1 から K まで動くということなので、z_{n-1}=j とおいて、j が 1 から K まで動くことにしよう。もちろん z_{n-1}=j も右辺に当てはめてあげる。
すると次の式が得られた。

$\alpha_k^{(n)}=p({\boldsymbol x}_n|{\boldsymbol z}_n=k)\sum_{j=1}^K \alpha_j^{(n-1)}p({\boldsymbol z}_n=k|{\boldsymbol z}_{n-1}=j) \hspace{4ex} (k=1,\cdots,K)$

x_n は観測値なので、今その与えられた値を i と書くことにすれば、p(x|z) たちの定義から、この式はさらに次のようになる。

$\alpha_k^{(n)}=\mu_{ik}\sum_{j=1}^K \alpha_j^{(n-1)}A_{jk} \hspace{4ex} (k=1,\cdots,K)$

さあ、どんどん簡単になってきた。ここで上の書き換え方針を思いだそう。
今、数式の中でΣがわざとらしいくらい目立っているので、2番目の行列積の形に当てはめたくなる。

行列積にするには、和を取っているインデックスの j が真ん中でペアになるようにすればいいんだけど、ベクトルは K×1 次元の行列とみなすんだったことを忘れずに(資料の他のページにはちゃんと書いてあるよ)。
というわけで、転置をうまく使って添え字の順番を入れ替えることで

$\alpha_k^{(n)}=\mu_{ik}\sum_{j=1}^K ({\boldsymbol A}^T)_{kj}\alpha_j^{(n-1)} \hspace{4ex} (k=1,\cdots,K)$

となる。
$\sum_{j=1}^K ({\boldsymbol A}^T)_{kj}\alpha_j^{(n-1)}$ がまさに行列積を表しているので、 ${\boldsymbol b}={\boldsymbol A}^T{\boldsymbol \alpha}^{(n-1)}$ とおけば、 b は K 次のベクトルで、 $b_k=\sum_{j=1}^K ({\boldsymbol A}^T)_{kj}\alpha_j^{(n-1)}$ であることがわかる。
上の式の該当部分を b_k に置き換えると、さらに簡単になる。

$\alpha_k^{(n)}=\mu_{ik}b_k \hspace{4ex} (k=1,\cdots,K)$

これを上の方針と見比べると、全ての項が k という一つのインデックスにしか依存していない(i と n は固定)ので、今度は要素積の形をしていることがわかる。
要素積は数式の記号がないので、とりあえず * と書くことにすると、最初の数式は実は次のように書き換えられることがわかった。

${\boldsymbol \alpha}^{(n)} = {\boldsymbol \mu}_i * ({\boldsymbol A}^T{\boldsymbol \alpha}^{(n-1)})$

これを実装するのは、もうかんたん。alpha, A, mu がそれぞれ適当なベクトルや行列であれば、

alpha[n] = mu[i,:] * numpy.dot(A.T, alpha[n-1])

でおしまい。

読み解きと書き替えは特に numpy に限った話ではないし、この「逐語訳」はほとんどの行列ライブラリでできるレベル。
もちろん R でも簡単に実装できる。

βの更新式 (PRML 式 13.38) も同様に書き換えれば、やはり誰でも実装できるので、練習問題でやってみるとおもしろいかも。