ML

Speculative Decoding & Draft Model Distillation

Training efficient draft models for speculative decoding: distilling small models to match large model token distributions, acceptance rate optimization, MedUSA heads, and Eagle-style draft architectures.

Advanced

Prerequisites

Table of Contents

  1. Learning Objectives
  2. Notation
  3. Core Intuition
  4. Speculative Decoding Overview
  5. Training Draft Models via Distillation
  6. Optimizing Acceptance Rate
  7. Medusa: Multi-Head Draft
  8. Eagle: Autoregressive Draft Heads
  9. Online Speculative Distillation
  10. Common Pitfalls
  11. Summary
  12. Exercises

Learning Objectives

  1. Explain speculative decoding and the role of the draft model.
  2. Derive the optimal distillation objective for maximizing acceptance rate.
  3. Describe Medusa's parallel multi-head drafting.
  4. Explain Eagle's feature-level autoregressive drafting.
  5. Design a distillation pipeline for creating draft models.

Notation

  • pp — target (large) model distribution
  • qq — draft (small) model distribution
  • γ\gamma — expected acceptance length
  • α\alpha — acceptance rate per token

Core Intuition

Speculative decoding uses a small "draft" model to quickly propose multiple tokens, then a large model verifies them in parallel. The draft model's quality determines the acceptance rate — how many proposed tokens the large model agrees with. Distillation is the IDEAL way to train draft models: make the small model's distribution match the large model's distribution as closely as possible, maximizing the acceptance rate.

Speculative Decoding Distillation

Draft model (distilled)Target model (verify)t1t2t3t4✗t5✗Acceptance: 63% | Speedup: 2.6×Draft quality 60% → 3/5 tokens acceptedDistill draft from target to maximize token acceptance rate
Quality
0.60
AcceptedRejected
Explore: Speculative decoding distillation trains a small draft model to match the target's token distribution — higher acceptance means fewer verification steps and greater speedup.

Speculative Decoding Overview

Algorithm:

  1. Draft model generates KK tokens autoregressively (fast, small model).
  2. Target model verifies all KK tokens in ONE forward pass (parallel).
  3. Accept longest prefix where draft matches target (rejection sampling).
  4. Expected speedup: γ/(1+draft_cost/target_cost)\gamma / (1 + \text{draft\_cost}/\text{target\_cost}) where γ\gamma is expected accepted length.

Acceptance probability for token ii:

αi=min⁡(1,p(xi∣x<i)q(xi∣x<i)).(1)\alpha_i = \min\left(1, \frac{p(x_i | x_{<i})}{q(x_i | x_{<i})}\right). \tag{1}

Expected accepted length:

γ=∑k=1K∏i=1kE[αi]≈11−αˉ,(2)\gamma = \sum_{k=1}^K \prod_{i=1}^k \mathbb{E}[\alpha_i] \approx \frac{1}{1 - \bar{\alpha}}, \tag{2}

where αˉ\bar{\alpha} is the average per-token acceptance rate.

For αˉ=0.8\bar{\alpha} = 0.8: Expected acceptance = 5 tokens. With 7B draft and 70B target: 3-4x speedup.

Training Draft Models via Distillation

Standard approach: Distill the large model into a smaller architecture:

L=Ex∼D∑tDKL(p(⋅∣x<t)∥qθ(⋅∣x<t)).(3)\mathcal{L} = \mathbb{E}_{x \sim \mathcal{D}} \sum_t D_{\text{KL}}(p(\cdot|x_{<t}) \| q_\theta(\cdot|x_{<t})). \tag{3}

Token-level KL minimization directly optimizes acceptance rate, because:

αˉ=1−12E[TV(p,q)]≥1−12DKL(p∥q).(4)\bar{\alpha} = 1 - \frac{1}{2}\mathbb{E}\left[\text{TV}(p, q)\right] \geq 1 - \sqrt{\frac{1}{2}D_{\text{KL}}(p\|q)}. \tag{4}

Lower KL → higher acceptance → more speedup.

On-policy distillation (using target model's own generations as training data):

  • Generate sequences from the target model.
  • Compute target model's next-token distributions.
  • Train draft on these distributions.
  • Better than off-policy (training data distribution = inference distribution).

Optimizing Acceptance Rate

Beyond KL minimization:

Top-k focused distillation: Acceptance mostly depends on high-probability tokens. Focus distillation on the top-kk entries of the target distribution:

Ltop-k=DKL(ptop-k∥qtop-k).(5)\mathcal{L}_{\text{top-k}} = D_{\text{KL}}(p_{\text{top-k}} \| q_{\text{top-k}}). \tag{5}

Temperature matching: Use same temperature for draft and target (usually T=0T=0 for greedy, or matched sampling temperature).

Domain-specific drafts: A small model fine-tuned on the deployment domain has higher acceptance than a generic small model.

Medusa: Multi-Head Draft

Cai et al. (2024): Add multiple prediction heads to the TARGET model itself:

headk:ht→pk(xt+k∣x≤t),(6)\text{head}_k: h_t \to p_k(x_{t+k} | x_{\leq t}), \tag{6}

where hth_t is the target model's hidden state at position tt.

Each head predicts kk tokens ahead (independently, in parallel).

Tree-based verification: Generate a tree of candidate continuations from the multiple heads. Verify all candidates in one forward pass using tree attention.

Training (distillation from target):

  • Freeze the target model backbone.
  • Train only the Medusa heads to predict future tokens.
  • Each head is distilled from the target's ground-truth next-token distribution.

Advantage: No separate draft model; heads are tiny (1 linear layer each). 2-3x speedup with minimal quality impact.

Eagle: Autoregressive Draft Heads

Li et al. (2024): Use the target model's FEATURES (not just last hidden state) for autoregressive drafting:

h^t+1=DraftNet(ht,embed(xt)),(7)\hat{h}_{t+1} = \text{DraftNet}(h_t, \text{embed}(x_t)), \tag{7}

where DraftNet is a lightweight transformer (1-2 layers) that predicts the next hidden state.

Key insight: Predicting in FEATURE space (then projecting to vocab) is easier than predicting in TOKEN space directly.

Architecture:

  1. DraftNet takes current hidden state + current token embedding.
  2. Produces predicted next hidden state.
  3. Target model's LM head converts to distribution.
  4. Autoregressive: feed predicted hidden state back for multiple steps.

Advantage over Medusa: Autoregressive drafting captures token dependencies. Higher acceptance rate (0.85+ vs 0.75 for Medusa).

Online Speculative Distillation

Continuously improve draft model during serving:

  1. Serve requests using speculative decoding.
  2. Collect rejected tokens (where draft disagreed with target).
  3. Fine-tune draft on these "hard" examples.
  4. Deploy updated draft.

The rejection signal is free — it comes from the verification step that's already happening.

Result: Draft model gradually improves on the actual deployment distribution, increasing acceptance rate over time.

Common Pitfalls

Pitfall 1. Training draft model on different data than the target will see at inference. Distribution mismatch → low acceptance rate. Always distill on representative data.

Pitfall 2. Making the draft model too small. Below a certain capacity, no amount of distillation can achieve high acceptance. Rule of thumb: draft should be at least 1/10th the target's size.

Pitfall 3. Using greedy decoding for draft when target uses sampling. The acceptance criterion requires distribution matching; greedy draft with sampling target has poor acceptance on non-argmax tokens.

Summary

  • Speculative decoding uses small draft model to propose tokens, large model verifies.
  • Draft distillation: Minimize KL between draft and target distributions.
  • Medusa: Multiple parallel heads on the target model itself (no separate model).
  • Eagle: Feature-space autoregressive drafting (highest acceptance rate).
  • Online distillation: Continuously improve draft from rejection signal.
  • Acceptance rate of 0.8+ achievable with good distillation → 3-4x speedup.

Exercises

Exercise 1. For acceptance rate α=0.75\alpha=0.75 and draft length K=5K=5: compute expected accepted tokens and overall speedup if draft is 10x faster than target.

Exercise 2. Derive equation 4 (the relationship between TV distance, KL divergence, and acceptance rate).

Exercise 3. For Medusa with 3 heads: how many candidate sequences are generated with top-3 per head? How many tokens in the verification tree?

Exercise 4. Compare the memory overhead of Medusa (3 linear heads) vs a separate 1B draft model for a 70B target.

Exercise 5. Design an online speculative distillation loop: specify the data buffer, update frequency, and stopping criterion for a production chatbot.

Next

Online & Mutual Distillation