谷歌发布Transformer架构的继任者Titans: 提出了一种新的神经长期记忆模块，它可以学习记忆历史背景，并帮助注意力在利用过去长期信息的同时关注当前背景。这种神经记忆具有快速并行训练的优势，同时还能保持快速推理

近年来，循环模型和注意力机制在深度学习领域得到了广泛应用。循环模型旨在将数据压缩成固定大小的记忆（隐藏状态），而注意力机制则允许模型关注整个上下文窗口，捕捉所有标记的直接依赖关系。然而，这种更精确的依赖关系建模带来了二次方的计算成本，限制了模型的上下文长度

谷歌发布《泰坦：在测试时学习记忆》（Titans: Learning to Memorize at Test Time）的论文提出了一种新型的神经长期记忆模块，该模块能够在利用长远历史信息的同时，让注意力机制专注于当前上下文。该神经记忆的优势在于可以快速并行训练，并保持快速推理。论文指出，由于上下文有限但依赖关系建模精确，注意力机制可以作为短期记忆；而神经记忆由于其记忆数据的能力，可以作为长期、更持久的记忆。基于这两个模块，论文引入了一种新的架构族——泰坦（Titans），并提出了三种变体，以探讨如何有效地将记忆融入架构中

核心创新

这篇论文的核心创新在于提出了一个能够在测试时学习记忆的神经长期记忆模块。这个模块的工作方式如下：

• • 记忆的获取： 该模块将训练过程视为在线学习问题，旨在将过去的信息压缩到其参数中。受人类记忆启发，该模块将“违反预期”的事件（即令人惊讶的输入）视为更值得记忆的。它通过计算神经网络相对于输入的梯度来衡量输入的“惊讶程度”，并使用这一指标来更新记忆。

• • 遗忘机制： 为了解决有限记忆的问题，该模块引入了一种自适应的遗忘机制，该机制考虑了记忆大小和数据惊讶程度，从而更好地管理记忆。

• • 记忆的结构： 论文探索了不同的记忆结构，发现深度记忆模块（即使用多层感知机）比线性模型更有效。

• • 记忆的检索： 该模块通过简单的正向传递（不更新权重）来检索与查询相对应的记忆。

Titans架构

基于长期神经记忆模块，论文提出了泰坦架构，该架构包含三个分支：

• • 核心分支（Core）： 使用注意力机制进行数据处理，关注有限的上下文窗口。

• • 长期记忆分支（Long-term Memory）： 使用神经长期记忆模块来存储和回忆历史信息。

• • 持久记忆分支（Persistent Memory）： 使用可学习但不依赖于数据的参数来编码任务相关知识。

论文提出了三种不同的泰坦变体：

1. 1. 上下文记忆（Memory as a Context，MAC）： 将长期记忆视为当前信息的上下文，使用注意力机制融合这些信息。

2. 2. 门控记忆（Memory as a Gate，MAG）： 使用门控机制将长期记忆与核心分支的信息融合。

3. 3. 层式记忆（Memory as a Layer，MAL）： 将长期记忆模块作为深度神经网络的一层。

实验结果

论文在语言建模、常识推理、长距离依赖、时间序列预测和DNA建模等任务上进行了广泛的实验，实验结果表明：

• • 超越基线： 泰坦架构在所有基准测试中都优于现代循环模型及其混合变体。

• • 优于Transformer： 泰坦架构在相同上下文窗口大小下优于Transformer模型，并且在长上下文任务中表现出更具竞争力的性能

• • 可扩展性： 与Transformer不同，泰坦架构可以有效地扩展到大于2M的上下文窗口大小

• • 深度记忆的重要性： 更深的长期记忆模块可以在更长的序列中获得更好的性能。

• • 组件的贡献： 论文进行了消融研究，验证了泰坦架构中每个组件的积极贡献，其中权重衰减（Weight Decay）、动量（Momentum）、卷积（Convolution）和持久记忆（Persistent Memory）贡献最大

Titans 已使用 PyTorch 和 JAX 实现，谷歌计划很快公开用于训练和评估模型的代码

参考：

https://arxiv.org/pdf/2501.00663v1