Draft-Model Distillation for Speculative Decoding

Speculative decoding runs a small draft model ahead of a large target model and verifies the draft's tokens in parallel, with a rejection rule that leaves the target's output distribution exactly unchanged (Leviathan et al., ICML 2023, Fast Inference from Transformers via Speculative Decoding, arXiv:2211.17192). Because the output is unchanged, the draft model's quality is not a quality lever at all. It is a latency lever, and the quantity it controls is the acceptance rate.

That makes draft selection a strange optimisation. The best small model you have is often not the best draft, because what matters is not how good its predictions are but how closely they match the target's, token for token, on the text the target is actually going to produce.

The objective is agreement, and distillation optimises it directly

If each drafted token is accepted independently with probability \(\alpha\) and the draft proposes \(\gamma\) tokens per step, the expected number of tokens accepted per verification is

which is superlinear in \(\alpha\) near 1: pushing acceptance from 0.7 to 0.8 buys more than pushing it from 0.5 to 0.6. Knowledge distillation is the natural tool because the loss it minimises, a divergence between draft and target distributions, is a direct surrogate for \(\alpha\). DistillSpec makes two design choices on top of that and reports 10 to 45% speedups over standard speculative decoding across benchmarks: generate the distillation data on policy from the draft itself, and choose the divergence to match the decoding strategy the system will serve (Zhou et al., ICLR 2024, DistillSpec: Improving Speculative Decoding via Knowledge Distillation, arXiv:2310.08461).

The on-policy choice is the one worth understanding. Acceptance is evaluated on sequences the draft proposes, so the distribution that matters is the draft's own, not the target's. Training on target-generated text optimises agreement on trajectories the draft will never visit, which is exactly the argument in on-policy distillation and exposure bias applied to a latency objective rather than a quality one.

Removing the second model entirely

If the draft must agree with the target, the cheapest way to agree is to be the target. Medusa attaches extra decoding heads to the frozen backbone and verifies several candidate continuations with tree attention, reporting more than 2.2x speedup with the backbone frozen and 2.3 to 2.8x when trained jointly (Cai et al., 2024, Medusa, arXiv:2401.10774). EAGLE instead drafts autoregressively in the target's second-to-top-layer feature space, which it argues is more predictable than token space, and reports draft accuracy around 0.8 (Li et al., 2024, EAGLE, arXiv:2401.15077).

Both are distillation in the sense that matters here: a small head is trained to imitate the target's own next-token behaviour. They trade the separate draft model's flexibility for guaranteed architectural agreement, and they relate closely to multi-token prediction as a training objective.

When it breaks

Acceptance is not the objective; wall-clock is. A larger draft accepts more and costs more per drafted token. The product of acceptance and draft cost is what to optimise, and teams that report only \(\alpha\) improvements routinely ship a regression.

A draft is tuned to a traffic mix. On-policy distillation on one prompt distribution produces agreement on that distribution. Code traffic and chat traffic have different acceptance rates against the same draft, so one measured speedup is not a system-wide number.

Serving temperature changes everything. The divergence that maximises acceptance under greedy decoding is not the one that maximises it under temperature sampling, which is why DistillSpec ties the divergence to the decoding strategy. A draft distilled for greedy serving degrades when a product turns the temperature up.

Batch size erodes the win. Speculative decoding pays off when decoding is memory-bandwidth-bound and verification is nearly free. At large batch sizes the target is compute-bound, the parallel verification is no longer free, and the speedup shrinks regardless of how good the draft is.

Every change to the target invalidates the draft. The draft is distilled against a specific target checkpoint. Update the target, and acceptance drops silently; the output stays correct, which is exactly why the regression is easy to miss. Pin the pair and re-measure acceptance on every target change.