给 RNN 装上“会增长的记忆”：解读 Google 的 Memory Caching

这篇论文来自 Ali Behrouz 团队，是他们 Titans / Atlas / Miras / Nested Learning 这条测试时记忆（test-time memorization）研究线上的最新一篇。它想回答一个老问题：为什么效率更高的递归模型，在长上下文和召回任务上始终打不过 Transformer？ 答案被归结为一个结构性瓶颈——RNN 只有一块固定大小的记忆，被迫不断遗忘。而本文给出的解法朴素得出奇：别让旧记忆消失，把它们存档下来。

一、问题背景：固定记忆 vs 增长记忆

要理解这篇论文，先要看清两类模型在"记忆"上的根本差异。

1.1 注意力 = 增长的关联记忆

标准注意力（Vaswani et al., 2017）对每个查询 $q_i$，要遍历全部历史 token：

$$y_i=\frac{1}{Z_i}\sum _{t=1}^i\exp \left(q_i^{\top }{k}_t\right)v_t$$

它本质上是一个容量随上下文增长的关联记忆——每个 token 都被原样缓存（即 KV-cache），因此检索时能直接访问序列任意位置。这正是它召回能力强的来源，但代价是 $\mathcal{O}(L^2)$ 的计算和推理时不断膨胀的显存。

1.2 线性注意力 / RNN = 固定的压缩记忆

线性注意力（Katharopoulos et al., 2020）把 $\exp (\cdot )$ 换成可分解核 $\varphi (\cdot )$，于是注意力退化成一个固定大小的记忆矩阵的递归更新：

$${\mathcal{M}}_t={\mathcal{M}}_{t-1}+v_t\varphi (k_t{)}^{\top },y_i=\frac{1}{Z_i}{\mathcal{M}}_i\varphi (q_i)$$

记忆只有一块、大小不变。序列一长，这块记忆就"溢出"——被迫遗忘旧信息。这是 RetNet、RWKV、DeltaNet、Titans 这一大类模型在 recall 密集任务上吃亏的共同根源。

1.3 测试时记忆视角

作者所在团队提出过一个统一框架：把序列模型的更新规则看作一个前向过程中动态进行的"在线学习/记忆"。最简形式下，记忆模块 $\mathcal{M}(\cdot )$ 在学习键值映射，每步求解一个带保留项的内部目标（attentional bias）：

$${\mathcal{M}}_{t+1}=\arg \underset{\mathcal{M}}{\min }\mathcal{L}(\mathcal{M}(k_t);v_t)+{\mathcal{R}}_{e_t}(\mathcal{M};{\mathcal{M}}_t)$$

这个视角是本文的关键支点：既然递归更新是一个优化过程，那么记忆的中间状态就是这个优化过程的检查点（checkpoint）。Memory Caching 做的事，本质上就是把这些检查点存下来。

二、Memory Caching 框架

2.1 核心机制：分段、缓存、聚合

做法分三步，直觉非常清晰：

标准 RNN 的输出是 $y_t={\mathcal{M}}_t(q_t)$；而 MC 给出一个聚合函数 $\mathrm{Agg}(\cdot )$：

$$y_t=\mathrm{Agg}(\{{\mathcal{M}}_{L^{(1)}}^{(1)}(\cdot ),\dots ,{\mathcal{M}}_{L^{(s-1)}}^{(s-1)}(\cdot )\};{\mathcal{M}}_t^{(s)}(\cdot );q_t)$$

2.2 为什么是 $\mathcal{O}(NL)$

这里有个值得敲黑板的点：记忆的写入（更新）完全没改，仍是 $\mathcal{O}(L)$；多出来的成本只在读取——每个 token 要对 $N$ 块缓存各做一次前向，于是每 token 检索成本 $\mathcal{O}(N)$，整体落在 $\mathcal{O}(NL)$，其中 $1\le N\le L$。这正是它能"几乎免费地"嫁接到任意 RNN 上的原因。两个端点也很自然：

• $N=1$：不缓存任何状态 → 退化成普通递归模型 $\mathcal{O}(L)$；
• $N=L$：每个 token 自成一段、全部缓存 → 逼近注意力的"直接访问" $\mathcal{O}(L^2)$。

三、四种聚合策略

聚合函数 $\mathrm{Agg}(\cdot )$ 怎么设计，是本文的主要技术贡献。论文给了四种，从最朴素到最精巧：

3.1 残差记忆（Residual Memory）

最简单——把在线记忆和所有缓存记忆的输出直接相加，相当于在记忆状态间加残差连接：

$$y_t=\underset{\text{在线记忆}}{{\mathcal{M}}_t^{(s)}(q_t)}+\underset{\text{缓存记忆}}{\sum _{i=1}^{s-1}{\mathcal{M}}_{L^{(i)}}^{(i)}(q_t)}$$

这里有个微妙之处：如果记忆是线性的（矩阵），这些缓存可以提前求和，公式会塌缩回一块固定记忆，理论上等于没改。但实验上它仍然有效——作者解释这是因为它充当了一个"保留算子（retention operator）"，强化了对远期信息的访问。

3.2 门控残差记忆（GRM）

针对上面的塌缩问题，引入输入相关的门控 ${\gamma }_t^{(i)}\in [0,1]$ 给每段加权：

$$y_t={\gamma }_t^{(s)}{\mathcal{M}}_t^{(s)}(q_t)+\sum _{i=1}^{s-1}{\gamma }_t^{(i)}{\mathcal{M}}_{L^{(i)}}^{(i)}(q_t)$$

精妙的是门控的定义方式。如果只让 $\gamma$ 是输入的线性投影，它就退化成"按位置过滤"——只看段的位置、不看段的内容。作者改用一个连接参数 $u_t=x_t{W}_u$，去和该段内容的均值池化表示做相似度：

$${\gamma }_t^{(i)}=⟨u_t,\mathrm{MeanPooling}(S^{(i)})⟩,\text{再过 softmax}$$

这样门控就成了上下文感知的检索：当前 token 该多看哪一段，取决于它和那一段内容的相似度。因为 $\gamma$ 逐 token、依赖输入，所以即使线性记忆也不会塌缩。

3.3 记忆汤（Memory Soup）

灵感来自 model soup（权重平均）。前面是"先各自检索、再加权输出"，记忆汤反过来——先把各段缓存记忆的参数加权平均成一块全新的、数据相关的记忆 ${\mathcal{M}}_t^{\ast }$，再用查询去读：

$${\theta }_{{\mathcal{M}}_t^{\ast }}=\{\sum _{i=1}^s{\gamma }_t^{(i)}{W}_1^{(i)},\dots ,\sum _{i=1}^s{\gamma }_t^{(i)}{W}_c^{(i)}\},y_t={\mathcal{M}}_t^{\ast }(q_t)$$

对线性记忆，它和 GRM 数学上等价（线性下"先平均参数"与"先各自输出再求和"一回事）。但对深度/非线性记忆（如 DLA、Titans）就分道扬镳了：它本质是给每个 token 临时插值出一个专属的非线性检索函数。

3.4 稀疏选择缓存（SSC）

前三种每个 token 都要看遍所有 $N$ 块缓存，超长序列开销大。SSC 借鉴 MoE，用一个路由器按相关性分数 $r_t^{(i)}=⟨u_t,\mathrm{MeanPooling}(S^{(i)})⟩$ 选出 Top-$k$ 块缓存（加当前在线记忆）来读：

$${\mathcal{R}}_t=\mathrm{Top}\text{–}k(\{r_t^{(i)}{\}}_{i=1}^{s-1}),y_t={\gamma }_t^{(s)}{\mathcal{M}}_t^{(s)}(q_t)+\sum _{i\in {\mathcal{R}}_t}{\gamma }_t^{(i)}{\mathcal{M}}_{L^{(i)}}^{(i)}(q_t)$$

工程上的好处很实在：段的均值池化可预计算、Top-$k$ 可并行，而且只需把"被选中"的记忆加载进加速器，训练和推理的显存都能省。作者把它解读成一个"稀疏的统一记忆"——写入只激活一小块参数（无干扰地存），读取激活更大一块（自适应地取）。

四、两个理论亮点

除了工程方法，论文还有两处漂亮的理论桥接，把 RNN、混合模型、注意力统一在了一个框架下。

4.1 用 MC 重新发现"门控注意力"

把段长设为 1、记忆设成"无值的向量记忆"（每个 token 自成一段、只存自己），套上 MC 推导出来的，恰好是门控全局 softmax 注意力：

$$y_t=\left(\sum _{i=1}^t\frac{\exp (u_t^{\top }{k}_i)}{{\sum }_{\ell }\exp (u_t^{\top }{k}_{\ell })}{v}_i^{\prime }\right)\otimes \sigma (x_t{W}_Q)$$

更进一步，他们指出"压缩器 + 全局注意力"这种近来很流行的混合架构，等价于段长为 1、缓存 checkpoint 的 MC。这就给"为什么混合模型有效"提供了一个解释：注意力本质上在强制缓存过去的输入，从而扩大了递归模型的有效记忆容量。 而当查询从 $q_t=1$ 放开到 $q_t=x_t{W}_Q$，MC 还能变成一种"自组装注意力"——每个查询可以自己构造送进注意力块的输入序列 $\{{\mathcal{M}}_{L^{(i)}}^{(i)}(q_t)\}$，而非固定不变。

⚠️ 作者诚实地注明：这个等价只在"过度简化"的版本下成立，一旦考虑归一化和前馈层，两者的表达力会有差异。

4.2 分段就是压缩与算力的旋钮

分段方式决定了 $N$，也就决定了压缩程度与计算成本的折衷。

设每段大小为 $C=L/N$，则总成本约 $\mathcal{O}(p\cdot L^2/C)$——这是个常数更小的 Transformer。另一种是对数分段：把 $L$ 写成二进制，按非零位取 2 的幂作段长（如 $37=(100101{)}_2\to 32,4,1$），此时 $N\le {\log }_{2}L$，成本压到 $\mathcal{O}(pL\log L)$。代价是对远期 token 的分辨率太低，召回能力受损。论文还把同期工作 Log-Linear Attention（Guo et al., 2025）重新表述成"GRM + 对数分段"的一个 MC 特例（记作 Log-Linear++），作为实验里的对照基线。

五、实验结果

作者在 760M（30B token）和 1.3B（100B token）两个规模上，把 MC 挂到 SWLA、DLA、Titans 三种递归架构上验证。

5.1 语言建模与常识推理

所有 MC 变体都稳定优于其基线。以 1.3B 规模为例（平均分，越高越好）：

整体看：GRM 效果最好，SSC 次之；MC 对 Titans 的平均增益约 +0.8%（论文原话）。

5.2 长上下文召回（NIAH）

大海捞针任务上 MC 的优势在长上下文处最明显。以最难的 S-NIAH-3（UUID 检索）为例：

MC 把基线在 8K 处从 42.8 拉到 69.0，提升巨大，并全面优于 Log-Linear；但在 16K 处仍不及 Transformer。

5.3 上下文检索：缩小差距，但未反超

这是最该如实呈现的结果。6 项检索任务平均分上：Transformer 仍以 41.0 居首，MC 最好的 Titans+GRM 是 40.5——是"逼近"而非"反超"。MC 把递归模型从 ~31 拉到 ~40，确实大幅缩小了差距。此外 LongBench、MQAR 上 MC 变体也都优于各自基线。效率方面，SSC 是"两全其美"的那个：开销接近基线 RNN，长序列上远比 Transformer 高效。

六、批判性评估：几点要保留的判断

技术读者会注意到结果是否被过度包装。以下几点建议如实看待：

1. 贡献是增量而非颠覆。 语言建模上 MC 对 Titans 的平均增益约 +0.8%，体量不大；在最关键的上下文召回任务上，Transformer 仍稳居第一。论文摘要本身也诚实地写明了"close the gap"而非"超越"。

2. 这是同一团队在自家框架内的延续。 大量自引（Titans、Atlas、Miras、Nested Learning），核心叙事建立在他们自己的"关联记忆/测试时记忆"范式上。作为对照基线的 Log-Linear++，也是他们把别人的方法重新表述进自家框架后再来比较——逻辑成立，但属于对自己有利的设定。

3. 有一处实验结果可疑。 消融表（Table 5）中"Shared $u$ and $q$"两行的三列指标全是 00.0 / 00.0 / 00.0。这要么是该配置导致训练彻底崩溃的真实发现，要么是 v1 预印本里尚未填入的占位符。鉴于这是 2026 年 2 月的 arXiv v1，更稳妥的态度是把它当作不确定项，而非直接转述为"共享投影会致命崩溃"的结论。

4. 效率收益强依赖分段选择。 $\mathcal{O}(NL)$ 在 $N$ 增大时会重新逼近 $\mathcal{O}(L^2)$；而且除 SSC 外的三种变体仍需存储全部 $N$ 块记忆状态，显存并非免费。真正缓解显存的只有 SSC。

5. 验证规模有限。 全部实验在学术规模（≤1.3B 参数、≤100B token）完成，能否迁移到前沿规模仍是未知数——这也是作者自己强调"作为概念验证（proof of concept）"的原因。

七、总结

Memory Caching 的核心，是把递归模型的"记忆状态"从一个被不断覆盖的单点，变成一条可以回看的检查点序列。它最有价值的地方有两层：

• 工程层面：给出一个 $\mathcal{O}(NL)$ 的旋钮，让 RNN 几乎免费地获得随长度增长的有效记忆，且四种聚合策略各有取舍（GRM 性能最佳、SSC 效率最佳）；
• 理论层面：把混合架构、门控注意力都纳入同一框架解释，为"为什么这些经验配方有效"提供了统一视角。

它不是 Transformer 的颠覆者，而是递归模型阵营里一次扎实的、诚实的能力补强——在长上下文召回这块短板上把差距显著缩小，同时守住了递归模型的效率优势。对关注高效长序列建模的工程师和研究者来说，这是一篇值得细读的方法论文章。