跳转至

Sparse Delta Memory: Scaling the State of Linear RNNs through Sparsity

会议: NeurIPS 2026
arXiv: 2607.07386
论文: https://arxiv.org/abs/2607.07386
代码: https://github.com/facebookresearch/sparse-delta-memory
领域: 其他
关键词: 线性循环神经网络, 稀疏增量记忆, 乘积键记忆, 长上下文检索, 线性注意力

一句话总结

提出稀疏增量记忆（Sparse Delta Memory, SDM），利用乘积键稀疏寻址将门控线性 RNN 的隐状态容量扩展数千倍，在严格同 FLOPs 和同参数量下大幅超越 Gated DeltaNet 并媲美甚至超越全注意力模型。

研究背景与动机

随着前沿大语言模型向智能体（Agent）和长程任务演进，代码工程、复杂研究推理以及长视频建模等应用要求架构具备维持跨长时程交互的能力。标准 Transformer 依赖 Softmax 全注意力，将历史交互信息保存在 KV 缓存（KV Cache）中，虽然保证了极强的长程检索能力，但其推理时的内存与计算开销随序列长度线性膨胀，导致超长序列下的上下文学习变得计算不可承受。线性 RNN（如 Mamba2、Gated DeltaNet 等）通过固定大小的循环状态实现了每 token 恒定的计算量与内存占用，但其压缩状态容量极度受限，在长上下文检索与上下文学习（In-Context Learning）任务中显著落后于 Softmax 注意力。

增大线性 RNN 的隐藏状态是提升召回能力的理论正道，然而现代门控线性 RNN 的状态更新均为稠密外积操作，计算复杂度直接正比于状态尺寸，导致状态扩张伴随着无法承受的 FLOPs 暴涨。这种状态容量与计算预算之间的强耦合，构成了线性序列模型扩展的根本瓶颈。本文切入这一核心矛盾，提出无需全量更新庞大状态矩阵，而是借助稀疏寻址将单步读写收敛到固定数量的活跃槽位中。核心 idea：将门控增量网络（Gated DeltaNet）的稠密状态更新替换为基于乘积键寻址的稀疏读写机制，在保持单 token 计算 FLOPs 和参数量严格不变的前提下，将隐状态容量提升三个数量级，并结合可学习的初始状态赋予模型强大的参数化先验知识。

方法详解

整体框架

Sparse Delta Memory（SDM）旨在解决线性 RNN 中状态尺寸与计算复杂度绑定的问题。模型整体采用滑动窗口注意力（SWA，窗口大小为 128）与全局 SDM 层以 3:1 比例交替堆叠的混合架构。对于全局层中的每个 SDM 头部，系统维护一个具有 \(N\) 个槽位、特征维度为 \(d_v\) 的显式大记忆表 \(\mathbf{M}_t \in \mathbb{R}^{N \times d_v}\)。在每个时间步，输入特征被投影为寻址向量，经由低复杂度的乘积键索引快速选出极少数的活跃槽位，仅对被选中的槽位执行门控增量写入并进行稀疏加权读取，其余槽位保持恒等，从而在常数级计算开销下完成对超大容量状态的精确演化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 token 特征 xt"] --> B["乘积键稀疏寻址<br/>投影并计算外和，选出 W 个写槽位与 R 个读槽位"]
    B --> C["门控增量写入<br/>对选中槽位应用衰减因子与输入门更新差值"]
    C --> D["稀疏加权读取<br/>聚合 R 个选中槽位特征并输出 yt"]
    D --> E["RMS-Norm 归一化与输出门控"]
    E --> F["多头投影混合输出 ot"]

关键设计

1. 乘积键稀疏寻址：以次线性开销在大容量记忆表中精确定位槽位

为摆脱稠密状态下单步计算随总槽位数 \(N\) 线性扩张的瓶颈，SDM 引入了类似 Product-Key Memory（PKM）的双半向量外和机制。输入向量 \(\mathbf{x}_t\) 首先经过投影矩阵生成寻址向量，并被对半拆分为两个长度均为 \(\sqrt{N}\) 的子向量 \(\mathbf{k}'_{1,t}, \mathbf{k}'_{2,t}\)（查询侧同样拆分为 \(\mathbf{q}'_{1,t}, \mathbf{q}'_{2,t}\)）。通过计算外和矩阵 \(\mathbf{k}'_{1,t} \oplus \mathbf{k}'_{2,t} \in \mathbb{R}^{\sqrt{N} \times \sqrt{N}}\)，隐式产生了对应 \(N\) 个槽位的全部相关度得分。得益于排序外和的代数性质，模型无需显式实例化完整的 \(N\) 个得分，仅需 \(O(\sqrt{N} \times d + W^2 + R^2)\) 的复杂度即可提取出得分最高的 \(W\) 个写入索引 \(\mathcal{I}^w_t\) 与 \(R\) 个读取索引 \(\mathcal{I}^r_t\)。在实际配置中设定 \(W=R=64 \ll N\)，从而让显式记忆规模能够扩展到百万级槽位而寻址耗时保持在微秒级。

2. 稀疏门控增量读写：局部更新下的干扰消除与常数级 FLOPs 约束

在获得稀疏索引后，SDM 并不对整个记忆表执行退化，而是仅针对选中的 \(W\) 个写入槽位 \(i \in \mathcal{I}^w_t\) 施加门控增量规则。该机制延续了 Delta 规则的核心思想，即在写入新特征前先减去槽位中已有的旧内容，防止特征范数发散与长程混叠；同时引入逐 token 的遗忘门 \(\alpha_t \in (0, 1)\) 与输入门 \(\beta_t \in [0, 1]\)：

\[\mathbf{\tilde{M}}_{t}[i] \leftarrow \alpha_{t} \cdot \mathbf{M}_{t-1}[i]\]
\[\mathbf{M}_{t}[i] \leftarrow \mathbf{\tilde{M}}_{t}[i] + \beta_{t} \cdot k_{t}^{(i)} \cdot \left(\mathbf{v}_{t} - \mathbf{\tilde{M}}_{t}[i]\right)\]

未被选中的绝大多数槽位直接沿用前一时刻状态 \(\mathbf{M}_t[i] = \mathbf{M}_{t-1}[i]\)。对于读取操作，输出仅由选中的 \(R\) 个槽位进行稀疏线性加权得到 \(\mathbf{y}_t = \sum_{i \in \mathcal{I}^r_t} q_{t}^{(i)} \mathbf{M}_t[i]\)。通过将 \(W\) 和 \(R\) 严格约束为基线 Gated DeltaNet 的特征维度（例如 64），整层的运算量完全等价于稠密 GDN，但状态容量却因解耦了 \(N\) 而实现成千上万倍的剧增。

3. 可学习初始状态：将预训练参数记忆与上下文动态记忆统一

在传统的线性 RNN 或 Softmax Transformer 中，循环状态或 KV 缓存通常以全零状态启动，其内容完全来源于上下文推理过程。然而 SDM 拥有高达数亿乃至数十亿数值的记忆槽位，如果仅仅将其置零，不仅浪费了极高的参数表达空间，也无法在序列初始阶段发挥记忆优势。本文将初始状态 \(\mathbf{M}_0\) 作为模型预训练的可学习参数直接参与端到端反向传播。这一设计不增加任何前向推理 FLOPs，却让记忆表充当了类似于前馈网络（FFN）参数的非结构化知识库，使模型在极短上下文时即可直接索引预训练掌握的事实知识，而在序列延伸时又能动态覆写更新。

损失函数 / 训练策略

模型采用标准自回归下一词预测交叉熵损失。为了在训练长序列时高效实现批次内并行计算，SDM 引入了分块（Chunk-wise）WY 并行表示法，将序列划分为大小为 \(C=128\) 的分块。在块内通过基于双指针合并算法的稀疏内积 CUDA 算子快速求解因果依赖三角系统；在块间沿时间步传递记忆状态。为解决超大记忆表反向传播时的显存占用难题，模型采用原地更新（In-place update）与逆向撤销策略，在反向时重放并逆转前向写入，使得每层的记忆检查点显存占用从 \(O((T/C) \times N \times d)\) 剧降到常数级（在 8B 模型、16k 长度训练下显存从 226 GB 降至 8 GB）。预训练遵循 Warmup-Stable-Decay（WSD）学习率调度，并在 1.4B 与 8B 规模下在 128k 上下文上执行两阶段微调。

实验关键数据

主实验

主实验在 1.4B 与 8B 参数规模下，与同为混合架构（3:1 SWA）的 Full Attention、Mamba2 和 Gated DeltaNet（GDN）展开严密对比。评估涵盖文本建模困惑度（Validation NLL）、常识与多任务推理（15 项评测集平均准确率），以及长上下文多针检索评测集 RULER（4k 至 131k 长度的平均严格匹配准确率）。

模型规模 全局层架构 验证集 NLL ↓ 常识/推理均分 (%) ↑ RULER 检索均分 (%) ↑ 相对 GDN RULER 提升
1.4B FullAttn (全注意力) 2.660 37.98 32.5 +12.5
1.4B Mamba2 2.686 38.00 17.9 -2.1
1.4B GDN (Gated DeltaNet) 2.683 37.93 20.0 基线
1.4B SDM (本文) 2.654 38.51 31.2 +11.2
8B FullAttn (全注意力) 2.285 56.65 61.2 +27.0
8B GDN (Gated DeltaNet) 2.298 55.70 34.2 基线
8B SDM (本文) 2.253 56.84 50.2 +16.0

在短文本和通用知识评测中，SDM 取得了所有架构中的最低 NLL；在 8B 规模下，其在 RULER 长程检索任务上相比同 FLOPs 的 GDN 绝对跃升 16.0%（从 34.2% 飙升至 50.2%），并大幅缩减了线性 RNN 与无界 KV Cache 全注意力模型之间的鸿沟。

消融实验

消融实验深入验证了记忆表状态容量与可学习初始状态 \(\mathbf{M}_0\) 对长上下文检索与建模能力的独立贡献（基于 1.4B 和 0.8B 规模）。

实验组别 模型架构 初始状态 \(\mathbf{M}_0\) 每层状态容量 代码验证 NLL ↓ 常识准确率 (%) ↑ RULER 召回 (%) ↑
初始状态消融 (1.4B) GDN 全零 (null) 0.2 MB 0.849 37.9 20.0
初始状态消融 (1.4B) GDN 可学习 (learned) 0.5 MB 0.850 38.0 20.6
初始状态消融 (1.4B) SDM 全零 (null) 211 MB 0.845 37.3 28.0
初始状态消融 (1.4B) SDM 可学习 (learned) 211 MB 0.822 38.5 31.2
状态容量消融 (0.8B) GDN 全零 (null) 0.2 MB 0.963 33.8 16.0
状态容量消融 (0.8B) SDM 可学习 (learned) 27 MB 0.947 33.7 20.2
状态容量消融 (0.8B) SDM 可学习 (learned) 108 MB 0.937 33.6 20.7
状态容量消融 (0.8B) SDM 可学习 (learned) 432 MB 0.914 34.6 21.5

关键发现

  • 状态容量是长程检索的核心使能因子：当 SDM 采用全零初始状态时，其 RULER 召回率依然达到 28.0%，远超 GDN 的 20.0%，证明即使没有预训练知识注入，纯粹依靠稀疏化带来的大状态容量便能从根本上解除线性 RNN 的信息瓶颈。
  • 小状态模型无法受益于可学习初始化：GDN 引入可学习 \(\mathbf{M}_0\) 后性能基本持平（准确率仅从 37.9% 微调至 38.0%），而 SDM 依托数百兆字节的参数空间，使得代码 NLL 大幅下降至 0.822，证明只有大容量记忆才具备承载预训练知识先验的条件。
  • 状态尺寸呈现平滑幂律增益：将 SDM 状态从 27 MB 扩展至 432 MB，代码建模 NLL 呈现 0.947 → 0.937 → 0.914 的单调下降，性能扩展具备高度可预测性。

亮点与洞察

  • 将稀疏寻址嫁接于增量更新：首次将 PKM 的双子向量外和高效索引算法迁移到循环状态的时间步递推中，成功打破了长期以来线性注意力模型中“状态扩容必伴随矩阵乘 FLOPs 暴涨”的固有僵局。
  • 兼顾常数级推理吞吐与极低衰减：在 100 万 token 的极端上下文外推评测中，随着输入位置延伸，SDM 的困惑度依然持续下降，彻底克服了固定状态循环神经网络在跨越数十万 token 后的灾难性容量饱和问题。
  • 反向重放与就地撤销策略：巧妙利用增量更新的代数可逆性，在反向传播时实时逆转稀疏更新，使万倍容量状态的模型训练显存开销收敛在合理范围内。

局限与展望

  • 显存与带宽受限：虽然计算 FLOPs 保持恒定，但数 GB 的超大隐藏状态必须驻留在 GPU 高带宽内存（HBM）而非片上高速 SRAM 中，导致受限于访存带宽，目前 8B 规模端到端训练吞吐量比极致优化的 GDN Triton 算子慢约 1.49 倍。
  • 多键多值极端检索仍有差距：在 RULER 的超高复杂度多键追踪（如 multikey_2）任务中，SDM 尽管超越同类线性 RNN，但仍落后于全量保留所有 token 的无界 KV Cache 全注意力模型。
  • 硬件专属稀疏核函数仍待深度打磨：当前实现虽验证了可行性，未来若能协同定制化片上异步访存与更高精度的低比特量化（如 FP8/INT4 状态），推理端吞吐潜力将得到更充分的释放。

相关工作与启发

  • vs Gated DeltaNet (GDN)：GDN 使用全量稠密外积更新 \(d_{qk} \times d_v\) 的小状态矩阵，状态容量通常不足 1 MB；SDM 将其改造为稀疏寻址大表，在单步计算 FLOPs 严格一致的条件下实现状态数千倍的跃迁。
  • vs Fast-Weight PKM (Zhao & Jones, 2026)：后者探索了 Test-Time-Training 结合稀疏 PKM，但未在等 FLOPs 条件下严格对齐，且仅局限于 100M 参数级的小模型；本文建立了严密的等 FLOPs 架构体系并成功外推扩展至 8B 模型与万亿级 token。
  • vs Softmax Full Attention：全注意力通过无损 KV Cache 实现了无上限的记忆召回，但带来线性的推理开销；SDM 证明通过常数级的稀疏读写显式记忆，可在保留线性推理速度的同时逼近全注意力的检索极限。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将高效乘积键稀疏化扩展成功引入门控增量线性注意力，破除了线性 RNN 的容量瓶颈。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 8B 规模万亿 token 预训练、128k 上下文微调、严格 isoFLOP 对比、百万 token 外推及深入消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 架构形式化推导清晰严密，从算法动机、数学机制到工程并行实现叙述极其完整。
  • 价值: ⭐⭐⭐⭐⭐ 为长上下文模型与无衰减循环记忆架构提供了极具实用潜力的新范式，推动了线性模型理论边界。

相关论文