Paper: Online Speculative Decoding (ICLR-2024)

Reviewer comment (Experiments):
In a few places in the paper, the author claims a translation between token acceptance rate and latency reduction. Is this done empirically or theoretically? Throughout the paper, the baseline seems to be against the offline distilled model and how the online model converges and eventually exceeds the performance of the offline distilled model, but the comparison did not include a vanilla model.

Author response:
**W1.1: Latency reduction calculation:** The reduction in latency is theoretically achieved by applying Equation 2 in the paper. The achievable speedup is decided by the token acceptance rate ($\alpha$) and the execution time ratio ($c$) of running the draft model and the target model. OSD focuses on improving the token acceptance rate and does not affect c. We implemented OSD in llama.cpp. Using LLaMA-1.1B as the draft model  and Vicuna-33B as the target model,  we got the measured and theoretical speedup on A100 below:

| Inference schemes | Baseline w/o speculative decoding | OSD, $\alpha$ = 0.5 | OSD, $\alpha$=0.6 | OSD, $\alpha$=0.7 | OSD, $\alpha$=0.8 | OSD, $\alpha$=0.9|
|--|--|--|--|--|--|--|
Measured time in ms/token (speedup) | 51.09 | 39.3 (1.28$\times$) | 35.48 (1.44$\times$) | 30.96 (1.65$\times$) | 25.42 (2.01$\times$) | 19.43 (2.63$\times$) |
Theoretical time  in ms/token (speedup) | 51.09 | 39.0 (1.31$\times$) | 32.12 (1.59$\times$) | 26.07 (1.96$\times$) | 20.77 (2.46$\times$) | 16.38 (3.12$\times$)|

**W1.2 Vanilla speculative decoding as the baseline:**
We define vanilla speculative decoding as the method where the unmodified, preliminary draft model is utilized to approximate the target model. **It is essential to note that this approach serves as the baseline for our comparisons.** The token acceptance rates associated with vanilla speculative decoding are listed in the 'Standard' column of Table 1. A significant contribution of this paper is the application of knowledge distillation to speculative decoding, a technique not researched in prior research.

**W1.3 Online speculative decoding baselines:** 
The experiment presented in Section 5.2 is **not** designed to demonstrate that the 'online model outperforms the offline model.' Instead, our objective is to illustrate the performance of the offline model (indicated by the red dashed line in Figure 2) as the maximum potential performance for the online model. This is due to the offline model being constructed assuming **complete awareness of the (future) data distribution** (note OSD does not). Additionally, another baseline we consider is the draft model that has been distilled using only 10% of the data (orange line in Figure 2). The distilled model does not update online. In this context, OSD exhibits superior performance, as it rapidly adapts to and learns from the underlying data distribution.

Reviewer comment (Experiments):
The author claims an expected improvement over vanilla speculative decoding but does not show it empirically.

Author response:
**W2: Compare with vanilla speculative decoding:** As mentioned above, we **indeed have compared** with the vanilla speculative decoding in the offline setting (vanilla speculative decoding and static model distilled with 10% data) as shown in Table 1 standard column.

Reviewer comment (Experiments):
The author showed the results of the online distilled model after two epochs. What's the performance like during the first two epochs of fine-tuning?

Author response:
**W4: Performance of finetuning one epoch:** We report the token acceptance rate $\alpha$ (LLaMA-160M as the draft model and Vicuna-7B as the target model) of the first two epochs in the offline setting below:
| Dataset | Spider | Gsm8k | Code-search-Python | Alpaca-finance |
|--|--|--|--|--|
| $\alpha$ without fine-tuning | 0.28 | 0.58 | 0.38 | 0.57 |
| $\alpha$ after fine-tuning one epoch | 0.71 | 0.72 | 0.58 | 0.65 |
| $\alpha$ after fine-tuning two epochs | 0.76 | 0.75 | 0.65 | 0.67 |

Reviewer comment (Experiments):
If we know that the performance improvement only shows after a certain amount of fine-tuning, does the real-world workload motivate this scenario? It's nice that the author considers the case of distribution shift, but the duration of each phase is also set arbitrarily and does not necessarily reflect the deployment scenario.

Author response:
**W5: Does the experiment motivate the real deployment scenario:** In Figure 2, we demonstrate that OSD is capable of rapidly increasing the token acceptance rate after observing less than 5K data points. We contend that this is a relatively small sample size in the context of real-world workloads. Regarding Figure 3, the primary aim of this figure is to illustrate OSD's quick adaptability to varying data distributions, even when distribution shift may be arbitrary,  as evident from its analysis of fewer than 2K records. We believe these findings provide robust evidence of OSD's effectiveness in real-world deployments, where the model would be tasked with processing hundreds of thousands of data points.

**W5: How does batch size affect speculative decoding?**
Thank you for your inquiry! Indeed, the batch size might impact the accuracy of Formula 2. This formula operates under the assumption that there are ample computational resources available to facilitate the verification step, as detailed in [1]. This assumption holds most of the time since LLM inference is typically constrained by memory bandwidth or capacity[2][3].  But still, investigating the intersection between speculative decoding and batching requests [2][4] constitutes an intriguing avenue for research. Notably, most existing studies [1][5][6] have concentrated on small batch sizes.  We consider it a potential area for future exploration.

**More comments on batch size and speculative decoding**
We would also like to further comment on the remark about batch size and speculative decoding. Speculative decoding is beneficial in systems that are not computationally bounded, which is partially determined by batch size as prior work on continuous batching [1] has shown. Exploring the tradeoff between speculative decoding and continuous batching is challenging, since the decision to speculate depends not only on the batch size, but also on how much computational resources are available and which of the requests from prior batches should be included if any of them failed to verify. We believe this topic merits its own paper and is orthogonal to ODS's contributions, among which is to show the efficacy of online speculation in systems that are not computationally-bounded.

Reviewer comment (Experiments):
All comparisons are against other knowledge distillation methods but lack an analysis against known speculative decoding techniques, such as vanilla speculative decoding and tree-based decoding.

Author response:
**W2: Compare with vanilla speculative decoding:** As mentioned above, we **indeed have compared** with the vanilla speculative decoding in the offline setting (vanilla speculative decoding and static model distilled with 10% data) as shown in Table 1 standard column.

Reviewer comment (Evaluation):
The authors claim there is spare compute as LLM serving is Memory Bandwidth bound. And based on this insight they propose OSD. However, concrete numbers regarding these are missing. Further the evaluation does not talk about runtime, only about token acceptance rates. Here is why I believe this is important: in my opinion/experiments for most Large LLMs we are on a roofline where we are memory bandwidth bound, and even the draft model is going to consume some amount of Memory Bandwidth when performing training. This could adversely affect the LLM being served, due to interference. Therefore, concrete numbers are going to be useful.

Author response:
**W1.1 OSD Runtime numbers:**
| Inference schemes | Baseline w/o speculative decoding | OSD, $\alpha$ = 0.5 | OSD, $\alpha$=0.6 | OSD, $\alpha$=0.7 | OSD, $\alpha$=0.8 | OSD, $\alpha$=0.9|
|--|--|--|--|--|--|--|
|Measured time in ms/token (speedup) | 51.09 | 39.9 (1.28$\times$) | 35.48 (1.44$\times$) | 30.96 (1.65$\times$) | 25.42 (2.01$\times$) | 19.43 (2.63$\times$) |
|Theoretical time  in ms/token (speedup) | 51.09 | 39.0 (1.31$\times$) | 32.12 (1.59$\times$) | 26.07 (1.96$\times$) | 20.77 (2.46$\times$) | 16.38 (3.12$\times$)|

Moreover, the observed speedup closely aligns with the theoretical expectations. The primary discrepancies can be attributed to two factors: (1) Slow Sampling: Speculative decoding needs additional sampling steps, as the draft model generates preliminary tokens. For optimal performance, the sampling process must be expedited. (2) To attain significant speedup, the execution time ratio (denoted as 'c') between the draft and target models should be minimized. However, in practical implementations, the overall execution time for the draft model is disproportionately affected by kernel launch overheads and Python-related delays, resulting in slower-than-anticipated performance. We are actively working on solving those engineering problems.

**W1.2 OSD might consume more memory bandwidth because of updating the draft model:**
As pointed out by the reviewer, LLM inference is memory bandwidth bound. When the input/output length is short, the memory operations are dominated by loading model parameters from GPU HBM to SRAM. We analyze the memory loading requirements of different inference techniques below (batch_size=1). 

We first introduce the notations used in the analysis. $M$/$m$: The total bytes of the target/draft model. $L$: inference length. $a_1$/$a_2$: The expected generation length for a single run of the target LLM of Vanilla speculative decoding(VSD)/OSD. $I$:  the interval to update the draft model.

On a high level, $\frac{L}{a} * M$ represents the bytes required to load the target model, while $L * m$ indicates the bytes needed for loading the draft model. For OSD, $m * \frac{L}{I}$ denotes the bytes necessary to load the draft model for updates.

|     | Memory Loading Formula | Memory Loading bytes of (LLaMA-160M, Vicuna-7B) pair, $L$=128, $a_1$=1.4, $a_2$=3.4 |
| -- | -- | -- |
| Inference without speculative decoding | $L * M$ | 1792 GB|
| Vanilla speculative decoding (VSD) | $\frac{L}{a_1} * M + L * m $ | 1320 GB |
| OSD | $\frac{L}{a_2} * M + L * m + m * \frac{L}{I} $ | 573 GB |

We applied Formula 2 from our paper to calculate $a_1$, $a_2$, using the token acceptance rates for VSD and OSD on the Spider dataset with the LLaMA-160M and Vicuna-7B models as the draft and target models, respectively. 
This resulted in $a_1 = 1.4$ and $a_2 = 3.4$. The memory sizes are $M$ = 14GB for the target model and $m$ = 0.32GB for the draft model. For OSD, the draft model is updated every 8 iterations ($I$=8). Using these values, we have estimated the memory loading bytes, presented in the right column.

As seen in table (1) Updating the draft model is not memory-intensive because the draft model is small. The majority of memory operations are still dominated by loading the target model. (2) OSD can significantly reduce memory bandwidth. This efficiency is achieved through a higher token acceptance rate, which consequently decreases the frequency of calling the larger model for verification.

Reviewer comment (Experiments):
My second concern is regarding data mixes. To be fair the authors have done a fair evaluation. However, I believe the evaluation is merely focussed on showing that OSD works. To me it is to some extent straightforward that as a model is fine-tuned on the same distribution it starts mimicking, therefore the offline evaluation is kind of straightforward. However, as the authors very well understand (from their online evaluation) it is not very straightforward. I am curious why the authors decided to have a separate model for each language. Is it a typical scenario for deploying speculative decoding? Further, can the authors report speculative decoding numbers on English language without filtering?

Author response:
**W2.1:** Could you explain “However, as the authors very well understand (from their online evaluation) it is not very straightforward.”?

**W2.2: Why use multiple draft models?** We propose that employing distinct draft models for queries on various topics can enhance the token acceptance rate. Unlike utilizing a single draft model for all topics, assigning specific draft models to individual topics narrows the range of query distributions each model must adapt to. This focused approach simplifies the learning process for each draft model, as they deal with a more limited set of queries. 
To substantiate this hypothesis, we measured and plotted the token acceptance rates using both strategies - a single universal draft model versus multiple topic-specific draft models in Figure 6 in the revised paper, to highlight the idea of customizing the draft model for different types of queries. As seen from the graph, across all topics, employing multiple draft models results in an increase in the token acceptance rate by 0.1 to 0.2. This aligns with our expectation that draft models benefit from a narrower query distribution, making it easier to learn and adapt. We leave it to future research to decide the optimal number of draft models and the best classification strategy.

**Please check the revised paper for Figure 6, please download the PDF instead of using the browser to open the PDF to avoid rendering issues. Thanks!**

Reviewer comment (Experiments):
I would really like to see where the authors think their approach will fail. Are there dataset mixes where this idea will fail? Can we evaluate straight up on LMSys-chat to see how it works without all the filtering?

Author response:
**W3: When will OSD fail?** The effectiveness of OSD can be impacted 
(1) **Rapid Changes in Query Distribution:** OSD may struggle in situations where the query distribution is in flux. If the distribution changes too quickly, the draft model faces challenges in adapting to such a diverse data range. This situation can lead to a stagnation or decrease in the token acceptance rate, diminishing the efficacy of OSD. 
(2) **Limitations in Compute-Bound Systems:** In systems that are already compute-bound, the trade-off that speculative decoding offers—extra computation in exchange for reduced memory operations—may not be beneficial. This is particularly true in batch-serving systems with very large batch sizes under high request loads, where the system's computational capacity is already fully utilized, rendering speculative decoding less effective. This constraint is more broadly a limitation of speculative decoding in general, rather than being specific to OSD.

Reviewer comment (Experiments):
Lack of comparison with SOTA works using "multiple draft models". One example [1].

Author response:
**W1: Could the authors compare the proposed online speculative decoding to the multi-head speculative decoding work?**
Thank you for mentioning Medusa. The main differences between OSD and Medusa are as follows: 
(1) Medusa needs extra training on a large model: The training process for Medusa requires more resources. Despite Medusa only fine-tuning the additional heads while keeping the original target model frozen, it still needs to run the forward pass on the entire target model. This aspect of the training process remains costly. Empirically, we compare the training time of Medusa heads of Vicuna-7B and the OSD distillation cost on LLaMA-160M as the draft model for Vicuna-7B in the following table:

|   | Extra parameters | Offline training time on a single A100 GPU for 1 epoch |
|--|--|--|
| Medusa | 0.44B | 7h |
| OSD | 0.16B | 32min |

(2) Medusa, being an offline algorithm, cannot adapt to shifts in data distribution, unlike the online nature of OSD: It requires pre-known target distributions and specific datasets to fine-tune its components. During distribution shifts, Medusa's static weights may lead to reduced token acceptance rates. Conversely, OSD is online and adapts quickly and automatically to distribution changes by periodically updating the draft model during non-peak hours, ensuring more responsive and efficient adjustments.

(3) Lastly, at this point, Medusa only supports decoder-only generative models like LLaMA while OSD is generally applicable to all transformer models.

Reviewer comment (Experiments):
The speedup is theoretically estimated. Lack of real-hardware evaluation.

Author response:
**Q2: The speedup is theoretically estimated. Lack of real-hardware evaluation. Could the authors show real hardware evaluation results?**
In comparison with vanilla LLM inference, we can achieve measured speedup using our online distilled draft models developed on top of the llama.cpp repo [2]. The experiments are conducted on a single A100 GPU with batch size = 1 and draft token length = 8,  results are attached below:
| Inference schemes | Baseline w/o speculative decoding | OSD, $\alpha$ = 0.5 | OSD, $\alpha$=0.6 | OSD, $\alpha$=0.7 | OSD, $\alpha$=0.8 | OSD, $\alpha$=0.9|
|--|--|--|--|--|--|--|
|Measured time in ms/token (speedup) | 51.09 | 39.9 (1.28$\times$) | 35.48 (1.44$\times$) | 30.96 (1.65$\times$) | 25.42 (2.01$\times$) | 19.43 (2.63$\times$) |
|Theoretical time  in ms/token (speedup) | 51.09 | 39.0 (1.31$\times$) | 32.12 (1.59$\times$) | 26.07 (1.96$\times$) | 20.77 (2.46$\times$) | 16.38 (3.12$\times$)|

From the table, OSD can obtain more than 2x speedup in comparison with vanilla LLM inference when we use TinyLLaMA-1.1B [3] as the student model and Vicuna 33B as the teacher model with an above 80% token acceptance rate. Moreover, the observed speedup closely aligns with the theoretical expectations. The primary discrepancies can be attributed to two factors: (1) Slow Sampling: Speculative decoding necessitates additional sampling steps, as the draft model generates preliminary tokens. For optimal performance, the sampling process must be expedited. (2) To attain significant speedup, the execution time ratio (denoted as 'c') between the draft and target models should be minimized. However, in practical implementations, the overall execution time for the draft model is disproportionately affected by kernel launch overheads and Python-related delays, resulting in slower-than-anticipated performance. We are actively working on solving those engineering problems.

[1] https://github.com/FasterDecoding/Medusa

[2] https://github.com/ggerganov/llama.cpp

[3] https://github.com/jzhang38/TinyLlama

Reviewer comment (Novelty):
Could the authors compare the proposed online speculative decoding to the multi-head speculative decoding work [1]? For example, can the proposed online update [1]? What are the potential challenges?

Author response:
**W2: For example, can the proposed online update [1]? What are the potential challenges?**
Certainly, OSD can be integrated into Medusa for online serving with concurrent model head updates. However, this approach presents certain challenges. Notably, training the Medusa heads demands more FLOPs compared to training the draft model, as estimated previously. Therefore, implementing Medusa with OSD would be more feasible when the system is under a lower load.

Reviewer comment (Experiments):
Could the authors show real hardware evaluation results?

Author response:
**Q2: The speedup is theoretically estimated. Lack of real-hardware evaluation. Could the authors show real hardware evaluation results?**
In comparison with vanilla LLM inference, we can achieve measured speedup using our online distilled draft models developed on top of the llama.cpp repo [2]. The experiments are conducted on a single A100 GPU with batch size = 1 and draft token length = 8,  results are attached below:
| Inference schemes | Baseline w/o speculative decoding | OSD, $\alpha$ = 0.5 | OSD, $\alpha$=0.6 | OSD, $\alpha$=0.7 | OSD, $\alpha$=0.8 | OSD, $\alpha$=0.9|
|--|--|--|--|--|--|--|
|Measured time in ms/token (speedup) | 51.09 | 39.9 (1.28$\times$) | 35.48 (1.44$\times$) | 30.96 (1.65$\times$) | 25.42 (2.01$\times$) | 19.43 (2.63$\times$) |
|Theoretical time  in ms/token (speedup) | 51.09 | 39.0 (1.31$\times$) | 32.12 (1.59$\times$) | 26.07 (1.96$\times$) | 20.77 (2.46$\times$) | 16.38 (3.12$\times$)|

From the table, OSD can obtain more than 2x speedup in comparison with vanilla LLM inference when we use TinyLLaMA-1.1B [3] as the student model and Vicuna 33B as the teacher model with an above 80% token acceptance rate. Moreover, the observed speedup closely aligns with the theoretical expectations. The primary discrepancies can be attributed to two factors: (1) Slow Sampling: Speculative decoding necessitates additional sampling steps, as the draft model generates preliminary tokens. For optimal performance, the sampling process must be expedited. (2) To attain significant speedup, the execution time ratio (denoted as 'c') between the draft and target models should be minimized. However, in practical implementations, the overall execution time for the draft model is disproportionately affected by kernel launch overheads and Python-related delays, resulting in slower-than-anticipated performance. We are actively working on solving those engineering problems.

[1] https://github.com/FasterDecoding/Medusa

[2] https://github.com/ggerganov/llama.cpp

[3] https://github.com/jzhang38/TinyLlama

Based on reviewers' requests, we measure  OSD offline distilled model on a single A100-80G across four evaluated datasets in the paper. We use teacher sampling with forward KD as the distillation method. Using TinyLLaMA-1.1B as the draft model and Vicuna-33B as the target model, we report the token latency and speedup compared with inference without speculative decoding in the table below. Inference without speculative decoding has a token latency of 51.09 ms/token and OSD can achieve 1.69$\times$ to 2.17$\times$ speedup across four evaluated datasets.

| Dataset | Spider | Gsm8k | Alpaca-Finance | Code-Python |
| -- | -- | -- | -- | -- |
|Measured time in ms/token (speedup) | 23.53 (2.17 $\times$) | 27.40 (1.89 $\times$) | 26.53 (1.92 $\times$) | 30.12 (1.69 $\times$)|