微信扫一扫,关注公众号

  • 科技行者

  • 算力行者

见证连接与计算的「力量」

首页 东京科学大学与帝国理工学院联手破解AI"记忆难题":让大模型不再为读长文章而烦恼

东京科学大学与帝国理工学院联手破解AI"记忆难题":让大模型不再为读长文章而烦恼

2026-05-27 10:16
分享至:
----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.-
2026-05-27 10:16 科技行者

这项由东京科学大学与英国帝国理工学院联合开展的研究,以预印本形式于2026年5月发布在arXiv平台,论文编号为arXiv:2605.18226。研究的核心问题说起来其实很接地气:当你让一个AI助手参考一大堆背景资料来回答你的问题时,它到底有多慢、有多健忘?

以现在大家熟悉的ChatGPT、Claude这类AI聊天工具为例,它们在回答问题之前,往往需要先"读完"你给它的一大段背景信息——可能是几百个例子、一本规则手册、或者一堆参考文件。读完这些之后,AI才开始生成回答。问题在于,这个"读"的过程不仅耗时耗资源,而且随着对话越来越长,AI会越来越"忘记"最初那段背景资料说的是什么,就好像一个人读了太多东西,脑子里塞满了新内容,把最开始读的那部分给挤没了。更麻烦的是,每次有新用户来问问题,AI都得把这段背景资料从头到尾再读一遍,哪怕这段资料一个字都没变过。

研究团队把这个现象称为两个结构性代价:其一是"前缀衰减"——随着AI生成的回答越来越长,背景资料的影响力越来越弱;其二是"推理低效"——背景资料越长,AI处理每一个新词时需要的计算量就越大,速度越来越慢,占用的内存也越来越多。Anthropic公司(Claude的开发者)甚至公开表示,他们的代码助手Claude Code整个系统都围绕着"提示词缓存"这一技术构建,目的就是减少这种重复读取的成本——这足以说明这个问题在工业界有多头疼。

为了解决这个问题,研究团队提出了一种他们称之为"注意力状态记忆"(Attention-State Memory,简称ASM)的新方法。这个方法的核心思路可以用图书馆的读书笔记来理解:与其让AI每次都把那本厚厚的参考书从头读到尾,不如提前让AI读一遍,把它读书时的关键感悟和要点整理成一本精华笔记本,之后的每次查询,AI只需翻翻这本笔记本就行了。这个方案既不需要修改AI本身的"大脑"(不需要重新训练模型),也不需要保留那本厚书(不需要在推理时保存完整的背景资料),只需要保留那本经过精心整理的笔记本。

一、背景知识:AI读长文章究竟有多累

要理解这项研究解决了什么问题,先得弄清楚现代AI在处理长背景资料时究竟经历了什么。

现代大语言模型(Large Language Model,LLM)在回答问题时,其内部有一个叫做"注意力机制"(Attention)的核心计算模块。你可以把这个机制理解成AI的"眼神扫描器"——每当AI要生成一个新词,这个扫描器就会回头扫描之前所有的文字,判断哪些部分对当前这个词最重要,然后综合所有这些信息来决定该写什么。

问题就出在这个"回头扫描"上。假设背景资料有一万个词,AI每生成一个新词,就要扫描这一万个词一遍。如果AI要生成五百个词的回答,那这一万个词就要被扫描五百遍。背景资料越长,每次扫描的范围就越大,速度就越慢,占用的显存就越多。这种关系是线性的——背景资料翻倍,计算量也翻倍。

目前学界和工业界应对这个问题主要有两条路。第一条路叫"前缀压缩"——把那本厚书变薄,把一万个词压缩成两千个词,然后继续让AI读这本精简版的书。这条路虽然节省了一些资源,但AI依然需要在每次生成新词时回头扫描那两千个词,计算量减少了,但没有根本消除;而且压缩过程往往会损失信息,就像把一本完整的教材缩写成摘要,难免会丢掉一些重要细节。第二条路叫"前缀内化"——通过专门的训练,让AI把那本书的内容"背"进自己的神经网络参数里,之后回答问题时完全不需要那本书了。这条路彻底解决了推理时的扫描问题,但代价极大——需要针对每一本参考书单独训练一次,耗时耗资源,而且参考书内容一旦更新,就得重新训练。

研究团队指出,这两条路各有无法绕开的根本局限:前者没有真正去掉推理时的扫描负担,后者需要梯度训练,成本太高。于是,他们设计了第三条路——既不在推理时扫描那本书,也不通过训练把书背进大脑,而是提前做好读书笔记,推理时只查笔记。

二、注意力状态记忆的数学基础:为什么笔记可以替代原书

研究团队能走这条第三条路,依赖于一个数学上的关键发现,他们把它称为"在线Softmax恒等式"(Online-Softmax Identity)。这个发现其实来自更早的研究(特别是FlashAttention等高效注意力计算方法),但研究团队把它用在了一个全新的场景上。

这个数学发现说的是:AI的注意力计算有一个特殊的可拆分性质。具体来说,如果AI要对一段很长的文字(比如一本书)做注意力计算,这个计算可以分成很多小块分别完成,最后再把结果合并起来,合并后的结果和一次性读完整本书的结果完全一致,没有任何精度损失。

研究团队从这个性质中发现了两个对他们的方法至关重要的机会。第一个叫"充分性":对于一个给定的查询(也就是AI正在生成的某个词),只要保存这个词对那本书做注意力计算时得到的"注意力状态"(具体来说是注意力输出值和一个归一化系数),就完全可以在未来需要时还原出这段计算结果,不需要再保存原来那本书了。这就好比说,你读完一本书之后,只要把你的读书感悟记录下来,以后回忆这本书对你的影响时,只需要翻感悟笔记,不需要再重读原书。

第二个叫"可组合性":如果你已经有了两段不同内容对同一个查询的注意力状态,可以用一个简单的公式把它们合并成一个注意力状态,这个合并后的状态等价于对这两段内容拼在一起做一次注意力计算的结果。这就意味着,你可以把一本很长的书分成很多章,分别做好笔记,再把这些笔记合并成一本综合笔记,效果和一次性读完整本书完全一样。

这两个性质加在一起,勾勒出了一个新的处理方案:与其每次推理时都让AI扫描那本书,不如提前针对一批代表性的查询,计算好它们对那本书的注意力状态,把这些状态存成笔记本,推理时直接查笔记本就好。这就是注意力状态记忆(ASM)的核心思想。

三、注意力状态记忆的具体工作方式:笔记是怎么做出来的

ASM的工作分为两个阶段:离线的"制作笔记"阶段,和在线的"查笔记"阶段。

制作笔记的过程分为收集和聚类两个步骤。首先是收集步骤:研究团队给AI提供一批"背景资料+示例问答"的组合,让AI做一次正向传播计算(也就是正常读一遍,但不做训练,不更新任何参数)。在这个过程中,每一层注意力模块里,每一个属于"回答部分"的词(也就是示例回答中的每个词),都会和背景资料做一次注意力计算,产生一个"注意力状态"。研究团队把这些词对应的查询向量(可以理解为这个词的语义特征)和它对应的注意力状态一起记录下来,形成一个原始数据集。

然后是聚类步骤:原始数据集可能包含成千上万条记录,直接存起来太占空间,查起来也太慢。研究团队使用K均值聚类算法(K-means),把这些记录按照查询向量的相似性分成K个组,每组计算出一个代表性的"质心"。对于每个质心,它存储的注意力状态是通过一种叫做"注意力感知聚合"的方式合并出来的——本质上就是用那个可组合性公式,把这一组里所有记录的注意力状态合并成一个。这样,最终的笔记本只有K条记录,K是一个可以自由设定的超参数,典型值是几千到一万多。

这个笔记本的结构很简单:每条记录包含三个东西——一个代表性的查询向量(用来在推理时匹配最相似的记录)、一个注意力输出值(存储了背景资料的精华)、以及一个归一化系数(用于后续的合并计算)。每一层注意力模块各自拥有独立的一本笔记本,全部笔记本的构建过程只需要做正向传播,不需要任何梯度计算或参数更新。

在线查笔记的过程也分为两步:检索和合并。检索时,AI对当前正在处理的词计算出一个查询向量,然后在当前层的笔记本里找余弦相似度最高的那条记录,把对应的注意力状态取出来。合并时,AI用"可组合性"公式,把检索到的笔记状态和AI当前对非背景部分(也就是用户的问题和已生成的回答部分)做注意力计算得到的状态合并在一起,得到最终的注意力输出。这个最终输出和"老老实实读完整本背景资料"得到的输出在理论上等价,而实际上因为用质心近似代替了真实的查询状态,会有一定近似误差,但实验证明这个误差对最终回答质量的影响很小。

四、让制作笔记更省内存:分块制作的技巧

研究团队还解决了一个实际工程问题:制作笔记本身虽然不需要训练,但如果背景资料很长(比如几万个词),在制作过程中仍然需要把整段背景资料一次性装进GPU显存,这对内存受限的设备来说是个麻烦。

解决方案直接来自那个"可组合性"性质:既然注意力状态可以分块计算再合并,那么制作笔记时也可以把背景资料分成若干块,每块单独做一次正向传播,分别得到各块的注意力状态,然后用合并公式把所有块的状态合并到一起。这样,GPU每次只需要处理一小块资料,内存峰值大幅降低,而最终得到的笔记本和一次性处理整段资料的效果完全一样。

研究团队在实验中验证了这个技巧的有效性。他们测试了一个16K词长度的背景资料,分别用一次性处理和分成四个4K块分别处理两种方式制作笔记,结果发现分块方式的准确率(78.5%)和一次性处理(79.0%)几乎完全相同,而每次GPU只需要处理4K词而不是16K词,内存压力降低到原来的四分之一。

五、让查笔记更快:分层检索结构

笔记本里的条目数K越大,笔记越详细,理论上效果越好,但查找时间也越长——如果要在一万条记录里找最相似的那一条,需要把所有一万条都比对一遍,时间复杂度是O(K)。研究团队通过引入分层索引结构把这个复杂度降到了O(log K)。

具体来说,他们把笔记本组织成两层结构:第一层是粗粒度的"大类目录",第二层是每个大类下的详细条目。查找时,先在第一层目录里找最相似的几个大类(研究中发现取前16个大类就足够保证准确率),然后只在这几个大类的条目里做精细搜索。这样,无论笔记本有多大,查找时间都不会随之线性增长,而是以对数速度增长,极大地限制了推理延迟随内存预算增大而增加的趋势。

六、针对现代AI架构的适配:GQA的处理

现代大模型大多使用一种叫做"分组查询注意力"(Grouped-Query Attention,GQA)的架构优化——简单来说,就是多个注意力"读头"共享同一套键值缓存,以此减少内存占用。研究团队也对ASM做了相应适配。

在GQA设置下,同一组内的多个查询头共享同一套背景资料的键值,这意味着它们对背景资料的注意力状态在结构上是耦合的。研究团队的处理方式是:把同一组内所有查询头的查询向量拼接成一个更长的向量,用这个拼接向量作为笔记本的检索键。这样,一次检索就能同时服务于整组查询头,效率更高,存储的质心数量也可以按键值头的数量而不是查询头的数量来设计,进一步节省空间。

七、实验结果:笔记本的效果究竟怎么样

研究团队在两类主流应用场景上做了全面测试:多样本上下文学习(In-Context Learning,ICL)和检索增强生成(Retrieval-Augmented Generation,RAG),使用的基础模型是LLaMA 3.1-8B。

在ICL测试中,他们使用了ManyICLBench这个专门评估多样本上下文学习的基准测试集,选取了其中五个任务(意图分类、对话关系识别、信息检索等)进行评测。测试的设置是把一个32K词长的前缀(包含大量示例问答对)压缩成不同大小的笔记本(从1K到16K条记录),然后看回答准确率如何。

结果显示,在1K到8K条记录的范围内,ASM的平均准确率持续高于或持平于直接使用同等内存的ICL方式——也就是说,用同样多的内存,ASM比ICL存储更多有用信息、给出更准确的回答。这个优势在内存预算较小时尤为明显:当每个词条目存储的是原始键值缓存(ICL方式)时,条目数少意味着只能保留前缀中很少的一部分词;而当每个条目存储的是覆盖整个前缀的注意力状态(ASM方式)时,即使条目数很少,每个条目也携带了前缀的全局信息。简单说,ICL方式的笔记本存的是原文片段,ASM的笔记本存的是读后感,读后感更浓缩,信息密度更高。

在延迟测试中,研究团队直接测量了每生成一个新词时注意力模块所花费的时间。ICL方式的延迟随前缀长度线性增长,毫无悬念。ASM方式的延迟则随条目数增长极为缓慢:在4K条记录时,ASM就已经比ICL快了,在16K条记录时,ASM比ICL快了1.8倍。以8K条记录这个常用配置为例,注意力延迟降低了1.36倍。

与另一种主流的KV缓存压缩方法KVzip相比,ASM在几乎所有设置下都表现更好。研究团队分析认为,KV缓存压缩方法在压缩过程中容易丢失那些对回答至关重要的"标签词"(比如分类任务里各个类别的名称),而ASM存储的是对整个前缀的综合注意力输出,这类关键信息更不容易被稀释掉。

对于两个推理类任务(数学计数和科学问答),ICL本身在这类任务上的提升就很有限,ASM的表现与ICL大体相当——这说明ASM忠实地继承了ICL能带来的所有收益,既不会无中生有地提升ICL本就无法提升的能力,也不会无缘无故地损失ICL能提供的那点收益。

在RAG测试中,他们使用的是RuleArena中的NBA规则推理任务,前缀是一本约两万词长的NBA球员交易规则手册。测试比较了三种设置:零样本(不提供规则手册,21.2%准确率)、完整ICL(提供全部两万词规则手册,24.1%准确率)、以及不同大小的ASM笔记本(1K到16K条记录)。

结果颇为亮眼:在4K条记录这个最优配置下,ASM达到了25.5%的准确率,不仅超过了零样本基线,还超过了完整提供规则手册的ICL基线。而4K条记录的内存占用只是完整保存两万词规则手册KV缓存的约20%。研究团队还观察到一个有趣现象:准确率不是随着条目数单调增长的,而是在某个中间值(4K)达到峰值,之后又下降。这说明笔记本的大小是一个需要针对具体任务调整的超参数,并不是越大越好,过多的质心会导致每个质心覆盖的样本过少,代表性反而下降。

八、换个模型还管用吗:跨架构泛化实验

一个自然的疑问是:这套方法是否只对LLaMA 3.1-8B这一个模型有效?研究团队在另外两个模型上做了验证测试:规模更小的LLaMA 3.2-3B,以及来自完全不同家族的Qwen3-8B。

在banking77(意图分类)任务上,这两个模型上的ASM同样在小内存预算下明显超越同等内存的ICL,随着内存增大两者逐渐趋同。在NBA规则推理任务上,两个模型的ASM都超越了各自的零样本基线和完整ICL基线,与LLaMA 3.1-8B上观察到的规律一致。

值得一提的是,不同模型上最优的笔记本大小有所不同——研究团队认为这与各模型注意力头的几何结构和GQA的分组方式有关,并不需要特别的架构适配,只需在验证集上做一次简单的超参数搜索即可确定。

九、方法的局限性:哪些场景下笔记本会失效

研究团队坦诚地指出了ASM的核心假设及其适用边界:这套方法依赖于"查询向量具有局部聚类结构"这一前提,即来自相同任务或相同场景的查询向量在空间上是聚集的,少量质心就能较好地代表整个查询分布。

这个假设在ICL和RAG场景下通常成立——用户的问题往往属于同一类任务,查询分布比较集中。但在一些极端情况下,比如一个跨越多个话题的长期多轮对话(时而聊技术问题,时而聊生活随感,时而讨论学术内容),查询分布可能非常分散,少量质心难以覆盖所有可能的查询模式。在这类场景下,ASM的检索准确率可能下降,方法效果会打折扣。研究团队把适配非稳态分布的查询场景列为未来的研究方向之一。

说到底,这项研究本质上是在回答一个工程师每天都在面对的问题:当AI需要反复参考同一份参考资料时,有没有办法让它第一次读完之后就不用再读第二遍,而且这个"一次读完"的成本要尽可能低?他们给出的答案是:可以,只要你在第一次读的时候足够聪明地做笔记,记录下的不是原文片段,而是经过提炼的"读后感",以后再有类似的问题时,翻翻笔记就够了。

这个答案的实用性在数字上得到了体现:同等内存预算下更高的回答准确率,以及最高1.8倍的推理加速。对于那些需要大量部署AI服务的企业来说,这意味着在不降低回答质量的前提下,可以用更少的显存服务更多的用户,或者让相同硬件上的AI响应得更快。对于普通用户来说,这意味着未来的AI助手在参考长篇资料时可能会更快、更不容易"忘记"最初的任务说明。

感兴趣的读者可以在arXiv上通过编号arXiv:2605.18226找到完整论文,代码也已在GitHub上公开,有技术背景的读者可以进一步探索。一个值得深思的问题是:如果AI可以把外部知识以"笔记"形式存储,那么"学习"和"记忆"之间的边界,在AI系统里究竟应该如何划定?

Q&A

Q1:注意力状态记忆(ASM)和普通的KV缓存有什么区别?

A:普通KV缓存存的是背景资料里每个词的键值对,推理时还是要把所有词都扫描一遍,内存和计算量随背景资料长度线性增长。ASM存的是对整段背景资料做注意力计算后得到的"综合输出状态",推理时不需要扫描背景资料,只需检索最相近的记录条目,检索成本以对数速度增长,远比线性慢。两者占用相同内存时,ASM能存下对更长背景资料的压缩感悟,而普通KV缓存只能保留背景资料的一个片段。

Q2:ASM制作笔记本需要重新训练模型吗?

A:完全不需要训练。制作笔记本的过程只需要对模型做正向传播(也就是让模型正常读一遍背景资料和示例数据),记录中间的注意力状态,再做一次聚类,整个过程没有任何梯度计算或参数更新。这意味着背景资料更新时,只需重新运行一次正向传播来更新笔记本,成本远低于需要重新微调模型的方法。

Q3:ASM在推理时完全不需要背景资料了吗?

A:是的,推理时完全不需要背景资料本身。用户的问题直接输入模型,模型在每一层注意力计算时从笔记本里检索出最相近的注意力状态,将其合并进当前的注意力输出,整个过程不需要背景资料的任何原始内容,原始背景资料的KV缓存可以完全丢弃,这是ASM在内存使用上能大幅低于传统方法的根本原因。

分享至
0赞

好文章,需要你的鼓励

推荐文章
----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.-