Memory Caching: RNNs with Growing Memory

Transformer 已成为近期序列建模领域绝大多数突破性进展的事实标准主干网络，这主要归功于其内存容量随上下文长度增长而线性扩展的特性。尽管这一特性在检索类任务中看似合理，但它会导致计算复杂度呈平方级增长（即 $O(L^2)$），因而促使近期研究致力于探索可行的、具有亚平方级（subquadratic）复杂度的循环替代方案。尽管此类循环架构在多个领域已展现出富有前景的初步成果，但在对召回能力要求较高的任务中，其性能仍普遍落后于 Transformer 模型；这一差距通常被归因于其固定大小的内存容量。本文提出“记忆缓存”（Memory Caching, MC）技术——一种简洁而高效的方法，通过缓存循环模型内部状态（即隐状态）的关键检查点（checkpoints），显著增强其建模能力。记忆缓存使循环神经网络（RNN）的有效内存容量得以随序列长度增长而动态扩展，从而在两类极端情形之间提供灵活的权衡：一端是传统 RNN 的固定内存容量（对应线性时间复杂度 $O(L)$），另一端则是 Transformer 的随上下文长度平方增长的内存开销（对应 $O(L^2)$ 复杂度）。我们共提出了四种 MC 变体，涵盖门控聚合机制与稀疏选择机制，并深入探讨了这些变体对线性记忆模块及深层记忆模块的影响。我们在语言建模和长上下文理解任务上的实验结果表明，MC 能显著提升循环模型的整体性能，充分验证了其有效性。而在上下文内召回（in-context recall）任务上的结果进一步显示：虽然 Transformer 仍保持最高准确率，但我们的各类 MC 变体表现出极具竞争力的性能——不仅大幅缩小了与 Transformer 的性能差距，而且全面优于当前最先进的循环模型。

recurrent模型在长序列建模中受限于固定大小的内存（O(1)或O(L)状态容量），导致其在recall-intensive任务（如in-context learning、长程依赖检索）上显著落后于Transformer；而Transformer虽具备随上下文长度线性增长的隐状态记忆（实际为二次计算开销），但难以扩展至超长序列。本文旨在验证：能否在保持RNN亚二次复杂度（O(L)）的前提下，使recurrent模型的有效记忆容量随序列长度自适应增长，从而弥合与Transformer在回忆密集型任务上的性能鸿沟？该问题属于序列建模基础架构演进中的关键权衡问题——并非全新，但此前缺乏系统性、轻量级且理论清晰的内存扩展机制。

提出Memory Caching（MC）：在RNN前向传播过程中，选择性地缓存历史时间步的隐藏状态（memory checkpoints），并在后续步骤中通过轻量级聚合（如gated fusion或sparse selection）动态读取，使有效记忆容量从O(1)提升至O(L)（随序列长度线性增长），同时维持总体计算复杂度为O(L)（非O(L²)）。其新意在于：不引入额外参数或注意力机制，不改变RNN核心循环结构，仅通过缓存+可学习门控/稀疏路由实现‘软扩展’内存，本质是将传统RNN的‘遗忘即丢弃’范式升级为‘遗忘即归档’，首次将cache作为可训练的记忆增强原语嵌入标准RNN流程。

实验覆盖语言建模（WikiText-103, PG-19）、长上下文理解（LongBench, Needle-in-a-Haystack）、以及in-context recall任务（如multi-hop retrieval）；MC在LSTM/GRU/SSM（如Mamba）等不同recurrent backbone上均取得一致提升，在needle-recall任务中达到Transformer 92–96%的准确率，显著优于FlashAttention-LSTM、RWKV、Hyena等SOTA recurrent模型；所有MC变体均无需修改训练目标，仅需微调；论文开源代码与checkpoint；值得深入的方向包括：cache压缩策略（量化/聚类）、与位置编码协同设计、在强化学习序列决策中的泛化性验证。

‘Transformers are RNNs’ (2023), ‘Linear Transformers Are Secretly Fast RNNs’ (2023), ‘Mamba: Linear-Time Sequence Modeling with Selective State Spaces’ (2023), ‘RWKV: Reinventing RNNs for the Transformer Era’ (2023), ‘StreamingLLM: Efficient Attention with Constant Memory’ (2023), ‘KV Cache Quantization for LLM Inference’ (2024)