记忆缓存：具有增长记忆的循环神经网络 Memory Caching: RNNs with Growing Memory
通过缓存RNN记忆状态检查点，使有效记忆容量随序列长度增长，在固定与二次复杂度间灵活权衡
前置知识
循环神经网络（RNN）的记忆状态
RNN通过隐藏状态 $M_t$ 压缩历史信息，该状态在每个时间步根据当前输入 $x_t$ 和前一状态 $M_{t-1}$ 更新。传统RNN维护固定大小的记忆，这意味着随着序列变长，模型被迫遗忘早期信息。记忆状态本质上是一个关联记忆，学习键 $k_t$ 和值 $v_t$ 之间的映射关系。在现代RNN变体中，这个更新过程可以被解释为一个在线优化过程。
理解RNN固定记忆的局限性是理解本文动机的关键——Memory Caching正是为了解决这个根本性问题而提出的。
注意力机制与KV缓存
Transformer的注意力机制通过存储所有历史token的键值对来实现动态记忆访问。给定查询 $q_i$，注意力计算涉及softmax归一化的加权求和。这种机制提供O(L²)的复杂度，但具有随序列增长的记忆容量。KV缓存是推理时存储过去token键值对的技术，虽然提升了检索能力，但带来了二次方的内存和计算开销。
本文的核心目标是在RNN的线性复杂度和Transformer的二次复杂度之间找到平衡点。
线性注意力
线性注意力通过用可分离核替换softmax算子，将注意力转化为高效的递归形式。这种公式化将复杂度从O(L²)降低到O(L)，但代价是记忆容量固定。线性注意力可以被视为具有矩阵值记忆模块的递归神经网络。
本文将Memory Caching应用于线性注意力及其深度变体，展示了MC如何提升这类高效RNN的记忆能力。
测试时学习/记忆
这是一个统一框架，将序列模型的更新规则解释为动态的上下文学习/记忆过程。模型作为关联记忆，通过优化内部目标来学习键值映射。不同的注意力偏置选择可以恢复线性注意力、DeltaNet等架构。
本文正是基于这一视角，将缓存的记忆状态视为优化过程的检查点，从而增强模型在长序列上的信息检索能力。
混合专家（MoE）
MoE是一种稀疏激活架构，使用路由器根据输入选择性地激活专家子网络。给定输入x，路由器计算每个专家的相关性分数，然后选择Top-k个专家进行处理。这种设计允许模型拥有大量参数但每个输入只使用一小部分。
本文的稀疏选择性缓存（SSC）变体借鉴了MoE的路由器设计，让每个token选择性地访问最相关的缓存记忆。
研究动机
循环神经网络（RNN）通过固定大小的隐藏状态压缩整个输入序列，这种设计在处理长序列时面临根本性的记忆瓶颈。具体而言，当序列长度增长时，固定容量的记忆被迫遗忘早期信息，这在召回密集型任务中表现尤为明显。实验数据显示，在16K长度的Needle-In-A-Haystack任务中，基础DLA模型的准确率从4K时的96.4%骤降至16K时的44.0%。类似地，在检索任务中，基础RNN模型与Transformer的性能差距随输入长度增加而显著扩大。另一方面，虽然Transformer通过存储所有历史token的KV缓存实现了增长记忆，但其O(L²)的计算复杂度和巨大的推理内存需求限制了其在超长序列上的应用。这种效率与记忆容量之间的根本性权衡促使研究者探索介于两者之间的次二次架构。
本文的目标是本文的核心目标是提出一种通用技术，使得循环神经网络的有效记忆容量能够随序列长度增长，同时保持可控的计算复杂度。具体而言，作者希望实现一个灵活的中间地带：在RNN的O(L)复杂度和Transformer的O(L²)复杂度之间进行插值，提供O(NL)的复杂度，其中N是段数且1≤N≤L。这种方法需要满足三个条件：首先，它应该能够无缝应用于各种RNN架构；其次，它应该提供多种聚合策略以适应不同的记忆模块类型；最后，它应该在语言建模、长上下文理解和检索任务上展示出一致的性能提升。
与已有工作不同的是，本文的独特切入点在于将循环神经网络的记忆更新过程重新解释为一个优化过程，其中缓存的记忆状态是这个优化的检查点。这一视角揭示了一个关键洞察：传统RNN在整个序列上维护单一记忆状态，相当于只保留了优化过程的最终结果，而丢弃了所有中间状态。与现有的混合架构不同，Memory Caching提出了一种更通用且模块化的框架——它不改变底层RNN的记忆更新规则，而是在检索阶段引入对所有缓存记忆的访问。此外，本文首次系统性地研究了不同聚合策略对不同类型记忆模块的影响。
核心方法
Memory Caching的核心直觉可以通过一个类比来理解：想象你在阅读一本长书，传统RNN就像是读完后只记住最后一页的笔记，而Memory Caching则像是每读完一章就做一份章节摘要并保存起来，后续需要查找信息时可以回顾所有章节摘要。技术上，给定输入序列，MC首先将序列分割为N个段，每个段有相应的长度。对于每个段，使用记忆模块压缩该段的信息，记忆更新遵循原有的RNN规则。关键创新在于检索阶段：与传统RNN仅使用当前在线记忆不同，MC同时使用所有缓存的记忆状态来计算输出。这种设计的巧妙之处在于，它将计算复杂度从RNN的O(L)和Transformer的O(L²)扩展为可控的O(NL)，其中N是段数。当N=1时退化为标准RNN，当N=L时接近Transformer的行为。
MC与已有方法的本质区别在于其分段压缩加全局检索的范式。传统混合架构需要修改模型架构本身，而MC是一种可插拔的技术，可以应用于任何现有的RNN架构而不改变其记忆更新规则。更重要的是，MC引入了四种创新的聚合策略来适应不同类型的记忆模块：残差记忆通过简单的求和聚合所有缓存记忆；门控残差记忆（GRM）引入输入依赖的门控参数使模型能够动态调整各缓存记忆的贡献；记忆汤受权重平均启发通过插值缓存记忆的参数构建输入依赖的新记忆模块；稀疏选择性缓存（SSC）借鉴MoE思想使用路由器选择Top-k最相关的缓存记忆。
方法步骤详情
Memory Caching的完整流程包括以下步骤：步骤1是序列分段，将输入序列分割为N个段，分段策略可以是固定大小或对数大小。步骤2是段内记忆压缩，对每个段使用RNN的记忆更新规则处理段内所有token。步骤3是缓存检查点，每个段处理完成后缓存该段的最终记忆状态。步骤4是聚合检索，对于当前段中的token计算其输出时需要同时使用在线记忆和所有缓存记忆。步骤5是稀疏优化（可选），对于SSC变体使用路由器选择最相关的缓存记忆。
技术新颖性
Memory Caching的技术新颖性体现在多个层面。首先从理论角度看，MC首次明确建立了循环模型记忆与优化过程检查点之间的联系，这为理解RNN的记忆机制提供了新的理论框架。论文证明了当段大小为1且使用无值向量记忆时，MC能够恢复门控全局softmax注意力块，这一发现不仅揭示了注意力机制与RNN之间的深层联系。其次从架构设计角度看，MC提出了记忆汤这一全新聚合策略，通过在参数空间而非输出空间插值记忆状态。第三从效率角度看，SSC变体通过MoE风格的选择性访问机制实现了内存消耗的显著降低。最后MC的通用性使其能够无缝应用于线性注意力、深度线性注意力、Titans等多种架构。
实验结果
实验结果全面验证了Memory Caching在多个维度上的有效性。在语言建模任务上，MC增强的模型在所有基准上都展现出一致的性能提升。以1.3B参数规模为例，Titans + GRM将困惑度从基础Titans的15.60降低到15.37，平均准确率从56.82%提升到58.33%，这是一个显著的+1.51%改进。DLA + GRM同样表现出色，困惑度从16.31降至16.08，平均准确率从53.72%提升到55.96%（+2.24%）。在Needle-In-A-Haystack任务中，MC的优势更加明显。在最困难的S-NIAH-3任务上，Titans + GRM在16K长度上达到32.2%的准确率，相比基础Titans的21.2%提升了11个百分点。在检索任务上，MC显著缩小了RNN与Transformer之间的差距。在FDA任务上，Titans + GRM在16K长度上达到52.6%的准确率，远超基础Titans的32.5%。训练效率方面，MC变体在保持性能提升的同时，训练吞吐量介于RNN和Transformer之间，SSC变体展现出最佳的效率-性能权衡。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 语言建模（WikiText困惑度） | 困惑度 | Titans+GRM: 15.37 | 基础Titans: 15.60 | 降低0.23（1.5%相对提升） |
| 常识推理（平均准确率） | 准确率 | Titans+GRM: 58.33% | 基础Titans: 56.82% | 提升1.51个百分点 |
| Needle-In-A-Haystack（S-NIAH-3，16K） | 准确率 | Titans+GRM: 32.2% | 基础Titans: 21.2% | 提升11.0个百分点 |
| 检索任务（FDA，16K长度） | 准确率 | Titans+GRM: 52.6% | 基础Titans: 32.5% | 提升20.1个百分点 |
| LongBench（多文档QA平均） | 准确率 | Titans+GRM: 29.1% | 基础Titans: 26.1% | 提升3.0个百分点 |
局限与改进
尽管Memory Caching展现出显著的效果，但仍存在若干局限性。首先关于复杂度权衡：MC的O(NL)复杂度虽然介于RNN和Transformer之间，但当段数N接近L时计算开销仍可能变得显著。其次关于记忆模块类型的限制：实验主要集中在线性注意力和深度线性注意力上，对于更复杂的非线性记忆模块还需要更多验证。第三关于分段策略的敏感性：论文显示对数分段虽然计算效率更高但在召回密集型任务上性能明显下降。第四关于训练与推理的差异：MC可以作为后训练技术应用但效果可能不如端到端训练的版本。最后实验规模相对较小（最大1.3B参数），在更大规模模型上的表现尚待验证。
独立分析的弱点
Memory Caching存在几个值得深入分析的弱点。弱点1是固定分段策略的局限性，当前MC使用固定大小或对数大小的分段但最优分段策略可能因序列内容而异，改进方向是设计自适应分段策略。弱点2是门控机制的表达能力，GRM使用简单的点积相似度计算门控参数，MeanPooling可能丢失段内位置信息，改进方向是使用更复杂的段表示。弱点3是Memory Soup的计算开销，对于深度记忆模块需要为每个token构建新的记忆模块，改进方向是探索参数空间的低秩近似。弱点4是SSC的Top-k选择缺乏梯度，当前SSC使用硬Top-k选择在训练时无法通过梯度下降优化，改进方向是使用Gumbel-Softmax实现可微分的选择。
未来方向
论文作者在结论中指出未来工作可以探索更具表达力的池化或路由机制来进一步提升性能。基于本文成果可以延伸出多个研究方向。方向1是动态记忆管理，借鉴操作系统中的虚拟内存管理思想设计更智能的缓存替换策略。方向2是层次化Memory Caching，在多个时间尺度上应用MC形成层次化的记忆结构。方向3是与高效注意力的结合，探索在不同效率约束下的最优架构配置。方向4是多模态扩展，将MC框架自然地扩展到视觉、音频等多模态序列建模。方向5是理论分析，需要更深入地理解MC如何影响模型的表达能力、泛化能力和记忆容量。
复现评估
论文的复现条件相对友好。论文没有明确提到代码开源但提供了详细的架构细节和实验设置描述，包括模型维度、学习率调度、批量大小等关键超参数。数据集使用了公开可用的FineWeb数据集和Long-Data-Collections，下游评估基准都是标准数据集。算力需求方面760M模型训练30B tokens，1.3B模型训练100B tokens需要相当的计算资源。复现难度中等，核心算法相对简单但需要正确实现不同的RNN基线作为基础架构。建议复现者首先在小规模上验证MC的基本效果然后再扩展到更大规模。
论文图表