Advertisement

论文翻译——Titans: Learning to Memorize at Test Time

阅读量:

1 引言

在连续十年的研究中, 人类持续投入了大量精力致力于探索如何高效利用循环模型与注意力机制。循环架构旨在将输入数据编码为固定长度的记忆单元(称为隐含状态), 而注意力机制则赋予网络能够聚焦于整个上下文窗口的能力, 从而捕获所有标记间的直接依赖关系。然而, 对依赖关系建模的提升不仅带来了额外的成本, 同时也限制了模型处理固定长度上下文的能力。为此, 我们开发出一种新型神经长期记忆模块,该模块能够学习历史化信息序列, 并指导注意力机制专注于当前情境下的内容, 同时有效继承过去积累的信息 。研究表明, 这种新型神经记忆系统展现出高效的并行训练能力与推理能力,其通过整合长期信息进行推理的能力远超现有架构 。从信息存储角度来看, 我们认为传统的注意力机制因有限的上下文窗口却能精确捕捉到依赖关系,因而充当着短期记忆的角色, 而基于强大信息编码能力的新一代神经记忆系统则可承担起长期维持与深度存储信息的任务 。基于这两种创新组件构建而成的泰坦架构系列中,我们提出了三种变体方案, 以解决不同场景下的信息整合问题。经过一系列语言建模、知识推理、基因分析及时间序列预测任务测试后发现: 相较于Transformer架构及其近期发展出来的线性循环模型家族成员, 泰坦系列模型在性能上更为优异; 并且能在超过200万个独立样本窗口下实现高效处理, 在模拟"针尖探海"类型任务中的准确性显著高于传统方法

备注:提到了注意力的短期记忆和长期的神经记忆(长期、持久)

“记忆的真正艺术是注意力的艺术!” ——塞缪尔·约翰逊,1787年

模块(Bietti等人, 2024),它们通过存储键值关联并计算查询与键之间的配对相似度来检索相关项。基于设计可知, Transformer的输出完全由当前上下文窗口中令牌的直接依赖关系决定。然而,这种精确建模依赖关系导致了关于上下文长度的二次时间和内存复杂度问题。在现实世界应用领域(如语言建模、视频理解、长序列预测等)中,随着场景复杂性的提升,上下文窗口可能变得极大,从而限制了Transformer在这些下游任务中的适用性

备注:query和上下文之间的成对相似度进行检索,说的非常形象。

为了解决 Transformers 的扩展性问题, 最近的研究致力于开发变种线性 Transformers(Kacham 等人 2024;Katharopoulos 等人 2020;S. Yang 等人 2024),其中他们将 softmax 替换为注意力核函数(有关详细信息,请参见第 2.1 节),从而显著降低了内存消耗。尽管这些变体在效率上有所提升, 并且能够处理更长的上下文, 但它们在性能上仍无法与标准 Transformer 相媲美, 因为这种替代方法使模型变成了一个线性循环网络, 其中数据被压缩为矩阵状态(Katharopoulos 等人 2020)。然而, 这种情况揭示了一个矛盾的事实:一方面, 这些线性模型通过增强扩展性和效率(基于线性和二次复杂度)提供了优势;另一方面, 非常长的上下文无法在矩阵状态中得到适当压缩(S. Wang 2024)。

此外,除了效率之外,大多数现有的架构——从霍普菲尔德网络(Hopfield 1982)到 LSTM(Jürgen Schmidhuber 和 Hochreiter 1997)以及 Transformer(Vaswani 等人 2017)——在处理泛化、长度外推和/或推理(Anil 等人 2022;Qin、Y. Zhong 和 Deng 2024)方面面临挑战,而这些都是许多现实世界中复杂任务不可分割的部分。尽管这些架构从人脑中汲取灵感,但它们都缺少以下关键要素:(1)学习过程的关键组成部分——例如短期记忆、长期记忆、元记忆、关注当前上下文等(Cowan 2008);(2)这些组件如何互联形成能够独立运行的系统;以及/或(3)从数据中主动学习并记忆过去历史的抽象能力。我们认为,在有效的学习范式中,类似于人脑,存在着不同的但相互关联的模块,每个模块负责学习过程中的一个关键组成部分

内存视角

认知过程是人类学习的核心环节,并非可分割的整体(Terry 2017)。若记忆系统未能正常运转,则人类及动物将受限于单一反应模式而非灵活适应机制 。由此可知,在机器学习领域中研究记忆机制始终是众多创新性工作的重要研究背景;例如,在Hopfield网络(Hopfield, 1982)、长短期记忆网络(LSTM, Jürgen Schmidhuber与Hochreiter, 1997)以及Transformer架构(Vaswani等人, 2017)等方面均可见到相关探讨。

参考相关文献(Okano, Hirano 和 Balaban 2000),大多数现有的架构认为记忆是由输入引发的神经更新过程,并将学习视为在给定目标下获取有效和有用记忆的行为。从这一角度来看,在时间点 t 接收到一个新的输入 x_t 后:(1) 模型会通过函数 f(M_{t-1}, x_t) 对当前的记忆 M_{t-1} 进行压缩更新;(2) 并通过函数 g(M_t, x_t) 对应查询到与当前输入相关的记忆信息。类似地,在Transformer架构中:(1) 模型会通过追加键-值矩阵对到内存中来实现对新信息的学习;(2) 并通过计算查询向量与内存键向量之间的相似度来检索与查询相关的内存值向量,并利用该相似度对相应的值向量进行加权以生成输出结果

这种观点有助于我们更好地理解现有技术范式及其关键差异,并为架构设计提供指导原则。例如,在Vaswani等人(2017)提出的Transformer模型与Katharopoulos等人(2020)提出的线性Transformer模型之间的主要区别在于内存结构与内存更新机制:前者通过保持所有原始信息(在上下文窗口内)而不进行数据压缩来实现信息存储;而后者则将历史数据编码为固定大小的矩阵形式以减少占用空间。尽管线性Transformer与线性RNN(包括状态空间模型)均采用了信息压缩策略,在内存更新阶段均实现了类似功能;但二者的内存结构存在显著差异:在线性RNN中采用向量形式而非矩阵形式的信息存储方式。因此这一观点最终引导我们提出了以下三个关键问题:(Q1)什么样的内存架构最为理想?(Q2)何种记忆更新机制最为合适?(Q3)什么样的记忆检索过程最为高效?

深入探讨人类记忆的本质,这一概念既不强调统一性...也不局限于单一功能...(Cowan, 2008)。实际上...由短期记忆、工作记忆和长期记忆等多种系统共同构成...每个系统均服务于特定功能...具有独特的神经架构并能独立运作。(Willingham, 1997)。基于此观察结果...我们提出以下问题:(Q4)如何构建一个高效的架构将不同联结型的记忆模块有机整合?最终发现存储过程是一个复杂的神经活动...需要将过去的信息进行编码并保存起来。假设仅仅依赖一个单一的向量或矩阵来进行编码与存储长期历史信息可能显得过于简化。(Q5) 是否有必要引入一个深度记忆模块以有效存储与回忆长时间信息?

本文旨在开发一个长期神经记忆模块以解决上述五个关键问题。该模块在测试阶段表现出色,并能够可靠地存储和检索信息。在此基础上,我们深入探讨了如何将其整合到架构中。

神经记忆(§3)。我们开发了一种深度神经长期记忆模块(NLM),该模块作为元上下文模型,在测试阶段能够通过其参数将数据进行存储与检索。受人类长期记忆系统(Mandler 2014)理论框架的启发,在设计该模块时我们特别关注了如何使令人惊讶的事件更容易被记住这一关键特性。为此我们基于输入信号计算其信息 Interestingness 的方法进行了优化,并将其整合到联想记忆损失函数中(有关详细计算方法,请参见§3.1)。为了实现对有限内存的有效管理我们提出了一个基于比例和数据惊讶度的动态遗忘机制(DMF),这一机制不仅继承了传统循环模型的核心功能还能显著提升系统的性能表现。(Dao 和 Gu 2024;Gu 和 Dao 2024;S. Yang 等人 2024)值得注意的是我们的实验表明这种动态遗忘机制本质上等价于通过小批量梯度下降算法结合动量项与权重衰减来优化元神经网络架构。(Yu Sun 等人 2024)在此基础上我们提出了一种高效可并行化的快速训练算法以实现深度神经长期记忆系统的最优性能。

在§4中讨论了设计长期神经记忆之后的问题时指出一个关键挑战是如何有效地将这些记忆整合到深度学习架构中。为此我们提出了泰坦这一深度模型家族它由三个主要超头构成:(1)核心组件该模块主要负责执行数据处理的核心流程(我们采用了有限窗口大小的注意力机制);(2)长期记忆单元这是我们设计的神经长期记忆模块用于存储和检索长期保存的信息;(3)持久性知识存储这是一组可学习的参数它们与时间无关但能够编码任务所需的知识。最后作为理论验证我们构建了泰坦的三个变体分别对应于基于上下文的架构基于层次化的架构以及带有门控分支结构的设计。

本节实验结果表明,在语言建模、常识推理、密集型检索、大规模检索、时间序列预测以及DNA建模等六个典型任务上进行了广泛的实证评估。通过系统分析发现,在全面基准测试中我们的TFT-3架构超越了当前主流循环架构及其组合形式(结合滑动窗口注意力版本),并且相较于Transformer,在相同长度的查询窗口上表现更为突出。值得注意的是,在采用完整上下文信息的情况下TFT-3依然保持着超越性能优势;而当仅考虑单个上下文窗口时其性能表现同样令人称道。这些成果是基于其设计使得其扩展性超越Transformer,在处理超过200万个查询窗口时仍能保持竞争力所实现

2 预备知识

本节将探讨本文中所使用的符号及其相关的背景概念。令输入为𝑥 ∈ R^{N×din},其中M代表神经网络或神经记忆模块;查询Q、键K以及值V分别代表注意力机制的不同组成部分;我们用S(i)表示第i个段,并通过S(j, i)来表示第i个段中的第j个标记(其中唯一的例外是带有时间下标的标记t)。为了简化讨论,在某些情况下我们将使用相同的符号来指代矩阵、向量或特定元素;例如,在这种情况下我们用S(j, i)来指代第i个段中的第j个标记。特别地,在涉及时间递归的情境下(如神经网络的状态),我们将保留t作为下标以区分不同时间步的信息。给定一个神经网络N以及一个输入样本x,请注意N(x)(分别N∗(x))用于表示带权(而不带权)的前向传播过程;此外,在某些情况下我们将单独关注某一层网络的行为并将其记作N(k),其中k代表该层的编号。在此背景下首先我们将介绍注意力机制及其高效变体的基本原理;随后回顾现代线性RNN模型的核心思想;最后从记忆视角探讨这些架构的特点并由此引出本研究的核心组件——Titans。

2.1 背景

从这一角度来看,在Equation 4所示的线性 Transformer 递归公式中等价于将键值对 (𝐾𝑡, 𝑉𝑡 ) 通过加法压缩并存储到矩阵值内存单元 M𝑡中。因此,在处理长序列数据时这种加法性质会导致内存溢出问题严重影响模型性能。为此研究人员已经集中探索了两个有潜力的方向:(1)引入遗忘机制:一些研究已经针对线性模型提出了自适应(数据依赖)遗忘门机制它可以在需要时动态调节记忆容量作为此类模型参考GLA(S. Yang B. Wang Shen 等人 2024)、LRU(Orvieto 等人 2023)、Griffin(De 等人 2024)、xLSTM(Beck 等人 2024)以及Mamba2(Dao 和 Gu 2024)。值得注意的是这些方法都与传统状态空间模型的离散版本具有内在关联。(2)优化更新操作:为了克服传统递归模型中内存更新操作所带来的加法性质的问题Widrow 和 Hoff (1988) 提出了Delta规则其中在执行更新操作之前系统会先清除旧的记忆内容以减少冗余计算。近年来S. Yang B. Wang Yu Zhang 等人 (2024) 提出了一个高效的可并行算法来加速训练过程。此外S. Yang Kautz 和 Hatamizadeh (2024) 进一步改进了DeltaNets通过引入遗忘门增强了其泛化能力

内存模块在神经网络设计中扮演着核心角色(Graves, Wayne 和 Danihelka, 2014;JH Schmidhuber, 1992;Jürgen Schmidhuber 和 Hochreiter, 1997;J. Zhang 等人, 2024)。将线性层视为键值(关联)内存系统的思想源自快速权重程序的概念,在其中动态快速程序被整合到循环神经网络中以实现可写性记忆(JH Schmidhuber, 1992)。赫布法则(Hebb, 2005)与德尔塔法则(Prados 和 Kak, 1989)作为快速权重程序中最常用的学习机制,在多个研究领域得到了广泛应用(Irie 等人, 2021;Munkhdalai 和 Sordoni, 2019;Munkhdalai 和 H. Yu, 2017;Schlag 等人, 2021;S. Yang 等人, 2024)。然而,在现有研究中,默认假设基于瞬时的惊讶信息输入序列...

我们深入探讨了我们的系统结构与现有模型之间的关联。在附录 A 中涉及了其他相关领域的研究。

3、测试时学习记忆

为了解决长期记忆能力不足的问题,并使模型具备学习、遗忘以及检索信息的能力,在本节内容里我们将构建一个神经长期记忆模块。这一模块被定义为一种能够在测试阶段学习记忆的元模型结构。其中,在第3.1节里我们将深入探讨神经记忆的动机与设计思路;而在第3.2节里则将重点研究架构设计如何实现高效并行训练过程;最后,在第3.3节里我们将通过引入持久记忆模块来增强整个架构性能。

3.1 长期记忆

为了构建一个神经长期记忆模块, 我们需要设计一个模型, 其能够嵌入到模型的参数中反映历史信息的关键特征。例如LLMs, 它们已被证明具备对训练数据的记忆能力(Leybzon 和 Kervadec 2024; Schwarzschild 等人 2024; Staab 等人 2024)。一种直观的方法是训练一个神经网络, 希望其能记住从训练过程中获得的知识。然而, 记忆常被视为一种负面特征, 因为它限制了模型的学习与推理能力(Bayat 等人 2024), 同时可能导致隐私泄露风险(Staab 等人 2024), 并在测试阶段表现得不如预期(测试性能低下)。此外, 记忆特定于训练数据在测试阶段可能无益, 因为测试数据通常来自分布外分布(分布外数据)。基于此, 我们认为:为了实现更好的泛化能力于测试阶段, 我们需要构建一个在线元模型系统, 其能动态地学习如何在测试时刻管理与遗忘相关数据特征。

备注:这就是知识的注入,能实现精准准入吗?

在学习过程中,衡量意外性的方法是一个重要的指标。核心思想在于将训练视为一个持续的在线学习过程,在这个过程中我们的目标是将过去的信息 𝑥₁至𝑥_{t−1}压缩至模块Mₜ的参数中。与人类相比,在人类的认知中,违反预期的事件更为显著(Mandler 2014)。受此启发地,在这种框架下我们可以将意外性简单定义为输入数据梯度的大小。具体而言,在线更新算法会根据该梯度值动态调整模型参数以适应变化的数据模式。较大的梯度值意味着输入数据与之前的历史数据存在更大的差异程度。基于这一原则,在每次迭代更新时我们都可以计算出当前状态下的异常程度,并据此优化模型的学习机制以提升长期记忆的质量。

备注:是否能够利用强化学习对历史信息进行压缩?或者是否有其他无损的信息压缩技术?

然而这种意外指标可能会使我们错过后续发生的重要信息尽管如此在经历了一系列的异常步骤后梯度可能会急剧减小从而导致模型陷入平坦区域即局部最小值并遗漏掉序列中某些重要的细节从人类的认知角度来看一个事件即使令人难忘也可能不会持续很长时间地让我们感到惊讶其原因在于最初的时刻足够令人震撼以至于能够长时间地抓住我们的注意力从而使我们记住整个时间段为了提升上述提到的意外指标即Equation 8我们将这一指标分解为两个部分第一个部分是过去异常我们用它来衡量最近一段时间内发生的异常次数和强度;第二个部分是当前输入的异常程度我们用它来评估最新的数据点是否超出了预期

有趣的是,在数学表达式中使用了类似于带动量的一阶优化器的概念。

我们的目标是实现意外指标的监控与评估。我们所讨论的意外指标基于损失函数 l(·;·), 这正是我们在测试中追求的目标。换句话说, 我们的记忆模块充当了一个元模型角色, 在此框架下它通过损失函数 l(·;·) 来学习相应的函数。

在本研究中,我们聚焦于记忆网络 ,其目标是将历史数据编码为键-值对的形式 。对于输入序列中的每一个元素𝑥𝑡,类似于Transformer架构(Vaswani等人提出),我们通过线性变换分别生成对应的键和值表示

在内循环阶段中,在**元模型(记忆)**框架下优化上述损失函数的过程中, 模型掌握如何在测试时记忆键值之间的映射关系. 类似于元学习模型(Nichol 2018; Zintgraf 等人 2019), 记忆模块的学习主要发生在内循环阶段. 因此, 上述损失函数中的超参数属于这一过程的关键组件. 具体而言, 在内层循环中我们优化M模块的权重, 在外层循环则优化整个架构体系中的其他可调参量.

遗失策略。在处理长度极长的序列(例如数百个甚至数千个标记)时,在决定哪些信息应被遗失这一点上具有重要意义——即使采用深度或极大规模的记忆矩阵也能发挥作用。为此我们设计了一种自适应式的遗失机制它允许模型自动丢弃无用或冗余的信息从而使有限的记忆容量能够得到更有效的利用也就是说对于接下来的一个标记𝑥𝑡我们将对更新规则进行相应的调整

内存体系结构方面,在本文中我们着重聚焦于具有𝐿M ≥ 1层的简单MLP作为我们长期记忆的基础构建块 。这种选择的主要原因是其能够有效促进长期记忆的设计理念,并通过巧妙的方式将其融入整体体系结构之中。然而我们的公式推导与体系结构创新为我们提供了一个系统性地开发出更具效率与效能的记忆数据处理机制的新途径。近期研究者们在这一领域已取得了一系列令人鼓舞的研究成果(Berges等2024;Cetin等2024;J. Zhang等2024)。基于现有研究成果我们可以设想一种创新性的研究路径即通过引入这些新型体系结构来替代传统的简单MLP从而探索出更加高效的数据处理方式

备注:需要研究下这篇论文

[10] Vincent-Pierre Berges, Barlas Oğuz, Daniel Haziza, Wen-tau Yih, Luke Zettlemoyer, and Gargi Gosh. "Memory Layers Scaled Up". In: arXiv preprint arXiv:2412.09764 (2024).

回忆提取方面,在前面我们探讨了一个可在测试期间学习长期记忆模块的设计与训练方法。核心问题在于:如何从存储的信息中提取所需内容?主要采用无权重更新(即推理)的前向传递策略来实现查询功能。形式上讲,在给定输入序列 x_1x_n 的情况下(其中 n 代表时间步长),通过线性变换矩阵 W_Q 将输入向量映射到查询空间以生成查询向量 q_t = x_t W_Q^T;随后以该查询向量作为索引,在内存库中快速定位并提取相关信息。

开发了一种基于记忆检索的知识管理系统,并强调了其重要性与实用性。例如,在各个专业领域中均可应用这一系统进行高效管理与查询操作。当用户输入查询时,默认采用该系统而非传统的向量检索方法以提高准确性,并进一步整合这些元数据信息以丰富当前上下文背景。

3.2 长期记忆训练的并行化方法

接下来我们详细阐述了在内层循环中进行权重计算并将其等价于仅通过矩阵乘法与总和运算实现这一过程。参照Yu Sun等人(2024)的研究工作其发现当采用恒定学习率的小批量梯度下降优化模型其前向传播过程等价于仅通过矩阵乘法实现这一结论在此基础上我们将序列划分为长度不小于1且等于b的块并采用以下形式表示小批量梯度下降:

3.3 持久性内存

从记忆的角度来看,在之前的讨论中指出

实际上,在𝑊𝐾和𝑊𝑉不受输入影响时

技术视角。基于带因果掩码的注意力机制对序列中的初始标记存在内在偏向,在大多数情况下会导致在大多数情况下注意力权重会异常活跃于初始标记。从技术层面而言,在序列开始处的这些可学习参数可以通过更为高效地分配注意力权重来缓解这一问题(Han 等人 2024;Xiao 等人 2024)。

(a) 存储介质(MAC)作为背景信息处理。我们对序列进行分段处理,并在每个时间段内应用完整的因果注意力机制。同样地,在前 𝑁𝑝 个标记位置设置了SSD分区,在随后的 𝑁𝑙 个标记位置则配置了AES-NI引导区。

备注:这里没太看明白,需要仔细研究下。

4 如何整合记忆?

一个尚未得到解答的重要问题是:如何有效且高效地将设计的神经记忆整合到深度学习架构中?如前所述,在探索记忆机制方面已有诸多研究工作取得了显著成果。从记忆机制的角度来看,在Transformer架构中K-Value矩阵对的作用可被恰当地理解为一种联想式信息存储单元。基于其在准确捕捉依赖关系方面的优势以及有限的窗口覆盖范围的特点,则将其视为短时记忆模块更为贴切合理。另一方面,在这一研究工作中我们提出了一种全新的神经记忆系统能够通过持续学习从数据中提取特征并将其嵌入到权重参数之中从而具备长期存储功能这种特性使其能够承担起长期记忆的角色在本节中我们重点介绍并详细探讨了三种不同的Titans变体及其性能特点

4.1 内存作为上下文

该架构具有两大显著优势:第一点在于其利用同时获取历史和当前上下文的能力,在处理当前数据时能够灵活判断是否有必要保留长期记忆的信息;第二点在于通过注意力模块实现对来自当前上下文的有效信息的筛选作用,在长期记忆中仅存储有意义的数据片段。这与传统方法不同,在传统方法中并非所有的标记都被认为是必要的资源;即使在测试阶段也无法简单地忽略这些资源。因此,在测试阶段依然需要动态更新神经网络的记忆权重参数以维持模型性能的稳定性。(i)在持久记忆参数方面保持固定不变;(ii)注意力模块的权重则由上下文学习器动态调整;(iii)长期记忆模块同样参与信息的学习与更新过程。也就是说,在测试阶段并不会停止对神经记忆权重的优化更新工作。”

备注:

模型在训练过程中积累了知识即持久记忆,在显存中存储着参数的数量为1亿个单元,并且这些参数所承载的知识非常丰富。用普通文本文件模拟这样的过程时会发现:每本红楼梦约96万字相当于大约占用了1.8MB的内存空间;而使用显存容量为1GB(即1,024MB)则可以存储大约6本这样的书籍所代表的知识总量。为了验证这一点,请问神经网络是否能够实现比现有编码方法更加高效且精确的信息提取能力呢?

(2)是否需要长期记忆。由注意力机制决定

如图4所示:记忆模块被定义为门(MAG)架构。该架构同样具有三个主要部分:第一部分是核心功能区;第二部分是辅助信息存储区域;第三部分是长期记忆系统。值得注意的是,在这个设计中只有长期记忆系统被整合到辅助区域中;其余的记忆功能则通过门控机制与核心模块相连接以确保数据的安全传输路径。经过测试发现其行为模式与Figure2一致

4.2 门控记忆

该设计中的总体注意力掩码如图3b所示。在此设计中,“滑动窗口注意力充当精确短期记忆” 而神经记忆模块则承担模型衰退记忆的任务。此体系结构也可视为一个多头架构 ,其中各头特征各异(X. Dong 等人 2024)。

4.3 内存作为一层

这一变体采用了深度神经网络的神经记忆模块(如图5所示)。在现有文献研究中较为常见的是采用混合架构通过融合循环模型与完整的滑动窗口注意力机制实现功能。给定输入𝑥我们有:y = f(x)

图 5展示了内存构成层 (MAL) 架构。在该架构中,内存层承担着在注意力机制之前缩减过去与当前情境信息的任务。

不依赖短期记忆的记忆。尽管在上述讨论中我们探讨了MAL作为LMM与注意力按顺序结合而成的概念,但MAL的一个简单变体则是将LMM视为一个不依赖短期记忆的序列模型.从记忆的角度来看,正如第一节所述,我们期望每个记忆系统模块都能够独立运作,即便其他组件受到影响.因此即使缺乏短期记忆(即不再具备注意力机制),长期记忆模块仍应具备强大的建模能力.在我们的实验研究中,我们将这种变体命名为泰坦型LMM或简称为TBM.附录C详细探讨了TBM与其他现代循环模型之间的联系。

4.4 架构细节

继最近的现代线性循环模型(Gu 和 Dao 2024;S. Yang、Kautz 和 Hatamizadeh 2024)之后,在每个查询、键和值投影之后引入了一维深度可分离卷积层。尽管这一设计在性能上未带来显著提升(虽然对性能没有显著影响),但这些卷积操作已被证明能够有效提升模型性能,并且在计算资源利用方面表现出了较高的效率

门控。我们还遵循最近的架构,在最终输出投影之前使用归一化和门控,并使用线性层(Mehta 等人 2023)。

定理 4.1

5 实验

本研究中, 我们重点关注泰坦架构中的三个关键变体: (1)基于上下文的记忆单元 (MAC), (2)以门机制形式存在的记忆单元 (MAG), 以及 (3)以层式结构呈现的记忆单元 (MAL). 此外, 我们还设计了一个独立的神经记忆模块. 我们选择将长期记忆单独作为一个模块的原因在于, 这一设计理念源于我们对学习机制的独特理解. 根据第1节所述, 学习被定义为系统获取高质量的记忆能力的过程. 因此, 我们期望我们的长期记忆系统能够在无需注意力的情况下有效地从大量数据中提取知识. 在每一种变体模型中, 我们考察了四个不同的规模选项: (i)17,000万个参数; (ii)34,000万个参数; (iii)40,000万个参数; 和 (iv)76,000万个参数. 前三个变体模型均基于来自FineWeb-Edu数据集 (Penedo等人2024) 的15亿级别标记样本进行训练, 而最后一个变体则采用了同一数据集中更大的样本量——3百亿级别标记样本

最近,在泰坦架构中进行不同组件消融实验的研究工作取得了进展。我们将其神经记忆模块作为基准模型进行分析,在每次实验中逐一替换或移除关键组件:首先是以线性记忆替代深度记忆;其次移除卷积层;再次从惊喜度量中去除动量相关参数;之后放弃权重衰减机制(亦即不采用遗忘机制);最后移除持久记忆模块。具体结果可见表 5 中的数据展示。其中各个组件均显著提升了系统性能水平,并且在评估指标对比中发现:权重衰减、动量、卷积和持久记忆这四个模块在整体性能提升方面表现最为突出

6 结论

本文提出了一种基于循环神经网络的长时记忆机制,并将其设计为一种元上下文学习器,在测试阶段能够自主构建和更新长期记忆信息。该神经记忆模块本质上是一个自适应循环模型,在训练过程中能够自动识别并增强那些更具代表性的或与已有知识更为接近的新标记信息。与现有的循环模型架构(如Transformer)相比,在记忆更新机制和存储能力方面表现出更强的表现力。为了实现这一目标,我们提出了Titans架构及其三个主要变体:其中我们建议将记忆模块作为(1)上下文提取器、(2)门控机制以及(3)层级化存储模块进行集成。通过一系列实验任务评估表明,在长上下文处理能力方面,Titans模型显著优于Transformer架构以及近期提出的现代线性循环模型。此外,在实验中我们还发现,在特定条件下Titans还能够扩展至超过2M个上下文窗口大小,并且在保持较高准确性的同时实现了对基线性能的超越。

泰坦模型是基于 PyTorch 和 JAX 开发的,我们打算尽快提供给训练与评估使用。

全部评论 (0)

还没有任何评论哟~