Online Speculative Decoding

📄 arXiv: 2310.07177v4 📥 PDF

作者: Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, Alvin Cheung, Zhijie Deng, Ion Stoica, Hao Zhang

分类: cs.AI, cs.CL, cs.LG

发布日期: 2023-10-11 (更新: 2024-06-10)

🔗 代码/项目: GITHUB

💡 一句话要点

提出在线推测解码以提升大语言模型推理效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 推测解码 在线学习 知识蒸馏 语言模型 推理优化 用户查询适应 模型更新

📋 核心要点

  1. 现有的推测解码方法在面对多样化文本输入时，草稿模型的预测准确性较低，限制了其有效性。
  2. 本文提出的在线推测解码通过持续更新草稿模型，适应用户查询数据的分布，从而提高预测准确性。
  3. 实验结果表明，在线推测解码显著提高了令牌接受率，并实现了1.42倍至2.17倍的延迟减少。

📝 摘要（中文）

推测解码是一种加速大型语言模型推理的关键技术，通过使用较小的草稿模型预测目标模型的输出。然而，当面临多样化的文本输入和草稿模型与目标模型之间显著的能力差距时，其有效性可能受到限制。为了解决这一挑战，本文提出了在线推测解码的概念，核心思想是根据观察到的用户查询数据持续更新多个草稿模型。通过适应查询分布，减轻草稿模型训练分布与查询分布之间的偏差，从而使草稿模型能够更准确地预测目标模型的输出。我们基于知识蒸馏开发了在线推测解码的原型，并使用合成和真实查询数据进行了评估。结果显示，令牌接受率显著提高0.1至0.65，延迟减少1.42倍至2.17倍。

🔬 方法详解

问题定义：本文旨在解决推测解码中草稿模型预测准确性不足的问题，尤其是在面对多样化文本输入时，草稿模型与目标模型之间的能力差距显著影响了推测效果。

核心思路：提出在线推测解码，通过持续更新草稿模型，使其能够适应用户查询数据的分布，从而提高预测的准确性，减少推理延迟。

技术框架：整体架构包括多个草稿模型的训练与更新模块，用户查询数据的实时采集与分析模块，以及基于知识蒸馏的输出预测模块。

关键创新：最重要的创新在于通过实时更新草稿模型，使其能够动态适应查询分布，显著提高了预测准确性，与传统静态草稿模型方法形成鲜明对比。

关键设计：在设计中，采用了知识蒸馏技术来优化草稿模型的训练过程，并通过特定的损失函数来平衡草稿模型与目标模型之间的输出差异。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，在线推测解码在令牌接受率上提高了0.1至0.65，延迟减少幅度达到1.42倍至2.17倍，展现出显著的性能提升，证明了该方法的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和智能助手等。通过提高大型语言模型的推理效率，在线推测解码可以显著提升用户体验，降低响应时间，具有广泛的实际价值和未来影响。

📄 摘要（原文）

Speculative decoding is a pivotal technique to accelerate the inference of large language models (LLMs) by employing a smaller draft model to predict the target model's outputs. However, its efficacy can be limited due to the low predictive accuracy of the draft model, particularly when faced with diverse text inputs and a significant capability gap between the draft and target models. We introduce online speculative decoding to address this challenge. The main idea is to continuously update the (multiple) draft model(s) on observed user query data. Adapting to query distribution mitigates the shifts between the training distribution of the draft model and the query distribution, enabling the draft model to more accurately predict the target model's outputs. We develop a prototype of online speculative decoding based on knowledge distillation and evaluate it using both synthetic and real query data. The results show a substantial increase in the token acceptance rate by 0.1 to 0.65, bringing 1.42x to 2.17x latency reduction. Our code is available at https://github.com/LiuXiaoxuanPKU/OSD.