VQ一下Key，Transformer的复杂度就变成线性了

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 月之暗面

研究方向 | NLP、神经网络

Efficient Transformer，泛指一切致力于降低 Transformer 的二次复杂度的工作，开始特指针对 Attention 的改进，后来更一般的思路，如傅立叶变换、线性 RNN 等，也被归入这个范畴。不得不说，为了降低 Transformer 的二次复杂度，各路大牛可谓是“八仙过海，各显神通”，各种神奇的思路“百花齐放”，笔者也从中学习到了不少理论知识。

然而，尽管 Efficient Transformer 在理论上是精彩的，但实际上该领域一直都是不愠不火的状态，并没有实际表现十分出色的模型，在 LLM 火爆的今天，甚至已经逐渐淡出了大家的视野，也淡出了笔者的兴趣范围。

不过，最近有一篇论文《Transformer-VQ: Linear-Time Transformers via Vector Quantization》[1]，却让笔者为之拍案叫绝。作者非常高明地洞察到，只需要对标准 Attention 的 Key 做一下 VQ（Vector Quantize），复杂度就会自动降低为线性！这种线性化思路保留了标准 Attention 的形式，是标准 Attention 到线性 Attention 的一个完美过渡，同时最大程度上保留了标准 Attention 的能力。

高效难题

说起来，本站也算是比较早关注 Efficient Transformer 相关工作了，最早可以追溯到 2019 年解读 Sparse Transformer 的一篇文章《为节约而生：从标准Attention到稀疏Attention》。此后，陆续写的关于 Efficient Transformer 的其他文章还有：

线性Attention的探索：Attention必须有个Softmax吗？

Performer：用随机投影将Attention的复杂度线性化

Nyströmformer：基于矩阵分解的线性化Attention方案 [2]

Transformer升级之路：从Performer到线性Attention

线性Transformer应该不是你要等的那个模型

FLASH：可能是近来最有意思的高效Transformer设计

Google新作试图“复活”RNN：RNN能否再次辉煌？

然而，正如本文开头所说，尽管 Efficient Transformer 已有不少工作，也曾被大家寄予厚望，但实际上该领域一直都没什么能“出圈”的作品，这其中的原因可能是：

1. 不少 Efficient Transformer 的提速以牺牲效果为代价；

2. 很多 Efficient Transformer 的复杂度降低仅仅是理论上的，实际使用提升不明显；

3. 有些 Efficient Transformer 难以用来训练 Causal LM，所以在 LLM 流行的今天就没有了用武之地；

4. Flash Attention 的出现表明即便是标准的 Transformer 仍有很大的提速空间。

VQ一下

那么，Transformer-VQ 为何又具备的“出圈”潜力？

简单来说，Transformer-VQ 就是对 Attention 的 Key 向量序列进行了“聚类”，并用所属类的类别中心近似原向量，然后 Attention 的复杂度就变成线性了。也就是说，Transformer-VQ 仅仅改变了 Key 的形似，其余部分（理论上）完全不变，所以这是一种对 Attention 改动非常小的线性化方案，也能非常清楚体现出线性化后损失的精度在哪里（即用类别中心近似原向量的差距）。

铺