译者 | 朱先忠
审校 | 重楼
本文要讨论什么内容?
Meta公司努力开发出的CoCoMix(Continuous Concept Mixing:连续概念混合,出自Jihoon等人2025年发表的论文,见【引文1】)框架实现了概念学习,即学习单词背后的概念而不是仅仅预测下一个标记,从而使其具有极强的可操控性和可解释性。
但是,依然存在一个核心问题:即使是概念上非常出色的模型,在训练之后的实际部署中,也可能难以应对细微差别或事实性的回忆挑战。你可以问一个看似简单的问题,比如:“在我们之前那场涉及200万个标记的对话中,我们在哪里讨论过匹诺曹那出了名的不断增长的鼻子?”无论LLM的概念能力多么强,如果答案超出了它的上下文范围,它就无法回答这个简单的问题。
那么问题来了:我们能否在推理过程中为这些智能LLM配备适应性强的“记忆”或性能提升呢?
1. 基础问题:Transformer
Transformer(出自Vaswani等人2017年发表的论文,见【引文2】)在现代人工智能领域已变得无处不在。自其取得突破性成功以来,它已成为各领域的首选架构。
回想2020年,人们对任何机器学习问题的默认反应往往是“把注意力集中到它上面”——令人惊讶的是,它确实有效,而且通常表现优于最先进的模型。视觉任务?请使用Transformer(Dosovitskiy等人2020年发表的论文,见【引文3】)。时间序列预测?又是Transformer(Zerveas等人于2021年发表的论文,见【引文4】)。自然语言处理?Transformer几乎定义了它(Rogers等人于2021年发表的论文,见【引文5】)。
但是,随着我们对大型模型的依赖加深以及计算预算的扩大,即使是这种“全能”架构也开始显示出其局限性——因此,人们开始努力进一步扩展其功能。
存在瓶颈?可以求助于注意力机制的“人人对话”方法。这种方法非常巧妙,但成本却极其高昂——想象一下,一个房间里有一百万人,每个人都必须记住与所有人的每一次对话。这限制了Transformer的“工作记忆”——使其记忆变得狭窄,难以进行理解海量文档所需的“长期回忆”,因为早期信息会逐渐消失。
除了上下文限制之外,普通的Transformer还面临另一个根本障碍:训练后缺乏适应性。虽然它们擅长运用海量预训练知识来预测下一个标记(这是一个复杂的推理和预测过程),但这与真正的学习并不相同。就像谷歌地图一样,虽然它会为你找到“最短路径”,但它会忘记前方有施工路段,让你冲破路障。而人类向导则会为你指引一条替代的小巷路线。
这种无法从当前正在处理的数据中“即时学习”的能力,对于需要不断适应或记忆训练集之外的新经验的任务来说,是一个严重的限制。
下面这张图解释了传统Transformer注意力机制的局限性。其中,顶部展示了一个小型且有限的“注意力窗口”,其中包含t1-t6个标记,表明像t203这样的旧标记已被“遗忘”。底部则展示了18个标记的全对全注意力机制,突出显示了其“O(n²)”的二次成本、计算消耗和资源密集度,并得出结论:它“无法扩展到百万级以上的上下文窗口”。
当前原始Transformer的众多问题中的两个(作者本人绘制)
解决方案?Titans框架!
研究人员没有仅仅针对单一限制,而是从更广阔的视角出发:像人脑这样的智能系统如何管理记忆并适应新情况?这并非关乎拥有一个庞大且随时可用的记忆体。它其实是一个更灵活的设置,其中不同的组件相互协调,以处理不同类型的信息和经验。
Meta公司研究人员新研发成功的Titans架构(出自Behrouz等人于2025年发表的论文,见【引文6】))就包含了这一点,它不是围绕单一的整体注意力模块构建的,而是围绕一个专门的记忆系统合作团队构建的,每个系统在理解和响应手头的任务方面都发挥着至关重要的作用。
(1)架构组件:内存模块
- 短期记忆(STM):这是一种敏锐、注重细节的“专家”。它的功能类似于你所知的注意力,但它不会被过去的一切(现在是LMM的工作)所淹没,而是专注于当下。这就像你记住了对方刚刚对你说的话,时间刚好够你回应。
- 长期记忆模块(LMM):这是最令人兴奋的新增功能。它旨在在推理过程中学习和适应——没错,就是在推理过程中,而且是即时方式!我所说的“适应”字面意思是它的参数会发生变化!想象一下,你多年来一直在理解一位朋友——不断积累经验,同时过滤掉不重要的事件。
- 持久记忆(PM):此模块保存着基础的、特定于任务的知识。这些是模型在主要训练过程中获得的可学习的基本见解。这些知识并非即时动态的,但为其他两个模块提供了必要的基础和背景。它就像你的个性、你的举止、走路或开车的能力,这些你不需要重新学习或改变。
三种记忆模块:短期记忆(STM)、长期记忆模块(LMM)和持久记忆(PM)
这张图展示了三种记忆模块:短期记忆,表现为一个在“STM/Attention”笔记本电脑前感到压力的人物,专注于当前情境;长期记忆,表现为一个在“LTM weights”笔记本电脑前面带微笑的人物,用羽毛笔更新自身,记录历史情境;持久记忆,表现为一个平静的人物,手持石碑,石碑上写着“预设相同权重”,体现了固定的、与数据无关的任务知识(作者本人绘制此图)。
(2)这些内存模块是如何实现的?
那么,这三者究竟是如何协同工作的呢?首先,STM本质上是标准的自注意力计算,它是原生Transformer的核心。它的“记忆”是训练过程中学习到的键值缓存和注意力矩阵。
另一方面,PM是一组可学习的参数,它们被添加到输入序列的前面,并在训练过程中学习,并充当模型在推理过程中无论如何都要遵循的“圣杯”。
到目前为止还算容易理解,嗯?接下来,让我们深入探讨其中的创新之处和真正令人兴奋的部分。虽然它是作为一个简单的MLP网络实现的,但它可以在测试时进行调整——LMM模块:
(3)Titans框架核心:自适应长期记忆(LMM)模块
等一下……测试时更新参数?这难道不是只在训练时才会做的事吗?这难道不是作弊吗?
当你听到“测试时训练”这个词时,你是否想到了这些问题?这些问题确实存在,但这并不意味着作弊。Titans利用在线学习和元学习的原理,实现快速、局部的更新,这些更新专门针对记忆,而非一般的任务改进。它在测试时不会查看外部标签来计算梯度和优化参数;相反,一切都保持独立:模型仅使用已知信息和当前感知信息进行内部调整。
在人类记忆中,常规和可预测的事件往往会逐渐消失,而意外或令人惊讶的时刻往往会持续存在(来自Mandler于2014年发表的论文,见【引文7】)。这是实现动态测试时间更新背后的核心思想。
LMM如何学习:关联损失函数
LMM充当联想记忆:它学习将“键”(线索)与“值”(信息)联系起来。对于每条新数据xt(MAG和MAL中的输入块,MAC中的STM(自注意力机制)输出):
- 键值提取:系统首先使用可学习的变换(Wk和Wv)将xt转换为特定键(kt)和关联值(vt)。
使用线性层将xt映射到kt和vt(作者本人绘制)
- 测试LMM:在当前状态下,对LMM进行“询问”:给定这个新密钥kt,你会预测什么值?我们将其预测称为pt。
Mt-1:当前LMM状态;kt:当前块的密钥(作者本人绘制)
- 计算损失:通过LMM预测的错误程度来衡量:
预测输出与“基本事实”之间的标准MSE损失(作者本人绘制)
梯度和“意外”信号
为了让LMM从这种损失中学习,我们加入了“惊喜信号”,它能够衡量模型在看到基本事实(vt)时的“惊讶”程度。这个“惊喜”在数学上定义为损失函数相对于LMM参数的梯度。
“意外”的测量,即模型与预测“正确的”vt的距离(作者本人绘制)
考虑到LMM的当前知识,较大的梯度意味着xt是非常“令人惊讶”或出乎意料的。
基本学习步骤:
LMM学习的最简单方法是稍微调整其参数以减少这种意外(即减少损失),就像梯度下降中的一步一样:
Mt:更新后的LMM参数;Mt-1:之前的LMM参数;lr:学习率(作者本人绘制)
细化惊喜:利用动量和遗忘实现更智能的学习
仅仅对眼前的“惊喜”做出反应是不够的。良好的记忆力不仅需要洞察趋势,还必须知道何时放弃旧的、不相关的信息。
智能学习方向(ΔΘMt):首先,LMM计算出调整其参数的最佳方向。这不仅基于当前的意外,还基于对近期意外的“记忆”。
参数变化是根据之前的变化和当前的意外情况计算得出的(作者本人绘制)
- ΔΘMt:LMM参数的建议变更。
- ηt*ΔΘMt-1:这是动量 ——它延续了上一步的学习趋势。ηt(数据相关)决定了过去动量的持续时间。
- θt*∇Loss_current_surprise:这是当前意外的影响。θt(数据相关)衡量其影响力。最终参数更新(ΘMt):然后,LMM更新其实际参数,将旧知识与新的学习方向相结合,并且至关重要的是,允许“遗忘”。
最终更新包括更新多少以及保留多少(作者本人绘制)
- ΘMt:从xt学习后的LMM的新参数。
- (1—at)*ΘMt-1:这是保留了多少旧LMM状态。at(依赖于数据,介于0和1之间)是遗忘因子 -如果at很高,则会遗忘更多的旧状态。
- ΔΘMt:上面计算的智能学习方向。
下图展示了LTM模块的更新过程。分块输入序列(例如STM输出)被投影到键向量和值向量中。键向量在LTM模块中进行前向传递,并与值向量一起计算损失。该损失的梯度(通过不带更新的后向传递)通过加权和与动量缓冲区中存储的先前更新相结合。组合后的更新经过“遗忘”门,该门确定LTM的新权重。
整个LMM更新过程的可视化(作者本人绘制)
简而言之:
LMM会查看当前数据的“意外情况”(∇Loss_current_surprise),将其与近期的学习趋势(momentumΔΘMt-1)相融合,然后更新其内部知识(ΘMt),并决定在此过程中保留或丢弃多少旧信息(at)。数据依赖门控(ηt,θt,at)使其能够动态自适应。
(4)Titans的架构蓝图:记忆在行动
谷歌研究人员探索了这三个内存模块的三种主要排列方式:
记忆作为情境(MAC)
在这种设置中,Titans为STM(标准自注意力块)创建了一个增强且更丰富的环境。
- 非常长的输入序列被分成多个段或块。
- 在处理时,模型将块映射到查询,并使用它来从LMM中检索相关的历史上下文(查询经过LMM,其输出是历史上下文)。
- 然后将这些检索到的历史标记与静态持久内存标记和当前段标记连接起来。
- 整个扩展序列(持久+历史+当前)被输入到STM(注意力)层,该层处理这个大型组合上下文中的关系。
- 注意力层的输出反映了对当前数据考虑到历史和任务知识的深刻理解,然后将其用作LMM动态参数更新过程的输入数据。
- 使用相同的注意力结果再次查询更新后的LMM,然后通过门控总和或部分连接将其响应与原始注意力相结合以产生最终结果。
类比:文本(序列)以页面(块)的形式呈现。对于每一页,不断学习的笔记记录者(LMM)会快速从过去的笔记中找到相关的摘要,并将其与重要的“规则手册”笔记(PM)进行混合。学生(STM/注意力)会阅读整篇文章——规则手册、相关的过去摘要以及当前页面——并根据从丰富的上下文中学习到的知识,准确地告诉笔记记录者当前页面上哪些要点对于将来的总结至关重要。
最终答案的形成是基于学生的详细阅读和笔记记录者更新后的记忆视角的综合考虑。
下图展示的是记忆即语境(MAC)架构图。分块输入向LMM查询历史语境。结合持久记忆,这为注意力模块(STM)构建了一个“更丰富的语境”。注意力模块的输出更新LMM,并与更新后的LMM输出相结合,形成最终结果。
MAC架构可视化(作者本人绘制)
记忆之门(MAG)
该设计以不同的方式并行集成存储器,通过动态门混合来自STM和LMM输出的信息。
- 输入序列(持久序列+当前序列)被输入到两条路径:一条路径通往STM,此处使用滑动窗口注意力机制实现,使其能够聚焦于每个符号周围的局部邻域。另一条并行路径通往LMM网络。
- LMM的输入用于动态更新其参数,并生成相应的输出表示。
- 然后,STM(滑动窗口注意力机制)的输出和LMM前向传递的输出会通过一个门控机制。该门控机制会学习在局部注意力机制的信息和LMM的状态/输出所代表的信息之间进行混合或选择。
- 该门控机制的输出是该令牌位置的最终输出。
打个比方:每收到一页纸,两个任务就会同时进行:学生(STM)专注于眼前的那一页,而独立的笔记记录者(LMM)则通读当前页面,并将其与目前为止所见的内容联系起来,默默地更新其整体理解。学生根据段落给出答案,而笔记记录者则根据其不断发展的宏观框架提出想法。
协调者(门控机制)只需将两个输出融合或从中挑选,即可得出最终答案。
以下给出内存即门控(MAG)架构示意图。分块输入序列与持久内存令牌组合。该增强序列同时查询LMM模块(也用于更新)和注意力模块(STM)。然后,LMM分支和STM分支的输出通过门控和进行组合,以生成最终输出。
MAG架构可视化(作者本人绘制)
记忆层(MAL)
在这个变体中,LMM网络被用作初始处理层,在序列到达注意力机制(STM)之前对其进行修改。
- 输入序列(仍然使用“持久+当前”方式)首先作为不同的层通过LMM网络。
- 当LMM网络逐块处理序列时,其参数会动态更新。它还会生成一个变换后的输出序列。
- 然后,将LMM层的转换输出序列用作后续STM(注意力)层(滑动窗口或窗口内完全注意力)的输入。
- Attention层的输出是该序列的模型的最终输出。
打个比方:首先,每一页新内容都会直接发送给主笔记员(LMM),LMM会处理所有内容,并不断总结和更新其总结风格。之后,这份(可能细节较少的)总结会被交给学生(STM)。学生只会看到并关注总结文本的局部内容,他们的答案完全基于主笔记员提供的内容。
下图展示的是记忆即层(MAL)架构图。一个以持久记忆标记开头的分块输入序列被送入LMM模块进行查询和更新。LMM的输出随后作为注意力模块(STM)的输入(查询),STM产生最终输出。
MAL架构可视化(作者本人绘制)
3. 我们从中获得了什么?结果和发现
那么,现在我们对Transformer之后的下一个革命性技术已经了如指掌,但它真的会那么伟大吗?谷歌的研究人员真的破解了能够记忆、适应并克服此前被认为不可能的挑战的模型密码吗?让我们逐一回顾一下这些新发现:
语言能力:不仅仅是文字
Titans框架的功能远不止于更准确地预测下一个单词。得益于其动态长期记忆模块(LMM),它展现出对语言和语境更深入、更直观的理解。与Transformer++等强大的基准模型以及一些最新的循环模型相比,Titans框架的表现始终优于它们,不仅在语言建模方面,在常识推理任务上也同样如此。
(来源:改编自Behrouz等人于2025年发表的论文,表1)
Titans在常识和推理任务上的表现(混合情况下使用:MAC、MAG、MAL;简单情况下使用:LMM)
大海捞针挑战
Titans框架的设计在S-NIAH任务中展现出卓越的性能连续性,这与RULER基准测试(Hsieh等人于2024年发表的论文,见【引文8】)相符,该基准测试旨在评估有效的上下文长度。Titans模型(包括独立的神经记忆模型LMM)即使在16K个标记的情况下也能保持强劲的检索率,而一些最先进的循环模型的准确率则随着序列长度的增加而急剧下降。
(来源:Behrouz等人于2025年发表的论文,表2)
Titans(混合情况下使用:MAC、MAG、MAL;简单情况下使用:LMM)在RULER的S-NIAH任务上的表现(出自Hsieh等人于2024年发表的论文,见【引文8】)
在BABILong中掌握复杂推理
检索事实是一回事。但如何运用多个事实,并跨越海量上下文进行推理呢?这才是真正的考验,也正是BABILong基准测试(来自YuryKuratov等人于2024年发表的论文,见【引文9】)的要求。Titans(尤其是MAC架构)不仅表现出色,而且超越了所有模型。即使是像GPT-4和Llama3.1-70B这样的大型模型,即使是那些能够访问外部工具或检索系统的模型,Titans最大的模型参数量也高达7.6亿!
除此之外,Titans(MAC混合架构)即使在1000万个标记的情况下也能达到70%的准确率。换个角度来看,这就像在整个《哈利·波特》系列中导航和寻找拼图碎片……甚至还多10倍。
(来源:Behrouz等人于2025年发表的论文,图6)
BABILong上不同LLM的准确率与序列长度图(出自Yury Kuratov等人于2024年发表的论文,见【引文9】)
内存深度与速度
研究人员探索了通过堆叠更多层来加深长期记忆模块(LMM)时会发生什么。结果如何?更深层的LMM显著提升了其存储和组织重要信息的能力,使其更不容易忘记关键细节,尤其是在大多数模型难以保持上下文的长序列中。
虽然LMM本身能够获得线性时间复杂度,从而高效处理大量输入,但更深的LMM确实需要付出一些代价:吞吐量降低,或者每秒处理的令牌更少。
下面的折线图展示了不同深度(L_M=1、2、3、4)的LMM模型的训练吞吐量(103标记/秒)与序列长度的关系。所有LMM变体均表现出几乎恒定的吞吐量,无论序列长度如何,这表明吞吐量呈线性增长。然而,较深的LMM(L_M=3和L_M=4)的吞吐量逐渐低于较浅的LMM(L_M=1和L_M=2),这表明随着内存深度的增加,效率会有所降低。
(来源:Behrouz等人于2025年发表的论文,图8)
不同LMM深度的序列长度与吞吐量
超越语言任务
另一个令人兴奋的事实是,同样的记忆机制在传统语言任务之外也能发挥作用。在时间序列预测(一个以混乱、变化的模式而闻名的领域)中,长期记忆模块(LMM)的表现足以匹敌高度专业化的模型,包括那些基于Mamba(之前的SOTA)的模型。
在DNA建模这项完全不同的任务中,该架构表现出了强大的效果。这种通用性实属不易,这表明,如果处理得当,记忆不仅有用,而且是跨领域的基础。
(来源:改编自Behrouz等人于2025年发表的论文,表3)
神经记忆(以LMM为模型)在各种时间序列数据集上的表现
(来源:Behrouz等人于2025年发表的论文,表4)
神经记忆模块(LMM作为模型)在基因组基准上的表现(Grešová等人于2023年发表的论文,见【引文10】)
4. 结论和最终想法
本次对Titans的深入研究就到此结束。探索这种架构真的非常有趣——看到研究超越了规模化,深入探究记忆和学习如何以更具适应性、更像人类的方式运作,令人耳目一新。
谷歌的基础性工作传承在此延续,从发明Transformer到现在重新思考AI如何在推理过程中学习。Titans仿佛是这种精神的自然演进。
话虽如此,如今的AI领域比2017年更加拥挤了。无论多么精彩的新想法,要成为主流都面临着更加艰难的道路。性能只是其中之一——效率、简洁性和社区影响力比以往任何时候都更加重要。
尽管如此,Titans有力地预示着未来模型将不再仅仅基于已知知识进行思考,而是能够真正地在实践中适应。无论这是否会成为下一个“只需关注”的时刻,这都是迈向更智能、更智慧的AI的充满希望的一步。
参考文献
【1】Tack,Jihoon等人,“使用连续概念进行LLM预训练。”(2025)arXiv预印本 arXiv:2502.08524。
【2】Vaswani,Ashish等人,“你只需要注意力。”(2017),神经信息处理系统的进展30。
【3】Dosovitskiy,Alexey等人,“一张图片胜过16×16个单词:用于大规模图像识别的Transformers。”(2020),arXiv预印本 arXiv:2010.11929。
【4】Zerveas,George等人,“基于Transformer的多元时间序列表示学习框架。”(2021),第27届ACM SIGKDD知识发现与数据挖掘会议论文集。
【5】Rogers,Anna等人,“BERTology入门:我们对BERT工作原理的了解。”(2021年),计算语言学协会汇刊8:842–866。
【6】Behrouz,Ali、Peilin Zhong和Vahab Mirrokni。“Titans:学习在考试时记忆。”(2024年),arXiv预印本 arXiv:2501.00663。
【7】Mandler,George。“情感与认知”(2014年)。心理学出版社,3–36。
【8】Hsieh,Cheng-Ping等人,“RULER:长上下文语言模型的真实上下文大小是多少?”,载于:第一届语言建模会议。2024年。
【9】Kuratov,Yury等人。“Babilong:用大海捞针的长上下文推理测试LLMS的极限。”(2024),神经信息处理系统进展,37:106519–106554。
【10】Grešová,Katarína等人,“基因组基准:基因组序列分类数据集集合。”(2023)BMC基因组数据,24.1:25。
译者介绍
朱先忠,51CTO社区编辑,51CTO专家博客、讲师,潍坊一所高校计算机教师,自由编程界老兵一枚。
原文标题:Can AI Truly Develop a Memory That Adapts Like Ours?,作者:Moulik Gupta