标准的AE由编码器(encoder)和解码器(decoder)两部分组成,。整个模型可以看作一个“压缩”与“解压”的过程:
首先编码器将真实数据(真实样本)压缩为低维隐空间中的一个隐向量,该向量可以看作输入的“象征”;然后解码器将这个隐向量解压,得到生成数据(生成样本)。
在训练过程中,会将生成样本与真实样本进行比较,朝着减小二者之间差异的方向去更新编码器和解码器的参数,最终目的是期望由真实样本压缩得到的隐向量能够尽可能地抓住输入的精髓,使得用其重建出的生成样本与真实样本尽可能接近。AE可应用于数据去噪、可视化降维以及数据生成、图像检索等方向。
编码器和解码器可由全连接层或卷积层组成
VAE是AE的升级版本,其结构也是由编码器和解码器组成。AE在生成数据时只会模仿而不会创造,无法直接生成任意的新样本,这是因为AE在生成样本时用到的隐向量其实是真实样本的压缩编码,也就是说每一个生成样本都需要有对应的真实样本,AE本身无法直接产生新的隐向量来生成新的样本。作为AE的重要升级,VAE的主要优势在于能够产生新的隐向量,进而生成有效的新样本。VAE能够生成新样本(即VAE与AE的最大区别)的原因是,VAE在编码过程中加入了一些限制,迫使编码器产生的隐向量的后验分布尽量接近某个特定分布(如正态分布)。VAE训练过程的优化目标包括重构误差和对后验分布的约束这两部分。VAE编码器的输出不再是隐空间中的向量,而是所属正态分布的均值和标准差,然后再根据均值与标准差来采样出隐向量。由于采样操作存在随机性,每一个输入图像经过VAE得到的生成图像不再是唯一的,只要是从隐空间的正态分布中采样得到的,生成的图像就是有效的。
在代码实现过程中,模型通过编码器生成正态分布的均值和方差,采取重参数化的方式进行采样
VQ-VAE相比VAE有两个重要的区别:首先VQ-VAE采用离散隐变量,而不是像VAE那样采用连续的隐变量;然后VQ-VAE需要单独训练一个基于自回归的模型如PixelCNN来学习先验(prior),而不是像VAE那样采用一个固定的先验(标准正态分布)。
VQ-VAE与VAE的最主要的区别是VQ-VAE采用离散隐变量,对于encoder的输出通过向量量化(vector quantisation,VQ)的方法来离散化。至于为啥采用离散编码,论文的slides给出了三个主要原因:
原来的VAE的后验分布是多元高斯分布,但对于VQ-VAE,经过VQ之后,后验分布可以看成是一个多类分布(categorical distribution),而且其概率分布为one-hot类型,此时,后验分布其实变成了一个确定的分布,因为确定的过程没有任何随机因素。此时就可以直接简单计算出后验分布和先验分布的KL散度就是一个常量,那么VQ-VAE的训练损失就剩下了一项重建误差。实际上VQ-VAE的训练过程就没有用到先验分布,所以后面我们需要单独训练一个先验模型来生成数据,这是VQ-VAE和VAE的第二个区别。VQ-VAE分成两个阶段来得到生成模型,可以避免VAE训练过程中容易出现的“posterior collapse”。
VQ-VAE还存在一个问题,那就是由于argmin操作不可导,所以重建误差的梯度就无法传导到encoder。论文采用straight-through estimator来解决这个问题,所谓straight-through estimator其实就是一种用来估计一些不可导函数梯度的方法,如下图所示,threshold function是不可导的,此时我们在计算梯度时,直接忽略它而采用上游得到的梯度,这个行为就认为threshold function是一个identity function一样。
虽然通过straight-through estimator方法,重建误差的梯度可以传导到encoder,但是embedding向量就接收不到重建误差带的梯度了,这也意味着embedding空间无法参与学习。为了让embedding空间参与训练,论文采用了一种简单的字典学习方法,即计算encoder的输出和对应的量化得到的embedding向量的L2误差。我们也可以采用另外一种方式:指数移动平均(exponential moving averages,EMA)来更新embedding空间。
除此之外,论文还额外增加一个训练loss:commitment loss,这个主要是约束encoder的输出和embedding空间保持一致,以避免encoder的输出变动较大(从一个embedding向量转向另外一个)。commitment loss也比较简单,直接计算encoder的输出
和对应的量化得到的embedding向量的L2误差。
综上,VQ-VAE共包含三个部分的训练loss:reconstruction loss,VQ loss,commitment loss。其中reconstruction loss作用在encoder和decoder上,VQ loss用来更新embedding空间(也可用EMA方式),而commitment loss用来约束encoder。
训练好VQ-VAE后,还需要训练一个先验模型来完成数据生成,对于图像来说,可以采用PixelCNN模型,实际上,PixelCNN不是唯一一种可用的拟合离散分布的模型。我们可以把它换成Transformer,甚至是diffusion模型。这里我们不再是学习生成原始的pixels,而是学习生成离散编码。首先,我们需要用已经训练好的VQ-VAE模型对训练图像推理,得到每张图像对应的离散编码;然后用一个PixelCNN来对离散编码进行建模,最后的预测层采用基于softmax的多分类,类别数为embedding空间的大小。那么,生成图像的过程就比较简单了,首先用训练好PixelCNN模型来采样一个离散编码样本,然后送入VQ-VAE的decoder中,得到生成的图像。
GAN是专门为了优化生成任务而提出的模型。生成模型的一大难点在于如何度量生成分布与真实分布的相似度。一般情况下,我们只知道这两个分布的采样结果,很难知道具体的分布表达式,因此难以找到合适的度量方法。GAN的思路是,把这个度量任务交给一个神经网络来做,这个网络被称为判别器(Discriminator)。GAN在训练阶段用对抗训练方式来交替优化生成器与判别器。整个模型的优化目标是
网络在参数更新过程中,对于判别网络:问题变成了有监督的二分类问题了(判别真样本集的loss+判别假样本集的loss)。对于生成网络:当通过原始的噪声数组Z生成了假样本后,把这些假样本的标签都设置为1,即认为这些假样本在生成网络训练的时候是真样本。因为此时是通过判别器来生成误差的,而误差回传的目的是使得生成器生成的假样本逐渐逼近为真样本(当假样本不真实,标签却为1时,判别器给出的误差会很大,这就迫使生成器进行很大的调整;反之,当假样本足够真实,标签为1时,判别器给出的误差就会减小,这就完成了假样本向真样本逐渐逼近的过程),起到迷惑判别器的目的。现在对于生成网络的训练,有了样本集(只有假样本集,没有真样本集),有了对应的label(全为1),有了误差,就可以开始训练了。在训练这个串接网络时,一个很重要的操作是固定判别网络的参数,不让判别网络参数更新(使用detach对其进行分离,将其变为叶子节点,避免其在G中向后回传梯度。),只是让判别网络将误差传到生成网络,更新生成网络的参数。在生成网络训练完后,可以根据用新的生成网络对先前的噪声Z生成新的假样本了,不出意外,这次生成的假样本会更真实。有了新的真假样本集(其实是新的假样本集),就又可以重复上述过程了。整个过程就叫单独交替训练。可以定义一个迭代次数,交替迭代到一定次数后停止即可。
■AE和VAE的联系与区别
AE和VAE二者都属于有向图模型,模型的目的都是对隐变量空间进行建模;但AE只会模仿而不会创造,VAE则可以根据随机生成的隐向量来生成新的样本。
AE的优化目标是最小化真实样本与对应的生成样本之间的重构误差,但在VAE中,除了考虑重构误差之外,还加入了对隐变量空间的约束目标。
AE中编码器的输出代表真实样本对应的隐向量,而VAE中编码器的输出可以看作由两个部分组成:一部分是隐向量所对应的分布的均值;另一部分是标准差。计算均值的编码器就相当于AE中的编码器,而计算标准差的编码器相当于为重构过程增加噪声,使得解码器能够对噪声更为鲁棒(当噪声为0时,VAE模型就退化成AE)。
在VAE的损失函数中,重构损失的目的是降低真实样本和生成样本之间的差异,而迫使隐向量后验分布接近正态分布,实际上是增加了生成样本的不确定性,两种损失相互对立。这与GAN一样,它们内部都存在对抗思想,只不过VAE是将两部分同步优化的,而GAN则是交替优化的。
与AE相同的是,GAN的优化目标只涉及生成样本和真实样本之间的比较,没有VAE中对后验分布的约束。不同的是,GAN设计了判别器,并用对抗训练方式绕过了对分布间距离的度量,且在判断样本真假时不需要真实样本与生成样本一一对应(而在AE/VAE中都需要二者一一对应才能计算重构误差)。
GAN没有像AE那样从学习到的隐向量后验分布中获得生成样本的能力,可能因此导致模式坍塌、训练不稳定等问题。另外,在AE/VAE中隐向量空间是数据的压缩编码所处的空间,隐向量用“精炼”的形式表达了输入数据的特征。如果我们想在抽象的语义层次上对数据进行操控,比如改变一张图像中人的发色,直接在原始数据空间中很难操作,而VAE在隐空间的表达学习能力,使得可以通过在隐变量空间上的插值或条件性嵌入等操作来实现对数据在语义层次上的操控。