东京初创公司Sakana AI的研究人员开发了一种新技术，使语言模型能够更有效地使用内存，帮助企业降低基于大型语言模型（LLMs）和其他基于Transformer的模型构建应用的成本。

这项技术被称为“通用Transformer Memroy”（An Evolved Universal Transformer Memory），它利用特殊的神经网络来优化LLM，保留重要的信息片段，并从其上下文中丢弃冗余细节。

优化Transformer Memory

Transformer模型的响应依赖于它们的“上下文窗口”中的内容——即用户输入的信息。

上下文窗口可以被视为模型的工作内存。调整上下文窗口的内容可以极大地影响模型的性能，这促使了“提示工程”领域的兴起。

当前模型支持非常长的上下文窗口，包含数十万甚至数百万个token（LLM的数值表示，用户在其提示中输入的词汇、词组、短语、概念和数字）。

这使用户能够在他们的提示中塞入更多信息。然而，更长的提示可能会导致更高的计算成本和更慢的性能。优化提示以移除不必要的token同时保留重要信息可以降低成本并提高速度。

当前的提示优化技术资源密集或需要用户手动测试不同配置以减小提示的大小。

神经注意力内存模块

Universal Transformer Memroy通过使用神经注意力内存模型（NAMMs），即简单的神经网络，来优化提示，决定是“记住”还是“忘记”LLM内存中存储的每个token。

研究人员表明：“这一新能力使得Transformer能够舍弃无用或冗余的细节，专注于最关键的信息，这对于需要长上下文推理的任务来说是至关重要的。”

Universal transformer memory （来源：Sakana AI）

NAMMs在训练期间独立于LLM进行训练，并在推断时与预训练模型结合使用，这使它们灵活且易于部署。然而，它们需要访问模型的内部激活，这意味着它们只能应用于开源模型。

与Sakana AI开发的其他技术一样，NAMMs通过进化算法而不是基于梯度的优化方法来训练。通过迭代地变异并通过试错选择性能最佳的模型，进化算法优化了NAMMs的效率和性能。这一点尤为重要，因为NAMMs试图实现一个非微分目标：保留或丢弃token。

NAMMs作用于LLMs的注意力层，这是Transformer架构的关键组成部分之一，它决定了模型上下文窗口中每个token的关系和重要性。基于注意力值，NAMMs决定哪些token应当保留，哪些可以从LLM的上下文窗口中丢弃。这种基于注意力的机制使得可以在不同模型上使用训练有素的NAMM，无需进一步修改。例如，一个在仅文本数据上训练的NAMM可以应用于视觉或多模态模型，无需额外训练。

神经注意力记忆模型 (NAMM) 检查注意力层，以确定哪些标记应该从上下文窗口中保留或丢弃（来源：Sakana AI）

测试结果

为了测试Universal Transformer Memroy的概念，研究人员在一个开源的Meta Llama 3-8B模型上训练了一个NAMM。他们的实验显示，使用NAMMs的Transformer基础模型在处理自然语言和编码问题的长序列上表现更好。同时，通过丢弃不必要的token，NAMM使LLM模型在执行任务时节省了高达75%的缓存内存。

“在我们的基准测试中，NAMMs为Llama 3-8BTransformer提供了明显的性能改进。”研究人员写道。此外，我们的内存系统还带来了显著的副作用，减少了每层的上下文大小，而从未针对内存效率进行过显式优化。”

NAMM 模型在提高模型性能的同时，还与领先的快速优化技术相竞争 （来源：Sakana AI）

他们还在70B版本的Llama以及为其他模态和任务设计的Transformer模型上测试了该模型，例如Llava（计算机视觉）和Decision Transformer（强化学习）。

研究人员表示：“即使在这些分布外的设置中，NAMMs通过丢弃诸如冗余视频帧和次优动作等token，保留了其优势，使其新的基础模型能够专注于最相关的信息以提高性能。”

任务依赖行为

另一个有趣的发现是，NAMMs会根据任务自动调整其行为。

例如，在编码任务中，模型会丢弃与注释和空白相关的连续token块，这些token不会影响代码的执行。另一方面，在自然语言任务中，模型会丢弃代表语法冗余的token，这些token不会影响序列的含义。研究人员发布了创建自己的NAMMs的代码。像Universal Transformer Memroy这样的技术对于处理数百万token的企业应用非常有用，可以从速度提升和成本降低中受益。训练有素的NAMM的可重用性也使其成为企业中不同应用的通用工具。

对于未来，研究人员建议更先进的技术，例如在训练LLMs期间使用NAMMs以进一步扩展它们的内存能力。研究人员表示：“这项工作只是开始挖掘们这一新类内存模型的潜力，我们预期它可能会为未来几代Transformer的进步提供许多新的机会。”