记忆是认知的关键组成部分，让人类能够从充斥我们生活的无尽噪音中选择性存储和提取重要信息。相比之下，大语言模型（Large Language Model，LLM）则缺乏这种能力，只能无差别地存储和处理所有过去的输入，这在长时间任务中对它们的性能和成本产生了严重影响。
就像人类大脑会选择性地保留重要信息并逐渐淡忘不重要的细节一样，人工智能系统也需要这样的智能记忆管理机制。否则，随着模型规模的不断扩大，其训练和部署对计算资源和内存的需求也会无限上升。
长期以来，研究人员一直在探索如何让 AI 系统具备更接近人类的记忆能力。传统的解决方案主要依赖预设的规则来管理模型的记忆，比如基于时间顺序或注意力分数（Attention Score）来选择性地保留或丢弃信息。
然而，这些方法往往过于机械，无法像人类记忆那样智能地区分信息的重要性，导致在提高效率的同时往往会损害模型的性能。
在此背景下，日本初创公司 Sakana AI 的研究团队提出了一种新的解决方案——神经注意力记忆模型（Neural Attention Memory Models，NAMMs）。
这一方案借鉴了自然进化在塑造人类记忆系统中发挥的关键作用，通过进化算法训练一个专门的神经网络来，这种方法能够像人类大脑一样主动选择和保留重要信息，从而在提高效率的同时提升模型的性能。
就像人类大脑会根据信息的长期使用价值来决定是否保留一样，NAMMs 通过对注意力模式的分析来评估信息的重要性。其核心机制包括三个关键组成部分：特征提取系统、记忆管理网络和进化优化策略。
图丨 NAMMs 执行过程中的三个主要步骤（来源：arXiv）
首先是特征提取机制。NAMMs 采用短时傅里叶变换（STFT）来处理注意力矩阵的列向量。具体来说，它使用大小为 32 的 Hann 窗口进行处理，生成 17 个复值频率的频谱图表示。
这种表示方式非常巧妙，因为它既保留了注意力值随时间变化的频率特征，又大大压缩了数据量。研究团队通过实验发现，这种频谱表示比直接使用原始注意力值或手工设计的特征更有效。
其次是向后注意力记忆（BAM）架构的设计。这是 NAMMs 的核心创新之一，它引入了一种特殊的注意力机制，允许 token 只能关注其在 KV 缓存中的“未来”相关内容。
这种设计的妙处在于，它能够建立 token 之间的竞争关系，使得模型能够学会保留最有信息量的 token。例如，当出现重复的句子或词语时，模型会倾向于保留最新的出现，因为它包含了更完整的上下文信息。
图丨 NAMMs 的反向掩码使每个 token 仅关注 KV 缓存中的未来相关词（来源：arXiv）
在优化策略上，研究团队采用了 CMA-ES（协方差矩阵自适应进化策略）算法。传统的梯度下降法难以处理记忆管理这种具有离散决策的问题，而 CMA-ES 通过模拟自然进化过程，能够直接优化非可微的目标函数。
具体来说，团队采用了增量进化的方式，从单个任务开始，逐步增加训练任务的数量，这种方式能够提供更好的正则化效果，提高模型的泛化能力。
研究团队选择 Llama 3-8b 作为基础模型训练了 NAMMs，并在 LongBench、InfiniteBench 以及 ChouBun 上进行了全面评估。结果显示，NAMM 为 Llama 3-8b Transformer 带来了明显的性能提升，在总体表现上超越了已有的 H2O 和 L2 这两种手工设计的内存管理方法。
例如，在 LongBench 基准测试中，NAMMs 不仅将 KV 缓存大小减少到原来的 25%，还实现了 11% 的性能提升。在 InfiniteBench 测试中模型性能从基线的 1.05% 提升到了 11%，同时将缓存大小减少到原来的 40%。
图丨 LongBench 基准测试结果（来源：arXiv）
NAMMs 的另一个重要特性是其出色的零样本迁移能力。研究团队发现，仅在语言任务上训练的 NAMMs 可以直接应用到其他架构和模态上。
比如，当应用到 Llava Next Video-7B 模型时，NAMMs 在 LongVideoBench 和 MLVU 基准测试中都取得了不错的表现，视觉任务的性能提升了 1%，同时将视频帧的缓存大小减少到原来的 72%。
在强化学习方面，使用决策转换器（Decision Transformer）时，NAMMs 在 D4RL 基准测试中实现了 9% 的性能提升，同时将缓存大小减少到原来的 81%。
深入分析 NAMMs 的工作机制，研究团队发现它学会了一种智能的记忆管理策略。通过观察不同层的记忆保留模式，发现模型在早期和中期层倾向于保留更多和更老的 token，这可能是因为这些层负责处理和聚合长距离信息。而在信息密度较高的代码任务中，模型则学会了保留相对更多的 token。
实际上，NAMMs 延续了 Sakana AI 此前的研究方法，即从自然界获取灵感，通过模拟自然进化的过程来优化 AI 系统。这一研究思路与该公司在模型合并和进化优化方面的技术积累具有内在的一致性。
同 Sakana AI 此前开发的自动化“进化”算法能够自主识别和合并优秀模型一样，NAMMs 也采用了进化算法来优化记忆管理系统，无需人工干预即可实现性能的持续提升。
其特殊的研发思路，已经为这家仅成立了一年的初创公司赢得了 2.1 亿美元的 A 轮融资，在这融资中，其估值已达 15 亿美元。
未来，研究团队可能会探索更复杂的记忆模型设计，比如考虑更细粒度的特征提取方法，或者研究如何将 NAMMs 与其他优化技术结合使用。
他们表示：“这项工作才刚刚开始挖掘我们新类记忆模型的潜力，我们预计这可能会为未来几代 Transformer 的发展提供许多新的机会。”
