AI领域迎来新突破：谷歌Titans架构横空出世
在Transformer架构主导AI领域八年之后，谷歌正式发布了其继任者——Titans。这一全新架构的核心在于将测试时计算应用于记忆层面，旨在解决大型语言模型（LLM）在处理长上下文时遇到的瓶颈。
论文一作Ali Behrouz指出，尽管注意力机制在LLM发展中至关重要，但其在处理长上下文方面存在局限性。Titans架构巧妙地将注意力机制与元上下文记忆相结合，使得模型能够在测试时学习记忆，从而将上下文窗口扩展到惊人的200万tokens。
图源：https://x.com/behrouz_ali/status/1878859086227255347
这意味着，谷歌Transformer迎来了它的「继任者」。
图源：https://x.com/mark_k/status/1878896628654022993
研究人员长期以来都在探索循环模型和注意力机制的有效结合。循环模型将数据压缩到固定大小的记忆中，而注意力机制则可以处理整个上下文窗口。然而，更准确的依赖建模往往伴随二次成本，导致模型只能处理有限的上下文长度。
为了解决这一问题，谷歌提出了长期神经记忆模块。该模块能够学习记忆历史上下文，并辅助注意力机制处理当前上下文。这种设计不仅具有快速并行化训练的优势，还能保持高效的推理速度。
谷歌认为，注意力机制虽然受限于上下文，但能更准确地建模依赖关系，可作为短期记忆；而神经记忆则能持久存储数据，发挥长期记忆的作用。基于这两个模块，谷歌推出了Titans架构，并将其分为三种变体：记忆作为上下文（MAC）、记忆作为门（MAG）和记忆作为层（MAL）。
实验结果表明，Titans在语言建模、常识推理、基因组学和时序预测任务中的表现均优于Transformer和现代线性循环模型。尤其在“大海捞针”测试中，Titans架构能够有效处理超过200万tokens的上下文，且准确性显著高于基准模型。
论文地址：https://arxiv.org/pdf/2501.00663v1
值得一提的是，论文作者之一Peilin Zhong毕业于清华姚班，并在哥伦比亚大学获得博士学位，现为谷歌研究科学家。
目前，Titans架构的非官方实现已在GitHub上开源。
GitHub 地址：https://github.com/lucidrains/titans-pytorch
学习测试时记忆：Titans的核心机制
Titans架构的核心在于其长期神经记忆模块，这是一种能够在测试时学习和记忆的元模型。
长期记忆
为了实现长期记忆，模型需要将历史信息编码到参数中。传统的神经网络记忆机制存在泛化能力受限、隐私问题等缺陷。因此，谷歌设计了一种在线元模型，使其能在测试时学习记忆或遗忘数据，从而提升泛化性能。
学习过程与意外指标
训练长期记忆的关键是将训练视为在线学习问题。模型需要将过去信息压缩到长期神经记忆模块中。受到人类记忆的启发，谷歌定义了模型的“意外”程度——即其相对于输入的梯度。梯度越大，表明输入数据与过去数据的偏差越大，因而更值得记忆。基于此，记忆更新规则如下：
为了解决重要信息因“意外”度持续降低而被遗忘的问题，谷歌将“意外”分解为两个指标：过去意外和瞬时意外，分别衡量最近过去的意外程度和传入数据的意外程度。
这些意外指标基于损失函数：
记忆模块基于该损失函数学习一个记忆函数，实现测试时的泛化。
在本文中，谷歌侧重于联想记忆，将过去数据存储为键值对。类似于Transformer，谷歌使用线性层将输入投影为键和值：
记忆模块学习键和值之间的关联，损失定义如下：
遗忘机制
为处理超长序列，Titans引入了自适应遗忘机制，使得记忆可以遗忘不再需要的信息，有效管理有限的记忆容量。更新规则如下：
记忆架构
谷歌选择了具有L_M≥1层的简单MLP作为长期记忆架构，希望借此鼓励长期记忆的设计与架构融合。本文的架构为更有效高效的数据记忆神经架构开辟了新的研究方向。
检索记忆
Titans通过前向传递（推理）来检索与查询对应的记忆。给定输入 x_t，模型使用线性层W_Q将其投影为查询q_t，并通过以下公式检索相应信息：
并行化长期记忆训练
长期记忆模块的训练在理论上需要庞大的计算量，因此需要并行化训练过程。通过将序列拆分成大小为b的块，使用小批量梯度下降，可以将过程转化为矩阵乘法：
当M_t = W_t为线性时，损失函数如下：
此外，该方法还可扩展以合并动量项。在动量的块式梯度下降中，动量项可以表示如下：
作为函数块的参数
为提升训练速度，谷歌没有让参数a_t、θ_t和η_t依赖于输入，而是让它们成为函数块。在实验中，谷歌在每个块中使用了相同的值，并表示，这种简化可能是未来工作的一个方向。
下图展示了如何并行并在使用矩阵乘法时完成神经记忆训练。
Titans架构的记忆融合
如何有效地将神经记忆融合到深度学习架构中是Titans面临的重要问题。谷歌认为，Transformer的K和V矩阵可以解释为联想记忆块，作为短期记忆处理当前上下文窗口。而神经记忆则能从数据中持续学习，作为长期记忆发挥作用。Titans通过三个变体来实现记忆的有效融合。
记忆作为上下文(MAC)
MAC变体将记忆视为当前信息的上下文。将长序列分割为固定大小的片段，使用过去的片段作为历史信息。模型使用当前上下文作为对记忆的查询，检索相应的历史信息。检索的信息与持久记忆参数作为注意力模块的输入：
给定一个长序列：
检索与S^(𝑡)对应的历史信息：
将历史信息与持久记忆参数一起作为注意力模块的输入：
MAC的优点在于，注意力模块可同时利用历史和当前上下文，并使长期记忆模块仅存储来自当前上下文的有用信息。
记忆作为门(MAG)
MAG变体使用滑动窗口注意力(SWA)和长期记忆，其中SWA充当短期记忆，而神经记忆则作为模型的衰减记忆。该架构可以视为多头架构的不同实现。
其中一个分支直接使用输入更新长期记忆，另一个分支使用SWA：
整体注意力掩码如下：
记忆作为层(MAL)
MAL变体使用深度神经网络，其中混合模型堆叠具有完整或滑动窗口注意力的循环模型。
给定输入 x，可得：
无注意力记忆
Titans还测试了仅使用长期记忆模块（LMM）的情况，以验证记忆系统的独立性。即使没有注意力机制，长期记忆模块仍然应具有较强的性能。
架构细节
Titans的所有模块都使用了残差连接，使用SiLU(.)激活函数，并对查询和键进行归一化。
在查询、键和值投影后，都融合了一个1D深度可分离卷积层，以提升性能。同时，在最终输出投影之前，利用线性层进行归一化和门控。
实验结果：Titans性能卓越
谷歌实验主要关注MAC、MAG和MAL三种变体，以及单独的神经记忆模块。每个模型都有四种尺寸，参数分别为170M、340M、400M和760M。
语言建模
在语言建模和常识推理任务中，神经记忆模块在困惑度和准确度方面表现最佳。Titans的三种变体均优于Samba和Gated DeltaNet-H2。
大海捞针
在“大海捞针”测试中，神经记忆模块表现出最佳结果。这归因于Titans的三个关键差异：神经记忆能够更好处理记忆容量，具有深度非线性记忆，并能通过遗忘机制有效管理记忆。
BABILong基准
在微调设置中，Titans(MAC)的表现优于所有模型，甚至包括GPT4等大型模型，以及基于Transformer的RMT等记忆模型。
深度记忆的影响
随着记忆深度的增加，模型在所有序列长度上实现了更好的困惑度。更深的记忆模块对序列长度的鲁棒性更强。随着参数量的增加，所有模型在较长的序列上都表现出更好的性能。
时序预测
Titans在时序预测任务中表现出色，神经记忆模块优于所有基线。
DNA 建模
Titans(LMM)在DNA建模任务中具有竞争力，在不同的下游基因组任务中表现良好。
效率
在训练吞吐量方面，Titans神经记忆模块比Mamba2和Gated DeltaNet稍慢，而Titans(MAL)比基线和神经记忆模块都快。
综上所述，谷歌的Titans架构通过引入长期神经记忆模块，成功地扩展了AI模型的上下文处理能力，并在多个任务中表现出卓越的性能，预示着AI领域的新纪元。
更多AI行业最新资讯新闻信息请关注AI人工智能网站--AITOP100平台--AI资讯专区：https://www.aitop100.cn/infomation/index.html
