MSA：面向1亿token终身记忆的端到端可训练稀疏注意力架构 MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens
稀疏注意力+文档级RoPE，端到端训练实现1亿token线性扩展
前置知识
注意力机制 (Attention Mechanism)
Transformer的核心组件，通过Query-Key-Value计算实现序列内元素间的信息交互。标准自注意力的计算复杂度为 $O(L^2)$，其中 $L$ 为序列长度。每个token通过Q与所有K的点积计算注意力权重，再与V加权求和得到输出。这种全局交互能力赋予了模型强大的表达能力，但也成为扩展到超长序列时的计算瓶颈。
MSA的核心创新就是用稀疏注意力替代标准的稠密注意力，理解注意力机制的QKV计算流程和复杂度瓶颈是理解本文的基础。
KV缓存 (Key-Value Cache)
在自回归生成过程中，已处理token的Key和Value向量会被缓存以避免重复计算。对于长度为 $L$ 的上下文，KV缓存存储所有历史token的K和V矩阵，占用内存与序列长度和层数成正比。当上下文扩展到百万甚至亿级token时，KV缓存的内存占用会成为实际部署的主要障碍。本文提出的压缩KV缓存策略正是针对这一痛点。
MSA通过chunk级均值池化压缩KV缓存，使得100M token的推理在2块A800 GPU上成为可能，理解KV缓存的存储开销是理解这一技术贡献的前提。
旋转位置编码 (RoPE)
Rotary Position Embedding是一种将位置信息编码到Query和Key向量中的方法。它通过旋转矩阵将绝对位置信息嵌入，使得点积自然反映相对位置关系。标准RoPE对整个序列使用连续递增的位置ID，当推理时序列长度超过训练长度时会出现分布外(OOD)问题，导致性能退化。MSA提出的文档级RoPE正是为解决这一外推问题。
MSA采用文档独立RoPE+全局RoPE的混合策略，是实现从64K训练到100M推理外推的关键技术，不理解RoPE就无法理解这一设计。
检索增强生成 (RAG)
Retrieval-Augmented Generation是一种通过外部知识库增强LLM能力的范式。它首先从大规模文档库中检索与查询相关的片段，然后将检索结果拼接到上下文中供模型生成回答。标准RAG流程包括：文本分块、向量编码、相似度检索、（可选重排序）、拼接生成。这种方法虽然可以扩展到大规模知识库，但检索和生成是解耦的，检索质量成为性能天花板。
RAG是MSA的主要对标方法。论文通过大量实验对比MSA与SOTA RAG系统，证明端到端可训练的稀疏注意力在精度和可扩展性上的优势。
Top-k选择 (Top-k Selection)
一种从候选集中选择得分最高的k个元素的操作。在MSA中，模型需要从所有文档中选择最相关的top-k个文档参与注意力计算。关键挑战在于标准的top-k操作不可微分，无法直接用于端到端训练。MSA通过将top-k集成到稀疏注意力中，并设计辅助损失函数来实现可微分的路由训练。
MSA的稀疏注意力机制的核心就是路由+top-k选择，理解这一操作及其可微性挑战是理解本文技术路线的关键。
连续预训练 (Continual Pre-training)
在已有预训练模型的基础上，使用特定领域或任务的数据继续训练。与从头预训练相比，CPT可以复用模型已有的知识和能力，同时注入新的技能。本文中MSA在Qwen3-4B-Instruct-2507基础上进行CPT，使用158.95B token的语料训练模型的检索能力，包括一个特殊的warmup阶段专门对齐路由器。
CPT是MSA训练流程的第一阶段，也是赋予模型检索能力的基础。实验显示去掉CPT会导致平均性能下降31.3%，证明其不可或缺。
研究动机
当前大语言模型的有效上下文长度通常限制在128K到1M token之间，远未达到人类终身记忆的规模。认知科学估计人类记忆的功能信息容量约为 $10^9$ 比特，假设每个token的有效语义密度为3-5比特，这相当于约2-3亿token的终身容量。然而，现有三类记忆方法都存在根本性缺陷：参数记忆（如LoRA、持续预训练）虽然精度高，但无法扩展且容易灾难性遗忘；外部存储记忆（如RAG、MemAgent）虽然可以扩展，但依赖离散语义表示，检索与生成解耦，精度有天花板；潜状态记忆中的KV缓存方法（如DSA）精度高但计算成本过高无法扩展到1亿token，线性注意力方法（如RWKV、DeltaNet）复杂度低但固定大小的状态必然导致信息丢失。以Qwen2.5-14B-1M为例，在256K token时准确率就从97%骤降至90%，在1M时进一步降至53%；而Qwen3-4B-Instruct在256K时准确率仅为48%，在1M时更是崩溃到25%。即使是最先进的Qwen3-Next-80B-A3B，在256K后也会出现显著退化。这些数据表明，现有方法无法同时满足高精度、可扩展、端到端可训练这三大需求。
本文的目标是本文的目标是设计一种端到端可训练、高效且大规模可扩展的记忆模型框架，使LLM能够处理终身级别的信息量——即数亿token的上下文，同时保持可接受的精度退化。具体而言，MSA希望实现以下几个关键目标：第一，实现训练和推理的近线性复杂度 $O(L)$，使得100M token的推理在2块A800 GPU上成为可能；第二，从16K到100M token的扩展过程中，精度退化控制在9%以内；第三，通过端到端训练实现检索和生成的联合优化，消除传统RAG中检索质量的性能天花板；第四，支持复杂多跳推理，能够跨文档片段整合证据链。
与已有工作不同的是，现有方法在可扩展高精度记忆和端到端可训练性两个维度上都存在根本性不足。在第一个维度上，提供强精度的方法受限于固定的上下文或状态容量，而容量可扩展的方法难以保证可靠性。在第二个维度上，没有现有范式提供完全可微分、联合优化的记忆流水线，同时保持架构兼容性、高精度和对灾难性遗忘的鲁棒性。MSA的独特切入角度在于：将top-k选择与稀疏注意力结合，创造一种既可微分又能实现稀疏检索的机制。通过文档级RoPE解耦位置语义与文档数量，使得模型可以在小上下文上训练（64K）却能在超大上下文上推理（100M）。这种「短训练、长推理」的设计理念，以及将KV缓存压缩与分层存储相结合的工程实现，构成了MSA区别于所有现有方法的核心差异化。
核心方法
MSA的核心直觉是：与其在推理时处理全部历史上下文（全注意力的 $O(L^2)$），不如先识别出最相关的少数文档，只对这些文档进行精细注意力计算。技术路线分为三个层次：首先是架构层面，MSA用基于文档的稀疏注意力替代标准稠密注意力，通过路由器（Router Projector）计算查询与文档的相关性分数，选择top-k个最相关文档参与注意力计算；其次是位置编码层面，采用文档独立RoPE为每个文档分配从0开始的位置ID，加上对活跃上下文使用全局RoPE，实现从短上下文训练到长上下文推理的外推；最后是训练层面，通过两阶段连续预训练（warmup对齐路由器 + 主训练优化生成）和两阶段后训练（8K基础SFT + 64K课程学习），逐步建立模型的检索和推理能力。整个框架在Qwen3-4B-Instruct-2507骨干上实现，保持与主流LLM的架构兼容性。
MSA与已有方法的本质区别在于将稀疏检索内化到注意力机制中，实现端到端可微分。传统RAG将检索和生成解耦：检索器使用模型无关的语义嵌入，生成器使用模型内部表示，两者之间存在优化间隙。MSA则通过引入专门的Router Projector（参数矩阵 $W^{KR}$ 和 $W^{QR}$），在模型的隐藏空间中直接计算查询-文档相关性。具体而言，对于每个注意力头 $h$，路由器生成专门的路由键 $K^R_{i,h} = H_i W^{KR}_h$ 和路由查询 $Q^R_{q,h} = H_q W^{QR}_h$，然后通过余弦相似度计算相关性分数。这种设计的关键创新在于：路由器的参数通过端到端训练与生成目标联合优化，辅助损失 $\mathcal{L}_{aux}$ 直接监督路由决策的正确性。同时，top-k选择虽然本身不可微，但通过在训练中引入对比学习目标，路由器学会了产生正确的选择。这使得MSA既保持了稀疏注意力的计算效率，又获得了端到端优化的精度优势。
方法步骤详情
MSA的方法包含训练和推理两个阶段的完整流程。训练阶段：（1）连续预训练——在158.95B token语料上进行，第一阶段warmup使用损失 $\mathcal{L} = 0.1\mathcal{L}_{LLM} + \mathcal{L}_{aux}$，学习率 $1 \times 10^{-4}$，重点对齐路由器投影器；第二阶段主训练切换为 $\mathcal{L} = \mathcal{L}_{LLM} + 0.1\mathcal{L}_{aux}$，学习率降至 $6 \times 10^{-6}$，优化生成能力。（2）后训练——第一阶段在8K上下文上进行SFT，建立基础指令跟随能力；第二阶段扩展到64K上下文，进行数据清洗和课程学习。推理阶段分为三个Stage：Stage 1离线编码——对整个文档库进行前向传播，生成K、V、KR矩阵，通过chunk级均值池化（chunk大小 $P=64$）压缩为 $\bar{K}$、$\bar{V}$、$\bar{K}^R$ 并缓存；Stage 2在线路由——计算查询的路由向量 $Q^R_q$，与缓存的路由键计算余弦相似度，在token维度取均值、文档维度取最大值后，通过top-k（$k=16$）选择最相关文档；Stage 3稀疏生成——将选中文档的压缩KV与查询的本地KV拼接，执行自回归生成。对于多跳推理，MSA引入Memory Interleave机制：模型先生成文档ID，加载原始文本扩展上下文，再进行下一轮检索，迭代直到信息充分后生成最终答案。
技术新颖性
MSA的技术新颖性体现在四个层面。第一，稀疏注意力与top-k的可微分集成——这是首次在注意力机制内部实现端到端可训练的稀疏检索。通过Router Projector在模型隐藏空间中进行路由决策，并用辅助对比损失 $\mathcal{L}_{aux}$ 监督，绕过了top-k操作的不可微性问题。第二，文档级RoPE的外推策略——通过为每个文档分配独立位置ID（从0开始），解耦了位置语义与文档总数，使得在64K上下文训练的模型可以在100M上下文上推理，位置ID不再随文档数量单调递增。第三，分层存储的Memory Parallel策略——将路由键（$\bar{K}^R$）常驻GPU显存用于快速检索，内容KV（$\bar{K}$、$\bar{V}$）卸载到CPU内存，仅在top-k选择后异步加载。这种分层设计使得100M token的KV缓存（约169GB）在2xA800 GPU（共160GB显存）上成为可能。第四，自适应Memory Interleave机制——模型在推理时自适应地决定需要多少轮检索，每轮检索的文档数量不固定，直到累积的证据足够回答问题。这种迭代检索-扩展-生成的循环有效支持了跨文档的多跳推理。
实验结果
MSA在多个实验维度上展现了卓越的性能。在QA任务上，MSA（4B参数）在9个基准测试中平均LLM Judge得分3.760，相比同骨干标准RAG（3.179）提升16.0%，相比带重排序的RAG（3.355）提升11.5%，相比HippoRAG2（3.227）提升14.8%。更令人印象深刻的是，MSA甚至超越了使用KaLMv2-Embedding（12B）+ Qwen3-235B的SOTA RAG系统（3.580），平均提升5.0%。在NIAH测试中，MSA从32K到1M token的准确率从98.77%仅降至94.84%（3.93个百分点的下降），而Qwen3-4B-Instruct从95%崩溃到25%，Qwen3-Next-80B-A3B从100%降至81%。在MS MARCO的可扩展性测试中，MSA从16K到100M token仅退化8.8%（4.023到3.669），而GPT-4.1和DeepSeek-V3.2虽然起始分数相近（约3.6-3.7），但随着上下文增长出现明显下降。消融实验进一步验证了各组件的重要性：去掉连续预训练导致平均性能下降31.3%（3.694到2.537），去掉原始文本注入导致37.1%的下降（3.694到2.325），去掉Memory Interleave导致5.3%下降（3.694到3.497），而从S1到S2的课程学习带来7.6%的提升（3.694到3.976）。在MS MARCO这种7.34M token的大规模语料上，课程学习的增益尤为显著，达到29.5%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MS MARCO v1 (7.34M tokens) | LLM Judge (0-5) | 4.141 | RAG: 3.011 / RAG+RR: 3.032 / SOTA(Qwen3-235B): 3.028 | vs同骨干最优: +36.9% / vs SOTA RAG最优: +36.8% |
| Natural Questions (1.47M tokens) | LLM Judge (0-5) | 3.545 | RAG: 3.452 / RAG+RR: 3.494 / SOTA(Llama3.3): 3.756 | vs同骨干最优: +1.5% / vs SOTA RAG最优: -5.6% |
| DuReader (277K tokens) | LLM Judge (0-5) | 4.155 | RAG: 3.726 / RAG+RR: 3.848 / SOTA(Qwen3-235B): 4.044 | vs同骨干最优: +8.0% / vs SOTA RAG最优: +2.7% |
| TriviaQA 10M (10M tokens) | LLM Judge (0-5) | 4.621 | RAG: 4.414 / RAG+RR: 4.391 / SOTA(Llama3.3): 4.740 | vs同骨干最优: +4.7% / vs SOTA RAG最优: -2.5% |
| NarrativeQA (538K tokens) | LLM Judge (0-5) | 3.395 | RAG: 2.860 / RAG+RR: 3.638 / SOTA(Qwen3-235B RR): 3.427 | vs同骨干最优: -6.7% / vs SOTA RAG最优: -0.9% |
| 2WikiMultiHopQA (722K tokens) | LLM Judge (0-5) | 4.280 | RAG: 3.136 / RAG+RR: 3.159 / SOTA(Qwen3-235B): 3.582 | vs同骨干最优: +35.5% / vs SOTA RAG最优: +19.5% |
| HotpotQA (1.35M tokens) | LLM Judge (0-5) | 4.061 | RAG: 3.787 / RAG+RR: 4.022 / SOTA(Qwen3-235B): 4.225 | vs同骨干最优: +1.0% / vs SOTA RAG最优: -3.9% |
| MuSiQue (1.41M tokens) | LLM Judge (0-5) | 2.211 | RAG: 1.928 / RAG+RR: 1.965 / SOTA(Qwen3-235B): 2.647 | vs同骨干最优: +12.6% / vs SOTA RAG最优: -16.5% |
| RULER NIAH (1M tokens) | Accuracy | 0.9484 | RL-MemoryAgent-14B: 0.9266 / Qwen3-Next-80B-A3B: 0.8078 | vs MemoryAgent: +2.4% / vs 80B模型: +17.4% |
局限与改进
论文在第7节明确承认了MSA当前的局限性：当任务需要对多个文档之间建立强耦合依赖时，MSA仍然存在不足。在证据分布广泛且高度关联的场景中，该方法难以仅通过内在记忆维持准确的结构对齐。Memory Interleave虽然是缓解这些问题的有潜力方向，但其有效性依赖于更高效和原则性的设计来更好地保留文档间关系。从独立分析来看，MSA还存在以下局限：第一，离线编码阶段的计算成本 $O(LG)$ 虽然可以摊销，但对于频繁更新的知识库，重新编码整个语料的开销仍然可观；第二，chunk级均值池化压缩不可避免地损失细粒度信息，在需要精确定位到句子级别的细粒度检索任务上可能表现不佳；第三，top-k=16的固定选择策略可能不适合所有场景——当答案需要综合更多文档时（如需要50个文档的综述任务），固定k值可能成为瓶颈；第四，MuSiQue上与235B模型16.5%的差距表明，在需要复杂推理的场景中，模型容量仍然是决定性因素，4B的骨干模型存在推理能力的天花板；第五，论文缺乏对Memory Interleave具体迭代次数和计算开销的详细分析。
独立分析的弱点
MSA存在几个值得深入分析的弱点。首先，分层存储策略虽然巧妙地解决了内存限制，但CPU-GPU之间的异步数据传输在实际高并发服务场景中可能成为瓶颈——当多个查询同时请求不同的top-k文档时，PCIe带宽可能不足以支撑。改进方向可以考虑使用NVMe SSD作为中间层，或者设计预测性预加载策略。其次，文档级RoPE虽然实现了外推，但每个文档内部从位置0开始编码可能丢失文档间的时序关系信息——在时间敏感的问答场景（如「A事件发生在B之前还是之后」）中可能产生问题。可以探索在文档间引入微小的位置偏移来保留时序信号。第三，辅助损失中的温度参数是固定的，但不同层的路由器可能需要不同的粒度——浅层可能需要更粗粒度的选择，深层需要更精细的区分。第四，当前实现仅在模型的后半层启用MSA路由，前半层依赖标准注意力，这意味着检索能力仅在高层语义表示中建立，可能限制了模型在需要底层特征匹配的任务上的表现。
未来方向
论文提出的Memory Interleave机制是一个极具潜力的研究方向，未来可以探索更高效的设计来更好地保留文档间关系。基于MSA的成果，有多个可延伸的研究方向：第一，动态k选择——让模型根据查询复杂度自适应地调整top-k值，简单查询用更少的文档，复杂查询用更多，而不是固定k=16；第二，层级化记忆——借鉴人类记忆的工作记忆和长期记忆分层结构，对文档库进行多粒度索引，实现从摘要到细节的渐进式检索；第三，流式记忆更新——设计增量编码机制，使得新文档可以在线加入记忆库而无需全量重新编码，这对于实时信息更新场景至关重要；第四，多模态扩展——将MSA框架扩展到图像、视频等多模态记忆，利用视觉-语言对齐实现跨模态的终身记忆；第五，记忆编辑——支持对记忆库中特定文档的选择性遗忘或修正，这对于隐私保护和知识更新都是必要的能力；第六，探索MSA在Agent系统中的应用，利用其终身记忆能力构建具有持久记忆的对话Agent和任务规划系统。
复现评估
从复现角度来看，MSA的复现具有中等难度。有利因素方面：骨干模型Qwen3-4B-Instruct-2507是公开可用的；论文详细描述了超参数设置（chunk大小64、top-k=16、两阶段学习率 $1 \times 10^{-4}$ 和 $6 \times 10^{-6}$）；训练数据量为158.95B token的语料，虽然是大规模但并非不可复现的量级。挑战因素方面：论文未明确说明是否开源代码和模型权重；100M token的推理需要2xA800 GPU（80GB显存），这对大多数研究者来说是较高的硬件门槛；连续预训练需要大量计算资源；158.95B token的训练语料的具体构成和清洗流程未完全公开；多个基准测试使用了LLM Judge作为评估指标，这引入了一定的评估随机性。总体而言，核心算法的复现在学术机构级别是可行的，但要完全复现论文中的所有结果（特别是100M token的扩展性实验）需要相当的算力投入。
论文图表
该图以折线图形式展示了不同模型在RULER NIAH基准上从32K到1M token的准确率变化。MSA从98.77%（32K）仅降至94.84%（1M），下降3.93个百分点。RL-MemoryAgent-14B从98.42%降至92.66%，下降5.76个百分点。Qwen3-Next-80B-A3B从100%降至80.78%，Qwen3-30B-A3B从99%降至80%，Qwen2.5-14B-1M从100%降至53%，Qwen3-4B-Instruct从95%降至25%。MSA展现了最平缓的下降曲线和最高的终端准确率。
NIAH是评估长上下文能力的经典测试，该图直观对比了MSA与多种基线方法在极端长度下的稳定性差异，验证了MSA在外推能力上的优势。