[论文] 记忆缓存：RNNs with 增长记忆

Source: arXiv - 2602.24281v1

（请提供需要翻译的正文内容，我将为您翻译成简体中文。）

概述

论文 “Memory Caching: RNNs with Growing Memory” 提出了一种轻量级的附加组件，使循环神经网络（RNN）能够在序列变长时扩展其有效记忆。通过检查点缓存隐藏状态，作者弥合了经典 RNN 的线性时间、固定大小记忆与 Transformer 的二次时间、不断增长记忆之间的差距——提供了一种可调的权衡，能够在当今硬件上部署。

关键贡献

• Memory Caching (MC) technique：一种简单的机制，用于存储和复用过去的隐藏状态检查点，能够在不改变 RNN 核心递归结构的前提下有效扩大其记忆容量。
• 四种 MC 变体：
  1. Plain caching – 对每个隐藏状态进行朴素存储。
  2. Gated aggregation – 学习对缓存状态的加权融合。
  3. Sparse selective caching – 根据学习得到的重要性分数，仅保留一部分检查点。
  4. Hybrid deep‑memory caching – 将 MC 与多层（深层）记忆模块相结合。
• 复杂度插值：MC 可以配置为在 (O(L))（类似 RNN）到 (O(L^2))（类似 Transformer）的时间复杂度之间运行，让实践者能够在延迟与准确性之间找到最佳平衡点。
• 实证验证：在语言建模基准（如 WikiText‑103）和长上下文推理任务上展示出持续的性能提升，缩小了与 Transformer 的性能差距，同时保持更低的成本。
• 开源实现：作者发布了代码和预训练检查点，开发者可以轻松将 MC 插入现有的 RNN 流程中。

方法论

1. Baseline RNN – 作者从标准循环架构（例如 LSTM 或 GRU）开始，该架构处理 token 序列 ({x_t}_{t=1}^L) 并产生隐藏状态 (h_t)。
2. Checkpointing – 在可配置的间隔（或当学习到的“重要性”信号突升时），将当前隐藏状态保存到缓存 (C = {c_1, …, c_K})。
3. Memory read‑out – 当 RNN 需要在步骤 (t) 产生输出时，它会查询缓存。
   • Plain MC 简单地将所有缓存状态拼接或求平均。
   • Gated aggregation 学习门 (g_k = \sigma(W_g c_k + b_g))，并计算 (\tilde{h}_t = \sum_k g_k c_k)。
   • Sparse selective MC 对评分函数 (s_k = f(c_k)) 进行 top‑k 选择，只保留最相关的检查点。
4. Integration – 将检索到的记忆 (\tilde{h}_t) 与当前隐藏状态合并（例如通过相加或小型前馈网络），再送入最终输出层。
5. Training – 整个系统保持端到端可微；缓存操作使用高效的张量索引实现，因此训练开销保持适中。

结果与发现

• 性能提升：所有 MC 变体都提升了困惑度和下游任务得分，门控聚合带来最强的提升。
• 效率：即使是完整缓存版本的成本也远低于 Transformer 的二次方成本，稀疏版本可调优至几乎与普通 RNN 同速。
• 内存‑准确率权衡：通过调整缓存大小或稀疏度，开发者可以实现所需的平衡——例如，10% 缓存可获得约 90% 的完整缓存增益，额外计算量不足 20%。

实际意义

• 可部署在边缘/低功耗设备上：MC 让你在保持 RNN 轻量递归特性的同时，处理更长的上下文（例如聊天记录、流式日志），而不会导致内存或延迟激增。
• 即插即用升级：现有的 LSTM/GRU 代码库只需几行包装代码即可采用 MC；无需重写整个模型或切换到 Transformer 结构。
• 成本效益的扩展：对于处理海量文本流的 SaaS 平台，MC 提供了折中方案——比普通 RNN 有更好的召回率，又比运行全尺度 Transformer 更便宜。
• 混合架构的潜力：MC 可以与近期的线性注意力 Transformer 结合，产生“记忆增强”混合模型，进一步突破上下文长度的限制。
• 研究复用：开源的缓存模块可以作为其他序列密集领域的构建块，例如 DNA 序列分析、时间序列预测，或需要长期状态的强化学习代理。

限制与未来工作

• 缓存管理开销：虽然作者已经将其保持在较低水平，但对于非常长的序列（数十万步）仍需仔细调节缓存大小和驱逐策略，以避免 GPU 内存突增。
• 任务特定调优：最佳稀疏度水平或门控架构在不同领域之间会有所差异；尚未找到一种“一刀切”的设置。
• 比较范围：实验主要聚焦于语言建模和召回任务；更广泛的基准（例如多模态视频字幕、代码生成）仍未涉及。
• 作者提出的未来方向包括：
  1. 学习动态缓存更新调度，
  2. 将 MC 与检索增强模型结合，
  3. 将该技术扩展到非 RNN 的递归结构，如 Neural ODE 或状态空间模型。

作者

• Ali Behrouz
• Zeman Li
• Yuan Deng
• Peilin Zhong
• Meisam Razaviyayn
• Vahab Mirrokni

论文信息

• arXiv ID: 2602.24281v1
• 分类: cs.LG, cs.AI
• 出版时间: 2026年2月27日
• PDF: 下载 PDF