谷歌 Titan 架构突破 Transformer 内存瓶颈

谷歌推出新型 Titan 架构

科技界正在热议谷歌推出的新型架构 Titan。该架构旨在挑战 Transformer 模型，尤其是在内存处理方面的局限性。由于该架构由谷歌内部团队开发，因此备受关注，被认为是 Transformer 模型的潜在继任者。

现有模型面临的内存挑战

传统的 LSTM 和 Transformer 模型虽然创新，但在模拟人类记忆方面面临挑战：

容量有限：数据通常被压缩成固定大小的隐藏状态，限制了可以保留的信息量。
计算开销大：虽然能够捕获长距离依赖关系，但计算成本会随着序列长度的增加而呈平方级增长，这使得它对于非常长的序列效率低下。
过度依赖训练数据：简单地记忆训练数据并不总是有助于现实世界的应用，因为测试数据可能超出训练分布。

Titan 的方法：神经启发式记忆模块

Titan 团队采取了不同的方法，试图将信息编码到神经网络的参数中。他们开发了一个在线元模型，该模型旨在学习如何在测试期间记住和忘记特定数据。该模型受到神经心理学原理的启发，包含以下关键要素：

惊喜作为触发器：意外事件更容易被记住。“惊喜”是通过内存模块输入的梯度来衡量的。梯度越大，输入就越出乎意料。
动量和遗忘机制：动量机制将短期惊喜累积到长期记忆中，而遗忘机制会擦除旧记忆，防止内存溢出。
基于多层感知机（MLP）的内存：内存模块由多个 MLP 层组成，使其能够存储数据的深度抽象，使其比传统的基于矩阵的内存更强大。

这种在线元学习方法有助于模型专注于学习如何适应新数据，而不是仅仅记忆训练数据。该模块还设计用于并行计算，从而提高了效率。

将内存模块集成到深度学习架构中

Titan 研究团队提出了三种将内存模块集成到深度学习架构中的变体：

MAC（内存作为上下文）：此方法将长期和持久内存（编码任务知识）组合为输入到注意力机制的上下文。
MAG（内存作为门）：此方法使用内存模块与跨两个分支的滑动窗口注意力机制的门控融合。
MAL（内存作为层）：在此方法中，内存模块被实现为独立的层，该层在将历史信息馈送到注意力机制之前对其进行压缩。

该团队发现，每种变体都有其优点和缺点。

Titan 的性能和优势

Titan 在包括语言建模、常识推理和时间序列预测在内的各种任务中都表现出卓越的性能。它超越了 Transformer 和 Mamba 等最先进的模型。值得注意的是，长期记忆模块（LMM）本身在多项任务中都优于基线模型，这展示了其在没有短期记忆（注意力）的情况下独立的学习能力。

在旨在查找长文本中细微线索的“大海捞针”测试中，即使序列长度从 2k 增加到 16k，Titan 仍保持了大约 90% 的准确率。该团队表示，标准测试并未完全展示 Titan 在处理长文本方面的优势。在需要从分布在极长文档中的事实进行推断的任务中，Titan 也优于 GPT4、Mamba 甚至带有 RAG 的 Llama3.1。

Titan 在时间序列预测和 DNA 序列建模等特定领域也表现出了令人印象深刻的性能。

Titan 背后的团队

这项研究由 Google Research NYC 算法和优化小组的一个团队进行，该团队目前不属于 Google DeepMind。

Ali Behrouz，康奈尔大学实习生，是该论文的第一作者。
钟沛林，清华大学校友，哥伦比亚大学博士毕业生，自 2021 年以来一直是谷歌的研究科学家。他以本科生身份在 STOC 2016 上发表了第一作者论文而闻名。
Vahab Mirrokni，谷歌研究员和副总裁，领导该团队。

该团队使用 Pytorch 和 Jax 开发了 Titan，并计划很快发布用于训练和评估的代码。