-Titans: Learning to Memorize at Test Time

Ali Behrouz† , Peilin Zhong† , and Vahab Mirrokni†

Google Research

摘要

超过十年来，人们已经进行了大量关于如何有效利用循环模型和注意力的研究 。 虽然循环模型旨在将数据压缩成固定大小的记忆（称为隐藏状态），但注意力允许关注整个上下文窗口，捕捉所有标记的直接依赖关系。 然而，这种更准确的依赖建模会带来二次成本，限制模型只能处理固定长度的上下文。 我们提出了一个新的神经长期记忆模块，该模块学会记忆历史上下文，并帮助注意力把焦点放在当前上下文，同时利用长记忆的信息。 我们展示了这种神经记忆在训练时具有快速可并行化的优势，同时保持了快速的推理速度。从记忆的角度看，我们认为 注意力由于其有限的上下文但准确的依赖建模，类似于短期记忆 ， 而神经记忆由于其能够记忆数据，更像是长期、更持久的记忆 。基于这两个模块，我们引入了一种新的架构系列，称为 Titans ，并提出了三种变体来解决如何有效地将记忆结合到这种架构中。我们在语言建模、常识推理、基因组学和时间序列任务的实验结果表明， Titans 比 Transformers 和最近的现代线性循环模型更有效 。相比基线， 它们在大海捞针任务中可以有效地扩展到超过 2M 的上下文窗口大小，并具有更高的准确性。

“ 记忆的真正艺术是注意力的艺术！ ” — 萨缪尔 · 约翰逊 ， 1787

目录

1 介绍

1.1 记忆视角

1.2 贡献和路线图

2 2. 初步工作

2.1 背景

3 学习如何在测试时记忆

3.1 长期记忆

3.2 如何并行化长期记忆训练

3.3 持久性内存

4 如何融入记忆？

4.1 记忆作为背景

4.2 门控记忆

4.3 内存作为一层

4.4 架构细节

5 实验

5.1 实验设置

5.2 语言建模

5.3 大海捞针

5.4 BABILong 基准测试

5.5 深度记忆的效应

5.6 时间序列预测训练吞吐量

5.7 DNA 建模

5.8 效率

5.9 消融研究

6 结论

参考文献

A 相关工作

A.1 线性递归模型

A.2 基于Transformer 的架构

A.3 测试时间训练和快速权重程序

B 语言模型和常识推理数据集

C 长期记忆模块 （LMM ） 作为序列模型

1 介绍

Transformer ，纯基于注意力的架构（ Vaswani 等， 2017 年），已经被确定为序列建模中的最先进模型，主要是由于它们具有上下文学习的能力和在规模上的学习能力（ Kaplan 等， 2020 年）。 Transformers 的主要构建模块 —— 注意力模块 —— 作为关联性记忆块 （ Bietti 等， 2024 年）发挥作用，它 们学习存储键值关联并通过计算查询 Q （即搜索信号）和键 K （即上下文）之间的成对相似度来检索它们 。因此， 根据设计， Transformer 的输出仅在当前上下文窗口中的 Token 的直接依赖条件下。然而，这种对依赖关系的准确建模会带来关于上下文长度的二次时间和内存复杂度 。在复杂的现实世界任务中（例如，语言建模（ N.F.Liu 等， 2024 年），视频理解（ C.-Y.Wu 等， 2019 年），长期时间序列预测（ H.Zhou 等， 2021 年）），上下文窗口可能变得非常大，使得 Transformers 在这些下游任务中的适用性有挑战性。

为了解决 Transformer 的可伸缩性问题，最近的研究旨在设计不同变体的线性 Transformer （ Kacham, Mirrokni 和 P. Zhong 2024; Katharopoulos 等， 2020 年 ; S. Yang ， B. Wang ， Shen 等， 2024 年），其中在 注意力机制中用核函数替代了 softmax （详见 §2.1 ），导致内存消耗显著减少。尽管线性 Transformer 具有高效性和能够扩展到更长的上下文，但与 Transformer 相比，线性 Transformer 在性能上并不具备竞争力，因为核技巧使模型成为了线性递归网络，在其中数据被压缩为矩阵值状态（ Katharopoulos 等， 2020 年）。然而，线性递归（或线性 Transformer ）模型带来了一个矛盾的事实：一方面，我们使用这些线性模型来提升可伸缩性和效率（线性 vs. 二次复杂性），这些优势在非常长的上下文中表现出来 ; 另一方面，一个非常长的上下文无法被适当地压缩为小的矢量值或矩阵值状态（ S. Wang 2024 ）。

此外，除了效率之外，大多数现有的架构，从 Hopfield 网络（ Hopfield 1982 ）到 LSTMs （ Jürgen Schmidhuber 和 Hochreiter 1997 ）和 Transformers （ Vaswani 等， 2017 年），在处理泛化、长度外推和 / 或推理（ Anil 等， 2022 年； Qin ， Y. Zhong 和 Deng 2024 年）时面临挑战，所有这些都是许多复杂的现实世界任务的不可分割的部分。虽然这些架构从人类大脑中获取灵感，但每种架构都缺少：（

1 ）学习过程的关键组件 —— 如短期记忆、长期记忆、元记忆、注意当前上下文等 （ Cowan ， 2008 年）；（ 2 ） 这些组件如何是相互关联的系统，可以独立运行 ；和 / 或（ 3 ） 积极从数据中学习并记忆过去历史的抽象能力 。我们认为，在一个有效的学习范式中，类似于人类大脑，有着独特而相互关联的模块，每个模块负责学习过程中的关键组件。

1.1 记忆视角

记忆是一种基本的心理过程，是人类学习不可分割的组成部分 （ Terry 2017 ）。没有正常运作的记忆系统，人类和动物将被限制在基本的反射和刻板行为中。因此，记忆已经成为机器学习文献中许多重要研究的灵感来源；例如， Hopfield 网络（ Hopfield 1982 ）、 LSTM （ Jürgen Schmidhuber 和 Hochreiter 1997 ）和 Transformer （ Vaswani 等， 2017 ）。

灵感来自于神经心理学文献中对记忆和学习的常见定义 （ Okano ， Hirano 和 Balaban 2000 ）， 大多数已有的体系结构将记忆视为输入引起的神经更新，并将学习定义为在给定目标的情况下获取有效和有用的记忆的过程。 从这个角度来看，循环神经网络（ RNNs ）（ Williams 和 Zipser 1989 ）可以定义为具有矢量值记忆模块 M （也称为隐藏状态）的模型，其有两个主要步骤：给定时间 t 的新输入 𝑥𝑡 ，模型（ 1 ）使用函数 𝑓 （ M𝑡−1,𝑥𝑡 ）来更新内存。

( 压缩版本 ) ；并使用函数 𝑔(M𝑡,𝑥𝑡) 检索输入的对应记忆（详见第 2.1 节）。类似地， Transformer 可以被看作是具有不断增长记忆和两个类似步骤的体系结构。也就是说，键和值矩阵对作为模型的记忆，并且模型： (1) 通过将键和值追加到记忆中来更新记忆（不压缩），并 (2) 通过查找查询向量和键向量的相似性来检索查询向量的对应记忆，然后用于权重值向量的输出 。

这一观点可以帮助我们更好地理解现有的范式，它们之间的关键区别，并设计更有效的架构。 例如， Transformer （ Vaswani 等人， 2017 ）和线性 Transformer （ Katharopoulos 等人， 2020 ）之间的主要区别在于记忆结构以及记忆更新步骤，其中线性 Transformer 将历史数据压缩为固定大小的矩阵值记忆，而 Transformer 在不进行任何压缩的情况下保留所有历史数据（在上下文长度内）。虽然线性 Transformer 和线性 RNN （包括状态空间模型）在记忆更新步骤中都会压缩信息，但关键区别在于记忆的结构。

线性 RNN （与线性 Transformer 相对）使用向量值记忆（与矩阵值记忆相对）。因此，这个观点激励我们提出以下问题：（ Q1 ）什么构成了良好的记忆结构？（ Q2 ）什么是适当的记忆更新机制？和（ Q3 ）什么是良好的记忆检索过程？

重温人类记忆的理解，它既不是一个单一的过程，也不仅仅有一个功能（ Cowan 2008 ）。事实上，记忆是一个由系统组成的联盟 - 例如，短期记忆，工作记忆和长期记忆 - 每个系统都有不同的功能和不同的神经结构，并且每个系统都可以独立操作 （ Willingham 1997 ） 。这一事实激励我们提出以下问题：（ Q4 ）如何设计一个包含不同连接的记忆模块的高效架构。最后，存储记忆是一个神经过程，需要编码和存储过去的抽象。单纯地假设一个单一的向量或者矩阵足以线性地编码数据来存储长期历史是过于简化的。（ Q5 ）是否需要一个深度记忆模块来有效存储 / 记住长期历史？

1.2 贡献和路线图

在本文中，我们旨在通过设计一个长期的神经记忆模块来回答上述五个问题，该模块可以在测试时高效有效地学习记忆。在基于其设计的基础上，我们讨论了如何将其融入到架构中。

神经记忆 （ §3 ）。我们提出了一种（深度）神经长期记忆，该记忆模块（作为元上下文模型）学会了如何在测试时将数据存储在其参数中。受人类长期记忆系统的启发（ Mandler 2014 ），我们设计了这个记忆模块，使得违背期望的事件（令人惊讶）更易记忆。为此，我们通过神经网络对输入关联记忆损失的梯度来衡量输入的意外性（具体细节见 §3.1 ）。为了更好地处理有限的记忆空间，我们引入了一个衰减机制，考虑到记忆容量与数据意外性之间的比例，以实现更好的记忆管理。我们展示了这种衰减机制实际上是现代递归模型中遗忘机制的泛化（ Dao 和 Gu 2024 ； Gu 和 Dao 2024 ； S. Yang 、 Kautz 和 Hatamizadeh 2024 ）。有趣的是，我们发现这种机制等价于用小批量梯度下降、动量和权重衰减优化元神经网络。基于将小批量梯度下降张量化以使用更多的矩阵乘法运算（ Yu Sun 等人 2024 ），我们提出了一种快速且可并行化的算法来训练我们的深度神经长期记忆。

Titans 架构 （ §4 ）。在设计长期神经记忆之后，一个重要的待解决问题是如何将记忆有效且高效地融入深度学习架构中。我们介绍了 Titans ，一系列由三个超头组成的深度模型：（ 1 ）核：该模块包含短期记忆，负责主要的数据处理流程（我们使用有限窗口大小的注意力）；（ 2 ）长期记忆：这个分支是我们的神经长期记忆模块，负责存储 / 回忆长期的过去；（ 3 ）持久记忆：这是一组可学习但与日期无关的参数，用于编码关于任务的知识。 最后，作为概念验证，我们提出了 Titans 的三种变体，将记忆作为：（ i ）上下文，（ ii ）层，以及（ iii ）门控分支加入其中。

实验结果 （ §5 ）。我们对语言建模、常识推理、回忆密集型、大海捞针、时间序列预测和 DNA 建模任务进行了实验评估。我们观察到，我们的 Titan 架构在一系列综合性基准测试中表现优于所有现代递归模型以及它们的混合变体（与滑动窗口注意力相结合）。此外，相比于使用相同上下文窗口的 Transformer 模型， Titan 表现更好，并且与使用全部上下文的 Transformer 模型具有竞争力。与 Transformer 相反， Titan 模型可以扩展到超过 2M 的上下文窗口大小。

2 初步工作

在本节中，我们讨论我们在论文中使用的符号和一些背景概念。我们让 𝑥 ∈ R𝑁×𝑑in 为输入， M 为神经网络（神经记忆模块）， Q ， K ， V 为注意力机制的查询，键和数值， M 为注意力掩码。在对序列进行分段时，我们使用 S(𝑖) 来指代。

第 𝑖 段落。在本文中，我们滥用符号，并使用下标来引用矩阵、向量或段的特定元素。例如，我们将 S(𝑗𝑖) 视为第 𝑖 段中的第 𝑗 个 Token 。唯一的例外是带有 𝑡 下标的情况，我们保留用于索引时间上的重复，或者神经网络在时间 𝑡 上的状态。给定一个神经网络 N 和一个数据样本 𝑥 ，我们使用 N (𝑥) （或 N∗(𝑥) ）来表示带有（或不带有）权重调整的前向传播。同样，我们滥用符号并使用 N (𝑘) 来引用神经网络的第 𝑘 层。在接下来的内容中，我们首先讨论注意力及其有效变体的背景，然后回顾现代线性 RNN 。最后，我们讨论这些架构的记忆视角，这促使我们设计 Titans 。 （原文没有给出 "Titans" 的解释，我没有进行翻译）

2.1 背景

注意力。 Transformer 变压器（ Vaswani et al. 2017 ）作为许多深度学习模型的事实支柱，基于注意力机制。给定输入 𝑥 ∈ R𝑁×𝑑in ，因果注意力基于输入相关的键、值和查询矩阵上的 softmax 计算输出 y ∈ R𝑁×𝑑in ：

WQ 、 WK 和 WV 属于可学习参数，它们都是 R𝑑in×𝑑in 的矩阵。尽管 transformers 在召回方面具有强大和有效的能力，但是为了计算输出，需要至少 𝑁 × 𝑑 个运算符，导致更大的内存消耗和对于更长序列的较低吞吐量。

高效的注意力机制。 为了改善长序列的 softmax 注意力在内存消耗和吞吐量方面的问题，各种研究聚焦于注意力的 I/O 感知实现（ Dao 2024; Dao, D. Fu, et al. 2022 ），通过稀疏化注意力矩阵（ B. Chen et al. 2021; Choromanski et al. 2021; Dai et al. 2019 ）、近似 softmax (Arora et al. 2024) ，或者开发基于核（线性）的注意力机制（ Aksenov et al. 2024; Kacham, Mirrokni, and P. Zhong 2024; Schlag, Irie, and Jürgen Schmidhuber 2021; S. Yang, B. Wang, Shen, et al. 2024 ）来设计更高效的注意力机制。在本部分中，我们将重点放在后者，即线性注意力机制上，其中标准注意力的 softmax 被替换为替代的核函数 𝜙(., .) ，使得 𝜙(𝑥,𝑦) = 𝜙(𝑥)𝜙(𝑦) 。因此，注意力可以表示为：

导致吞吐量更高，因为术语 𝐾𝑗=1 𝜙(𝐾𝑗) 和 𝐾ℓ=1 𝜙(𝐾ℓ) 在每个步骤中被重复使用。当选择内核为单位矩阵时（ Yutao Sun 等人， 2023 年），上述表述也可以用递归格式写成：

线性注意力机制可有效支持推理。

现代线性模型及其记忆视角。正如之前讨论的，人们可以将学习定义为获取有效和有用记忆的过程。基于这一点，可以将循环神经网络（ RNNs ）的隐藏状态视为记忆单元，模型旨在将信息压缩到其中。因此，在一般形式的循环神经网络中，隐藏状态可以被视为记忆单元，循环过程可以拆分为在记忆单元中的读取和写入操作。也就是说，我们让 𝑥 ∈ R𝑁×𝑑in 表示输入， M ∈ R𝑑 表示记忆单元，而 y ∈ R𝑑in 表示输出，那么循环神经网络的一般形式被定义为：

𝑓(.) 是读取函数， 𝑔(.) 是写入函数。请注意，这里的 M𝑡 的下标表示时间 𝑡 时内存的状态。从这个角度来看，线性 Transformer 的递推公式（见方程 4 ）等同于将键和值（ 𝐾𝑡,𝑉𝑡 ）进行加性压缩并写入到矩阵值记忆单元 M𝑡 中。因此，当处理长上下文数据时，这个加性过程的性质会导致内存溢出，严重影响模型性能。为了解决这个问题，研究集中在两个有前途的方向上： (1) 添加遗忘机制：一些研究提出了适应性（数据相关）的遗忘门机制用于线性模型，当需要时能够擦除记忆。作为这类模型的例子，我们提到 GLA （ S. Yang, B. Wang, Shen 等人， 2024 年）、 LRU （ Orvieto 等人， 2023 年）、 Griffin （ De 等人， 2024 年）、 xLSTM （ Beck 等人， 2024 年）和 Mamba2 （ Dao 和 Gu ， 2024 年），后者还与传统状态空间模型的离散版本相连（ Gu 和 Dao ， 2024 年）。（ 2 ）改进写入操作：为了克服传统递归模型中内存写入操作的加性性质， Widrow 和 Hoff （ 1988 年）提出了 Delta Rule ，即在添加记忆（即键值对）之前，模型首先移除它的过去值。为了增强可并行化训练和扩展性， 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 年）和 delta 学习规则（ Prados 和 Kak 1989 年）是快速权重程序中最流行的学习规则，已在各种研究中进行了广泛探讨（ Irie 、 Schlag 等人 2021 年； Munkhdalai 、 Sordoni 等人 2019 年； Munkhdalai 和 H. Yu 2017 年； Schlag 、 Irie 和 Jürgen Schmidhuber 2021 年； JH Schmidhuber 1992 年； S. Yang 、 Kautz 和 Hatamizadeh 2024 年； S. Yang 、 B. Wang 、 Yu Zhang 等人 2024 年）。然而，所有这些模型都基于瞬时惊喜，缺乏序列中的 Token 流动（参见第 3.1 节），并且大多数模型缺乏遗忘门，导致内存管理不佳。

我们在附录 C 中进一步讨论了我们的架构与最近模型之间的联系。附录 A 中还讨论了其他相关工作。

3 学习如何在测试时记忆

为了克服长期记忆的缺乏并使模型能够学习、遗忘和检索信息，在本节中，我们提出了一个神经长期记忆模块，它是一个在测试时学会记忆的元模型 。在 3.1 节中，我们首先讨论了神经记忆的动机和设计。在 3.2 节中，我们讨论了我们的架构设计如何受益于快速且可并行化的训练。最后，在 3.3 节中，我们使用持久记忆模块增强了我们的架构，其中我们使用可学习但与数据无关的参数来学习任务的元信息。

3.1 长期记忆

设计神经长期记忆模块时，我们需要一个能够将过去历史的抽象编码到其参数中的模型。 一个例子是已经被证明能够记住其训练数据的 LLM 模型（ Leybzon 和 Kervadec ， 2024 ； Schwarzschild 等， 2024 ； Staab 等， 2024 ）。因此，一个简单的想法是训练一个神经网络，并期望它能够记住其训练数据。然而，记忆几乎一直被认为是神经网络中不期望出现的现象，因为它限制了模型的泛化能力（ Bayat 等， 2024 ），引起了隐私方面的担忧（ Staab 等， 2024 ），从而导致在测试时性能不佳。此外，训练数据的记忆在测试时可能并没有帮助，因为数据可能是来自分布之外的。我们认为，我们需要一个在线元模型，在测试时学习如何记住 / 忘记数据。在这种设置下，模型正在学习一种能够进行记忆的函数，但它没有过度适应训练数据，在测试时表现出更好的泛化能力。

学习过程和惊奇度量。 训练长期记忆的关键思想是将其训练视为在线学习问题，在其中我们的目标是将过去的信息 𝑥1, . . . ,𝑥𝑡−1 压缩成我们的长期神经记忆模块 M𝑡 的参数。如前所述，对于违背期望的事件（即令人惊讶的事件）人类更容易记住（ Mandler 2014 ）。受此启发，模型的惊奇度可以简单定义为其相对于输入的梯度。梯度越大，输入数据与过去数据的差异就越大。因此，使用这个惊奇度分数，我们可以更新记忆为：

这种惊喜指标可能导致错过在一个重大惊喜时刻之后出现的重要信息。也就是说，在经历了几个令人惊喜的步骤之后，梯度可能变得极小，导致停留在一个平坦区域（即局部最小值），并且错过了序列中的某些部分的信息。从人类记忆的角度来看，一个事件在很长一段时间内可能不会始终让我们感到惊讶，尽管它是令人难忘的。原因是最初的瞬间足够惊喜，从而在很长的时间范围内引起我们的注意，导致我们记住了整个时间范围。为了改进上述惊喜指标（方程式 8 ），我们将惊喜指标分为（ 1 ）过去的惊喜，用于衡量最近过去的惊喜量；和（ 2 ）瞬时惊喜，用于衡量即将到来数据的惊喜量。

过去的惊奇 |{z} 瞬间的惊奇 𝑆𝑡

有趣的是，这个公式与具有动量的梯度下降类似，其中动量元素在哪里。 因此，在这个公式中，动量充当跨时间（序列长度）的惊喜记忆。在这个公式中，项 𝜂𝑡 是一个数据相关的惊喜衰减（ 𝑥𝑡 的函数），控制着随时间衰减的惊喜程度，而项 𝜃𝑡 则控制着多少瞬时惊喜应以数据相关的方式纳入最终的惊喜度量中。这种数据依赖性在这个设计中特别重要：虽然以前的标记的惊喜可能需要影响下一个标记的惊喜，但只有当所有标记都是相关的并且处于相同的上下文中时才有效。因此，数据相关的 𝜂 可以控制记忆是否需要 : (1) 通过设置 𝜂𝑡 → 0 忽略上次的惊喜（可能是由于上下文的变化），或者 (2) 通过设置 𝜂𝑡 → 1 全部纳入上次的惊喜（可能是因为该标记与其最近的过去标记高度相关）。

我们的目标。 我们上述的惊喜度量基于损失函数 ℓ(.; .) ，这是我们的记忆在测试时学习行动的目标。也就是说，我们的记忆模块是一个元模型，它基于损失函数 ℓ(.; .) 学习一个函数。

在这项工作中，我们关注关联记忆，我们的目标是将过去的数据存储为键和值对。给定 𝑥𝑡 ，类似于 Transformers （ Vaswani 等人， 2017 年），我们使用两个线性层将 𝑥𝑡 投影为一个键和一个值：

其中 WK 和 WV ∈ Rdin×din 。接下来，我们希望我们的 memory 模块能够学习 key 和 values 之间的关联。为此，我们将损失定义如下：  我们希望我们的记忆模块学习键与值之间的关联。为此，我们定义损失如下：

通过在我们的元模型（内存）的内循环中优化上述损失函数，模型学习如何在测试时记忆键和值之间的映射。注意，类似于元学习模型（ Nichol 2018; Zintgraf et al. ）。 Note: The translation has been provided in the original text, no translation is needed.

2019 年，内循环中进行了记忆的训练，因此上述损失函数中的参数 𝑊𝐾 和 𝑊𝑉 是超参数。因此，在内循环中，我们优化 M 的权重，而在外循环中，我们优化整个架构的其他参数。

遗忘机制。 在处理非常大的序列时（例如，数百万个标记），管理应该遗忘哪些过去的信息是至关重要的，即使有一个深度或非常大的矩阵值内存也是如此。为此，我们使用了一种自适应遗忘机制，使得内存可以遗忘不再需要的信息，以更好地管理内存的有限容量。也就是说，给定下一个标记 𝑥𝑡 ，我们将更新规则修改为：

𝛼𝑡 ∈ [0, 1] 是灵活控制记忆的门控机制，即决定应该忘记多少信息。例如，它可以通过 𝛼𝑡 → 0 更新记忆而不影响过去的抽象，也可以通过 𝛼𝑡 → 1 清空整个记忆。在本节后面，我们将展示这种权重衰减机制与现代循环神经网络中的门控机制（ Dao and Gu 2024; Orvieto et al. 2023 ）密切相关。

记忆架构。 在这篇论文中，我们将简单的 MLP （多层感知器）作为我们长期记忆的架构，主要原因是我们想更好地激发长期记忆设计的动机以及将其融入架构的方式。然而，我们的公式和架构设计开辟了设计神经架构以更有效和高效地记忆数据的新研究方向。最近，已经有一系列有希望的工作来设计这样的架构（ Berges 等人， 2024 年； Cetin 等人， 2024 年； J. Zhang 等人， 2024 年），将它们融入我们的框架（即，用这些架构替换简单的 MLP ）可能是一个有趣的未来工作。

当使用向量值或矩阵值的存储器（ De 等人， 2024 年 ; Orvieto 等人， 2023 年 ; S. Yang ， B. Wang ， Shen 等人， 2024 年）时，存储器模块会将过去的数据进行压缩，并将其适应到一条线上。也就是说，从元学习或在线学习的角度来看（ Yu Sun 等人， 2024 年），使用矩阵值存储器 M = 𝑊 ∈ R𝑑in×𝑑in 等效于优化 ℓ(𝑊𝑡−1;𝑥𝑡) = ∥𝑊𝑡−1k𝑡 − v𝑡 ∥22 ，这是一个在线线性回归目标，因此最优解假设历史数据之间的依赖关系是线性的。 另一方面，我们提出，深度存储器模块（即。。。）

在理论结果上与至少包含两层的 MLP 严格比线性模型更具表现力（ Hornik, Stinchcombe, and White 1989 ）一致，我们在第 5.5 节中展示深度记忆模块在实践中更加有效。

在上文中，我们讨论了如何设计和训练一个长期记忆模块，该模块在测试时学习记忆。一个关键的问题是：如何从记忆中检索信息？我们仅使用前向传播而不进行权重更新（即推断）来检索一个与查询对应的记忆。形式上，给定一个输入。

在这里，我们使用线性层 𝑊𝑄 来投影输入，即， q𝑡 = 𝑥𝑡𝑊𝑄 ，并通过以下方式从记忆 𝑦𝑡 中检索相应（或有用的）信息：

图 1 ：神经记忆训练可以在并行和使用矩阵乘法的情况下进行。

3.2 如何并行化长期记忆训练

如上所述，我们的长期记忆模块的设计相当于通过优化联想记忆损失函数 ℓ(M𝑡−1;𝑥𝑡) = ∥M𝑡−1 (k𝑡) − v𝑡 ∥22 来训练元模型，其中使用了动量和权重衰减的梯度下降法。因此，理论上，长期记忆模块的训练需要 O(𝑁) 个浮点运算操作，其中 𝑁 为序列长度。然而，实际上，我们需要并行化训练过程，并充分利用硬件加速器（如 TPUs 、 GPUs ），我们需要将过程张量化并使用更多的矩阵乘积运算。

接下来，我们将展示通过使用小批量梯度下降、数据相关的学习率和权重衰减来计算内部循环中的权重，可以进行重新构建，只使用矩阵乘法和求和运算。我们基于 Yu Sun 等人（ 2024 年）的工作进行了改进，该工作显示了使用小批量梯度下降（具有恒定学习率）优化模型的前向传播可以使用矩阵乘法计算。我们可以将序列分割成大小为 𝑏 ≥ 1 的块，并将小批量梯度下降写成以下形式：

其中， 𝑡' = 𝑡 − mod(𝑡,𝑏) ，而 𝛽𝑖 = Î𝑖𝑗=1(1 − 𝛼𝑗) 。为了简单起见，我们关注第一个块，即 𝑡 = 𝑏 ，因此 𝑡' = 0 。此外，我们说明了当 M𝑡 = 𝑊𝑡 是线性时的过程。具有 𝑁𝑝 ≥ 2 的 MLPs 的过程类似。使用我们的损失函数，我们有：

𝑖=1𝛽𝑖 的中文翻译为 𝑖=1𝛽𝑖 。

Θ𝑏 =diag 𝜃1 𝜃2 . . . 𝜃𝑏 和 B𝑏 在 𝛽𝛽𝑏𝑖 上被类似地定义。注意，我们不需要为 𝑘 = 1, . . . , 𝑁 /𝑏 存储所有 Θ𝑘𝑏 和 B𝑘𝑏 ，而是为每个块存储这些矩阵，从而使用更少的内存。接下来，我们扩展这个表示，以便我们也可以包含动量项。在带有动量的分块梯度下降中，如果我们考虑动量项，我们有：

𝑢𝑡 = ∇ℓ (𝑀𝑡′;𝑥𝑡) 。注意，我们可以同时计算所有 𝑢𝑡 ，因此方程 18 是一个具有 𝑢𝑡 作为输入、 𝑆𝑡 作为隐藏状态和 𝜂𝑡 作为输入相关转换值的线性递归。因此，我们可以使用并行的关联扫描（ J.T. Smith, Warrington 和 Linderman 2023 ）来计算此块中的 𝑆𝑡s 。

参数作为块的函数。我 们可以将参数 𝛼𝑡 ， 𝜃𝑡 和 𝜂𝑡 设置为其块的函数，而不是使其依赖于输入（即使其为 Token𝑥𝑡 的函数）。尽管失去了一定的表达能力，但这种形式可以帮助加快训练速度。在这种情况下，我们在每个块中使用相同的 𝛼 ， 𝜃 和 𝜂 值。因此，在方程 17 中，我们可以使用单一标量来存储 Θ 。同样，在方程 18 中我们也可以加快运算速度。也就是说，当 𝜂 和 𝜃 在每个块内是可学习但时间不变的时候，该方程变成了一个线性时不变系统（ LTI ），可以通过全局卷积来计算（ Gu ， Goel 和 Re 2022 ）。在我们的实验中，我们将这些参数设置为 Token 的函数。然而，这种简化（即作为块的函数）可以成为未来以更高效的方式训练更大模型的研究兴趣所在。

图 2 ：内存作为上下文（ MAC ）架构。该架构包括三个分支：（ 1 ）核心，（ 2 ）上下文。

长期记忆（ long-term memory ），以及持久记忆（ persistent memory ）。 核心分支将相应的长期记忆和持久记忆与输入序列连接起来。接下来，注意力在序列上进行操作，并决定哪部分信息应存储在长期记忆中。在测试时，与上下文记忆相关的参数仍在学习中，与核心分支相关的参数负责上下文内的学习，而持久记忆的参数负责存储关于任务的知识并且是固定的。

3.3 持久性内存

我们的长期记忆也可以被视为上下文记忆，这意味着输出完全取决于上下文。 因此，除了我们的长期记忆外，我们还使用一组可学习但与输入无关的参数作为任务相关的记忆。这种类型的记忆在文献中被称为持久性或元记忆（ X. Dong 等， 2024 年； Sukhbaatar ， Grave 等， 2019 年）。给定 𝑁𝑝≥1 ，我们使用可学习的参数 𝑃=𝑝1 𝑝2 ...𝑝𝑁𝑝 ，并将其附加到我们序列的开头：即，给定一个上下文窗口大小为 𝑁 ，我们将输入修改为：

其中 || 表示连接。接下来，我们从三个方面讨论了持久性内存的动机：

记忆视角。 正如前面讨论过的，我们的神经长期记忆是一种上下文记忆，其中所有参数都是依赖于输入的。然而，一个有效的记忆系统还需要输入无关的参数来存储任务知识的抽象。也就是说，掌握一个任务需要记忆如何完成任务的知识，而这些参数负责存储这样的知识。

前馈网络视角。 在 Transformer 架构中，注意力模块之后是全连接层，据显示这些层类似于注意力权重，但具有数据无关的参数。也就是说， Sukhbaatar 、 Grave 等人（ 2019 年）表明，将全连接层中的 ReLU 替换为 Softmax 会产生类似注意力权重的结果，其中权重是数据无关的：

实际上，当 𝑊𝐾 和 𝑊𝑉 与 𝐾 和 𝑉 矩阵在注意力模块中作为输入无关时，它们的作用类似于 𝐾 和 𝑉 矩阵。希望持久的记忆权重具有相同的功能，这意味着在序列的第一部分使用它们将导致具有输入无关的注意力权重（ Sukhbaatar, Grave 等， 2019 年）。

技术观点 。带有因果掩膜的注意力对序列中的初始标记具有隐式偏见，因此注意力权重几乎总是对初始标记高度活跃，导致性能下降。从技术角度来看，序列开始处的这些可学习参数可以通过更有效地重新分配注意力权重来减轻这种影响（ Han 等， 2024 年； Xiao 等， 2024 年）。

记忆作为上下文（ MAC ）。我们将序列分割并在每个窗口中使用全因果注意力。再次，第一

𝑁𝑝 tokens 是持久性内存，下一个 𝑁𝑙 是长期记忆 tokens 。

内存作为门控（ MAG ）。我们使用滑动窗口注意力（ SWA ）作为短期记忆，我们的神经记忆模块作为长期记忆，通过门控结合在一起。

图 3 ：不同类型的泰坦的注意力蒙版。

4 如何融入记忆？

一个至今未解答的重要问题是：如何将设计的神经记忆有效而高效地融入到深度学习结构中？ 正如前文所讨论的，从记忆的角度来看， transformers 中的 K 和 V 矩阵可以被解释为一个联想记忆模块。由于它们对依赖关系的准确建模以及有限的上下文窗口，我们将其解释为短期记忆模块，关注当前的上下文窗口大小。另一方面，我们的神经记忆能够不断学习来自数据并将其存储在权重中，能够扮演长期记忆的角色。在本节中，我们旨在通过提出 Titans 的三个不同变体来回答上述问题。在我们的实验证明中，我们展示了每个变体都具有自己的优点 / 缺点，并且在非常长的上下文中也可以在效率和效果之间实现权衡。

4.1 记忆作为背景

在第一种架构设计中（见图 2 ），我们将内存视为当前信息的上下文。也就是说，给定一个长序列 𝑥 ∈ R𝑁×𝑑in ，我们首先将序列分成固定大小的段 S(𝑖) ，其中 𝑖 = 1, . . . , 𝑁 /𝐶 。给定传入的段落 S(𝑡) ，我们将其视为当前上下文和其以前的段落视为历史信息。因此，让 M𝑡−1 代表段落 S(𝑡) 之前的长期记忆状态，我们将输入的上下文作为查询用于从长期记忆 M𝑡−1 中检索相应的信息。也就是说，我们检索与 S(𝑡) 对应的过去信息为：

q𝑡 =S(𝑡)𝑊𝑄 。接下来，我们将使用这些历史信息以及我们持久性内存参数作为输入序列传递给注意力模块。

整个序列上的注意力图结构如图 3a 所示。然后我们使用 𝑦𝑡 来更新下一个片段和最终输出的长期记忆模块：

在上述过程中，请注意，我们通过前向传播来更新 M𝑡−1 的权重。

这种架构有两个关键优势： (1) 注意力机制结合历史和当前上下文，能够决定在给定当前数据时是否需要长期记忆的信息。 (2) 注意力模块有助于

图 4 ：记忆作为门（ MAG ）架构。这个架构同样由（ 1 ）核心、（ 2 ）语境记忆和（ 3 ）持久记忆三个分支组成。然而，它仅将持久记忆纳入上下文，并使用一个门控机制将记忆与核心分支结合起来。在测试时，行为与图 2 相同。

长期记忆只能存储来自当前上下文的有用信息 。也就是说，每个片段中并不是所有的标记都是有用的，而且记住所有标记可能导致内存溢出。因此，注意机制帮助记忆理解哪些信息是有用的，更好地管理内存容量。（ 3 ）在测试时：（ i ）持久性记忆参数是固定的，因为它们编码了关于任务的知识，不应该更改；（ ii ）注意模块的权重是上下文学习器；（ iii ）长期记忆模块仍在学习（记忆）测试时的信息。也就是说，我们甚至在测试时更新神经记忆的权重，因为这些权重编码了很长时间的抽象信息。

4.2 门控记忆

在下一个变体中（见图 4 ），在一个分支中，我们直接使用输入数据来更新长期记忆，在第二个分支中，我们使用滑动窗口注意力（ SWA ）：

滑动窗口注意力机制带有前缀，记作 SW-Attn∗( 参见图 3b) 。需要注意的是，与以前的设计不同，我们不再对输入数据进行分段。同时，我们滥用符号并使用 M(𝑥) 来指代序列中所有 Token 递归完成后记忆的最终输出。在上述方程中， ⊗ 可以是任何非线性门控操作。在我们的实验中，

我们使用可学习的矢量值权重来对输出 𝑦 和 M(𝑥˜) 进行归一化，然后接着使用非线性函数 𝜎(.) 。

这个设计的整体关注掩码在图 3b 中显示。在这个设计中，滑动窗口注意力作为精确的短期记忆，而神经记忆模块作为模型的衰减记忆。这种架构设计也可以看作是一个多头架构，其中头部的结构不同（ X. Dong 等人， 2024 年）。

4.3 内存作为一层

最后一种变体使用深度神经网络的神经内存层（ Memory As a Layer ， MAL ）（见图 5 ）。这种架构设计在文献中更为常见，其中混合模型将递归模型与完全或滑动窗口的注意力机制相叠加。给定输入 𝑥 ，我们有：

图 5 ：内存作为一层（ MAL ）架构。在这种架构中，内存层负责在注意力模块之前压缩过去和当前的上下文。

SW-Attn 表示滑动窗口注意力。 这种设计的主要缺点是模型的能力受到每个层的限制，因此无法利用注意力和神经记忆模块的互补数据处理优势。在我们的实验中，为了评估这种设计中的记忆，我们使用了类似 H3 （ D. Y. Fu 等人， 2023 年）的架构，其中我们用我们的神经记忆模块（ LMM ）替换了序列模型。

记忆无需注意力 。尽管在上面我们讨论了 MAL 作为 LMM 和注意力的顺序组合，但 MAL 的一个简单变体是将 LMM 作为一个没有任何注意力的序列模型。从记忆的角度来看，就像我们在第一节中讨论的那样，我们期望记忆系统的每个部分都能独立工作，即使其他组件受到干扰。因此，即使没有短期记忆（即注意力），长期记忆模块仍然应该是一个强大的模型。在我们的实验证明中，我们将这种变体称为 LMM 或 Titans （ LMM ）。关于 Titans 和其他现代循环模型的连接，我们在附录 C 中提供了额外的讨论。

4.4 架构细节

为了简化和展示的目的，我们避免讨论实现细节，比如使用残差连接、线性层门控和归一化。 在所有块中，我们使用残差连接。在我们的实现中，我们使用 SiLU(.) 激活（ Elfwing, Uchibe 和 Doya 2018 ）作为计算查询、键和值的非线性激活函数，并使用 ℓ2 范数对查询和键进行归一化。

卷积。 在最近的现代线性循环模型（ Gu 和 Dao 2024; S. Yang ， Kautz 和 Hatamizadeh 2024 ）之后，我们在每个查询、键和值投影后都加入了 1D 的深度可分离卷积层。尽管对性能影响不大，但这些 1D 卷积已经显示出性能的提升，并且计算效率也很高。

门控。 我们还遵循最近使用归一化和门控的架构，该架构在最终的输出投影之前使用线性层（ Mehta 等人， 2023 年）。

定理 4.1. 与变压器、对角线线性递归模型和 DeltaNet 相反，它们都限于 TC0 （ Merrill ， Petty 和 Sabharwal 2024 ），泰坦能够解决超出 TC 0 范围的问题，这意味着泰坦在状态跟踪任务中理论上比变压器和大多数现代线性递归模型更具表达能力。

5 实验

接下来，我们评估 Titans 及其变体在语言建模、常识推理、大海捞针、 DNA 建模和时间序列预测任务中的表现 。更详细地，在本节中，我们回答以下实证问题：（ 1 ）相对于基线，在下游任务中 Titans 的表现如何？（见 §5.2 、 §5.6 和 §5.7 ）；（ 2 ） Titans 的实际上下文长度是多少？（见 §5.3 和 §5.4 ）；（ 3 ） Titans 在上下文长度方面如何扩展？（见 §5.8 ）；（ 4 ）记忆深度如何影响性能和效率？（见 §5.5 ）；以及（ 5 ） Titans 的每个组件在其性能中的贡献是什么？（见 §5.9 ）。

5.1 实验设置

模型。 在我们的实验中，我们关注的是 Titans 的三个变种，我们将其称为：带有 （ 1 ）内存作为上下文（ MAC ）的 Titans ，（ 2 ）带有内存作为门（ MAG ）的 Titans ，以及（ 3 ）带有内存作为层（ MAL ）的 Titans ，以及（ 4 ）仅有神经内存模块。 我们将长期记忆作为一个独立的模块使用的原因是基于我们对学习的定义。正如在第 1 节中讨论的那样，我们将学习定义为获取有效和有用的记忆的过程。因此，我们期望我们的长期记忆能够在没有注意力的情况下从数据中有效地学习。对于这些模型的每一个，我们考虑了四个不同的规模：（ i ） 170M ，（ ii ） 340M ，（ iii ） 400M 和（ iv ） 760M 参数。前三个模型是使用从 FineWeb-Edu 数据集（ Penedo 等人， 2024 年）采样的 15B 个 Token 进行训练，而最后一个模型则是使用来自同一数据集的 30B 个 Token 进行训练。

基准。我们将我们的模型与最先进的线性循环模型、变压器模型和混合模型进行比较。

(recurrent + attention) 更具体地说，在语言任务中，我们与 Transformer++ (Touvron et al. 2023) 进行了比较。

RetNet （孙宇涛等， 2023 年）、门控线性注意力（ GLA ）（杨帅、王斌、沈等， 2024 年）、 Mamba （顾道， 2024 年）、 Mamba2 （道和顾， 2024 年）、 DeltaNet （杨帅、王斌、张宇等， 2024 年）、 TTT （孙宇等， 2024 年）和门控 DeltaNet （杨帅、考茨和哈塔米扎德， 2024 年）。在海量数据中寻找少量关键数据的任务中，我们还与 GPT4 （ Achiam 等， 2023 年）、 Llama3 与 RAG （ Touvron 等， 2023 年）、 RecurrentGemma2-9B （ Botev 等， 2024 年）和 Mistral （姜等， 2023 年）模型进行了比较，这些模型都包含在基准测试中（ Yuri Kuratov 等， 2024 年）。在时间序列任务中，我们与基于 Mamba 的模型（ Behrouz 、 Santacatterina 和 Zabih ， 2024 年）、基于 Transformer 的模型（刘扬等， 2023 年；聂等， 2022 年；张云豪和闫， 2023 年）和线性模型（ Das 等， 2023 年；李等， 2023 年；吴等， 2023 年；曾等， 2023 年）进行比较。

在训练中，我们遵循了 S. Yang ， Kautz 和 Hatamizadeh （ 2024 ）的培训程序，并使用了 32K 的 LLama 2 分词器，训练长度为 4K 个标记。我们采用了学习率为 4𝑒-4 的 AdamW 优化器，采用余弦退火调度，批量大小为 0.5M 标记，权重衰减为 0.1 。

5.2 语言建模

我们首先关注语言建模和常识推理任务中的困惑。 表 1 报告了 Titans' 变体以及基线模型，分别采用 340M 、 400M 和 760M 三种不同规模的结果。在不涉及混合模型的情况下，包括 Transformer++ 在内，我们的神经记忆模块在困惑度和准确性方面表现最佳。将我们的神经记忆模块与 TTT 进行比较，后者也是基于梯度的递归模型，可以展示出我们的权重衰减以及动量的重要性。正如前面讨论的那样，权重衰减可以被解释为一种门控机制，用于在需要时遗忘过去的数据。此外，动量可以通过为惊喜度量提供额外的存储器来帮助我们更好地管理记忆。虽然一些基线模型也利用了门控机制，例如 Mamba 、 Mamba2 和 Gated DeltaNet ，但我们的神经记忆模块的卓越性能表明了惊喜机制和具有深度和非线性记忆的重要性。我们在 5.5 节进一步讨论后者。

比较混合模型，我们发现 Titan 的三种变体（ MAC ， MAG 和 MAL ）胜过 Samba （ Mamba + attention ）和 Gated DeltaNet-H2 （ Gated DeltaNet + attention ）。 我们将 Titan （ MAL ）的优越性能归因于神经记忆模块的强大，因为其架构设计和使用的注意力模块都是相同的。比较 Titan （ MAG ）和（ MAC ），我们发现虽然它们的性能相近，但 MAC 在处理数据中更长依赖性时表现更好。有趣的是， MAG 和 MAC 都胜过 MAL 变体，由于它们使用相同的模块，我们将这归因于这些模型的架构设计。这一发现尤为重要，因为文献中目前的混合模型（除了 Hymba （ X. Dong 等， 2024 年））都使用 MAL 风格的循环模型和注意力的组合。

5.3 大海捞针

将模型扩展到更长的上下文窗口并不总是等效于对于非常长的序列有效 （ Hsieh 等， 2024 年）。针对大海捞针（ NIAH ）任务旨在衡量模型的实际有效上下文长度。在这个任务中，我们评估模型在从长的干扰文本中检索信息（即 “ 针 ” ）方面的表现（即，表 1 ： Titans 和基于递归和 Transformer 的基准在语言建模和常识推理任务上的性能。混合模型用 ∗ 标记。简单和混合模型中的最佳结果已突出显示）。

在这部分中，我们使用来自 RULER 基准的单一 NIAH （ S-NIAH ）任务（ Hsieh 等， 2024 ），并在长度为 2K 、 4K 、 8K 和 16K 的序列上评估 Titans 和基线模型。结果报告于表 2 中。神经记忆模块在所有三个任务中相对于基线模型取得了最好的结果。我们将这种卓越的性能归因于 Titans 与现有序列模型的三个关键差异：（ 1 ）与 TTT 相比，我们的神经记忆可以通过使用动量和遗忘机制（即权重衰减）更好地处理内存容量。因此，随着序列长度的增加，神经记忆的性能不会下降，呈现一致的趋势；（ 2 ）与具有门控（遗忘）机制的 Mamba2 相比， Titans 具有深层非线性记忆，从而实现更好的内存管理。另外，与我们的神经记忆和 DeltaNet 不同， Mamba2 无法删除记忆，因此表 2 ： Titans 和基线模型在 RULER 基准的 S-NIAH 任务上的性能。在简单模型和混合模型中，最佳结果已经被突出显示。

图 6 ： Titans 在 BABILong 基准测试中的表现。 Titans （ MAC ）胜过所有基准模型，包括极其庞大的模型，例如 GPT4 。

我们可以看到，在增加序列长度时，性能明显下降； (3) 与 DeltaNet 相比，虽然它能够使用增量规则去除记忆，但它无法抹去记忆，缺乏遗忘机制。最后，如预期所见，当使用 Titans 变种时，我们可以看到与 MAC 对应的最好结果相当或更好。

5.4 BABILong 基准测试

在前一节中，我们讨论了一个简单的 NIAH 任务的结果，其中需要检索一个单独的针头。尽管 Titans 相比基准表现更好，但它们在非常长的序列上的优势仍然被隐藏起来。为此，在本节中，我们使用了来自 BABILong 基准测试（ Yuri Kuratov 等人， 2024 年）的一个更难的任务，其中模型需要在极长的文档中推理出分布的事实。我们遵循基准测试中的原始实验设置和训练过程。有两个设置： (1) 少样本设置，我们使用大型预训练模型； (2) 微调设置，我们微调 Titans 的 MAC 变体，以便与其他微调基线进行比较。少样本设置的结果报告在图 6a 中。在这个设置中，我们可以看到 Titans 胜过所有其他基线，例如 Mamba2.8B （ Gu and Dao 2024 ）， RWKV-6-7B （ Peng ， Goldstein 等人， 2024 年）， RecurrentGemma-9B （ Botev 等人， 2024 年）， Gemma-9B （ Team 等人， 2024 年）， Llama3.1-8B （ Touvron 等人， 2023 年）， GPT-4 和 GPT4o-mini （ Achiam 等人， 2023 年）。这些结果是在 Titans （ MAC ）的参数数量远少于基线的情况下取得的。

在微调设置中，我们将 Titans 的小规模微调版本（ MAC ）与以下进行比较：（ i ）像 Mamba （ Gu 和 Dao 2024 ）、 RMT （ Bulatov ， Yury Kuratov 和 Burtsev 2022 ）这样的具有几乎与 Titans 相同参数数量的小模型的微调版本；（ ii ）带有检索增强生成（ RAG ）的大模型，例如 Llama3.18B （ Touvron 等人 2023 ）；以及（ iii ）像 GPT-4 （ Achiam 等人 2023 ）、 GPT4o-mini 、 Qwen2.5-72B （ A. Yang 等人 2024 ）和 Llama3.1-70B （ Touvron 等人 2023 ）这样的极大模型。基准结果由（ Yuri Kuratov 等人 2024 ）报告。 Titans 和基准的结果在图 6b 中报告。 Titans 胜过所有模型，甚至胜过像 GPT4 这样的极大模型。此外，与基于 Transformer 的带有记忆模型（如 RMT ）相比， Titans 显示出更好的性能，主要是由于其强大的内存。也就是说， RMT 将历史数据压缩为 16 大小的矢量值内存，而 Titans 具有上下文在线内存学习器，能够将过去编码到模型的参数中。有趣的是，即使是

(a) 170M 参数 (b) 360M 参数 (c) 760M 参数

图 7 ：记忆深度对困惑度的影响。更深的长期记忆有利于更长序列的扩展。

表 3 ：长期预测的表现。最佳结果已经突出显示。

Llama3.1-8B 模型与 RAG 相比，在参数减少了约 70 倍的情况下表现更差。

5.5 深度记忆的效应

在本节中，我们评估了深度记忆在墙钟训练时间和模型性能方面的影响。 为此，我们专注于我们神经记忆模块的不同变体，其中 𝐿M = 1, 2, 3, 4 。我们还将 Mamba 作为模型性能的基准。为了公平比较，我们对所有模型使用相同的训练过程，并在