【扩散模型Diffusion Model系列】0-从VAE开始（隐变量模型、KL散度、最大化似然与AIGC的关系）

VAE

VAE(Variational AutoEncoder)，变分自编码器，是一种无监督学习算法，被用于压缩、特征提取和生成式任务。相比于GAN(Generative Adversarial Network)，VAE在数学上有着更加良好的性质，有利于理论的分析和实现。

1 生成式模型的目标——KL散度和最大化似然MLE

生成式模型(Generative Model)的目标是学习一个模型，从一个简单的分布$p(x)$中采样出数据$x$，通过生成模型$f(x)$来逼近真实数据的分布$p_{data}(x)$，并生成样本，实现了上面这一点即使我们所希望的结果。

自然，我们可以想到，生成模型最本质的目标就是最小化模型生成的样本分布$p_{\theta }(x)$和真实样本分布$p_{data}(x)$之间的KL散度：
$$\begin{array}{rl}& \underset{\theta }{\mathrm{argmin}}{D}_{KL}(p_{data}(x)||p_{\theta }(x))\\ =& \underset{\theta }{\mathrm{argmin}}\int p_{data}(x)log\frac{p_{data}(x)}{p_{\theta }(x)}\\ =& \underset{\theta }{\mathrm{argmax}}\int p_{data}(x)log{p}_{\theta }(x)【p_{data}(x)无参数优化】\\ =& \underset{\theta }{\mathrm{argmax}}{E}_{x\sim p_{data}(x)}\left[\log p_{\theta }(x)\right]【期望的定义】\\ \approx & \underset{\theta }{\mathrm{argmax}}\frac{1}{m}\sum _{i=1}^m\log p_{\theta }(x_i)【从数据集中采样m个，估算期望，对应于训练过程】\\ =& \underset{\theta }{\mathrm{argmax}}\prod _{i=1}^m{p}_{\theta }(x_i)【最大化似然】\end{array}$$

2 从AE到VAE

显然上述的生成式模型并不专门针对VAE，任何一个输出和输入相同分布的模型都可以得到此结论，那么不得不提的就是AE(AutoEncoder)，诸如MAE、DAE、VQVAE等。

AE的目标是最小化重构误差，即重构误差越小，则表示模型生成的数据和真实数据的分布越接近，和上述描述的生成式模型目标一致，但AE之所以不能用于生成式模型，是因为AE的Bottleneck的分布实际上是未知的，我们无法凭空采样一个符合bottleneck分布的数据，所以AE不能直接用于生成式模型。

AE和VAE实际上都可以被视为一个隐变量模型$p(x|z)$，认为在真实数据分布之后，存在着一个隐变量$z$，其分布为$p(z)$，$x$和$z$之间存在一个隐变量连接，即$p_{\varphi }(x|z)$。
例如可以将所有的矩形视为一个真实分布$p(x)$，而所有的长和宽的分布视为$p(z)$，那么显然，当我们从$p(z)$采样一个长宽$z即z\sim p(z)$时，事实上也采样到了一个矩形，这是因为我们认为存在明确的$p_{\varphi }(x|z)$，即矩形的宽和高和矩形的分布存在一个连接。

在AE中，$z$是bottleneck特征向量，很好地表征了原始数据的特征，因此可以利用Decoder即$p_{\theta }(x|z)$进行复原，理论上如果我们可以采样到$z$，那么就可以进行复原，但事实上我们不知道$z$的分布，因此我们无法用AE进行生成式。

而在VAE中，我们希望通过Encoder的学习，将真实的后验分布$p_{\varphi }(z|x)$进行近似，即$p_{\theta }(z|x)$，并且希望后验分布$p_{\varphi }(z|x)$服从于正态分布$N(0,I)$，这样的话，在优化足够好的Encoder，即$p_{\theta }(z|x)\approx N(0,I)$时，我们有：
$$\begin{array}{rl}p(z)=& \int p_{\varphi }(z|x)p(x)dx=\int p_{\theta }(z|x)p(x)dx\\ =& \int N(0,I)p(x)dx=N(0,I)\int p(x)dx=N(0,I)\end{array}$$
这样的话，我们就可以轻松地从正态分布中采样$z\sim p(z)$，为此我们必须考虑对“AE的bottleneck”进行修改，从而让$p_{\theta }(z|x)$的分布近似于$N(0,I)$，这也是为什么VAE输出的是正态分布的参数$\mu ,{\sigma }^2$。

理论上，我们通过重参数技巧$x=\mu +\sigma ϵ,ϵ\sim N(0,I)$，即可实现输出为$N(\mu ,{\sigma }^2)$，且将采样这一不可导的操作转为可导。

若是不对编码器$p_{\theta }(z|x)$加以限制，只使用MSE进行训练，VAE会逐渐退化为AE，因为网络一定会倾向于将${\sigma }^2\to 0$,因为这最有利于重建，那么我们最直接的想法就是使用另外2个MSE，强迫$\mu \to 0,{\sigma }^2\to I$，但这样3个MSE之间的比例就会十分难以调整，容易顾此失彼，因此，我们继续从MLE出发，继续推导VAE的损失函数。

3 VAE的损失函数

承接第一节，我们已经确认了生成式网络的最终目标就是最大化$p_{\theta }(x)$的似然，而正如常识所知，直接最大化$p_{\theta }(x)$太过困难，我们采用隐变量模型建构，那么公式如下:
$$\begin{array}{rl}log{p}_{\theta }(x)& =log{p}_{\theta }(x)\int p_{\varphi }(z|x)dz\\ & =\int p_{\varphi }(z|x)log{p}_{\theta }(x)dz\\ & =\int p_{\varphi }(z|x)log\frac{p_{\theta }(x,z)}{p_{\theta }(z|x)}dz【条件概率的定义】\\ & =\int p_{\varphi }(z|x)log\frac{p_{\theta }(x,z)p_{\varphi }(z|x)}{p_{\theta }(z|x)p_{\varphi }(z|x)}dz\\ & =\int p_{\varphi }(z|x)log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}dz+\int p_{\varphi }(z|x)log\frac{p_{\varphi }(z|x)}{p_{\theta }(z|x)}dz\\ & =E_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]+D_{KL}(p_{\varphi }(z|x)||p_{\theta }(z|x))\\ & \ge E_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]【KL散度\ge 0，可利用-lnx\ge 1-x证明】\end{array}$$
最终我们可认为损失函数为:
$$\underset{\theta }{\mathrm{argmax}}log{p}_{\theta }(x)=\underset{\theta }{\mathrm{argmax}}{E}_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]$$
$$L=-E_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]$$
对于上式我们可以有2种理解：

1. 最大化$p_{\theta }(x)$转化为了最大化下界ELBO(Evidence Lower Bound)，因此我们只需要去优化$E_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]$
2. $E_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]=log{p}_{\theta }(x)-D_{KL}(p_{\varphi }(z|x)||p_{\theta }(z|x))$，最小化损失函数$L$即最大化$E_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]$时，会最大化似然$log{p}_{\theta }(x)$，即让生成图片更真实的同时；最小化Encoder建模的$p_{\theta }(z|x)$和真实隐变量后验分布$p_{\varphi }(z|x)$之间的KL散度(当然是事实上是二者trade off)

若是我们使得$p_{\theta }(z|x)\to N(0,I)$，即大功告成。于是我们继续分解:
$$\begin{array}{rl}{E}_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]& =\int p_{\varphi }(z|x)log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}dz\\ & =\int p_{\varphi }(z|x)log\frac{p_{\theta }(x|z)p(z)}{p_{\varphi }(z|x)}dz\\ & =\int p_{\varphi }(z|x)log{p}_{\theta }(x|z)dz+\int p_{\varphi }(z|x)log\frac{p(z)}{p_{\varphi }(z|x)}dz\\ & =E_{z\sim p_{\varphi }(z|x)}[log{p}_{\theta }(x|z)]-D_{KL}(p_{\varphi }(z|x)||p(z))\\ & \approx E_{z\sim p_{\theta }(z|x)}[log{p}_{\theta }(x|z)]-D_{KL}(p_{\theta }(z|x)||p(z))\end{array}$$
其中$E_{z\sim p_{\theta }(z|x)}[log{p}_{\theta }(x|z)]$为最大似然，我们假设最终为正态分布，最大似然就完全等价于最小化重建损失MSE
而$D_{KL}(p_{\theta }(z|x)||p(z))$则为正则项，用于约束Encoder的输出，具体公式如下：
$$\begin{array}{rl}{D}_{KL}(p_{\theta }(z|x)||p(z))& =-\int p_{\varphi }(z|x)log\frac{p(z)}{p_{\varphi }(z|x)}dz\\ & =\int p_{\varphi }(z|x)[\frac{z^2}{2}-log\frac{1}{\sqrt{2\pi }}-\frac{(z-{\mu }_{\theta }(x){)}^2}{2{{\sigma }_{\theta }(x)}^2}+log\frac{1}{\sqrt{2\pi {{\sigma }_{\theta }(x)}^2}}]dz\\ & =\frac{1}{2}\int p_{\varphi }(z|x)[z^2-(\frac{z-{\mu }_{\theta }(x)}{{\sigma }_{\theta }(x)}{)}^2-log{{\sigma }_{\theta }(x)}^2]dz\\ & =\frac{1}{2}[-1+{{\mu }_{\theta }(x)}^2+{{\sigma }_{\theta }(x)}^2-log{{\sigma }_{\theta }(x)}^2]【E(z^2)={\mu }^2+{\sigma }^2，用于解答z^2和(\frac{z-\mu }{\sigma }{)}^2】\end{array}$$
综上，我们得到了VAE的损失函数如下：
$$\begin{array}{rl}{L}_{vae}& =-E_{z\sim p_{\varphi }(z|x)}[log\frac{p_{\theta }(x,z)}{p_{\varphi }(z|x)}]\\ & =-E_{z\sim p_{\theta }(z|x)}[log{p}_{\theta }(x|z)]+D_{KL}(p_{\theta }(z|x)||p(z))\\ & =MSE(x,p_{\theta }(x,ϵ))+\frac{1}{2}[-1+{{\mu }_{\theta }(x)}^2+{{\sigma }_{\theta }(x)}^2-log{{\sigma }_{\theta }(x)}^2]\end{array}$$
具体实现上，即是Encoder后接两层Linear，分别预测${\mu }_{\theta }(x)和{\sigma }_{\theta }(x{)}^2$，然后通过重参数化技巧，采样一个$x^{\prime }={\mu }_{\theta }(x)+{\sigma }_{\theta }(x)ϵ$输入Decoder，重建x，当然在细节上，我们可以选择预测$log{\sigma }_{\theta }(x{)}^2$，从而避免了网络输出为负的情况。

4 结语

现在准备开始写Diffusion Model的博客，算是一个总结，也算是对学习知识的回顾，学到现在真的得到了太多人博客的帮忙，希望自己也能成其中的一员。

Reference:

苏剑林.《变分自编码器（一）：原来是这么一回事》

苗思奇.《机器学习方法—优雅的模型（一）：变分自编码器（VAE）》