EM算法与变分推断 —— 数学推导

---

EM算法

对于概率图模型中包含有隐变量的情况，可以使用EM算法进行参数估计。隐变量是指不可观测的变量，但其参与到了样本的生成过程。例如在混合高斯模型中，样本$x$的生成过程为首先确定其所属的类别$z$，之后根据其类别选择相应的高斯分布$\mathcal{N}({\mu }_z,{\sigma }_k)$，生成样本$x$。在该生成过程中，样本所属的类别$z$即为一个隐变量。

本文综合了一些相关资料，主要聚焦于EM算法及变分推断的推导过程。

EM算法推导 方法1

对于EM算法，可从边际似然函数的角度出发进行推导：在包含隐变量的图模型中，模型参数为$\theta$，则样本$x$的边际似然函数为：
$$p(x;\theta )=\sum _{z}p(x,z;\theta )\text{\hspace{1em}(1)}$$

则对于包含$N$个样本的训练集 D = { x ( 1 ) , x ( 2 ) , . . . , x ( N ) } D=\{x^{(1)},x^{(2)},...,x^{(N)}\} D={x(1),x(2),...,x(N)}，训练集$D$的对数边际似然函数为：
$$\begin{array}{rl}L(D;\theta )& =\frac{1}{N}\sum _{n=1}^N\log p(x^{(n)};\theta )\\ & =\frac{1}{N}\sum _{n=1}^N\log \sum _{z}p(x^{(n)},z;\theta )\end{array}\text{\hspace{1em}(2)}$$

我们的最终目标是希望最大化整个数据集的对数边际似然函数$L(D;\theta )$，然而(2)式中涉及分布$p(x,z;\theta )$的求和（或积分），难以计算。因此可以引入一个关于隐变量$z$的近似分布$q(z)$，称为变分函数，并利用其对边际似然函数进行推导：
$$\begin{array}{rl}\log p(x;\theta )& =\log \sum _{z}p(x,z;\theta )\\ & =\log \sum _{z}q(z)\frac{p(x,z;\theta )}{q(z)}\\ & \ge \sum _{z}q(z)\log \frac{p(x,z;\theta )}{q(z)}\\ & =ELBO(q,x;\theta )\end{array}\text{\hspace{1em}(3)}$$上述推导中运用了Jensen不等式，即对于凹函数$g(x)$，有$g(E[x])\ge E[g(x)]$；$ELBO(q,x;\theta )$为$\log p(x;\theta )$的下界，称为证据下界 (Evidence LowerBound)

回到问题本身，我们的目标是最大化$\log p(x;\theta )$，而由(3)式可知其下界为$ELBO(q,x;\theta )$，因此EM算法的思想是寻找一个最优的近似分布$q^{\ast }(z)$，使$ELBO(q^{\ast },x;\theta )$达到当前参数$\theta$固定情况下的最大值，然后改变参数$\theta$使得$ELBO(q,x;\theta )$在当前$q^{\ast }(z)$固定的情况下达到最大，之后不断地重复这两个步骤。

对于如何寻找最优的$q^{\ast }(z)$：根据Jensen不等式的等号成立条件可知，对于(3)式，仅当$q(z)=p(z|x;\theta )$时等号成立，此时$ELBO(q,x;\theta )$达到最大值$\log p(x;\theta )$

最终EM算法的步骤如下：
E-step：固定参数${\theta }_t$，即固定$\log p(x;{\theta }_t)$，令$q_{t+1}(z)=p(z|x;{\theta }_t)$，使得$\log p(x;{\theta }_t)=ELBO(q_{t+1},x;{\theta }_t)$
M-step：固定$q_{t+1}(z)$，找到一组参数使得$ELBO(q,x;\theta )$达到最大，即${\theta }_{t+1}=\underset{\theta }{\mathrm{argmax}}ELBO(q_{t+1},x;\theta )$

EM算法推导 方法2

对于对数边际似然函数$\log p(x;\theta )$可采用另一种分解方式，由此产生第二种推导方法：
$$\begin{array}{rl}\log p(x;\theta )& =\sum _{z}q(z)\log p(x;\theta )因为\sum _{z}q(z)=1\\ & =\sum _{z}q(z)\log [\frac{p(x,z;\theta )}{p(z|x;\theta )}\cdot \frac{q(z)}{q(z)}]\\ & =\sum _{z}q(z)\log \frac{p(x,z;\theta )}{q(z)}+\sum _{q}q(z)\log \frac{q(z)}{p(z|x;\theta )}\\ & =ELBO(q,x;\theta )+KL[q(z)||p(z|x;\theta )]\end{array}\text{\hspace{1em}(4)}$$其中$KL[q(z)||p(z|x;\theta )]$为分布$q(z)$和分布$p(z|x;\theta )$的KL散度。KL散度可以衡量两个分布的近似程度，其值是非负的。两个分布越相近，KL散度就越小；当两个分布完全相同时KL散度等于0。

由于KL散度的非负性，因此可以得出和 推导方法1 中一样的结论，即$ELBO(q,x;\theta )$为$\log p(x;\theta )$的下界。

这里可以采取与 推导方法1 相同的思路对EM算法进行推导：我们的目标是最大化$\log p(x;\theta )$，根据(4)式，可以将问题转为最大化$ELBO(q,x;\theta )+KL[q(z)||p(z|x;\theta )]$，因此可以先固定模型的参数$\theta$，找到一个最优的近似分布$q^{\ast }(z)$使得$ELBO(q^{\ast },x;\theta )$达到最大，由(4)式可知当$KL[q(z)||p(z|x;\theta )]=0$，即$q(z)=p(z|x;\theta )$时，$ELBO(q,x;\theta )=\log p(x;\theta )$，取得最大值；之后固定$q^{\ast }(z)$，找到最优的参数$\theta$使得$ELBO(q^{\ast },x;\theta )最大$，不断重复这两个步骤。

下图直观地展示了EM算法的迭代过程：在E-step中更新$q_{t+1}$使得$ELBO(q_{t+1},x;{\theta }_t)=\log p(x;{\theta }_t)$，即KL项等于0；在M-step中更新模型的参数${\theta }_{t+1}$使得$ELBO(q_{t+1},x;{\theta }_{t+1})\ge ELBO(q_{t+1},x;{\theta }_t)$，这样必然会导致$\log p(x;{\theta }_{t+1})\ge \log p(x;{\theta }_t)$，通过$n$次迭代使得对数边际似然函数$\log p(x;\theta )$不断增大直至收敛，从而得到最终模型的参数${\theta }_n$

变分推断

EM算法中一个很重要的环节是在E-step中寻找隐变量$z$的近似分布$q^{\ast }(z)$，在上述的推导过程中，我们通过令$KL[q(z)||p(z|x;\theta )]=0$，即 $q(z)=p(z|x;\theta )，$得到隐变量$z$最优的近似分布$p^{\ast }(z)$，然而这只局限于分布$p(z|x;\theta )$容易计算的情况，如混合高斯模型等。对于复杂的模型，$p(z|x;\theta )$几乎无法计算，于是可以借助变分推断的方法去近似求得$q^{\ast }(z)$

变分推断推导

变分推断是利用变分法进行推断的一种应用，对于EM算法而言，其变分推断的目标是寻找到隐变量$z$的一个简单的近似分布$q^{\ast }(z)$来代替难以求解的分布$p(z|x;\theta )$，问题定义如下：
$$q^{\ast }(z)=\underset{q(z)\in \mathcal{Q}}{\mathrm{argmin}}KL[q(z)||p(z|x;\theta )]\text{\hspace{1em}(5)}$$其中$\mathcal{Q}$是一个分布族。利用(4)式对(5)式进行推导，剔除了与优化目标无关的$\log p(x;\theta )$，将问题转化为了最大化$ELBO(q,x;\theta )$：
$$\begin{array}{rl}{q}^{\ast }(z)& =\underset{q(z)\in \mathcal{Q}}{\mathrm{argmin}}KL[q(z)||p(z|x;\theta )]\\ & =\underset{q(z)\in \mathcal{Q}}{\mathrm{argmin}}[\log p(x;\theta )-ELBO(q,x;\theta )]\\ & =\underset{q(z)\in \mathcal{Q}}{\mathrm{argmax}}ELBO(q,x;\theta )\end{array}\text{\hspace{1em}(6)}$$对于隐变量$z$，可以将其拆分成$n$个子变量$z_1,z_2,...,z_n$，显然这些子变量并不一定满足独立性，即：
$$p(z)\ne \prod _i^n{p}_i(z_i)\text{\hspace{1em}(7)}$$于是我们可以在变分推断中，选用mean-field（平均场）作为候选分布族$\mathcal{Q}$，构造出$n$个近似分布$q_1(z_1)...q_n(z_n)$使得这$n$个子变量相互独立，即：
$$q(z)=\prod _i^n{q}_i(z_i)\text{\hspace{1em}(8)}$$之后利用$q(z)$对问题进行推导：
$$\begin{array}{rl}ELBO(q,x;\theta )& =\sum _{z}q(z)\log \frac{p(x,z;\theta )}{q(z)}\\ & =\sum _{z}q(z)\log p(x,z;\theta )-\sum _{z}q(z)\log q(z)\\ & =\underset{\text{Part 1}}{\underbrace{\sum _z\prod _i^n{q}_i(z_i)\log p(x,z;\theta )}}-\underset{\text{Part 2}}{\underbrace{\sum _z\prod _i^n{q}_i(z_i)\sum _j^n\log q_j(z_j)}}\end{array}\text{\hspace{1em}(9)}$$
接下来分别对Part 1和Part 2进行推导，其它相关资料大多是用积分形式进行推导的，这里我转化为了离散的形式，感觉比积分形式更加清晰一些。这里对于下述推导中一些关键的符号表示进行说明，以便理解：

• ${\sum }_z$：表示对隐变量 z = { z 1 , z 2 , . . . , z n } z=\{z_1,z_2,...,z_n\} z={z1​,z2​,...,zn​}的所有可能取值组合情况进行求和
• ${\sum }_{z_j}$：表示对子隐变量$z_j$的所有可能取值进行求和
• ${\sum }_{z_{\backslash }{z}_j}$：表示对除了$z_j$的其它$n-1$个子隐变量的所有可能取值组合情况进行求和

对于Part 1进行如下推导：(思路是将其中一个子变量近似分布$q_j(z_j)$提出来)
$$\begin{array}{rl}\sum _z\prod _i^n{q}_i(z_i)\log p(x,z;\theta )& =\sum _{z_j}{q}_j(z_j)\sum _{z_{\backslash }{z}_j}\prod _{i\ne j}^n{q}_i(z_i)\log p(x,z;\theta )\\ & =\sum _{z_j}{q}_j(z_j){\mathbb{E}}_{q_{\backslash }{q}_j(z_{\backslash }{z}_j)}\log p(x,z;\theta )\end{array}\text{\hspace{1em}(10)}$$对于Part 2进行如下推导：
$$\begin{array}{rl}\sum _z\prod _i^n{q}_i(z_i)\sum _j^n\log q_j(z_j)& =\sum _j^n\sum _z\log q_j(z_j)\prod _i^n{q}_i(z_i)\\ & =\sum _j^n\sum _{z_j}\log q_j(z_j)\cdot \sum _{z_{\backslash }{z}_j}\prod _i^n{q}_i(z_i)\\ & =\sum _j^n\sum _{z_j}\log q_j(z_j)\cdot q_j(z_j)\underset{所有组合的概率和=1}{\underbrace{\sum _{z_{\backslash }{z}_j}\prod _{i\ne j}^n{q}_i(z_i)}}\\ & =\sum _j^n\sum _{z_j}{q}_j(z_j)\log q_j(z_j)\end{array}\text{\hspace{1em}(11)}$$将(10)和(11)代入到(9)式中得到：
$$ELBO(q,x;\theta )=\sum _{z_j}{q}_j(z_j){\mathbb{E}}_{q_{\backslash }{q}_j(z_{\backslash }{z}_j)}\log p(x,z;\theta )-\sum _j^n\sum _{z_j}{q}_j(z_j)\log q_j(z_j)\text{\hspace{1em}(12)}$$由(12)式可得针对某个子变量$z_k$，其$ELBO$函数为：
$$\begin{array}{rl}ELBO(q_k,x;\theta )& =\sum _{z_k}{q}_k(z_k){\mathbb{E}}_{q_{\backslash }{q}_k(z_{\backslash }{z}_k)}\log p(x,z;\theta )-\sum _{z_k}{q}_k(z_k)\log q_k(z_k)\\ & =\sum _{z_k}{q}_k(z_k)[\underset{可转变为\log 以化简为KL散度}{\underbrace{{\mathbb{E}}_{q_{\backslash }{q}_k(z_{\backslash }{z}_k)}\log p(x,z;\theta )}}-\log q_k(z_k)]\end{array}\text{\hspace{1em}(13)}$$设$\log {\tilde{p}}_k(x,z_k;\theta )={\mathbb{E}}_{q_{\backslash }{q}_k(z_{\backslash }{z}_k)}\log p(x,z;\theta )$，则(13)式可推导为：
$$\begin{array}{rl}ELBO(q_k,x;\theta )& =\sum _{z_k}{q}_k(z_k)[\log {\tilde{p}}_k(x,z_k;\theta )-\log q_k(z_k)]\\ & =-\sum _{z_k}{q}_k(z_k)\log \frac{q_k(z_k)}{{\tilde{p}}_k(x,z_k;\theta )}\\ & =-KL[q_k(z_k)||{\tilde{p}}_k(x,z_k;\theta )]\end{array}\text{\hspace{1em}(14)}$$因此最终问题转变为了最大化$-KL[q_k(z_k)||{\tilde{p}}_k(x,z_k;\theta )]$，即得到$q_k(z_k)$的最优表达式，其中 k ∈ { 1 , 2 , . . . , n } k\in \{1,2,...,n\} k∈{1,2,...,n}：
$$q_k(z_k)=\exp [{\mathbb{E}}_{q_{\backslash }{q}_k(z_{\backslash }{z}_k)}\log p(x,z;\theta )]\text{\hspace{1em}(15)}$$根据(15)式，可利用坐标上升法对每个自变量的近似分布$q_1(z_1),...,q_n(z_n)$进行迭代优化，直至收敛。

本文参考资料

1. 邱锡鹏，神经网络与深度学习，机械工业出版社，https://nndl.github.io/, 2020.
2. 论文篇：Latent Dirichlet Allocation（LDA）（三）