Titans：结合注意力机制用于短期记忆，以及神经记忆用于长期记忆的 Transformer 模型

Titans：结合注意力机制用于短期记忆，以及神经记忆用于长期记忆的 Transformer 模型！

🕙发布时间：2025-02-25

更多LLM架构文章：LLM架构专栏
近日热文：
1. 全网最全的神经网络数学原理（代码和公式）直观解释
2. 大模型进化史：从Transformer到DeepSeek-R1的AI变革之路
3. 2W8000字深度剖析25种RAG变体：全网最全~没有之一
4. 3W6000字了解大模型LLM：部署、优化与框架
知乎【柏企】
公众号【柏企科技说】【柏企阅文】

论文

Titans 论文地址
Titans 是一种全新的神经架构，它将注意力机制与长期记忆模块相结合，能够高效处理大上下文窗口。注意力机制捕捉短期依赖关系，而神经记忆则存储历史上下文以供长期使用。这种设计实现了快速的并行训练和推理，能够扩展到超过 200 万个标记，并且在语言建模、常识推理、基因组学和时间序列分析等任务中，精度有所提高。Titans 的表现优于 Transformer 和现代线性循环模型，在需要从大型数据集中检索特定信息的任务中表现尤为突出。

预备知识

Transformer 依赖于注意力机制，该机制计算标记之间的依赖关系，但对于长序列而言，内存和计算成本很高。为了解决这个问题，人们开发了诸如稀疏化、softmax 近似和线性（基于内核）注意力等高效注意力方法，这些方法降低了复杂度并提高了吞吐量。线性注意力模型用内核函数取代 softmax，支持循环公式，实现高效推理。然而，它们的累加内存结构在处理长上下文时可能会导致溢出。

为了解决这一问题，出现了两种方法：一种是添加遗忘机制，比如在 GLA、LRU、Griffin、xLSTM 和 Mamba2 等模型中使用自适应门；另一种是受 Delta 规则启发改进写入操作，在添加新值之前删除过去的值。最近的模型还注重可并行化训练，并引入遗忘门来增强内存管理。

测试时学习记忆

长期记忆

论文作者探讨了神经长期记忆模块的设计，该模块能够以动态的、依赖数据的方式有效地记忆和遗忘信息。与传统模型因记忆训练数据而有过拟合风险不同，这种方法使用在线元学习框架，在测试时平衡记忆和泛化能力。该模型基于一个“惊喜度量”来更新其记忆，这个度量衡量输入数据与过去数据的差异程度，同时考虑了瞬时惊喜和过去的惊喜（类似于带动量的梯度下降）。

一种由门控参数控制的自适应遗忘机制，通过决定保留或丢弃多少过去的信息来管理内存容量。该架构使用多层感知器（MLP）来捕捉非线性关系。键和值由输入数据生成，记忆模块使用在线线性回归目标来学习它们之间的关联。

信息通过查询投影前向传递来检索，且不会更新记忆权重。

如何并行化长期记忆训练

长期记忆模块作为一个元模型，使用带动量的梯度下降和权重衰减进行训练。为了在硬件加速器上优化性能，通过张量操作和矩阵乘法对训练过程进行并行化。序列被划分为多个块，在每个块内应用小批量梯度下降。

学习率、权重衰减和动量衰减等参数可以根据数据进行调整，或者简化为在每个块内保持恒定，这样在一定程度上牺牲了一些表现力，但能加快训练速度。当参数在块内不随时间变化时，系统表现为线性时不变系统，可以使用全局卷积进行计算。

持久记忆

长期记忆模块作为上下文记忆，完全依赖输入数据。为了补充这一点，论文作者使用了一组可学习的、与输入无关的参数（持久记忆或元记忆）来存储特定任务的知识。这些参数被附加到输入序列的开头。

从内存的角度来看，持久记忆存储任务知识的抽象，有助于更好地掌握任务。从前馈网络的角度来看，这些参数类似于 Transformer 前馈层中的注意力权重，作用类似于键值对，但与输入无关。从技术角度来看，它们有助于减轻注意力机制对序列中初始标记的偏向，更有效地重新分配注意力权重，从而提高性能。

如何整合记忆？

论文作者提出了三种将记忆整合到 Transformer 架构中的方法。

记忆作为上下文

第一种架构通过将长序列划分为固定大小的段来处理长序列。对于每个输入段，模型使用该段作为查询，从长期记忆中检索相关的历史信息。检索到的信息与持久记忆参数相结合，被输入到注意力模块中。

注意力模块输出一个优化后的表示，用于更新模型的最终输出并更新长期记忆。这种设计有三个优点：

1. 注意力机制决定是否需要历史上下文。
2. 注意力有助于识别和仅存储有用信息，防止内存溢出。
3. 持久记忆保持固定以保留任务知识，注意力权重在上下文中自适应调整，并且长期记忆在测试期间继续学习，更新其对过去信息的抽象。

门控记忆

在这种架构中，输入数据通过两个并行分支进行处理。一个分支直接更新长期记忆，而另一个分支以前缀的形式使用持久记忆参数应用滑动窗口注意力（SWA）。与前一种设计不同，这里的输入不进行分段。

SWA 的输出和长期记忆的输出通过非线性门控机制进行组合。SWA 充当精确的短期记忆，而神经记忆模块作为逐渐衰减的长期记忆。这种设置可以看作是一种多头架构，每个分支都充当一种独特的注意力机制。

内存作为一层

在这种架构中，神经记忆模块用作深度神经网络中的一层。输入与持久记忆参数相结合，首先由神经记忆模块处理，其输出随后通过 SWA。这种顺序设置很常见，但限制了模型同时利用注意力和记忆模块互补优势的能力。

一种更简单的变体完全不使用注意力，仅依靠长期记忆模块（LMM）作为序列模型。这种方法测试了没有短期记忆组件的记忆系统的独立有效性。

架构细节

1. 所有模块都有残差连接；
2. 激活函数为 SiLU；
3. 每次投影后使用 1D 深度可分离卷积神经网络（CNN）。

实验

Titans 的神经记忆模块在不同模型规模（3.4 亿、4 亿、7.6 亿参数）的语言建模和常识推理任务中，表现均优于基线模型。

在“大海捞针”（NIAH）任务中，Titans 在序列长度为 2K、4K、8K 和 16K 时，表现优于所有基线模型。

## 推荐阅读
1. DeepSeek-R1的顿悟时刻是如何出现的？ 背后的数学原理
2. 微调 DeepSeek LLM：使用监督微调（SFT）与 Hugging Face 数据
3. 使用 DeepSeek-R1 等推理模型将 RAG 转换为 RAT
4. DeepSeek R1：了解GRPO和多阶段训练
5. 深度探索：DeepSeek-R1 如何从零开始训练
6. DeepSeek 发布 Janus Pro 7B 多模态模型，免费又强大！

本文由mdnice多平台发布