Transformer与RNN合体,谷歌打下显存门槛,解锁超长上下文
3 小时前 / 阅读约6分钟
来源:36kr
谷歌提出Memory Caching技术,通过缓存RNN历史状态快照,解决大模型处理长文本时的内存瓶颈,降低推理资源门槛,提升长上下文能力,缩小与Transformer差距。

最近,谷歌跟内存干上了。

上个月,谷歌的 TurboQuant 研究曾经引发过一场行业地震,其宣称能直接把大模型最吃显存的 KV Cache 压缩几倍,让市场开始担忧未来对内存的需求,引发了一波内存股暴跌,后续还有不小的学术争议。

这个星期,又有一篇谷歌论文引发了 AI 圈的关注,作者表示他们提出的方法解决了大模型处理长文本时的「内存瓶颈」(又来了),但实施的是完全不同的技术路线。

他们新开一条道路,通过对于大模型架构的机制创新,赋予了 RNN 「可生长的记忆容量」,找到了一种兼顾 Transformer 与 RNN 优势的新方法

该方法可以让 AI 处理更长的文本,解锁「超长上下文」能力,与此同时极大地降低了推理的资源门槛。

人们纷纷表示欢迎:大模型生产环境要的就是这个。

当前的大模型几乎都是基于 Transformer 架构打造的,它占据统治地位,主要归功于其可增长的记忆容量(注意力机制的计算和空间复杂度随上下文长度呈二次方

增长)。这使得 Transformer 极其擅长长上下文的信息召回(Recall)。

然而,这种二次方复杂度也导致了严重的算力和显存瓶颈,使得处理超长文本的成本极高。

为解决这个问题,社区一直在复兴 RNN、线性注意力模型(Linear Attention)以及状态空间模型(SSM,如 Mamba 等)。这些循环架构的优势在于记忆容量固定(复杂度为

),推理速度快且显存占用低。然而,它们的致命弱点在于:无论读了多少文本,都必须把所有过去的信息压缩到一个「固定大小」的隐藏状态(Hidden State)中。

这种「信息漏斗」导致它们在密集召回任务(Recall-intensive tasks,比如从极长的文档中精准提取某个细节)中表现往往远不如 Transformer。

为打破僵局,来自 Google Research 的研究团队提出了一种名为 Memory Caching(MC)的技术,据说简单而有效。

论文:《Memory Caching: RNNs with Growing Memory》

论文链接:https://arxiv.org/abs/2602.24281

在这项研究的视角中,存在一个架构光谱:一端是 Transformer(无压缩,Token 级缓存),另一端是传统的 RNN(全压缩,单一记忆)。而「记忆缓存(MC)」则解锁了介于两者之间的新形态:将成组 Token 压缩并缓存到长期记忆状态中,然后在需要时进行检索。

Transformer 会缓存每一个单独的标记,而 RNN 则考虑一个固定大小的记忆,并将上下文中的一切压缩到记忆的参数中。那么,如果我们把 RNN 的历史记忆也缓存下来会怎样?

简单来说,与其让 RNN 只维护一个不断被覆盖和更新的「当前状态」,不如定期对这些隐藏状态进行「快照打卡」(Caching checkpoints)。这样,在进行信息检索时,模型不仅可以查看当前的「在线记忆」,还能直接调取「缓存记忆」中的历史快照,瞬间找回过去的相关信息。

在研究的过程中,作为概念验证,研究人员提出了三种变体,基于过去的信息如何组合在一起。

1、门控残差记忆(Gated Residual Memory):使用查询从过去检索相关信息,然后执行类似注意力的池化来组合检索到的信息。实际上,RNN 的记忆在增长,因此解码成本也在增长:

2、记忆汤(Memory Soup):另一种结合过去记忆的方式,是直接结合记忆的权重,而不是针对特定查询的输出。在这种情况下,我们需要对过去记忆的权重执行类似注意力的池化操作,然后对池化后的记忆执行一次检索。同样,这种变体相对于上下文长度具有不断增长的有效记忆,因此解码成本也在增长。

3、稀疏选择性缓存(Sparse Selective Caching,SSC):到目前为止,似乎没有免费的午餐,我们需要在不断增长的有效记忆和每 token 恒定的解码成本之间做出选择。

于是作者提出了 SSC,这是一种类似于 MoBA 的专家混合模型,在序列维度上稀疏地选择过去缓存记忆的一个子集,从而引出一个模型,其有效记忆在增长,但其每 token 解码成本保持相对恒定:

那么效果如何呢?

该方法可以作为一种通用框架,插入到各种现有的循环架构中,如线性注意力模型,或作者之前提出的深度内存模型 Titans 等。

实验结果表明了其强大的有效性:

模型在语言建模和常识推理任务上的表现。

作者在 13 亿参数的模型上进行了语言建模、召回密集型、长上下文以及 needle-in-a-haystack 等实验,结果显示 MC 相较于基础模型提供了改进。其中包括:

长上下文能力提升:在语言建模和长上下文理解任务中,加入了 MC 机制的循环模型性能得到了全面提升。

缩小与 Transformer 的差距:在最具挑战性的「上下文内召回(in-context recall)」任务中,加入了 MC 的模型击败了目前最先进(SOTA)的其他循环模型。

依然存在上限:虽然 MC 极大地弥补了 RNN 的召回短板,并极大缩小了与 Transformer 之间的性能差距,但论文也指出,在纯粹的密集召回任务上限上,Transformer 依然保持着最佳的准确率。

总的来说,该研究利用一种极其优雅的算法直觉(缓存历史状态快照),解决了一个长期存在的理论难题,让非 Transformer 架构在实用性上又向前迈进了一大步。

尽管在极限的密集召回上,它依然尚未彻底超越 Transformer,但新路已经铺就,随着 RNN、SSM 等架构的持续进化,Transformer 一家独大的现状或许要有所改变了?