CS236 Deep Generative Models Lecture 6

课程主页：https://deepgenerativemodels.github.io/

课件资料：https://github.com/Subhajit135/CS236_DGM，https://github.com/deepgenerativemodels/notes

视频地址：https://www.bilibili.com/video/av81625948?from=search&seid=4636291486683332935

这里回顾CS236 Lecture 6的课程内容，这一讲介绍了EM算法以及ELBO。

Partially observed data

假设我们有联合分布

$p(\mathbf{X}, \mathbf{Z} ; \theta)$

$\mathbf X$为可观测变量，$\mathbf Z$为不可观测变量，数据集为$\mathcal{D}=\left\{\mathbf{x}^{(1)}, \cdots, \mathbf{x}^{(M)}\right\}$，那么数据的对数似然为

$\begin{aligned} \ell(\theta ; \mathcal{D}) &=\log \prod_{\mathbf{x} \in \mathcal{D}} p(\mathbf{x} ; \theta)\\ &=\sum_{\mathbf{x} \in \mathcal{D}} \log p(\mathbf{x} ; \theta) \\ &=\sum_{\mathbf{x} \in \mathcal{D}} \log \sum_{\mathbf{z}} p(\mathbf{x}, \mathbf{z} ; \theta) \end{aligned}$

注意到下式很难计算

$\sum_{z} p_{\theta}(x, z)$

所以后续主要讨论如何对其处理。

First attempt: Naive Monte Carlo

利用恒等变形

$\sum_{\text {All possible values of } \mathbf{z}} p_{\theta}(\mathbf{x}, \mathbf{z})=|\mathcal{Z}| \sum_{\mathbf{z} \in \mathcal{Z}} \frac{1}{|\mathcal{Z}|} p_{\theta}(\mathbf{x}, \mathbf{z})=|\mathcal{Z}| \mathbb{E}_{\mathbf{z} \sim \text {Uniform}(\mathcal{Z})}\left[p_{\theta}(\mathbf{x}, \mathbf{z})\right]$

所以可以使用蒙特卡洛模拟的方法估计上式

$\sum_{z} p_{\theta}(\mathbf{x}, \mathbf{z}) \approx|\mathcal{Z}| \frac{1}{k} \sum_{j=1}^{k} p_{\theta}\left(\mathbf{x}, \mathbf{z}^{(j)}\right)$

The EM Algorithm

利用KL散度，我们有恒等式

$\begin{aligned} D_{K L}(q(\mathbf{z}) \| p(\mathbf{z} | \mathbf{x} ; \theta)) &=\sum_{\mathbf{z}} q(\mathbf{z}) \log \frac{q(\mathbf{z})}{p(\mathbf{z} | \mathbf{x} ; \theta)} \\ &=-\sum_{\mathbf{z}} q(\mathbf{z}) \log p(\mathbf{z} | \mathbf{x} ; \theta)+\underbrace{\sum_{\mathbf{z}} q(\mathbf{z}) \log q(\mathbf{z})}_{-H(q)} \\ &=-\sum_{\mathbf{z}} q(\mathbf{z}) \log p(\mathbf{z} | \mathbf{x} ; \theta)-H(q) \\ &=-\sum_{\mathbf{z}} q(\mathbf{z}) \log \frac{p(\mathbf{z}, \mathbf{x} ; \theta)}{p(\mathbf{x} ; \theta)}-H(q) \\ &=-\sum_{\mathbf{z}} q(\mathbf{z}) \log p(\mathbf{z}, \mathbf{x} ; \theta)+\sum_{\mathbf{z}} q(\mathbf{z}) \log p(\mathbf{x} ; \theta)-H(q) \\ &=-\sum_{\mathbf{z}} q(\mathbf{z}) \log p(\mathbf{z}, \mathbf{x} ; \theta)+\log p(\mathbf{x} ; \theta)-H(q) \geq 0 \end{aligned}$

即

$D_{K L}(q(\mathbf{z}) \| p(\mathbf{z} | \mathbf{x} ; \theta))=-\sum_{\mathbf{z}} q(\mathbf{z}) \log p(\mathbf{z}, \mathbf{x} ; \theta)+\log p(\mathbf{x} ; \theta)-H(q) \geq 0$

即对任意$q$，成立Evidence lower bound (ELBO)

$\log p(\mathbf{x} ; \theta) \geq \sum_{\mathbf{z}} q(\mathbf{z}) \log p(\mathbf{z}, \mathbf{x} ; \theta)+H(q)$

当且仅当$q=p(\mathbf{z} | \mathbf{x} ; \theta)$时等号成立

$\log p(\mathbf{x} ; \theta)=\sum_{\mathbf{z}} q(\mathbf{z}) \log p(\mathbf{z}, \mathbf{x} ; \theta)+H(q)$

记

$H(q)+\sum_{z} q(\mathbf{z}) \log p(\mathbf{z}, \mathbf{x} ; \theta)=\log p(\mathbf{x} ; \theta)-D_{K L}(q \| p(\mathbf{z} | \mathbf{x} ; \theta)) \triangleq F[q, \theta]$

我们如果想最大化$\log p(\mathbf{x} ; \theta)$，那么最大化其下界$F[q, \theta]$可以取得较好的效果，EM算法的本质是在$F[q, \theta]$上使用坐标上升：

初始化$\theta^{(0)}$
$q^{(0)}=\arg \max _{q} F\left[q, \theta^{(0)}\right]=p\left(\mathbf{z} | \mathbf{x} ; \theta^{(0)}\right)$
$\theta^{(1)}=\arg \max _{\theta} F\left[q^{(0)}, \theta\right]=\arg \max _{\theta} \sum_{z} q^{(0)}(\mathbf{z}) \log p(\mathbf{z}, \mathbf{x} ; \theta)$
$q^{(1)}=\arg \max _{q} F\left[q, \theta^{(1)}\right]$
…

不难看出我们有

$\begin{aligned} \log p\left(\mathbf{x} ; \theta^{(0)}\right)=F\left[q^{(0)}, \theta^{(0)}\right] \leq F\left[q^{(0)}, \theta^{(1)}\right] &=\log p\left(\mathbf{x} ; \theta^{(1)}\right)-D_{K L}\left(q^{(0)} \| p\left(\mathbf{z} | \mathbf{x} ; \theta^{(1)}\right)\right) \\ & \leq \log p\left(\mathbf{x} ; \theta^{(1)}\right)=F\left[q^{(1)}, \theta^{(1)}\right]=\cdots \end{aligned}$

The Variational EM Algorithm

还有一种变分EM算法：

$D_{K L}(q(\mathbf{z} ; \phi) \| p(\mathbf{z} | \mathbf{x} ; \theta))=-\sum_{\mathbf{z}} q(\mathbf{z} ; \phi) \log p(\mathbf{z}, \mathbf{x} ; \theta)+\log p(\mathbf{x} ; \theta)-H(q(\mathbf{z} ; \phi))$

重写为

$H(q(\mathbf{z} ; \phi))+\sum_{\mathbf{z}} q(\mathbf{z} ; \phi) \log p(\mathbf{z}, \mathbf{x} ; \theta)=\log p(\mathbf{x} ; \theta)-D_{K L}(q(\mathbf{z} ; \phi) \| p(\mathbf{z} | \mathbf{x} ; \theta))=F[\phi, \theta]$

变分EM对$F[\phi, \theta]$做坐标上升

初始化$\theta^{(0)}$
$\phi^{(1)}=\arg \max _{\phi} F\left[\phi, \theta^{(0)}\right] \neq p\left(\mathbf{z} | \mathbf{x} ; \theta^{(0)}\right)$
$\theta^{(1)}=\arg \max _{\theta} F\left[\phi^{(1)}, \theta\right]$（可以使用MC）
$\phi^{(2)}=\arg \max _{\phi} F\left[\phi, \theta^{(1)}\right] \neq p\left(\mathbf{z} | \mathbf{x} ; \theta^{(1)}\right)$
…

注意上述方法不一定保证收敛到局部最优。

The Evidence Lower bound

回顾之前介绍的ELBO：

$\begin{aligned} \log p(\mathbf{x} ; \theta) & \geq \sum_{\mathbf{z}} q(\mathbf{z} ; \phi) \log p(\mathbf{z}, \mathbf{x} ; \theta)+H(q(\mathbf{z} ; \phi))=\underbrace{\mathcal{L}(\mathbf{x} ; \theta, \phi)}_{\text {ELBO }} \\ &=\mathcal{L}(\mathbf{x} ; \theta, \phi)+D_{K L}(q(\mathbf{z} ; \phi) \| p(\mathbf{z} | \mathbf{x} ; \theta)) \end{aligned}$

图示如下：

The Evidence Lower bound applied to the entire dataset

将上述不等式应用到数据集中得到

$\ell(\theta ; \mathcal{D})=\sum_{\mathbf{x}^{\prime} \in \mathcal{D}} \log p\left(\mathbf{x}^{i} ; \theta\right) \geq \sum_{\mathbf{x}^{\prime} \in \mathcal{D}} \mathcal{L}\left(\mathbf{x}^{i} ; \theta, \phi^{i}\right)$

所以

$\max _{\theta} \ell(\theta ; \mathcal{D}) \geq \max _{\theta, \phi^{1}, \cdots, \phi^{M}} \sum_{\mathbf{x}^{i} \in \mathcal{D}} \mathcal{L}\left(\mathbf{x}^{i} ; \theta, \phi^{i}\right)$

Learning via stochastic variational inference (SVI)

现在讨论如何对上述不等式求最大值，首先我们有

$\begin{aligned} \mathcal{L}\left(\mathbf{x}^{i} ; \theta, \phi^{i}\right) &=\sum_{\mathbf{z}} q\left(\mathbf{z} ; \phi^{i}\right) \log p\left(\mathbf{z}, \mathbf{x}^{i} ; \theta\right)+H\left(q\left(\mathbf{z} ; \phi^{i}\right)\right) \\ &=E_{q\left(\mathbf{z} ; \phi^{i}\right)}\left[\log p\left(\mathbf{z}, \mathbf{x}^{i} ; \theta\right)-\log q\left(\mathbf{z} ; \phi^{i}\right)\right] \end{aligned}$

所以我们可以使用蒙特卡洛方法以及梯度下降更新上述参数。

$E_{q(\mathbf{z} ; \phi)}[\log p(\mathbf{z}, \mathbf{x} ; \theta)-\log q(\mathbf{z} ; \phi)] \approx \frac{1}{k} \sum_{k} \left(\log p\left(\mathbf{z}^{k}, \mathbf{x} ; \theta\right)-\log q\left(\mathbf{z}^{k} ; \phi\right)\right)$

显然我们有

$\begin{aligned} \nabla_{\theta} E_{q(z ; \phi)}[\log p(z, x ; \theta)-\log q(z ; \phi)] &=E_{q(z ; \phi)}\left[\nabla_{\theta} \log p(z, x ; \theta)\right] \\ & \approx \frac{1}{k} \sum_{k} \nabla_{\theta} \log p\left(z^{k}, x ; \theta\right) \end{aligned}$

但是注意到$\nabla_{\phi} \mathcal{L}(\mathbf{x} ; \theta, \phi)$比较复杂，无法直接写出。