让ChatGPT持续对话不丢失记忆的新方法

研究者们为ChatGPT等大型语言模型中存在的一项棘手难题找到了简明有效的解决办法，这一问题原本会严重影响模型的性能。

在涉及连续对话多个回合的人工智能对话中，推动诸如ChatGPT之类聊天机器人的强大的大型语言机器学习模型有时会突然出现故障，造成机器人性能急剧下降。

来自麻省理工学院及其他机构的研究团队发现了这一问题的一个出人意料的原因，并提出了一个简便的解决办法，这使得聊天机器人能够持续对话而不会崩溃或变慢。

他们的方法对许多大型语言模型中心的关键缓存（相当于对话的记忆存储）进行了调整。在某些情况下，当缓存需要存储的信息量超出了其实际容量时，最早的数据便会被替换，这可能会导致模型失效。

通过保留内存中最初的一些数据点，研究者们的方案让聊天机器人可以无论对话进行多长时间都能持续不断地交流。

这种被称为StreamingLLM的方法，甚至当对话延续到超过四百万个单词时，也能保持模型的高效率。与另一种通过不断重计算之前对话的一部分来避免系统崩溃的方法相比较，StreamingLLM的表现速度快了22倍以上。

这将使得聊天机器人能全天进行长对话，无需重复启动，有效提供写作、编辑、生成代码等任务的人工智能助手。

“如今，借助这种方法，我们能够持续性地部署这些大型语言模型。通过创建一个我们随时可以沟通、并能根据近期对话内容做出回应的聊天机器人，我们能在许多新的应用领域中利用这些聊天机器人，”电气工程和计算机科学（EECS）研究生、该论文主要作者Guangxuan Xiao表示。

Xiao的合作作者包括，他的导师Song Han，EECS副教授、MIT-IBM Watson AI实验室成员、NVIDIA杰出科学家；以及Meta AI研究科学家Yuandong Tian；卡内基梅隆大学助理教授Beidi Chen；和Meta AI研究科学家、论文的资深作者Mike Lewis。这项研究成果将在国际学习表征大会上发表。

一个令人费解的现象

大型语言模型将数据，比如用户查询中的词汇，编码成称作“代币”的表征形式。许多模型采用所谓的“注意力机制”，利用这些代币来创造新文本。

一般来说，AI聊天机器人会根据它所见过的文本来撰写新内容；因此，它会把新近的代币存储在一个名为KV缓存的记忆库中，稍后加以使用。注意力机制会创建一个包含缓存所有代币的网格图，“注意力地图”，在这个地图上标明每个词汇和其他词汇之间的相关性强弱。

理解这些关系帮助此类大型语言模型能够创作出貌似人类的文本。

但是，当缓存体量变得巨大时，注意力地图也随之扩张至庞大无比，导致计算速度下降。

而且，如果编码内容所需的代币超出缓存上限，模型的表现就会受损。例如，一个广为使用的模型能够存储4096个代币，但一篇学术论文可能包含约10000个代币。

研究人员为了规避这些问题，采取了“滑动缓存”策略，淘汰最旧的代币，以便加入新代币。然而，一旦第一个代币被移出，模型的性能往往会骤降，这直接影响新生成词汇的质量。

在这篇新论文中，研究者们发现，如果他们在滑动缓存中保持第一个代币不变，即使超出了缓存大小，模型仍能保持其性能。

但这个发现看上去似乎违反常理。一本小说中的第一个词很少与最后一个词有直接联系，为何第一个词对于模型产生新词至关重要呢？

研究者们在他们的新论文中还揭示了这一现象背后的原因。

注意力集中点

一些模型在其注意力机制中使用了Softmax运算，该运算会给每个词元打分，反映它们相互之间的关联程度。Softmax运算要求所有的注意力得分加起来总和必须为1。由于大部分词元之间的关系并不密切，因此它们的注意力得分非常低。模型把任何剩余的注意力得分都赋给了第一个词元。

研究人员把这个第一个词元称之为“注意力集中点”。

“我们需要一个注意力集中点，而模型选择第一个词元作为这个集中点，因为它对所有其他词元都是可见的。我们发现，为了维持模型的动态性，我们必须不断地把这个注意力集中点保留在缓存中。”韩这样说道。

在开发StreamingLLM时，研究者发现在缓存的最开始放置四个注意力集中点词元，会获得最优的性能表现。

他们还注意到，每个词元的位置编码在新增词元和淘汰旧词元时必须保持不变。例如，如果第五个词元被淘汰了，第六个词元即使变成了缓存里的第五个也必须保持其为第六的编码。

结合这两种思路，他们让StreamingLLM能够保持连贯的对话，并且在性能上超过了那些使用重计算技术的流行方法。

比如，在缓存容量为256词元时，使用重计算技术的方法解码一个新词元需要63毫秒，而StreamingLLM仅需要31毫秒。但如果缓存容量增加到4096词元，重计算技术需要1411毫秒来解码一个新词元，StreamingLLM仍然只需要65毫秒。

“StreamingLLM采用了创新的注意力池技术，解决了在处理长达400万令牌文本时性能和内存稳定性的问题。”新加坡国立大学计算机科学总统青年教授杨勇表示，虽然他并未参与这项工作。“这项技术的能力不仅印象深刻，它还具有改变游戏规则的潜力，让StreamingLLM能够在众多AI领域得到应用。StreamingLLM在性能和多样性上的表现，预示着它将成为推动我们如何运用AI生成应用的革新技术。”

同意这一观点的还有卡内基梅隆大学的机器学习和计算机科学助理教授陈天琪，他也未参与这项研究。他说：“StreamingLLM使我们能够顺畅扩展大型语言模型的对话长度。我们已经成功地利用它在iPhone上部署Mistral模型。”

研究团队还在模型训练过程中探索了注意力池的应用，他们在所有训练样本前加入了几个占位符。

他们发现，配合注意力池进行训练的模型，仅在缓存中使用一个注意力池，就能维持性能，而通常预训练模型要稳定性能则需四个。

不过，尽管StreamingLLM使模型可以进行持续对话，模型仍无法记忆未被存入缓存的词汇。未来，研究人员计划通过探索检索被删除令牌的方法，或让模型记住之前的对话来克服这一限制。

StreamingLLM已经集成到了NVIDIA的大语言模型优化库[TensorRT-LLM]中。

这项工作部分得到了麻省理工学院-IBM沃森人工智能实验室、麻省理工学院科学中心和美国国家科学基金会的支持。