简单理解VQGAN
TL; DR:与 VQVAE 类似,隐层压缩表征+自回归生成的两阶段图像生成方法。增加感知损失和对抗损失,提高压缩表征模型解码出图片的清晰度。还可以通过编码并预置条件表征,实现条件生成。
隐层压缩表征+自回归生成:如何改进?
我们在之前介绍 vqvae 的文章中提到,vqvae 本身实际是一种图像的压缩表征模型,其编码器可以将真实的像素图片压缩为隐空间的小特征图,而其解码器可以将隐空间的特征小图重构为真实的像素图片。作为一种压缩表征模型,vqvae 本身并不能进行图像生成,需要根据训好的 vqvae 再训练一个自回归模型(如 PixelCNN),来对隐空间特征图进行采样,再用 vqvae 的解码器解码为新的真实图片。
vqvae 这种隐层压缩表征+自回归生成的两阶段图像生成方法是否还存在改进空间呢?答案当然是肯定的。实际上当时压缩表征和自回归生成两个阶段都有可以借鉴的改进思路。
首先,在隐层表征阶段,简单的 L2 重构损失只约束重构结果与原图像素值之间的差异,而无法约束图片感知效果的差异,会导致生成图片非常模糊,要考虑如何改进。比如增加更全面的监督信号的损失函数,如 LPIPS、对抗损失等。
然后,在自回归生成阶段,PixelCNN 还是太过简单了,要考虑是否存在更强力的模型。在当时,LLM 还没有展现出绝对的统治力,但是以 OpenAI 的 GPT 系列为代表的语言模型已经取得很大进展,并且其自回归的形式也十分适合用于隐空间特征图的采样。另外,当时 ViT 等 Transformer 在视觉领域的工作也已经出来了,相较于 CNN,Transformer 不存在局部性的归纳偏置,对图像全局构图的理解能力更好,并且表达能力更强,在大数据量下性能更好。
vqgan 这篇工作,正是基于上述几个点对 vqvae 做出了改进。
隐层感知压缩表征
在隐层压缩表征的训练上,vqvae 使用了 reconstruction loss、embedding loss 和 commitment loss 三部分损失。vqgan 整体上沿用了 vqvae 的思路,但是为了使 codebook 的语义感知更加丰富,同时避免单独的 L2 损失导致的解码图片模糊的问题,vqgan 做出了两个改进。
一是增加了感知损失,即 LPIPS。简单来说就是将重构图片过一个训练好的 VGG 网络,提取其各层特征图,要求其与原图的各层特征也尽可能接近。而不只是像 L2 损失那样,简单地要求生成结果的像素值与原图接近。这样一来,除了纹理特征之外,高层语义特征的一致性也得到了监督。
二是增加了对抗损失,具体来说,是一个 PatchGAN。PatchGAN 不止是要求辨别器分辨出真实图片和生成图片,而是更细粒度地要求分辨出每一个小图像块是真实的还是生成的。
从下图的对比中可以看到,增加了这两个损失的监督,VQGAN 的解码结果明显比 VQVAE 更加清晰锐利,没有那么模糊。
基于Transformer的自回归生成
无条件生成
训练好 vqgan 之后,需要训练一个自回归模型来进行采样。这一阶段的训练其实就是常见的自回归语言模型,根据已有的 token,去预测下一个 token,即 p ( s i ∣ s < i ) p(s_i|s_{<i}) p(si∣s<i) ,去最大化这个 likelihood p ( s ) = ∏ i p ( s i ∣ s < i ) p(s)=\prod_{i}p(s_i|s_{<i}) p(s)=∏ip(si∣s<i) 。
条件生成
对于给定条件 c c c 的条件生成,vqgan 的做法就像文本生成领域其他的 decoder-only 模型类似,就是将条件放到最前面即可,此时 likelihood 为 p ( s ∣ c ) = ∏ i p ( s i ∣ s < i , c ) p(s|c)=\prod_ip(s_i|s_{<i},c) p(s∣c)=∏ip(si∣s<i,c) 。注意对于条件的编码,需要重新训练一个新的 vqgan。
总结
vqgan 利用 Transformer 和感知损失等技术,从生成质量和可控生成上对 vqvae 进行了很有意义的改进。vqvae、vqgan这一路压缩表征的方法在 LDM、MultiModal Transformer 中都有很多的应用。