The Machine Learning Practitioner\u2019s Guide to Speculative Decoding
Image by Author
Mastering LLM Inference Speed with Speculative Decoding
Large language models meticulously construct text token by token. Each token generation necessitates a complete forward pass through the model, demanding the loading of billions of parameters from memory. This process inherently introduces latency into applications and escalates inference costs. Speculative decoding revolutionizes this paradigm by employing a compact draft model to generate a sequence of tokens. Subsequently, a larger, target model verifies these draft tokens in parallel. The outcome is output quality comparable to standard generation, coupled with a remarkable 2–3× inference speedup, and often even greater acceleration.
Understanding the Bottlenecks in Large Language Model Inference
Before dissecting speculative decoding, it's crucial to grasp the underlying challenges and foundational concepts that explain its efficacy.
The Challenge of Sequential Generation
Large language models operate autoregressively, generating text one token at a time. Each subsequent token is contingent upon the entire preceding sequence. The newly generated token is then appended to the input and fed back into the model for the next step. A token can represent a complete word, a word fragment, or even a single character, dictated by the model's tokenizer. The autoregressive generation process unfolds as follows: 1. The model ingests input tokens. 2. It executes a full forward pass through all its layers. 3. It predicts the probability distribution for the next token. 4. It samples or selects the most probable token. 5. This token is appended to the input sequence. 6. The cycle repeats from step 1. Consider generating the sentence “The scientist discovered a new species” (approximately six tokens). The model must perform six complete forward passes sequentially to achieve this.
The Memory Bandwidth Constraint
While one might assume computation is the primary bottleneck, given the sheer scale of these models, the reality is often different. Modern GPUs and TPUs boast immense computational power, but their memory bandwidth is considerably more constrained. The critical issue is that each forward pass requires the entire model's weights to be loaded from memory into the computation cores. For substantial models, this can involve transferring terabytes of data per generated token. Consequently, the GPU's compute cores frequently remain idle, awaiting data. This state is known as being memory-bound.
Nuances of Token Prediction Difficulty
It's noteworthy that not all tokens present an equal challenge to predict. Observe the following text:
After the sequence “The discovery was made in the”, predicting “Amazon” is relatively straightforward, as it appeared earlier in the context. However, following “The scientist discovered a new”, predicting “species” demands a deeper comprehension of semantic context and common research outcomes. The pivotal insight here is that if certain tokens are easier to predict, a smaller, faster model might adeptly handle them.The scientist discovered a new species in the Amazon. The discovery was made in the Amazon rainforest.
The Mechanics of Speculative Decoding
Speculative decoding draws inspiration from computer architecture's speculative execution technique. This involves performing operations in advance, assuming they will be necessary, and then verifying and discarding them if they prove incorrect. At a conceptual level, the objective is to mitigate sequential bottlenecks by decoupling rapid guessing from precise verification. This is achieved by:
- Utilizing a small, high-speed draft model to predict multiple tokens ahead.
- Employing a larger target model to validate all these predicted tokens concurrently in a single forward pass.
Stage 1: Token Speculation (Draft Generation)
The smaller, faster draft model generates several candidate tokens, typically predicting three to ten tokens ahead. While this model may not possess the absolute accuracy of the primary model, its speed advantage is significant. Visualize this as an agile assistant offering astute predictions of what comes next. Incidentally, speculative decoding is also referred to as assisted generation and is integrated within the Hugging Face Transformers library.
Stage 2: Parallel Verification
Following the draft model's token generation, the crucial next step is verification. Recall that a standard forward pass in the large model yields a single token. In this phase, however, a single forward pass through the larger target model is performed, using the entire sequence of draft tokens as input. Leveraging the inherent capabilities of transformer models, this unified forward pass generates probability distributions for the subsequent token at each position within the sequence. This enables simultaneous verification of all draft tokens. The computational load at this stage is comparable to a single standard forward pass, yet it facilitates the validation of multiple tokens.
Stage 3: Rejection Sampling for Precision
The final stage involves deciding which draft tokens to accept or reject. This is accomplished through a probabilistic method known as rejection sampling, which rigorously ensures that the output distribution precisely matches that of the target model's standard generation. For each draft token's position, a comparison is made between:
- P(draft): The probability assigned by the draft model to its chosen token.
- P(target): The probability assigned by the target model to the same token.
| 12345678910 | For each draft token in sequence: if P(target) >= P(draft): Accept the token (target agrees or is more confident) else: Accept with probability P(target)/P(draft) if rejected: Discard this token and all following draft tokens Generate one new token from target model Break and start next speculation round |
Illustrative example: Assume the draft model proposes the sequence “discovered a breakthrough”.
Token 1: “discovered”
- P(draft) = 0.6
- P(target) = 0.8
- Since 0.8 ≥ 0.6 → ACCEPT
Token 2: “a”
- P(draft) = 0.7
- P(target) = 0.75
- Since 0.75 ≥ 0.7 → ACCEPT
Token 3: “breakthrough”
- P(draft) = 0.5
- P(target) = 0.2
- Since 0.2 < 0.5, this token is flagged as questionable.
- The system rejects it and all subsequent draft tokens.
- The target model then generates its own token: “new”.
In this scenario, two draft tokens were accepted, and one new token was generated, yielding three tokens with effectively one target model forward pass (in addition to the draft token generation).
Consider the scenario where the draft model proposes K tokens. What occurs when all K draft tokens are accepted?
When the target model accepts all K draft tokens, the process generates a total of K+1 tokens in that iteration. The target model validates the K draft tokens and concurrently generates one additional token beyond them. For instance, if K=5 and all drafts are accepted, you obtain six tokens from a single target forward pass. This represents the optimal outcome: K+1 tokens per iteration compared to a single token in standard generation. The algorithm then proceeds with the extended sequence as the new input.
Key Performance Metrics for Speculative Decoding
To ascertain the effectiveness of speculative decoding for your specific application, diligent tracking of these metrics is essential.
Acceptance Rate (α)
This metric quantifies the probability that the target model accepts a draft token and stands as the paramount indicator of performance.
Example: If you draft five tokens per round and average three acceptances, your α = 0.6.
- High acceptance rate (α ≥ 0.7): Signifies excellent speedup; your draft model is optimally matched.
- Medium acceptance rate (α = 0.5–0.7): Indicates good speedup, making the technique worthwhile.
- Low acceptance rate (α < 0.5): Suggests poor speedup; consider an alternative draft model.
Speculative Token Count (γ)
This parameter dictates how many tokens your draft model proposes in each round and is fully configurable. The optimal γ selection is contingent upon your acceptance rate:
- High α: Employ a larger γ (7–10 tokens) to maximize speedup.
- Low α: Opt for a smaller γ (3–5 tokens) to prevent wasted computation.
Acceptance Length (τ)
This metric represents the average number of tokens actually accepted per round. A theoretical formula governs this:
Empirical benchmarks demonstrate that speculative decoding can achieve 2–3× speedups with robust acceptance rates (α ≥ 0.6, γ ≥ 5). Tasks heavily reliant on input grounding, such as translation or summarization, exhibit greater speedups, while more creative generation tasks may show less pronounced benefits.
Practical Implementation of Speculative Decoding
Let's implement speculative decoding leveraging the Hugging Face Transformers library. We will utilize Google's Gemma models: the 7B model as our target and the 2B model as our draft. Experimentation with various target and draft model pairings is encouraged, with the fundamental principle being that the target model is larger and more sophisticated, while the draft model is considerably smaller. Follow along with this Colab notebook for a hands-on experience.
Step 1: Installing Necessary Dependencies
Begin by installing the Hugging Face Transformers library along with PyTorch for streamlined model inference.
| 1 | pip install transformers torch accelerate huggingface_hub |
This command installs all essential components for efficient loading and execution of large language models.
Step 2: Loading the Models
Next, load both the target and draft models. A critical prerequisite is that both models must share the same tokenizer.
| 12345678910111213141516171819202122232425262728293031 | from transformers import AutoModelForCausalLM, AutoTokenizerimport torch # Choose your models - draft should be much smaller than targettarget_model_name = "google/gemma-7b-it"draft_model_name = "google/gemma-2b-it" # Set devicedevice = "cuda" if torch.cuda.is_available() else "cpu"print(f"Using device: {device}") # Load tokenizer (must be the same for both models)tokenizer = AutoTokenizer.from_pretrained(target_model_name) # Load target model (the large, high-quality model)print("Loading target model...")target_model = AutoModelForCausalLM.from_pretrained( target_model_name, torch_dtype=torch.float16, # Use fp16 for faster inference device_map="auto") # Load draft model (the small, fast model)print("Loading draft model...")draft_model = AutoModelForCausalLM.from_pretrained( draft_model_name, torch_dtype=torch.float16, device_map="auto") print("Models loaded successfully!") |
Accessing gated models like Gemma necessitates logging into Hugging Face. First, obtain a Hugging Face API token by navigating to huggingface.co/settings/tokens and creating a new access token with at least "read" permissions.
Option 1 (Recommended in Colab): Execute the subsequent code in a new cell and provide your token when prompted:
| 12 | from huggingface_hub import loginlogin() |
Option 2 (Environment Variable): Before running any code interacting with Hugging Face, configure the HF_TOKEN environment variable. For example:
| 12 | import osos.environ["HF_TOKEN"] = "YOUR_HF_TOKEN_HERE" |
For gated models, ensure you accept the license or terms of use on the Hugging Face model page prior to download. Once authenticated and approved, the model can be downloaded and utilized.
Step 3: Preparing Your Input Data
Craft a prompt and tokenize it. The tokenizer transforms text into numerical IDs for model processing.
| 12345678 | # Create a promptprompt = "Quantum entanglement is a phenomenon where" # Tokenize the inputinputs = tokenizer(prompt, return_tensors="pt").to(device) print(f"Input prompt: {prompt}")print(f"Input token count: {inputs['input_ids'].shape[1]}") |
The tokenizer segments the prompt into tokens, serving as the foundational context for subsequent generation.
Step 4: Establishing Baseline Autoregressive Inference
Begin by establishing a baseline through standard text generation. This serves as a crucial benchmark for evaluating the speedup achieved with speculative decoding.
| 12345678910111213141516171819 | import time # Standard generation (no speculation)print("\ --- Standard Generation (Baseline) ---")start_time = time.time() baseline_output = target_model.generate( **inputs, max_new_tokens=50, do_sample=False, pad_token_id=tokenizer.eos_token_id) baseline_time = time.time() - start_timebaseline_text = tokenizer.decode(baseline_output[0], skip_special_tokens=True) print(f"Generated text:\ {baseline_text}\ ")print(f"Time taken: {baseline_time:.2f} seconds")print(f"Tokens per second: {50/baseline_time:.2f}") |
Step 5: Harnessing Speculative Decoding for Generation
Activate speculative decoding by incorporating the assistant_model and num_assistant_tokens parameters into the generation call. These parameters instruct the target model to leverage the draft model for generating a specified number of tokens per speculation round.
| 12345678910111213141516171819202122232425262728 | import timeimport warnings # Import warnings module # Speculative decoding - just add assistant_model parameter!print("\ --- Speculative Decoding ---")start_time = time.time() with warnings.catch_warnings(): warnings.simplefilter("ignore") # Ignore all warnings within this block speculative_output = target_model.generate( **inputs, max_new_tokens=50, do_sample=True, # Set to False for greedy decoding pad_token_id=tokenizer.eos_token_id, assistant_model=draft_model, # This enables speculative decoding! num_assistant_tokens=10 ) speculative_time = time.time() - start_timespeculative_text = tokenizer.decode(speculative_output[0], skip_special_tokens=True) print(f"Generated text:\ {speculative_text}\ ")print(f"Time taken: {speculative_time:.2f} seconds")print(f"Tokens per second: {50/speculative_time:.2f}") # Calculate speedupspeedup = baseline_time / speculative_timeprint(f"\ Speedup: {speedup:.2f}x faster!") |
Expect to observe an approximate 2× performance enhancement. The magnitude of this speedup is intrinsically linked to the chosen target-draft model pairing. In essence, the draft model proposes candidate tokens, and the target model efficiently validates multiple options concurrently, drastically reducing the necessity for sequential forward passes through the larger model.
Strategic Application of Speculative Decoding
Based on extensive research and practical deployments, speculative decoding demonstrates optimal performance in the following scenarios:
Optimal Use Cases
- Speculative decoding excels at accelerating input-grounded tasks such as translation, summarization, and transcription.
- It is particularly effective when employing greedy decoding, consistently selecting the most probable token.
- Beneficial for low-temperature sampling, ensuring focused and predictable outputs.
- A valuable technique when the primary model barely fits within GPU memory constraints.
- It effectively reduces latency in production environments where adding more GPUs is infeasible.
Scenarios Where Speculative Decoding May Not Be Optimal
- Speculative decoding imposes an increased memory overhead, as both models must reside in memory.
- Its efficacy diminishes with high-temperature sampling, common in creative writing applications.
- Performance gains are notably reduced if the draft model is poorly matched to the target model.
- Benefits are marginal for very small target models that already occupy minimal memory.
Let us conclude with guidance on selecting an effective draft model that yields substantial improvements in inference times.
Selecting an Effective Draft Model
The success of speculative decoding hinges critically on the judicious selection of the draft model. An ill-chosen draft model may yield minimal speedup or even degrade performance. The ideal draft model possesses the following attributes:
- Identical tokenizer to the target model. This is an absolute requirement.
- At least 10× fewer parameters than the target model. An overly large draft model will slow down token generation, negating the intended benefit.
- Similar training data distribution to maximize the acceptance rate.
- Ideally, belonging to the same architecture family as the target model.
For specialized applications, consider fine-tuning a smaller model to closely emulate your target model's behavior. This can significantly enhance acceptance rates. The process involves:
- Gathering outputs from your target model using representative inputs.
- Fine-tuning a smaller model to accurately predict these gathered outputs.
This investment in fine-tuning pays dividends when consistent high performance is paramount in production environments. Consult Get 3× Faster LLM Inference with Speculative Decoding Using the Right Draft Model for deeper insights.
Concluding Thoughts on Speculative Decoding
Speculative decoding presents a pragmatic methodology for accelerating large language model inference without compromising output quality. By leveraging a smaller draft model to propose multiple tokens and subsequently verifying them in parallel with the target model, significant 2–3× speedups, or even greater, can be realized. This technique is effective because it directly addresses the memory-bound nature of LLM inference, minimizing the frequency of loading extensive model parameters. While performance is influenced by factors like draft model quality and acceptance rate, speculative decoding stands as a valuable technique for production systems where latency and cost efficiency are critical. Explore these supplementary resources for further understanding:
- Fast Inference from Transformers via Speculative Decoding
- Accelerating Large Language Model Decoding with Speculative Sampling
- Looking back at speculative decoding
Future articles will explore additional inference optimization techniques, offering further strategies to enhance the speed and cost-effectiveness of your large language model applications.