简单得令人尴尬的FSQ：“四舍五入”超越了VQ-VAE

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 月之暗面

研究方向 | NLP、神经网络

正如 “XXX is all you need” 一样，有不少论文都以“简单得令人尴尬”命名（An Embarrassingly Simple XXX），但在笔者看来，这些论文大多数都是噱头多于实力。不过，笔者最近阅读到的一篇论文，真的让人不由得发出“简单得令人尴尬”的感叹～

论文的标题是《Finite Scalar Quantization: VQ-VAE Made Simple》[1]，顾名思义，这是一篇旨在用 FSQ（Finite Scalar Quantization）简化 VQ-VAE 的工作。随着生成模型、多模态 LLM 的逐渐流行，VQ-VAE 及其后续工作也作为“图像的 Tokenizer” 而“水涨船高”。然而，VQ-VAE 的训练本身也存在一些问题，而 FSQ 这篇论文则声称通过更简单的“四舍五入”就可以达到同样的目的，并且有着效果更好、收敛更快、训练更稳的优点。

FSQ 真有这么神奇？接下来我们一起学习一下。

首先，我们来了解一下 “VQ”。VQ 全称是 “Vector Quantize”，可以翻译为“向量量子化”或者“向量量化”，是指将无限、连续的编码向量映射为有限、离散的整数数字的一种技术。如果我们将 VQ 应用在自编码器的中间层，那么可以在压缩输入大小的同时，让编码结果成为一个离散的整数序列。

假设自编码器的重构损失能够让我们满意，那么这个整数序列就是原始图像的等价物，所有关于原始图像的操作都可以转化为整数序列上的操作。比如我们想训练图像生成模型，就只需要训练整数序列生成模型，而这跟本文生成等价，所以我们可以用它来训练一个 GPT，模型和流程都跟文本一模一样，训练完成后，我们就可以从 GPT 模型中采样整数序列，然后送到解码器中得到图像，从而完完成了图像生成模型的构建。

说白了，“VQ +自编码器”将任意输入都转化为跟文本一致的整数序列，统一了不同模态数据的输入形式，同时也统一了它们的处理和生成模型。

而这样的一个带有 VQ 功能的自编码器，就被称为 “VQ-VAE”。

早在四年前的文章《VQ-VAE的简明介绍：量子化自编码器》[2] 中我们就介绍过了 VQ-VAE，尽管被冠以 “VAE（Variational AutoEncoder）”之名，但它实际上跟 VAE 没啥关系，如上一节所说，它只是一个带有 VQ 功能的 AE（AutoEncoder）。

既然是 AE，那么有 encoder 和 decoder，一个普通的 AE 是这样的：

VQ-VAE 则稍微复杂一些：

让我们来逐步解释一下。首先，第一步是相同的，输入到 encoder 中，输出编码向量。然而，我们并不是直接将输入到 decoder 中，而是先维护一个编码向量表（Codebook），从中选取与最相近的一个送入到 decoder 中进行重构。由于编码表是有限的，所以我们也可以将实际的编码结果理解为一个整数（即与最相近的的），这就是 VQ-VAE 中 “VQ” 的含义。

当然，实际应用中，为了保证重构的清晰度，encoder 的输出可能是多个向量，每个向量经历同样的量化步骤变成一个整数，所以结果就是一张原本在连续实数空间的图片，被编码 VQ-VAE 编码为了一个整数的序列，这跟文本 Tokenizer 的作用是类似的，所以就有了“图像的 Tokenizer” 的说法。

梯度

然而，由于整个前向计算的流程中出现了，所以梯度无法回传到 encoder，这意味着我们无法优化 encoder。此时常见的手段是 Gumbel Softmax，但 Gumbel Softmax [3] 的效果通常也是次优的，所以作者巧妙地借助了 Straight-Through 为 VQ-VAE 设计了更好的梯度。可以说，这是 VQ-VAE 最精彩的内容，它告诉我们什么 “Attention is all you need” 都是虚的，“Gradient” 才是真正的 “all we need”！

具体来说，VQ-VAE 利用了深度学习框架基本上都自带的 stop_gradient（即公式中的）函数来自定义梯度，所有经过的输入，都会保持同样的输出，但梯度被强迫为零。所以对于式（2）中的，我们有

这样一来，送入 decoder 的还是量化过后的，但优化器求梯度时用的是，而是 encoder 出来的，所以 encoder 也能够被优化了。这个操作就叫做 “Straight-Through Estimator（STE）”，是为神经网络的不可导模块设计梯度的常用技巧之一。

由于基于梯度的优化器依然是当前的主流，所以直接设计梯度往往比设计 loss 更贴近本质，当然通常也更难、更让人由衷地赞叹。

Loss

不顾，事情还没完，此时有两个问题：1、现在 encoder 是有梯度了，但是编码表却没了梯度；2、虽然可以随意定义梯度，但不是胡乱定义一个梯度都可以成功优化模型的。从（3）可以看成，要使它在数学上严格成立，那么唯一的解是，这告诉我们如果 STE 是合理的，那么与至少是相近的。于是为了梯度的合理性，同时也为了优化编码表，我们还可以补充一项辅助 loss：

这样既可以迫使与接近，又可以让也拥有了梯度，一举两得！但细想之下，还是有点美中不足：理论上 encoder 和 decoder 的重构 loss 已经足够优化了，所以额外引入的一项应该主要用来优化，而不应该反过来明显影响。为此，我们再次利用技巧，不难证明式（4）的梯度等价于

第一项把的梯度停掉了，剩下的梯度，第二项则反过来，目前两项是 1:1 的权重求和，意味着两项相同程度地相互影响，而刚才我们说了，这辅助 loss 应该主要用来优化而不是，所以我们引入，将辅助 loss 改为

然后再加到重构 loss 中，就得到了 VQ-VAE 总的 loss 了

除此之外，的优化还有另外的方案：首先将式（6）的置零，这样一来就又没有梯度了；然后我们观察到，VQ-VAE 的 VQ 操作其实跟 K-Means 聚类是有点相似的，相当于是个聚类中心。根据我们对 K-Means 的了解，聚类中心等于该类的所有向量的平均，所以的一种优化方案就是的滑动平均

这等价于指定使用 SGD 优化这一项 loss（其他项可以用 Adam 等）。该方案被 VQ-VAE-2 [4] 所使用。

FSQ

可能有些读者疑惑，本文的主题不是 FSQ 吗？前面介绍 VQ-VAE 的篇幅是不是有点多了？事实上，由于 FSQ 完全对得起“简单得令人尴尬”这个评价，相比 VQ-VAE，介绍 FSQ 只需要“寥寥几行”，所以 VQ-VAE 不写长点，这篇博客就没几个字了哈～当然，将 VQ-VAE 写详细一点，也能让大家更深刻体会到 FSQ 的简单。

准确来说，FSQ 只是用来替代 VQ-VAE 中的 “VQ” 的，它的离散化思路非常非常简单，就是“四舍五入”。首先，假设我们有一个标量，我们定义：

这里的是一个超参数，就是 sigmoid 函数（原论文用了，笔者认为用 sigmoid 更科学），就是四舍五入为一个整数，所以不难看出，即 FSQ 运算将输出限制在了个整数之中，从而实现了离散化。当然，多数情况下一个标量还不够，对于，每一维可以执行 FSQ 运行，于是

即维向量被离散为个整数之一。但要注意，操作同样是没有梯度的（或者说梯度为零），不过经过 VQ-VAE 的铺垫，有些读者可能已经猜到接下来要做什么了：同样是利用 STE 技巧

即反向传播用之前的求梯度。由于前后本身是数值近似的，所以 FSQ 不需要额外 loss 来迫使近似的出现，也没有额外的编码表需要更新，FSQ 的简洁可见一斑！

▲ VQ 与 FSQ 对比（来自原论文）

实验

如果将 VQ 理解为直接将编码向量聚类为个不同的类别，那么 FSQ 就是将编码向量归纳出个属性，每个属性划分为了个等级，从而直接表达了个不同的整数。当然，从最一般的考虑，每个属性的等级数也可以是不相同的，从而不同的组合数为。

按照原论文的建议（之前的 LFQ [5] 则相当于），所以如果要对齐 VQ-VAE 的编码数量的话，对于 FSQ 来说应该有，即 FSQ 对编码向量的维度是有限制的（一般就只是个位数），并且通常是远小于 VQ-VAE 的编码维度（一般是三位数），这个直接后果是当编码总数比较小（从而也比较小）时，FSQ 的效果通常不如 VQ：

▲ 不同编码表大小下 VQ 与 FSQ 的效果差异

从图上可以看到，当编码表大小在 1000 左右时，FSQ 与 VQ 的效果接近；当编码表大小明显超过 1000 时，FSQ 占优；反之当编码表大小明显小于 1000 时，则 VQ 占优，这跟笔者自己的实验结果相近。笔者的参考代码为：

https://github.com/bojone/FSQ

其他实验就是比较常规的证明 FSQ 比 VQ 更优异的各种任务实验了，读者自行阅读原论文就好。

思考

从形式上来看，假设，那么 VQ 就好比是“一个类的分类器”，而 FSQ 则是“ 个级的打分器”，不管是从参数量、几何直观或者表达能力来看，其实 FSQ 都不如 VQ，但为什么 FSQ 有机会取得比 VQ 更好的结果呢？笔者认为有两方面的原因。

第一个原因，是 encoder 和 decoder 太强。虽然 FSQ 本身弱一些，但是 encoder 和 decoder 都足够强了，所以基于神经网络的万能拟合能力假设，FSQ 相对于 VQ 的劣势，完全可以在 encoder 和 decoder 中弥补过来。而在的设定下，两者的离散化程度都是一样的，也就是说 encoder 与 decoder 之间的“信息瓶颈”是一样的，因此 FSQ 本身的问题就显得微不足道了。

第二个原因，是 VQ 的“队友”（梯度）太弱。VQ 的经典问题是编码表坍缩：当编码表增大时，编码表并没有被充分利用起来，反而由于恶性竞争导致编码表聚集到一块了，经典表现就是一个 5000 的编码表，最终效果还不如 500 的编码表。

归根结底，这是梯度不够合理所致，尽管 VQ 已经巧妙地设计了梯度，但对于这种硬指派的运算，基于梯度的优化都存在“赢者通吃”问题，这是坍缩的根本原因，而 FSQ 的运算并不涉及到指派，它是直接取的近似值。当然，往大了讲，其实跟 VQ 类似的 K-Means 经常也有聚类中心坍缩的问题，可见难优化已经是一个老大难的问题了。因此与其说 FSQ 太强，倒不如说是 VQ 的“队友”太弱。

从以上两点分析可以看出，FSQ 要想超过 VQ，除了编码表要足够大之外，还有 encoder 与 decoder 要足够复杂，但这并非总能满足，比如有些场景下，我们希望模型的每一层输出都被量化，这时候平摊下来的 encoder 和 decoder 未必足够复杂，此时 FSQ 本身的不足就成为效果的瓶颈了。

此外，VQ 之后的向量维度没有变化，可以是任意多维，而 FSQ 之前的向量必须投影到维，这是很严重的降维，当我们需要用到投影之前的高维度近似向量时，就很难靠 FSQ 之后的低维向量简单恢复过来。

所以，如果单纯是作为“图像的 Tokenzier”，那么 FSQ 或许已经可以取代 VQ，但这并不意味着任意场景下 VQ 都可以被 FSQ 取代。