新型LLM优化技术降低内存成本

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

这项技术名为“通用 Transformer 内存”（universal transformer memory）。它使用特殊的神经网络优化 LLM，使其保留重要的信息片段，同时从上下文中丢弃冗余细节。

优化 Transformer 内存

作为 LLM 基础架构的 Transformer 模型，其响应取决于“上下文窗口”中的内容，也就是模型从用户那里接收到的输入。

上下文窗口可以视为模型的工作内存。调整上下文窗口的内容，可能会对模型性能产生巨大影响，并由此催生了“prompt engineering”这一完整领域。

目前的模型支持非常长的上下文窗口，可以容纳数十万甚至数百万个 token。token 是 LLM 对用户通过 prompt 输入的单词、词语片段、短语、概念和数字所采用的数值表示。

这使用户能够在 prompt 中塞入更多信息。然而，prompt 越长，计算成本可能越高，运行速度也可能越慢。通过优化 prompt，在保留重要信息的同时删除不必要的 token，可以降低成本并提升速度。

目前的 prompt 优化技术往往需要消耗大量资源，或者要求用户手动测试不同配置，才能缩短 prompt。

神经注意力内存模块

通用 Transformer 内存使用神经注意力内存模型（neural attention memory models，NAMM）优化 prompt。NAMM 是一种简单的神经网络，负责决定要“记住”还是“忘记”存储在 LLM 内存中的每个 token。

研究人员写道：“这项新能力让 Transformer 可以丢弃无用或冗余的细节，专注于最关键的信息。我们发现，对于需要长上下文推理的任务而言，这一点至关重要。”

通用 Transformer 内存（来源：Sakana AI）

NAMM 与 LLM 分开训练，并在推理阶段与预训练模型结合，因此具有很强的灵活性，也很容易部署。不过，它们需要访问模型内部的激活值，这意味着 NAMM 只能应用于开源模型。

与 Sakana AI 开发的其他技术一样，NAMM 使用进化算法训练，而不是基于梯度的优化方法。进化算法通过反复进行模型变异，并在不断试错的过程中选出表现最好的模型，从而针对效率和性能优化 NAMM。这一点尤其重要，因为 NAMM 试图实现的是一个不可微目标：保留或丢弃 token。

NAMM 在 LLM 的注意力层上运行。注意力层是 Transformer 架构的关键组件之一，负责确定模型上下文窗口中各个 token 之间的关系及其重要性。NAMM 根据注意力值判断哪些 token 应该保留，哪些可以从 LLM 的上下文窗口中丢弃。这种基于注意力的机制，使训练完成的 NAMM 无须进一步修改便可以应用到不同模型上。例如，使用纯文本数据训练的 NAMM，无须额外训练便可以应用于视觉模型或多模态模型。

神经注意力内存模型（NAMM）通过检查注意力层，判断上下文窗口中的哪些 token 应该保留或丢弃（来源：Sakana AI）

通用内存的实际表现

为了验证通用 Transformer 内存概念在实际应用中的效果，研究人员基于开源的 Meta Llama 3-8B 模型训练了一个 NAMM。实验表明，在使用 NAMM 后，基于 Transformer 的模型处理超长序列中的自然语言和编程问题时，表现更加出色。与此同时，NAMM 通过丢弃不必要的 token，让 LLM 在执行任务时最多可以节省 75% 的缓存内存。

研究人员写道：“在我们的各项基准测试中，NAMM 都为 Llama 3-8B Transformer 带来了明显的性能提升。此外，尽管我们的内存系统从未针对内存效率进行显式优化，它仍然带来了显著的附加收益，缩小了每一层的上下文规模。”

NAMM 在提升模型性能的同时，也能与领先的 prompt 优化技术竞争（来源：Sakana AI）

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

研究人员写道：“即使面对这些分布外场景，NAMM 仍然能够通过丢弃冗余视频帧和次优动作等 token 来保持优势，让新的基础模型可以专注于最相关的信息，从而提升性能。”

根据任务调整行为

另一项有趣的发现是，NAMM 会根据任务自动调整自己的行为。

例如，在编程任务中，模型会丢弃连续成块的 token，这些 token 对应的是不会影响代码执行的注释和空白字符。

而在自然语言任务中，模型会丢弃代表语法冗余、不会影响序列含义的 token。

研究人员已经发布了用于创建自定义 NAMM 的代码。通用 Transformer 内存等技术，对于需要处理数百万 token，并且能够从速度提升和成本降低中受益的企业应用非常有用。训练完成的 NAMM 还可以重复使用，因此它能够成为一种通用工具，应用于企业内部的不同场景。

展望未来，研究人员提出了一些更先进的技术方向，例如在 LLM 训练期间使用 NAMM，进一步扩展模型的内存能力。

研究人员写道：“这项工作才刚刚开始发掘这类新型内存模型的潜力。我们预计，它可能为推动下一代 Transformer 的发展带来许多新的机会。”