©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:
▲ 图二:不断“翻书”