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

TL;DR

The paper tackles the high inference cost of large language models by introducing Early-exiting Speculative Decoding (EESD), which uses an Early-exiting layer on the initial N layers to draft tokens, a self-distillation process to boost draft quality, and a Thompson Sampling-based controller to adaptively determine drafting length. It presents Beta-TS and Calibration-TS variants to optimize drafting steps per instance, achieving substantial end-to-end speedups on 13B and 70B models with reduced training and deployment requirements. Comprehensive experiments across multiple benchmarks show EESD outperforming prior speculative decoding methods, including notable gains on CodeLLaMA-2-13B and robust performance with fewer trainable parameters. The work also analyzes component Contributions, model-structure choices, and practical enhancements like tree attention, offering a solid path toward faster, scalable LLM inference without sacrificing output fidelity.

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.

Figures (6)

• Figure 1: Experimental results using LLaMA-2-70B on the Gsm8k. (a) Speedup comparison with Medusa medusa and Self-SD DBLP:journals/corr/abs-2309-08168. EESD achieves a highest speedup with a best tradeoff between draft token generation speed and acceptance rate. (b) Generation costs (seconds) with different drafting steps (K) in randomly select five samples from Gsm8k. The optimal value of K varies across different samples, indicating that a fixed K value for all samples is not ideal.
• Figure 2: The framework of EESD which consists of three components: (1) Early-exiting layer which generate draft tokens efficiently and effectively; (2) Self-distillation which distills knowledge from the LLM (the target model); (3) TS control mechanism which can predict the optimal timing of terminating the draft token generation in each round. We divide the LLM (the target model) into two parts: the first-N layers and the last-M layers.
• Figure 3: We evaluate the speedup in generating 512 tokens using the EESD method at varying K values.
• Figure 4: Effect of the different first-N layers. We valuate EESD (+Beta-TS) across varying N values of the first-N layers.
• Figure 5: Effect of varying the number of Early-exiting layers.