EM算应用：两硬币、三硬币

参考了很多资料弄明白了EM算法在两硬币和三硬币中的算法以及代码。

0 EM算法

关于EM算法定义、使用场景、EM算法推导、收敛性证明直接看《统计学习方法》即可。
EM算法应用于有隐变量的场景。例如投掷2枚硬币，每次记录抛掷的硬币是A还是B，记录是正面还是反面。根据投掷结果计算2枚硬币的正面的概率是很好计算的。如果在记录过程中，没有记录抛掷的硬币是A还是B，这就是包含隐变量。这个时候没有解析解，只能通过迭代的方法求解。这也就是EM算法发挥作用的地方。

1 三硬币模型

有三枚硬币A，B，C，先投掷A，如果是正面就投掷B，如果是反面就投掷C，若我们只能观测到最后的投掷结果（B或者C的结果）而不能知道投掷的过程，如何估算三颗硬币的正面率？

1.1 定义

• 参数$\theta =(\pi ,p,q)$,$\pi ,p,q$分别表示硬币A,B,C是正面的概率。
• 观测到了n次投掷结果，记为$Y=(y_1,y_2,...y_n)$，每次投掷10次，$y_i\in [0,10]$,表示本次试验中硬币正面朝上的次数
• 隐变量设为$Z=(z_1,z_2,...z_n)$，$z_i=0$表示硬币A正面朝上，将选择硬币B投掷；$z_i=1$表示硬币A背面朝上，将选择硬币C投掷。

1.2 推导

1.2.1 E步

EM算法的E步是为隐变量建期望。

$$\begin{gathered} Q(\theta ,{\theta }^{(i)})={\sum }_{z}P(z|Y,{\theta }^{(i)})logP(Y,z|\theta ) \\ ={\sum }_{j=1}^n{\sum }_{k=1}^{K}P(z=l_k|y_j,{\theta }^{(i)})logP(y_j,z=l_k|\theta ) \\ ={\sum }_{j=1}^n[P(z=0|y_j,{\theta }^{(i)})logP(y_j,z=0|\theta )+P(z=1|y_j,{\theta }^{(i)})logP(y_j,z=1|\theta )] \end{gathered}$$
这里我们认为z的取值有k个。

我们先看下$P(z=0|y_j,{\theta }^{(i)})$应该怎么计算，这是一个条件概率，可以转为联合概率与全概率相除。
$$\begin{gathered} P(z=0|y_j,{\theta }^{(i)})=\frac{P(z=0,y_j|{\theta }^{(i)})}{{\sum }_{z}P(z=z,y_j|{\theta }^{(i)})} \\ =\frac{P(z=0,y_j|{\theta }^{(i)})}{P(z=0,y_j|{\theta }^{(i)})+P(z=1,y_j|{\theta }^{(i)})} \\ =\frac{{\pi }^{(i)}(p^{(i)}{)}^{(y_j)}(1-p^{(i)}{)}^{10-y_j}}{{\pi }^{(i)}(p^{(i)}{)}^{(y_j)}(1-p^{(i)}{)}^{10-y_j}+(1-{\pi }^{(i)})(q^{(i)}{)}^{(y_j)}(1-q^{(i)}{)}^{10-y_j}} \\ ={\mu }_j^{(i+1)} \end{gathered}$$

10-$y_j$表示硬币反面朝上的次数。在有些公式中写的是1-$y_j$。这是因为要解决的问题是投一次硬币的问题。或者说一次试验只投一次硬币。

我们再看$P(y_j,z=0|\theta )$的计算方式，它表示用硬币B，并且正面朝上的次数为$y_j$。
$P(y_j,z=0|\theta )=\pi p^{y_j}(1-p{)}^{10-y_j}$
$P(y_j,z=1|\theta )=(1-\pi )q^{y_j}(1-q{)}^{10-y_j}$

那么我们把这三个式子代入到Q函数中得到：
$$\begin{gathered} Q(\theta ,{\theta }^{(i)})={\sum }_{j=1}^n[{\mu }_j^{(i+1)}log(\pi p^{y_j}(1-p{)}^{10-y_j})+(1-{\mu }_j^{(i+1)})log((1-\pi )q^{y_j}(1-q{)}^{10-y_j}) \\ ={\sum }_{j=1}^n[{\mu }_j^{(i+1)}(log\pi +y_{j}log+(10-y_j)log(1-p))+(1-{\mu }_j^{(i+1)})(log(1-\pi )+y_{j}logq+(10-y_j)log(1-q))] \end{gathered}$$

这样就将Q函数中的隐变量消除了，用模型参数θ来表示

1.2.1 M步

EM算法的M步是求下一轮参数。
有了Q函数，分别对θ中的具体参数求导，推导出参数更新公式。

$\pi =\frac{1}{n}{\sum }_{j=1}^n{\mu }_j^{(i+1)}$
$p=\frac{{\sum }_{j=1}^n{\mu }_j^{(i+1)}{y}_j}{{\sum }_{j=1}^{n}10\ast {\mu }_j^{(i+1)}}$

$q=\frac{{\sum }_{j=1}^n(1-{\mu }_j^{(i+1)})y_j}{{\sum }_{j=1}^{n}10(1-\ast {\mu }_j^{(i+1)})}$

1.3 代码

class EM3Coins:
    def __init__(self, thetas, data):
        self.thetas = thetas
        self.data = data

    def em_algo(self,  max_iter=30, eps=1e-3):
        pi, p, q = self.thetas
        pi = pi[:, None]
        thetas = np.array([p, q])
        ll_old = -np.infty
        n = len(self.data)
        sum = np.sum(self.data[0])
        for i in range(max_iter):
            like = pi * np.array([np.prod(np.power(theta, self.data), axis=1) for theta in thetas])
            p_coin = like / np.sum(like, axis=0)  # 2x5
            new_pi = 1/n * np.sum(p_coin, axis=1)
            p_coin_data = self.data[:, 0].dot(p_coin.T)  # 2x5
            new_thetas = p_coin_data / (sum * np.sum(p_coin, axis=1))
            new_thetas = np.array([[theta, 1 - theta] for theta in new_thetas])
            ll_new = np.sum(p_coin * np.log(like))
            if np.abs(ll_new - ll_old) < eps:
                break
            ll_old = ll_new
            thetas = new_thetas
            pi = new_pi[:,None]
        return pi.T, thetas
def main2():
    observed_data = np.array([(5, 5), (9, 1), (8, 2), (4, 6), (7, 3)])
    theta = np.array([[0.5, 0.5], [0.8, 0.2], [0.6, 0.4]])
    em = EM3Coins(thetas=theta, data=observed_data)
    pi, thetas = em.em_algo()
    print(pi)
    print(thetas)

最后得到pi=[[0.51950478 0.48049522]],p=[0.79406571 0.20593429],q=[0.51505001 0.48494999]

2 两硬币

假设现在有两个硬币A和B，我们想要知道两枚硬币各自为正面的概率即模型的参数。我们先随机从A,B中选一枚硬币，然后扔10次并记录下相应的结果，H代表正面T代表反面。对以上的步骤重复进行5次。如果在记录的过程中我们记录下来每次是哪一枚硬币（即知道每次选的是A还是B），那可以直接根据结果进行估计（见下图a）。

但是如果数据中没记录每次投掷的硬币是A还是B（隐变量），只观测到5次循环共50次投币的结果，这时就没法直接估计A和B的正面概率。这时就该轮到EM算法大显身手了，EM算法特别适用于这种含有隐变量的参数求解问题（见下图b）。

先初始化输入参数，如上图1步给了一个初始值0.6（A硬币正面的概率），0.5（B硬币正面的概率）。接下来先进行E步（对隐变量求期望），如上图2步：以第一条数据为例，5H5T，为A的概率为 $p_A=0.6^5\times 0.4^5$ ,为B的概率 $p_B=0.5^5\times 0.5^5$ ，归一化后得P(A)=0.45,P(B)=0.55,剩下几条数据同理可得。而后通过M-step可计算重新迭代的概率值。如上图第一次迭代后
，循环上面的E、M步骤直至收敛我们就可以得到最终的答案，如上图进过10次迭代后得到了最终的结果。

用术语描述。

2.1 定义

• 参数$\theta =(p,q)$,$p,q$分别表示硬币A,B是正面的概率。
• 观测到了n次投掷结果，记为$Y=(y_1,y_2,...y_n)$，每次投掷10次，$y_i\in [0,10]$,表示本次试验中硬币正面朝上的次数
• 隐变量设为$Z=(z_1,z_2,...z_n)$，$z_i=0$表示选择硬币A投掷；$z_i=1$表示选择硬币B投掷。

2.2 推导

2.2.1 E步

EM算法的E步是为隐变量建期望。

$$\begin{gathered} Q(\theta ,{\theta }^{(i)})={\sum }_{z}P(z|Y,{\theta }^{(i)})logP(Y,z|\theta ) \\ ={\sum }_{j=1}^n{\sum }_{k=1}^{K}P(z=l_k|y_j,{\theta }^{(i)})logP(y_j,z=l_k|\theta ) \\ ={\sum }_{j=1}^n[P(z=0|y_j,{\theta }^{(i)})logP(y_j,z=0|\theta )+P(z=1|y_j,{\theta }^{(i)})logP(y_j,z=1|\theta )] \end{gathered}$$
这里我们认为z的取值有k个。

我们先看下$P(z=0|y_j,{\theta }^{(i)})$应该怎么计算，这是一个条件概率，可以转为联合概率与全概率相除。
$$\begin{gathered} P(z=0|y_j,{\theta }^{(i)})=\frac{P(z=0,y_j|{\theta }^{(i)})}{{\sum }_{z}P(z=z,y_j|{\theta }^{(i)})} \\ =\frac{P(z=0,y_j|{\theta }^{(i)})}{P(z=0,y_j|{\theta }^{(i)})+P(z=1,y_j|{\theta }^{(i)})} \\ =\frac{(p^{(i)}{)}^{(y_j)}(1-p^{(i)}{)}^{10-y_j}}{(p^{(i)}{)}^{(y_j)}(1-p^{(i)}{)}^{10-y_j}+(q^{(i)}{)}^{(y_j)}(1-q^{(i)}{)}^{10-y_j}} \\ ={\mu }_j^{(i+1)} \end{gathered}$$

10-$y_j$表示硬币反面朝上的次数。在有些公式中写的是1-$y_j$。这是因为要解决的问题是投一次硬币的问题。或者说一次试验只投一次硬币。

我们再看$P(y_j,z=0|\theta )$的计算方式，它表示用硬币A，并且正面朝上的次数为$y_j$。
$P(y_j,z=0|\theta )=p^{y_j}(1-p{)}^{10-y_j}$
$P(y_j,z=1|\theta )=q^{y_j}(1-q{)}^{10-y_j}$

那么我们把这三个式子代入到Q函数中得到：
$$\begin{gathered} Q(\theta ,{\theta }^{(i)})={\sum }_{j=1}^n[{\mu }_j^{(i+1)}log(p^{y_j}(1-p{)}^{10-y_j})+log((1-\pi )q^{y_j}(1-q{)}^{10-y_j}) \\ ={\sum }_{j=1}^n[{\mu }_j^{(i+1)}(y_{j}log+(10-y_j)log(1-p))+(1-{\mu }_j^{(i+1)})(y_{j}logq+(10-y_j)log(1-q))] \end{gathered}$$

这样就将Q函数中的隐变量消除了，用模型参数θ来表示

2.2.1 M步

EM算法的M步是求下一轮参数。
有了Q函数，分别对θ中的具体参数求导，推导出参数更新公式。过程不再详细写了，参考三硬币模型，应该可以推出来。
$p=\frac{{\sum }_{j=1}^n{\mu }_j{y}_j}{{\sum }_{j=1}^{n}10{\mu }_j}$
$q=\frac{{\sum }_{j=1}^n(1-{\mu }_j)y_j}{{\sum }_{j=1}^{n}10(1-{\mu }_j)}$

2.3 代码

class EM2Coins:

    def __init__(self, thetas, data):
        self.thetas = thetas
        self.data = data  # 5x2

    def em_algo(self, max_iter=30, eps=1e-3):
        ll_old = -np.infty
        sum = np.sum(self.data[0])
        for i in range(max_iter):
            like = np.array([np.prod(np.power(theta, self.data), axis=1)for theta in self.thetas])  # 2x5
            p_coin = like / np.sum(like, axis=0)  # 2x5,表示概率 p_coin[i][j]表示第j次试验选择硬币i的概率
            p_coin_data = self.data[:, 0].dot(p_coin.T)
            new_thetas = p_coin_data / (sum * np.sum(p_coin, axis=1))
            # 这里不够优雅
            new_thetas = np.array([[theta, 1 - theta] for theta in new_thetas])

            #p_coin_data = np.array([p[:, None] * self.data for p in p_coin])
            #new_thetas = np.array([pd.sum(0) / pd.sum() for pd in p_coin_data])
            #print(new_thetas)

            ll_new = np.sum(p_coin * np.log(like))
            if np.abs(ll_new - ll_old) < eps:
                break
            ll_old = ll_new
            self.thetas = new_thetas
        return self.thetas
 def main():
    observed_data = np.array([(5, 5), (9, 1), (8, 2), (4, 6), (7, 3)])
    theta = np.array([[0.4, 0.6], [0.6, 0.4], [0.5, 0.5]])
    em = EM2Coins(thetas=theta, data=observed_data)
    theta = em.em_algo()
    print(theta)

最终结果：p=[0.7967724 0.2032276 ], q=[0.51961361 0.48038639]

3 参考链接

为了弄明白代码中的每一步的含义，为了弄清楚E和M分别做了什么事情，参考了很多内容。主要参考内容是《统计学习方法》、《机器学习公式推导与代码实现》、知乎链接1、链接2