我在Performer中发现了Transformer-VQ的踪迹

2024-01-15 12:01

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 月之暗面

研究方向 | NLP、神经网络

前些天我们在《VQ一下Key，Transformer的复杂度就变成线性了》介绍了“Transformer-VQ”，这是通过将 Key 序列做 VQ（Vector Quantize）变换来实现 Attention 复杂度线性化的方案。

诚然，Transformer-VQ 提供了标准 Attention 到线性 Attentino 的一个非常漂亮的过渡，给人一种“大道至简”的美感，但熟悉 VQ 的读者应该能感觉到，当编码表大小或者模型参数量进一步增加时，VQ 很可能会成为效果提升的瓶颈，因为它通过 STE（Straight-Through Estimator）估计的梯度大概率是次优的（FSQ 的实验结果也算是提供了一些佐证）。

此外，Transformer-VQ 为了使训练效率也线性化所做的梯度截断，也可能成为将来的效果瓶颈之一。

为此，笔者花了一些时间思考可以替代掉 VQ 的线性化思路。从 Transformer-VQ 的形式中，笔者联想到了 Performer，继而“顺藤摸瓜”地发现原来 Performer 可以视为 Soft 版的 Transformer-VQ。进一步地，笔者尝试类比 Performer 的推导方法来重新导出 Transformer-VQ，为其后的优化提供一些参考结果。

前情回顾

首先，让我们花一些时间回顾一下 Transformer-VQ。设

，Transformer-VQ 的关键，是对做了如下 VQ 近似：

这里的都是矩阵，其中是可训练的参数，则定义为：

说白了，VQ 就是用与最相近的那个来近似。在这个近似之下，我们有（简单起见，以 Encoder 为例）

了解线性 Attention 的读者很容易认出来，最后一个式子的运算就是线性复杂度的，它就是本文的主角之一 Transformer-VQ（的分子，还有分母同理）。

没有很复杂的推导，线性 Attention 就出来了，这就给我们一种感觉，仿佛我们是在对 Key 做近似的“不经意间”就将 Attention 的复杂度降为了线性，美感十足。因此，再次回到了我们已经提过多次的评价——Transformer-VQ 提供了标准 Attention 到线性 Attentino 的一个非常漂亮的过渡。

似曾相识

Transformer-VQ 的让笔者联想到了之前的文章《Transformer升级之路：从Performer到线性Attention》。在那篇文章中，笔者对 Performer 的结果做了一些简化，然后断言线性 Attention 的的最佳激活函数是，而 Transformer-VQ 同样出现了，所以它们之间也许有着某种相关性。

为了挖掘这种联系，让我们请出 Performer，它基于一个漂亮的近似：

由于最后还要对所有的注意力归一化，所以去掉上式中的都不会影响最终结果，同时，如果假设的模长都相等（参考 JL 引理），那么的指数都减去也不会影响结果。于是，Performer 等价于用以下的格式做：

对比最后一个式子和（3），就会发现它们有诸多相似之处：不就相当于编码表？不就相当于？至于最后的，它以为 logits 做 softmax，突出的不就是与最相近的那个？而 softmax 的极限就是 one hot，所以这不正好对应着 Transformer-VQ 的矩阵？因此，这不能说一模一样，但也有六七分相似了。

依样葫芦

当然，上述结果更多的是一种形象的类比而不是等价性，因为Performer本质上基于完全不同的近似思路，比如它里边的是随机采样并固定下来的，这意味它们作为中心向量的近似程度其实是很差的。

但这种类似引发了一个思考：能否模仿 Performer 的思路来重新推导一遍 Transformer-VQ 呢？即像式（4）一样，先构造一个精确相等的结果，然后再转化为采样近似来得到线性版本。

经过几天的思考，笔者发现了一种可以构造出期望推导的方案。首先，我们借助狄拉克函数 [1] 写出

这是纯粹有狄拉克函数的定义给出的恒等式，还没涉及到任何精巧的运算或者近似。然而，当我们将它代入 Attention（的分子）时，出现了一些有意思的结果：

最后一个等号，不就正好是线性 Attention 的形式？！当然，由于需要对积分，所以上式跟《Transformer升级之路：5、作为无限维的线性Attention》[2] 一样，都是“无限维”的线性 Attention，暂时只有形式上的价值。

通常来说，我们会将理解为正态分布在的极限，这也意味着具有条件分布的意义。不过，从生成模型的角度来看，狄拉克函数就是单点分布，说白了就是把训练集背下来，所以它没有抽象和泛化能力。为了缓解这一点，我们将用 GMM [3]（Gaussian Mixture Model，高斯混合模型）来近似：

代入式（7），然后取的极限，我们就得到

这就得到一个有限维的线性 Attention。如果将对齐 Transformer-VQ 的 one hot 分布的定义，那么得到的结果就是 Transformer-VQ 的式（3）。

文章小结

本文介绍了笔者的一个发现：早期的线性 Attention 工作 “Peformer” 可以视为一个 “Soft” 版的 Transformer-VQ。然后，在这个观察上进一步得到了 Transformer-VQ 的一个新推导：利用狄拉克函数将标准 Attention 转化为无限维线性 Attention，然后加上 GMM 近似就可以得到 Transformer-VQ。