Auto-Encoding Variational Bayes论文解读

论文链接：Auto-Encoding Variational Bayes

本文基本按照原文的思路用自己的话对其进行重述和讲解

动机

考虑以下问题情境：假设有独立同分布的数据集$X={x^{(i)}},i=1,2,…,N$

数据点$x^{(i)}$的生成分为两步：(1) 从先验分布$p{\theta^*}(z)$中生成$z^{(i)}$； (2) 从条件分布$p{\theta^*}(x|z)$中生成$x^{(i)}$

其中隐变量$z^{(i)}$和分布参数$\theta^*$都是未知的，此时常常会遇到以下两种情况：

难解性：似然的积分 $p\theta(x)=\int p{\theta}(z)p{\theta}(x|z)\mathrm{d}z$ 难以计算，使得似然无法估计或求微分；后验概率 $p{\theta}(z|x)=\frac{p{\theta}(x|z)p{\theta}(z)}{p_{\theta}(x)}$ 难以计算，导致不能用EM算法；平均场变分推理要求的积分难以计算
大数据集：数据集非常大的情况下基于采样的方法（例如蒙特卡罗EM）效率很低，因此我们希望使用minibatch进行参数更新

我们主要关心在上述情况下的三个相关问题：

如何对模型参数 $\theta$ 进行高效的极大似然估计(ML)或极大后验概率估计(MAP)
如何对隐变量$z$进行高效的近似后验推断 $p_{\theta}(z|x)$
如何对边缘似然$\log p_{\theta}(x)$进行高效的推断

SGVB和AEVB

为了解决上述问题，我们引入一个识别模型$q{\phi}(z|x)$，它是真实后验$p{\theta}(z|x)$的近似

我们将$q{\phi}(z|x)$视作一个Probablistic Encoder，而将$p{\theta}(x|z)$视作一个Probablistic Decoder

在变分推断中似然$\log p_{\theta}(x^{(i)})$可表示为

$\log p_{\theta}(x^{(i)})=D_{KL}(q_{\phi}(z|x^{(i)})\|p_{\theta}(z|x^{(i)}))+L(\theta,\phi;x^{(i)}) \tag{1}$

其中$L(\theta,\phi;x^{(i)})$是变分下界，表示为

$\log p_{\theta}(x^{(i)})\geq L(\theta,\phi;x^{(i)})=E_{q_{\phi}(z|x^{(i)})}\left[\log p_{\theta}(x^{(i)},z)-\log q_{\phi}(z|x^{(i)})\right] \tag{2}$

还可以表示为

$\begin{aligned} L(\theta,\phi;x^{(i)}) &=E_{q_{\phi}(z|x^{(i)})}\left[\log p_{\theta}(x^{(i)}|z)+\log p_{\theta}(z)-\log q_{\phi}(z|x^{(i)})\right]\\ &=E_{q_{\phi}(z|x^{(i)})}\left[\log p_{\theta}(x^{(i)}|z)\right]-D_{KL}(q_{\phi}(z|x^{(i)})\|p_{\theta}(z)) \end{aligned} \tag{3}$

考虑用一个随机噪声变量$\epsilon\sim p(\epsilon)$和一个确定性的可微函数$g{\phi}(\epsilon,x)$来重参数化(reparameterize)随机变量$\hat{z}\sim q{\phi}(z|x)$，即

$\hat{z}=g_{\phi}(\epsilon,x)\quad with\quad \epsilon\sim p(\epsilon) \tag{4}$

此时用蒙特卡罗估计$f(z)$的期望，就发现参数$\phi$从随机抽样中分离出来了，使得抽样变得可行

$E_{q_{\phi}(z|x^{(i)})}[f(z)]=E_{p(\epsilon)}\left[f(g_{\phi}(\epsilon, x^{(i)}))\right]\simeq \frac{1}{L}\sum_{l=1}^L f(g_{\phi}(\epsilon^{(l)}, x^{(i)})) \tag{5}$

将这个思路代入式(2)中，就得到了随机梯度变分贝叶斯(Stochastic Gradient Variational Bayes, SGVB)算法，即可以将梯度法应用于参数$\phi$的优化

$\hat{L}^A(\theta,\phi;x^{(i)})=\frac{1}{L}\sum_{l=1}^L \log p_{\theta}(x^{(i)},z^{(i,l)})-\log q_{\phi}(z^{(i,l)}|x^{(i)}) \\ \text{其中}\ z^{(i,l)}=g_{\phi}(\epsilon^{(i,l)}, x^{(i)}) \tag{6}$

一般来说，式(3)中的KL散度项$D{KL}(q{\phi}(z|x^{(i)})|p_{\theta}(z))$可以被解析地积分（论文附录B有解释），因此不用采样，于是根据式(3)有SGVB的另一个版本，且方差更小

$\hat{L}^B(\theta,\phi;x^{(i)})=-D_{KL}(q_{\phi}(z|x^{(i)})\|p_{\theta}(z)) +\frac{1}{L}\sum_{l=1}^L \log p_{\theta}(x^{(i)}|z^{(i,l)}) \\ \text{其中}\ z^{(i,l)}=g_{\phi}(\epsilon^{(i,l)}, x^{(i)}) \tag{7}$

给定一个大小为$N$的数据集，我们可以使用minibatch构建整个数据集的变分下界，即

$L(\theta,\phi;X)\simeq \hat{L}(\theta,\phi;X^M)=\frac{N}{M}\sum_{i=1}^M\hat{L}(\theta,\phi;x^{(i)})$

其中minibatch $X^M$是从数据集$X$中随机抽取的$M$个数据构成的，实验表明，若$M$足够大，则对每个数据点$x^{(i)}$只需要令$L=1$即可

式(7)的后项显然就是Auto-Encoder的Negative reconstruction error，而前项KL散度则可视为正则项因此该算法也可称为自编码器变分贝叶斯(Auto-Encoding Variational Bayes, AEVB)

即输入样本$x$和噪声$\epsilon$，$g{\phi}(\cdot)$将其映射为隐变量$z$，对应Ecoder，之后$p{\theta}(x|z)$将$z$重构为$x’$，对应Decoder

AEVB

重参数化

上面提到$z\sim q{\phi}(z|x)$（或$z\sim q{\phi}(z)$）可以使用辅助随机变量$\epsilon\sim p(\epsilon)$和确定性函数$z=g_{\phi}(\epsilon,x)$进行重参数化

重参数化使得随机抽样与参数$\phi$无关，从而使得抽样变得可行

举个例子，假设$z\sim p(z)=N(\mu, \sigma^2)$，令$z=\mu+\sigma\epsilon$且$\epsilon \sim N(0,1)$，于是有

$E_{N(z;\mu,\sigma^2)}[f(z)]=E_{N(0,1)}[f(\mu+\sigma\epsilon)]\simeq \frac{1}{L}\sum_{l=1}^L f(\mu+\sigma\epsilon^{(l)})$

一般的，针对给定的$q{\phi}(z|x)$有几个选择$g{\phi}(\cdot)$和$p(\epsilon)$的基本方法

逆累积分布(ICDF)易处理：此时令$\epsilon \sim U(0, I)$，令$g{\phi}(\epsilon,x)$为$q{\phi}(z|x)$的逆累积分布函数。例如：指数、柯西、瑞利分布等
位置-尺度(location-scale)族：此时令$\epsilon$为标准分布（location=0, scale=1），令$g_{\phi}(\cdot)=location+scale\cdot \epsilon$。例如：高斯、均匀、拉普拉斯、t-分布等
组合分布：一般可以将随机变量描述为$\epsilon$的不同变换。例如：Gamma、Beta、F分布等

Variational Auto-Encoder

这里将给出一个使用神经网络实现AEVB的例子，称为变分自编码器VAE

假设先验$p{\theta}(z)=N(z;0,I)$，近似后验$q{\phi}(z|x)$有很多可选形式，这里我们假设真实后验$p_{\theta}(z|x)$近似于具有对角方差的高斯分布，则

$\log q_{\phi}(z|x^{(i)})=\log N(z;\mu^{(i)},(\sigma^{(i)})^2I)$

其中均值$\mu^{(i)}$和方差$\sigma^{(i)}$是Encoder的输出，即它们是数据点$x^{(i)}$和参数$\phi$的函数

由于先验和近似后验假设都是高斯分布，因此式(7)中KL散度项可以直接计算（论文附录B），于是有

$\hat{L}(\theta,\phi;x^{(i)})\simeq \frac{1}{2}\sum_{j=1}^J\left(1+\log((\sigma_j^{(i)})^2)-(\mu_j^{(i)})^2+(\sigma_j^{(i)})^2\right)+\frac{1}{L}\sum_{l=1}^L \log p_{\theta}(x^{(i)}|z^{(i,l)}) \\ \text{其中}\ z^{(i,l)}=g_{\phi}(\epsilon^{(i,l)}, x^{(i)}), \quad \epsilon^{(l)}\sim N(0,I)\tag{7}$

tf2+keras有官方的VAE实现示例：Variational AutoEncoder Keras实现