隐变量模型、Auto-Encoder、VAE、VQVAE的学习

1 生成模型以及隐变量模型角度理解

(1) 生成模型是什么都东西？

生成模型的一般定义：给定从真实分布$p(x)$采样得到的数据$x\sim p(x)$，训练得到一个由$\theta$控制逼近真实分布的$p_{\theta }(x)$，则称$p_{\theta }(x)$为生成模型。

(2) 隐变量模型是什么？为什么要用隐变量模型呢？

但实际上我们并不好逼近$p(x)$：1）首先是它的分布过于复杂，而神经网络往往只能拟合一些简单分布；2）其次我们也并不清楚$p(x)$到底是什么分布。从$x$下手不好下手，那么我们可以从更好下手的$z$下手，我们有$p(x)=\int p(x|z)p(z)dz$，其中$p(z)$可以是任意我们指定的分布，这样我们就可以通过中间变量，或者说隐变量$z$去间接求出真实分布了，而这样的模型又叫做隐变量模型。

(3) 如何去度量模型呢？

我们已知$p(z)$【因为是我们自定义的一种分布】，那么只要让神经网络去拟合$p_{\theta }(x|z)$即可，度量拟合好坏的指标可以是他们之间的KL散度，既$KL(p(x|z)||p_{\theta }(x|z))$，但实际上这个也不好求，因为必须要遍历所有的z，而且并不是所有z都有对应的x的。

根据联合分布$p(x,z)=p(x|z)p(z)$，我们也可以去求该分布的KL散度，因为$p(z)$是先验分布，所以求该KL散度也等价于拟合了$p(x|z)$了。
$$\begin{array}{rl}KL(p(x,z)||p_{\theta }(x,z))& =\int \int p(x,z)log\frac{p(x,z)}{p_{\theta }(x,z)}dzdx\\ & =\int p(x)\int p(z|x)log\frac{p(x)p(z|x)}{p_{\theta }(x,z)}dzdx\\ & =E_{x\sim p(x)}[\int p(z|x)log\frac{p(x)p(z|x)}{p_{\theta }(x,z)}dz]\\ & =E_{x\sim p(x)}[logp(x)\int p(z|x)dz+\int p(z|x)log\frac{p(z|x)}{p_{\theta }(x,z)}dz]\end{array}$$
对于第一项，$\int p(z|x)dz=1$,所以有$E_{x\sim p(x)}[logp(x)]=C$（因为p(x)是我们的训练集，相当于固定了，那么求期望当然也是常数）

所以我们有
$$L=KL(p(x,z)||p_{\theta }(x,z))-C=E_{x\sim p(x)}[\int p(z|x)log\frac{p(z|x)}{p_{\theta }(x,z)}dz]$$
最小化$L$，相当于最小化$KL(p(x,z)||p_{\theta }(x,z))$; 也可以理解为最大化$-L$,那么$-L$被称为变分下限(ELBO)。对$L$再进行推导：

$$\begin{array}{rl}L& =E_{x\sim p(x)}[\int p(z|x)log\frac{p(z|x)}{p_{\theta }(x,z)}dz]\\ & =E_{x\sim p(x)}[\int p(z|x)log\frac{p(z|x)}{p_{(}z)p_{\theta }(x|z)}dz\\ & =E_{x\sim p(x)}[-\int p(z|x)log{p}_{\theta }(x|z)dz+\int p(z|x)log\frac{p(z|x)}{p(z)}dz]\\ & =E_{x\sim p(x)}[E_{x\sim p(z|x)}[-log{p}_{\theta }(x|z)]]+E_{x\sim p(x)}[KL(p(z|x)||p(z))]\end{array}$$

对于第一项，是在最小化对数似然函数$-log{p}_{\theta }(x|z)$，也就是在最大化$log{p}_{\theta }(x|z)$,该似然函数叫做decoder。

对于第二项，可见引入了真实后验分布$p(z|x)$，并且希望真实后验分布尽可能逼近先验分布$p(z)$，这样就能使得采样的$z$放入decoder中可以更好地生成图像。我们也需要神经网络去拟合$p_{\theta }(z|x)$，而$p_{\theta }(z|x)$又叫encoder。

至此我们有了隐变量模型的大体框架：用神经网络拟合函数$p_{\theta }(z|x)$和$p_{\theta }(x|z)$，使得最小化$L$（也就是最大化对数似然函数$log{p}_{\theta }(x|z)$以及使得先验分布更接近于真实后验分布，这样就能使得先验分布采样的$z$能和真实的$x$对应起来）。

(4) 为什么可以用KL散度作为优化目标呢？

实际上KL散度和极大似然估计本质上是等价的，我们再对$-L$进行变形：
$$\begin{array}{rl}-L& =C-KL(p(x,z)||p_{\theta }(x,z))\\ & =E_{x\sim p(x)}[\int p(z|x)log\frac{p_{\theta }(x,z)}{p(z|x)}dz]\\ & =E_{x\sim p(x)}[\int p(z|x)log\frac{p_{\theta }(x)p_{\theta }(z|x)}{p(z|x)}dz]\\ & =E_{x\sim p(x)}[log{p}_{\theta }(x)]-KL(p(z|x)||p_{\theta }(z|x))\end{array}$$

第一项就是对数似然，而第二项是大于等于0的数，所以对数似然大于等于变分下限$-L$。严格来说，该KL散度不可能取到0，因为对变分下限优化的目的是让$p_{\theta }(z|x)$与$p(z)$相近，而$p(z|x)$与$p(z)$再形式上不宜与，所以后面的VAE或者扩散模型，损失函数是没有办法真的训练一个最大似然的。

(5）学习过程中比较懵逼的地方

模型优化的目标之一是$p_{\theta }(z|x)$与$p(z)$更接近，为什么上面的优化公式没有体现呢？或者体现在哪个地方我没看出来。

解释如下：

2 自编码器AE

2.1 简单理解AE

通过encoder对数据的维度进行压缩，然后通过decoder升维，得到$x^{\prime }$，$x^{\prime }$存在的价值就是试图让他和$x$相同，那么我们便可以理解为encoder压缩提取的特征真的是最本质的特征。

简单用MSE Loss来构成重建损失：
$$L_{AE}(\theta ,\varphi )=\frac{1}{n}\sum _{i=1}^n(x_i-f(g(x_i)){)}^2$$

2.2 概率角度理解AE

AE的思想就是隐变量$z$和$x$能唯一对应，编码器和解码器的降维和升维相当于高维空间和低维流形之间的映射，属于单隐变量模型（diffusion是多隐变量模型），且隐变量空间连续。

因为唯一对应，所以我们可以认为真实后验分布$p(z|x)=\delta (z-C(x))$, 右式是狄拉克函数，这里的含义是$z$和$C(x)$唯一对应，$C$即为编码器encoder。$C_{\theta }(x)$可以以很简单的用神经网络拟合。

关于似然函数用神经网络去拟合高斯分布 $p_{\theta }(x|z)=N(G_{\theta }(z),{\sigma }^2)$，$G$为decoder，方差设置为定值，所以最后我们用输出的均值替代输出$x$，也就解决了采样问题。

由第一部分我们可以得到$L=E_{x\sim p(x)}[E_{x\sim p(z|x)}[-log{p}_{\theta }(x|z)]]+E_{x\sim p(x)}[KL(p(z|x)||p(z))]$

第一项可以化简为均方根误差

而第二项可以化简为一个与$x,\theta$无关的常数

3 变分自编码器VAE

Auto-Encoder的思想是$x$和$z$是唯一对应的，但生成的本质实际上是一个分布映射到另一个分布，我们不可能将分布里的所有数据点都一一对应上,所以就会出现下面这张图出现的情况(李宏毅ML18讲)，当分布中某一个点没有对应到观测样本，并且我们的模型通常都不是线性的，所以很多地方其实是不可知会生成什么图像的。

3.1 理解VAE

1) VAE和AE的不同点是什么？

VAE则是加上了一层噪声，使得噪声的分布内都能生成同一张图像，这样原本没有对应上的点此时会变成两个分布的重叠部分，生成的图像则会根据两个重叠程度生成合理的图。更本质的理解就是，VAE中每个$x$对应着随机变量$z$，也就是对应着一个分布而不是对应确定值。

2) VAE两种不同视角

若根据混合高斯模型理论，任意复杂的分布$p(x)$均可以由多个高斯分布去进行拟合，既$p(x)={\int }_{z}p(z)p(x|z)dz$ , 其中$p(z)$可以理解为对应的高斯分布的权值, $p(x|z)$理解为该$z$的条件下，对应的高斯分布。

若根据隐变量模型，$z$为隐变量，则$p(x)$可以通过边际化联合概率求得，既$p(x)={\int }_{z}p(x,z)dz$，再由乘法公式$p(x,z)=p(x)p(z|x)=p(z)p(x|z)$代入。

我们的目标是最大化似然函数 ${\sum }_{i=1}^{N}logp(x)={\sum }_{i=1}^{N}log{\int }_{z}p(x,z)dz$，但积分是很难求的。但我们可以用更简单的形式去等价替代。

3.1.1 证据下界(Evidence Lower Bound, ELBO)

ELBO的名字来源是因为$p(x)$在贝叶斯公式里$p(z|x)=\frac{p(x|z)p(z)}{p(x)}$是证据(evidence)，所以叫ELBO。 首先要明确的是： encoder用$q_{\varphi }$表示，decoder用$p_{\theta }$表示，$q_{\varphi }(z)$记为$z$的概率密度函数。

我们推导：
$$\begin{array}{rl}logp(x)& =log{\int }_z{p}_{\theta }(x,z)dz\\ & =log{\int }_z{q}_{\varphi }(z)\frac{p_{\theta }(x,z)}{q_{\varphi }(z)}dz\\ & =log{E}_{z\sim q_{\varphi }(z)}[\frac{p_{\theta }(x,z)}{q_{\varphi }(z)}]\\ & \ge E_{z\sim q_{\varphi }(z)}[log\frac{p_{\theta }(x,z)}{q_{\varphi }(z)}]\\ & ={\int }_z{q}_{\varphi }(z)log{p}_{\theta }(x,z)dz-{\int }_z{q}_{\varphi }(z)log{q}_{\varphi }(z)dz\\ & =E_{z\sim q_{\varphi }(z)}[log{p}_{\theta }(x,z)]-E_{z\sim q_{\varphi }(z)}[log{q}_{\varphi }(z)]\\ & =L(q,\theta )\end{array}$$
其中不等式是根据詹森不等式(琴生不等式)推导出来的，这样我们就可以得到ELBO，记为$L(q,\theta )$

由前面的乘法公式，我们可得联合分布有两种表达形式，既$$p(x,z)=p(x)p(z|x)=p(z)p(x|z)$$
$p(z),p(z|x)$分别为z的先验和后验概率

先代入第一种形式$p(x,z)=p(x)p(z|x)$(省略期望下标$z\sim q_{\varphi }(z)$)：
$$\begin{array}{rl}L(q,\theta )& =E[logp(x)]+E[log{p}_{\theta }(z|x)]-E[log{q}_{\varphi }(z)]\\ & =E[logp(x)]-KL(q_{\varphi }(z)||p_{\theta }(z|x))\\ & =logp(x)-KL(q_{\varphi }(z)||p_{\theta }(z|x))\end{array}$$
第一项由于$p_{\theta }(x)$与隐变量$z$无关所以可以省略掉期望，由第一项得，这里是最大化对数似然； 由第二项得，这里是最小化$q_{\varphi }(z)$和$p_{\theta }(z|x)$的KL散度，也就是
$$logp(x)=L(q,\theta )+KL(q_{\varphi }(z)||p_{\theta }(z|x))$$
只要我们能够将后验$p_{\theta }(z|x)$和$q_{\varphi }(z)$互相逼近，那么证据下界也就等价于似然函数，那么最大化ELBO等价于最大化似然函数。

我们代入第二种形式

由以上可知，当$p_{\theta }(z|x)$和$q_{\varphi }(z)$互相逼近时，ELBO等价于似然函数，所以我们这里设$q_{\varphi }(z)=q_{\varphi }(z|x)=p_{\theta }(z|x)$ (为了区分所以设置为$q(z|x)$)
$$\begin{array}{rl}L(q,\theta )& =E[logp(z)]+E[log{p}_{\theta }(x|z)]-E[logq(z)]\\ & =E[logp(z)]+E[log{p}_{\theta }(x|z)]-E[log{q}_{\varphi }(z|x)]\\ & =E[log{p}_{\theta }(x|z)]-KL(q_{\varphi }(z|x)||p(z))\end{array}$$

第一项，$p_{\theta }(x|z)$可以看作时由$z$生成$x$，也就是decoder； 由于$z$为随机变量(某种分布），所以需要对$z$求期望，而求期望则需要用到后验$q(z|x)$，后验相当于encoder；
第二项，是$q_{\varphi }(z|x)$和$p(z)$的KL散度，KL散度大于等于0，而这里要最大化ELBO，所以要最小化KL散度，也就是让这两个分布更接近，这里相当于约束项。

代入的两项可以得到的信息
由代入后验$p_{\theta }(x|z)$我们可以得知，我们可以用最大化ELBO来代替最大化似然函数，前提是$KL(q_{\varphi }(z)||p_{\theta }(z|x))\to 0$，也就是z的概率密度函数和z的后验概率相近。
由代入条件概率$p_{\theta }(z|x)$我们可以得知，ELBO包含了$p_{\theta }(x|z)$以及$q_{\varphi }(z|x)$,且有约束项$KL(q_{\varphi }(z|x)||p(z))$来使得上面的前提成立，也就是z的概率密度函数和z的先验密度相近，此处z的概率密度函数等同于z的后验概率，也就是确保确定值$x$通过神经网络后计算出来的随机变量$z$所处的分布，和先验分布$p(z)$等价，这样由$z$生成$x$就确保真实。

3.1.2 Encoder和Decoder

$z$是一个随机变量，在VAE中假设$z$服从高斯分布，既先验$P(Z)\sim N(0,1)$

3.1.2.1 Encoder

既然$p(z)$服从高斯分布，那么后验$q_{\varphi }(z|x)$也服从高斯分布，但均值和方差${\mu }_{\varphi },{\Sigma }_{\varphi }$是未知的。 且$Z$是一个随机变量，不能由确定变量$x$直接通过神经网络预测得到，所以通过预测均值和方差间接预测随机变量$z$。既
$$\begin{gathered} {\mu }_{\varphi }(x)=encoder(x) \\ {\Sigma }_{\varphi }(x)=encoder(x) \end{gathered}$$

那么如何衡量两者的好坏呢？预测均值和方差相当于预测分布，所以使用KL散度来衡量好坏，且两个高斯分布的KL散度是有已知公式的，如下
$$KL(N({\mu }_{\varphi },{\Sigma }_{\varphi })||N(0,1))=\frac{1}{2}({\sigma }_{\varphi }^2+{\mu }_{\varphi }-1-log{\sigma }_{\varphi }^2)$$
直觉上理解就是，当均值为0，方差为1时，此项会最小，所以约束其方差为1。

3.1.2.1 Decoder

条件概率$p_{\theta }(x|z)$是从随机变量$z$通过decoder生成随机变量$x$，同理随机变量是不可以直接生成的，所以也通过让神经网络预测均值和方差的方法，只不过这里将方差设置为一个常数（大多数方法也会如此实现），所以只是预测均值即可。既
$${\mu }_{\sigma }=decoder(z)$$
但是这个期望是没有办法直接解析求得的，因为后验概率$q_{\varphi }(z|x)$有一个神经网络的存在，而求期望就是求积分。但可以通过马尔可夫链方式求解近似，就是从后验概率分布随机采样很多个$z$代入求均值，既
$$E_{z\sim q_{\varphi }(z|x)}[log{p}_{\theta }(x|z)]=\frac{1}{L}\sum _{i=1}^L[log{p}_{\theta }(x|z^i)]$$
$L$为超参数，根据经验这里设置为1，但是$z$的采样也是一个随机过程，既$z\sim N({\mu }_{\varphi },{\Sigma }_{\varphi })$，是不可导的（因为同一个输入可能有不同输出），所以VAE作者这里使用参数重整化的方法，从$N(0,1)$采样标准高斯噪声， 然后通过神经网络预测均值和方差，将噪声乘以方差的开根，再加上均值，得到$z$。

3.1.3 汇总

根据高斯公式：
$$\begin{array}{rl}{p}_{\theta }(x|z)& =-\frac{1}{2}exp((x-{\mu }_{\theta }{)}^2)\end{array}$$

所以ELBO公式有：
$$\begin{array}{rl}L(p,\theta )& =E[log{p}_{\theta }(x|z)]-KL(q_{\varphi }(z|x)||p(z))\\ & =-\frac{1}{L}\sum _{i=1}^L(x-{\mu }_{\theta }{)}^2-({\sigma }_{\varphi }^2+{\mu }_{\varphi }-1-log{\sigma }_{\varphi }^2)\end{array}$$

3.2 概率角度理解

VAE认为一个$x$对应一个分布，这个分布又能采样一个点生成$z$，然后通过解码器生成$x$。设先验分布标准高斯分布$N(0,I)$，真实后验分布为$p(z|x)=N(\mu ,{\sigma }^2)$的高斯分布，其中$\mu =C(x)$，其中$C$为encoder，用神经网络拟合就是$p_{\theta }(z|x)=N(C_{\theta }(x),{\sigma }^2)$。

又回到之前推导的公式$L=E_{x\sim p(x)}[E_{x\sim p(z|x)}[-log{p}_{\theta }(x|z)]]+E_{x\sim p(x)}[KL(p(z|x)||p(z))]$

第一项的推导和之前一样可以推导为均方根误差
第二项的推导为：

这也是3.1 给出的那个约束的loss。
所以AE和VAE的区别就是后验分布和先验分布的假设，AE是将后验分布假设为$p(z|x)=\delta (z-C(x))$，且将先验分布假设为$p(z)=\frac{1}{D}$。 VAE则假设为高斯分布（实际上会使用重参数技巧，采样自正态分布）

4 VQVAE

参考：
[1] https://lilianweng.github.io/posts/2018-08-12-vae/
[2] https://zhuanlan.zhihu.com/p/640000410
[3] https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/VQVAE
[4] https://www.spaces.ac.cn/archives/6760
[5] https://zhuanlan.zhihu.com/p/633744455

4.1 AE、VAE和VQVAE

已知AE由输入$x$，经过encoder得到确定值$z$，再经过decoder还原$x^{\prime }$，$x、z$是一一对应的，所以严格来说，AE并不是生成模型，而是一个数据压缩模型。AE无法做到生成任务是因为$z$的分布是不规则的(既我们不能知道$z$服从什么分布)。

VAE就是对AE的一种改进，通过约束$z$满足正态分布，这样解码器就认得编码器编码出的向量，也就能实现生成任务。注意，AE和VAE生成的都是连续变量。

VAE生成的图像质量并不高，VQVAE的作者认为VAE不好的原因是因为编码器生成的连续变量，而如果生成的是离散变量那么生成的图像质量会更好。比如画一个人，他是男(0)或者是女(1)，我们可以用离散的向量进行编码，而不是连续向量(0.7为男，0.3为女）。

VQVAE采用NLP里的embedding的思想，embedding可以看作是特殊的连续向量，将编码器输出的变量$z$与嵌入空间(embedding space, codebook)使用K最邻近算法(距离使用均方根差，然后使用argmin找到与该向量距离最小的索引）找出最接近的嵌入向量索引，然后根据索引用该嵌入向量替代原来的编码器的输出。

现有的问题就是：

1. VAE约束$z$满足正态分布形式，VQVAE又尝试将$z$离散化，这样又会丢失掉编码空间的规范性，而离散化变量是不容易采样的，离散化后VQVAE不就不能生成图像了嘛？素嘟，VQVAE实际上是一个AE，只不过其编码出来的是离散的变量，VQVAE不是一个生成模型，他是一个图像压缩模型，而生成模型是其他的生成模型（pixelCNN、diffussion(构成stabe diffusion)等)

现在尝试理解以下问题：

• 如何输出离散变量？
• 如何优化encoder、decoder？
• 如何优化codebook（不是fixed，而是trainable的）

4.2 输出离散变量

假设输入$x$为 [B,C,N,N], 编码器encoder的输出$z$为 [B,D,M,M] , codebook的尺寸为 [K,D]，表示有K个嵌入向量，每个嵌入向量的维度为D。在计算距离时，先让$z$和codebook扩维（好使用广播机制），既让$z$ 变为 [B,1,D,M,M], 而 codebook变为 [1,K,D,1,1], 然后使用均方根差计算距离，并在第二维（K）使用argmin 找到每个$z$ 对应嵌入向量最小距离的索引，然后得到索引之后我们就可以找到codebook中替代原编码输出的$z_q$了。索引的尺寸为 [B,M,M] 既每个输出都对应一个嵌入向量索引，少了D维是因为这就是嵌入向量的深度。所以$z_q$的尺寸为 [B,D,M,M]

4.3 优化encoder和decoder

argmin是一个离散的操作，所以是不可导的。在AE中，我们的reconstrucion loss为：
$$L=||x-decoder(z)|{|}^2$$

VQVAE中变为：
$$L=||x-decoder(z_p)|{|}^2$$
而这项损失的梯度是不能传导到encoder的，因为$z$到$z_p$之间的变化是离散的。

VQVAE这里使用一种叫Straight-Through Estimator的方法，既设计函数：
$$L=||x-decoder(z-sg(z_p-z))|{|}^2$$
sg表示stop-gradient函数，在前向传播时为sg(x)=x，在反向传播时为0，也就是前向传播还是这么用，反向传播就用了，直接把求得的梯度去给encoder使用。
所以前向传播时： $x-decoder(z_p)$ ， 反向传播时 $L=||x-decoder(z)|{|}^2$

在pytorch中很容易实现：

l = x -  decoder(z - (zp-z).detach()) 

4.4 优化codebook

由前面提到，codebook也是可训练的，那么codebook优化的目标是什么呢？ codebook优化的目标应该是： 尽可能使得codebook中的每一个嵌入向量都能尽可能表示编码器输出的每一类。比如[青年]这个的嵌入向量应该能表示14-48岁的所有编码器输出。

同样使均方根误差来计算，且使用sg函数将loss分成两部分，并配置不一样的权值（作者认为codebook应该学习的比encoder快，原论文参数为$\beta =1,\alpha =0.5$，实验可知，$\beta \in (0.1,2)$结果都差不多。
$$L_e=\beta ||sg(z)-z_p|{|}^2+\alpha ||z-sg(z_p)|{|}^2$$
第一项又叫VQ（vector quantisation)误差，主要作用是优化codebook； 第二项又叫专注（commitment)误差，主要作用是使encoder输出不要偏离codebook太远。

4.5 小结

至此，VQVAE的思想就简单理解了。

5 参考

[1] https://zhuanlan.zhihu.com/p/611466195
[2] https://lilianweng.github.io/posts/2018-08-12-vae/
[3] https://www.youtube.com/watch?v=8zomhgKrsmQ
[4] https://www.zhangzhenhu.com/aigc/变分自编码器.html