VQ-VAE-2 (2019)1¶
摘要¶
本文探索了将向量量化变分自编码器(VQ-VAE)模型应用于大规模图像生成的可行性. 为此, 作者扩展并增强了VQ-VAE中的自回归先验(autoregressive prior), 由此在保持高一致性与高逼真度(high-fidelity)的同时生成了更优质的合成样本. 模型仅依赖简单的前馈编码器和解码器网络, 因而在对编码或解码速度要求苛刻的场景中尤具吸引力. 此外, VQ-VAE只需在压缩的潜空间中对自回归模型进行采样, 与在像素空间采样相比——尤其是在处理大尺寸图像时——速度提升达一个数量级. 研究结果表明, 通过多尺度分层架构并在潜码上施加强有力的先验, 该模型在生成样本质量上可与当前在ImageNet等复杂数据集上表现最优的生成对抗网络(Generative Adversarial Networks, GANs)相抗衡, 同时避免了GANs常见的模式崩溃及多样性不足等问题.
简介¶
近年来, 深度生成模型取得了显著进展. 这部分得益于架构创新以及计算能力的提升, 使得在更大规模的数据量和模型规模上进行训练成为可能. 这些模型生成的样本若不仔细辨别, 很难与真实数据区分开来, 它们的应用涵盖超分辨率, 领域编辑, 艺术化操作, 以及文本到语音和音乐生成等场景.
本文将生成模型区分为两大类: 基于似然的模型(包括VAE, 流模型和自动回归模型); 以及隐式生成模型, 如GAN. 不同模型在样本质量, 多样性, 生成速度等方面各有取舍.
隐式生成模型和显式生成模型
显式生成模型和隐式生成模型的区分核心在于模型是否为数据分布提供了一个可计算的概率密度函数.
对于VAE, 模型结构中会显式假设一个先验分布\(p(z)\)和条件生成分布\(p_{\theta}(x|z)\), 并通过变分下界来近似最大化对数似然. 数据分布\(p_{\theta}(x)\)虽然直接计算困难, 但可ELBO是它的下界, 优化过程可以显式地评估和最大化两项\(\log p_\theta(x)\ge \mathbb{E}_{q_\phi(z\mid x)}[\log p_\theta(x\mid z)]-D_{\mathrm{KL}}(q_\phi(z\mid x)\Vert p(z))\), 从而间接地估算并提升模型对真实数据分布地拟合程度. 因此VAE是有一个"可写出, 可评估"地似然估计过程, 故称为显式生成模型.
对于GAN, 它的生成器只定义了一个从噪声到样本空间的转换, 但是没有给出样本\(x\)出现的概率密度函数\(p_G(x)\)的解析表达式. 判别器通过区分"真样本"和"假样本"来引导生成器学习, 但是始终无法显式计算或者评估生成分布的似然, 只能靠样本对抗训练来逼近真实分布. 因此, GAN只能隐式地学习到一个分布, 无法像VAE那样直接输出或者计算\(\log p(x)\), 故称为隐式生成模型.
GANs通过极小极大目标函数进行优化, 其中生成器神经网络将随机噪声映射到图像, 判别器通过把生成样本判定为真或假来定义生成器的损失函数. 随着数据量和模型规模的扩大, GANs现已能够生成高质量, 高分辨率图像. 然而, 众所周知, 这些模型的样本并未充分覆盖真实分布的多样性. 此外, 目前尚缺乏在测试集上评估过拟合程度的令人满意的泛化度量, 因而GANs的评估依然具有挑战性. 在模型比较与选择时, 研究者通常依赖图像样本或图像质量的代理指标, 例如InceptionScore(IS, 评估生成图像类别分布的置信度与多样性)和FréchetInceptionDistance(FID, 度量生成分布与真实分布在Inception特征空间中的Fréchet距离).
这里说的是IS, FID这些指标只能简介衡量生成图像质量的好坏.
相反, 基于似然的方法优化训练数据的负对数似然(NLL, 用于衡量模型对数据拟合程度的度量). 该目标便于模型比较并衡量对未见数据的泛化能力. 此外, 由于模型最大化对训练集所有样本赋予的概率, 理论上它们能够覆盖数据的所有模态, 并且不会遭遇GAN常见的模态崩溃和多样性缺失问题. 尽管具有这些优势, 在像素空间直接最大化似然具有挑战性. 首先, 像素空间的NLL并不总是样本质量的良好度量, 因此无法可靠地用于不同模型类别间的比较. 其次, 这些模型没有内在激励去关注全局结构. 例如, 通过引入诸如多尺度的归纳偏置或对图像的主要比特平面(bit planes, 按像素强度二进制位分层得到的图像平面)建模, 可以缓解部分问题.
基于似然的方法和GAN的区别
在最大似然训练里面, 我们直接最小化负对数似然\(L_{\text{NLL}}=-\sum_{i=1}^N \log p_{\theta}(x_i)\), 这等价于最小化前向KL散度\(D_{KL}(p_{data}||p_{\theta})=\sum p_{data}(x)\log \frac{p_{data}(x)}{p_{\theta}(x)}dx\). 由于\(\log \frac{p_{data}}{p_{\theta}}\)在\(p_{\theta}(x)\rightarrow 0\)的时候会发散到\(+\infty\), 模型若在任何真实样本\(x\)附近赋予极小概率就会遭遇无限大的惩罚, 因此最优策略在数据分布支持集的每个模态上到留出足够的概率质量.
GAN的生成器目标与此不同, 当判别器最优且训练稳定的时候, 生成器在逼近的是\(D_{KL}(p_{\theta}||p_{data})\)或者JS散度近似项. 反向KL的特点是当\(p_{\theta}(x)\rightarrow 0\)而\(p_{data}(x)\)仍为正的时候惩罚很小, 但是在\(p_{\theta}(x)\)把概率放到\(p_{data}(x)=0\)的区域时惩罚趋近于无限大. 换言之, 生成器更倾向于挑选几个易于拟合且能骗过判别器的高密度模态, 同时把剩余模态的概率压到0以避免在错误位置浪费质量, 这就是"模态奔溃".
基于似然方法的缺陷
像素空间的似然更像一种"低级信道编码损失": 它严格评估模型在\(256^{H\times W\times C}\)这一巨大离散空间内对每个像素具体取值的匹配程度, 而人眼更关心形状, 语义与整体协调性, 二者关注焦点天然错位, 由此带来两层偏差. 一是"量化噪声放大": 即便模型已精准复现全局轮廓, 但若在每个像素仍有1–2级灰度误差, \(-\log p_\theta(x)\)会骤增数百甚至上千比特, 然而肉眼几乎察觉不到差异. 二是"误导性锐化": 模型若把概率压缩到少数像素取值, 可显著降低NLL, 却易在纹理处产生过锐伪影, 视觉体验反而变差——NLL与感知质量并不单调对应. 所以说, 像素空间NLL给出的数字大小, 并不能在任意两种生成模型之间提供一个可信的优劣排序.
更根本的问题在于像素似然缺乏驱动模型关注长程依赖的内在激励. 梯度几乎完全由局部误差产生, 网络只需让局部纹理由真度极高, 便能取得很低的整体NLL, 而对物体结构, 透视一致性或跨区域配色并无直接奖励. 解决思路之一是注入多尺度归纳偏置: 先在低分辨率或高频截断的图像上训练与生成, 使网络必须先复原大尺度轮廓, 再逐级细化; 生成过程自底向上, 低频先定型, 高频随后补充, 从而把关注点强行拉向全局. 另一种做法是比特平面分层建模: 将8bit图像拆为八张二值平面, 先预测高位平面以锁定暗亮分区与物体形状, 再条件于此生成低位细节, 让"结构先行, 纹理随后"成为优化路径. 实践证明, 这类层级化或多尺度策略能在不牺牲像素似然的前提下, 显著提升主观感知质量, 也让模型在不同类别间的比较更具可解释性.
高频截断
高频截断指的是在图像领域, 人为将表示快速亮度变化的高频分量置零或者舍弃, 只保留低频(缓慢变化)的部分. 在生成模型中, 以"高频截断图像"作为第一阶段目标, 网络只能靠低频信息获得梯度 , 因此被迫优先学习全局布局, 随后再完整像素空间微调, 逐层补回高频细节, 便实现了"先结构, 后纹理的优化路径".
比特平面
在数字图像中, 每个像素通常使用8位二进制数来表示亮度(灰度图)或者每个颜色通道的强度, 例如x=11001010_2, 比特平面就是把这8位按照权重拆成8张二值图像, 第k个比特平面包含所有像素的第k位. 在这些平面中, 高位平面指的是靠近最高有效位(MSB)的一两层, 什么是MSB ? 例如, 每个像素写为二进制串x=b_7b_6b_5b_4b_3b_2b_1b_0, 那么, b_7就是最高有效位, b_0是最低有效位. 高位平面携带的信息量大, 一旦预测正确, 剩余条件熵迅速下降.
本文借鉴有损压缩的思想, 以减轻生成模型对可忽略信息进行建模的负担. 实际上, 像JPEG等传统压缩技术已证明, 在几乎不影响主观图像质量的前提下, 可以去除超过80%的数据. 按照相关工作提出的方法, 作者首先通过对自编码器的中间表示进行向量量化(vector-quantization,VQ,一种将连续向量映射为离散码本索引的过程), 将图像压缩到离散潜在空间. 该潜在表示的大小仅为原始图像的1/30, 但解码器仍能以极小失真重构图像. 对这一离散表示的先验分布, 作者采用集成自注意力机制的最先进PixelCNN, 即PixelSnail进行建模. 从该先验中采样后, 解码得到的图像在视觉质量和结构连贯性上与重构结果相当(参见图1). 此外, 在离散潜在空间上训练与采样均较像素空间快约30倍, 使得在更高分辨率图像上进行训练成为可能. 最后, 所使用的编码器与解码器沿袭原始VQ-VAE的简洁与高效, 因此该方法在需要对大型图像进行快速, 低开销编解码的场景下具有吸引力.
和VQ-GAN有点像, 都是在量化空间操作, 而不是在像素空间中操作. 并且都有一个自回归模型(在VQ-VAE-2中是PixelCNN, 在VQ-GAN中是Transformer). 不同点是: (1) 训练目标和损失函数. VQ-VAE-2用的是原版VQ-VAE的三项损失函数, 而VQ-GAN在这个基础上加入了感知损失和对抗损失. (2) 编码器. VQ-VAE使用的是层次化的codebook, 顶层负责全局o廓, 底层补充局部纹理. VQ-GAN放弃层次化, 改用单层但是感受野更大的编码器. (3) 自回归生成模型. VQ-VAE-2采用的是PixelCNN, VQ-GAN采用的是transformer.
方法¶
该方法采取两阶段策略: 首先训练一个分层VQ-VAE(图2a)将图像编码到离散潜在空间; 然后在由全部数据诱导的离散潜在空间上拟合强大的PixelCNN先验模型.
学习层次潜向量¶
与原始VQ-VAE不同, 本工作采用向量量化码的层次结构来建模大尺寸图像. 其核心动机是将纹理等局部信息与物体形状, 几何等全局信息分离建模, 因而可针对每一层的先验模型专门捕获该层存在的特定相关性. 图2a展示了多尺度分层编码器结构: 顶层潜向量负责建模全局信息, 而在顶层潜向量条件下的底层潜向量用于表示局部细节(见图3). 需要指出, 若底层潜在码不依赖顶层潜在码, 则顶层潜在码必须编码像素的全部细节. 为避免此问题, 每一层均直接依赖像素, 鼓励各潜在映射编码互补信息, 从而降低解码器的重构误差. 更多细节参见算法1.
Algorithm 1 VQ-VAE training (stage 1)
Require: Functions E_top, E_bottom, D, x (batch of training images)
1: h_top ← E_top(x)
2: e_top ← Quantize(h_top) ▷ quantize with top codebook eq 1
3: h_bottom ← E_bottom(x, e_top)
4: e_bottom ← Quantize(h_bottom) ▷ quantize with bottom codebook eq 1
5: x̂ ← D(e_top, e_bottom)
6: θ ← Update(L(x, x̂)) ▷ Loss according to eq 2
对最后一句话的解读
如果底层潜变量\(z_{\text{bottom}}\)不依赖于顶层潜变量\(z_{\text{top}}\), 那么顶层潜变量就需要从像素中包揽所有全局和局部的信息(包括纹理等细节). 为了避免让顶层潜变量承担过多的负担, 作者才因此在设计上让每一层既能从上层潜变量得到条件信息, 又能直接从像素获取额外特征, 从而形成互补.
对于256×256图像, 作者采用两级潜向量层次结构. 如图2a所示, 编码器首先将图像下采样4倍, 得到64×64的表示, 并将其量化为底层潜在映射. 接着, 另一组残差块(residual blocks)再将表示缩小2倍, 量化后得到顶层32×32潜在映射. 解码器同样是前馈网络, 以所有量化层次的潜向量为输入, 由若干残差块和多层带步长的反卷积(strided transposed convolutions)组成, 将表示逐步上采样回原始图像尺寸.
在潜码上学习先验分布¶
为了进一步压缩图像, 并在测试阶段能够直接从第一阶段学到的模型中采样, 我们在潜变量\(z\)上学习先验分布\(p_{\varphi}(z)\). 用神经网络从训练数据中拟合先验\(p_{\varphi}(z)\)已是惯例, 实践证明它能显著提升潜变量模型如PixelCNN的表现. 这一过程同时缩小了边际后验\(q_{\phi}(z)\)与先验\(p_{\varphi}(z)\)之间的差距, 因此在测试时从先验\(p_{\varphi}(z)\)采样得到的潜变量更符合解码器在训练期间见过的分布\(q_{\phi}(z)\), 进而生成更连贯的输出. 从信息论角度看, 把先验\(p_{\varphi}(z)\)拟合到学习到的后验\(q_{\phi}(z)\)可视为对潜空间的无损压缩: 通过更贴近边际后验分布\(q_{\phi}(z)\)的概率模型PixelCNN重新编码潜变量, 平均码长趋近于香农熵\(H(Z)\). 若真实熵与先验的负对数似然\(-\!\log p_{\varphi}(z)\)之间的差距越小, 则冗余越少, 由这些潜变量解码得到的图像样本就会越逼真.
VQ-VAE只是一种离散化的工具, PixelCNN负责在这些符号上建模
- 先验: 或者说是真实先验, \(p(z)\)
- 旧先验: 强行设置的极简先验, \(p_0(z)\)
- 边际后验: 编码器学到的先验, \(q_{\phi}(z)\)
- 新先验/辅助先验: PixelCNN学习到的先验, \(p_{\varphi}(z)\)
VQ-VAE本身只做两件事: **(1)把像素映射成离散码本索引; (2)把整张索引图一次性解码回像素. 它并不会告诉你应当如何挑选码字来生成可信图像. 要"从零"合成新图, 必须额外学习潜空间的概率模型——即新先验**\(p_{\varphi}(z)\). PixelCNN, PixelSnail以及VQ-GAN中的Transformer先验正负责这一环节: 它们通过自回归链式分解\(\prod_{i,j}p_{\varphi}(z_{i,j}\mid z_{<i,<j})\), 一边拟合边际后验\(q_{\phi}(z)\), 一边在生成时逐位置输出条件分布并随机或贪婪采样. 缺少这一步, 只能随意拼接码本索引, 采样极易落入编码器未见过的区域, 最终图像往往崩坏.
当VQ-VAE训练完成后, 编码器已把每张训练图像\(x\)压缩成潜向量\(z\). 这些\(z\)在潜空间中的统计规律由边际后验\(q_{\phi}(z)\)描述; 它通常既非独立, 也不服从简单高斯, 而蕴含丰富的空间与语义依赖. 如果推理阶段仍从旧先验\(p_{0}(z)\)(如各位置独立均匀)采样, 就极易落到\(q_{\phi}\)的低密度区, 解码器便输出噪声或伪影. 自回归先验的核心工作, 正是把\(q_{\phi}(z)\)逼近成新的\(p_{\varphi}(z)\), 并在生成时通过条件概率一步步指导采样, 确保每个\(z_{i,j}\)都落在编码器熟悉的高密度区域.
因此, VQ-VAE的一个重要动机就在于为自回归模型提供带有语义颗粒度的局部token: 先按扫描顺序逐个生成VQ-token, 拼成完整索引网格\(z\), 再将其一次性送入解码器, 还原出最终图像.
信息论角度
在信息论里, 若潜变量真实分布是\(q_\phi(z)\), 用任意概率模型\(p_\varphi(z)\)去编码\(z\), 平均码长等于交叉熵
其中\(H(Z)=-\mathbb E_{q_\phi(z)}[\log q_\phi(z)]\)是香农熵, 而\(\mathrm{KL}(q_\phi\parallel p_\varphi)\)衡量两分布差异. 若用PixelCNN把\(p_\varphi\)训练得几乎与\(q_\phi\)重合, 则\(\mathrm{KL}\)项极小, 平均码长便逼近理论下界\(H(Z)\). 这意味着编码里不再携带多余比特——也就是冗余最小.
与此同时, 当\(\mathrm{KL}(q_\phi\parallel p_\varphi)\)变小, \(-\!\log p_\varphi(z)\)在样本上的期望就接近真实熵\(H(Z)\). 这不仅提高压缩效率, 还直接影响生成质量:测试阶段我们从\(p_\varphi\)采样潜变量并送入解码器. 若\(p_\varphi\)与训练时的\(q_\phi\)几乎一致, 解码器面对的是"熟悉"的潜在表示, 因而能够重现训练中学到的对应像素模式; 反之, 若两者差距大, 采样点落入解码器未充分见过的区域, 生成结果就会失真或不连贯.
在VQ-VAE框架中, 该辅助先验\(p_{\varphi}(z)\)在后处理的第二阶段由PixelCNN这类强大的自回归(autoregressive)神经网络进行建模. 顶层潜在图的先验负责捕捉全局结构信息, 因此作者为其加入多头自注意力(self-attention)层, 以利用更大的感受野来建模图像中相距较远的空间相关性. 相比之下, 面向局部信息的底层条件先验需在更高分辨率上运行; 若直接沿用顶层的自注意力结构, 显存开销将难以承受. 因此, 对于这一局部先验, 作者发现利用来自顶层先验的大型条件堆叠即可获得良好性能(见图2b). 层次化的分解还使得训练更大模型成为可能: 作者分别训练每一层先验, 从而充分发挥硬件加速器的计算与内存资源.
本文的顶层先验网络针对32×32维度的潜在变量进行建模. PixelCNN中的残差门控卷积层每隔5层便插入一次因果多头自注意力层, 以扩大感受野并捕捉长程依赖关系. 为正则化模型, 作者在每个残差块之后以及每个注意力矩阵的logits上均施加dropout. 实验证明, 在PixelCNN堆叠之上再加入由1×1卷积构成的深层残差网络, 能在几乎不增加训练时间和显存占用的前提下进一步提高对数似然. 底层条件先验作用于64×64分辨率的潜在变量, 在内存与计算成本方面显著更高. 如前所述, 这一层主要编码局部特征, 在已有顶层先验的条件下并不需要巨大的感受野. 因此, 作者为底层采用了不包含自注意力层的较轻量网络, 同时发现引入深层残差条件堆叠对提升性能大有裨益.
多样性-质量折中¶
与GAN相比, 在最大似然目标(MLE)下训练的概率模型被迫覆盖整个训练数据分布. 原因在于, MLE目标可写成数据分布与模型分布之间的正向KL散度(forward KL-divergence,Kullback–Leibler散度的一种形式); 若模型对某个训练样本赋予零概率, 该散度将发散至无穷. 虽然后者能够覆盖数据分布的全部模态这一性质颇具吸引力, 但相较于对抗式建模, 基于似然的模型需要拟合数据中的所有模态, 因此任务更加艰巨. 进一步地, 自回归模型的祖先采样(ancestral sampling, 即按序列顺序逐步采样)在实践中会沿长序列累积误差, 最终导致生成样本质量下降. 近期的GAN框架提出了自动化的样本选择流程, 以在多样性与质量之间进行权衡. 本文亦基于如下直觉提出一种自动化的多样性–质量折中方法: 样本越接近真实数据流形(datamanifold), 越可能被预训练分类器判定为正确类别. 具体而言, 作者利用一个在ImageNet上训练的分类网络, 依据其为正确类别输出的概率为模型生成的样本打分.
-
Razavi, A., Oord, A. van den, & Vinyals, O. (2019). Generating diverse high-fidelity images with VQ-VAE-2 (No. arXiv:1906.00446). arXiv. https://doi.org/10.48550/arXiv.1906.00446 ↩