Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding

Existing LLMs typically follow an autoregressive decoding process. Given a prompt sequence $x_{1},...,x_{t}$, the model calculates the probability distribution of the next token $p(x|x_{1},...,x_{t})$. We present a greedy decoding process in Algorithm 1. In practice, instead of choosing the token with the highest probability (as in greedy decoding), we can sample tokens based on their probability distribution, which introduces some randomness and generates more diverse outputs.

Ideally, the computational cost of autoregressive decoding is comparable to that of sequence-level forward processing for an equivalent number of tokens.²²2 In fact, due to the causal nature of language modeling, autoregressive decoding could potentially save some attention computation. However, this decoding process is significantly bounded by the memory bandwidth of the device. When decoding each token, all the model parameters need to pass through the accelerator chip. So the model size divided by the memory bandwidth gives a hard ceiling on the decoding speed, resulting in a much longer inference time.

To mitigate the inherent inefficiency of autoregressive decoding, speculative decoding can be employed to enhance the inference speed of LLMs. This strategy involves two models: an LLM that we want to optimize, and a draft model that runs faster, albeit potentially at a lower quality. Speculative decoding can be explained as a two-stage process: (1) drafting: the draft model first generates $K$ draft tokens from a given prompt sequence $x_{1},...,x_{i}$, denoted as $x_{i+1},...,x_{i+K}$. (2) verification: following the drafting stage, the original LLM is then employed to validate these draft tokens. This validation is accomplished in a single forward pass, where the LLM predicts the probability distributions for each draft token and assesses whether they align with the draft tokens. Once a draft token $x_{j}$ is not validated, we use the original LLM’s prediction to override $x_{j}$, and start the next round of drafting beginning from token $x_{j+1}$.

The above process is based on the observation that computing the forward pass of a short continuation of tokens in parallel is not much slower than that of a single token. Consequently, the verification stage could be significantly more efficient than decoding tokens using the original LLM in standard autoregressive decoding.

In contrast to existing methods that use a standalone draft model to obtain draft tokens, our paper proposes a novel ‘self-speculative’ approach. We employ the original LLM itself for both the drafting and verification stages. During the drafting stage, the LLM selectively skips some of its intermediate layers so as to generate draft tokens quicker. Subsequently, these draft tokens are verified by the original LLM. Algorithm 2 presents a detailed description of the greedy decoding process. A complete sampling-based decoding process is elaborated in Appendix K.

Despite the simplicity of the main idea of self-speculative decoding, it poses several challenges:

Challenge 1: First, it is non-trivial to determine which layers and the number of layers to skip during drafting. If an excessive number of layers are skipped, the quality of the draft could be significantly compromised. This could result in a low acceptance rate in the verification stage, consequently increasing the overall inference time. On the other hand, if fewer layers are skipped, it ensures a higher acceptance, but also caps the maximum speedup that could be achieved.

Challenge 2: It is hard to decide when to stop the generation of draft tokens. As shown in Figure 2, the choice of the number of draft tokens to generate significantly influences the end-to-end speedup. In speculative decoding, if a draft token is rejected, all subsequent draft tokens will be discarded. Therefore, generating an excessive number of draft tokens could lead to unnecessary computation, thereby increasing the inference time.

In Sections 3.3 and 3.4, we will detail our approach to address these two challenges respectively.

The selection of skipped layers is essential in our method, shaping the configuration of the draft model and, consequently, the speedup achieved via self-speculative decoding. This selection process must carefully balance two key factors: the draft model’s ‘effectiveness’ and ‘efficiency.’ Both are intrinsically linked to the selection of skipped layers and significantly influence the overall performance of our method. Specifically:

(1) The ‘effectiveness’ of the draft model is quantified by the acceptance rate, which measures the agreement between the draft and verify models. However, we note that an exclusive focus on optimizing the acceptance rate could lead to no layers being skipped, i.e. the draft model identical to the verify model, resulting in an acceptance rate of 100%, but without any speedup.

(2) On the other hand, the ‘efficiency’ of the draft model can be quantified by the number of layers in the draft model. Indeed, minimizing the number of layers can reduce the inference latency of the draft model, but an extreme setup where all layers are skipped would result in the draft model generating low-quality tokens. This would drastically lower the acceptance rate, negating any potential speedup.

In this section, we frame the layer selection process as an optimization problem, our primary objective is to optimize the average inference time per verified token. This metric provides a comprehensive measure of the end-to-end inference time, including both drafting and verifying stages, normalized by the number of verified tokens.

Objective Function. This metric is a function of the selection of layers to be skipped in the draft model. The function, represented as $f(\boldsymbol{z})$, takes the selection of layers ($\boldsymbol{z}$) as input and returns the average inference time per verified token on a development set. Here, $\boldsymbol{z}$ is a vector that represents the layers to be skipped.

Optimization Problem. The optimization problem’s goal is to find the input $\boldsymbol{z}^{*}$ that minimizes the objective function $f(\boldsymbol{z})$. This problem can be formally expressed as:

While a brute force search could find the globally optimal solution for smaller models with a manageable solving space, it becomes prohibitively expensive for larger language models with many layers (e.g., $L=160$ for LLaMA-2-70B).

To tackle this, we employ Bayesian optimization Jones et al. (1998). As shown in Figure 3, it iteratively selects new inputs $\boldsymbol{z}^{*}$ for evaluation, based on a surrogate model of the objective function, i.e. Gaussian process Rasmussen et al. (2006), and an acquisition function. The latter balances exploration (testing inputs where the model’s prediction is uncertain) and exploitation (testing inputs where the model anticipates a favorable result). This procedure continues until a predetermined number of iterations is reached. We use the obtained $\boldsymbol{z}^{*}$ to accelerate text generation, and $\boldsymbol{z}^{*}$ is fixed for each model (i.e., generating draft model at model-level) without further updating. When the draft model’s target tasks vary significantly, task-level optimization is more appropriate to achieve good performance. Specifically, building the development set of optimization process from a single data source to mitigate inter-task interference.

In addition, while we here adopt skipping intermediate layers as a simple yet effective strategy to expedite the drafting stage, our scheme can be integrated with quantization Dettmers et al. (2022) and sparsification Sun et al. (2023) to further reduce resource consumption, as detailed in Appendix H.

Our self-speculative decoding approach incorporates an adaptive draft-exiting mechanism to enhance computational efficiency during the drafting stage. In speculative decoding, if a draft token is rejected, all subsequent draft tokens will be discarded accordingly. The draft-exiting mechanism prevents the wasteful allocation of computational resources toward draft tokens that are less likely to be accepted in the verification stage.

Specifically, it compares the predicted probability of each draft token against a threshold $\gamma$. If the predicted probability falls below $\gamma$ such that $p(x_{t+1}|x_{1},...,x_{t})<\gamma$, indicating low confidence, it immediately stops drafting. This approach ensures a better use of computing by focusing on the generation and verification of high-quality tokens, thereby improving the overall efficiency.

Moreover, it is worth noting that a static threshold may not accurately reflect the actual acceptance rate between the drafting and verification stages. For example, more challenging examples with a lower acceptance rate would be better served by a higher $\gamma$. To avoid the need for case-by-case threshold determination, we use an adaptive threshold that adjusts dynamically according to an updating rule, thereby allowing for an accurate reflection of the acceptance rate and better handling of examples in different difficulties. We denote the acceptance rate ($AR$) at $e$-th drafting stage as $AR_{e}$. Consequently, the update rule is defined as follows:

where $\alpha$ represents a target acceptance rate, $\epsilon$ is the update step-size, and $\beta_{1}$ and $\beta_{2}$ are factors designed to mitigate fluctuations of $\gamma$ and $AR$ respectively. Notably, when $e$ is 1, $\beta_{1}=0$. We update $\gamma$ after each verification stage. This updating rule ensures that the acceptance rate remains in close proximity to a target acceptance rate $\alpha$.

We evaluate a diverse range of models including LLaMA-2-13B, LLaMA-2-13B-Chat, CodeLLaMA-13B, and LLaMA-2-70B. Detailed setup can be found in Appendix B. We perform Bayesian optimization³³3https://github.com/bayesian-optimization/BayesianOptimization (MIT License) is used. (BO) for 1000 iterations to select the skipped layers in the drafting stage⁴⁴4Appendix B reports the offline BO time at model-level.. Results of tuning the number of BO iterations are reported in Appendix E. The datasets include CNN/Daily Mail (CNN/DM), Extreme Summarization (XSum), and HumanEval. These tasks cover the evaluation of text and code generation capabilities. Appendix I shows effectiveness on more diverse tasks such as solving math problems and open-domain chitchat. We perform 1-shot evaluation for CNN/DM and XSum, and compare the ROUGE-2 Lin (2004). We compare pass@1 and pass@10 Kulal et al. (2019) for HumanEval. We randomly sample 1000 instances from the testset for CNN/DM and XSum.

We evaluate the performance of our decoding scheme, denoted as ‘Self-Speculative’, with both greedy decoding (temperature = 0.0) and random sampling (temperature = 0.2/0.6), across text and code generation. The baseline is ‘Autoregressive’, which uses the original model to perform standard autoregressive decoding. The experiments involve various scales of LLaMA-2 and its fine-tuned models. The results can be found in Tables 1 and 2. We visualize the layer skipping distribution for different models in Appendix C.

For text generation tasks, Table 1 shows that our method, when applied with temperature settings of 0.0 and 0.2 achieves considerable speedups ranging from 1.210$\times$ to 1.992$\times$. Another important observation from these results is the minimal to nonexistent loss in ROUGE-2⁵⁵5 We attribute any slight differences observed in the case of greedy decoding to numerical rounding errors. , which verifies a core advantage of our decoding scheme, namely consistent output quality. In Section J.1, we compare the ROUGE-2 with other mainstream LLMs and evaluate the various metrics (ROUGE-1 and ROUGE-L) to quantitatively show our output quality. Moreover, the case study in Section J.2 qualitatively presents consistent output examples. In particular, our approach can be effectively applied on LLaMA-2-13B-Chat, a fine-tuned LLaMA-2-13B for conversation scenarios, indicating the compatibility of our method with fine-tuned models. This effectively addresses the dependency of the original speculative decoding method on high-quality draft models, which can be challenging to train and obtain, especially for fine-tuned models. Furthermore, the higher speedup achieved on LLaMA-2-70B suggests that larger models introduce more redundancy. This allows the drafting stage to skip a larger percentage of intermediate layers, thereby enhancing the overall speedup.

We also tested CodeLLaMA-13B, another fine-tuned LLaMA-2-13B for code generation. We used the HumanEval benchmark. Table 2 shows that we achieve speedups of 1.345$\times$ and 1.456$\times$, respectively, while maintaining similar task scores in terms of pass@1 and pass@10. This further validates the compatibility of our scheme for coding.

To investigate the impact of skipped layer selection, we conduct experiments on LLaMA-2-13B, which comprises 80 layers. Throughout the BO process, we track the number of layers skipped, denoted as $||\boldsymbol{z}^{*}||$, and the speedup relative to the autoregressive baseline. Figure 4 shows the results, where the dashed line indicates the maximum speedup for runs that skip the same number of layers.

These results reveal that: (1) The peak end-to-end speedup is observed when about half of the layers are skipped during the drafting stage; (2) The specific combination of layers skipped also plays a significant role. In particular, an inappropriate combination of skipped layers can actually result in a decrease in the end-to-end inference speed. (3) There is a noticeable drop in speedup when more than 42 layers are skipped. This suggests that the quality of drafting significantly deteriorates when an excessive number of layers are omitted.

These findings indicate the importance of layer selection in the implementation of self-speculative decoding. However, alternative layer skipping strategies do not achieve satisfactory speedup compared to BO, as detailed in Appendix D.

Performance degradation in drafting may be compensated by adopting aggressive skipping strategy and further training the draft model on a small amount of data, as described in Appendix F. This finding aligns with the Sheared-LLaMA Xia et al. (2023), which shows the effectiveness of pruning followed by fine-tuning on a small corpus.

Here we explore the effectiveness of the adaptive draft exit mechanism, specifically whether a threshold is needed and whether a static threshold is sufficient. Our settings are LLaMA-2-13B, CNN/DM, and greedy decoding.

Fixed Number of Draft Tokens. We evaluated a self-speculative decoding variant where the draft model generates a constant number $K$ of tokens. As Table 3 shows, speedup initially increases with $K$, then decreases. This pattern arises as a large $K$ (e.g., $K=8$) produces many tokens likely to fail verification, wasting computation in drafting.

While an appropriate $K$, e.g. $K=4$, can partially alleviate this issue, a static setting limits the draft model’s potential and achieves only a modest speedup (1.44$\times$). This is because a static $K$ does not adapt to the complexity of different instances. Ideally, we should use a larger $K$ for simpler instances and a smaller $K$ for challenging ones.

Moreover, Table 3 reveals that the acceptance rate and speedup do not have a direct proportionality. For instance, when $K=2$, the acceptance rate peaks at 0.924, yet the speedup is only 1.37$\times$. This discrepancy is due to the overly conservative $K=2$, which underutilizes the draft model’s capacity. By generating fewer draft tokens, we miss opportunities to produce more valid draft tokens, thus limiting the overall speedup.

Draft-Exiting with Static Threshold. Another variant is to stop generating draft tokens if the confidence score falls below a predefined static threshold. Table 4 shows that different static thresholds have large differences in acceleration (1.38$\times$~1.58$\times$). This highlights the importance of adaptively determining the appropriate threshold to optimize the speedup. Specifically, we observe that high thresholds ($\gamma$=0.8) tend to underestimate the capabilities of the drafting model. Despite a high acceptance rate ($AR$=0.935), this does not necessarily result in the best speedup due to a reduced number of drafting tokens (1.55$\times$). Conversely, a lower threshold ($\gamma$=0.2) tends to overestimate the drafting model’s capabilities, leading to a significantly lower acceptance rate ($AR$=0.749), wasting computational resources during the drafting stage, and thereby leading to slower inference speed (1.38$\times$).

Draft-Exiting with Adaptive Threshold. To address the issue of optimal threshold determination, we propose an adaptive draft-exiting mechanism. Specifically, we evaluate the acceptance rate and compare it to a target acceptance rate. The threshold is updated with an updating rule depicted in Section 3.3. Table 4 shows that the speedup achieved by our adaptive threshold update method (1.57$\times$) is comparable to, if not superior to, the speedup achieved with careful tuning of static thresholds. This indicates that dynamic threshold updating yields efficient and stable inference acceleration. This is mainly because the acceptance rate gets closer to the target $AR$ by adjusting the $\gamma$ value in a timely manner for instances of varying difficulties. Also, Appendix G reveals the adaptive draft-exiting is insensitive to changes in $K$.

Figure 5 captures the evolution of the draft-exiting threshold across different models and datasets over approximately 5,000 drafting and verification iterations. We can observe that our strategy adeptly adjusts the threshold, facilitating effective acceleration. However, the substantial variability across different models and datasets confirms the limitation of static threshold settings and the necessity for adaptive updates.

Table 5 presents a computation breakdown comparing the baseline with our ‘Self-Speculative’ decoding method. Our approach exhibits a significant speedup in average inference time per token compared to ‘Autoregressive’. This speedup is primarily attributed to two key techniques: the selection of skipped layers and the adaptive draft-exiting. Notably, the drafting stage consumes the majority of inference latency, highlighting the need for draft model optimizations to improve overall inference speedup. Importantly, our adaptive exit mechanism ($\gamma$ update) incurs negligible computational cost as it does not involve neural network calculations.