让AI真正拥有记忆 MSA如何把上下文扩展到1亿Token
原标题：
MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens
论文链接：
https://arxiv.org/abs/2603.23516
项目主页：
https://evermind.ai/blogs/breaking-the-100m-token-limit-msa-architecture-achieves-efficient-end-to-end-long-term-memory-for-llms
代码地址：
https://github.com/EverMind-AI/MSA
这是一篇来自EverMind的论文，提出了一种可插拔的稀疏注意力机制，用优雅的工程设计绕过了长上下文的"不可能三角"——让我们从头讲清楚它到底做了什么。
"记忆"对AI意味着什么
当你和ChatGPT或Claude聊天时，它们是怎么"记住"对话的？答案很简单：把所有对话文字一股脑塞进模型，模型在回答时可以看到全部内容。这段能被模型看到的文字长度，就叫做上下文窗口（Context Window）。
问题在于，上下文窗口并不是越大越好——它有一个非常致命的成本问题。Transformer里最核心的注意力机制（Attention）在计算时，时间和显存消耗都与输入长度的平方成正比。
也就是说，输入长度翻倍，计算量变成原来的4倍；长度变10倍，计算量暴增100倍。
现有方案与MSA
面对这个问题，研究者们已经探索了三条主要路线，但每条路都有明显的缺陷：
RAG（检索增强生成）
优点：知识库无限大，扩展方便
缺点：检索模块与生成模块分离，无法端到端优化；难以捕捉隐式关联
隐状态压缩（如Mamba, RWKV）
优点：推理成本固定，不随长度增加
缺点：灾难性遗忘——压缩会丢信息，精度随长度急剧下降
全量注意力扩展
优点：信息无损，精度最高
缺点：O(n²)二次方成本，100万Token时显存需求高达数百GB
再来看MSA（文本）
优点：1亿Token，精度接近全量，成本可控
缺点：需要额外训练；工程实现复杂
MSA的核心野心是：既要RAG的可扩展性，又要全量注意力的精准度——同时把这两件事统一到一个可以端到端训练的神经网络里。
MSA的核心思想
RAG的思路是：先检索，再生成。两个步骤是分开的，检索模块不会因为生成效果好不好而得到反馈、自动改进。
MSA想做的事情是：把"检索"这个动作内化到Attention层里，让它成为神经网络的一部分，可以随整个模型一起被梯度下降所优化。
这个想法本身并不新鲜，但实现起来有一个根本难题：Attention在选择"看哪些Token"时，需要对所有候选Token打分，这本身就是O(n)的操作，无法避免。
MSA的解法是引入两层路由机制，用粗到细的方式大幅减少需要精确计算的Token数量。
第一层路由负责主题级别的粗筛：把整个知识库按照文档或段落切成很多小块，用一个轻量的注意力模块对所有块打一个粗略的相关性分数，从1亿Token的海量知识里挑出几千个最相关的文档。这一步成本很低，因为每个"块"只用一个简单的向量来代表，不需要精确计算。
第二层路由负责Token级别的精选：在第一步选出的文档里，通过Pooling（池化）操作和Top-k选择，进一步把候选范围缩小到真正关键的几千个Token。这几千个Token的KV向量，才会被加载进来，参与最终的精确注意力计算。
这样一来，最终实际参与计算的Token数量从1亿降到了几千，计算成本就变得完全可控了。
四项工程创新
可微分的稀疏化：让"选择"能被梯度训练
传统Top-k操作是不可微分的（就像"四舍五入"一样，无法传播梯度）。MSA通过特殊设计，使得"选哪些Token"这个动作本身也能参与反向传播，从而让整个检索-生成流程真正做到端到端训练。模型会自动学会"什么样的问题该去检索什么样的记忆"。
Document-wise RoPE：位置编码的聪明解耦
RoPE（旋转位置编码）是现代LLM中常用的位置编码方式，但它有一个问题：如果训练时只见过64K以内的位置，推理时突然出现第5000万个位置，模型会完全懵掉。MSA的解法是：每个文档内部的位置从0开始独立计算，而不是在全局1亿Token的序列里连续编号。这样，模型在训练时学到的"文档内部位置关系"，在推理时完全可以复用，实现了从64K训练到1亿Token推理的零代价外推。
KV缓存的分层存储：GPU+CPU协同
1亿Token的KV缓存如果全放在GPU显存里，需要几百GB——这是不现实的。MSA的解法是精妙的分层：路由用的"索引键"（体积小）放在GPU显存，供快速检索；内容KV（体积大）放在CPU内存，只在被选中时才传输到GPU。这个设计让仅仅两张A800 GPU（各80GB显存）就能支撑1亿Token的推理，把硬件需求降到了现实可操作的范围。
Memory Interleave：支持多跳推理
有些复杂问题不是一次检索就能解决的，比如"找到A项目的负责人，然后查他最近三个月的工作记录"。MSA支持多轮"生成式检索→上下文扩展"循环——模型先生成一个中间答案（比如"A项目负责人是张三"），再以此为新的查询去做下一轮检索，最终拼出完整答案。这让模型具备了多跳推理能力，而不仅仅是简单的单轮问答。
实验结果
论文的核心实验是大海捞针（NIAH，Needle-in-a-Haystack）测试：在海量文本中藏一个关键事实（比如"猫的最爱食物是三文鱼"），然后问模型这个事实是什么。这是衡量长上下文能力最直接的基准。
从32K到1M Token，MSA的精度从98.77%降到94.84%，只跌了不到4个百分点。相比之下，同底座的Qwen3-4B在1M Token时直接崩溃到24.69%（几乎等同于随机猜测），连旗舰级的Qwen3-80B也掉到了80.78%。
在MS MARCO等问答数据集上，MSA的表现同样亮眼：从16K扩展到1亿Token，性能降幅不超过9%，同时超越了同底座的RAG方案，甚至超过了业界最优的RAG系统。
总结
MSA的核心贡献可以用一句话概括：把RAG"外挂检索"变成Attention层的"内置稀疏化"，通过端到端训练让模型自己学会检索，从而在几乎不损失精度的前提下，把上下文扩展到了人类工程能力的极限——1亿Token。
它代表的方向，是让AI的"记忆"从一个工程补丁，变成神经网络本身的一部分。这条路还很长，但MSA是一个有说服力的早期里程碑。