[论文解读] Online Speculative Decoding

Xiaoxuan Liu, Lanxiang Hu+5 more|arXiv (Cornell University)|Oct 11, 2023

Topic ModelingComputer Science被引用 3

一句话总结

本文提出在线推测解码（online speculative decoding），一种利用大语言模型服务集群中的空闲计算资源，持续基于实时用户查询数据对草稿模型进行再训练的方法。通过应用在线知识蒸馏，草稿模型能实时适应查询分布，使token接受率提升10%–65%，推理延迟降低1.22倍至3.06倍。

ABSTRACT

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.

研究动机与目标

• 解决推测解码中因分布偏移和能力差距导致的草稿模型推测准确率低下的问题。
• 利用大语言模型服务集群中未被充分利用的计算资源（空闲FLOPs）实现草稿模型的持续再训练。
• 在不增加推理成本的前提下，实现草稿模型对动态变化的查询分布的实时适应。
• 通过在用户查询数据上应用在线知识蒸馏，提升在线大语言模型服务中的token接受率并降低延迟。
• 证明在数据量有限的情况下，在查询特定数据上进行在线再训练的表现优于静态的离线微调。

提出的方法

• 利用大语言模型服务集群中的空闲计算能力（空闲FLOPs）对草稿模型进行在线再训练。
• 应用广义知识蒸馏（GKD）根据草稿模型与目标模型输出之间的差异更新草稿模型。
• 维护一个错误推测的缓冲区，并仅在空闲FLOPs可用时执行梯度更新，以最小化开销。
• 训练多个专用草稿模型，每个模型针对特定语言或主题聚类进行调优，以提升特定分布下的准确性。
• 在在线蒸馏过程中，使用草稿模型与目标模型输出分布之间的距离度量来计算损失。
• 机会式训练确保对推理延迟影响最小，同时支持持续的模型适应。

Figure 1: Online speculative decoding overview. For each prompt, the draft model suggests multiple tokens in a single step. The target model then verifies these tokens, accepting some and rejecting others. If the student proposes incorrect tokens, both the draft and target distributions are stored i

实验结果

研究问题

• RQ1 能否通过使用真实用户查询数据对草稿模型进行在线再训练，提升推测解码中的推测准确率？
• RQ2 当仅有少量查询数据可用时，在线推测解码与离线微调相比表现如何？
• RQ3 在大语言模型服务集群中，空闲FLOPs在多大程度上可被重新用于草稿模型训练，且不影响推理性能？
• RQ4 对查询分布偏移的实时适应是否能带来token接受率的持续提升？
• RQ5 多个针对特定语言或主题训练的专用草稿模型，是否能优于单一通用草稿模型？

主要发现

• 在线推测解码在多样化数据集上使token接受率提升0.1至0.65（10%–65%），显著提升推测效率。
• 该方法实现1.22倍至3.06倍的推理延迟降低，其中在真实世界的LMSYS-chat数据上收益最高。
• 在真实的大语言模型聊天对话中，即使每种语言的数据点少于2,000个，OSD也能使接受率提升0.1至0.2。
• 在主题特定的对话中，所有主题的接受率均超过0.6，社交与计算机科学类讨论的接受率接近0.8。
• 高频词如SELECT、FROM、WHERE和+在在线蒸馏后，其精确率和召回率均有显著提升。
• 尽管仅使用了少量数据，OSD的性能仍可匹配或超越在70%–100%查询数据上微调的模型。

Figure 3: Online acceptance rate ( $\alpha$ ) across different datasets. The x-axis represents the number of records that OSD has processed. Alpha is averaged over the most recent 50 records.