[Paper Review] Online Speculative Decoding

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

Topic ModelingComputer Science3 citations

TL;DR

This paper proposes online speculative decoding, a method that continuously retrains draft models on real-time user query data using idle computational resources in LLM serving clusters. By applying online knowledge distillation, the draft model adapts to the query distribution in real time, increasing token acceptance rates by 10–65% and reducing inference latency by 1.22× to 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.

Motivation & Objective

• Address the low speculation accuracy of draft models in speculative decoding due to distribution shifts and capability gaps.
• Leverage underutilized computational resources (spare FLOPs) in LLM serving clusters for continuous draft model retraining.
• Enable real-time adaptation of draft models to evolving query distributions without increasing inference cost.
• Improve token acceptance rate and reduce latency in online LLM serving through online knowledge distillation on user query data.
• Demonstrate that online retraining on query-specific data outperforms static, offline fine-tuning even with limited data.

Proposed method

• Utilize idle computational capacity (spare FLOPs) in LLM serving clusters to perform online retraining of draft models.
• Apply Generalized Knowledge Distillation (GKD) to update draft models based on discrepancies between draft and target model outputs.
• Maintain a buffer of incorrect speculations and perform gradient updates only when spare FLOPs are available, minimizing overhead.
• Train multiple specialized draft models, each tuned to a specific language or topic cluster, for improved distribution-specific accuracy.
• Use distance metrics between draft and target model distributions to compute loss during online distillation.
• Opportunistic training ensures minimal impact on inference latency while enabling continuous model adaptation.

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

Experimental results

Research questions

• RQ1 Can online retraining of draft models using real user query data improve speculation accuracy in speculative decoding?
• RQ2 How does online speculative decoding perform compared to offline fine-tuning when only a small fraction of query data is available?
• RQ3 To what extent can spare FLOPs in an LLM serving cluster be repurposed for draft model training without affecting inference performance?
• RQ4 Does real-time adaptation to query distribution shifts lead to sustained improvements in token acceptance rate?
• RQ5 Can multiple specialized draft models, each trained on a specific language or topic, outperform a single general-purpose draft model?

Key findings

• Online speculative decoding increases token acceptance rate by 0.1 to 0.65 (10–65%) on diverse datasets, significantly improving speculation efficiency.
• The method achieves a 1.22× to 3.06× reduction in inference latency, with the highest gains observed on real-world LMSYS-chat data.
• On real LLM chat conversations, OSD improves acceptance rates by 0.1 to 0.2 even with fewer than 2,000 data points per language.
• For topic-specific conversations, acceptance rates exceed 0.6 across all topics, with Social and Computer Science discussions reaching near 0.8.
• High-frequency tokens such as SELECT, FROM, WHERE, and + show marked improvements in both precision and recall after online distillation.
• OSD matches or exceeds the performance of models fine-tuned on 70–100% of the query data, despite using only a fraction of the data.

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.