突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元

人的智能由推理能力与长期记忆能力构成。近年来,大模型的推理能力快速发展,但其长期记忆能力始终受限于上下文长度,难以突破。历史上,多种技术路线都试图解决这一问题,但均未能同时满足扩展性(Scality)、精度(Precision)与效率(Efficiency)的要求,形成了一个“不可能三角”。近期,论文《MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens》提出了一种新的解决方案。

该工作提出了一种名为MSA(Memory Sparse Attention)的新型记忆架构。它通过记忆稀疏注意力机制、支持超长上下文外推的文档级旋转位置编码(Document-wise RoPE)、KV缓存压缩与内存并行,以及支持复杂推理的记忆交错(Memory Interleave)机制,构建了一个可扩展至1亿Token的大模型长时记忆框架。在主流的长文本问答、大海捞针等评测中,该方法取得了领先的结果。尤为重要的是,当上下文长度从16K增加到100M时,模型性能仅下降9%,展现出极强的扩展能力。

该方法可视为大模型的一个记忆插件,为解决长期记忆问题提供了新的思路。在智能体应用爆发的时代,这项工作有望成为开启“记忆即服务”(Memory-as-a-Service)新范式的重要里程碑。

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元

  • GitHub 链接:https://github.com/EverMind-AI/MSA
  • 论文链接:https://zenodo.org/records/19103670

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元
图 1:MSA-4B 在长文本问答任务中的得分随上下文长度衰减情况显著优于基线模型(论文原图)

1. LLM 长时记忆的“不可能三角”

近年来,大型语言模型(LLM)的能力边界不断拓宽,但在模拟人生(Life Long)级别的长时、细粒度记忆方面,始终面临巨大挑战。无论是需要通读并理解长篇小说的文学分析,还是要求在多轮对话中保持人格一致性的数字孪生,抑或是需要追溯漫长历史记录的智能体系统,都对模型的有效上下文长度提出了极高要求。然而,主流LLM受限于全注意力机制(Full Attention)的二次方复杂度,其有效上下文窗口长期被限制在百万(1M)Token量级,与人类一生数亿Token的记忆容量相去甚远。

为突破这一瓶颈,学界和业界探索了三条主要技术路线,但每条路线都在解决部分问题的同时引入了新的局限,形成了一个难以调和的“不可能三角”:

  1. 参数化记忆:通过持续训练或微调将知识固化进模型参数。此方法精度高,但扩展性差,更新成本高昂且易发生灾难性遗忘。
  2. 外部存储记忆:以检索增强生成(RAG)为代表,将记忆外置于向量数据库。此方法扩展性好,但其“检索-生成”两阶段分离的非端到端特性,导致检索精度成为性能瓶颈,难以进行深度语义对齐。
  3. 潜状态记忆:利用模型内部的隐藏状态(如KV缓存)作为工作记忆。此方法语义保真度高,但面临效率与容量的直接冲突。基于KV缓存局部保留的方法(如StreamingLLM)精度高但扩展性受限;而基于线性注意力的方法(如RWKV)虽实现了线性复杂度,却在超长上下文中因有损压缩导致精度严重下降。

正是在此背景下,《MSA》一文提出了一个极具雄心的目标:设计一个端到端可训练的、能以线性复杂度扩展至亿级Token、同时保持高精度的潜状态记忆框架。MSA的出现,旨在正面挑战并打破上述“不可能三角”,为LLM赋予真正意义上的“终身记忆”。

2. MSA 架构深度剖析:四大创新支柱

MSA的革命性并非源于单一技术的突破,而是一套环环相扣、系统性的架构创新。这套“创新栈”协同工作,共同构成了其高性能的基石。

2.1 核心基石:记忆稀疏注意力 (Memory Sparse Attention)

MSA的核心思想是在Transformer的注意力层引入一种可微分的、基于内容的稀疏化机制。它不再让模型在推理时关注所有历史记忆,而是设计了一个高效的“路由”(Routing)模块,动态选择最相关的记忆子集参与计算。

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元
图 2:记忆稀疏注意力层架构(论文原图)

上图是MSA技术实现的核心,展示了一个高度优化的Transformer层如何将海量外部记忆(Docs)与当前查询(Query)高效融合。其流程可分解为左右两个协同工作的模块。

左侧是标准的Transformer外壳。 MSA层被包裹在一个标准的Pre-Norm Transformer模块中。它取代了传统的自注意力(Self-Attention)层,其输出与输入通过残差连接(图中的⊕符号)相加,随后经过RMSNorm归一化和前馈网络(FFN)层处理。这一设计确保了MSA可以作为即插即用的模块,无缝集成到现有的大模型架构中。

右侧是MSA的“双重路由”稀疏注意力机制。 这是创新的核心,通过一个精巧的“双重路由”机制,实现了从海量文档中“优中选优”的过程,避免了对所有记忆进行暴力全量计算。

  • 输入与投影:模型接收外部知识库(Docs)和当前查询(Query)。知识库被投影成四份独立的表征:查询(Q)、键(K)、值(V),以及一个专门用于路由的键(K^R)。
  • 第一重路由(主题级筛选):模型首先对整个知识库进行粗粒度的“主题筛选”。它通过一个注意力机制和词元级平均池化(Token-wise mean pooling)操作,计算当前Query与所有文档“主题”(由K^R代表)的关联度,输出一个主题概率分布,以决定哪些文档在宏观上与当前查询最相关。同时,辅助损失(L_aux)在此处被计算,以确保路由的稀疏性和有效性,防止所有查询都涌向少数几个热门文档,保证记忆的均衡利用。
  • 第二重路由(词元级精筛):在第一步锁定的相关文档内部,模型进行更精细的“内容筛选”。通过池化(Pooling)和Top-k操作,在词元级别上计算查询与具体内容块的相似度,从相关文档中进一步挑选出最关键的句子或段落(图中绿色高亮的小方块)。
  • 最终注意力计算:经过“双重路由”后,只有被选中的、最精华的K和V向量才会被加载到注意力计算中。这些经过稀疏化处理的记忆K/V,与原始Query的K/V进行拼接(Concat),共同进入最终的Multi-head Attention层,生成最终输出。

这一机制的独创性在于,它将RAG系统中的“检索”步骤,内化为了一个可端到端训练的神经网络模块。与依赖外部、固定的相似性度量(如向量余弦距离)的RAG不同,MSA的路由器是在训练过程中与生成任务共同优化的(通过辅助的对比学习损失L_aux),这意味着它能学会一种更符合模型内部“世界观”、与最终任务目标更对齐的检索策略。这从根本上解决了RAG“检索”与“生成”目标不一致的核心痛点,是其实现高精度的关键。

2.2 扩展性关键:文档级旋转位置编码 (Document-wise RoPE)

2.2 位置编码解耦:实现稳定外推的核心

实现从短训练文本(如64K)到亿级推理文本的成功外推,核心挑战在于位置信息的处理。若采用传统的全局位置编码,当推理时的序列长度远超训练范围时,位置索引会发生剧烈偏移,导致模型性能急剧下降。

MSA为此提出了一个简洁而高效的解决方案:为每个独立的记忆单元(文档)分配一套独立的旋转位置编码(RoPE)。这意味着,无论记忆库中有多少文档,模型在处理每个单元时,其内部的“位置坐标系”都是从零开始且稳定不变的。这种设计将文档的内部相对位置与其在全局记忆中的绝对位置解耦,使得模型在训练时习得的位置感知能力,能够无损地泛化到包含海量文档的推理场景。这正是MSA实现惊人长度外推能力的理论基础。

2.3 工程化落地:KV缓存压缩与内存并行

理论上的可行性必须通过工程实现才能转化为现实。在亿级Token的尺度下,即便经过压缩,KV缓存的存储需求也高达上百GB,远超单个GPU节点的显存容量。MSA通过一套精巧的“内存并行”策略解决了这一物理瓶颈。

  • 分层存储:MSA观察到,在路由阶段,模型仅需要体积较小的路由键(K_R)进行快速匹配;而占据绝大部分空间的内容键值(K, V),仅在完成Top-k文档选择后才需要。因此,它设计了一套分层存储系统:
    • GPU常驻路由键:将所有文档的K_R分布式存储在多张GPU的高速显存中,确保低延迟的全局检索。
    • CPU卸载内容键值:将庞大的K和V矩阵存储在大容量的CPU内存中。
  • 异步调度:当GPU完成路由计算、确定Top-k文档后,系统再异步地从CPU内存中调取所需的内容KV到GPU,用于最终的生成计算。

这种“GPU快查、CPU慢取”的策略,优雅地将存储瓶颈从有限的GPU显存转移到了海量的CPU内存,使得在标准硬件上运行亿级Token的推理成为可能。这不仅是工程上的创举,更是该技术走向实际应用的关键前提。

2.4 复杂推理能力:记忆交错

对于需要整合多个分散在不同文档中的证据才能回答的复杂问题(即多跳推理),单次的“检索-生成”循环往往力不从心。为此,MSA引入了记忆交错机制。

该机制允许模型进行多轮次的“生成式检索 → 上下文扩展”循环。在第一轮,模型根据原始问题,首先生成它认为最相关的文档ID序列;随后,系统获取这些文档的原文,并将其追加到原始问题之后,形成一个新的、更丰富的“中间问题”;在下一轮,模型基于这个新问题,再次生成新的文档ID……如此循环往复,直到模型认为积累的证据足够充分,便停止生成文档ID,转而生成最终答案。

这种迭代式的推理链,模拟了人类“发现线索A → 顺藤摸瓜找到线索B → 整合AB形成完整证据链”的思考过程。它赋予了MSA动态规划其信息搜集路径的能力,是其在多跳问答任务上表现出色的重要原因。

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元
图3:三阶段推理流程图——包含记忆交错机制的推理过程(论文原图)

这张图以一个具体的多跳问答案例(“埃里克·瓦茨的父亲何时出生?”)为例,完整展示了记忆交错机制的三阶段工作流程:

  1. 第一阶段:全局记忆编码(图左)。这是一个离线预处理步骤。整个知识语料库被一次性编码,生成庞大的KV缓存,即全局“记忆”。此记忆库通常存储在成本更低的CPU内存或SSD中,供实时查询调用。
  2. 第二阶段:路由与上下文组装(图中)。这是推理的核心循环,图中以三步迭代展示了证据链的构建过程。
    • 在第1轮迭代中,模型根据初始问题,检索到第一个证据块:“埃里克·瓦茨是比尔·瓦茨的儿子”。此时信息不足以直接回答,模型生成引用标记,将文档内容纳入上下文。
    • 在第2轮迭代中,上下文已扩展。模型基于新形成的查询(“比尔·瓦茨何时出生?”),再次检索,找到包含比尔·瓦茨出生日期的证据块,并将其追加到上下文中。
  3. 第三阶段:稀疏生成。当证据链完整后,模型输出特殊标记宣告证据搜集结束,随后整合所有上下文证据,生成最终答案:“1939年5月5日”。

这张图直观揭示了记忆交错机制的本质:它将“推理”与“检索”深度交织,使模型能够从一个模糊的初始问题出发,通过逐步发现、逐步聚焦的方式,最终锁定精确答案。这种能力对于解决真实世界中答案分散在多个文档中的复杂问题,具有重要价值。

3. 实验数据再解读:MSA的价值验证

论文通过一系列详尽的实验,从多个维度验证了MSA架构的有效性。

3.1 惊人的扩展性与鲁棒性

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元
图4:在“大海捞针”测试中,MSA在上下文从32K扩展至1M时,准确率仅从99%下降至95%,表现出极强的稳定性。相比之下,其他长上下文模型在256K后性能显著衰减(论文原图)。

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元
图5:在MS MARCO问答测试中,将记忆规模从16K扩展至100M(跨越4个数量级),MSA的性能评分衰减率不足9%。这直观证明了其架构在抵抗大规模无关信息干扰方面的卓越鲁棒性(论文原图)。

3.2 端到端优化的威力

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元
图6:在9个QA基准测试的平均分上,4B参数的MSA模型显著优于基于同样4B底座构建的、包含重排器的复杂RAG系统,甚至在多个数据集上超过了由顶级检索器和235B参数大模型组成的RAG系统。这充分证明了MSA端到端优化带来的高精度优势(论文原图)。

3.3 各组件的不可或缺性

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元

4. 总结:MSA 的独创性与核心价值

综合上述分析,MSA 架构的核心独创性与价值可总结如下:

突破记忆瓶颈!MSA架构实现100M Token长时记忆,开启大模型终身记忆新纪元

总而言之,MSA 的真正价值,并不仅仅是发布了一个性能强大的长上下文模型,而是为 AI 记忆领域提供了一套全新的、经过完整验证的、兼具扩展性、精度与效率的技术框架。它证明,我们不必在 RAG 的“低精度”和全注意力的“高成本”之间做出妥协。通过将稀疏化思想与神经网络的端到端学习能力相结合,构建一个独立的、可扩展的、与 LLM 兼容的“记忆层”是完全可行的。

这为未来 AI 生态的发展描绘了一幅新的蓝图:记忆可以作为一种独立的、可插拔的服务,与各种推理核心(LLM)自由组合。用户的数据和“记忆资产”将不再被锁定在单一的模型或平台中。从这个角度看,MSA 不仅是一项重要的学术成果,更可能成为开启“记忆即服务”(Memory-as-a-Service)新阶段的里程碑。

5. 背景信息:EverMind 与“发现式 AI”愿景

为更完整地理解 MSA 研究背后的驱动力,有必要了解其出品方 EverMind 的宏观战略背景。EverMind 是盛大集团在 AI 领域长期孵化的核心团队之一,其使命是攻克 AI 的长期记忆难题,推动 AI 向自我演化(Self Evolving)的方向发展。

根据相关专访信息,盛大集团的 AI 战略聚焦于构建一个更具开创性的“发现式 AI(Discoverative AI)”生态。其终极目标是让 AI 辅助人类发现新知识、解决根本性问题,而非仅仅模仿和重组已有信息。在这一愿景中,两大技术基石被置于核心地位:

  • MiroMind:专注于推理。该团队致力于通过可验证推理等技术路径,提升模型的可靠性与主动发现能力。
  • EverMind:专注于记忆。该团队的使命是为 AI 打造一个可无限扩展、高保真、且独立于特定模型的长期记忆系统。稳定可靠的记忆是进行复杂推理与知识创造的基础。

因此,EverMind 与 MiroMind 共同构成了“发现式 AI”蓝图的核心驱动力,分别对应着“记忆”与“推理”两大认知支柱。本文所介绍的 MSA 架构,正是 EverMind 团队践行其理念的核心技术成果。其设计不仅是对现有长文本技术瓶颈的突破,也体现了在构建独立、可控的 AI 基础设施上的长期投入。


关注“鲸栖”小程序,掌握最新AI资讯

本文来自网络搜集,不代表鲸林向海立场,如有侵权,联系删除。转载请注明出处:http://www.itsolotime.com/archives/26458

(0)
上一篇 2小时前
下一篇 2小时前

相关推荐