PaperlayerView exact revision on arXiv

Source-linked AI summary

Online Speculative Decoding

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

arXiv:2310.07177v4cs.AIcs.CLcs.LG

TL;DR

Speculative decoding can lose effectiveness on diverse inputs when the draft model has limited predictive accuracy relative to the target. Online speculative decoding periodically updates draft models from target-model corrections on observed queries. Across synthetic and real data, it increases token acceptance by 10–65% and reduces latency by 1.4–2.1×.

  • Problem

    Speculative decoding’s effectiveness depends on draft-model accuracy, which can be low for diverse inputs given the capability gap between draft and target models.

    Source: §1 · p. 1

  • Method

    OSD periodically fine-tunes compact draft models using target-model corrections on observed queries, adapting them to the query distribution.

    Source: §1 · pp. 1–2

  • Results

    Token acceptance increases by 10–65%, with a 1.4–2.1× reduction in latency.

    Source: §1 · p. 2

  • Takeaways & Limitations

    OSD can be combined with static draft-model construction and match the accuracy achieved when all query data are available beforehand.

    Source: §1 · p. 2

    The optimal replay-buffer eviction or retention strategy remains for future exploration, while update interval I can be adjusted to system load and query-distribution change.

    Source: §4.2.1 · p. 4

Abstract

from arXiv · show

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.

1. Introduction

Speculative decoding can reduce LLM serving latency, but diverse inputs and a capability gap can limit draft-model accuracy. OSD uses target-model corrections to update draft models online and route queries to domain-matched drafts.

1. Introduction

  • Speculative decoding reduces latency by verifying multiple draft-proposed tokens in parallel, but diverse inputs and the draft–target capability gap can limit acceptance.Using a larger, more accurate draft model can increase latency and defeat the technique’s purpose.Source: §1 · p. 1
  • OSD periodically fine-tunes the draft model using target-model corrections, aiming to reduce query latency while preserving the draft’s compact size.Knowledge distillation aligns the draft and target probability distributions using token-level information from proposals and verification.Source: §1 · pp. 1–2
  • Because service queries often share domain-specific distributions, OSD continuously adapts draft models to incoming query data.The paper describes background fine-tuning applied periodically or updates using spare serving-system compute.Source: §1 · p. 2
  • OSD routes each query to a draft model tuned for its domain, narrowing the distribution each draft must handle.The paper reports training separate drafts for distinct languages or topics.Source: §1 · p. 2
  • OSD improves token acceptance rate by up to 10-65% on diverse datasets, translating into a 1.4-2.1× latency reduction.The paper also reports that OSD can match the accuracy achieved when all query data is available beforehand.Source: §1 · p. 2

2. Related Work

Prior speculative-decoding approaches improve draft accuracy through varied designs, while generally keeping drafts static after deployment. OSD instead adapts drafts to evolving query distributions and combines with existing draft-construction methods.

2. Related Work

  • Existing speculative-decoding methods improve accuracy with staged, retrieval-augmented, multiple-draft, or multi-candidate approaches, while another line uses auxiliary target-model modules.These approaches predominantly assume a static draft after deployment; OSD adapts to evolving query distributions and is described as orthogonal to them.Source: §2 · p. 2
  • Generalized knowledge distillation combines forward and reverse KL objectives using teacher- and student-sampled data, extending distillation for autoregressive LLMs.OSD applies distillation to speculative decoding and dynamically adjusts draft models; DistillSpec was developed concurrently and addresses a narrower part of OSD’s framework.Source: §2 · p. 2

3. Background

Speculative decoding uses a small draft model to propose tokens and a larger target model to verify them in parallel. Its acceptance rate relates to generated length and speedup, while OSD addresses the static-draft assumption by adapting to query data.

3. Background

  • A draft model proposes k tokens, and the target model evaluates their probabilities in parallel before the process repeats with the resulting sequence.The resulting samples are described as following the target model’s distribution.Source: §3 · p. 3
  • The acceptance rate α measures the expected probability that a proposed token is accepted and influences expected output length and speculative-decoding speedup.The speedup analysis assumes evaluating k + 1 target-model tokens takes roughly as long as generating one token.Source: §3 · p. 3
  • Speculative decoding can expose draft inaccuracies and target-model corrections, enabling draft refinement; prior approaches kept drafts static during deployment.The paper motivates OSD as a way to adapt drafts online.Source: §3 · p. 3

4. Online Speculative Decoding

Section 4 frames online speculative decoding as online knowledge distillation, using target-model corrections to adapt draft models during serving. It also describes routing queries to domain-specific draft models.

4.1. Knowledge Distillation for Speculative Decoding

  • Draft-model distillation aligns predictive distributions using forward KL, reverse KL, or their mixture; the best distance varies with task and relative model capacity.Speculative decoding provides proposed and verified token probabilities for constructing the distillation objectives.Source: §4.1 · pp. 3–4
  • Training can sample from a mixture of target and draft outputs, omit gradients through sampling, and estimate misalignment with a separately chosen distance measure.This separates the sampling distribution from the choice of KL, reverse KL, or JSD objective.Source: §4.1 · pp. 3–4

4.2.1. ALGORITHM

  • OSD starts from a warmup-trained draft model, records target logits at draft errors during serving, and periodically updates the draft from a replay buffer.The update interval is adjustable, and the loss choice depends on the draft-target pair and input data.Source: §4.2.1 · p. 4
  • Replay-buffer eviction and retention remain open choices; this study uses no eviction and releases the buffer after each model update.The update interval can be adjusted to system load and the rate of query-distribution change.Source: §4.2.1 · p. 4

4.2.2. LATENCY & FLOPS ANALYSIS

  • Against standard speculative decoding, expected speedups were 2.42×, 1.43×, 1.64×, and 1.22× for Vicuna-7B (LLaMA-160M draft), and 3.06×, 1.76×, 2.72×, and 1.55× for Flan-T5-XL 3B (T5-small 80M draft).The values are reported across four evaluated datasets for each model pair.Source: §4.2.2 · p. 5
  • Target-to-draft FLOPs ratios were 18.75 for (LLaMA-160M, Vicuna-7B) and 12.6 for (T5-small 80M, Flan-T5-XL 3B), while draft updates are not memory-intensive.The target model continues to dominate memory operations; higher acceptance reduces calls to it for verification.Source: §4.2.2 · p. 5

4.3. Draft Model Customization and Routing

  • Queries can be routed by language metadata or topic clusters from a 67M BERT model; the optimal routing queries and customized-model count or size remain future work.The authors report that a single draft model performs worse than domain-tailored drafts and note that query distributions shift over time.Source: §4.3 · p. 5

5. Experiments

Experiments test online speculative decoding across offline and online settings, multiple datasets, and customized draft models. They report higher acceptance rates and lower latency, while also examining distribution shifts, learned tokens, memory costs, and Medusa integration.

5.1. Offline Evaluation

  • Teacher sampling achieves the highest acceptance rate for both Vicuna-7B and FLAN-T5-XL.Reverse KL and JSD parallel or underperform forward KL; the best distance measure or sampling method varies by task and model.Source: Table 1 · p. 6, §5.1 · p. 6

5.2. Online Evaluation

  • Online acceptance starts below the offline-distilled model, then rises swiftly to comparable levels and sometimes exceeds offline test rates.Online rates are moving averages of the latest 50 requests, while offline rates use the full test set.Source: §5.2 · p. 6
  • At distribution boundaries, acceptance dips, then rebounds; after a shift, OSD catches up to or surpasses static distillation trained with 70%–100% of the online data.The comparison initially aligns with or trails the 30%-data setting; boundaries occur around 2K, 4K, and 6K records.Source: §5.2 · pp. 6–7, Figure 4 · p. 7

5.4. Latency Measurement

  • 2.63× is the maximum measured speedup when acceptance exceeds 0.9; at equal acceptance, TinyLLaMA-1.1B with Vicuna-33B is faster than LLaMA-160M with Vicuna-7B.The reported latency difference factor is c = 0.08 for the former pair and c = 0.13 for the latter.Source: Table 2 · p. 8, §5.4 · p. 8
  • 1.42×–2.17× speedup across four evaluated datasets with OSD’s distilled model.The measurements use two draft/target configurations; the reported maximum speedup is 2.63× when acceptance exceeds 0.9.Source: Table 3 · p. 8, §5.4 · p. 8

5.3. Real Workloads & Customized Draft Model

  • 0.1–0.2 higher acceptance with separate topic-specific drafts than with one universal draft across topics.Language-specific drafts also gain 0.1–0.2 with under 2K data points; topic acceptance exceeds 0.6, with Social and Computer near 0.9. Five preloaded drafts total 800M parameters, about 10% of the 7B target’s size, versus 2% overhead for one draft.Source: §5.3 · pp. 7–8, Figure 6 · p. 9, Figure 5 · p. 9

5.5. Qualitative Analysis

  • Distillation improves precision and recall for the 100 most frequent target-generated tokens, which account for 72.2% of appearances.Across datasets, improved tokens align with the data distribution; Spider examples include SELECT and WHERE.Source: §5.5 · p. 8, Figure 7 · p. 9

5.6. OSD and Medusa

  • 2.01× speedup for Medusa-v1 with OSD on Spider, versus 1.34× for Medusa alone.On Arena, Medusa-v1 without updates reaches 2.03×; standard speculative decoding with OSD reaches 1.51× and improves Medusa-v1 by 0.35×.Source: §5.6 · pp. 8–9, Table 4 · p. 9

6. Conclusion

Online speculative decoding continuously enhances the draft model across varying data distributions, and experiments on synthetic and real data show faster adaptation and higher token acceptance.

6. Conclusion

  • Online speculative decoding continuously enhances the draft model based on varying data distributions.The conclusion frames speculative decoding’s efficiency as depending on how well the draft model approximates the target.Source: §6 · p. 9
  • Experiments on synthetic and real data show that online speculative decoding swiftly adapts to new data distributions.The reported evaluations include both synthetic and real data.Source: §6 · p. 9
  • The experiments show that online speculative decoding significantly enhances token acceptance rate.The conclusion identifies improved acceptance as the reported outcome of adaptation.Source: §6 · p. 9

A.1. Speculative Decoding Sampling

Speculative decoding accepts draft proposals probabilistically, then samples a correction token when proposals fail; the resulting sequence follows the target model’s distribution.

A.1. Speculative Decoding Sampling

  • A proposal y_i is accepted when a uniform draw u is at most p(y_i|x, y_<i)/q_θ(y_i|x, y_<i); otherwise decoding exits the proposal loop.The process tests proposals in order from i = 1 to k.Source: §A.1 · p. 12
  • After the accepted proposals, an additional token y_(a+1) is sampled from a normalized distribution.Here, a is the number of accepted tokens; norm(·) makes probabilities over the vocabulary sum to 1.Source: §A.1 · p. 12
  • The sequence of accepted proposals and the additional token strictly follows the target LLM distribution, and each run generates at least one new token.The generated tokens are appended to the context and the process repeats until EOS.Source: §A.1 · p. 12

A.2. Speculative Decoding Speedup

Speculative-decoding speedup depends on acceptance rate, proposal length, and draft-to-target run time; its relationship with acceptance rate is superlinear.

A.2. Speculative Decoding Speedup

  • Doubling acceptance yields over 2× speedup, while small α can make speculative decoding slower than standard inference.Figure 8 describes the relationship as superlinear; its c is the draft-to-target single-run time ratio and k is proposal length.Source: Figure 8 · p. 12
  • The expected improvement factor for offline speculative decoding is expressed in terms of α, k, and c.The analysis defines c as the time for one draft-model run divided by the time for one target-model run.Source: §A.2 · p. 12
  • The speedup analysis compares average time to produce one token under speculative decoding with standard decoding time.The passages express speculative token time using the target-model run time T and the expected number of generated tokens per run.Source: §A.2 · p. 12

A.3. Latency Analysis

Improving acceptance from α1 to α2 reduces average token generation time and yields a speedup over standard speculative decoding; the analysis omits draft-update latency under stated conditions.

A.3. Latency Analysis

  • Improving acceptance from α1 to α2 lowers average generation time per token.The analysis expresses the resulting time using T, the generation time for standard decoding.Source: §A.3 · p. 12
  • The resulting speedup factor is given relative to standard speculative decoding as a ratio of sums of powers of α2 and α1.The expression compares the acceptance rates before and after improvement.Source: §A.3 · p. 12
  • The analysis omits draft-update latency because update computation is marginal relative to large-model computation and updates occur during moderate request loads.The authors also note that updates may be masked or run concurrently on separate devices.Source: §A.3 · p. 12

A.4. Flops Analysis

The FLOPs analysis accounts for draft proposals, target-model verification, and periodic draft updates, then compares the resulting costs with target-model inference. It estimates that draft updates require substantially fewer FLOPs than large-model inference.

A.4. Flops Analysis

  • Draft-model updates require significantly fewer FLOPs than large-model inference, while the cited LMSYS-Chat-1M workload analysis estimates notably low GPU computation utilization.The workload estimate uses a 30B model served with eight A100 GPUs; the authors note that Arena may not represent an efficient or high-traffic service.Source: §A.4 · p. 13
  • The cost model counts draft-token proposals, target-model verification, and updates across I requests, approximating training at three times inference cost per token.The simplified total FLOPs expression combines draft forward work with target forward work.Source: §A.4 · p. 13
  • The analysis defines the target-to-draft FLOPs proportion as a way to compare the models’ computational demands.The supplied passage states the proportion but does not provide its displayed formula.Source: §A.4 · p. 13
  • 18.75 for LLaMA-160M/Vicuna7B at 0.71 average acceptance and 12.6 for T5-small 80M/Flan-T5-XL 3B at 0.76 are the estimated target-to-draft FLOPs ratios.Both estimates assume an average of five proposed tokens per run.Source: §A.4 · p. 13

A.5. Bandwidth Analysis

These analyses examine bandwidth requirements, adaptation across mixed and shifting query distributions, and speedup under different evaluation conditions. They report acceptance-rate behavior during adaptation and compare measured with theoretical speedup, including factors that affect realized latency.

A.5. Bandwidth Analysis

  • For short inputs, memory operations are dominated by loading model parameters; the bandwidth analysis estimates loading for target and draft models and for draft updates.Its estimate uses Spider acceptance rates for LLaMA-160M and Vicuna-7B, with the draft updated every eight iterations.Source: §A.5 · p. 13, Table 5 · p. 14

A.6. Data Mix

  • During the mixed-data experiment, Gsm8k acceptance rises with exposure; Alpaca-finance acceptance initially stays consistent, then rises, while Gsm8k remains consistent in the latter half.The authors report no signs of forgetting the learned Gsm8k knowledge in this experiment.Source: §A.6 · p. 14, Figure 9 · p. 14

A.7. Real Workloads

  • The real-workload evaluation samples 10K records from LMSYS-Chat-1M, which spans 25 models and over 150 languages, and uses only Vicuna conversations.The dataset includes interactions from April to August 2023.Source: §A.7 · p. 14

A.8. Various Degrees of Distribution Shifts

  • After an abrupt Gsm8k-to-Spider shift, acceptance dips near the transition but quickly recovers; under gradual shift, the decline is less discernible.The gradual-shift simulation changes the dataset mixture progressively, while the abrupt case switches between datasets without a transition.Source: §A.8 · pp. 14–15, Figure 11 · p. 15

A.9. Theoretical and Measured Latency

  • Measured speedup closely aligns with theoretical expectations, while slow sampling and draft-model kernel-launch and Python overhead contribute to discrepancies.The authors also note that significant speedup requires minimizing the draft-to-target execution-time ratio.Source: §A.9 · p. 15, Table 7 · p. 15

A.9.1. GENERATION LENGTH AND SPEEDUP

  • Measured speedup increases as generation length grows from 16 to 2048 tokens, but improving token acceptance rate has a stronger effect than extending generation length.The experiment fixes batch size at 1, prefix length at 512, and draft token length at 5.Source: §A.9.1 · p. 16, Table 8 · p. 16

A.9.2. DRAFT MODEL EXECUTION TIME

  • With draft updates every eight iterations, target-model forward execution takes 208 times longer for the 160M/7B pair and 82.3 times longer for the 1.1B/33B pair.Execution times were measured on one A100-80G with prompt length 64 and generation length 128.Source: §A.9.2 · p. 16
Loading 2310.07177v4…