VAE一步一步详细推导

VAE

最近在推DM，发现VAE又忘了，重新看一遍Tutorial on Variational Autoencoders 挺费事的，这一次干脆把每一步的推导过程和需要注意的细节都尽量记录下来，加深印象也方便下次快速捡起来。

引言

目的是从样本集$X$估计其概率分布$P(X)$，最直观的就是先对$P(X)$做一些假设，如正态分布、均匀分布等，然后用最大似然估计(Max Likelihood Estimation, MLE)：

$$MLE=\sum _i^{n}logP(x;\theta )$$

或者（Maximum a Posteriori，MAP）估计参数$\theta$

$$MAP=\sum _i^{n}logP(\theta |x)=\sum _i^n(logP(x|\theta )+logP(\theta )-logP(x))$$

在任务复杂的情况下，假设目标服从某个简单分布一般不符合实际情况，比如二值图像不符合高斯分布，邻域内像素之间的关联性等，因此，这种方法不太可行。

MAP和MLE，参数估计两大方法各有特点，如果考虑先验和evidence的话MAP更合适。具体分析网上文章很多。

Latent variable

一个分布经过足够复杂的高维抽象后可以映射到另一个分布，如下图所示。基于此，换个思路，我们可以把一个简单易采样的分布P(z)，通过一个足够复杂且可微的函数g，映射到另一个我们想要的分布g(P(z))。

VAE把这个简单分布定义为隐变量$z$，且$P(z)\sim N(0,I)$，那根据全概率定理改写公式：

$$\begin{array}{rl}P(X)& =\int p(x,z)dz\\ & =\int p(x|z)p(z)dz\\ & =\sum _{i}p(x|z_i)p(z_i)\\ & \approx \frac{1}{n}\sum _{i}p(x|z_i)\end{array}$$
这里，$p(x|z)$可以当作一个decoder，$z\to X$，也可以写作$f(z;\theta )$。

即正态分布经过一个decode网络生成一个目标图像。为了便于求导（如果是二值图像可以假设为Bernoulli），假设
$$P(X|z;\theta )\sim N(X|f(z;\theta ),{\sigma }^2\ast I)$$

这样操作后，把对P(X)采样的问题，转换成了对p(z)的映射问题。

ELBO

接上一步，对z进行采样，然后估计P(x|z)会碰到问题：

• z是连续值，大部分z对应的p(x|z)可能是等于0的，即无效采样；
• z和X不能很好对应起来

结果就是拟合效率低下甚至无法学习。为此，VAE设计了一个编码器Q(z|x)。样本x输入Q得到z，这个z和x就对应起来了，接着z输入P生成的$x_{}{}^{\prime }$应该是最接近x的，这样就极大的约束了采样空间，提高了拟合效率。至此，VAE中涉及的P、Q、z、x就到齐了，下一步是把它们联系起来。直观的方法就是希望从两个方向估计的Q(z)足够接近：

$$D[Q(z)\mid \mid P(z\mid X)]=E_{z\sim Q}[logQ(z)-logP(z\mid X)]$$
右侧用贝叶斯公式展开：
$$D[Q(z)||P(z|X)]={\mathbb{E}}_{z\sim Q}[\log Q(z)-\log P(X|z)-\log P(z)]+\log P(X)$$
$\log P(X)$是我们想要最大化的目标，去掉常数项，移项，部分替换成KL散度后：
$$\log P(X)-D[Q(z)||P(z|X)]={\mathbb{E}}_{z\sim Q}[\log P(X|z)]-D[Q(z)||P(z)]$$
Q(z)替换成Q(z|x)，目的是帮助建立z和x的联系，辅助推断x，形成一个编码器的结构（输入x，输出z）,至此，Q的形式就和本小节开始时定义的形式就对应上了。

$$\log P(X)-D[Q(z|X)||P(z|X)]={\mathbb{E}}_{z\sim Q}[\log P(X|z)]-D[Q(z|X)||P(z)]$$

至此左侧为优化目标，而右侧则是可以迭代优化的等价项。因为散度是非负的，所以：
$$\log P(X)\ge {\mathbb{E}}_{z\sim Q}[\log P(X|z)]-D[Q(z|X)||P(z)]$$
最终，ELBO的形式如下：
$$ELBO={\mathbb{E}}_{z\sim Q}[\log P(X|z)]-D[Q(z|X)||P(z)]=E_{Q(z|x)}[\log (\frac{p(x,z)}{Q(z|x)})]$$

优化技巧

最大化ELBO，等价于右侧最大化第一项，以及最小化第二项。

第一项，回想之前的关于p(x|z)的假设：
$$P(X|z;\theta )\sim N(X|f(z;\theta ),{\sigma }^2\ast I)$$
那么：

$$\begin{array}{rl}\log (p(x|z))& =\log (\frac{1}{\sqrt{2\pi {\sigma }^2}}\ast ex{p}^{-\frac{(x-f(z;\theta ){)}^2}{2{\sigma }^2}})\\ & =\log (\frac{1}{\sqrt{2\pi {\sigma }^2}})-\frac{(x-f(z;\theta ){)}^2}{2{\sigma }^2}\end{array}$$
因为$\sigma$在这里是一个定值，因此
$$\log P(X|z)\simeq C-\frac{1}{2}\Vert X-f(z;\theta ){\Vert }^2/{\sigma }^2$$

第二项，对于两个高斯分布，KL散度的公式可以用解析法得到：
$$D[N({\mu }_0,{\Sigma }_0)||N({\mu }_1,{\Sigma }_1)]=\frac{1}{2}(\text{tr}({\Sigma }_1^{-1}{\Sigma }_0)+({\mu }_1-{\mu }_0{)}^T{\Sigma }_1^{-1}({\mu }_1-{\mu }_0)-k+\log \frac{\det ({\Sigma }_1)}{\det ({\Sigma }_0)})$$
至此，目标函数的两个部分就都i出来了。用流程图表示，就是文章里面最经典的那张图的左侧部分：

$\sigma$是一个常数，类似于loss的加权系数

Reparameterization

目前，我们推导出了待优化的目标函数形式，其流程如上图左。仔细一看会发现，P和Q中间的数据流$N(,)$是不可微的，这就会导致梯度流无法传导。这里，作者用了一个技巧，较Reparameterization（重参数）。

简单的说，如果一个分布$z\sim N(\mu ,{\sigma }^2)$，则我们可以通过另一个分布$ϵ\sim N(0,1)$用如下公式进行采样：
$$z=\mu +\sigma \ast ϵ$$
基于此，我们改写ELBO：
$${\mathbb{E}}_{ϵ\sim N(0,I)}\left[\log P(X|z=f(z;\theta )+\sigma ϵ)\right]-D[Q(z|X)||P(z)]$$
对应上图右侧部分。

一些思考

关于$P(x|z)$和$Q(Z)$的假设

VAE假设P为正态分布，目的是方便计算并且正态分布大部分情况下都还算是比较合理的，也可以选择其他分布，如Bernoulli（在$log(B)$的形式下是可微的）或者其他更复杂的形式比如混合高斯等。Q同样如此。

关于VAE生成质量

影响生成质量的有几点：

• 图像的pix之间或者局部和整体之间都存在复杂的关联性，对图像的分布做各项独立的高斯分布假设不太合理；
• VAE显示建模一个连续分布，连续分布的采样空间远大于离散分布，因此需要更庞大且多样的数据进行拟合；
• ELBO的两个部分是一个对抗和均衡的过程，因此，在KL的误差范围内所有的特征都可能会被L2平滑。这样一定程度能够保证生成图像的连续性，坏处就是会模糊。

关于VAE的一些改进方法

• VQVAE。实际是AE+PixelCNN，连续空间离散化，提高生成质量。
• Perceptual Loss。更复杂的loss，从content和sytle层面出发。
• Inverse Autoregressive Flow。提升分布假设的灵活性，更好的匹配真实分布。
• 和GAN的结合方法，比如IntroVAE。
• 其他的，NVAE等

不知是否有没发现的错漏，到时候再改正。

Tutorial on Variational Autoencoders

An Introduction to Variational Autoencoders

Understanding Variational Autoencoders (VAEs) 这个写的挺好的，之前还看了几篇大牛写的忘了链接了，下次找到了再补上