Amazon Science 与 Booking.com 用知识蒸馏加 Medusa 将 LLM 推理提速最高 20 倍

旅游平台每天为数百万用户提供服务，回答关于目的地的问题、提供酒店推荐，并帮助用户规划行程。近年来，许多平台转向使用大语言模型（LLM）来改进这些服务。得益于其规模——通常拥有数百亿参数——这些LLM能够给出准确且与上下文相关的旅行建议。遗憾的是，如此庞大的模型计算成本也很高，而且这一成本会随客户查询数量的增加而上升。更甚者，这类大模型对于网页和移动应用而言可能过于缓慢，在这些场景中，哪怕100毫秒的延迟也可能让公司损失客户和收入。为应对这些扩展性和延迟方面的挑战，我们探索了两种互补的优化技术。一种是称为知识蒸馏（KD）的压缩技术，即由一个缓慢但强大的模型为一个更快的轻量级模型生成训练目标。另一种称为Medusa，是一个将顺序的下一词元预测——LLM中常见的瓶颈——转化为并行计算的框架。我们与领先的在线旅行社Booking.com合作，将KD和Medusa技术应用于为其用户体验各方面提供支持的LLM，旨在改善延迟或成本效益。我们尤其聚焦于三项功能：AI行程规划器（AI Trip Planner）：一个对话式助手，必须在不到一秒内从对话中提取旅行偏好（延迟）；智能筛选器（Smart Filters）：一项将自然语言搜索查询实时转换为结构化信息筛选条件的功能（延迟）；住宿描述生成（Accommodation Description Generation）：一项以数百万房源规模生成房源描述的功能（成本效益）。 我们证明，将KD与Medusa结合使用，能够在成本不到三分之一的情况下保持最佳开源LLM的质量，同时运行速度最高提升至20倍。我们的框架确保旅游平台或其他依赖LLM的应用能够在真实场景中提供快速、高性价比且准确的响应。我们在去年的计算语言学协会（ACL）会议上展示了这些技术，并获得了Industry Track最佳论文奖。我们的方法：Medusa加知识蒸馏。Medusa是一个旨在降低聊天和代码生成等实时应用中LLM延迟的框架。它在运行时——即推理时——而非训练期间提升计算效率。通常，LLM一次生成一个词元——即单词或单词片段。Medusa通过并行预测多个词元来加速这一过程。Medusa的做法是向基础LLM添加称为Medusa头的预测模块，每个并行解码的词元对应一个头。LLM所运行的GPU硬件针对并行性进行了优化，因此能够以极小的额外开销处理并行词元解码。Medusa可带来显著的加速，通常为2至4倍，且不损害生成质量。知识蒸馏通过将知识从更大的“教师”模型迁移到更小的“学生”模型，来解决大型、参数密集型模型的部署挑战。通常，教师模型首先在高质量标注数据上进行训练或微调，然后用于在额外的无标注数据上生成预测。随后将人工标注数据和机器标注数据结合起来训练一个更小的学生模型，该模型学习模仿教师的行为。 学生模型保留了教师的大部分性能，但效率显著更高，使其更适合高流量消费级网站等真实场景。在我们的工作中，我们将KD和Medusa集成到一个统一的流水线中，从预训练基础模型开始。与标准KD一样，我们在标注数据集上微调教师模型，并用它在无标注的领域相关文本上生成预测。然而在学生训练阶段，我们将Medusa头附加到学生模型上，并在教师预测与标注数据组成的合并语料上训练它。这样便产生了一个单一的紧凑模型，它继承了教师的知识，同时获得了多词元预测能力。实验我们选择TinyLlama——一个拥有11亿参数的紧凑模型——作为学生模型，并选择拥有700亿参数的Llama-3.1作为教师模型。我们进行了四次递进式实验，每一次都建立在前一次的基础上。我们首先针对三个用例——AI行程规划、智能筛选和住宿描述生成——分别在任务特定数据上微调TinyLlama。即使没有进一步优化，微调后的模型也明显比少样本GPT-4o更准确，这证实了领域特定的微调弥补了它们紧凑规模的不足。然而，它们的性能仍不及大得多的教师模型。应用知识蒸馏——在人工标注和教师标注数据的混合上训练学生模型——将这一差距缩小了90%以上，使小模型接近与教师持平。更小的模型规模还使推理速度大约提升至6倍，但这对于实时使用来说仍然不够快。然而，添加Medusa头后，推理速度又额外提升了2至3.6倍。 我们测试了两种Medusa变体：在Medusa-1中，我们先微调基础LLM，然后冻结其权重并单独微调Medusa头；在Medusa-2中，我们在单个阶段中训练所有内容。两者都实现了相当的加速，且没有明显的质量损失，但Medusa-2更简单的流水线使其对管理多个模型的生产团队来说更实用。结果将KD与Medusa-2结合同时带来了两方面收益：我们的最终模型在质量上匹配700亿参数教师模型，同时根据用例不同，运行速度快10到20倍。对于三个用例中的每一个，我们都使用上述流水线，以领域特定数据微调了一个专用模型。对于AI行程规划器，我们展示了相比GPT-3.5用户参与度提升2.9%。对于智能筛选，我们实现了99%的查询响应速度快15倍。对于AI生成的住宿描述，我们展示了用户“有帮助”评分提高1.4%，而成本不到GPT-4的三分之一。感谢我们的AWS同事Aleksandra Dokic以及Booking.com团队：Emmanouil Stergiadis、Eran Fainman、Ilya Gusev和Moran Beladev。