VAE解读

VAE(Variational AutoEncoder)变分自编码器

• 要讲扩散模型，不得不提VAE。VAE和GAN一样，都是从隐变量$Z$生成目标数据$X$。

• 它们假设隐变量服从某种常见的概率分布（比如正态分布），然后希望训练一个模型$X=g(Z)$，这个模型将原来的概率分布映射到训练集的概率分布，也就是分布的变换。注意，VAE和GAN的本质都是概率分布的映射。

• tip:关于隐变量解释：

  • 在统计里，隐变量是不可观测的随机变量，我们通常通过可观测变量的样本对隐变量作出推断。
  • 举例1：
    • 高斯混合模型中，GMM中隐变量指的是每个observation对应的高斯component，由于产生过程是不可观测的（或者说隐藏的），故得名隐变量。我们可以通过收集样本对隐变量的后验概率进行推断，然后用估计的后验概率来对数据进行聚类。
  • 举例2：
    • 一个人拿着n个袋子，里面有m种颜色不同的球。现在这个人随机地抓球，规则如下：1. 先随机挑一个袋子2. 从这个袋子中随机挑一个球如果你站在这个人旁边，你目睹了整个过程：这个人选了哪个袋子、抓出来的球是什么颜色的。然后你把每次选择的袋子和抓出来的球的颜色都记录下来（样本观察值），那个人不停地抓，你不停地记。最终你就可以通过你的记录，推测出每个袋子里每种球颜色的大致比例。并且你记录的越多，推测的就越准（中心极限定理）。然而，抓球的人觉得这样很不爽，于是决定不告诉你他从哪个袋子里抓的球，只告诉你抓出来的球的颜色是什么。这时候，“选袋子”的过程由于你看不见，其实就相当于是一个隐变量。隐变量在很多地方都是能够出现的。现在我们经常说的隐变量主要强调它的“latent”。所以广义上的隐变量主要就是指“不能被直接观察到，但是对系统的状态和能观察到的输出存在影响的一种东西”。所以说，很多人在研究隐变量。以及设计出各种更优(比如如可解释、可计算距离、可定义运算等性质)的隐变量的表示。
      

• 换句话说，大致意思就是先用某种分布随机生成一组隐变量，然后这个隐变量会经过一个生成器生成一组目标数据。VAE和GAN都希望这组数据的分布$\hat{X}$和目标分布$X$尽量接近。

• 但是这种方法本质上是难以实现，因为“尽量接近”并没有一个确定的关于$\hat{X}$和$X$的相似度的评判标准。换句话说，这种方法的难度就在于，必须去猜测“它们的分布相等吗”这个问题，而缺少真正interpretable的价值判断。

  • `为什么用KL散度不行？
  • KL散度的公式:$$D_{KL}(A||B)=\sum _i{P}_A(x_i)\log (\frac{P_A(X_i)}{P_B(x_i)})=\sum _i[P_A(x_i)\log (P_A(x_i))-P_A(x_i)\log (P_B(x_i))]$$
  • 从公式可以一目了然，因为KL散度是针对两个已知的概率分布求相似度的，而$\hat{X}$和$X$的概率分布目前都是未知。

• GAN的做法就是直接把这个度量标准也学过来就行，相当生猛。但是这样做的问题在于依然不interpretable，非常不优雅。VAE的做法就优雅很多了，我们先来看VAE是怎么做的，理解了VAE以后再去理解Diffussion就很自然了。

到底什么是生成模型

• 拿到一批sample（称为$X$），想要用$X$学到它的分布$P(X)$ ，这样就能同时学到没被sample到的数据了，用这个分布$P(X)$就能随意采样，然后获得生成结果。但是这个分布很难实现。所以绕个弯，整一个隐变量$Z$，这东西可以生成$Z$。不妨假设$Z$满足正态分布，那就可以先从正态分布里面随便取一个$Z$，，然后用$Z$和$X$的关系算出$P(X)$，公式：$$P(X)=\sum _{Z}P(X|Z)P(Z)$$
• 公式中$X$是结果，$Z$是因，所以$P(X|Z)$为似然估计，$P(Z)$是先验
• 换句话说，就是不直接求$P(X)$，而是造一个别的变量（好听的名字叫“隐变量”），获得这个隐变量和$X$的关系，也能得到$P(X)$，这样也可以得到一个很棒的生成模型。
• 接下来就是结合自编码器来实现重构，保证有效信息没有丢失，再加上一系列的推导，最后把模型实现。

VAE核心

• 在整个 VAE 模型中，我们并没有去使用 $P(Z)$（先验分布,隐变量空间的分布）是正态分布的假设，我们用的是假设 $P(Z|X)$（后验分布）是正态分布。

• 具体来说，给定一个真实样本 $X_k$，我们假设存在一个专属于 $X_k$ 的分布 $P(Z|X_k)$（后验分布），并进一步假设这个分布是（独立的、多元的）正态分布。

• 为什么要强调“专属”呢？

  • 因为我们后面要训练一个生成器 $X=g(Z)$，希望能够把从分布 $P(Z|X_k)$ 采样出来的一个 $Z_k$ 还原为 $X_k$。
  • 如果假设$P(Z)$是正态分布，然后从$p(Z)$中采样一个$Z$，那么我们怎么知道这个$Z$对应于哪个真实的$X$呢？
  • 现在$p(Z|X_k)$专属于$X_k$，我们有理由说从这个分布采样出来的$Z$应该要还原到$X_k$中去。

• 这时，每一个$X_k$都配上了一个专属的正态分布，，才方便后面的生成器做还原。但这样有多少个$X$就有多少个正态分布了（$k$个$X$sample，就有$k$个正态分布$P(X_k|Z)$）。

• 我们知道正态分布有两组参数：均值$\mu$和方差${\sigma }_2$（多元的话，它们都是向量），那怎么找出专属于$X_k$的正态分布$P(Z|X_k)$的均值和方差呢？好像并没有什么直接的思路。那好吧，那我就用神经网络来拟合出来吧！这就是神经网络时代的哲学：难算的我们都用神经网络来拟合。

  • 这里的拟合与GAN不同，本质上是在学习$X_k$和$Z$的关系，而非学习比较$X$与$\hat{X}$的标准

• 于是我们构建两个神经网络${\mu }_k=f_1(X_k),\log {\sigma }_k^2=f_2(X_k)$来算它们了，。我们选择拟合$\log {\sigma }_k^2$而不是直接拟合${\sigma }_k^2$，是因为${\sigma }_k^2$总是非负的，需要加激活函数处理，而拟合$\log {\sigma }_k^2$不需要加激活函数，因为它可正可负。

• 到这里，就能知道专属于$X_k$的均值和方差了，也就知道它的正态分布长什么样了，然后从这个专属分布中采样一个$Z_k$出来，然后经过一个生成器得到${\hat{X}}_k=g(Z_k)$，现在我们可以放心地最小化$D({\hat{X}}_k,X_k{)}^2$，因为$Z_k$是从专属$X_k$的分布中采样出来的，这个生成器应该要把开始的$X_k$还原回来。
  

• 图中均值方差计算模块就是两个神经网络：${\mu }_k=f_1(X_k),\log {\sigma }_k^2=f_2(X_k)$

• 均值和方差的计算本质上都是encoder，VAE其实利用了两个encoder去分别学习均值和方差

• 生成器是通过最小化$D({\hat{X}}_k,X_k{)}^2$来训练的，最终会使得$X_k$和${\hat{X}}_k$趋向一致

• 因为$Z_k$是通过重新采样过的，，而不是直接通过均值和方差encoder算出来的，所以这个生成器的输入$Z$是有噪声的，也就是重构过程受噪声影响

• 显然噪声会增加重构的难度，不过好在这个噪声强度（也就是方差决定的）通过一个神经网络算出来的，所以最终模型为了重构得更好(也就是最终会使得$X_k$和${\hat{X}}_k$趋向一致)，肯定会想尽办法让方差为$0$。而方差为$0$的话，也就没有随机性了，所以不管怎么采样其实都只是得到确定的结果（也就是均值，根据方差公式就可以得出每个采样都是均值），只拟合一个当然比拟合多个要容易，而均值是通过另外一个神经网络算出来的，模型则会慢慢退化成普通的AutoEncoder，噪声不再起作用，VAE就变成AE了

• VAE就是让所有的$P(Z|X)$趋于标准的正态分布$N(0,I)$，这样就防止了噪声为零，同时保证了模型具有生成能力。

• 怎么理解“保证了生成能力”呢？

  • 根据定义：$$P(Z)=\sum _{X}P(Z|X)P(X)=\sum _{X}N(0,I)P(X)=N(0,I)\sum _{X}P(X)=N(0,I)$$
  • 其中${\sum }_{X}P(X)=1$，$P(Z|X)\sim N(0,I)$

• 这样我们的先验假设得证：$P(Z)$是标准的正态分布。然后就可以从$N(0,I)$中采样来生成图像了
  

• 那怎么让所有的$p(Z|X)$都向$N(0,I)$看齐呢？

  • 最直接的方法应该是在重构误差的基础上中加入额外的loss：$$L_{\mu }=\Vert f_1(X_k){\Vert }^2$$ $$L_{{\sigma }^2}=\Vert f{2}_{(}{X}_k){\Vert }^2$$
  • 因为它们分别代表了均值${\mu }_k$和方差的对数$\log {\sigma }_k^2$，达到$N(0,I)$就是希望二者尽量接近于$0$。

• 不过，这又会面临着这两个损失的比例要怎么选取的问题，选取得不好，生成的图像会比较模糊。所以，原论文直接算了一般（各分量独立的）正态分布与标准正态分布的KL散度$KL(N(\mu ,{\sigma }^2)\Vert N(0,I))$作为这个额外的loss，计算结果为:
  KaTeX parse error: No such environment: equation at position 7: \begin{̲e̲q̲u̲a̲t̲i̲o̲n̲}̲ \begin{aligned…

• 下面分别里面的三项：

  • 第一项：$-{\sum }_i^{d}P(x)\log {\sigma }^2$实际上就是$\log {\mu }^2$乘以概率密度求和为1，所以就是$-log{\sigma }^2$
  • 第二项：正态分布的二阶矩，正态分布的二阶矩为${\mu }^2+{\sigma }^2$
  • 第三项：-方差除以方差，所以等与$-1$
  • 所以$KL(N(\mu ,{\sigma }^2)\Vert N(0,1))=\frac{1}{2}(-\log {\sigma }^2+{\mu }^2+{\sigma }^2-1)$

• 即$$Los{s}_{\mu ,{\sigma }^2}=\frac{1}{2}\sum _{i=1}^d(-\log {\sigma }_i^2+{\mu }_i^2+{\sigma }_i^2-1)$$

• 这里的$d$是隐变量$Z$的维度，而${\mu }_i$和${\sigma }_i^2$分别代表一般正态分布的均值向量和方差向量的第$i$个分量。直接用这个式子做补充$loss$，就不用考虑均值损失和方差损失的相对比例问题了。显然，这个$loss$也可以分两部分理解：
  $$Los{s}_{\mu ,{\sigma }^2}=L_{\mu }+L_{\sigma }^2$$
  $$L_{\mu }=\frac{1}{2}\sum _{i=1}^d=\frac{1}{2}\Vert f_1(X){\Vert }^2$$
  $$L_{\sigma }^2=\frac{1}{2}\sum _{i=1}^d\left(-\log {\sigma }_i^2+{\sigma }_i^2-1\right)$$

  VAE的本质是什么？

• 在VAE中，它的Encoder有两个，一个用来计算均值，一个用来计算方差，这已经让人意外了：Encoder不是用来Encode的，是用来算均值和方差的，还是用神经网络计算的

• 它本质上就是在我们常规的自编码器的基础上，对encoder的结果（在VAE中对应着计算均值的网络）加上了“高斯噪声”（正态分布的随机采样），使得结果decoder能够对噪声有鲁棒性；而那个额外的KL loss（目的是让均值为0，方差为1），事实上就是相当于对encoder的一个正则项，希望encoder出来的东西均有零均值，为了防止噪声消失，将所有$P(Z|X)$趋近于标准正态分布，将encoder的均值尽量降为 0，而将方差尽量保持住。这样一来，当decoder训练的不好的时候，整个体系就可以降低噪声；当decoder逐渐拟合的时候，就会增加噪声。

• 那另外一个encoder（对应着计算方差的网络）的作用呢？它是用来动态调节噪声的强度的。直觉上来想，当decoder还没有训练好时（重构误差远大于KL loss），就会适当降低噪声（KL loss增加），使得拟合起来容易一些（重构误差开始下降）；反之，如果decoder训练得还不错时（重构误差小于KL loss），这时候噪声就会增加（KL loss减少），使得拟合更加困难了（重构误差又开始增加），这时候decoder就要想办法提高它的生成能力了。

我们希望X经过编码后，Z的分布都具有零均值和单位方差，这个“希望”是通过加入了KL loss来实现的,如果现在多了类别信息Y，我们可以希望同一个类的样本都有一个专属的均值${\mu }^Y$（方差不变，还是单位方差），这个${\mu }^Y$让模型自己训练出来。这样的话，有多少个类就有多少个正态分布，而在生成的时候，我们就可以通过控制均值来控制生成图像的类别。事实上，这样可能也是在VAE的基础上加入最少的代码来实现CVAE的方案了，因为这个“新希望”也只需通过修改KL loss实现:
$$Los{s}_{\mu ,{\sigma }^2}=\frac{1}{2}\sum _{i=1}^d(-\log {\sigma }_i^2+({\mu }_i-{\mu }_i^Y{)}^2+{\sigma }_i^2-1)$$

参考文献

苏剑林