谷歌提出创新神经记忆架构，突破Transformer长上下文限制

让 AI 模型拥有人类的记忆能力一直是学界关注的重要课题。传统的深度学习模型虽然在许多任务上取得了显著成效，但在处理需要长期记忆的任务时往往力不从心。就像人类可以轻松记住数天前看过的文章重点，但目前的 AI 模型在处理长文本时却经常丢三落四，这种差距一直是困扰研究人员的关键问题。

近日，谷歌研究院的研究团队在这一领域取得了重要突破，Ali Behrouz、钟沛林（Peilin Zhong）和 Vahab Mirrokni 在预印本平台 arXiv 发表论文，提出了一种名为 Titans 的创新神经记忆架构，其核心是一个能在测试时动态学习和更新的神经长期记忆模块。

目前主流的 Transformer 架构在处理文本时面临着明显的局限。虽然它能够准确捕捉短距离的文本关系，但其注意力机制的计算复杂度会随着序列长度呈二次增长，这使得处理长文本时的计算成本变得难以承受。为解决这一问题，研究者们提出了 RetNet、Mamba 等线性递归模型。这些模型虽然将计算复杂度降至线性水平，但由于需要将信息压缩到固定大小的状态中，在处理超长序列时常常出现严重的信息丢失。

论文作者 Ali Behrouz 在 X 上解释了 Titans 的设计理念：“我们从人类记忆的视角重新思考这个问题。人类大脑会优先记住那些违反预期的事件，但有趣的是，一个事件虽然可能值得长期记住，它的'惊讶度'却会随时间推移而减弱。”基于这一对人类记忆特性的观察，研究团队开发出了 Titans 独特的记忆更新机制。

具体来说，Titans 包含三个核心组件，分别对应人类记忆系统的不同方面。

首先是基于注意力机制的核心短期记忆，负责处理当前上下文窗口内的信息，类似于人类的工作记忆。

其次是 Titans 最具创新性的神经长期记忆模块，它能在测试时动态学习和更新记忆，通过神经网络的梯度来衡量信息的重要性。这个模块将“惊讶度”分为瞬时惊讶度（momentary surprise）和过去惊讶度（past surprise）两个维度，前者衡量当前输入与已有记忆的差异程度，后者评估近期历史信息的重要性。其计算公式如下：

St = ηtSt-1 - θtl(Mt-1; xt)

其中，ηtSt-1 代表过去惊讶，它通过一个数据依赖的衰减率 ηt 来控制历史信息的保留程度。当系统判断当前上下文与过去相关时，ηt 会接近 1，保持对历史信息的重视；当需要切换到新的上下文时，ηt 会接近 0，允许系统“忘记”不再相关的历史信息。

而 -θtl(Mt-1; xt) 则代表瞬时惊讶，它通过计算损失函数关于输入的梯度来量化当前输入 xt 与已有记忆 Mt-1 之间的差异程度。θt 参数控制了系统对新信息的敏感度。这种设计让模型能够像人类一样，对违反预期的信息保持特别的关注。

在具体实现中，研究团队使用了一个关联记忆损失函数：

l(Mt-1; xt) = ||Mt-1(kt) - vt||^2

其中，输入 xt 会被转换为键值对：kt = xtWK 和 vt = xtWV。这种设计让模型能够学习和存储信息之间的关联关系，类似于人类记忆中的联想机制。

第三个组件是持久记忆，这是一组与数据无关的可学习参数，专门用于存储完成特定任务所需的基础知识，类似于人类的程序性记忆。

从效率角度来看，研究团队还特别优化了 Titans 的训练过程。Behrouz 在推文中提到：“我们通过 TTT（Sun et al., 2024）扩展了 mini-batch 梯度下降的可并行化对偶形式，并通过额外的矩阵乘法来实现权重衰减。”通过将 mini-batch 梯度下降的前向传播重新表述为矩阵乘法操作，他们实现了高效的并行训练。

然后，研究团队提出了三种将这些组件整合的变体架构：

• MAC（Memory as Context）：这种方案将记忆作为上下文信息处理，对输入序列进行分段，并使用过去的记忆状态来提取相应的记忆信息。

• MAG（Memory as Gate）：这种方案使用滑动窗口注意力机制作为短期记忆，同时使用神经记忆模块作为长期记忆，通过门控机制将两者结合。

• MAL（Memory as Layer）：这种方案将神经记忆作为深度神经网络的一个层，更接近传统的混合模型设计。

实验结果显示，Titans 在多个测试基准上都表现出色。

在语言建模任务上，拥有 760M 参数的 Titans(MAC) 在 WikiText 上达到了 19.93 的困惑度，显著优于同等规模的 Transformer++（25.21） 和 Mamba2（22.94）。在常识推理任务上，Titans 在包括 PIQA、HellaSwag、WinoGrande 等 9 个基准测试中的平均准确率达到 52.51%，超过了现有的最好成绩。

值得一提的是，在 BABILong 基准测试中，即使是参数量较小的 Titans 也展现出出色的性能。在需要在超长文档中进行推理的任务中，Titans(MAC) 不仅战胜了 Mamba2、RWKV 等现代模型，甚至超越了参数量大得多的 GPT-4 和 Llama3-70B。实验证明，Titans 能够有效处理超过 200 万个 token 的上下文窗口，这一突破将为长文本处理开辟新的可能。

研究团队表示，Titans 的 PyTorch 和 JAX 实现代码将很快开源。随着这一技术的进一步发展和应用，我们有望看到更多能够处理超长文本的高效 AI 系统涌现，这对于文档分析、长文本理解、知识检索等领域都具有重要意义。

当然，这项研究仍有进一步探索的空间。如何设计更高效的神经记忆架构、如何在更大规模的模型中应用这一技术，都是未来值得关注的方向。但毫无疑问，Titans 的出现为解决 AI 系统的长期依赖问题提供了一个极具前景的新方向。

参考资料：

1.https://arxiv.org/abs/2501.00663

2.https://x.com/behrouz_ali/status/1878859673555624373

运营/排版：何晨龙