SimVQ:使用一个线性层解决矢量量化模型中的表示坍缩问题

摘自:https://spaces.ac.cn/archives/10519

音频表征工作

SimVQ: Addressing Representation Collapse in Vector Quantized Models with One Linear Layer

Github: https://github.com/youngsheen/SimVQ

论文提出只在 VQ 的编码表多加一个线性变换[W],无需其他改动,就能达到加速收敛、提升编码利用率、降低重构损失等效果,相当简单有效。

普通AE和VQ-VAE的数学形式:

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

改动

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

在编码表多乘了一个矩阵W,其他原封不动。

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

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

实验

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

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

分析:

直观来想,VQ对编码表的更新是比较“孤立”的,比如某个样本z被VQ为q,那么这个样本的梯度就只会影响q,不会影响编码表里的其他向量;但SimVQ不同,它不单会更新q,还会更新W,从几何意义上看,W就相当于编码表的基底,一旦更新W,那么整个编码表就会更新了。所以说,SimVQ使得整个编码表的“联动”更为密切,从而更有机会找到更优的解,而不是陷入“各自为政”的局部最优。

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

然而,物极必反,如果全体编码都使劲往高利用率方向走,那么反而可能会导致编码表坍缩(codebook collapse),因此SimVQ默认采用了一个很保守的策略:只更新W,所有的q在随机初始化后就不更新了,这样一来就几乎杜绝了编码表坍缩的可能性。好消息是,在适当的编码维度下,实验显示q,W都更新和只更新W的表现都差不多,所以读者可以按照自己的偏好选择具体的形式。

延伸:

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

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

这是因为不少人认为,深度学习模型之所以“好训”,过参数化有不可忽视的作用,因此贸然去掉已经充分验证的模型的过参数化风险很大。这里的“好训”,主要是指梯度下降这种理论上容易陷入局部最优的方法居然经常可以找到一个实际表现很好的解,这本身就是一件很不可思议的事情。还有《On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization》等工作,表明过参数化隐式地加速了训练,作用类似于SGD中的动量。

最后,VQ本质上可以理解为一种稀疏训练方案,所以SimVQ所带来的启发和改动,也许还能用于其他稀疏训练模型,比如MoE(Mixture of Experts)。当前的MoE训练方案中,Expert之间的更新也是比较独立的,只有被Router选中的Expert才会更新参数,那么是不是有可能像SimVQ一样,所有的Expert后都接一个共享参数的线性变换,用来提高Expert的利用效率?当然MoE本身跟VQ也有很多不同之处,这还只是个猜测。

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注