大语言模型（LLM）在落地应用时，高昂的计算成本始终是企业面临的核心挑战。特别是随着上下文窗口不断扩展，动辄数十万乃至数百万token的输入，带来了巨大的计算开销与响应延迟，让许多团队望而却步。近期，东京初创公司Sakana AI的研究团队提出了一项创新方案——名为“通用Transformer记忆”（An Evolved Universal Transformer Memory）的技术，有望将此类成本大幅降低，最高可达75%。

该技术的核心思路并不复杂：通过一个专用的神经网络模块，教会模型在上下文中“聪明地”保留关键信息，同时果断丢弃冗余细节。如此一来，模型仅需处理真正有价值的内容，效率和性能自然能够得到同步提升。

优化Transformer内存机制

Transformer模型的响应质量，很大程度上取决于其“上下文窗口”中存储的内容——也就是用户输入的信息。可以将上下文窗口理解为模型的工作内存。窗口中存放什么、如何存放，直接影响着模型的输出效果，这也正是“提示工程”这一领域兴起的原因。

当前的主流模型已支持超长上下文窗口，能够容纳数十万甚至数百万个token——token是LLM对用户输入中词汇、词组、短语、概念和数字的数值化表示。长窗口允许用户向提示中塞入更多信息，但代价也很明显：更长的提示意味着更高的计算成本和更慢的响应速度。因此，优化提示内容、剔除不必要的token，同时保留关键信息，已经成为降本增效的关键路径。

现有的提示优化技术要么消耗大量资源，要么需要用户手动反复测试不同配置来缩小提示规模，效率并不理想。

神经注意力记忆模块

通用Transformer记忆的突破之处在于引入了一种名为“神经注意力记忆模型”（Neural Attention Memory Models，简称NAMMs）的简单神经网络。它的任务十分明确：逐一检查LLM内存中的每个token，决定是“记住”还是“遗忘”。

研究人员指出，这一新能力使Transformer能够主动舍弃无用的冗余细节，将注意力集中在最关键的信息上，这对于需要长上下文推理的任务而言至关重要。

NAMMs的训练独立于LLM之外，在推理时才与预训练模型结合使用，因此部署起来非常灵活。不过，它需要访问模型内部的激活值——这意味着目前只能应用于开源模型。

与Sakana AI此前开发的多项技术类似，NAMMs采用进化算法，而非基于梯度的优化方法。通过迭代变异和试错，挑选出性能最优的模型，进化算法不断优化NAMMs的效率和表现。这一点尤其关键，因为NAMMs要实现的是一种非微分目标：保留还是丢弃token，这一决策本身无法用梯度直接优化。

NAMMs作用于LLM的注意力层——这是Transformer架构的核心组件之一，负责判断上下文窗口中每个token之间的关系和重要性。基于注意力值，NAMMs决定哪些token该保留，哪些可以从上下文窗口中移除。这种基于注意力的机制还有一个显著优势：训练好的NAMMs可以跨模型复用，无需额外调整。例如，一个仅在文本数据上训练的NAMM，可以直接应用于视觉或多模态模型。

测试结果

为验证通用Transformer记忆的实际效果，研究团队在开源的Meta Llama 3-8B模型上训练了一个NAMM。实验数据显示，配备NAMMs的Transformer基础模型在处理自然语言和编码问题的长序列任务时，表现显著优于原始模型。同时，通过丢弃不必要的token，NAMM帮助LLM在执行任务时节省了高达75%的缓存内存。

研究者在论文中写道：“在我们的基准测试中，NAMMs为Llama 3-8B Transformer带来了明显的性能改进。此外，我们的内存系统还产生了一个显著的‘副作用’：每层的上下文大小被自动缩减，而我们在训练时从未明确优化过内存效率。”

团队还进一步在70B版本的Llama上进行了测试，并扩展到其他模态和任务的Transformer模型，例如LLaVA（计算机视觉）和Decision Transformer（强化学习）。

研究人员表示：“即使在这些分布外的场景中，NAMMs通过丢弃冗余视频帧、次优动作等token，依然保持了其优势，使基础模型能够专注于最相关的信息，从而提升性能。”

任务依赖行为

一个特别有趣的发现是，NAMMs会根据任务类型自动调整其行为策略。

例如，在编码任务中，模型倾向于丢弃与注释和空白相关的连续token块——这些内容不影响代码的执行。而在自然语言任务中，模型则主要丢弃语法上冗余的token，这些token不影响序列的整体含义。这种任务感知的自主调节能力，让NAMMs显得格外智能。

研究团队已经发布了创建自有NAMMs的代码。对于处理数百万token的企业级应用，类似通用Transformer记忆这样的技术，能够直接转化为速度提升和成本降低。而训练好的NAMMs的跨模型可复用性，更使其有望成为企业中不同应用场景的通用优化工具。

至于未来方向，研究人员建议更先进的技术路径，例如在训练LLM的过程中就引入NAMMs，进一步扩展模型的内存能力。研究者在论文中总结道：“这项工作只是挖掘这类新型内存模型潜力的起点，我们预期它可能为未来几代Transformer的进步提供许多新的机会。”