TLDR: HiSpec is a new framework for speculative decoding in Large Language Models (LLMs) that significantly improves inference throughput by addressing the bottleneck of token verification. It uses “early-exit” models for low-overhead intermediate verification, allowing inaccurate tokens to be discarded sooner. HiSpec also reuses key-value caches and hidden states for efficiency and maintains accuracy through periodic full-model verification. It achieves an average throughput improvement of 1.28x and up to 2.01x compared to baseline methods without sacrificing accuracy.
Large Language Models (LLMs) have transformed many applications, but their deployment often faces a critical challenge: balancing throughput, accuracy, and latency. While larger models offer superior accuracy, they typically come with increased latency and reduced throughput. Speculative decoding emerged as an efficient technique to tackle this, using a smaller ‘draft’ model to quickly propose tokens, which a larger ‘target’ model then verifies to maintain accuracy.
However, the verification step itself often becomes the bottleneck, significantly slowing down the overall process. For instance, verifying tokens can be four times slower than generating them when a 3B model drafts for a 70B target model. Most prior research has focused on accelerating the draft generation, leaving the ‘verification wall’ largely unaddressed. Existing attempts at ‘intermediate verification’ – where an auxiliary model discards inaccurate draft tokens early – have faced issues like substantial training overheads, increased memory footprint, and compromised accuracy due to reliance on approximations.
Introducing HiSpec: Hierarchical Speculative Decoding
A new framework, Hierarchical Speculative Decoding (HiSpec), proposes an innovative solution to this verification bottleneck. HiSpec leverages ‘early-exit’ (EE) models for low-overhead intermediate verification. EE models are specially trained LLMs that allow tokens to ‘exit early’ by skipping subsequent layers if a confident prediction can be made. This unique feature makes them ideal for intermediate verification without incurring the heavy computational and memory overheads seen in previous methods.
HiSpec’s design is centered on efficiency and accuracy. It reuses key-value (KV) caches and hidden states across the draft, intermediate verifier, and target models, minimizing redundant computations. To ensure the final output quality matches that of the full target model, HiSpec also periodically performs a full-model verification. This hierarchical approach allows for early rejection of inaccurate tokens, saving significant computational resources and accelerating subsequent draft generation.
How HiSpec Works
The framework operates by using different layers of an EE model for different roles. A lower exit layer acts as the draft model, quickly proposing tokens. An intermediate exit layer serves as the intermediate verifier, tentatively accepting or rejecting these tokens. The full model (or final layer) acts as the target verifier for ultimate accuracy. When tokens are rejected early by the intermediate verifier, the system saves time by not processing those inaccurate tokens through the deeper, more computationally expensive layers of the target model.
A crucial aspect of HiSpec is the strategic positioning of the intermediate verifier. Research shows that earlier model layers (up to one-fourth the depth) are highly effective, correctly generating a significant portion of the response. HiSpec typically positions its intermediate verifier around one-fourth of the model’s depth, with the draft generation occurring at about one-eighth of the depth. This ‘sweet spot’ maximizes throughput while maintaining a high token acceptance rate.
Furthermore, HiSpec employs a dynamic policy for balancing intermediate and full-model verification. Instead of fixed intervals, it waits until a sufficient number of tokens (e.g., four) have been tentatively accepted by the intermediate verifier before invoking the full target model. This ensures accuracy without excessively slowing down the process.
Also Read:
- Enhancing Diffusion LLM Performance with Adaptive Block Sizing
- Adaptive KV Cache for Faster, Better Language Models
Performance and Impact
Evaluations across various benchmarks, including real-user conversations (ShareGPT), text summarization (CNN/DM, XSum), code generation (HumanEval), and mathematical reasoning (GSM8K), demonstrate HiSpec’s effectiveness. It improves throughput by an average of 1.28 times and up to 2.01 times compared to baseline single-layer speculative decoding methods, all without compromising accuracy. HiSpec consistently outperforms state-of-the-art prior works that primarily focus on accelerating draft token generation, highlighting the significant benefits of addressing the verification bottleneck.
The framework is also highly adaptable, applicable to both pre-trained and post-training modified EE models, facilitating its widespread adoption. By providing a robust and efficient method for LLM inference, HiSpec marks a significant step forward in making large language models more practical and scalable for diverse applications. You can read the full research paper here.
