跳转至

FSQ (2023)1

概要

作者提出在VQ-VAE的潜在表示中使用一种简单的有限标量量化(FSQ)方案替代向量量化(VQ). 具体而言, 作者首先将VAE表示降至少数(通常不足10)维, 并对每一维离散化为一组固定取值, 从而得到由这些集合直积隐式定义的码本. 通过恰当选择维度数量及每维可取值个数, 可使码本规模与VQ保持一致. 基于此类离散表示, 可训练此前针对VQ-VAE表示而设计的同类模型, 例如用于图像生成的自回归与掩码Transformer模型, 多模态生成模型以及密集预测的计算机视觉任务. 具体地, 作者在图像生成中将FSQ与MaskGIT结合, 在深度估计, 图像上色和全景分割中将FSQ与UViM结合. 尽管FSQ设计更为简洁, 但在所有任务中仍取得了具有竞争力的性能. 作者强调, FSQ不存在码本坍缩问题, 且无需VQ中用于学习高表达力离散表示的复杂机制(如承诺损失, 码本重新初始化, 码本拆分, 熵惩罚等). 代码已发布于GitHub.

简介

向量量化(VQ)最早由Gray提出, 近年来随着神经网络离散表示学习的兴起而再度焕发生机. 受VQ-VAE成功的启发, Esser等人(VQ-GAN)与Villegas等人分别表明: 在以GAN损失训练的VQ-VAE表示之上训练自回归Transformer, 能够构建功能强大的图像与视频生成模型. 同时, VQ已成为图像与音频表征学习的流行组件, 并被视为构建下一代多模态大型语言模型的有前景基础模块.

在训练VQ-VAE时, 目标是学习一个码本\(C\), 其元素能够对输入数据(通常为图像)形成压缩且具备语义的信息表示. 在前向过程中, 图像\(x\)被编码为表示\(z\)(通常是一系列特征向量), 然后\(z\)中的每个向量都会量化为——即替换为——码本\(C\)中与之最接近的向量. 由于量化操作不可微分, Van Den Oord等人在潜在表示中使用VQ训练VAE时采用了直通估计器(STE), 将解码器输入的梯度直接复制到编码器输出, 从而为编码器提供梯度. 然而, 这样依旧无法为码本向量产生梯度; 因此, 作者进一步引入了两个辅助损失: 一方面将码字向量拉向对应的(未量化)表示向量, 另一方面促使表示向量靠近其匹配的码字.

上述公式难以优化, 并导致码本利用不足这一已知问题: 当增大\(C\)的规模时, 许多码字会闲置. 随后出现的研究尝试通过多种技巧改进这一点, 例如重新初始化整个码本或部分码字, 采用随机化量化策略等.

本文旨在简化原始VQ-VAE公式, 其目标包括: i)移除辅助损失; ii)通过设计实现高码本利用率; iii)在功能层面保持与VQ一致, 从而能够作为VQ的直接替代方案.

为此, 作者借鉴神经压缩领域的研究, 该领域通常通过标量量化获得离散码. 早期工作表明, 只需将表示\(z\)中的每个标量元素独立四舍五入到最近整数即可完成量化. 当前大多数压缩方法采用无界标量量化, 即整数范围并不受编码器限制而是通过约束表示的熵来控制码率. 亦有研究选择限制量化器的取值范围以实现压缩.

标量量化

把连续向量\(z=(z_1,\dots,z_d)\)中的每个标量分量\(\,z_i\)各自独立地四舍五入到最近的整数: \(\tilde z_i=\text{round}(z_i)\). 这样做既简单又高效, 因此在神经压缩(neural compression)文献里很常见.

接着作者区分了两种主流做法:

  • 无界(unbounded)标量量化

    这里的"无界"指的是: 编码器不强行限定\(\tilde z_i\)必须落在某个固定区间, 理论上它可以取任意整数. 为了仍然得到可控的码率(bit rate), 研究者会在训练时给\(\tilde z\)加一个"熵约束"或率失真(rate–distortion)正则项, 使得整体熵\(H(\tilde z)\)保持在目标范围内——相当于让网络自己学会把大多数概率质量集中在少数码字上, 以便熵编码器能用较少比特描述它们.

  • 有界(bounded)标量量化

    另一条思路是在量化时直接把取值范围截断(例如规定\(\tilde z_i\in[-Q,Q]\)或只允许\(K\)个等级). 这样做的好处是熵编码更简单, 码本有限; 缺点是需要在"量化误差"和"可表示范围"之间手动权衡.

作者称这种方法为有限标量量化(FSQ). 其核心洞见在于: 只要为每个通道仔细设定取值范围, 就能构造出几乎任意大小的隐式码本. 设向量\(z\)含\(d\)个通道. 若将每个分量\(z_i\)映射到\(L\)个离散值(例如先做\(zi↦⌊L/2⌋\tanh(zi)\), 再四舍五入取整), 便得到量化后的\(\hat{z}\). 由此\(\hat{z}\)必属于\(L^d\)个唯一向量之一. 例如图1给出了\(d=3\), \(L=3\)的FSQ示例, 对应码本\(C = \{(-1,-1,-1), (-1,-1,0), (-1,-1,1), …, (1,1,1)\}\), 其大小为\(|C| = L^d = 27\).

图1: FSQ(左): 最后一层编码器将特征投影到$d$维空间(示例中$d=3$). 对编码器输出$z$的每一维设定$L$个离散取值(示例中$L=3$), 再将其四舍五入为整数, 得到量化向量$\hat{z}$, 即位于此超立方体内的最近点. VQ(右): 最后一层编码器同样投影到$d$维空间(示例中$d=7$, 在VQ中$d$通常更大). 生成的向量$z$通过最近邻查找, 被替换为码本中最接近的向量$\hat{z}$.

为了在舍入操作中获得梯度, 作者借助与VQ-VAE相同的直通估计器(STE,straight-through estimator). 因此, 当FSQ被置于以重构损失训练的自编码器中时, 梯度能够回传到编码器, 迫使模型将信息分散到多个量化区间(quantization bins), 以降低重构误差. 最终即可在无任何辅助损失的情况下,得到一个能够充分利用全部码字的量化器.

据作者所知, 除压缩任务外, FSQ尚未被用于视觉领域, 而在该场景中VQ仍占据主导地位. 本文通过将FSQ与强大的transformer/语言模型结合, 旨在改变这一局面. 总结而言, 本文的贡献如下:

  1. 作者证明, 在不同架构, 数据集和任务上, FSQ(有界标量量化)能够作为VQ(向量量化)的即插即用替代方案. 他们将FSQ应用于MaskGIT(一种基于mask预测的图像生成模型), 以及UViM框架中的深度估计, 图像着色和全景分割任务. 结果显示, 相应指标仅下降0.5–3%, 视觉效果几乎保持一致. 需要强调的是, 这两类模型在设计上差异显著, 包括卷积式vs.transformer自编码器, 掩码式vs.全自回归transformer, 仅解码器vs.编码器-解码器transformer等.
  2. 作者分析了VQ与FSQ之间的权衡, 从压缩角度刻画了两种方法在码本规模上的扩展特性与表示复杂度. 结果表明, FSQ能有效利用更大的码本以获得更优的重构指标与更高的样本质量. 在无需任何辅助损失的情况下, FSQ的码本利用率极高(大多数模型约为100%).
  3. 作者指出, 在码本规模较大时, VQ完整公式的高度灵活性几乎无法带来额外收益, 甚至表现劣于简单的FSQ方法. 这主要归因于VQ的优化难度较高, 而FSQ可被视为对标准VQ框架做出的两点改动: (a) 限制编码器输出范围; (b) 固定码本大小C. 同时, FSQ的(隐式)码本C维度显著小于VQ, 通常d < 10, 而VQ常见d ≥ 512.

图2: 左: 简化版VQ——比较VQ与FSQ在实现与优化上的异同. 右: 通过函数f对z进行范围限制,随后对输出进行舍入; 图示为单通道,且量化等级L=5.

方法

作者首先给出高层次的直觉. VQ在VQ-VAE的高维潜在空间中学习一种Voronoi分区(即最近邻分割), 从而对输入空间(如图像)进行复杂的非线性划分. 相比之下, FSQ仅在维度远低得多的空间中采用简单, 固定的网格分区. 直观而言, 这种做法之所以可行, 是因为在典型应用中VAE本身具有较高的模型容量(见Sec. 2), 因此VQ所需的非线性能够被编码器和解码器"吸收", 从而FSQ仍能实现与VQ相当复杂度的输入空间划分.

对上面这段话的解读

VQ把高维潜向量\(z\)映射到一个可以学习的码本. 每个码本向量定义了一个voronoi单元, 在潜空中形成了复杂, 非线性的最近邻分区. 由于潜在空间和输入图像通过非线性编码器一一对应, 这种复杂分区也对应输入图像的复杂划分. FSQ的做法则完全不同, 它先把潜空间的维度降低到很小, 再再每个维度上使用固定, 等距的网格(整数格点)做独立量化. 看似粗糙的网格为何仍行之有效呢? 关键在于典型的VAE拥有足够大的模型容量.

模型容量

在机器学习中, 模型容量指的是一个模型能够表示或者逼近的函数族的复杂程度的大小. 对于VAE而言, 编码器\(q_{\phi}(z|x)\)和解码器\(p_{\theta}(x|z)\)通常构建为深而宽的卷积或者Transformer, 其参数量轻易就能够达到上千亿的级别. 这样的模型容量允许他们学习到高度非线性的映射. 因此, 即使FSQ只提供了低维, 规则的网络, 编码器也能"事先"把数据折叠到合的位置, 然后解码器复杂展开. 换言之, 高容量使得量化器本身不必承担全部非线性表达, 许多细节被网络权重隐式吸收.

因此, 那些原本由VQ复杂Voronoi单元实的非线性, 现在被吸收到编码器/解码器的参数里面. 结果就是, 即便FSQ自身的分区在数学上只是简单的正方体网格, 整个VAE-FSQ系统对输入空间的有效划分复杂度依旧可以和VQ-VAE媲美, 但是实现和优化都更简单, 稳定, 码本利用率也更高.

有限标量量化

给定\(\mathbf{z}\in\mathbb{R}^d\), 目标是将其量化到有限的codewords集合. 为此, 作者首先施加边界函数\(f\), 随后取整得到 \(\hat{\mathbf{z}}=\operatorname{round}(f(\mathbf{z}))\) 其中\(f:\mathbf{z}\mapsto\lfloor L/2\rfloor\tanh(\mathbf{z})\)使\(\hat{\mathbf{z}}\)的每个通道仅取\(L\)个离散值. 因而\(\hat{\mathbf{z}}\)落在隐式codebook \(C\)中, \(C\)由各通道离散集合的笛卡尔积构成, 且 \(|C|=L^d\) \(C\)中的向量可直接枚举, 由此建立\(\hat{\mathbf{z}}\)与\(\{1,\ldots,L^d\}\)的一一映射. 因此, 在任何原本使用VQ的神经网络场景(如训练transformer)中, 只需相应调整VQ前后层的维度, 便可用FSQ替换VQ. 若第\(i\)个通道映射到\(L_i\)个值, 则 \(|C|=\prod_{i=1}^d L_i\)

图1(左)与图2展示了FSQ示意. 由于量化通过取整实现, 当\(L\)为偶数时需采用非对称的\(f\); 论文在附录A.1给出了通用\(f\)的代码. 为了在反向传播中穿过取整操作, 作者使用STE, 将梯度恒设为1. 在主流ML框架中, round_ste可通过"stop gradient"(sg)操作实现: \(\operatorname{round}_{\text{ste}}:x\mapsto x+\operatorname{sg}(\operatorname{round}(x)-x)\)

每个通道的量化级数不一定是要相同的.

超参数

FSQ具有如下超参数: 通道数\(d\)以及每个通道的量化级数\(L=[L_1,\dots,L_d]\). 在大多数实验中, 为了实现公平比较, 作者根据欲用FSQ替换的VQ码本来选择目标码本大小\(|C|\). 然而, 多种\(d\)与\(L_i\)的组合都可以逼近给定的\(|C|\)(即满足\(\prod_i L_i\approx|C|\)的任何\(L\)都是候选). 作者在研究中考察了这些配置, 发现并非所有选择都能得到最优结果. 不过, 他们提出了一个在全部任务中均表现良好的简单启发式: 对所有\(i\)设\(L_i\ge5\). 表1列出了常见目标\(|C|\)对应的\(L\)取值.

表1: 推荐的FSQ离散级数L组合, 以近似匹配给定的码本大小|C|

参数数量统计

作者指出, FSQ的参数量少于VQ, 因为在VQ中需学习一个大小为\(|C|·d的\)码本. 例如, 当\(|C|=2^{12}=4096\)且\(d=512\)时, 需要额外引入约\(2M\)参数, 而FSQ则无需这些参数. 此外, 对于FSQ, \(d\)通常远小于VQ(如表1所示, 在相同\(|C|\)下FSQ仅需\(d=5\)), 因此训练FSQ时编码器的最终dense层(全连接层)也包含更少参数. 为补偿这一差异, 作者尝试在VAE编码器末端及解码器开端增加更多dense层, 但并未获得额外收益. 因此, 在本文的所有模型中, 在同等码本大小下FSQ始终拥有更少的参数量.

  1. Mentzer, F., Minnen, D., Agustsson, E., & Tschannen, M. (2023). Finite scalar quantization: VQ-VAE made simple (No. arXiv:2309.15505). arXiv. https://doi.org/10.48550/arXiv.2309.15505 ↩