Speculative Decoding via Early-exiting for Faster LLM Inference with T hompson 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.

## introduction large language models (llms) excel in various nlp tasks due to their immense parameters and complex network @xcite @xcite . 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) @xcite @xcite . 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 @xcite , several advanced models such as medusa @xcite and self-sd @xcite 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 these tokens tends to be subpar foot_0 . 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 @xcite . 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: our investigations indicate that this framework strikes an excellent balance between the quality and speed of draft token generation. ## methodology the overall model architecture of eesd is illustrated in figure 2 . essentially, our model is com- early-exiting layer we divide the llm (the target model) into two parts: the first-n layers and the last-m layers. ## conclusion 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.