专栏名称: PaperWeekly

PaperWeekly是一个分享知识和交流学问的学术组织，关注的领域是自然语言处理的各个方向。我们热爱知识，分享知识，希望通过我们大家的努力为自然语言处理的发展做出一点点贡献。我们每周会分享一期特定话题的论文笔记和本周值得读的相关论文。

矢量量化（VQ）的旋转技巧：梯度直通估计的一般推广

PaperWeekly · 公众号 · 科研 · 2024-10-30 20:28

正文

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

随着多模态 LLM 的方兴未艾，VQ（Vector Quantization）的地位也“水涨船高”，它可以作为视觉乃至任意模态的 Tokenizer，将多模态数据统一到自回归生成框架中。

遗憾的是，自 VQ-VAE 首次提出 VQ 以来，其理论并没有显著进步，像编码表的坍缩或利用率低等问题至今仍亟待解决，取而代之的是 FSQ 等替代方案被提出，成为了 VQ 有力的“竞争对手”。

然而，FSQ 并不能在任何场景下都替代 VQ，所以 VQ 本身的改进依然是有价值的。近日笔者读到了《Restructuring Vector Quantization with the Rotation Trick》[1]，它提出了一种旋转技巧，声称能改善 VQ 的一系列问题，本文就让我们一起来品鉴一下。

回顾

早在五年前的博文《VQ-VAE的简明介绍：量子化自编码器》[2] 中我们就介绍过了 VQ-VAE，后来在《简单得令人尴尬的FSQ：“四舍五入”超越了VQ-VAE》介绍 FSQ 的时候，也再次仔细地温习了 VQ-VAE，还不了解的读者可以先阅读这两篇文章。

VQ-VAE 虽然被冠以 VAE 之名，但它实际上只是一个 AE，并没有 VAE 的生成能力。

它跟普通 AE 的区别是，它的编码结果是一个离散序列而非连续型向量，即它可以将连续型或离散型的数据编码为一个离散序列，并且允许解码器通过这个离散离散来重构原始输入，这就如同文本的 Tokenizer——将输入转换为另一个离散序列，然后允许通过这个离散序列来恢复原始文本——所以它被视作任意模态的 Tokenizer。

用公式来说，普通的 AE 是：

而 VQ-VAE 则是：

其中 “VQ” 主要就是指从变换到的过程，它将z映射成之一，这些就称为编码表（Codebook），也是可学习的向量。而训练 VQ-VAE 的“神之一手”，就是这一步，它称为梯度的“直通估计器（Straight-Through Estimator，STE）”。

STE

直通估计的出现，是因为从到的变换包含了不可导的运算，所以没法直接将梯度传播到编码器中，换句话说编码器是没法训练的。为此，VQ-VAE 想了一个技巧，它利用 stop_gradient 算子和与的最邻近特性，在反向传播时用替换，也就是。

此时，前向计算等价于不存在，所以，即送入 Deocder 的是，而求梯度时的梯度等于 0，所以，所以梯度可以绕过不可导算子直达编码器，这就是“直通估计器”。

不过这样一来，编码器是能优化了，但编码表却不能优化了，所以 VQ-VAE 往损失函数中加入了来优化编码表，其意图类似 K-Means，希望等于所有与它最邻近的 z 的中心。最后的，则希望编码器也主动配合来促进这种聚类特性。

从梯度的链式法则角度看，我们有：

注意这里都是向量，所以也都是向量，而则是一个矩阵。由于到的不可导性，所以问题卡在没有良好定义，而 STE 则相当于假设了（单位矩阵），所以。这个设置自然有一定的合理性，但有没有什么改进空间呢？

直观上来看，STE 导致的结果是，对于属于同一个的所有，它们的梯度都是相同的，而跟的距离远近无关，这似乎就是一个可改进的地方：我们是否可以定义更一般的，使得它跟的差异大小有关呢？为了达到这个目的，我们先将 STE 推广成：

其中是一个矩阵。再次根据前向传播不存在、反向传播梯度为零的原则，可以得出，这就相当于定义了。

旋转

那怎么选择呢？文章开头所提的论文提出了一个参考方案，基于从到的旋转变换来构建，即论文标题中的 “Rotation Trick”。

具体来说，原论文考虑了的简单情形，此时自动为零，从而简化成。为了找到矩阵，我们先将都归一化为单位向量，那么就可以构建一个从到的旋转变换。

具体的构造方式我们在《从一个单位向量变换到另一个单位向量的正交矩阵》已经探讨过，答案是：

其中是的夹角。利用这个结果，我们可以写出

为了提高计算的效率，我们通常选择利用矩阵乘法的结合律先计算和，但要注意我们实际上需要的是，所以要注意先停掉的梯度再去计算。

从几何意义上来看，，使得相对于的几何性质，跟相对于的几何性质是完全一致的，比如与的夹角等于与的夹角，它们的模长之比也相等，这些性质自然是有理论上的优雅性，但它是否真的能改善 VQ-VAE 的性能呢？接下来让我们转到实验部分。

实验

论文在相同的配置下对比了旧版 STE 和旋转技巧，发现旋转技巧的表现可谓“惊艳”：

▲ VQ-VAE + 旋转技巧的表现

▲ VQ-GAN + 旋转技巧的表现

简单来说，就是该高的地方（编码表利用率、IS）高、该低的地方（重构误差、Loss、FID）低，完全符合理想模型的特性了。论文的代码也已经开源，有兴趣的读者可以自行试跑一下。

Github：

https://github.com/cfifty/rotation_trick

思考

那这是不是意味着所有的 VQ-VAE/VQ-GAN，都可以无脑上旋转技巧了呢？笔者在以前自己写的能跑通的 VQ-VAE 代码加上了旋转技巧，发现效果反而变得更差了，具体表现是重构损失变得更高，编码表损失则更低了。

经过简单分析，笔者发现问题出在这个选择上，原本的 STE 则是，这里旋转矩阵跟单位矩阵的尺度是相当的，所以旋转技巧尺度上多出了。

如果初始化时（笔者写的 VQ-VAE 正好是这样），那么旋转技巧加持下重构损失的梯度就会比 STE 加持下重构损失的梯度小很多，于是对于编码器来说这一项的梯度占了主导。

换句话说，初始阶段相当于只在优化

，这会导致，即编码表坍缩，这就能解释编码表损失降低、重构损失增加的现象了。

所以，从 STE 切换到旋转技巧大概率至少需要重新调一下。笔者简单看了一下论文的开源代码，里边应该是利用初始 Encoder 的 K-Means 来初始化编码表的，这样一来与的数量级不至于差太远，从而可以比较顺畅地切换。

不过，即便精调了，笔者也没在自己的 VQ-VAE 代码上调出更优的效果，所以笔者对旋转技巧的有效性保持观望态度。抛开实践不说，理论方面笔者也理解不了旋转技巧的有效性。

原文的分析是，当与很相近时，就很接近，此时是合理的，而当与距离较远，比如位于类别的边界附近时，与的差距较大，即明显偏离，于是处于“乱飞”的状态，有助于冲破“牢笼”而迈向新的类别，从而提高编码表的利用率。但很显然，这个解释让人觉得很“没底”。

此外，旋转技巧还有一个问题，就是它确立了一个具有超然地位的中心位置——原点。不难理解，VQ 操作本身类似于 K-Means 聚类，而 K-Means 是无中心的，它具有平移不变性，而旋转则需要一个中心（原点），所以旋转技巧实际上跟 VQ 本意有点相悖。

当然，VQ 也可以改为按余弦值来找最邻近，这更契合旋转技巧，但也无法解释为什么旋转技巧对基于欧氏距离的 VQ 也有帮助。总的来说，旋转技巧起作用的根本原因，依旧是值得深思的问题。

最后，可能有读者疑问：既然 VQ 有这么多问题，为什么还要研究 VQ 呢？为什么不用更简单的 FSQ 呢？

笔者认为，诸如 FSQ 等替代品，并不是在任何场景都能取代 VQ，比如《VQ一下Key，Transformer的复杂度就变成线性了》介绍的Transformer-VQ，就很难用 FSQ 来替代 VQ，因为它是每一层都要 VQ 一下，这样分配下来相当于说 VQ 的模型很小，而 FSQ 测试下来只有当模型足够大时表现才比 VQ 好。