一种让 AI 聊天机器人整天交谈而不会崩溃的新方式

当人类与人工智能的对话涉及多轮连续对话时，驱动ChatGPT等聊天机器人的强大大型语言机器学习模型有时会开始崩溃，导致机器人的性能迅速下降。

来自麻省理工学院和其他地方的一组研究人员已经查明了这个问题的一个令人惊讶的原因，并开发了一个简单的解决方案，使聊天机器人能够保持不间断的对话而不会崩溃或减慢速度。

他们的方法涉及对许多大型语言模型核心的键值缓存（类似于对话记忆）进行调整。在某些方法中，当此缓存需要保存的信息超过其容量时，第一批数据就会被淘汰。这可能会导致模型失败。

通过确保这些前几个数据点保留在内存中，研究人员的方法允许聊天机器人无论对话持续多长时间都能继续聊天。

这种方法称为 StreamingLLM，即使对话持续超过 400 万字，模型也能保持高效。与另一种通过不断重新计算过去部分对话来避免崩溃的方法相比，StreamingLLM 的执行速度提高了 22 倍以上。

这可以使聊天机器人在整个工作日进行长时间的对话，而无需不断重启，从而使高效的人工智能助手能够完成文案、编辑或生成代码等任务。

“现在，通过这种方法，我们可以持久地部署这些大型语言模型。通过制作一个我们可以随时与之聊天的聊天机器人，并且始终可以根据我们最近的对话回复我们，我们可以在一些新的应用程序中使用这些聊天机器人，“电气工程和计算机科学（EECS）研究生、StreamingLLM论文的主要作者Guangxuan Xiao说。

Xiao的合著者包括他的导师Song Han，EECS副教授，麻省理工学院-IBM Watson人工智能实验室成员，NVIDIA杰出科学家;以及 Meta AI 的研究科学家田远东;卡内基梅隆大学助理教授陈贝迪;以及资深作者 Mike Lewis，Meta AI 的研究科学家。这项工作将在学习表征国际会议上发表。

令人费解的现象

大型语言模型将数据（如用户查询中的单词）编码为称为标记的表示形式。许多模型采用所谓的注意力机制，使用这些标记来生成新文本。

通常，AI 聊天机器人会根据刚刚看到的文本编写新文本，因此它将最近的令牌存储在内存中，称为 KV 缓存，以供以后使用。注意力机制构建一个网格，其中包含缓存中的所有标记，这是一个“注意力图”，用于映射每个标记或单词彼此标记之间的关系。

了解这些关系是使大型语言模型能够生成类似人类的文本的一项功能。

但是，当缓存变得非常大时，注意力图可能会变得更加庞大，从而减慢计算速度。

此外，如果编码内容需要的令牌数超过缓存可以容纳的令牌数，则模型的性能会下降。例如，一个流行的模型可以存储 4,096 个代币，但一篇学术论文中大约有 10,000 个代币。

为了解决这些问题，研究人员采用了一种“滑动缓存”，将最旧的代币凸出以添加新代币。但是，一旦第一个标记被逐出，模型的性能通常会直线下降，从而迅速降低新生成的单词的质量。

在这篇新论文中，研究人员意识到，如果他们将第一个令牌保留在滑动缓存中，即使超过缓存大小，模型也将保持其性能。

但这没有任何意义。小说中的第一个单词可能与最后一个单词无关，那么为什么第一个单词对于模型生成最新单词如此重要呢？

在他们的新论文中，研究人员还发现了这种现象的原因。

注意力汇

一些模型在其注意力机制中使用 Softmax 操作，该机制为每个令牌分配一个分数，表示它与彼此令牌的关系程度。Softmax 操作要求所有注意力分数的总和为 1。由于大多数代币没有很强的相关性，因此它们的注意力分数非常低。该模型在第一个令牌中转储任何剩余的注意力分数。

研究人员将第一个代币称为“注意力汇”。

“我们需要一个注意力汇，模型决定使用第一个令牌作为注意力汇，因为它是全局可见的——其他每个令牌都可以看到它。我们发现，我们必须始终将注意力汇保持在缓存中，以保持模型动态，“Han 说。

在构建 StreamingLLM 时，研究人员发现，在滑动缓存的开头有四个注意力接收器令牌可以带来最佳性能。

他们还发现，每个标记的位置编码必须保持不变，即使添加了新标记并淘汰了其他标记。如果令牌 5 被撞出，则令牌 6 必须保持编码为 6，即使它现在是缓存中的第五个令牌。

通过结合这两个想法，它们使 StreamingLLM 能够保持连续的对话，同时优于使用重新计算的流行方法。

例如，当缓存有 256 个令牌时，重新计算方法需要 63 毫秒才能解码新令牌，而 StreamingLLM 需要 31 毫秒。但是，如果缓存大小增长到 4,096 个令牌，则新令牌的重新计算需要 1,411 毫秒，而 StreamingLLM 只需要 65 毫秒。

“StreamingLLM的创新方法以注意力接收器机制为中心，确保了稳定的内存使用和性能，即使在处理长达400万个令牌的文本时也是如此，”新加坡国立大学计算机科学教授Yang You说，他没有参与这项工作。“这种能力不仅令人印象深刻;它具有变革性，使 StreamingLLM 能够应用于广泛的 AI 应用程序。StreamingLLM的性能和多功能性标志着它是一项非常有前途的技术，有望彻底改变我们处理人工智能驱动的生成应用程序的方式。

卡内基梅隆大学（Carnegie Mellon University）机器学习和计算机科学系的助理教授陈天琪（Tianqi Chen）也没有参与这项研究，他对此表示赞同，他说：“流式LLM可以平滑地延长大型语言模型的对话长度。我们一直在使用它来实现 Mistral 模型在 iPhone 上的部署，并取得了巨大成功。

研究人员还通过在所有训练样本中预置几个占位符标记来探索在模型训练期间注意力汇的使用。

他们发现，使用注意力接收器进行训练允许模型在其缓存中只有一个注意力接收器的情况下保持性能，而不是通常需要四个注意力接收器来稳定预训练模型的性能。

但是，虽然 StreamingLLM 使模型能够进行连续的对话，但模型无法记住未存储在缓存中的单词。未来，研究人员计划通过研究检索已被驱逐的令牌或使模型能够记住先前对话的方法来针对这一限制。

StreamingLLM 已整合到 NVIDIA 的大型语言模型优化库 TensorRT-LLM 中。

这项工作部分由麻省理工学院-IBM Watson人工智能实验室、麻省理工学院科学中心和美国国家科学基金会资助。