总部位于东京的创业公司Sakana AI的研究人员开发了一种新技术,使语言模型能够更有效地利用内存,帮助企业降低在大型语言模型 (LLM) 和其他基于Transformer的模型之上构建应用程序的成本。

这项技术名为“通用Transformer内存(Universal Transformer Memory)”,利用特殊的神经网络来优化LLM,保留重要的资讯,并从上下文中移除冗余的细节。

Transformer模型是LLM的核心,其回应依赖于“上下文窗口”的内容,也就是从用户那里接收到的输入内容。

上下文窗口可视为模型的工作内存。调整上下文窗口的内容会对模型的性能产生重大影响,这也催生了一个完整的“提示工程”领域。

目前的模型支持非常长的上下文窗口,可容纳数十万甚至数百万个token(LLM将用户输入的文本、词汇部分、短语、概念及数字转换成的数值表示)。

这使得用户能够将更多资讯塞进他们的提示中。然而,更长的提示会导致更高的运算成本和更慢的性能。优化提示以删除不必要的token,同时保留重要的资讯,可以降低成本并提高速度。

目前的提示优化技术既耗费资源,又需要用户手动测试不同的配置,以减少提示的大小。

通用Transformer内存通过神经注意力记忆模块(Neural Attention Memory Models, NAMMs)来优化提示,这是一种简单的神经网络,可以决定是否“记住”或“遗忘”存储在LLM内存中的每个token。

“这项新功能让Transformer能够移除无用或冗余的细节,并专注于最关键的资讯,对于需要长上下文推理的任务来说,这尤为重要,”研究人员表示。

Universal transformer memory。图片来源:Sakana AI)

NAMM是在与LLM分开的环境中训练的,并于推理时与预训练模型结合,这使其具有灵活性且易于部署。然而,NAMM需要访问模型的内部启动,这意味着它只能应用于开源模型。

与Sakana AI研发的其他技术一样,NAMM是通过演划算法训练的,而非基于梯度的优化方法。通过反复变异与选择最具表现力的模型,演划算法以试错方式优化NAMM的效率与性能。这尤其重要,因为NAMM正试图实现一个不可微分的目标:保留或丢弃token。

NAMM在Transformer的注意力层运行,这是Transformer架构的关键组成部分,负责确定模型上下文窗口中每个token的关系及重要性。根据注意力值,NAMM决定哪些token应保留,哪些应从模型的上下文窗口中移除。这种基于注意力的机制使得训练过的NAMM可以在不同模型间使用,而无需进一步修改。例如,针对纯文本数据训练的NAMM,可以应用于视觉或多模态模型,而无需额外训练。

类神经注意力内存模块 (NAMM) 会查看注意力层,以决定应从上下文窗口中保留或丢弃哪些token。图片来源:Sakana AI

为测试通用Transformer内存的概念,研究人员在开放源码的Meta Llama 3-8B模型之上训练了一个NAMM。他们的实验表明,使用NAMM,基于Transformer的模型在非常长的串行上的自然语言和程序代码问题上表现更好。同时,通过丢弃不必要的token,NAMM使LLM模型在执行任务时节省了高达75%的缓存内存。

研究人员写道:“在我们的基准测试中,NAMM为Llama 3-8B变压器提供了明显的性能改进。”“此外,我们的内存系统产生了显著的额外好处,减少了每一层的上下文大小,同时从未针对内存效率进行明确的优化。”

NAMM模型在提升模型性能的同时,也能与领先的提示优化技术相抗衡。图片来源:Sakana AI

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

研究人员写道:“即使在这些分布外的设置中,NAMM仍然通过丢弃诸如冗余的视频帧和次佳的动作等token来保持其优势,使其新的基础模型能够专注于最相关的资讯以提高性能。”

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

例如,对于程序代码编写任务,模型会丢弃对应于不影响程序代码执行的注解和空白的连续token块。

另一方面,在自然语言任务中,模型会移除表达语法冗余的token,这些token并不影响串行的意义。

研究人员已发布用于创建自有NAMM的程序代码。像通用Transformer内存这样的技术,对于处理数百万个token且可以从速度提升和成本降低中受益的企业应用程序非常有用。经过训练的NAMM的可重复使用性也使其成为在企业中的不同应用程序中使用的多功能工具。

对于未来,研究人员建议使用更先进的技术,例如在LLM的训练期间使用NAMM以进一步扩展其内存功能。

“我们的新型记忆模型才刚刚开始展现其潜力,我们期待这将为未来的Transformer时代带来更多突破,”研究人员表示。