谷歌新架构终结Transformer，长序列处理王者诞生？清华姚班校友新作

新智元报道
编辑：泽正 英智
【新智元导读】谷歌团队提出的Titans架构通过引入神经长期记忆模块，突破了传统Transformer架构在长序列处理中的局限。该架构通过创新的记忆整合和遗忘机制，在语言建模、常识推理、时间序列预测等任务中展现了显著的性能提升，在长上下文任务中的优势突出。
Transformer后继者终于现身！
2017年，Attention Is All You Need首次引入注意力机制，成为现代LLM诞生标志。
虽然Transformer依然大行其道，但其架构的缺陷却饱受诟病，尤其是无法扩展更长上下文。
正所谓，「风浪越大，鱼越贵！」
近日，谷歌研究团队迎难而上，提出支持200K上下文处理窗口的新架构——Titans。
最重要的是，轻松扩展到2M上下文的Titans架构，要比Transformer和线性RNN更加有效。


论文链接：https://arxiv.org/abs/2501.00663


Titans是什么
研究者认为大多数现有架构将记忆视为由输入引起的神经更新，并将学习定义为在给定目标的情况下有效获取有用记忆的过程。
从这个角度来看，循环神经网络（RNN）可以被定义为具有向量值记忆模块ℳ（也称为隐藏状态）的模型，其主要步骤包括：在时间t给定新输入时，
（1）使用函数更新记忆（带有压缩）；
（2）使用函数检索输入的相应记忆。
类似地，Transformer可以被视为具有不断增长的记忆和两个相似步骤的架构。即：
（1）通过将键和值附加到记忆中来更新记忆（无压缩）；
（2）通过查找查询向量与键向量的相似性来检索查询向量的相应记忆，然后将其用于加权值向量以生成输出。
由于记忆分为短期记忆、工作记忆和长期记忆，而其中每个部分都相互独立地服务于不同的场景，也具有不同的神经结构。
受此启发，研究者提出了两个问题：
1. 如何设计一个高效架构，将不同且相互关联的记忆模块整合起来？
2. 是否需要一个深度记忆模块，以有效存储和记住长期历史信息？
本研究旨在通过设计一个长期神经记忆模块来解决上述问题，神经长期记忆模块的设计受到人类长期记忆系统的启发，能存储和检索过去的信息。
该模块不是无差别地记住所有信息，而是会通过「惊讶度」来选择性地记住那些重要或令人惊讶的信息。
并且其记忆不是静态的，可以根据新的信息动态更新。这种动态更新机制类似于人类的学习过程，使得模型能够不断适应新的数据和任务需求。
为了更好地管理有限的内存，模块引入了衰减机制。该机制根据记忆的大小和数据的惊讶程度来调整记忆的权重，从而优化内存管理。
长期神经记忆模块设计完成后，面临的一个关键问题是如何把记忆高效地整合进深度学习架构。研究者提出了Titans架构，由三个模块构成：
1. 核心模块（Core）：包含短期记忆，负责主要的数据处理流程，采用具有有限窗口大小的注意力机制。
2. 长期记忆模块（Long-term Memory）：此模块是研究者设计的神经长期记忆模块，负责存储和记住远距离的历史信息。
3.持久记忆模块（Persistent Memory）：这是一组可学习但与数据无关的参数，主要用于对任务知识进行编码，为模型提供先验知识储备。
在此基础上，研究者提出了Titans架构的三种变体：
记忆作为上下文（MAC）架构
核心分支把对应的长期记忆、持久记忆和当前输入信息拼接在一起，然后用注意力机制来处理上下文，并决定哪些信息应存储在长期记忆中。
在测试时，与上下文记忆对应的参数仍在学习，与核心分支对应的参数负责上下文学习，而持久记忆的参数则负责存储任务相关知识，因此是固定的。


记忆作为门控（MAG）架构
在此架构中，一个分支用输入数据更新长期记忆，另一个分支使用滑动窗口注意力（SWA），最后将两者结果通过门控机制组合。
在此设计中，滑动窗口注意力充当精确的短期记忆，而神经记忆模块则作为模型的衰减记忆。这种架构设计也可视为一种多头架构，其中头的结构各不相同。
与MAC架构不同的是，MAG架构仅将持久记忆融入上下文，并通过门控机制将记忆与核心分支结合。门控机制决定了来自持久记忆的信息在多大程度上影响核心分支的处理结果。


记忆作为层（MAL）架构
将神经记忆模块作为深度神经网络的一层，结合滑动窗口注意力机制。记忆层的核心功能是对过去和当前的上下文信息进行压缩处理，之后将处理结果传递给注意力模块。


在测试时去学习记忆
神经长期记忆模块
对于神经网络来讲，记忆能力通常反而会限制模型的泛化能力，并可能引发隐私问题，导致在测试时性能下降。
此外，由于测试数据可能属于分布外数据，训练数据的记忆在测试时可能就并没有什么效用。
因此，研究者认为，训练长期记忆的关键思想是将其训练视为一个在线学习问题，学会在测试时如何记住或忘记数据。在这种设置中，模型学习的是一个能够记忆的函数，但不会过拟合训练数据，从而在测试时实现更好的泛化。
学习过程与目标函数：由于令人惊讶的事件对人类来说更易记住。受此启发，作者将「惊讶度」定义为相对于输入的梯度。梯度越大，输入数据与过去数据的差异就越大。因此，利用这个惊讶度，我们可以更新记忆如下：


如此就能将过去的信息压缩到长期神经记忆模块的参数中。
然而，这种惊讶度的度量方法也可能会导致错过一些重要信息。也就是说，在若干个惊讶步骤之后，梯度可能变得非常小，从而导致陷入平坦区域（即局部最小值），错失序列中的某些信息。
而从人类记忆的角度来看，某个事件虽然值得记住，但可能不会在很长时间内一直让我们感到惊讶。
因此为了改进上述惊讶度度量，作者将惊讶度度量分为两部分：（1）过去的惊讶，衡量最近过去的惊讶度；（2）瞬时惊讶，衡量即将到来的数据的惊讶度。


在此公式中，项是数据依赖的惊喜衰减，控制着惊喜随时间如何衰减；而项 θt 则控制着应以数据依赖的方式将多少瞬时惊喜纳入最终的惊喜度量中。
这种数据依赖性在此设计中尤为重要：虽然前一个标记的惊喜可能影响下一个标记的惊喜，但这主要在所有标记相关且处于同一上下文时才有效。
因此，数据依赖的η可以控制记忆是否需要：
（1）通过设置→0忽略上一次的惊喜（可能由于上下文的变化）
（2）通过设置→1完全纳入上一次的惊喜（可能因为该标记与其最近的过去标记高度相关）。
在本工作中，作者专注于关联记忆，旨在将过去的数据存储为键值对。即给定，类似于Transformers，使用两个线性层将其投影为键和值：


接着，作者希望记忆模块能够学习键和值之间的关联。为此，定义损失函数如下：


通过在元模型的内循环中优化上述损失函数，模型就可以学会如何在测试时记忆键与值之间的映射。
遗忘机制：在处理非常大的序列（例如，数百万个标记）时，明确哪些过去信息应该被遗忘至关重要。为此，作者使用了一种自适应遗忘机制，允许内存遗忘不再需要的信息，从而更好地管理内存的有限容量。也就是说，给定下一个标记，然后将更新规则修改为：


其中是灵活控制记忆的门控机制；即决定应遗忘多少信息。例如，它可以通过让来更新记忆而不影响过去的抽象，并可以通过让来清除整个记忆。
检索记忆：作者简单地使用不更新权重的前向传递（即推理）来检索与查询对应的记忆。形式上，给定输入，使用线性层 WQ 来投影输入，即，并通过以下方式从记忆中检索相应（或有用的）信息：


实验结果
在语言建模及常识推理任务中，对340M、400M、760M等不同参数规模下的Titans变体与多种基线模型进行对比。非混合模型里，Titans (LMM) 在困惑度和准确率上表现优异。
混合模型对比中，Titans的三个变体均比基线模型更好。MAC和MAG整体性能高于MAL，能更好地整合注意力和记忆模块。






在S-NIAH任务里，基于RULER基准测试，对2K、4K、8K 和 16K长度序列予以评估。神经记忆模块相较基线模型优势显著。在Titans变体中，MAC性能最佳。


在BABILong基准测试中，Titans (MAC) 展现了卓越的性能，能够有效扩展到超过200万的上下文窗口，超越了GPT-4、Llama3+RAG和Llama3-70B等大模型。
Titans (MAC) 的参数量远少于基线模型，展现出在长序列推理方面的高效性和强大能力。在微调设置环节，Titans（MAC）的表现更为出色。


研究发现，增加记忆深度可提升模型在较长序列上的性能，并改善困惑度，但训练速度会因此降低，呈现出性能与效率之间的权衡。


通过在Simba框架中替换Mamba模块，并在ETT、ECL、Traffic和Weather等基准数据集上测试，神经记忆模块超越了所有的基线模型。这表明其在处理时间序列任务中的潜在优势。


在DNA建模任务中，Titans架构也展示了其强大的长序列处理能力。实验结果表明，Titans架构在这些任务中能够有效地利用历史信息，从而提高模型的性能。


消融研究表明，神经记忆模块的所有组件对模型性能均有积极贡献，特别是权重衰减和动量。MAC和MAG在语言建模和常识推理上表现相近，但MAC在长上下文任务中表现最佳。


Titans架构通过引入神经长期记忆模块，显著提高了模型在处理长序列数据时的性能和效率。
实验结果表明，Titans架构在语言建模、常识推理、时间序列预测和DNA建模等任务中均表现出色，特别是在处理超200万上下文窗口任务中，能够有效地利用历史信息，提高模型的准确性。
作者介绍
Ali Behrouz


Ali Behrouz目前是康奈尔大学计算机科学系的二年级博士生，同时也是Google Research的研究实习生。在加入康奈尔之前，他在哥伦比亚大学师从Margo Seltzer教授，取得计算机科学硕士学位。
他对多样化的研究主题感兴趣，目前正致力于深度学习架构、图表示学习、医疗保健中的机器学习以及计算神经科学等领域的研究。
Peilin Zhong


Peilin Zhong现为谷歌纽约市算法与优化团队的研究科学家，该团队由Vahab Mirrokni领导。他在哥伦比亚大学获得博士学位（导师为Alex Andoni、Cliff Stein和Mihalis Yannakakis）。此前，他是清华大学交叉信息研究院（姚班）的本科生。2016年，他以第一作者发表的论文被顶会STOC 2016接收，创下首次有中国本科生在STOC上发表一作论文的记录。
他对理论计算机科学有广泛的兴趣，主要集中在算法的设计与分析上。一些特别感兴趣的领域包括并行和大规模并行算法、隐私算法、草图算法、流算法、图算法、机器学习、高维几何、度量嵌入、数值线性代数、聚类以及其他与大规模数据计算相关的算法。
Vahab Mirrokni


Vahab Mirrokni在纽约的谷歌研究院领导算法与优化团队。团队包括市场算法、图挖掘和大规模优化小组。此外，他还在纽约大学库朗研究所担任兼职副教授，讲授互联网算法与经济学。
参考资料：
https://arxiv.org/abs/2501.00663
https://x.com/behrouz_ali/status/1878859086227255347