香港科技大学（广州）的研究者如何让AI记忆力翻倍

这项由香港科技大学广州主导的研究成果发表于2026年第43届国际机器学习大会ICML 2026会议地点为韩国首尔论文收录于PMLR第306卷。论文预印本编号为arXiv:2605.05838有兴趣深入了解的读者可通过该编号查阅完整原文。一、先从一个让人抓狂的场景说起你有没有经历过这种情况跟朋友聊了很久聊到最后对方居然把你开头说的事忘得一干二净。你重复了三遍对方还是摸不着头脑。现代的大型语言模型就是驱动ChatGPT这类工具的核心引擎在处理超长文本时也面临着类似的困境——它们在阅读和处理长篇内容时往往记不住开头的细节越到后面早期的信息就越模糊。这个问题在AI研究圈里被称为长序列建模难题而解决它的主流方法之一是使用一种叫做线性注意力机制的技术。不过这种技术虽然在处理速度上非常高效却在记忆能力上有明显短板。香港科技大学广州的研究团队针对这一短板提出了一种名为动量DeltaNetMomentum DeltaNet简称MDN的新模型并在ICML 2026上公开了这项成果。二、传统的AI记忆方式像是一个只会看眼前的学生要理解这项研究先得了解AI是怎么读文章的。传统的Transformer模型也就是绝大多数现代语言AI的核心采用一种叫做自注意力机制的方式工作。用一个直观的比喻来说每当模型读到一个新词它都会回头看看之前读过的所有词逐一衡量它们和当前词的关联程度然后综合所有相关内容生成理解。这种方式非常全面但有一个严重的代价——随着文章变长需要回看的内容呈平方级别增长。文章长度翻倍计算量要变成原来的四倍。这就好比一个学生每读一个新字都要把全书翻回去重读一遍效率极低。为了解决这个效率问题线性注意力机制应运而生。它的思路是不再每次都回头看所有内容而是维护一个快速记忆本在论文中叫做快速权重Fast Weight。每读一个新词就把信息压缩汇总写进这个记忆本之后只查这个本子就够了。这样计算量就从平方增长变成了线性增长大大提高了效率。然而这个记忆本有个根本性的缺陷它的容量是固定的而且采用一种叫做随机梯度下降SGD的简单更新方式来写入信息。这种方式就像每次往杯子里加水时都会把之前的水晃出去一些——信息会快速衰减早期写入的内容很难被稳定保留。换句话说这类模型的记忆本写入策略太简单导致它在需要回溯早期信息时经常出错。三、动量一个从物理学借来的智慧研究团队的核心灵感来自优化算法领域的一个经典技巧——动量Momentum。在物理学里一个滚动的球因为有惯性不会因为路面微小的坑洼而突然停下而是会沿着整体趋势继续滚动。优化算法中的动量机制模仿了这种惯性它不只看当前这一步的更新方向还会把之前积累的历史更新方向纳入考量形成一个更稳定、更有方向感的合力。在这个研究里团队把这种动量机制引入到了线性注意力的记忆写入过程中。具体来说他们在原有的快速记忆本St旁边额外维护了一个动量记忆本Mt。每次更新快速记忆本时不再简单地根据当前这一步的误差直接写入而是先参考动量记忆本里积累的历史趋势再决定怎么更新。这个过程可以用炒股来理解。简单的SGD方式就像一个只看今天股价涨跌的短线投机者今天涨就买今天跌就卖策略极不稳定。而动量方式则像一个追踪趋势的中线投资者他不只看今天还参考过去一段时间的走势因此能做出更稳定、更有预见性的决策。对于AI记忆本的写入而言动量机制让模型能更好地抵抗当前输入的噪声干扰把真正重要的长期规律稳定地写进记忆里。从研究团队建立的数学框架来看这个带有动量的更新过程可以写成两个协同工作的方程动量记忆本的更新式为 Mt μt × Mt-1 - ηt × kt × (误差项)?而快速记忆本的更新式为 St αt × St-1 - βt × Mt。其中μt控制历史动量的保留程度βt控制动量对快速记忆本的影响力度αt是遗忘控制系数ηt是缩放因子。这些系数都是根据输入内容动态调整的让模型能针对不同的内容自适应地调节记忆策略。四、并行计算的难题如何让看历史变得高效理解了动量机制的原理接下来面临的是一个非常实际的工程挑战训练大型语言模型需要同时处理大量数据而动量机制天生就是串行的——第t步的动量必须等第t-1步的动量计算完才能算这种一步接一步的计算方式让并行处理极为困难。此前一些研究团队如TTT、Titans等为了让动量能够并行计算采用了一种折中方案把序列切成若干小块在每个小块内部做近似处理假装小块内的所有时间步是同时发生的。这种方案带来的问题是模型在训练时的行为和在实际推理生成文本时的行为出现了不一致——训练时假设了一些在推理时并不存在的信息导致效果打折扣。还有一种滑动窗口方案如Atlas模型虽然保持了一定的因果性但只能看到有限的历史上下文。香港科技大学广州的团队采用了一种全新的数学技巧来彻底解决这个问题。他们的关键发现是动量更新公式中有一个复杂的嵌套求和结构乍看之下无法并行但通过一种几何重排序的变换可以把这个嵌套结构展开成一个可以并行计算的形式。这个变换的核心思路可以用以下直觉来理解原本的嵌套求和像是在一个下三角区域里按行扫描而经过重排序之后同样的区域被按列扫描。两种扫描方式覆盖完全相同的点但按列扫描可以把每一列的贡献独立计算天然适合并行化。论文中用数学符号将这个变换表示为∑?∑?≤? a?·b? ∑?∑?≥? a?·b?这个等式看似简单却是整个并行化方案的核心支柱。基于这个变换研究团队推导出了一套完整的分块并行Chunkwise Parallel算法。整个序列被切成若干固定大小的块chunk每个块内部的计算可以高度并行块与块之间通过精心设计的状态传递机制串联起来。并且每个块内需要计算的各种系数μ、α、c、b、Γ等都可以在对数域log-domain内通过前缀求和cumsum高效完成复杂度仅为O(log C)其中C是块的大小。这样整个训练过程既保留了严格的因果关系每个时间步只看自己和之前的内容又实现了高效的并行计算从根本上解决了动量机制与并行训练之间的矛盾。五、第二个维度用物理系统的视角看记忆稳定性研究团队不仅设计了算法还从理论层面深入分析了动量机制对模型行为的影响这部分分析用到了动力学系统的语言。可以把一个线性注意力模型的记忆更新过程想象成一个弹簧系统。传统的一阶系统没有动量就像一个单纯的阻尼弹簧给它一个扰动它会以固定的速率衰减回到原点运动轨迹只在一个维度上。而引入动量之后系统变成了一个二阶系统就像一个连接了两根弹簧的质量块它的运动方式更加丰富可以产生振荡螺旋形衰减而不仅仅是单调地衰减。在数学上这种差异体现为状态转移矩阵的特征值的分布。一阶系统的特征值只能分布在实数轴上只有正负之分没有虚部。而二阶系统引入了复数特征值也就是带有虚部的特征值。这些复数特征值意味着系统能够表达相位信息记住的不只是多少还有节奏和韵律。这从本质上扩展了模型的表达能力。然而复数特征值是一把双刃剑。如果特征值落在复平面的左半部分负实部系统会产生符号翻转的行为——记忆本的某些方向会在正负之间不断震荡导致快速记忆和动量记忆之间的协作被破坏最终引发训练过程中的数值爆炸表现为出现NaN错误。为了解决这个问题研究团队设计了一套约束机制将特征值的分布限制在复平面的右半部分正实部包含第一和第四象限。具体实现方式是通过一个巧妙的参数化方案引入一个角度参数θt arctan(ηt × s)然后令αmax cos?(θt)βmax sin?(θt)这自然地保证了αmax βmax 1从而在几何上将特征值约束在稳定区域内。同时对动量系数μt的下界进行约束μlog ≥ -2避免动量过小导致机制失效。这套参数化方案既不引入太多额外参数控制各系数的投影矩阵维度远小于模型主体又从理论上保证了训练稳定性。六、MDN模型长什么样把动量机制和稳定性约束组合在一起就得到了动量DeltaNetMDN模型。它的整体架构沿袭了GDNGated DeltaNet和Comba两个前辈模型的框架但在几个关键位置做了改动。每一层MDN的核心计算单元接收四类输入查询向量q用于提问从记忆里读取信息、键向量k用于索引决定往记忆哪里写、值向量v实际要写入的内容、以及四个控制系数α遗忘控制、β误差写入力度、μ动量保留程度、η缩放因子。这些控制系数都是由输入内容经过轻量级的线性变换动态生成的让模型能根据当前上下文自适应地调节记忆策略。训练时MDN使用分块并行算法高效处理整个序列生成文字时MDN切换为逐步递推的方式Eq.4-5所示的递推公式确保训练和推理行为完全一致不存在训练-推理不匹配的问题。这一点是相较于TTT、Titans等采用块级更新方案的模型的重要优势。在实现层面研究团队用Triton一种面向GPU的底层编程语言编写了高度优化的核函数。针对反向传播训练时计算梯度的过程团队采用了一个聪明的策略在前向传播时保存修正值Ve而不是保存完整的快速记忆本S和动量记忆本M在反向传播时再从Ve重建S和M。这样的设计在保持较低内存占用的同时显著提升了训练吞吐量。七、实验结果记得更多答得更好研究团队在多个层面对MDN进行了系统性评测评测规模涵盖4亿参数400M和13亿参数1.3B两个量级与Transformer、Mamba2、GDN、Comba、KDA五个基线模型进行了全面对比。所有模型在完全相同的数据SlimPajama数据集的1000亿token子集和训练配置下训练确保对比的公平性。在合成记忆测试中研究团队使用了一个叫做多查询关联回忆MQAR的任务。这个任务的设定相当直接给模型看一串键值对比如L对应2I对应3N对应0……然后提问L对应什么A对应什么要求模型准确回忆。随着序列长度和键值对数量的增加任务难度急剧提升。测试结果显示MDN在128维和256维两个配置下都在不同序列长度和键值对数量组合上取得了接近或超越KDA一个专门针对关联回忆优化的模型的成绩远超Mamba2和GDN。在语言建模的基础指标上MDN在400M规模下的LAMBADA困惑度一种衡量语言预测准确度的指标数值越低越好达到了41.62明显优于Transformer的54.36、Mamba2的60.42和GDN的45.63在1.3B规模下同样以14.87的成绩领先所有对比模型。在常识推理任务包括HellaSwag、LAMBADA准确率、ARC-Easy、ARC-Challenge、PIQA、WinoGrande、BoolQ、SciQ等八项任务的平均得分上MDN在400M规模下以49.42分排名线性模型第一在1.3B规模下以58.82分同样排名第一超过了包括Transformer在内的所有对比模型。在上下文检索能力测试包括FDA、SWDE、SQuAD、NQ、TQA、DROP六项任务均使用2000个token的输入长度上MDN的优势更为突出。400M规模下MDN的平均检索得分为26.76比第二名KDA的24.47高出超过两个百分点FDA单项一种从长文档中抽取结构化信息的任务得分高达28.07远超KDA的18.44。这一显著提升印证了动量机制对于提升记忆检索能力的有效性。在长上下文建模能力评测中研究团队使用LongBench基准以16000个token的输入长度测试1.3B模型。MDN在代码补全LCC、RBP两项、文本摘要GovReport、QMSum等等任务上均取得了各模型中的最高分整体平均得分20.18超过第二名GDN的19.28和KDA的18.62。在大海捞针Needle-In-A-Haystack测试中研究团队把一段特定信息针藏在长达2000至16000个token的长文本中要求模型找到并回答相关问题。MDN在所有测试场景中均有改善尤其在训练长度4k token之外的超长上下文8k和16k中表现出色。在最具挑战性的多针设置文中藏了多条相关信息中8k长度下MDN在多键检索MK、多查询MQ、多值MV三个子任务上分别达到38.60、35.15和27.60比各自最强基线模型高出13.40、11.45和8.95分。在效率方面MDN的推理解码延迟与GDN和Comba几乎相同均远低于Transformer保持了线性模型在推理阶段的速度优势。训练吞吐量方面由于MDN需要同时维护两个状态快速记忆本S和动量记忆本M其训练速度略低于GDN和Comba但通过物化修正值策略已经达到了与Mamba2和KDA接近的水平内存占用也保持在合理范围内约44.89GB仅略高于Comba的42.34GB和GDN的39.38GB。八、消融实验每一个设计决策都有它的道理为了验证MDN各个设计组件的必要性研究团队在400M模型上进行了系统性的消融实验即逐一去掉某个组件看效果如何变化。去掉输出修正Output Correction即Comba引入的查询修正机制qt qt - dkt后模型的检索得分从26.76下降到25.52但仍优于GDN和Comba说明动量机制本身就能带来独立的改善。去掉动量机制将μ设为0退化为GDN配合αmax/βmax约束的版本后检索得分大幅下降到20.12且LAMBADA困惑度升至47.01清晰地说明动量是MDN性能提升的核心驱动力。去掉对μlog下界的约束即不对μ进行最小值截断会导致训练在第70步左右出现NaN崩溃证明了这个约束对于训练稳定性的重要性。去掉αmax约束会导致训练在第一步就出现NaN说明这个约束对于防止系统从一开始就不稳定是不可或缺的。将ηt的激活函数从tanh(·)1改为2·sigmoid(·)会导致LAMBADA困惑度上升到49.10检索得分下降到25.54说明激活函数的选择对模型行为有实质影响。对μlog下界的敏感性分析显示μlog_min -2是最优设置能同时取得最低的LAMBADA困惑度和最高的检索得分。提高下界如设为-1虽然仍能正常训练但检索性能有所下降说明过于限制动量的范围会削弱模型的表达能力。在混合模型实验中将部分线性注意力层替换为标准全注意力层当线性层与全注意力层的比例为7:1时MDN-H在语言建模平均分上进一步提升同时保持了有竞争力的检索能力暗示MDN可以有效减少对全注意力层的依赖是构建高效混合架构的潜力单元。研究团队还分析了快速记忆本在解码过程中的变化幅度?St 每步快速记忆本变化量的Frobenius范数均值。结果显示MDN在大多数解码步骤中的记忆变化幅度显著高于Comba和GDN说明MDN的记忆本在推理过程中更活跃——它会根据读入的新内容更积极地更新自身状态而不是趋于静止这种活跃的记忆动态与MDN在检索和推理任务上的优异表现相互呼应。说到底这项研究讲的是一个关于如何记得更好的故事。在AI领域记忆能力的提升不只是技术参数的改善它意味着未来的AI助手能更准确地记住你之前说过的话能在阅读一份很长的合同或报告时不遗漏关键细节能更好地理解跨越长篇幅的故事情节和逻辑推理。当然这项研究也有它坦诚承认的局限。由于算力限制实验仅进行到13亿参数规模7B乃至更大规模的验证还有待完成。训练吞吐量相较于最优化的一阶线性注意力模型仍有差距内核优化和内存效率还有提升空间。此外混合架构的层数比例和位置安排也还有更多探索余地。这项研究提醒我们很多时候提升AI能力的钥匙就藏在那些已被充分研究的数学工具里只是还没有人找到合适的方式把它们用对地方。有兴趣深入了解这些技术细节的读者可以通过arXiv编号2605.05838查阅完整论文。QAQ1动量DeltaNetMDN和普通的线性注意力模型相比最大的区别是什么A普通线性注意力模型如GDN、Mamba2更新记忆本时只看当前这一步的信息类似于只看今天股价就做交易决策。MDN额外维护了一个动量记忆本每次更新时参考历史积累的趋势因此对噪声干扰更有抵抗力能把更重要的长期信息稳定写入记忆在需要回溯早期内容的检索和推理任务上表现更好。Q2MDN训练速度比其他线性注意力模型慢还值得用吗AMDN的训练吞吐量确实略低于高度优化的GDN和Comba但已与Mamba2和KDA持平。而MDN的推理生成文字速度与GDN、Comba几乎相同远快于Transformer。考虑到MDN在检索、推理和长上下文理解方面的显著提升这个训练效率的小幅代价对于追求更强记忆能力的应用场景是值得的。Q3MDN的稳定性约束为什么必要不加会怎样A动量机制在数学上会引入复数特征值如果特征值落在复平面左半部分记忆本会出现正负震荡的符号翻转现象导致快速记忆和动量记忆之间的协作被破坏。实验验证了不加约束的后果非常直接去掉αmax约束会让训练在第一步就崩溃出现NaN去掉μlog下界约束则会在训练约70步后崩溃。稳定性约束是MDN能够正常训练的基础保障。