DiVeE：一种非常简洁的VQ训练方案
原文链接 | 2025-10-08 | 苏剑林
对于坚持离散化路线的研究人员来说，VQ（Vector Quantization）是视觉理解和生成的关键部分，担任着视觉中的“Tokenizer”的角色。它提出在2017年的论文《Neural Discrete Representation Learning》，笔者在2019年的博客《VQ-VAE的简明介绍：量子化自编码器》也介绍过它。
然而，这么多年过去了，我们可以发现VQ的训练技术几乎没有变化，都是STE（Straight-Through Estimator）加额外的Aux Loss。STE倒是没啥问题，它可以说是给离散化运算设计梯度的标准方式了，但Aux Loss的存在总让人有种不够端到端的感觉，同时还引入了额外的超参要调。
幸运的是，这个局面可能要结束了，上周的论文《DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick》提出了一个新的STE技巧，它最大亮点是不需要Aux Loss，这让它显得特别简洁漂亮！
离散编码#
老规矩，我们还是先回顾一下VQ现有的训练方案。首先要指出的是，VQ（Vector Quantizatio）本身实际上是一个很古老的概念，可以追溯到上世纪80年代，原意是将向量聚类并用对应的聚类中心代替，从而实现数据压缩。
但我们这里说的VQ，主要是指论文《Neural Discrete Representation Learning》所提的VQ-VAE中的VQ。当然，VQ的定义本身没啥变化，都是向量到聚类中心的映射，VQ-VAE一文的核心是提供了一套对latent变量进行VQ然后解码重构的端到端训练方案，其难点是VQ这一步是离散化操作，没有现成的梯度，需要为它设计梯度。
用公式来说，普通的AE（AutoEncoder）是：
\[\begin{equation}z = encoder(x),\quad \hat{x}=decoder(z),\quad \mathcal{L}=\Vert x - \hat{x}\Vert^2 \end{equation}\]
其中\(x\)是原始输入，\(z\)是编码向量，\(\hat{x}\)是重构结果。VQ-VAE想要做的是，基于VQ的思想，将\(z\)变成编码表\(E=\{e_1,e_2,\cdots,e_K\}\)中的一个：
\[\begin{equation}q = \newcommand{argmin}{\mathop{\text{argmin}}}\argmin_{e\in E} \Vert z - e\Vert\end{equation}\]
然后\(decoder\)以\(q\)为输入来进行重构。由于\(q\)是编码表的下标一一对应，所以\(q\)实际上是\(x\)的一个整数编码。当然，为了保证重构效果，实际使用时肯定不止编码成单个向量，而是多个向量，VQ之后就变成了多个整数。所以，VQ-VAE想要做的事情，就是将输入编码成一个整数序列，这跟文本的Tokenizer异曲同工。
梯度设计#
现在我们要训练的模块包括\(encoder\)、\(decoder\)和编码表\(E\)，由于VQ操作带有\(\argmin\)运算，梯度到\(q\)这里就断了，无法传到\(encoder\)中去。
VQ-VAE使用了一个名为STE的技巧，它说送入\(decoder\)的还是VQ后的\(q\)，但反向传播求梯度的时候，我们用的是VQ前的\(z\)，这样就可以把梯度传回给\(encoder\)，它可以用stop_gradient算子（\(\newcommand{sg}{\mathop{\text{sg}}}\sg\)）实现：
\[\begin{equation}z = encoder(x),\quad q = \argmin_{e\in E} \Vert z - e\Vert,\quad z_q = z + \sg[q - z],\quad \hat{x} = decoder(z_q)\end{equation}\]
简单来说，STE实现的效果是\(z_q=q\)但\(\nabla z_q = \nabla z\)，这样\(encoder\)就有梯度了，但\(q\)没梯度了，没法优化编码表了。为了解决这个问题，VQ-VAE补了两项Aux Loss：
\[\begin{equation}\mathcal{L} = \Vert x - \hat{x}\Vert^2 + \beta\Vert q - \sg[z]\Vert^2 + \gamma\Vert z - \sg[q]\Vert^2 \end{equation}\]
这两项Loss分别代表\(q\)向\(z\)靠近和\(z\)向\(q\)靠近，跟VQ的原始思想一致。STE跟这两项Aux Loss组合起来，就构成了标准的VQ-VAE。此外还有一个简单的变体，就是直接设\(\beta=0\)，但用\(z\)的滑动平均来更新编码表，这也等价于指定用SGD来更新\(q\)的Aux Loss。
这里顺便指出一点，VQ-VAE虽然被原论文冠以“VAE”之名，但它实际上就是个AE，所以叫“VQ-AE”原则上更贴切，只不过这个名字已经被叫开了，我们只好沿用它。后来的VQGAN，则是在VQ-VAE基础上叠加GAN Loss等技巧，提高了重构的清晰度。
替代作品#
对于笔者来说，这两个额外的Aux Loss是让人很难受的地方。估计业界跟笔者同样感受的人也不少，所以时不时也有一些相关的改进工作出来。
其中，最为“釜底抽薪”的做法是换用VQ以外的离散化方案，比如《简单得令人尴尬的FSQ：“四舍五入”超越了VQ-VAE》介绍的FSQ，它就不需要Aux Loss。如果说VQ是对高维向量进行聚类，那么FSQ就是对低维向量做“四舍五入”，从而实现离散化。然而，笔者在这篇文章评价过，FSQ并不能在任意场景都取代VQ，所以改进VQ本身依然是有价值的。
在提出DiVeQ之前，其实原作者还提出过一个名为“NSVQ”的方案，向“废除”Aux Loss迈出了一小步，它将\(z_q\)改为
\[\begin{equation}z_q = z + \Vert q - z\Vert \times \frac{\varepsilon}{\Vert \varepsilon\Vert},\qquad \varepsilon\sim\mathcal{N}(0, I)\label{eq:nsvq}\end{equation}\]
这里\(\varepsilon\)是跟\(z,q\)同大小的向量，其分量服从标准正态分布。替换为这个新的\(z_q\)之后，由于\(\Vert q - z\Vert\)的可导性，\(q\)也是有梯度的，所以原则上不需要Aux Loss就能训练编码表了。NSVQ的几何意义很直观，就是在“以\(z\)为中心、\(\Vert q-z\Vert\)为半径的圆”上均匀采样，缺点是此时送入\(decoder\)的并不是\(q\)，而推理阶段我们关心的是\(q\)的重构效果，所以NSVQ的训练和推理具有不一致性。
主角登场#
从NSVQ出发，如果想保持前向为\(q\)，但同时保留\(\Vert q - z\Vert\)带来的梯度，那么很容易提出改进版：
\[\begin{equation}z_q = z + \Vert q - z\Vert \times \sg\left[\frac{q - z}{\Vert q - z\Vert}\right]\label{eq:diveq0}\end{equation}\]
它在前向的时候严格有\(z_q = q\)，但反向时保留了\(z\)和\(\Vert q - z\Vert\)的梯度。这便是DiVeQ论文附录中的“DiVeQ-detach”，而正文的DiVeQ，则算是式\(\eqref{eq:diveq0}\)和\(\eqref{eq:nsvq}\)的某种插值
\[\begin{equation}z_q = z + \Vert q - z\Vert \times \sg\left[\frac{q - z + \varepsilon}{\Vert q - z + \varepsilon\Vert}\right],\qquad \varepsilon\sim\mathcal{N}(0, \sigma^2 I)\label{eq:diveq}\end{equation}\]
很明显，当\(\sigma=0\)时，结果是“DiVeQ-detach”，当\(\sigma\to\infty\)时，结果是“NSVQ”。原论文附录对\(\sigma\)做了搜索，大致结论是\(\sigma^2 = 10^{-3}\)是一个普遍较优的选择。
论文的实验结果显示，尽管式\(\eqref{eq:diveq}\)引入了随机性，同时也有一定程度的训练推理不一致性，但相比式\(\eqref{eq:diveq0}\)它的效果更好。不过，就笔者的审美而言，效果不能以牺牲优雅为代价，所以式\(\eqref{eq:diveq0}\)的“DiVeQ-detach”才是笔者心中的理想方案，下面分析中的DiVeQ都是指代“DiVeQ-detach”。
理论分析#
遗憾的是，原论文并没有多少理论分析，所以这一节笔者尝试对DiVeQ的有效性、它跟VQ原本训练方案的相关性等做一个基本的分析。首先，考虑式\(\eqref{eq:diveq0}\)的一般形式
\[\begin{equation}z_q = z + r(q, z) \times \sg\left[\frac{q - z}{r(q, z)}\right]\end{equation}\]
其中\(r(q,z)\)是任意关于\(q,z\)的可导的标量函数，它可以视为\(q,z\)的任意的距离函数。我们将损失函数记为\(\mathcal{L}(z_q)\)，那么它的微分是
\[\begin{equation}d\mathcal{L} = \langle\nabla_{z_q} \mathcal{L},d z_q\rangle = \left\langle\nabla_{z_q} \mathcal{L},dz + dr \times\frac{q-z}{r}\right\rangle = \langle\nabla_{z_q} \mathcal{L},d z\rangle + \langle\nabla_{z_q} \mathcal{L}, q-z\rangle d(\ln r)\end{equation}\]
其中\(\langle\nabla_{z_q} \mathcal{L},d z\rangle\)是原本VQ就有的，DiVeQ则多出了额外的\(\langle\nabla_{z_q} \mathcal{L}, q-z\rangle d(\ln r)\)，或者说，相当于引入了Aux Loss \(\sg[\langle\nabla_{z_q} \mathcal{L}, q-z\rangle] \ln r\)，如果\(r\)代表\(q,z\)的某种距离函数，那么它就是在拉近\(q,z\)的距离，跟VQ引入的Aux Loss异曲同工，这就成功从理论上解释了DiVeQ。
但别高兴太早，这个解释成立的前提是系数\(\langle\nabla_{z_q} \mathcal{L}, q-z\rangle > 0\)，否则反而在拉大距离了。为了证明这一点，我们考虑损失函数\(\mathcal{L}(z)\)在\(z_q\)处的一阶近似
\[\begin{equation}\mathcal{L}(z) \approx \mathcal{L}(z_q) + \langle\nabla_{z_q} \mathcal{L}, z - z_q\rangle = \mathcal{L}(z_q) + \langle\nabla_{z_q} \mathcal{L}, z - q\rangle\end{equation}\]
即\(\langle\nabla_{z_q} \mathcal{L}, q-z\rangle\approx \mathcal{L}(z_q) - \mathcal{L}(z)\)。注意\(z,z_q\)分别是VQ前、后的特征，VQ是一个损失信息的过程，所以用\(z\)来做目标任务（比如重构）会比\(z_q\)容易，因此只要开始收敛，可以预期\(z\)的损失函数更低，即\(\mathcal{L}(z_q) - \mathcal{L}(z) > 0\)，这样我们就证明了\(\langle\nabla_{z_q} \mathcal{L}, q-z\rangle > 0\)大概率是成立的。
改进方向#
严格来讲，\(\langle\nabla_{z_q} \mathcal{L}, q-z\rangle > 0\)只能说是DiVeQ有效的必要条件，要充分论证它的有效性还需要证明这个系数“恰到好处”。由于\(r(q,z)\)的任意性，我们只能具体函数具体分析。我们考虑\(r(q,z)=\Vert q-z\Vert^{\alpha}\)，那么它等效于引入如下Aux Loss
\[\begin{equation}\sg[\langle\nabla_{z_q} \mathcal{L}, q-z\rangle] \ln \Vert q-z\Vert^{\alpha}\approx \sg[\mathcal{L}(z_q) - \mathcal{L}(z)]\times \alpha\ln \Vert q-z\Vert\end{equation}\]
系数\(\mathcal{L}(z_q) - \mathcal{L}(z)\)跟主损失\(\mathcal{L}(z_q)\)是齐次的，这意味着它能够较好地自适应主损失的Scale，还能根据VQ前后的效果差距来调节Aux Loss权重。至于\(\alpha\)取多少好，笔者认为要看实验，个人尝试调了一下，发现\(\alpha=1\)确实普遍表现较好。有兴趣的读者也可以自行尝试调整\(\alpha\)，甚至尝试换用别的\(r(q, z)\)。
需要指出的是，DiVeQ只是提供了一种新的免Aux Loss的VQ训练方案，原则上它不解决VQ的其他问题，比如编码表利用率低、编码表坍缩等，原本在“STE + Aux Loss”场景下性质有效的增强技巧，都可以考虑叠加到DiVeQ上。原论文就将DiVeQ跟SFVQ结合起来，提出SF-DiVeQ，用来缓解编码表坍缩等问题。
不过，个人感觉SFVQ有点繁琐，所以这里就不打算展开介绍了，而且作者选择叠加SFVQ，更多是因为SFVQ本就是他之前的作品，属于一脉相承下来了。笔者更喜欢的是《VQ的又一技巧：给编码表加一个线性变换》中介绍的线性变换技巧，即在编码表后面再加一个线性变换，实测这样也能明显增强DiVeQ的效果。
文章小结#
本文介绍了VQ（Vector Quantization）的一种新的训练方案，它只需通过STE实现，不需要加额外的Aux Loss，从而显得特别简洁优雅。