“Transformer八子”之一，让AI大模型学会了主动“忘记”

还记得第一个“人工智能科学家”（The AI Scientist）吗？只需 15 美元，它就可以写出一篇超过顶级机器学习会议接受门槛的论文。

从生成新颖的研究想法、编写代码、执行实验、可视化结果，到通过撰写完整的科学论文来描述其研究结果，然后运行模拟审查流程进行评估，该“人工智能科学家”实现了全自动的科学发现。

据介绍，该“人工智能科学家”由日本人工智能初创公司 Sakana AI 打造，其中一位联合创始人为“Transformer 八子”之一 Llion Jones，投资方包括 Lux Capital、Khosla Ventures 等知名机构以及谷歌大神 Jeff Dean、Hugging Face 创始人 Clem Delangue 等个人投资者。

Sakana AI 公布了他们的另一项具有创新性的研究成果，即一个名为“神经注意力记忆模型”（NAMM）的革命性 Transformer 记忆系统，其灵感来自人类记忆如何选择性地保留和修剪信息。

主要特点如下：

一种新型记忆系统：神经注意力记忆模型（NAMM）优化了 Transformer 存储和检索信息的方式，释放出前所未有的效率和性能。
超强效果：有了 NAMM，Transformer 在各种语言和编码任务中都能取得优异成绩，同时所需的内存更少。
跨领域掌握：NAMM 仅在语言方面接受过训练，无需额外训练即可应用于视觉、强化学习和其他领域。

想象一下，Transformer 不仅能“记住”最重要的事情，还能主动“忘记”多余的细节，从而产生更智能、更快速、适应性更强的模型。

该研究解决了 Transformer 基础模型缺乏选择性存储信息能力的问题。NAMM 作为一种新型记忆方式，极大地提高了 Transformer 效率和性能，让其在多种任务中表现优异，还拥有跨领域掌握的能力。

Transformer 基础模型中的内存

Transformer 架构已成为深度学习的黄金标准，在现代基础模型设计中的应用无处不在，表现出卓越的性能和可扩展性。Transformer 的输出完全以输入 token 的最近上下文为条件，对于语言模型（LM）来说，最近上下文通常对应于前面的单词窗口。因此，这种上下文可以被视为 Transformer 的“工作记忆”，其中包含与其当前应用相关的最新输入。

这种工作记忆中包含的信息已被证明会对 Transformer 的性能产生相当大的影响。例如，即使只是通过提示工程仔细修改输入文本，也能让 LM 释放出全新的能力，执行训练数据之外的任务。

不过，提供处理长上下文的能力也会立即影响训练和推理成本，因为现代 Transformer 越来越耗费资源和成本。最近的许多方法都提出了部分抵消这些成本的方法，即通过精心设计的策略，研究在记忆上下文中丢弃 token 子集的效果。结果，这些方法在提高效率方面取得了初步成功，但却牺牲了原始模型的性能。

通过进化学习记忆框架

与此形成鲜明对比的是，他们的工作通过引入 NAMM，摆脱了以往依赖固定规则或手工策略的方法。NAMM 是一种简单的神经网络分类器，经过训练后可决定对存储在内存中的每个给定 token 进行“记忆”还是“遗忘”。这一新功能允许 Transformer 摒弃无用或多余的细节，而专注于最关键的信息，他们发现这对于需要长上下文推理的任务来说至关重要。

然而，训练 NAMM 是一项重大挑战，因为他们的记忆模型所做的任何决定都是二选一：每个 token 要么保留在记忆中，要么永远丢失。这就给问题引入了一个不可分的方面，使得使用梯度优化的传统训练技术变得不合适。

而进化则不需要梯度，因此在这些情况下表现出色。通过迭代突变和试错选择 SOTA 模型，进化算法使他们能够优化 NAMMs 的效率和性能，即使面对无差别操作也不例外。

图｜通过进化优化来优化 NAMM，迭代变异和选择网络参数，从而利用他们的新记忆系统获得最佳语言建模性能。

NAMM 背后的一个关键要素在于其对注意力矩阵的使用，而注意力矩阵是任何 Transformer 的任何层所共有的关键组件。这些矩阵编码了每个 token 相对于其他 token 的重要性，是决定遗忘哪些 token 的理想输入。由于这些特性，他们只需依赖注意力矩阵，就能直接在模型各层应用单一的 NAMM，甚至无需任何进一步的训练就能将相同的 NAMM 移植到其他 Transformer 上。这种无与伦比的转移特性不仅限于 LM，也适用于处理完全不同的输入模式和问题设置的基础模型（例如，计算机视觉、机器人控制）。

在技术上，NAMM 的执行主要分为三个步骤：

处理注意力序列——将内存中每个 token 的注意力值转换为频谱图：这是一种基于频率的表征，在音频、医学和地震学等领域已得到广泛应用。
压缩信息——然后使用元素指数移动平均法（EMA）对生成的表示进行压缩：将数据浓缩为每个 token 的注意力值历史的紧凑、固定大小的特征摘要。
决定记住什么——然后，NAMM 将这些特征作为其学习的神经网络分类器的输入：输出分数以决定哪些 token 需要“遗忘”，并允许 Transformer 专注于与其任务最相关的信息。

图｜NAMM 执行过程中三个主要步骤的示意图：将注意力序列处理为频谱图（左图），用 EMA 压缩信息（中图），计算分数以决定记住什么（右图）。

语言及其他领域的应用

他们在 Llama 3 8b 基本模型的基础上训练 NAMM，并在 LongBench、InfiniteBench 和 ChouBun 上全面评估了这种强大的内存增强 LM：这三个基准测试评估了 LM 处理超长输入文本信息的能力，以回答自然语言和编码问题，总计 36 个不同的任务。他们将 NAMM 与 H₂O 和 L₂（两种之前手工设计的内存管理方法）进行了比较。

在所有基准测试中，NAMM 的性能明显优于 Llama 3 8b Transformer 。此外，他们的内存系统还带来了显著的消极作用，减少了每一层的上下文大小，同时从未明确优化内存效率。虽然先前的基线系统也显著减少了上下文大小，但这些效率的提高往往是以性能为代价的——这与它们所宣称的目标一致，即保留而非提高原有的全上下文性能。

研究表明，他们的约束（conditioning）方法具有通用性，可以实现对全新基础模型的零样本迁移。特别是，他们在大型 Llama 70B LM 以及针对不同模式设计的 Transformer （如 Llava Next Video 和 Decision Transformer）上评估了 NAMM，以解决计算机视觉和强化学习任务。即使在这些非分布式环境中，NAMM 也能通过丢弃多余视频帧和次优动作等 token 保持其优势，从而使其新的基础模型能够专注于最相关的信息，从而提高性能。