Decoding Speculative Decoding

2024-02-02

arXiv.org

AI Summary

This work addresses the critical challenge of draft model selection in speculative decoding, revealing that inference latency—not language modeling capability—is the primary bottleneck limiting acceleration. We therefore propose a hardware-aware draft model design paradigm: abandoning conventional language-modeling–centric selection criteria in favor of optimizing explicitly for low latency and high throughput via lightweight architecture customization and co-optimization with the inference engine. Evaluated across over 350 configurations on LLaMA-65B and OPT-66B, our approach achieves a 111% throughput improvement for draft models, substantially outperforming existing methods. Furthermore, we demonstrate strong generalization across diverse foundation models—including LLaMA-1, LLaMA-2, LLaMA-3.1—and multiple supervised fine-tuned variants. Our framework establishes a new pathway toward efficient, general-purpose, and hardware-adapted large language model inference.

Technology Category

• Machine Learning: Hardware-aware ML
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Speculative Decoding is a widely used technique to speed up inference for Large Language Models (LLMs) without sacrificing quality. When performing inference, speculative decoding uses a smaller draft model to generate speculative tokens and then uses the target LLM to verify those draft tokens. The speedup provided by speculative decoding heavily depends on the choice of the draft model. In this work, we perform a detailed study comprising over 350 experiments with LLaMA-65B and OPT-66B using speculative decoding and delineate the factors that affect the performance gain provided by speculative decoding. Our experiments indicate that the performance of speculative decoding depends heavily on the latency of the draft model, and the draft model's capability in language modeling does not correlate strongly with its performance in speculative decoding. Based on these insights we explore a new design space for draft models and design hardware-efficient draft models for speculative decoding. Our newly designed draft model can provide 111% higher throughput than existing draft models and our approach generalizes further to all LLaMA models (1/2/3.1) and supervised fine-tuned models.

Problem

• Optimizing draft model selection for speculative decoding
• Analyzing factors affecting speculative decoding performance
• Designing hardware-efficient draft models for LLMs

Innovation

• Hardware-efficient draft models
• Speculative decoding optimization
• Throughput improvement draft models