混合モデルをEMアルゴリズムで解く&Rで実装する①

# はじめに

社会学や生物学では線形混合モデルと呼ばれるモデルが広く使われています。この線形混合モデルとは簡単にいうと一般的な線形モデルの拡張モデルであり、変量効果と呼ばれるパラメータを持つことが特徴としてあげられます(対する通常のパラメータは固定効果と呼ばれる)。変量効果とは簡単にいえば確率変数に従い生成されるパラメータであり、本当の意味でのパラメータではありません。真の意味のパラメータは変量効果が従う確率変数のパラメータであり、パラメータを規定するパラメータという解層構造を取っています。このパラメータを推定する方法はいくつかありますが、このような階層構造を取るが故に、EMアルゴリズムと相性が良いことが知られています。そこで今回は混合モデルの中でも生物学(育種)でよく使われているGBLUPにターゲットを絞り、EMアルゴリズムでパラメータを推定してみようと思います。まずは第一回としてEMアルゴリズムのGBLUPバージョンを導出していきます(かなり粗い計算をしているので間違いがあるかもしれません、ご了承ください)。
なお、私の所属している研究室ではprmlの輪読および実装を行っています。Qiitaにもまとめていますのでもしよろしければ参考にしてください([PRML　演習問題　解答集　まとめ](https://qiita.com/Lab_of_Biomet/items/15e38ca34fafa8176d89))

- [混合モデルとGBLUP](#混合モデルとgblup)
- [EMアルゴリズム](#emアルゴリズム)
- [GBLUPのEMアルゴリズム](#gblupのemアルゴリズム)


# 混合モデルとGBLUP
まずはじめに混合モデルとGBLUPの定義を記します。線形混合モデルは一般に次のように表されます。

```math
\displaylines{
\mathbf{y}=X\mathbf{b}+Z\mathbf{u}+\mathbf{e}\\
\mathbf{u}\sim N(\mathbf{0},G), \mathbf{e}\sim N(\mathbf{0},R)
}
```
ここで$\mathbf{y}$は従属変数、$\mathbf{b}$は固定効果、$\mathbf{u}$は変量効果、$\mathbf{e}$は誤差項を表しており、$\mathbf{b},\mathbf{e}$はそれぞれ$G,R$という分散共分散行列に従うと仮定します。ここで$G,R$が分かっている場合は、次の混合モデル方程式を解くことで$\mathbf{b},\mathbf{u}$に対する推定値(それぞれBLUE,BLUPという)を求めることができます

```math
\begin{bmatrix}
X^TR^{-1}X & X^TR^{-1}Z \\
Z^TR^{-1}X & Z^TR^{-1}Z+G^{-1}
\end{bmatrix}
\begin{bmatrix}
\mathbf{b}  \\
\mathbf{u}
\end{bmatrix}
=
\begin{bmatrix}
X^TR^{-1}y  \\
Z^TR^{-1}y
\end{bmatrix}
```
ただ一般にGとRは未知数であるためこれを解くにはEMアルゴリズムで解くなどの工夫が必要です。
次にGBLUPの定義ですが、次のように定義されます。変量効果と誤差の分散共分散行列に変化が入ります。
```math
\displaylines{
\mathbf{y}=X\mathbf{b}+\mathbf{u}+\mathbf{e}\\
\mathbf{u}\sim N(\mathbf{0},K\sigma^2_g), \mathbf{e}\sim N(\mathbf{0},I_n \sigma^2)
}
```
ここで$K$はゲノム関係行列と呼ばれゲノム情報を元に計算される個体間の遺伝的な距離を表します。このモデルにおいて真のパラメータは$\mathbf{b},\sigma_g,\sigma$の3つとなります。

# EMアルゴリズム
EMアルゴリズムとは尤度法によりパラメータを導出することが難しい場合に、潜在的な確率変数を導入し解を導出する逐次的アルゴリズムの一種です。
通常、最尤法では次のような対数尤度関数をパラメータ$\theta$に対して最大化します。
```math
L(\theta|D)=\log p(D|\theta)
```
EMアルゴリズムではこの対数尤度関数に対して潜在変数$u$を導入し、次の①②番を繰り返しながら真のパラメータを求めます

### ①期待値の計算(E step)

まず現段階におけるパラメータの推定値$\theta^{old}$とデータ$D$に条件づいた潜在変数$u$の確率分布$p(u|D,\theta^{old})$を求めます。次に潜在変数$u$とデータ$D$の同時分布の対数尤度関数$\log p(D,u|\theta)$の$p(u|D,\theta^{old})$に関するの期待値$\int \log p(D,u|\theta) p(u|D,\theta^{old}) du$を計算します。これは$\thetaと\theta^{old}$の関数となっていますのでこれを$L(\theta, \theta^{old})$とします。

### ②期待値の最大化(M step)
次に$L(\theta, \theta^{old})$を最大化する$\theta$を求め$\theta^{new}$とします。つまり
```math
\theta^{new}=\arg \max_{\theta} L(\theta,\theta^{old})
```
を求めるということです。
___
$\theta^{old}$を$\theta^{new}$で上書きするとそれに基づいてまた①②番を実行することができます。これを$\theta$が収束するまで繰り返せば真値にたどり着くことができる、というのがEMアルゴリズムです。

GBLUPの場合は$\theta=(\mathbf{b},\sigma_g,\sigma)^T$であり、潜在変数として$u$を用いればこのEMアルゴリズムを通して真のパラメータを求めることができます。次からは実際にGBLUPで使われているモデルに基づいてEMアルゴリズムの具体形を求めていきます。

# GBLUPのEMアルゴリズム

EMアルゴリズムを回すためにはまず$p(u|D,\theta)$と対数尤度関数$\log p(D,u|\theta)$の具体形を求める必要があります。先に$p(u|D,\theta)$を考えてみます。

**①$p(u|D,\theta)$の導出**

まず$p(u,D|\theta)=p(u|D,\theta)p(D,\theta)$および$p(u,D|\theta)=p(D|u,\theta)p(u|\theta)$が成り立つので次が言えます。
```math
\begin{align}
p(u|D,\theta)&=\frac{p(D|u,\theta)p(u|\theta)}{p(D,\theta)}\\
& \propto p(D|u,\theta)p(u|\theta)
\end{align}
```
ここで$p(D|u,\theta),p(u|\theta)$はそれぞれ以下の形を取ることが容易にわかります。
```math
p(D|u,\theta)=\frac{1}{\sqrt{(2\pi)^n|I_{n}\sigma^2|}}\exp \Bigl(-\frac{1}{2}(\mathbf{y}-X\mathbf{b}-\mathbf{u})^T(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b}-\mathbf{u})\Bigr)
```

```math
p(u|\theta)=\frac{1}{\sqrt{(2\pi)^n|K\sigma^2_g|}}\exp \Bigl(-\frac{1}{2}\mathbf{u}^T(K\sigma^2_g)^{-1}\mathbf{u}\Bigr)
```
よって
```math
p(u|D,\theta) \propto \exp \Bigl(-\frac{1}{2}(\mathbf{y}-X\mathbf{b}-\mathbf{u})^T(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b}-\mathbf{u})-\frac{1}{2}\mathbf{u}^T(K\sigma^2_g)^{-1}\mathbf{u}\Bigr)
```
さて、これは$\mathbf{u}$についての確率分布であるので指数関数の中身を$u$について整理していきます。行列の一般論から次が成り立つことが分かります。
```math
-\frac{1}{2}\mathbf{x}^TA\mathbf{x}+\mathbf{x}^T\mathbf{m}=-\frac{1}{2}(\mathbf{x}-A^{-1}\mathbf{m})^TA(\mathbf{x}-A^{-1}\mathbf{m})+\frac{1}{2}\mathbf{m}^TA\mathbf{m}
```
$p(u|D,\theta)$の指数の中身をこの左辺のように変形すると

```math
-\frac{1}{2}\mathbf{u}^T\big((I_{n}\sigma^2)^{-1}+(K\sigma^2_g)^{-1}\big)\mathbf{u}+\mathbf{u}^T(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b})
```
ここで$A^{-1}=(I_{n}\sigma^2)^{-1}+(K\sigma^2_g)^{-1},m=(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b})$とすると$p(u|D,\theta)$が次のように求まります。
```math
p(u|D,\theta)=\frac{1}{\sqrt{(2\pi)^n|A|}}\exp\Bigl(-\frac{1}{2}(\mathbf{u}-A\mathbf{m})^TA^{-1}(\mathbf{u}-A\mathbf{m})\Bigr)
```

**②$\log p(D,u|\theta)$の導出**
まず$p(D,u|\theta)=p(D|u,\theta)p(u|\theta)$より次がわかります。
```math
\log p(D,u|\theta)=\log p(D|u,\theta)+\log p(u|\theta)
```
ここで$p(D|u,\theta),p(u|\theta)$はすでに求まっているのでこれを代入すると

```math
\log p(D|u,\theta)=-\frac{1}{2}\log \Bigl((2\pi)^n|I_{n}\sigma^2|\Bigr)-\frac{1}{2}(\mathbf{y}-X\mathbf{b}-\mathbf{u})^T(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b}-\mathbf{u})
```
```math
\log p(u|\theta) = -\frac{1}{2}\log \Bigl((2\pi)^n|K\sigma^2_g|\Bigr)-\frac{1}{2}\mathbf{u}^T(K\sigma^2_g)^{-1}\mathbf{u}
```
これらは最終的に微分されるため定数項は適当にまとめることができます。それぞれ$C_1,C_2$とすると次を得ます。

```math
\log p(D|u,\theta)= -n\log \sigma -\frac{1}{2}(\mathbf{y}-X\mathbf{b}-\mathbf{u})^T(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b}-\mathbf{u})+C_1
```
```math
\log p(u|\theta)= -n\log \sigma_g-\frac{1}{2}\mathbf{u}^T(K\sigma^2_g)^{-1}\mathbf{u}+C_2
```
これらをまとめ、$C=C_1+C_2$とすると
```math
\log p(D,u|\theta)=-n\log \sigma -n\log \sigma_g -\frac{1}{2}(\mathbf{y}-X\mathbf{b}-\mathbf{u})^T(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b}-\mathbf{u})-\frac{1}{2}\mathbf{u}^T(K\sigma^2_g)^{-1}\mathbf{u}+C
```
先ほどと同様の変形をし、次を得ます。
```math
\log p(D,u|\theta)=-n\log \sigma -n\log \sigma_g -\frac{1}{2}\mathbf{u}^TA^{-1}\mathbf{u}+\mathbf{u}^Tm-\frac{1}{2}(\mathbf{y}-X\mathbf{b})^T(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b})+C
```
以上で$\log p(D,u|\theta)$の準備が整いました。
___

次に$L(\theta,\theta^{old})=\int \log p(D,u|\theta) p(u|D,\theta^{old}) du$を計算します。行列に対する一般論から$x\sim N(\mu,\Sigma)$と任意のベクトル$m$、行列$A$ついて次が成り立つことがわかります。
```math
\begin{align}
E[x^Tm]&=m^T\mu\\
E[x^TAx]&=tr(A\Sigma)+\mu^TA\mu
\end{align}
```
よって、次の計算結果を得ます。
```math
L(\theta,\theta^{old})=-n\log \sigma -n\log \sigma_g -\frac{1}{2}(\mathbf{y}-X\mathbf{b})^T(I_{n}\sigma^2)^{-1}(\mathbf{y}-X\mathbf{b})-\frac{1}{2}tr(A^{-1}A_{old})
-\frac{1}{2}(A_{old}\mathbf{m_{old}})^TA^{-1}(A_{old}\mathbf{m_{old}})+\mathbf{m}^TA_{old}\mathbf{m_{old}}+C
```
さて、ここでもう少し式の整理を行います。
```math
\begin{align}
(I_{n}\sigma^2)^{-1}&=\sigma^{-2}I_{n}\\
m&=\sigma^{-2}(\mathbf{y}-X\mathbf{b})\\
A^{-1}&=\sigma^{-2}I_{n}+\sigma^{-2}_gK^{-1}
\end{align}
```
以上より
```math
L(\theta,\theta^{old})=-n\log \sigma -n\log \sigma_g -\frac{1}{2\sigma^2}(\mathbf{y}-X\mathbf{b})^T(\mathbf{y}-X\mathbf{b})-\frac{1}{2}tr(A^{-1}A_{old})
-\frac{1}{2}(A_{old}\mathbf{m_{old}})^TA^{-1}(A_{old}\mathbf{m_{old}})+\mathbf{m}^TA_{old}\mathbf{m_{old}}+C
```
ここまでがE stepになります。次に$L(\theta,\theta^{old})$の最大化を行います。

```math
\frac{\partial}{\partial \mathbf{b}}L(\theta,\theta^{old})=\sigma^{-2}X^T(\mathbf{y}-X\mathbf{b_new})-\sigma^{-2}X^T(A_{old}m_{old})=0
```
これを整理して
```math
\mathbf{b_{new}}=(X^TX)^{-1}X^T(\mathbf{y}-A_{old}\mathbf{m_{old}})
```
次に$\sigma_g$について微分を行います。やや微分が大変なのでまず項ごとに微分を考えることにします。
まず$tr(A^{-1}A_{old})$は次のように変形をすることができます。
```math
\begin{align}
tr(A^{-1}A_{old})&=tr(\sigma^{-2}A_{old}+\sigma^{-2}_gK^{-1}A_{old})\\
&=tr(\sigma^{-2}A_{old})+tr(\sigma^{-2}_gK^{-1}A_{old})\\
&=\sigma^{-2}tr(A_{old})+\sigma^{-2}_gtr(K^{-1}A_{old})
\end{align}
```
よって、
```math
\frac{\partial}{\partial \sigma_g}tr(A^{-1}A_{old})=-\frac{2}{\sigma_g^3}tr(K^{-1}A_{old})
```
次に$(A_{old}\mathbf{m_{old}})^TA^{-1}(A_{old}\mathbf{m_{old}})$は次のように変形することができます。
```math
\begin{align}
(A_{old}\mathbf{m_{old}})^TA^{-1}(A_{old}\mathbf{m_{old}})&=(A_{old}\mathbf{m_{old}})^T(\sigma^{-2}I_{n}+\sigma^{-2}_gK^{-1})(A_{old}\mathbf{m_{old}})\\
&=\sigma^{-2}(A_{old}\mathbf{m_{old}})^T(A_{old}\mathbf{m_{old}})+\sigma^{-2}_g(A_{old}\mathbf{m_{old}})^TK^{-1}(A_{old}\mathbf{m_{old}})
\end{align}
```
従って
```math
\frac{\partial}{\partial \sigma_g}(A_{old}\mathbf{m_{old}})^TA^{-1}(A_{old}\mathbf{m_{old}})=-\frac{2}{\sigma_g^3}(A_{old}\mathbf{m_{old}})^TK^{-1}(A_{old}\mathbf{m_{old}})
```
以上より
```math
\frac{\partial}{\partial \sigma_g}L(\theta,\theta^{old})=-\frac{n}{\sigma_{g_{new}}}+\frac{1}{\sigma_{g_{new}}^3}tr(K^{-1}A_{old})+\frac{1}{\sigma_{g_{new}}^3}(A_{old}\mathbf{m_{old}})^TK^{-1}(A_{old}\mathbf{m_{old}})=0
```
これを$\sigma_{g_{new}}$について解くと
```math
\sigma_{g_{new}}^2=\frac{1}{n}\Bigl(tr(K^{-1}A_{old})+(A_{old}\mathbf{m_{old}})^TK^{-1}(A_{old}\mathbf{m_{old}})\Bigr)
```


最後に$\sigma$について微分を行います。
```math
\frac{\partial}{\partial \sigma}L(\theta,\theta^{old})=-\frac{n}{\sigma_{new}}+\frac{1}{\sigma_{new}^3}(\mathbf{y}-X\mathbf{b})^T(\mathbf{y}-X\mathbf{b})+\frac{1}{\sigma_{new}^3}tr(K^{-1}A_{old})+\frac{1}{\sigma_{new}3}(A_{old}\mathbf{m_{old}})^TK^{-1}(A_{old}\mathbf{m_{old}})-\frac{2}{\sigma_{new}^3}(\mathbf{y}-X\mathbf{b})^TA_{old}\mathbf{m_{old}}=0
```
ここで$\mathbf{b}$が出てきますが、これは$b_{new}$を指しています。これを整理すると
```math
\sigma_{new}^2=\frac{1}{n}\Bigl((\mathbf{y}-X\mathbf{b_{new}})^T(\mathbf{y}-X\mathbf{b_{new}})+tr(K^{-1}A_{old})+(A_{old}\mathbf{m_{old}})^TK^{-1}(A_{old}\mathbf{m_{old}})-2(\mathbf{y}-X\mathbf{b_{new}})^TA_{old}\mathbf{m_{old}}\Bigr)
```

以上より$L(\theta,\theta^{old})$を最大化する$\theta_{new}$について次の三式を得ることができました。
```math
\mathbf{b_{new}}=(X^TX)^{-1}X^T(\mathbf{y}-A_{old}\mathbf{m_{old}})
```
```math
\sigma_{g_{new}}^2=\frac{1}{n}\Bigl(tr(K^{-1}A_{old})+(A_{old}\mathbf{m_{old}})^TK^{-1}(A_{old}\mathbf{m_{old}})\Bigr)
```
```math
\sigma_{new}^2=\frac{1}{n}\Bigl((\mathbf{y}-X\mathbf{b_{new}})^T(\mathbf{y}-X\mathbf{b_{new}})+tr(K^{-1}A_{old})+(A_{old}\mathbf{m_{old}})^TK^{-1}(A_{old}\mathbf{m_{old}})-2(\mathbf{y}-X\mathbf{b_{new}})^TA_{old}\mathbf{m_{old}}\Bigr)
```
復習ですがここで$A_{old},\mathbf{m_{old}}$はそれぞれ次を指します。
```math
A_{old}=(\sigma_{old}^{-2}I_{n}+\sigma^{-2}_{g_{old}}K^{-1})^{-1}
```
```math
\mathbf{m_{old}}=\sigma_{old}^{-2}(\mathbf{y}-X\mathbf{b_{old}})
```

ここまでわかってしまえばあとは$\mathbf{b},\sigma,\sigma_g$に対して適当に初期を定めて$A_{old},\mathbf{m_{old}}$を計算して、それをもとに$\mathbf{b_{new}},\sigma_{new},\sigma_{g_{new}}$を計算して、次にこれを$\mathbf{b_{old}},\sigma_{old},\sigma_{g_{old}}$に代入してそれをもとに$A_{old},\mathbf{m_{old}}$を計算して…を収束するまで繰り返せば求めたいパラメータが求まるというわけです。

次回はこれを用いて実際にGBLUPをRで実装していきます。

[混合モデルをEMアルゴリズムで解く&Rで実装する②](https://qiita.com/KimuraSoh/items/11aeaa24d15b8bf56849)