News

MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens

原文链接：https://arxiv.org/abs/2603.23516 作者：Yu Chen, Runkai Chen, Sheng Yi, Xinda Zhao, Xiaohong Li, Jianjin Zhang, Jun Sun, Chuanrui Hu, Yunyun Han, Lidong Bing, Yafeng Deng, Tianqiao Chen 机构：Evermind / Shanda Group / Peking University 发布日期：2026-03-27

速查卡

项目	内容
一句话总结	端到端可训练的稀疏注意力框架，从 16K 扩展到 1 亿 token 时精度下降不到 9%
大白话版	给 AI 装了一个”超级记忆库”，能记住相当于人类一辈子的信息量，而且找东西还很快很准
核心数字	<9% 精度下降（16K→100M）、2×A800 GPU 推理 100M token、线性训练/推理复杂度
评级	A — 必读级突破：首次实现端到端可训练的亿级 token 记忆，且精度损失极小
代码	暂未开源（论文刚发布）
关键词	Memory Sparse Attention, 稀疏注意力, Document-wise RoPE, KV Cache 压缩, Memory Interleaving, 终生记忆

核心 Insight

这篇论文的核心洞察是：长期记忆和即时推理是两个不同的能力，应该在架构层面解耦。

当前 LLM 的长上下文方案都试图让模型”一视同仁”地处理所有 token——无论是 1 分钟前的对话还是 1 年前的记忆。MSA 的关键突破在于认识到，对于长期记忆中的大量文档，模型不需要对每个 token 进行全注意力计算，只需要：

高效找到相关文档（稀疏路由）
用压缩的 KV Cache 参与注意力计算
保持端到端可微分以持续优化

这就像人类记忆的工作方式——我们不会同时回忆一辈子的所有经历，而是根据当前问题快速检索相关记忆片段，然后用这些片段辅助思考。

为什么这个想法 work？

三个关键直觉：

文档级稀疏性： 对于一个查询，100 万个文档中可能只有几十个是相关的。传统全注意力浪费了 99.99% 的计算在无关文档上。MSA 通过学习到的路由器（Router）先筛选 Top-k 文档，再做注意力——计算量从 O(N²) 降到接近 O(N)。
Document-wise RoPE 解耦位置编码： 传统的全局位置编码在推理时文档数远超训练时会崩溃（位置 ID 漂移）。MSA 给每个文档独立分配从 0 开始的位置编码，使得训练在 64K 上下文上的模型可以外推到 1 亿 token——因为每个文档”看起来”和训练时一样。
端到端可训练性： 不像 RAG（检索和生成是分开优化的两个系统），MSA 的路由、压缩、注意力全部在同一个梯度流中联合优化，检索质量直接影响生成质量，梯度可以传播回去改进检索。

方法详解

整体架构

输入查询 → [Router Q/K Projector] → 相关性评分 → Top-k 文档选择
                                                    ↓
记忆库文档 → [独立编码 + KV 压缩] → 压缩 KV Cache →[拼接选中文档 KV + 查询 KV]
                                                    ↓
                                          [稀疏注意力计算] → 输出生成

关键技术组件

组件 1: 稀疏注意力机制（Sparse Attention）

做什么： 从海量记忆文档中高效选出最相关的 Top-k 个，只让这些参与注意力计算。

怎么做：

每个文档 d_i 通过主干模型生成标准 K、V 矩阵
同时通过额外的 Router K Projector 生成专用路由键 K^R
对文档做分块均值池化压缩 KV Cache
查询侧同样有 Router Q Projector 生成路由查询 Q^R
计算余弦相似度作为相关性分数
选择 Top-k 文档，将其压缩 KV 与查询的本地 KV 拼接后做注意力

关键公式：

$S_{ij} = \max_{\text{token } t}\left(\underset{\text{head } h}{\text{mean}}\left(\cos\left((Q^R_{q,h})_t, \bar{K}^R_{ij,h}\right)\right)\right)$

其中：

$Q^R_{q,h}$ 是查询的路由向量
$\bar{K}^R_{ij,h}$ 是第 i 个文档第 j 个块的压缩路由键
先跨注意力头取均值，再跨查询 token 取最大值

直觉解释： 这个路由机制就像图书馆的索引系统——每本书（文档）有一个标签（路由键），读者（查询）带着一个搜索词（路由查询），系统快速比对标签找到最相关的几本书，然后读者只需要翻阅这几本。

重要设计选择： MSA 路由仅在模型的后半层启用。实验表明前几层的隐藏状态还没有形成足够的高层语义抽象，路由在这些层效果不好。前半层对每个文档独立处理，不做跨文档检索。

组件 2: Document-wise RoPE（并行与全局位置编码）

做什么： 解决”训练短、推理长”的位置编码外推问题。

怎么做：

记忆文档： 每个文档独立分配位置 ID（都从 0 开始）
活跃上下文（查询+生成）： 使用全局 RoPE，位置 ID 从 k 开始（k = Top-k 检索数）

直觉解释： 想象一个图书馆——每本书内部的页码是独立的（Document-wise），但读者的阅读位置相对于他面前摊开的几本书是连续的（Global）。不管图书馆有 100 本还是 100 万本书，每本书的内部结构不变——这就是外推能力的来源。

数值例子： 假设训练时看 10 个文档（64K token），推理时有 100 万个文档（100M token）。传统全局 RoPE 会给第 100 万个文档分配远超训练范围的位置 ID → 崩溃。MSA 中每个文档仍然是从 0 开始的短序列 → 和训练时一样 → 性能稳定。

组件 3: Memory Interleaving（记忆交错机制）

做什么： 解决跨多个分散文档的多跳推理问题。

怎么做： 当答案需要综合多个文档的信息时，MSA 通过迭代地同步和整合不同记忆段的 KV Cache，保持跨文档的依赖关系。

为什么重要： 标准的稀疏注意力只能看到选中文档内部的关系，但无法捕捉”文档 A 的线索 + 文档 B 的线索 → 答案”这种跨文档推理。Memory Interleaving 通过交错注入不同文档的信息来解决这个问题。

训练策略

数据： 去重语料库 158.95 亿 token 的持续预训练
训练目标： 生成式检索（模型自回归生成相关文档 ID）+ 标准 LLM 损失
辅助损失： 监督每层路由器的检索决策，确保注意力分配到正确的证据文档
训练上下文长度： 仅 64K token（但可外推到 100M）
算力需求： 推理仅需 2×A800 GPU

与现有方法的关键区别

维度	RAG	线性注意力(RWKV/DeltaNet)	MSA
端到端可训练	❌（检索和生成分离）	✅	✅
终生记忆扩展	✅	❌（固定状态大小）	✅
精度保持	中等	低（极长时严重退化）	高（<9% 下降）
灾难性遗忘	低	高	无
计算复杂度	O(L)	O(L)	O(L)
与主流 LLM 兼容	中等	低	高

实验结果

主实验：长文本 QA

论文在 MS MARCO 数据集上评估了不同记忆长度下的 QA 精度：

方法	16K	128K	1M	10M	100M
Qwen2.5-14B-1M	基线	显著下降	严重退化	—	—
Qwen3-30B/80B-A3B	基线	下降	退化	—	—
MemAgent-14B	基线	轻微下降	中等退化	严重退化	—
SOTA RAG	基线	稳定	缓慢下降	中等下降	明显下降
MSA	基线	稳定	稳定	轻微下降	<9% 下降

解读：

传统长上下文模型在超过 1M token 后性能崩溃
即使是 SOTA RAG 系统，在 100M token 规模下也有明显性能损失
MSA 是唯一能在 100M token 下保持接近基线性能的方法

Needle-In-A-Haystack 测试

MSA 在 NIAH 测试中达到 SOTA，展示了对上下文退化的优越鲁棒性——即使”针”被埋在 1 亿 token 的”干草堆”中，MSA 仍能高精度找到。

可扩展性分析

论文的图 1 是最有说服力的数据：从 16K 到 100M token，MSA 的精度曲线几乎是平的，而其他所有方法的曲线都呈现不同程度的下降或崩溃。这种”近乎恒定”的精度-规模关系在此前的任何工作中都未曾实现。

复现评估

维度	评分(1-5)	详细说明
数据可得性	⭐⭐⭐⭐	MS MARCO 等标准数据集公开可用；158.95B token 预训练语料未公开
代码可得性	⭐⭐	论文刚发布，代码暂未开源
算力需求	⭐⭐⭐⭐	推理仅需 2×A800（非常友好）；预训练需要更多资源
工程复杂度	⭐⭐⭐	Router Projector 和 Document-wise RoPE 需要定制实现，但概念清晰
预期收益	⭐⭐⭐⭐⭐	对 Agent 长期记忆、数字孪生等场景有变革性价值

复现建议： 等待官方代码开源。如果需要快速验证概念，可以在现有 Transformer 框架上实现 Document-wise RoPE + Top-k 路由的简化版本。

批判性分析

局限性（论文承认的 + 我们发现的）

论文自述的局限：

目前主要在 QA 任务上评估，尚未在更广泛的下游任务（如摘要、对话、代码生成）上验证
预训练需要大量数据和计算资源

我们额外发现的问题：

文档粒度固定： MSA 以”文档”为最小检索单位——如果一个长文档只有一小段相关，仍然需要检索整个文档。更细粒度的（段落级或句子级）检索可能进一步提升效率。
路由器的可靠性： Top-k 路由是否总能找到正确文档？论文展示了平均性能，但未详细分析路由器的错误模式——漏掉关键文档的概率是多少？
动态记忆更新： 论文聚焦于”读”记忆的场景，对于”写”（新增文档）和”改”（更新已有文档）的效率和一致性讨论有限。

改进方向

多粒度检索： 结合文档级和段落级路由，对长文档做更精细的稀疏选择。预期可以在不增加太多计算的情况下进一步提升精度。难度：中等。
在线记忆更新： 设计增量式的 KV Cache 更新机制，使新文档可以高效添加到记忆库而不需要重新编码所有文档。难度：较高。
与 Agent 系统集成： MSA 的记忆能力与当前的 Agent 框架（如 OpenClaw）结合，可以实现真正具有”终生记忆”的 Agent。难度：工程挑战为主。

独立观察

与 DeepSeek V4 泄露的 Engram Memory 对比： DeepSeek V4 泄露规格中提到”Engram Memory — 条件 O(1) 记忆查找”。MSA 和 Engram 在目标上一致（终生记忆），但技术路径不同——MSA 是注意力层面的稀疏化，Engram 可能是参数级别的记忆嵌入。两种方法的对比将是 2026 年长上下文研究的重要主题。
对 Agent 时代的意义： 36Kr 报道的”Agent 时代算力荒”核心问题之一就是 Agent 的长期记忆导致 token 消耗爆炸。MSA 通过将 100M token 记忆的推理成本降到 2×A800 可承受的水平，为解决这一问题提供了技术路径。
Shanda Group / 盛大集团： 值得注意的是这篇论文来自陈天桥的盛大集团旗下 Evermind AI——盛大从游戏公司转型为 AI 基础研究公司，这篇论文的质量说明转型是认真的。

对领域的影响

短期：MSA 可能成为 Agent 长期记忆的首选技术方案，推动”有记忆的 Agent”从概念走向产品。

中期：如果 Document-wise RoPE 被证明在更广泛任务上有效，可能成为所有长上下文模型的标准组件。

长期：端到端可训练的超长记忆模型可能最终替代 RAG——不再需要外部检索系统，模型自身就拥有高效的”终生记忆”。