简单有效！Vector Quantization的又一技巧：给编码表加一个线性变换

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

在《VQ的旋转技巧：梯度直通估计的一般推广》中，我们介绍了VQ（Vector Quantization）的Rotation Trick，它的思想是通过推广 VQ 的 STE（Straight-Through Estimator）来为 VQ 设计更好的梯度，从而缓解 VQ 的编码表坍缩、编码表利用率低等问题。

无独有偶，近期发布在 arXiv 上的论文《Addressing Representation Collapse in Vector Quantized Models with One Linear Layer》[1] 提出了改善 VQ 的另一个技巧：给编码表加一个线性变换。这个技巧单纯改变了编码表的参数化方式，不改变 VQ 背后的理论框架，但实测效果非常优异，称得上是简单有效的经典案例。

基础

由于在《VQ-VAE的简明介绍：量子化自编码器》[2]、《简单得令人尴尬的FSQ：“四舍五入”超越了VQ-VAE》等文章中我们已经多次介绍了 VQ 和 VQ-VAE 了，所以这里不再娓娓道来，直接给出普通 AE 和 VQ-VAE 的数学形式：

再次强调老生常谈的一点：VQ-VAE 不是 VAE，它只是一个加上了 VQ 的 AE，没有 VAE 的生成能力。而 VQ 则是将任意向量映射为编码表中与它最邻近的向量的操作，这个操作本身具有不可导的特性，所以通过 STE 来为 encoder 设计了梯度，并且新增了 这两项损失，来为编码表提供梯度，同时也起到规整 encoder 表征的作用。

改动

论文将自己所提方法称为 SimVQ，但没有解释 Sim 是什么含义，笔者猜测 Sim 是 Simple 的缩写，因为 SimVQ 的改动确实太 Simple 了：

没错，就只是在编码表多乘了一个矩阵 ，其他原封不动。

如果原本就是用式（2）训练 VQ 的，那么 SimVQ 可以直接简单上；如果原本是用 EMA 来更新编码表的（即 ，然后用另外的滑动平均过程来更新编码表，这是 VQ-VAE-2 及后续一些模型的做法，在数学上等价于用 SGD 来优化编码表损失，而其他损失则可以用 Adam 等非 SGD 优化器），那么则需要取消这个操作，重新引入 项来端到端优化。

可能马上有读者质疑：这不就是将编码表的参数化从 改为 吗？ 可以合并成一个矩阵，等价于一个新的 ，按道理不改变模型的理论能力？是的，SimVQ 对模型能力来说是不变的，但对 SGD、Adam 来说却是变的，它会改变优化器的学习过程，从而影响学习结果的好坏。

实验

进一步思考和分析之前，我们先看看 SimVQ 的实验效果。SimVQ 做了视觉和音频的实验，比较有代表性的是 Table 1：

▲ SimVQ 的实验效果

根据论文的描述，SimVQ 的代码就是在第一行 VQGAN 的代码上改的，改动就只有往 VQ 层插入了个线性变换，然后提升就非常显著了，不仅在相同编码表大小下达到了最优的重构质量，还能通过增加编码表大小进一步提高重构质量，这足以体现 SimVQ 的魅力——简单且有效。

笔者也在自己之前写的 VQ-VAE 代码上做了尝试，实测显示这个线性变换的加入，明显加速了 VQ-VAE 的收敛速度，并且最终的重构损失也有所降低。笔者还实验了 取对角阵的变体，这时候就相当于每个编码向量都 element-wise 地与一个参数向量（全一初始化）相乘，结果显示这样的变体也能起到相近的效果，介乎 VQ 与 SimVQ 之间。

分析

直观来想，VQ 对编码表的更新是比较“孤立”的，比如某个样本 被 VQ 为 ，那么这个样本的梯度就只会影响 ，不会影响编码表里的其他向量；但 SimVQ 不同，它不单会更新 ，还会更新 ，从几何意义上看， 就相当于编码表的基底，一旦更新 ，那么整个编码表就会更新了。

所以说，SimVQ 使得整个编码表的“联动”更为密切，从而更有机会找到更优的解，而不是陷入“各自为政”的局部最优。

那为什么 SimVQ 能提高编码表的利用率呢？这个其实也不难理解。再次根据 是编码表基底的解释，如果编码表利用率过低，那么 就会出现“各向异性”，即基底偏向于那些被利用起来的编码，可是一旦基底发生这种变化，那么它的线性组合应该也是偏向于被利用起来的编码，从而利用率不会太低。

说白了，可学习的基底会自动让自己的利用率变高，从而让整个编码表的利用率都提高起来。

我们也可以从数学公式角度来描述这个过程。假设优化器为 SGD，那么 VQ 中编码 的更新为

这样如果当前批次中 没有被选中，那么 为零，当前编码表就不更新了。但如果 被参数化为 ，那么

可以看到：

1.  是基于全体被选中的编码的梯度之和来更新的，所以它自然会更倾向于高利用率方向；

2. 由于 的存在，不管编码i有没有被选中，它的更新都几乎不会为零；

3.  相当于是高利用率方向的投影，它使得每个编码都往高利用率方向走。

然而，物极必反，如果全体编码都使劲往高利用率方向走，那么反而可能会导致编码表坍缩（codebook collapse），因此 SimVQ 默认采用了一个很保守的策略：只更新 ，所有的 在随机初始化后就不更新了，这样一来就几乎杜绝了编码表坍缩的可能性。

好消息是，在适当的编码维度下，实验显示 都更新和只更新 的表现都差不多，所以读者可以按照自己的偏好选择具体的形式。

延伸

抛开 VQ 的背景，像 SimVQ 这种引入额外的参数但又在数学上等价，即不改变模型的理论拟合能力，只改变优化过程的动力学的做法，我们称为“过参数化（Overparameterization）”。

过参数化在神经网络中并不鲜见，比如现在模型的主流架构是 Pre Norm 即 ，RMSNorm 最后所乘的 向量通常都是过参数化的，因为 的第一层通常就是线性变换，比如 Attention 是线性变换投影到 Q、K、V，FFN 是线性变换来升维，等等，这些模型在推理阶段 向量完全可以合并到 的线性变换中，但鲜有看到在训练阶段就把 去掉的做法。

这是因为不少人认为，深度学习模型之所以“好训”，过参数化有不可忽视的作用，因此贸然去掉已经充分验证的模型的过参数化风险很大。这里的“好训”，主要是指梯度下降这种理论上容易陷入局部最优的方法居然经常可以找到一个实际表现很好的解，这本身就是一件很不可思议的事情。

还有《On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization》[3] 等工作，表明过参数化隐式地加速了训练，作用类似于 SGD 中的动量。

最后，VQ 本质上可以理解为一种稀疏训练方案，所以 SimVQ 所带来的启发和改动，也许还能用于其他稀疏训练模型，比如 MoE（Mixture of Experts）。

当前的 MoE 训练方案中，Expert 之间的更新也是比较独立的，只有被 Router 选中的 Expert 才会更新参数，那么是不是有可能像 SimVQ 一样，所有的 Expert 后都接一个共享参数的线性变换，用来提高 Expert 的利用效率？当然 MoE 本身跟 VQ 也有很多不同之处，这还只是个猜测。

小结

本文介绍了 VQ（Vector Quantization）的另一个训练技巧——SimVQ——只在 VQ 的编码表多加一个线性变换，无需其他改动，就能达到加速收敛、提升编码利用率、降低重构损失等效果，相当简单有效。

（文：PaperWeekly）