[论文速览] AE,VAE,VQ-VAE,VQ-GAN,FSQ
Pre
ref: 《An Introduction to Autoencoders》
ref: https://zhuanlan.zhihu.com/p/388620573
ref: https://www.spaces.ac.cn/archives/5253
ref: https://zhuanlan.zhihu.com/p/388620573
ref: https://www.spaces.ac.cn/archives/6760
ref: https://zhuanlan.zhihu.com/p/68339385
ref: https://zhuanlan.zhihu.com/p/515214329
ref: https://spaces.ac.cn/archives/9826
约定:以下用x代表输入,dist代表距离度量,enc代表编码器,dec代表解码器
关键词:autoencoder, representations, quantization, image synthesis
阅读理由:被MAGE里用到的VQGAN吸引来的,决定梳理一下autoencoder系列的发展历程
AE(AutoEncoder)
title:
accepted:
paper:
code:
最早由 Rumelhart, Hinton, Williams 在1986年的《Learning Internal Representations by Error Propagation》提出
Idea
将输入编码为压缩的表达,并将其解码并重建回原样。无监督,泛化性强,可做特征提取或数据降维等

Discuss
AE应能完美重建输入,但不希望它只是将输入原样吐出,通常可以通过创建bottleneck,并以某种方式加入正则,让它学到有用且有意义的表达。
bottleneck: 下采样将特征降维,
正则: \(argmin_{enc,dec}\; dist(x, dec(enc(x))) + \lambda \sum_i \theta^2_i\)
VAE(Variational AutoEncoder)
title: Auto-Encoding Variational Bayes
accepted: ICLR 2014
paper: https://arxiv.org/abs/1312.6114
code: https://github.com/AntixK/PyTorch-VAE (non-official)
Idea
跟GAN一样,都希望进行分布之间的变换,跟AE的差别在于enc输出不再是单个隐向量z,而是z的一个分布
Background
实际上AE也可以从隐空间采样去做生成,但是它的隐空间不连续,不同标签之间的隐向量存在空隙,去做插值就会发现效果很差

Method
对一批x,用X描述,如果能得到它的分布\(p(X)\),便可通过采样得到所有可能的x,想法很好但很难实现。
退一步,将分布改成
那么这就是一个用Z生成X的模型,可以先采样z,再根据它生成对应的x,也很不错
通常假设后验分布 \(p(Z|X)\) (下图中央绿色的)符合正态分布,对于给定的每个真实样本 \(x_k\),都去学习一个它专属的后验分布 \(p(z|x_k)\)。这样 \(p(z|x_k)\) 采样出来的 \(z_k\) 跟样本 \(x_k\) 是一一对应的,有理由相信它能用于重建 \(x_k\)
实际中 \(p(z|x_k)\) 通过两个网络分别去计算均值跟方差,有示意图如下:

vae为每个样本构造专属的正态分布,然后采样来重构
看上图,由于目标仍然是重构x,但因为噪声(正态分布的方差)的存在会干扰重建,如果不干预,模型会倾向于将方差预测为0,退化为AE,这样重建结果仅由均值确定,结果唯一,失去了多样性
因此实际上VAE还约束后验分布 \(p(Z|X)\) 接近标准正态,防止噪声为零,此时根据定义:
这样 \(p(Z)\) 也是标准正态分布,可以直接采样去生成图片 \(\hat{x_k} = g(z_k)\)。总结一下
\(p(z|x_k)\) 是图中间绿色的正态分布1~6, \(p(Z)\) 是采样出来的z_1~z_6的分布,跟 \(p(z|x_k)\) 一样, \(p(X|Z)\) 指生成器,之后仍然按照公式1去生成 \(x_k\)。也就是说训练的时候用x通过后验分布 \(p(Z|X)\) 去得到z,而推理时就可以直接采样出z去做生成
为了让 \(p(Z|X)\) 接近 \(\mathcal{N}(0,I)\),最简单可以分别加个损失:
分别约束均值跟方差对数 \(\log \sigma^2_k\) 为0,但二者的权重不好选取,生成图片容易模糊,原论文使用KL散度作为loss:
上面公式4里,d是隐变量Z的维度,该loss又可拆分成两个部分:
重参数技巧(reparameterization trick),因为 \(p(Z|X)\) 由模型拟合,采样 \(z_k\) 的操作不可导,无法直接优化模型,但通过:
可知 \(\frac{z-\mu}{\sigma}=\epsilon\) 服从标准正态分布,那么从 \(\mathcal{N}(\mu,\sigma^2)\) 中采样一个z,相当于从 \(\mathcal{N}(0,I)\) 中采样一个 \(\epsilon\) ,然后根据 \(z=\mu + \varepsilon \times \sigma\) 就可得到z
Discuss
VAE虽然也称是AE的一种,但跟AE不太一样,它的enc用于计算均值跟方差,而不是对输入进行编码。在实际实现中,就是在计算均值的enc输出中加入噪声,对应从 \(\mathcal{N}(\mu,\sigma^2)\) 采样的过程,噪声的存在使得学到的隐空间连续、稠密,或者说对噪声更有鲁棒性(更多的噪声可以对应到图片上)

vae的本质结构
由于重构是希望去掉噪声,而KL loss又希望加入噪声,这实际上存在对抗关系,但二者是共同进化,跟GAN不同
根据ref3, \(KL\Big(p(x)\Big\Vert q(x)\Big)\) 是个泛函,VAE的V体现在推导过程中用到了变分法。
至于条件信息的加入,即CVAE,可通过令每个类有一个专属的均值实现,有点AdaIN的感觉

一个简单的cvae结构
对应修改KL loss如下:
VAE形式优美,但是生成的图片相对模糊
VQ-VAE(Vector Quantized-Variational AutoEncoder)
title: Neural Discrete Representation Learning
accepted: NeurIPS 2017
paper: https://arxiv.org/abs/1711.00937
code: None
Idea
维护一个codebook,通过将CNN编码的特征 \(z_e\) 与codebook里的向量进行比较,从而将 \(z_e\) 转变为离散的 \(z_q\) 并使用 \(z_q\) 去重建图片。整体类似将编码特征进行聚类,减小特征向量的种类,以便于使用自回归模型去生成图片
Motivation&Solution
- 将图片逐像素看作离散的句子可以用自回归模型进行生成,但对于大图片开销很大 —— 降维再建模
- 逐像素方案严格区分两个像素值,即便十分接近,不符合人类直觉——像素100跟99相差不大 —— 聚类,减小类别数
Background
VAE 想学到 连续(continuous)的隐空间,而 VQ-VAE 想学到 离散(discrete)的隐空间。
PixelRNN、PixelCNN之类的自回归模型将图片看作离散的,比如32x32x3的图可看作长度为32x32x3=3072的句子,词汇表大小为256,用语言模型的方法逐像素地生成图片,即自回归
一方面递归顺序很重要,“从左往右再从上到下”、“从上到下再从左往右”、“先中间再四周”或者是其他顺序,对生成结果都有很大影响。另一方面是对采样过程加速的研究,由于逐像素生成,速度很慢,而且序列太长不容易捕捉依赖。还有一点是邻近的像素值99/100并没有太大区别(视觉上),但作为分类任务损失会很大。
为了解决这些问题,VQ-VAE希望先降维,然后再对编码向量用PixelCNN建模
Method
Overview

图1 左边描述VQ-VAE的流程,右边是嵌入空间的可视化,编码器输出z(x)被映射到最近点e_2,梯度\nabla_z L(红色的) 会使得编码器改变输出,可能会改变下一个前向的configuration
K在论文中取512,离散向量维度32x32x1,使用Exponential Moving Averages进行更新
图片先经过enc,得到特征图\(z_e(x)\),\(shape = H' \times W' \times D\):
将 \(H' \times W'\) 个\(D\)维向量分别去codebook里找最近的 \(e_i\),其index构成了青色的 \(q(z|x)\)
按照 \(q(z|x)\) 抽出对应的 \(e_i\),组合成紫色的 \(z_q(x)\),作为dec输入,得到结果\(\hat{x}\)
Learning
loss:
其中sg 是stop gradient 的缩写,表示不对 sg[] 方括号里的变量计算梯度,误差不向此变量传递。论文中\(\beta=0.25\)
第一部分是重构误差,与普通AE不一样,因为使用了量子化,解码器的输入变为量子化后的字典向量 \(z_q\),这一项将同时更新编码器和解码器。
第二部分计算编码器得到的潜在向量与字典向量的距离,并将其作为辅助误差项。此误差项只向字典向量 \(z_q\) 传递,让 \(z_q\) 靠近 \(z_e\) 不更新编码器和解码器。
第三部分与第二部分一致,也是计算潜在向量与字典向量的距离。不过这里对字典向量 \(z_e\) 使用了 stop gradient 约束,让 \(z_e\) 靠近 \(z_q\),以防止 \(z_e\) 在不同的 隐向量 \(e_i\) 间波动(fluctuating),使得此误差项只向编码器反向传递。
整体流程就是:
该过程也可以看作是聚类,将enc结果那些相近的特征用codebook相应的embedding替代。但这其实跟VAE关系不大,更像是一个编码离散向量的AE。
Prior
本文在离散隐变量z上的的先验分布 \(p(z)\) 是一个类别分布(categorical distribution)。在训练VQ-VAE时,先验保持恒定和一致(kept constant and uniform)。然后用PixelCNN建模先验分布,以自回归的方式生成z(跟VAE的直接采样不同),(而且这个z应该是one-hot,相当于整数),然后映射为 \(z_q\),再通过原始采样法用dec解码为图片
作者认为将来可以研究将先验跟VQ-VAE共同训练,可能会提升结果。
[注] categorical distribution: 类别分布,将一个小球放入 k 个桶,每个桶里有球的概率不同。采样出的变量 \(X\) 代表着 k 个桶内的小球个数,所以是一个One-hot向量,只有一个桶里面有球。
ref: https://zhuanlan.zhihu.com/p/59550457
[注] ancestral sampling: 祖先/原始采样法,\(p(x_t \mid x_{t - 1}, \dots, x_1)\) ,即根据之前节点的取值来采样下一个节点,符合自回归生成的思想
ref: https://blog.csdn.net/jiqiujia/article/details/47980693
ref: https://stats.stackexchange.com/questions/532030/why-is-ancestral-sampling-used-in-autoregressive-models
VQ-VAE-2(Vector Quantized-Variational AutoEncoder-2)
title: Generating Diverse High-Fidelity Images with VQ-VAE-2
accepted: NeurIPS 2019
paper: https://arxiv.org/abs/1906.00446
code: None
Idea
VQ-VAE的升级版,使用分层的VQ-VAE,更好地提取信息,并加入rejection sampling来权衡质量跟多样性。
Method

图2 VA-VAE架构 (a)输入256x256,压缩到64x64跟32x32的隐空间,编码器根据这俩特征图进行重建 (b)多阶段图片生成,顶层的PixelCNN建模于类别标签上,而底层则以类标签跟上层编码(first level code)为条件
VQ-VAE-2是VQ-VAE的升级版,主要变化是:
- 对VQ-VAE的encoder和decoder进行分层, bottom层对局部特征进行建模, top层对全局特征进行建模; 在top层中加入了self attention,更有效地提取全局信息,同时缓解了生成图像模糊的问题,保留了 AE 不会模式坍塌,生成样本多样化的优点。
- 自回归的采样会累积误差, 因此加入了rejection sampling来权衡质量跟多样性。即生成图片后, 通过一个ImageNet上预训练的分类器判断生成的图片类别是否符合要求, 不是就丢弃
Metrics&Results
Negative Log-Likelihood(NLL) 它能给出评价泛化性的客观标准,可用于检测过拟合,而其它常用的指标如FID跟 Inception Score 完全忽略了泛化性的问题,同样的问题也出现在近来提出的 Precision-Recall 和 Classification Accuracy Scores

表1 NLL跟MSE值,训练跟验证集上指标的微小差异说明显眼网络跟VQ-VAE都没有过拟合
注意这些NLL值仅在使用相同预训练的VQ-VAE编码器和解码器的先验模型之间是可比较的。

图7 多样性-质量权衡的定量评估
Precision - Recall 它被提出来替代FID跟IS评估GAN表现,希望显式的量化覆盖率(recall)跟质量(precision)之间的权衡。图7-b显示VQ-VAE比BigGAN-deep精准度稍低,但召回率高。
Classification Accuracy Score(CAS) 只用候选模型的样本训练一个ImageNet的分类器,然后在测试集的真实图片上评估它的分类精度,以此评判样本质量和多样性。

表2 CAS指标
第二行对于VQ-VAE,分类器只在样本上训练,缺乏高频信号、噪声等(由于压缩)。第三行用VQ-VAE去重建测试图片,然后分类器在这些图片上的效果就好多了,因为这样缩小了 domain gap。
FID and Inception Score 评估GAN最常用的俩指标,虽然有缺点,但还是贴出了响应结果。
图7表明rejection sampling能提高VA-VAE的FID跟IS指标,FID从30降到了10(应该是说图a里那条蓝色实线),而对于BigGan-deep,rejection sampling的效果好于BigGAN论文中的截断方法(图a里的critic)
图7a还能看出inception classifier对于轻微的模糊或其他VQ-VAE重建带来的扰动十分敏感,看淡紫色跟青色的俩点,FID分别是10跟2。为此也计算了 VQ-VAE samples 和 the reconstructions (which we denote as FID*) 之间的FID,表明 inception 网络统计数据比FID更接近真实图片。这段搞不明白了,怎么突然来了个inception classifier,如果是之前的分类器,怎么去算FID了??
总之应该是说VQVAE重建的图片没有原图那么清晰,对于FID的计算会有比较大的影响
Conclution
提出了使用带一个强力自回归模型作先验的VQ-VAE在生成多样且高分辨率图片的简单方法,其中编解码器架构都保持跟VQ-VAE一样简单轻量,差别在于对增加的分辨率使用了分级的多尺度隐特征。生成样本的忠实度可比肩最好的GAN,同时在几个类别上更加多样
并且相信自己的实验能证明隐空间上的自回归建模对于训练大规模生成模型是一种简单且有效的办法。说是轻量简单,但论文没有提及训练的细节,代码也没开源,无从得知
VQ-GAN
title: Taming Transformers for High-Resolution Image Synthesis
accepted: CVPR 2021
paper: https://arxiv.org/abs/2012.09841
code: https://github.com/CompVis/taming-transformers
Authors: Patrick Esser, Robin Rombach, Björn Ommer
Affiliation: IWR,海德堡大学(Heidelberg University)图像处理合作实验室
Keywords: Transformers, image synthesis, convolutional neural networks (CNNs), high-resolution images, autoregressive models
Idea
VQVAE的变体,用Transformer建模先验分布而不是原本的PixelCNN,属于两阶段方法。先用CNN(VQGAN)去学习上下文丰富的codebook,然后训练Transformer(gpt-2)建模这些离散的code,并利用学到的CNN dec根据它生成的code去生成图片
Motivation&Solution
- CNN更多地是关注到局部特征,计算量也相对小;而Transformer则采用长距离注意力,拥有着全局感受野,但缺乏对局部信息缺乏偏置先验 —— 将二者结合
- Transformer计算效率受到注意力机制的制约,处理不了太大的图片 —— 用离散的token来表示图片而不是原本的像素,并结合滑动窗口注意力降低计算量
Background
CNN更多地是关注到局部特征,计算量也相对小;而Transformer则采用长距离注意力,拥有着全局感受野,但缺乏对局部信息缺乏偏置先验,因此在处理高像素密度的图像时,计算效率上存在问题,难以处理有百万像素的高分辨率图片。
The Transformer Family 略
Convolutional Approaches 略
Two-Stage Approaches 指先学习数据的表征,然后再用其做生成。比如VQVAE、VQVAE-2等,最大能处理的图片最大也就192x192,为了保持学到的离散表征尽可能跟像素保持空间不变,又出现了 shallow VQVAE,它使用更小的感受野。而本文表明强力的第一阶段很有必要,捕捉学到的表征中尽可能多的上下文对于用Transformer做高效高分辨率图片合成很关键。
之前的Transformer最大只能生成64x64的图片,通过降低图片尺寸以减小计算量,然而这是牺牲图像质量来换取效率。文章提出的方法则通过CNNs先学习图像元素的词汇表,并使用Transformers来高效地合成高分辨率图像。这个想法不仅具有合理性,而且产生的效果优美。
Method(Model)
Overview
结合CNNs对图像的inductive bias和Transformers的表达能力,在不牺牲计算效率的情况下实现了高分辨率图像的合成。

图1 本文方法使Transformer能够合成高分辨率图片,这个是1280x460像素

图2 用VQGAN学习codebook,再用自回归transformer架构建模code的组合。离散codebook提供了高压缩率跟高感知质量
如图2所示,实际上模型分为两个部分,首先VQGAN这个模型类似VQVAE,包括enc、dec,都是基于卷积的,另一部分是小改过的GPT2,用于学习生成离散token的分布,结果扔给dec就能做生成,然后再使用PatchGAN作判别器评估生成图片,因此属于GAN
Learning an Effective Codebook of Image Constituents for Use in Transformers
(学一个图片内容的有效编码给Transformer用,即VQGAN的训练)
用CNN去学习一个codebook将图片离散化,这方面跟VQ-VAE是一样的。提到用codebook抽取的向量去重建实际上是不可导的,因此VQ-VAE用到了 straight-through gradient 技术,即将dec的梯度直接复制给enc,这样codebook就可以用VQ-VAE那个损失算:
Learning a Perceptually Rich Codebook 将上述损失第一项 \(\|x-\hat{x}\|_{2}^{2}\) 的L2损失换成了感知损失(LPIPS),然后加入对抗训练思想,整了个patchGAN充当判别器:
整体损失如下:
其中那个 \(\lambda\) 是自适应的:
里面 \({\mathcal{L}}_{\mathrm{rec}}\) 是感知重建损失, \(\mathcal{L}_{\mathrm{VQ}}\) 的第一项。 \(\nabla_{G_{L}}\) 表示输入,也就是解码器最后一层的梯度。 \(\delta\) 用于数值稳定,防止除零吧。
Learning the Composition of Images with Transformers
(用Transformer拟合图片(离散token)的分布,也就是自回归地生成token以供dec解码,因此Transformer要处理的输入尺寸其实不大)
这部分主要训练Transformer,用之前训好的CNN对图片编码,离散化后得到一个token序列,然后随机遮掩掉一些,然后用Transformer学习去预测被遮掩的,损失为交叉熵。
Latent Transformers 跟VQ-VAE好像一样,enc处理结果利用codebook换成对应的离散向量表示,再以某个顺序转换为序列 \(s\),以自回归的方式做图片生成
Conditioned Synthesis 条件图片合成,条件c可以是单一的类别标签,甚至是一张图片,加入c后学习序列的似然:
如果条件信息c具有空间范围(spatial extent),可以学习另一个VQGAN,再获得一个基于索引的表示(index-based representation) \(r \in \{0,\cdots,|\mathcal{Z}_c-1\}\) ,这 \(\mathcal{Z}_c\) 是新学到的一个codebook。那么将r放到s前面去并限制负对数似然对 \(p(s_{i}|s_{<i},r)\) 的计算就好了。
看起来像是设置了预先条件的祖先采样,应该是训练第二阶段的那个Transformer来实现

图3 滑动注意力窗口
Generating High-Resolution Images 通过调整VQGAN的下采样块可以调节输入Transformer的图片尺寸,但下采样次数过多时会导致重建质量下滑。本文通过图3所示的滑动窗口来应用transformer,确保得到的上下文仍足以准确还原图片
Experiment
Training Detail
设定codebook里向量维度 \(|\mathcal{Z}| = 1024\),图片大小256x256,并且预测的序列长度为 \(16\cdot 16\),也就是说编码后的特征图大小为 \(16\cdot 16\),这是12GB显存GPU能训练 GPT2-medium 架构(307Mparameters) 的最大可行长度
Dataset
无条件图片建模:ImageNet (IN), Restricted ImageNet(RIN), LSUN Churches and Towers (LSUN-CT)
条件图片建模:D-RIN(基于深度图的RIN), S-FLCKR(基于予以布局的Flickr landscape图片)
Results
Attention Is All You Need in the Latent Space
作者提出问题:只作用于低分辨率像素上的Transformer方法能否体现出对卷积方法的优势,用4个下采样块,做条件生成可能得多一个,然后在同样的表征上训练Transformer跟PixelSANIL去比较

表1在不同数据集跟模型尺寸上比较两个架构
看表1,三列结果分别是Transformer跟PixelCNN家族的P-SNAIL在相同训练步数、时间跟P-SNAIl采用固定训练时间的结果。说明Transformer结果总是更好,只是慢了些(两倍),如果训练一样的步数,效果还能更进一步。
A Unified Model for Image Synthesis Tasks
跟上一个实验设置一样,但进行条件图片合成

图4 Transformer统一多种图片生成任务 顶行:在ImageNet上无条件训练的补全(Completions from unconditional training on ImageNet.) 第二行:RIN上面的Depth-to-Image 第三行:ADE20K上的语义引导合成 第四行:DeepFashion上姿态引导的人类合成 底行:RIN上面的类别条件样 (Class-conditional samples)

图5 S-FLCKR上语义布局产生的结果,分辨率从上到下分别是1280x832, 1024x416, 1280x240

图6 将滑动窗口方法用到各种条件图片合成任务上 1:RIn的Depth-to-Image 2:IN上的随机超分辨率(Stochastic superresolution) 3/4:S-FLCKR的语义合成 5:IN上边缘引导的(Edge-guided)合成
看图4,5,6,VQGAN可以通用于多种条件图片合成任务
Building ContextRich Vocabularies
问题:上下文丰富的词汇表有多重要?固定Transformer,而通过改动VQ-GAN的下采样块来调整编码到第一阶段表征的上下文数量

图7 在HQ-Faces (CelebA-HQ and FFHQ)数据集上评估高效codebook的重要性,固定序列长度 $|s|=16\times 16=256$ 全局一致的结构只能用右边(f16)那上下文丰富的词汇表来建模 所有样本都使用temperature=1.0,top-100采样 最后一行是相对于f1的加速比,f1在像素上操作,用GTX Titan X产生一张图要7258秒
由于固定了Transformer,输入只能16x16,因此设下采样次数为f,那么每次对图片裁切 16f x 16f 大小的块,下采样后作为输入
图7表明了强大的VQGAN增加了Transformer的有效感受野,过小的感受野(f较小)模型无法捕捉到一致的结构。也就是说作者这里认为较多的下采样次数,得到的特征上下文越丰富,因为感受野大。此外还做了量化的实验,略
Benchmarking Image Synthesis Results

表2 语义图片合成的FID值 *使用本文的评估标准在每个数据集的验证集划分上得到的

表3 人脸图片合成的FID值
作者强调虽然表2表3中没取得SOTA,但本文提供了统一的模型,消除了完全对抗式方法跟似然方法的gap。其中表2采用top-100,表3在top-{100; 200; 300; 400; 500}中找最好的值
Class-Conditional Synthesis on ImageNet

表4 ImageNet上类别条件合成的FID值 用了跟VQVAE-2一样的rejection sampling,分类器基于ResNet-101 mixed k 表示从不同的top-k值采样,k in {100; 200; 250; 300; 350; 400; 500; 600; 800; 973}
表4跟一些SOTA的自回归方法比类别条件合成,这里没有标黑,总之比大部分方法优秀,而且在高拒绝率下能接近SOTA的质量。拒绝率指根据分类器丢弃掉的样本占总生成样本量的比例?
How good is the VQGAN? 通过codebook得到的Reconstruction FIDs提供了一种估计,指示生成模型在它上面训练后能取得的FID,为了量化VQGAN跟其他离散的、没使用感知损失跟对抗损失的VAEs(VQVAE-2,DALL-E)

表5 ImageNet上的FID,(FID/val)指原始验证集,(FID/train)指训练集,*表示用Gumbel-Softmax重参数化去训练的
表5显示VQGAN超越了VQVAE-2跟DALL-E的性能,并且压缩率还更高(codebook size小)。另一方面,更大的VQGAN(更大的codebook size或更长的编码(dim Z))进一步提高了性能。使用VQVAE-2那种层次codebook能取得最好的重建FID,尽管序列非常长,并不实际。对应的图片看图12(略)
Conclusion
通过将图片表示为感知丰富的图像成分的组合(composition of perceptually rich image constituents)降低了Transformer计算量。用CNN建模constituents,再用Transformer建模他们的组合能充分利用这两种架构。
Critique
- 遣词造句都很罕见,造成阅读困难,很不适应,比如 quantised token: a composition of perceptually rich image constituents from a codebook
- 方法其实比较直接,但有些的细节不是很清晰
- 实验内容很丰富,参数也都有给出,代码还开源
- 图片真的好多,40多页的图。虽然正文说12G的显卡就好,但附录提到在A100上训了45.8天,batchsize=16,模型很大,ImageNet上预训练的权重文件有16GB!
Unknown
- 超大分辨率图片是根据token一次性生成的,还是一部分一部分地生成再拼接?
FSQ
Pre
title: Finite Scalar Quantization: VQ-VAE Made Simple
accepted: arXiv2023
paper: https://arxiv.org/abs/2309.15505
code: https://paperswithcode.com/paper/finite-scalar-quantization-vq-vae-made-simple
关键词: Quantization
阅读理由: 新VQ? Google Research vq的发展史
Idea
作为VQ方法的简单替代
Motivation&Solution
- VQ系列方法需要额外的辅助损失(\(sg[\cdot]\))
- VQ的codebook编码利用率低下
Background
Vector quantization (VQ) 最早在1984年提出,近来在离散表征学习中复兴,在图片和音频领域都很流行。

图1 FSQ最后一层编码器输出3维向量(d=3),每一维度又有三个元素(L=3),并且直接取整为整型。VQ是输出向量到codebook里找最近的来替代。
作者说灵感来自 neural compression literature ,其中的离散编码都是通过标量量化(scalar quantization)获得的,但大多数的方法都是无边界的(unbounded)量化,整数的范围不由编码器限制,只取决于表征的熵。
Method(Model)
finite scalar quantization
而本文的方法是有限标量量化(finite scalar quantization, FSQ)。重点在于通过仔细挑选如何约束(bound)每个通道,能够获取任何期望尺寸的隐式codebook:假设有一个向量\(z\),它有\(d\)个通道(d维?),如果将其每个元素\(z_i\)映射到\(L\)个值(例如通过 \(z_i \mapsto \left \lfloor \frac{L}{2} \right \rfloor \tanh(z_i)\) 再四舍五入取整),这样就有了量化的\(\hat{z} \in \mathcal{C}\),图1那个FSQ就相当于有一个codebook \(\mathcal{C}\),且\(|\mathcal{C}|=L^d=27\)

图2 左:VQ和FSQ的对比 右:用f约束z,然后取整,L=5
如果将\(\tanh(\cdot)\)换成sigmoid,可以将量化后的值限定为自然数:
其中\(\sigma(x)=1/(1+e^{-x})\),\(Round\)表示取整,最终量化取值范围是L个整数\(\text{FSQ}(z)\in\{0,1,\cdots,L-1\}^d, \; z\in \mathbb{R}^d\)
如图2,FSQ大大简化了VQ,没有codebook,也不需要辅助损失,不过为了在取整后能获得梯度,跟VQ-VAE一样也用了STE(Straight-Through Estimator):
也就是\(Round\)操作的梯度为1,直接用取整前的梯度作为取整结果的梯度。
hyperparameters
FSQ有两个超参数:通道数\(d\)和每个通道的等级数\(\mathcal{L} = [L_1,\ldots, L_d]\),作者认为能满足 \(\prod_iL_i \approx |\mathcal{C}|\) 的超参都是可选的(也就是所有可能的量化变量数就相当于codebook大小,决定了能容纳的信息数)。不过并非所有选择结果都会好,作者给出了简单的heuristic:\(L_i \geq 5, \;\forall i\)
parameter count
FSQ参数远少于VQ,例如经典的一个VQ有\(|\mathcal{C}|=2^{12}=4096,\; d=512\),参数一共2M,而且FSQ的\(d\)也比VQ的更小,意味着编码器最后一层参数也更少,试着为VAE编码器和解码器开头加点dense层来补偿FSQ的参数减少,不过没啥用。
Experiment
tradeoff study

图3 128x128 ImageNet上VQ和FSQ的特性和权衡。FSQ的FID与codebook大小正相关。在codebook大小超过$2^10$时,FSQ有着更高的采样FID和codebook利用率,而VQ反而性能恶化。
同时FSQ有更高的压缩成本,而VQ的成本会在上升后又下降,变得再次容易建模。压缩成本与 sampling FID 负相关。
review of maskgit and uvim
MaskGIT中,作者训练了VQ-GAN,然后冻结它,再训练 masked Transformer BERT-style 来预测量化表征:给定一个表征\(\hat{z}\),一部分的token被随机遮掩,将结果的\(\hat{z}_M\)和 class token 一起喂入Transformer来预测每个 masked token 的分布。推理时只输入 class token 和 masked token,先基于预测置信度(prediction confidence)挑一些位置的token进行采样,结果用来替换相应的 masked token,然后再次运行模型,直到所有token都预测出来。
UViM是一个能处理计算机视觉多种预测任务的通用架构,第一阶段训练一个基于Transformer的VQ-VAE来建模目标任务的标签空间,第二阶段就训练一个编解码的Transformer来预测VQ-VAE编码的量化token。推理时基于输入自回归地用Transformer采样编码,再用VQ-VAE解码器解码。
characteristics and trade-offs for vq and fsq representations
在 128x128 ImageNet上训练MaskGIT,并尝试了不同配置的VQ和FSQ,指标:
- Reconstruction 量化自编码器的FID,表示二阶段Transformer是否完美建模了数据
- Codebook Usage 在编码验证集时至少用过一次的codewords所占的比例
- Sampling FID 解码Transformer以类条件采样的表征\(\hat{z}\)
- Compression Cost 表示将底层表征离散化的难度(建模复杂度)
早期实验发现FSQ有更高的召回率和较低的精度,受扩散模型文献启发,给MaskGIT加了 classifier
free guidance (CFG): 训练时用MASK token随机替换10%的类标签,让模型学习无条件分布。推理时将logtis插值, \(l_c, \;l_\emptyset\) 分别表示建模于类标签上的logits和无条件的logits,推理时 \(l' = l_c + \alpha(l_c - l_\emptyset)\) 来得到新logits,\(\alpha\) 就是CFG权重,直观上它会将预测的分布推向无条件的一边。(不是很理解,这有啥用)
results

表1 给定$\mathcal{C}$,要近似它所推荐的FSQ等级$\mathcal{L}$

图4 256x256 ImageNet上的MASKGIT结果 上方展示了每个MaskGIT模型最佳的classifier-free guidance (CFG)设置。作为参考展示了基于扩散的ADM模型;左下方是不同CFG权重下的精度召回率比较;右下方是不同CFG权重下的Sampling FID。ADM是水平线,因为1.5的CFG与其他人差太多,该横轴范围下无法展示。$\dagger$ 表示用 ADM TensorFlow Suite 来评估

图5 非特挑(Non-cherry-picked)样本,上面FSQ,下面VQ,跑了4个imagenet类别(330, 320, 510, 454)。俩模型正如图4的指标所示,很接近。

图6 UViM做深度估计的样本。VQ和FSQ仍然接近,不进行分割的VQ边缘粗糙。
Semantics 附录的实验并未发现VQ和FSQ的编码具有语义,两个量化器表现相似。
Precision-Recall trade-offs 精度是样本质量的测量,召回率则表示样本覆盖的真实分布的比例。图4的结果也说明调一调参都能得到差不多的结果。

表2 三个任务上UViM的结果。FSQ有竞争力,总体略差,$\dagger$ 表示UViM GitHub evaluation suite
表2里试着禁用了 codebook splitting(不知道是啥),总之FSQ不受影响,稳定。
appendix

图7 全景分割(panoptic segmentation)和图像着色(colorization)可视化
Conclusion
希望将来的工作能探索FSQ的更多应用。
Critique
附录直接放出了代码,很简洁,但 in jax。
RQ-VAE
Pre
title: Autoregressive Image Generation using Residual Quantization.
source: CVPR 2022
paper: https://ieeexplore.ieee.org/document/9879532/
code: https://github.com/kakaobrain/rq-vae-transformer
ref: https://zhuanlan.zhihu.com/p/666422373
关键词: Residual, Quantization
阅读理由: tokuni nai
Motivation
- 为了减少自回归模型的计算成本,较短的序列长度很重要
- 先前的VQ不能权衡缩短编码序列和图像失真

图1 256x256的条件生成样例
Idea
单个向量量化成深度为d的数组(原本vq都是二维的,它增加了一个维度),第二阶段使用空间Transformer+深度Transformer进行预测。同样用了GAN,更该叫RQ-GAN
Method(Model)

图2 RQ-VAE 和 RQ-Transformer 构成的两阶段生成框架。

公式3,4
公式3~4体现了RQ,即残差量化,其中\(k_d\)是向量\(z\)在深度d处的编码,\(r_d\)则是\(k_{d+1}\)对应的残差,\(r_0 = z\),\(d=1,\ldots,D\),那么\(z\)对应的量化向量则为:
从\(k_1\)到\(k_D\),体现出了一种从粗到细的思想,也就是加大D,\(\hat{z}\)就能够更加接近\(z\),换言之,通过增加D可以增强RQ-VAE的性能。
Experiment

图3 模型生成的图片

表1 无条件图片生成的FID比较

图4 1.4B参数RQ-Transformer的采样速度

表2 ImageNet上类别条件图片生成的FID和IS比较

表4 ImageNet验证集上重建图片FID与codebook大小的关系

图5 从粗到细近似的例子。第一张是原始图片,其余随着d增加,越来越清晰。
本文作者:心有所向,日复一日,必有精进
本文链接:https://www.cnblogs.com/Stareven233/p/17239337.html
版权声明:本作品采用知识共享署名-非商业性使用-禁止演绎 2.5 中国大陆许可协议进行许可。
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
2023-12-04 [pytorch] 余弦退火+warmup实现调研
2023-12-04 [论文速览] Randomized Quantization@ A Generic Augmentation for Data Agnostic Self-supervised Learning