主要观点总结
本文介绍了谷歌研究院的研究团队在AI模型的长期记忆能力方面取得的突破。他们提出了一种名为Titans的创新神经记忆架构，该架构包含三个核心组件，分别对应人类记忆系统的不同方面。Titans架构具有动态学习和更新的神经长期记忆模块，能够处理长文本任务。实验结果显示，Titans在多个测试基准上表现出色，有望为长文本处理开辟新的可能。
关键观点总结
关键观点1: 谷歌研究院团队提出了名为Titans的创新神经记忆架构，具有动态学习和更新的能力。
Titans架构包括基于注意力机制的核心短期记忆、神经长期记忆模块和持久记忆三个核心组件，分别对应人类记忆系统的不同方面。
关键观点2: Titans架构通过独特的记忆更新机制，能够处理长文本任务，并在多个测试基准上表现出色。
实验结果显示，Titans在常识推理任务、语言建模任务以及BABILong基准测试中均展现出出色的性能，并有效处理超过200万个token的上下文窗口。
关键观点3: Titans架构的设计理念是借鉴人类的记忆特性，如优先记住违反预期的事件，并根据信息的“惊讶度”进行记忆更新。
研究团队提出了三种将核心组件整合的变体架构，并展示了其在不同任务中的应用效果。
正文
让 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 - θt∇ℓ(Mt-1; xt)其中，ηtSt-1 代表过去惊讶，它通过一个数据依赖的衰减率 ηt 来控制历史信息的保留程度。当系统判断当前上下文与过去相关时，ηt 会接近 1，保持对历史信息的重视；当需要切换到新的上下文时，ηt 会接近 0，允许系统“忘记”不再相关的历史信息。而 -θt∇ℓ(Mt-1; xt) 则代表瞬时惊讶，它通过计算损失函数关于输入的梯度来量化当前输入 xt 与已有记忆 Mt-1 之间的差异程度。θt 参数控制了系统对新信息的敏感度。这种设计让模型能够像人类一样，对违反预期的信息保持特别的关注。在具体实现中，研究团队使用了一个关联记忆损失函数：ℓ(Mt-1; xt) = ||Mt-1(kt) - vt||²₂其中，输入 xt 会被转换为键值对：kt = xtWK 和 vt = xtWV。这种设计让模型能够学习和存储信息之间的关联关系，类似于人类记忆中的联想机制。第三个组件是持久记忆，这是一组与数据无关的可学习参数，专门用于存储完成特定任务所需的基础知识，类似于人类的程序性记忆。从效率角度来看，研究团队还特别优化了 Titans 的训练过程。Behrouz 在推文中提到：“我们通过 TTT（Sun et al., 2024）扩展了 mini-batch 梯度下降的可并行化对偶形式，并通过额外的矩阵乘法来实现权重衰减。”通过将 mini-batch 梯度下降的前向传播重新表述为矩阵乘法操作，他们实现了高效的并行训练。图丨展示如何并行使用 matmuls 进行神经记忆训练的示意图（来源：arXiv）然后，研究团队提出了三种将这些组件整合的变体架构：- MAC（Memory as Context）：这种方案将记忆作为上下文信息处理，对输入序列进行分段，并使用过去的记忆状态来提取相应的记忆信息。
图丨作为上下文（MAC）架构的内存（来源：arXiv）- MAG（Memory as Gate）：这种方案使用滑动窗口注意力机制作为短期记忆，同时使用神经记忆模块作为长期记忆，通过门控机制将两者结合。
图丨作为门的内存（MAG）架构（来源：arXiv）- MAL（Memory as Layer）：这种方案将神经记忆作为深度神经网络的一个层，更接近传统的混合模型设计。
图丨作为层的内存（MAL）架构（来源：arXiv）实验结果显示，Titans 在多个测试基准上都表现出色。在语言建模任务上，拥有 760M 参数的 Titans(MAC) 在 WikiText 上达到了 19.93 的困惑度，显著优于同等规模的 Transformer++（25.21） 和 Mamba2（22.94）。在常识推理任务上，Titans 在包括 PIQA、HellaSwag、WinoGrande 等 9 个基准测试中的平均准确率达到 52.51%，超过了现有的最好成绩。图丨Titans 和基于循环和 Transformer 的基线在语言建模和常识推理任务中的性能。混合模型用 ∗ 标记。简单和混合模型中的最佳结果被突出显示。（来源：arXiv）值得一提的是，在 BABILong 基准测试中，即使是参数量较小的 Titans 也展现出出色的性能。在需要在超长文档中进行推理的任务中，Titans(MAC) 不仅战胜了 Mamba2、RWKV 等现代模型，甚至超越了参数量大得多的 GPT-4 和 Llama3-70B。实验证明，Titans 能够有效处理超过 200 万个 token 的上下文窗口，这一突破将为长文本处理开辟新的可能。图丨Titans 和基线在 BABILong 基准上的性能（来源：arXiv）研究团队表示，Titans 的 PyTorch 和 JAX 实现代码将很快开源。随着这一技术的进一步发展和应用，我们有望看到更多能够处理超长文本的高效 AI 系统涌现，这对于文档分析、长文本理解、知识检索等领域都具有重要意义。当然，这项研究仍有进一步探索的空间。如何设计更高效的神经记忆架构、如何在更大规模的模型中应用这一技术，都是未来值得关注的方向。但毫无疑问，Titans 的出现为解决 AI 系统的长期依赖问题提供了一个极具前景的新方向。参考资料：
1.https://arxiv.org/abs/2501.00663
2.https://x.com/behrouz_ali/status/1878859673555624373
运营/排版：何晨龙
