谷歌Titans模型论文解析，Transformer迎来变革拐点—

一、引入 Titans 模型

我们将深入探讨谷歌研究院的一篇新论文《Titans: Learning to Memorize at Test Time》，该论文介绍了一种名为 Titans 的新模型架构。
在这里插入图片描述

Titans 在缓解 Transformer 二次方成本问题的同时，展现出了令人期待的成果。Titans 模型的设计灵感来源于人类大脑的记忆运作方式。论文中有一句有趣的话提到，记忆是一种基本的心理过程，是人类学习不可或缺的组成部分。如果没有正常运作的记忆系统，人类和动物将局限于基本的反射和刻板行为。

在这里插入图片描述

二、深度神经长期记忆模块

Titans 论文的一个关键贡献是深度神经长期记忆模块（deep neural long-term memory module）。我们先来了解一下它是什么，然后再探讨它是如何融入 Titans 模型的。

与循环神经网络中将记忆编码为固定向量不同，神经长期记忆模块本身就是一个模型，是一个多层神经网络，它将过去的历史抽象编码到自身的参数中。要训练这样一个模型，一种思路是训练它记忆训练数据。然而，众所周知，记忆会限制模型的泛化能力，可能导致性能不佳。

三、避免过拟合的记忆机制

研究人员设计了一种巧妙的方法，来创建一个能够记忆，但又不会在训练数据上过拟合的模型。这种方法的灵感来自人类记忆的一个类比。当我们遇到令人惊讶的事件时，我们更有可能记住它。神经长期记忆模块的学习过程正是基于这一理念设计的。

3.1 对 “惊讶” 进行建模

Updating the neural memory using a surprise element

通过论文中的上述定义，我们可以了解研究人员是如何对 “惊讶” 进行建模的。 Mt代表时刻t的神经长期记忆模块。它通过上一时刻的参数以及一个被建模为梯度的 “惊讶” 元素进行更新。如果梯度较大，说明模型对输入更 “惊讶”，这会导致模型权重得到更显著的更新。

然而，这个定义仍不理想，因为模型可能会错过 “惊讶” 时刻之后紧接着发生的重要信息。

3.2 对 “过去的惊讶” 进行建模

从人类的角度来看，一个令人惊讶的事件不会在很长一段时间内持续让我们感到惊讶，尽管它仍然令人难忘。我们通常会适应这个令人惊讶的事件。不过，该事件可能足够令人惊讶，从而在很长一段时间内吸引我们的注意力，使我们记住整个时间段。

Updating the neural memory using both past and momentary surprise

通过论文中的上述定义，我们可以了解到改进后的建模方式，其中包括对 “过去的惊讶” 的建模。现在，我们使用上一时刻的权重状态以及一个 “惊讶” 分量St来更新神经长期记忆的权重。“惊讶” 分量现在是随时间测量的，它由带有衰减因子的上一时刻 “惊讶” 和我们在上一节讨论过的当下 “惊讶” 组成。

还有一个尚未建模的因素是遗忘。

3.3 对遗忘进行建模

Allowing the model to forget past information

在处理非常长的序列（例如数百万个标记 tokens）时，管理哪些过去的信息应该被遗忘至关重要。通过论文中的上述定义，我们可以看到最终的建模方式。这些定义与上一节的定义相同，只是我们添加了一个自适应遗忘机制，用 alpha 表示，也称为门控机制（gating mechanism）。这使得记忆能够遗忘不再需要的信息。

3.4 损失函数

Neural memory loss function

损失函数由上述公式定义。该损失旨在通过将过去的数据存储为键值对来模拟联想记忆，并训练模型在键和值之间进行映射。与 Transformer 类似，线性层将输入投影为键和值。然后，损失值loss衡量记忆模块学习键值对之间关联的程度。

需要说明的是，模型不是一次性处理整个序列，而是逐步处理，在其权重中积累记忆信息。

四、Titans 架构

4.1 Titans 架构 1—— 记忆作为上下文（MAC）

论文提出了几种不同的架构。第一种称为记忆作为上下文（MAC）。我们可以通过论文中的下图来了解它。

Titans Memory as a Context Architecture

该模型由 3 个组件组成：持久记忆、核心和上下文记忆。给定一个长输入序列，我们将其划分为较小的序列。核心组件中显示的序列实际上是一个分块后的序列。

持久记忆

我们首先引入持久记忆，在序列开头添加与数据无关的可学习权重。实际上，这些是可学习的标记，用于存储不受输入影响的全局信息。序列还会经过一个注意力块，在序列开头添加可学习标记有助于避免注意力分散，因为由于偏差，注意力权重往往会过度关注初始标记。英伟达的 Hymba 模型添加可学习元标记的做法也取得了成功，我后续会介绍。

上下文记忆

上下文记忆（Contextual memory）由深度神经长期记忆模块构成，因为它依赖于上下文。我们使用神经长期记忆模块的前向传递，以分块后的输入序列作为输入，从长期记忆中检索记忆标记（memory tokens）。由于即使在测试阶段神经记忆也在不断更新，当我们处理长序列的一个分块时，神经记忆已经更新了同一序列先前分块的信息。检索到的记忆标记会在持久记忆标记之后添加到输入序列中。

核心组件

核心组件将所有信息整合在一起，使用一个注意力块，该注意力块的输入是扩展后的序列，其中包含来自持久记忆、上下文记忆和输入本身的信息。通过这种方式，模型可以利用多种类型的记忆。注意力块的输出用于更新神经记忆。注意力机制有助于根据给定的输入确定是否使用长期记忆。此外，注意力机制有助于长期记忆仅存储当前上下文中有用的信息。最终输出基于注意力块的输出和神经记忆的输出确定。

4.2 Titans 架构 2—— 记忆作为门控（MAG）

Titans 的下一个架构版本称为记忆作为门控（MAG: Memory as a Gate），我们可以通过论文中的下图来了解它。

Titans Memory as a Gate Architecture

在这个版本中，我们同样有代表持久记忆、核心和上下文记忆的 3 个分支。与前一个版本不同的是，这里的序列没有分块，输入序列是完整的输入序列。这是通过在注意力块中使用滑动窗口注意力实现的。持久记忆的可学习权重同样添加到序列开头。但与之前不同的是，神经记忆不会为注意力块的上下文提供数据。相反，神经记忆根据输入序列进行更新，其输出通过门控机制与核心分支相结合。

4.3 Titans 架构 3—— 记忆作为层（MAL）

Titans 架构的第三个变体称为记忆作为层（MAL）。我们可以通过论文中的下图来了解它。

Titans Memory as a Layer Architecture

与前一个版本类似，这里的序列也没有分块，并且我们使用滑动窗口注意力。

在这个版本中，我们将神经记忆用作模型的一层，输入序列与可学习权重首先经过神经记忆，然后再经过注意力块。这种设计允许堆叠多个神经记忆模块和注意力块，就像通常堆叠 Transformer 层一样。然而，这种顺序设计使得模型的能力受到每一层能力的限制。这与 Hymba 论文中的观察结果类似，出于同样的原因，Hymba 论文中的不同组件采用并行而非顺序放置。