Sakana AI ：进化的通用 Transformer 存储器 - 链载Ai

接回主题，想象一下 Transformer 不仅能记住最重要的事情，还能主动忘记多余的细节，从而产生更智能、更快速、适应性更强的模型。这正是团队的工作所取得的成果，也是它改变游戏规则的原因

：（那小编猜一下，那一定是多头自注意力机制下了很活

）

一种新型记忆：神经注意力记忆模型 (NAMM) 优化了变压器存储和检索信息的方式，实现了前所未有的效率和性能。
增强的结果：借助 NAMM，Transformer 可以在各种语言和编码任务中取得卓越的结果，同时所需的内存更少。
跨领域掌握：仅经过语言训练，NAMM 即可应用于视觉、强化学习和其他领域，无需额外训练。

上图展示了，经过语言训练的完全进化的 NAMM 可以零样本转移到其他转换器，甚至可以跨输入模式和其它任务域。

团队假设塑造 transformer 的潜在记忆 KV 缓存会带来新的机会来提高它们在下游任务中的能力。甚至允许基础模型在测试时学习上下文中的全新技能。此外，与 prompt 工程不同，直接管理 Transformer 的内存可以独立地为每个潜在级别提供不同的上下文，这样各个层和注意力头就可以专注于最相关的信息以满足其特定需求。

受这些考虑的启发，团队搬出了神经注意力记忆模型 (NAMM)，引入了一类经过进化训练的新网络，以学习一种高效的记忆系统，从而最大限度地提高预训练 Transformer 的下游性能。进化到本质上克服了内存管理操作的不可微操作性，其结果为二元（选择要保留/丢弃的标记，

强啊），这使得基于自然进化在塑造人类记忆方面发挥的关键作用，它似乎类似地根据信息的终生有用性有选择地合并和主动修剪信息。

有点像文本水印影子，关于文本水印技术，在抬头有引用，感兴趣大家去看看技术。

注意力和转换器

转换器是专门为高效处理输入序列而设计的神经网络架构。这些模型将一串标记（例如，单词嵌入、图像块、机器人状态等）作为 Embedding 输入，并在其层内产生一组长度相同的潜在标记。

多头点积注意力或简称为自注意力，是现代 Transformer 的特征，它促进了 token 表示之间的有效信息共享。注意力层进行一组并行计算，每个计算称为注意力头，将 token 映射到查询、键和值向量集合 d，这些向量沿着矩阵中的序列维度组织键值对 (Q, K, V) 并且该层的输出计算如下：

这里，M 表示与注意力矩阵相乘的可选掩码，通常会强制执行自回归条件，使得每个 token 无法关注其未来。注意力层的解释来自注意力矩阵的元素 A 即每个键之间的点积归一化。直观地讲，这些值中的每一个都可以理解为 token 的相对重要性。

频谱图的特征提取

这里团队提出短时傅里叶变换 (STFT) ，这是一种预处理一维非平稳信号典型技术。该技术已在音频、生物医学、地震和许多其他类型的模态特征提取方面得到广泛应用。STFT 对信号进行时间卷积，通过离散傅里叶变换将每个卷积窗口移至频域，从而生成原始输入的频谱图表示。

STFT 的计算过程可以分为以下几个步骤：

选择窗口：选择一个固定大小的窗口 w，通常是一个窗口或其他类型的窗口。
滑动窗口：将窗口在信号上滑动，每次滑动一个固定的距离（称为步长）。
计算傅里叶变换：对于每个窗口位置，将窗口内的信号段 v[t′] 乘以窗口函数w[t−t′]，然后对结果进行傅里叶变换。
生成频谱图：将每个窗口位置的傅里叶变换结果ω[n] 组合成一个二维的频谱图。

STFT 的公式通过将信号在时间 t 附近的 T 个样本提取出来，并对这些样本进行傅里叶变换，从而生成一个描述信号频率成分随时间变化的图。这个图可以用来提取信号的特征，用于进一步的分析或处理。

NAMM 原理 & 架构

NAMM (Neural Attention Memory Models) 是简单的神经网络分类器，经过训练可以决定是否“记住”或“忘记”存储在内存中的每个给定标记。这项新功能允许转换器丢弃无用或多余的细节，并专注于最关键的信息，我们发现这对于需要长上下文推理的任务至关重要。

然而训练 NAMM 带来了一个关键挑战，因为我们的内存模型做出的任何决定都有二元结果：每个 token 要么保留在内存中，要么永远丢失。这给问题带来了不可微分的方面，使得使用基于梯度的优化的传统训练技术不再适用。

另一方面，进化不需要任何梯度，能够在这些情况下表现出色。通过反复试验，如下图，对比 Alex (

一种神经卷积网络) 提供了多种进化算法可以迭代变异并选择性能最佳的模型，使团队能够优化 NAMM 的效率和性能，即使面对不可微分的操作也是如此。

NAMM 背后的一个关键因素在于它们使用了注意力矩阵，这是任何 Transformer 的任何层都共有的关键组件。这些矩阵编码了每个 token 相对于其他 token 的重要性，使其成为决定要忘记哪些 token 的理想输入。由于有这些属性，团队通过依靠注意力矩阵，可以直接在模型的各个层上应用单个 NAMM，甚至可以将同一个 NAMM 转移到其他 Transformer，而无需进一步训练。

这种无与伦比的迁移属性不仅限于 LM，还适用于处理完全不同的输入模式和问题设置（例如计算机视觉、机器人控制）的基础模型

（这就是端对端学习）。

从技术角度来看，NAMM 的执行主要分为三个步骤：

处理注意力序列：内存中每个标记的注意力值被转换为声谱图（一种基于频率的表示形式，已在音频、医学和地震学等领域得到很好的确立高手想法）。
压缩信息：然后逐个元素指数移动平均数（EMA）压缩生成的表示（将数据压缩为每个标记注意力值历史的紧凑、固定大小的特征摘要，白话理解成 Embedding）。
决定要记住什么：NAMM 使用这些特征作为其学习到的神经网络分类器的输入（用一个分数来决定要忘记哪些噪声标记，并允许转换器专注于与其任务最相关的信息，白话理解注意力QKV的套路，可以看下小编这篇中间解析 Transformer，一文了解 2024 China 谷歌 IO AI 技术点）。

如下图，NAMM 执行中的上面说的三个主要步骤：从左图将注意力序列处理成频谱图然后用 EMA 压缩信息计算分数，最后到右边决定要记住的内容。

通过分析记住和忘记的 token，团队发现 NAMM 会记住 Transformer 各个层中的不同 token。在向前反馈中，早期层中，NAMM 会保留全局信息，例如任务前言和关键字。而在后期层中，NAMM 似乎忘记了许多这些 token，这些 token 的信息已经被整合，而更关注文本中包含的局部细节和概念。

回到特征提取图谱，每个 token 的频谱图表示为 ω，由注意层处理，然后进行简单的线性运算以输出其相对分数。反向掩蔽引入了不对称性，确保每个标记只能关注其未来的相关标记。

另外不得不提反向注意力，反向注意力记忆网络通过首先直接应用自注意力层来处理这些表示，该层采用了反自回归后向掩蔽，旨在促进记忆中标记之间的不对称交互。请看下图。

反向注意力记忆模型（Backward Attention Memory Models, BAM）BAM 的设计目的是通过反向掩码（backward masking）来实现 token 之间的信息共享。反向掩码确保每个 toke n只关注其在 KV 缓存中的未来相关 token。这种设计引入了一种有目的的非对称关系，使得模型能够区分旧 token 和新 token

。这种机制有助于在保留最相关信息的同时，避免完全丢弃任何信息。

然后将自注意力的输出馈送到单个最终线性层以获得最终分数。在进行一些初步测试之后，采用了一些重要的附加设计选择。首先，出于效率考虑，在注意力机制中使用单个头，并且不使用层归一化。其次，注意力层产生的输出是频谱图特征维数的两倍。这些输出通过残差和乘性交互在最后的线性层之前集成回主网络中。

此外，团队发现 NAMM 的行为因任务而异。在编码任务中，裁剪后的标记大多是连续的文本块，对应于空格、注释和不必要的样板代码段（

如下图，其中删除线就是需要被裁剪掉的。绿色区域是普通文本，橘红区域是 Coding）。相反，在自然语言任务中，NAMM 会忘记英语语法中存在的语法冗余导致的句子中间的许多标记，从而使 LM 能够专注于重要的名称和关键概念。

回顾一下，整体文献的改进主要是几个大措施达成的。

记忆管理网络：NAMMs 通过引入一个学习型网络来管理记忆，优化了 Transformer 的性能和效率。这个网络能够根据不同层和注意力头的需要，提供不同的潜在上下文，专注于最相关信息。
进化优化：NAMMs 在预训练的 Transformer 模型上进行进化，以学习高效的记忆系统，从而最大化下游任务的性能。进化算法克服了记忆管理操作的非可微性，使得基于梯度的优化不再适用。
通用性：NAMMs 仅依赖于注意力矩阵中的值，使其能够通用应用于任何使用自注意力的模型。这种特性使得 NAMMs 可以在不同的 Transformer 架构之间进行零样本迁移。
性能提升：在多个长上下文基准测试中，NAMMs 实现了显著的性能提升，同时将模型的输入上下文减少到原始大小的一小部分。
零样本迁移能力：NAMMs 仅在语言任务上进行训练，但可以零样本迁移到全新的Transformer 架构，甚至跨输入模态，其好处扩展到视觉和强化学习领域。
特征提取：NAMMs 使用基于短时傅里叶变换（STFT）的特征提取方法，从注意力矩阵中提取频谱图表示，以压缩时间轴上的信息。
内存模型设计：NAMMs 设计了一个小的神经网络来输出每个KV缓存中的 token 的选择分数，并通过反向注意力记忆模型（BAM）实现跨token的信息共享。
增量进化：NAMMs 在长上下文语言建模任务的子集上进行进化优化，使用 CMA-ES优化算法，并在 Llama 3 8B 基础模型上应用 NAMMs。

NAMM 效果评估

团队在Llama 3 8b 基础模型上训练 NAMM ，并在数据集 LongBench 、InfiniteBench 和团队自己的 ChouBun 上全面评估这个强大的记忆增强型 LM。这三个基准评估了 LM 处理非常长的输入文本信息以回答自然语言和编码问题的能力，总共 36 个不同的任务。

将 NAMM 与 H₂O 和 L₂ 行了比较。在基准测试中，NAMM 为 Llama 3 8b Transformer 提供了明显的性能改进。此外，内存上还带来了明显的附带好处，减少了每层的上下文大小，但这都是从未明确针对内存效率进行优化。虽然之前的基准也显著减少了上下文大小，但这些效率提升往往是以性能为代价的，这也团队目标一致，即保留而不是提高原始的全上下文性能

。

另外还评估了为不同模态设计的转换器，例如 Llava Next Video 和 Decision Transformer 两者都是以处理计算机视觉和强化学习任务。NAMM 仍通过丢弃冗余视频帧和次优操作等标记来保留其优势，从而使其新的基础模型专注于最相关的信息以提高性能。

结论

引入了一种神经注意力记忆模型来学习基于 Transformer 的基础模型的记忆系统。通过在预训练语言模型的基础上发展 NAMM，证明了它们在三种语言的各种长上下文任务中的有效性，超越了以前手工设计的内存管理方法所能达到的水平。团队还表明，模型本质上具有跨架构、输入模式和任务域的零样本可迁移性。展望未来，这项工作才刚刚开始挖掘新一类记忆模型的潜力，预计这可能会为未来几代 Transformer 的发展提供许多新的机会。