时空之章：将Attention视为平方复杂度的RNN

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

近年来，RNN 由于其线性的训练和推理效率，重新吸引了不少研究人员和用户的兴趣，隐约有“文艺复兴”之势，其代表作有 RWKV [1]、RetNet [2]、Mamba [3] 等。当将 RNN 用于语言模型时，其典型特点就是每步生成都是常数的空间复杂度和时间复杂度，从整个序列看来就是常数的空间复杂度和线性的时间复杂度。

当然，任何事情都有两面性，相比于 Attention 动态增长的 KV Cache，RNN 的常数空间复杂度通常也让人怀疑记忆容量有限，在 Long Context 上的效果很难比得上 Attention。

在这篇文章中，我们表明 Causal Attention 可以重写成 RNN 的形式，并且它的每一步生成理论上也能够以 的空间复杂度进行（代价是时间复杂度非常高，远超平方级）。这表明 Attention 的优势（如果有的话）是靠计算堆出来的，而不是直觉上的堆内存，它跟 RNN 一样本质上都是常数量级的记忆容量（记忆瓶颈）。

超越线性的RNN

RNN 的支持者通常会给出一个看上去让人难以反驳的观点：想想你的大脑是 RNN 还是 Attention？ 

直觉来想，RNN 推理的空间复杂度是常数，而 Attention的 KV cache 是动态增长的，再考虑到人的脑容量是有限的，从这一点来看不得不说确实 RNN 更接近人脑。

然而，即便可以合理地认为脑容量限制了人每步推理的空间复杂度是常数，但它并没有限制每步的时间复杂度是常数，又或者换个说法，即便人的每步时间复杂度是常数，但人处理长度为 L 的序列时未必只扫描一遍序列（比如“翻书”），所以总的推理步数可能明显超出L，从而导致了非线性的时间复杂度。 

考虑到这一点，笔者“突发奇想”：是否可以一般化地考虑常数空间复杂度、非线性时间复杂度的 RNN 模型，来补足主流 RNN 的所没有的能力（比如上面说的翻书）？对于语言模型任务，假设样本是 a b c d e，那么训练任务就是输入 a b c d，预测 b c d e，常见的 RNN 如下图：

▲ 图一：常见RNN

这种 RNN 的问题就是没有翻书能力，每个输入读完就丢了。而 Attention 的特点就是每读一个 token，就完整地翻一遍历史，虽然这个做法可能存在效率问题，但它无疑是引入翻书能力的最简单粗暴的方式。而为了给 RNN 补上翻书能力，我们完全可以模仿 Attention 的做法来使用 RNN：

▲ 图二：不断“翻书”