谷歌新架构一战成名，打破Transformer记忆瓶颈，姚班校友钟沛林新作

梦晨 发自 凹非寺

量子位 | 公众号 QbitAI

想挑战 Transformer 的新架构有很多，来自谷歌的“正统”继承者 Titan 架构更受关注。

英伟达把测试时间计算（Test-time Computing）称为大模型的第三个 Scaling Law。

OpenAI 把它用在推理（Reasoning），谷歌这次把它用在了记忆（Memory）。

一作Ali Behrouz表示：

Titans 比 Transformers 和现代线性 RNN 更高效，并且可以有效地扩展到超过 200 万上下文窗口，性能比 GPT4、Llama3 等大模型更好。

他还解释了这篇研究的动机，团队认为 Transformer 中的注意力机制表现为短期记忆，因此还需要一个能记住很久以前信息的神经记忆模块。

新的长期记忆模块

提到记忆，大家可能会想到 LSTM、Transformer 等经典模型，它们从不同角度模拟了人脑记忆，但仍有局限性：

要么将数据压缩到固定大小的隐状态，容量有限

要么可以捕捉长程依赖，但计算开销随序列长度平方级增长

并且，仅仅记住训练数据在实际使用时可能没有帮助，因为测试数据可能在分布外。

为此，Titans 团队打算将过去信息编码到神经网络的参数中，训练了一个在线元模型（Online meta-model），该模型学习如何在测试时记住/忘记特定数据。

他们从神经心理学中汲取灵感，设计了一个神经长期记忆模块，它借鉴了人脑原理：

意料之外的事件（即“惊喜”）更容易被记住。‍

惊喜程度由记忆模块对输入的梯度来衡量，梯度越大说明输入越出人意料。

引入动量机制和遗忘机制，前者将短期内的惊喜累积起来形成长期记忆，后者可以擦除不再需要的旧记忆，防止记忆溢出。

记忆模块由多层 MLP 组成，可以存储深层次的数据抽象，比传统的矩阵记忆更强大。

这种在线元学习范式，避免了模型记住无用的训练数据细节，而是学到了如何根据新数据调整自己，具有更好的泛化能力。

另外，团队还验证了这个模块可以并行计算。

如何将这个强大的记忆模块融入深度学习架构中呢？

为此，Titans 提出了三种变体：

MAC，记忆作为上下文

将长期记忆和持久记忆（编码任务知识的不变参数）作为当前输入的上下文，一起输入给 attention。

MAG，记忆作为门

在记忆模块和滑动窗口 attention 两个分支上进行门控融合。

MAL，记忆作为层

将记忆模块作为独立的一层，压缩历史信息后再输入给 attention。

在实验中，发现每种方法都有自己的优缺点。

Titans 在语言建模、常识推理、时间序列预测等任务上全面超越 Transformer 和 Mamba 等各路架构的 SOTA 模型。

并且仅靠长期记忆模块（LMM，Long-term Memory Module）本身，就在多个任务上击败基线。

证明了即使没有短期记忆（也就是 Attention），该机制也具备独立学习的能力。

在长文本中寻找细粒度线索的“大海捞针”测试中，序列长度从 2k 增加到 16k，准确率保持在 90% 左右。

但团队认为，这些通用的测试已经体现不出 Titans 在长文本上的优势。

在另一项需要对分布在极长文档中的事实做推理的任务中，Titans 表现超过了 GPT4 、Mamba 等，以及 Llama3.1 + RAG 的系统。

另外在时间序列预测、DNA 序列建模等特定任务中，Titans 也取得不错的表现。

三位作者来自 Google Research NYC 算法和优化团队，目前还没被合并到 Google DeepMind。

一作是 Ali Behrouz 来自康奈尔大学的实习生。

钟沛林是清华姚班校友，博士毕业于哥伦比亚大学，2021 年起加入谷歌任研究科学家。

2016 年，钟沛林本科期间的一作论文被顶会 STOC 2016 接收，是首次有中国本科生在 STOC 上发表一作论文。

领队的 Vahab Mirrokni 是 Google Fellow 以及 VP。

团队表示 Titians 是用 Pytorch 和 Jax 中实现的，打算很快提供用于训练和评估模型的代码。

论文地址：https://arxiv.org/abs/2501.00663v1

参考链接：[1]https://x.com/behrouz_ali/status/1878859086227255347