Speculative Decoding via Early-exiting for Faster LLM Inference with Thompson Sampling Control Mechanism (2024)

Abstract: The recent advancements in large language models (LLMs) have been extraordinary, yet the escalating inference costs associated with them present challenges in real-world applications. To address these challenges, we propose a novel approach called Early-exiting Speculative Decoding (EESD) with lossless acceleration. Specifically, EESD utilizes a segment of the LLM to generate draft tokens, incorporating Early-exiting structures after the first N layers. To enhance the quality of draft tokens, a self-distillation method is integrated. This early-exiting design not only reduces deployment and training costs but also significantly accelerates the token generation speed. Moreover, we introduce a novel sampling mechanism that leverages Thompson Sampling to regulate the generation processes, automatically determining the quantity of draft tokens in each round. The original LLM is then employed to validate these draft tokens through a single forward pass, and thus guarantees that the final output text maintains a distribution consistent with vanilla auto-regressive decoding. The experimental results on both 13B and 70B models demonstrate that our approach decodes tokens at a markedly accelerated rate compared to prior methods, showing the effectiveness of our approach.

heading: 1 Introduction
text: Large Language Models (LLMs) excel in various NLP tasks due to their immense parameters and complex network (OpenAI, 2023; Chowdhery et al., 2023; Touvron et al., 2023a,b). However, these models generate tokens one-by-one in an auto-regressive manner during inference, making the generation exceedingly resource-intensive and time-consuming. To overcome this bottleneck, researchers have introduced an effective decoding technique - Speculative Decoding (SD) (Leviathan et al., 2023; Chen et al., 2023; Miao et al., 2023).
∗ Jingang Wang is the corresponding author.
SD essentially introduces two models, a small model (the draft model) which is used to concurrently generate multiple draft tokens, and the original LLM (the target model) which is employed for draft token verification. In this way, SD maintains the same performance as the auto-regressive decoding while boosting the inference speed.
Compared to vanilla Speculative Decoding (Leviathan et al., 2023; Chen et al., 2023), several advanced models such as Medusa (Cai et al., 2023) and Self-SD (Zhang et al., 2023b) have been introduced, which only require deploying one LLM instead of two models, resulting in fewer resources required for both training and deployment. While these approaches achieve promising results, there are two main limitations. First, they fail to optimize the trade-off between the quality and speed of the draft token generation. For example, as shown in Figure 1a, while Medusa can generate draft tokens rapidly, the quality of
3027
these tokens tends to be subpar1. On the other hand, Self-SD manages to produce high quality draft tokens but does so at a much slower speed, resulting in lower overall speedup. Second, most SD approaches commence verification after generating a pre-defined length of draft tokens (referred to as drafting steps K). The choice of K significantly influences the acceleration of the inference process. Typically, larger drafting steps result in faster end-to-end generation, but there is a potential trade-off as the acceptance rate may decrease if the quality of the longer draft sequence is not high. As illustrated in Figure 1b, the optimal value of K varies across different examples. This variation suggests that utilizing a fixed K may not yield the most effective strategy. Instead, an adaptive method is preferable to determine when to terminate the drafting process.
To address these challenges, in this paper, we propose a novel Early-Exiting Speculative Decoding method, named EESD, to facilitate efficient and qualitative generation of draft tokens. Specifically, EESD introduce an Early-exiting layer that is superimposed on the first-N layers of the LLM, which has shown powerful predictive potential in previous research (Bae et al., 2023; Schuster et al., 2022). A self-distillation method is further employed to enhance the learning of the Early-exiting layer. To identify the optimal drafting steps, we reformulate the task of determining the length of draft token generation as a multi-armed bandit (MAB) problem, and propose a novel Control Mechanism based on Thompson Sampling (TS) that is well-studied for estimating unknown parameters and facilitating optimal decision making. Comprehensive evaluations on both 13B and 70B models demonstrate the superior performances of our approach over several baselines. The main contributions of this paper are summarized as follows:
• We introduce a novel Early-exiting framework for generating draft tokens, which allows a single LLM to fulfill the drafting and verification stages. We train it using self-distillation. Our investigations indicate that this framework strikes an excellent balance between the quality and speed of draft token generation.
• We conceptualize the generation length of draft tokens as a MAB problem and propose a
1We use the acceptance rate to represent the quality of the draft tokens, which is percent of draft tokens are accepted by the target model during the verification.
novel control mechanism based on Thompson Sampling, which leverages sampling to devise an optimal strategy.
• We conducted extensive experiments on three benchmarks. The results affirm that EESD can significantly improve the model’s inference speed, outperforming existing SD methods.

heading: 2 Related Work
text: LLM Compression The central objective of model compression is to alleviate computational demands and enhance inference speed. The research on LLM compression mainly includes three directions, including knowledge distillation (Zhang et al., 2023a; Li et al., 2023; Gu et al., 2023), network pruning (Ma et al., 2023; Xia et al., 2023; Frantar and Alistarh, 2023) and quantization (Xiao et al., 2023; Liu et al., 2023; Frantar et al., 2022; Lin et al., 2023; Gong et al., 2023). The methods mentioned above work by reducing the model’s footprint, thereby decreasing memory demand and enhancing the speed of inference. However, these methods sacrifices a degree of LLM’s capability.
Efficient Decoding Leviathan et al. (2023); Chen et al. (2023) propose a method that uses a small model to generate draft tokens and then uses LLM for verification, which accelerates the decoding process while guaranteeing lossless outputs, named as Speculative Decoding. However, some researchers suggest that the extra small model is not essential for SD. For instance, Medusa (Cai et al., 2023) generates draft tokens by leveraging additional parameters instead of small model, while Self-SD (Zhang et al., 2023b) uses the substructure of LLM to generate draft tokens. In addition, He et al. (2023) unveil a method that replaces the generation of draft tokens with a large text database. In the other hand, some researchers introduce an Early-exiting method. This method dynamically modifies the depth of the decoder for each token generation, making predictions at an intermediate layer (Teerapittayanon et al., 2016; Elbayad et al., 2020). Furthermore, Bae et al. (2023) propose a new Earlyexiting method that incorporates a shallow-deep module and synchronized parallel decoding.

heading: 3 Methodology
text: The overall model architecture of EESD is illustrated in Figure 2. Essentially, our model is com-
posed of three key components. 1) the Earlyexiting layer that built on top of the first few layers of the LLM as a draft model; 2) the self-distillation method to enhance the learning of the draft model and boost the text generation quality; and 3) the Thompson Sampling control mechanism that adaptively determines the length of the draft tokens conditioned on its quality. We present the details of these components in the following sections.

heading: 3.1 Early-exiting Layer
text: Most previous methods (Bae et al., 2023; Kavehzadeh et al., 2023) use non-continuous subnetwork of original LLM (target model) as their draft model. In this work, we utilize the continuous first-N layers approach, which yields one significant advantage: the kv-cache of the draft model and the target model can share the first-N layers, thereby trimming redundant computation. Concreately, we formulated an Early-exiting Layer with the computation process as elucidated below,
p(yt) = softmax(W TTransformere(HNt )), (1)
where HNt represents the hidden state of the N-th layer, which is calculated from the first-N layers of the origin LLM. And t represents t-th token. p(yt) is obtained from HNt through one layer of transformer. For LLaMA model, we also add RMSNorm layer before the output prediction head, which is RMSNorm(Transformere(HNt )).
As mentioned above, we incorporate an learnable Transformer layer subsequent to the first-N layers, and train this layer and W in Eq.(1) (RMSNorm parameters are trained as well for LLaMA
model). In order to speed up the model convergence, we initialize the Transformere and W with the last layer and predict head of the original LLM respectively. Since the training is confined to only a single Transformer layer and W , with the first-N layers being frozen, this approach dramatically reduce the computational resources.

heading: 3.2 Self-Distillation
text: To further enhance the effectiveness of the draft model, we employ self-distillation to learn the knowledge from the LLM. The key idea is that there is a large amount of valuable data used during the training of the LLM. However, it is usually impossible to obtain these original data as most of them are not directly accessible. Therefore, we propose to bridge the gap with self-distillation, which guides the learning of the early-exiting layer by transfering the knowledge from the generated text of the LLM. Specifically, Liu et al. (2023) suggest that a hybrid generation strategy of greedy and sampling is effective, and we adopt this approach for text generation from LLM. It is worth noting that the text generated by the LLM may contain certain lower-quality samples. We thus retain a subset of the open-source data for training purposes. The parameters of the Early-exiting Layer are trained utilizing an amalgamation of data generated by the LLM and open-source datasets, with the crossentropy loss between the prediction of it and the ground truth of mixed datasets.

heading: 3.3 Thompson Sampling Control Mechanism
text: The above methods can effectively improve the quality and speed of draft token generation. However, as we mentioned in the introduction, a predetermined drafting step is not a good strategy. Consequently, we view the controlling draft model generation as a MAB problem. Specifically, we view it as a Bernoulli process, where the model independently determines whether to continue draft token generating, denoted as PB(θ). And the probability θ is uncertain, related to the input sample. The Thompson Sampling (TS) method can better estimate unknown variables through balancing exploration and exploitation (Slivkins, 2019). As illustrated in Algorithm 1, we utilize TS algorithm to adaptively estimate θ. The crux of the TS algorithm involves modeling uncertain parameters θ as a posterior distribution using Bayesian theory, i.e. P (θ|D), with D representing observed environmental samples. The core of this algorithm lies in designing a reasonable posterior distribution, which we will detail in the following chapters.
TS with Beta Distribution Considering that the sample adheres to the Bernoulli distribution, we adopt the conjugate distribution approach and select the Beta distribution as posterior distribution. This setup means the prior and posterior distributions share the same distribution function but with different parameters, which greatly simplifies the computational process. The probability density function of the Beta distribution is as follows:
Beta(θ;α, β) = 1 B(α, β) θα−1(1− θ)β−1, (2)
where B(α, β) is a standardization function. The Beta distribution has two parameters, α and β, so Φ = {α, β} and Φ0 = {α0, β0} in Algorithm 1. According to Bayesian inference, we can update parameters α and β according to the following formula, which is the 18-th step in Algorithm 1,
αt = α{t−|Qv|} + r, (3)
βt = β{t−|Qv|} + (n− r), (4)
where r represents the number of successful experiments in the observed samples, and n represents the total number of experiments. And in Algorithm 1, the value of r is set to |Qv| − 1, indicating that the draft model should continue to generate on this set of tokens (i.e., χ = 1). And the value of n is set to min(|Qv|+ 1, |Qd|), indicating the number of tokens that have been validated by the target
model. Because we stop verifying when we encounter an inconsistent token, subsequent tokens are not considered.
Algorithm 1 TS Control Algorithm Require: Target Model Mt; Draft Model Md; Max Gen-
eration Length L; Hyperparameters Φ0; Input Prompt {x0, ..., xn}. 1: Initialize prior probability P (θ|Φ0) according to user-set hyperparameters Φ0. 2: Initialize the result set Qg ← {x0, ..., xn} and t← 0. 3: while t < L do 4: Initialize the draft model result set and i, Qd ← Null, i← 0. 5: while t + i < L do 6: Get new token xi ←Md(Qg ⋃ Qd). 7: Add token xi to set Qd. 8: Sample θt+i from P (θ|Φt;D). 9: Sample χ ∈ {0, 1} from Bernoulli distribution
PB(θt+i). 10: i← i+ 1. 11: if χ = 0 then 12: Break 13: end if 14: end while 15: Verify the results Qd by Target Model Mt and get Qv received by Mt, Qv ⊆ Qd. 16: Add the set Qv to Qg ,
Qg ← Qg ⋃
Qv . 17: Update t according to length of Qv , t← t+ |Qv|. 18: Update the parameters Φt of the posterior distribution. 19: end while 20: return Qg
TS with Calibration In the prior segment, we introduce a TS algorithm with Beta distribution to improve the estimation of θ. However, according to MAB theory, initial phases are more explorationfocused, which may result in less accurate θ estimations (Ou et al., 2019; Peng et al., 2019). To alleviate this issue, we propose a novel hybrid method that combines Model Prediction and Sampling Prediction. We rely more on model prediction to mitigate inaccuracies from initial exploration. As the sampling prediction begins to converge later, we calibrate the model prediction with sampling prediction to achieve a more precise θ.
We train a single-layer to predict the value of θ. The computation formula for this is as follows,
θt+iM = Sigmoid(Wp(W iH (T ) t , H (D) t+i )), (5)
where t is the number of tokens that have already been generated, and i is the number of new tokens generated by the draft model in the current loop. H
(T ) t represents the hidden state of the LLM (target model) at position t in the last layer, while the corresponding H(D)t+i represents the hidden state of
the draft model at position t + i in the last layer. W i is the transformation matrix at the i-th position for target model, and considering that i might be very large, we have restricted the number of W i ∈ Rd×d, i.e. i = min(i, 10). We sample a portion of training dataset to train the parameters {W 1,W 2, ...,W 10} and Wp ∈ R2×2d. The labels for this data are acquired by comparing the tokens produced by both the target and draft models, signifying the true value of χ. Following this, we update the parameters using cross-entropy loss.
According to the central limit theorem, when the sample size is sufficiently large, the sample mean adheres to a Gaussian distribution. Therefore, we make an assumption that sample mean χ̃ of χ in one drafting round follows a Gaussian distribution, i.e. χ̃ ∼ N (µ, σ2S). As supported by Bayesian theory, when the random variable follows a Gaussian distribution with a known variance but an unknown mean and the prior distribution is also a Gaussian distribution, it satisfies the conjugate distribution. Consequently, we define µ to follow a Gaussian distribution with the model’s predict score as mean and predict error as variance, µ ∼ N (θM , σ2M ). In Algorithm 1, we set Φ0 = {σM , σS , θ̂0}, where σM , σS and θ̂0 are hyperparameters set by the user. In Step 8 of Algorithm 1, we sample θ value from Gaussian distribution, θt+i ∼ N (µt+i, σ2t+i), and compute the values of µ and σ using the formula provided,
µt+i = σ2S
nσ2M + σ 2 S
θt+iM + nσ2M
nσ2M + σ 2 S
θ̂t, (6)
1
σ2t+1 =
1
σ2M +
n
σ2S , (7)
Where n is verification times. We update parameter Φ based on the following formula in step 18,
θ̂t = θ̂{t−|Qv|} ∗ (t− |Qv|+ 1) + |Qv|
t+ 1 , (8)
For a more detail, please refer to Appendix B.

heading: 4 Experiment

heading: 4.1 Setup
text: Training stage We randomly extract 100,000 samples from the SlimPajama (Soboleva et al., 2023) to train LLaMA-2-70B, LLaMA-2-13B and CodeLLaMA-2-13B. And use the ShareGPT2
2https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna _unfiltered
dataset to train LLaMA-2-70B-chat and Vicuna13B (Zheng et al., 2023a). We choose the first-5 layers as the draft model for 70B models and first3 for 13B models. For fair comparison, we train our model and Medusa (Cai et al., 2023) using the same data and set Medusa head at 4. More training details can be found in Appendix C
Evaluation stage We conduct experiments on three benchmarks under 1-shot setting: Gsm8k (Cobbe et al., 2021), XSum (Narayan et al., 2018) and Humaneval (Chen et al., 2021). And we also evaluate EESD on the MT-bench (Zheng et al., 2023b). We randomly select 500 instances from the test set for evaluation. And we set final output length at 512 and batch size at 1. We set the drafting step K at 10 for Vanilla SD (Chen et al., 2023) and Self SD (Zhang et al., 2023b). The reported results are the average of 10 different runs. We have only conducted on greedy generation3, as the findings from the top-p sampling exhibit similar trends.
Metrics We propose a Harmonic Mean (HM) to assess the quality of the draft tokens and strategy for generating them, while the specific calculation formula is as S = 2∗vd∗rdvd+rd ∗ 100%, where vd indicates the percent of draft tokens that are accepted by the target model, and rd represents the proportion of tokens that come from the draft model. More detailed explanation in Appendix D. Due to the verification process, all baseline and EESD methods can assure that the generation results are identical to the original LLM, hence we only need to compare their speedup.
All experiments are conducted on NVIDIA A100-80GB GPUs.

heading: 4.2 Main Results
text: We report evaluation results for Gsm8k and XSum in Table 1, for Humaneval in Table 2 and for MTbench in Table 3. As shown in Table 1, 2 and 3, it is clear that EESD significantly outperforms the previous methods on both 13B and 70B models, especially on LLaMA-2-70B, which demonstrates the effectiveness of our approach. There are several key observations from these results. First, we observe that EESD can yield 2.45× times speedup on CodeLLaMA-2-13B for coding task, suggesting
3For all experiments, we only generate one top-1 draft token candidate in Step 6 of Algorithm 1, and retain the result consistent with the target model’s top-1 token on verifying at Step 15 of Algorithm 1.
our method exhibits particular effectiveness within this domain. Second, compared to Vanilla SD and Medusa, EESD shows superior results with fewer training and deployment parameters. For instance, EESD achieves up to 2.13× and 1.80× times faster speeds on llama-2-70b model with just 1.12B parameters being trained. While we introduce an additional training process as compared to Self-SD, we manage to significantly improve speed effectiveness, utilizing minimal training resources. Third, we discover that a stronger capability of the draft model, indicated by a higher HM value, does not necessarily result in higher speedup. It is essential
to consider the generation speed of draft tokens, and our approach can strike an optimal balance between the two to achieve higher speedup (detailed in Appendix D).

heading: 5 Analysis and Discussion

heading: 5.1 Ablation Study
text: To elucidate the impact of different components within our approach, we conduct a series of ablation studies. In Table 4, we exhibit experimental
4https://huggingface.co/TinyLlama/TinyLlama-1.1Bintermediate-step-1431k-3T
5https://huggingface.co/double7/vicuna-68m
results, and several significant insights can be inferred. First, we notice a substantial decrement in the model’s performance when we replace the TS control with a fixed K value, which signifies the effectiveness of our proposed method for managing the generation of draft tokens. Second, similar to the prior approach, we introduce a trainable lm head just after the first-N layers, dispensing with the Early-exiting Layer. However, such a modification result in a significant decline in the model’s performance, strongly indicating the fundamental role of the Early-exiting Layer in maintaining the quality of draft tokens. Third, a noteworthy observation is that our approach attains commendable results solely with only open-source data, especially on XSum. Furthermore, the performance can be improved with the addition of self-distillation, demonstrating the utility of data generated by original LLM. Fourth, within the Cali-TS approach, the role of sample prediction surpasses that of model prediction, and an integration of both can yield more optimal results.

heading: 5.2 Can the TS control mechanism predict the optimal drafting steps?
text: To investigate the ability of the TS control mechanism to automatically determine the quantity of draft tokens in each round, we conducted experiments on the effects of varying drafting steps K. As illustrated in Figure 3, the optimal K value differs across models and datasets, but the TS with Beta distribution consistently slightly exceeds the effect of the optimal K value. Moreover, with the boost from the model prediction, the TS with calibration can achieve a better acceleration effect. The experiment confirmed that the TS control mechanism can adaptively predict the optimal length for generating
draft tokens in each round.

heading: 5.3 The generality of TS control mechanism
text: To verify the generality and effectiveness of the proposed TS control mechanism, we further apply it to other SD models instead of a pre-defined K. The results are reported in Table 5. According to the results, we can observe that TS control mechanism could be easily integrated into other SD methods to lift their performances. Note that the results in Table 5 are different from the results of w/o TS in ablation study. In ablation study, we set K at 10, which is not a superior setting, and as shown in Figure 3, K=5 is a better setting for the EESD of LLaMA-2-13B. However, for vanilla SD and self SD, K=10 is a suitable setting.

heading: 5.4 Effect of the first-N layers
text: Our experiments explore the impact of varying the number of first-N layers. As shown in Figure 4, using more layers improves the quality of the draft tokens, as measured by a higher HM value. However, end-to-end speedup does not correspondingly increase along with draft quality. This suggests that the extra time required to generate draft tokens with more layers offsets some of the end-to-end speedup.
The results indicate that layer augmentation only leads to slight improvements in the quality of draft tokens. Therefore, utilizing fewer layers for generating draft tokens proves to be an effective strategy. In addition, for larger models, such as 70B, the value of N needs to be slightly larger. And it is empirically suggested that N should be 5%-10% of the total number of LLM layers.

heading: 5.5 One Transformer layer is best for Early-exiting layer?
text: As demonstrated in Table 4, it has been proven that adding one Transformer layer after the First-N layers significantly improves the draft model’s performance. To further investigate, we evaluate the effect of increasing the number of Transformer layers. Figure 5 illustrates that augmenting the number of Transformer layers following the First-N layers does yield an improvement in draft token quality.
Model Method Seq. Loaded Trainable Num. Batch TimeParams Params of GPU Size per batch
However, because degree of this improvement is relatively small, it results in a reduction in the overall end-to-end speedup. Therefore, the experiment indicates that a single Transformer layer is enough to ensure the quality of draft tokens, and perfectly balance the quality and generation speed of draft tokens to achieve the optimal end-to-end speedup.

heading: 5.6 Training Efficiency
text: We compare training efficiency of three methods, which are tested on NVIDIA A100-80G GPUs. We set batch size to 64 and used SlimPajama datasets to train these models. As shown in Table 6, EESD only requires loading the parameters of First-N layers and Early-exiting layer, while Medusa requires loading all parameters of LLM and Medusa heads. Notably, during training, although both Medusa and EESD only update a portion of parameters, Medusa requires each sample to be computed across the whole LLM network. In contrast, EESD only needs to compute across First-N layers. Consequently, compared to Medusa and Vanilla SD, EESD significantly reduces in both training time and memory consumption.

heading: 5.7 Implement Tree Attention
text: Tree attention has been a prevalent technique in inference acceleration (Miao et al., 2023; Spector and Ré, 2023). This technique functions by structuring numerous draft token candidates within a
tree framework, allowing the LLM to concurrently verify several potential draft sequences through parallel decoding. It significantly increases the acceptance rate, thereby augmenting the overall speed of end-to-end generation. As shown in Table 7, we can easily implement the tree attention mechanism to EESD, resulting in significant increases in speed. It can achieve up to 2.48× and 1.96× times speedup on LLaMA-2-70B, as well as up to 2.18× and 2.12× times speedup on LLaMA-2-13B.

heading: 5.8 Breakdown of Computation
text: Table 8 presents an analysis of the computational time required for EESD generating on 200 instances randomly selected from XSum. The results indicate that Cali-TS exhibits higher time consumption compared to Beta-TS during the sampling phase. However, Cali-TS significantly diminishes the time usage in the drafting and verification stages, due to its superior control over the draft token generation process. Consequently, Cali-TS can yield a lower total time consumption.

heading: 6 Conclusion
text: In this work, we propose EESD, a novel method designed to lossless accelerate LLM by leveraging its first-N layers for generating draft tokens and employing Thompson Sampling to regulate this process. Specifically, we introduce an Earlyexiting layer after first-N layers and train it using self-distillation, which strike an optimal balance between efficiency and performance of draft token generation. Furthermore, we devise a novel hybrid method that effectively combines model prediction and sampling prediction, resulting in remarkable generation speed enhancement. After conducting exhaustive experiments, the results demonstrate that EESD not only achieves a significant inference speedup but also substantially reduces both training time and memory consumption, compared to previous speculative decoding methods.
Limitations
In this section, we discuss the limitations of our work as follows. First, while we have given an empirical suggestion for the setting of N value in the First-N layers, we have not thoroughly studied the function of these first layers and how they affect the final outputs. We believe that a more detailed investigation of this is helpful for choosing the optimal N value. As such, we will conduct this research in future work. Second, we propose a model for predicting whether the draft token is consistent with the LLM’s token in Section 3.3. However, this model has a large number of parameters, which is not very friendly for training and deployment. Therefore, we plan to refine the model’s structure to improve its efficiency in future work.

heading: Acknowledgements
text: Jingang Wang is funded by Beijing Nova Program(Grant NO. 20220484098). We sincerely thank all reviewers for their valuable comments and suggestions, which are crucial for improving our work.

heading: A Thompson Sampling with Beta Distribution
text: Given that the samples follow a Bernoulli distribution, we can infer using Bayes’ theorem that when the prior distribution is Beta, the posterior distribution is also Beta. This phenomenon, known as a conjugate distribution, means that the prior and posterior share the same distribution function but with different parameters. Using a conjugate distribution greatly simplifies the computational derivation, leading us to select the Beta distribution as the prior. As shown in Algorithm 2, we implement Thompson Sampling algorithm with Beta distribution to iteratively estimate the value of θ. The prior distribution of θ is Beta, and as previously described, the posterior distribution of θ is also Beta.
Algorithm 2 Thompson Sampling with Beta Distribution Algorithm Require: Target Model Mt; Draft Model Md; Max Gener-
ation Length L; Hyperparameters α0,β0; Input Prompt {x0, ..., xn}.
1: Initialize prior probability Beta(θ;α0, β0). 2: Initialize the result set Qg ← {x0, ..., xn} and t← 0. 3: while t < L do 4: Initialize the draft model result set and i, Qd ← Null, i← 0. 5: while t + i < L do 6: Get new token xi ←Md(Qg ⋃ Qd). 7: Add token xi to set Qd. 8: Sample θt+i from Beta(θ;αt, βt). 9: Sample χ ∈ {0, 1} from Bernoulli distribution
PB(θt+i). 10: i← i+ 1. 11: if χ = 0 then 12: Break 13: end if 14: end while 15: Verify the results Qd by Target Model Mt and get Qv received by Mt, Qv ⊆ Qd. 16: Add the set Qv to Qg ,
Qg ← Qg ⋃
Qv . 17: Update t according to length of Qv , t← t+ |Qv|. 18: Calculate r ← |Qv| − 1. 19: Calculate n← min(|Qv|+ 1, |Qd|) 20: Update αt, αt ← α{t−|Qv|} + r. 21: Update βt, βt ← β{t−|Qv|} + (n− r). 22: end while 23: return Qg

heading: B Thompson Sampling with Calibration
text: In the previous section, we unveiled a Thompson Sampling algorithm with Beta distribution (BetaTS). This method progressively update the parameters of Beta distribution to enhance the precision of the estimated θ. However, based on the MultiArm Bandit (MAB) theory, the early phase is more
exploration-oriented, this predisposition can lead to a less accurate initial estimation of the θ. To escalate the efficiency of the Thompson Sampling method, we further propose a hybrid approach of model prediction and sampling prediction. In early stage, we rely more on model prediction to curtail the inaccuracies introduced by exploration. In later stages, as the sampling prediction converges, we calibrate the model prediction with the result of sampling to obtain an accurate estimate of θ. The details are illustrated in Algorithm 3.
Algorithm 3 Thompson Sampling with Calibration Algorithm Require: Target Model Mt; Draft Model Md; Max Gener-
ation Length L; Hyperparameters σM ,σS ,µ0,σ0; Input Prompt {x0, ..., xn}. 1: Initialize prior probabilityN (µ0, σ20). 2: Initialize the result set Qg ← {x0, ..., xn} , n← 0 and
t← 0. 3: while t < L do 4: Initialize the draft model result set and i, Qd ← Null, i← 0. 5: while t + i < L do 6: Get new token xi ←Md(Qg ⋃ Qd). 7: Add token xi to set Qd. 8: Get model predict score, θt+iM ←
Sigmoid(MLP (W iH (T ) t , H (D) t+i )).
9: Calculate µt+i ← σ 2 S
nσ2 M +σ2 S θt+iM + nσ2M nσ2 M +σ2 S θ̂t.
10: Calculate σ2t+i ← σ 2 Mσ 2 S
σ2 S +nσ2 M .
11: Sample θt+i from Gaussian distribution N (µt+i, σ2t+i). 12: Sample χ ∈ {0, 1} from Bernoulli distribution PB(θt+i). 13: i← i+ 1. 14: if χ = 0 then 15: Break 16: end if 17: end while 18: Verify the results Qd by Target Model Mt and get Qv received by Mt, Qv ⊆ Qd. 19: Add the set Qv to Qg ,
Qg ← Qg ⋃
Qv . 20: Update t according to length of Qv , t← t+ |Qv|. 21: Update n← n+ 1. 22: Update θ̂t, θ̂t ← θ̂{t−|Qv|}∗(t−|Qv|+1)+|Qv|t+1 . 23: end while 24: return Qg
According to the central limit theorem, when the sample size is sufficiently large, the sample mean adheres to a Gaussian distribution. Therefore, we make an assumption that sample mean χ̃ of χ in one drafting iteration follows a Gaussian distribution, i.e χ̃ = χ1+χ2+χ3+...+χkk ∼ N (µ, σ2S). And According to the central limit theorem, σS equals σχ/sqrt(k), where σχ is the standard deviation of the random variable χ and we assume σχ
is known. In this case, we make the assumption that k is a fixed value and pre-determined by the user, which guarantees that χ̃ is independently and identically distributed. Furthermore, we employ a model to estimate the value of µ. It can be posited that µ follows a Gaussian distribution, characterized by the model’s predicted value as the mean and the model’s predicted error as the variance, i.e. µ ∼ N (θM , σ2M ), where θM is model’s predicted value. Here, we presume that the model’s predicted error is a known entity. As supported by Bayesian theory, when the random variable follows a Gaussian distribution with a known variance but an unknown mean and the prior distribution is also a Gaussian distribution, it satisfies the conjugate distribution. Therefore, the posterior distribution is following,
P (µ|D) ∝ P (D|µ, σ2S)P (µ|θM , σ2M )
∝ N ( σ 2 S
nσ2M + σ 2 S
θt+iM + nσ2M
nσ2M + σ 2 S
θ̂t, σ2Mσ 2 S
σ2S + nσ 2 M
)
(9)
where n is the number of verification by LLM, {σM , σS , θ̂0} is pre-determined by the user and θ̂t is the observed sample mean of the random variable χ̃. Due to k is a constant, we set θ̂t to observed sample mean of the random variable χ, i.e.
θ̂t = θ̂{t−|Qv|} ∗ (t− |Qv|+ 1) + |Qv|
t+ 1 , (10)
Carefully thinking Eq.(6) and Eq.(9), we can observe that when n is small, the mean µ tends to θt+iM , and when n is large, µ tends to θ̂t. This achieves the reduction of uncertainty in sampling prediction through model prediction in the early exploration stage. In the later stage, as the number of observed samples increases, the accuracy of θ̂t markedly enhances. Concurrently, µ draws closer to θ̂t, thereby yielding more accurate prediction. Therefore, our proposed method of mixing model prediction and sampling prediction can outperform the Beta-TS algorithm.

heading: C Training Details
text: We implement all experiments with the deep learning framework PyTorch on NVIDIA A100-80G GPUs. We set the learning rate to 1e-3 and the batch size to 64, for training EESD and Medusa. The hyperparameter settings we adopt are shown in Table 9

heading: D Harmonic Mean Metrics
text: We believe that two indicators, the acceptance rate and the proportion of draft tokens, will affect the end-to-end acceleration effect. The acceptance rate, denoted as vd, indicates the percentage of draft tokens that are accepted by the target model. It is calculated as follows,
vd = Nright
Nall_draft , (11)
where Nright is the number of draft tokens that are accepted by the target model, and Nall_draft denotes the total count of tokens generated by the draft model. The proportion of draft tokens, denoted as rd, represents the proportion of tokens that come from the draft model, which is calculated as follows,
rd = Nright
L , (12)
where L is the total number of tokens in the final output sequence.
Once we have computed the aforementioned two metrics, we can infer the speedup. The inference time of end-to-end generation can be calculated according to the following formula,
T = rd ∗ L vd Td + (1− rd) ∗ L ∗ Tt, (13)
Where Td represents the time taken by the draft model to generate one token, and Tt represents the time taken by the target model to generate one token. We can compute speed of the method by sp = LT and speedup by speedup = sp Tt
. Therefore, by integrating Eq.(13), we obtain the following formula,
speedup = vd
(α− vd) ∗ rd + vd , (14)
where α equals TdTt . To achieve speedup greater than 1.00×, the term α− vd must be less than zero, given that both vd and rd are positive values. The speedup of the method is influenced by the factors α, vd and rd. Under the premise that α < vd, ideally, α should be as low as possible, while vd and rd should be as high as possible. Both vd and rd are influenced by the quality of the draft tokens and the strategy for generating the draft tokens. A well-devised strategy for draft token generation can increase the values of vd and rd, but the upper bound is restricted by the inherent quality of the draft tokens themselves. In our experiments, we observe that as the increasing of the drafting steps
K, the acceptance rate vd usually decreases, while the proportion of draft tokens rd increases in most cases. Therefore, we hope to come up with a metric that can help us understand the optimality of K. And then, we use the harmonic mean of vd and rd to assess it.
As shown in Table 10, 11, 12 and 13, we present the vd and rd scores of each baseline method as well as our method across three benchmark evaluations, as detailed explanations of Table 1, 2 and 3.
E Inference Time and Speed up
We present the comprehensive results of end-to-end inference time and tokens generated per second
for the auto-regressive method, three speculative decoding methods and EESD. These results provide detailed explanations of the data shown in Table 1, 2 and 3. These results, gathered from evaluations using the Gsm8k benchmark, are detailed in Table 10. Furthermore, we provide the additional results from the XSum dataset in Table 11, the Humaneval dataset in Table 12 and the MT-bench dataset in Table 13.

heading: F Case Study
text: As shown in Figure 6, we demonstrated two Xsum samples. As described in Section 4.1, we adopt a 1-shot setting and use a greedy generation strategy.
We observe that during the generation process of EESD, those draft tokens that are inconsistent with the original LLM’s output will be discarded. This ensures that the final result generated by EESD is the same as auto-regression. Furthermore, we find that for samples with a high draft token acceptance rate, the EESD tends to generate longer draft sequence in one drafting round, as shown in example 1. Conversely, for samples with a lower acceptance rate, the EESD displays a tendency to generate shorter draft sequence, minimizing the quantity of discarded draft tokens, as shown in example 2. The examples in Figure 6 shows that our method is effective in adaptively determining the length of draft token generation, leading to significant improvement in the final end-to-end generation speed.

heading: Future Work
text: As such, we will conduct this research in future work. Second, we propose a model for predicting whether the draft token is consistent with the LLM’s token in Section 3. 3. However, this model has a large number of parameters, which is not very friendly for training and deployment. Therefore, we plan to refine the model’s structure to improve its efficiency in future work