Sakana AI推出LLM记忆管理技术NAMMs，可将内存成本降低75%

记忆是认知的关键组成部分，让人类能够从充斥我们生活的无尽噪音中选择性存储和提取重要信息。相比之下，大语言模型（Large Language Model，LLM）则缺乏这种能力，只能无差别地存储和处理所有过去的输入，这在长时间任务中对它们的性能和成本产生了严重影响。

就像人类大脑会选择性地保留重要信息并逐渐淡忘不重要的细节一样，人工智能系统也需要这样的智能记忆管理机制。否则，随着模型规模的不断扩大，其训练和部署对计算资源和内存的需求也会无限上升。

长期以来，研究人员一直在探索如何让 AI 系统具备更接近人类的记忆能力。传统的解决方案主要依赖预设的规则来管理模型的记忆，比如基于时间顺序或注意力分数（Attention Score）来选择性地保留或丢弃信息。

然而，这些方法往往过于机械，无法像人类记忆那样智能地区分信息的重要性，导致在提高效率的同时往往会损害模型的性能。

在此背景下，日本初创公司 Sakana AI 的研究团队提出了一种新的解决方案——神经注意力记忆模型（Neural Attention Memory Models，NAMMs）。

这一方案借鉴了自然进化在塑造人类记忆系统中发挥的关键作用，通过进化算法训练一个专门的神经网络来，这种方法能够像人类大脑一样主动选择和保留重要信息，从而在提高效率的同时提升模型的性能。

就像人类大脑会根据信息的长期使用价值来决定是否保留一样，NAMMs 通过对注意力模式的分析来评估信息的重要性。其核心机制包括三个关键组成部分：特征提取系统、记忆管理网络和进化优化策略。

（来源：arXiv）

首先是特征提取机制。NAMMs 采用短时傅里叶变换（STFT）来处理注意力矩阵的列向量。具体来说，它使用大小为 32 的 Hann 窗口进行处理，生成 17 个复值频率的频谱图表示。

这种表示方式非常巧妙，因为它既保留了注意力值随时间变化的频率特征，又大大压缩了数据量。研究团队通过实验发现，这种频谱表示比直接使用原始注意力值或手工设计的特征更有效。

其次是向后注意力记忆（BAM）架构的设计。这是 NAMMs 的核心创新之一，它引入了一种特殊的注意力机制，允许 token 只能关注其在 KV 缓存中的“未来”相关内容。

这种设计的妙处在于，它能够建立 token 之间的竞争关系，使得模型能够学会保留最有信息量的 token。例如，当出现重复的句子或词语时，模型会倾向于保留最新的出现，因为它包含了更完整的上下文信息。

（来源：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）