Speculative Decoding Deep Dive II: Medusa, MTP

Three ways to produce the draft: a separate small model (classic), extra heads on the target (Medusa), or heads baked in during pretraining (MTP).

Classic solution

A small autoregressive model proposes $ABCD$; the target verifies all four in one forward (~2–3x). The original recipe (Leviathan et al., 2023; Chen et al., 2023) needs no extra training: draft with a smaller same-family model so $q$ and $p$ share tokens and context.

• Same tokenizer — else $q(x)$ and $p(x)$ aren't over the same $x$.
• Same chat template — a format mismatch is a context mismatch: $q(\cdot \mid c)\ne p(\cdot \mid c)$.
• Same architecture / pretraining — so $q\approx p$; a random small model gets $\alpha \approx 0$.

If none exists, distill or fine-tune the draft (DistillSpec). Medusa and MTP drop the separate model.

The draft is autoregressive, so cost grows with speculated length $\gamma$:

$$T_{\mathrm{draft}}\propto \gamma$$

Draft 1: upside near $2\times$. Draft 32: slow, and the tail is usually wrong. A few tokens is the structural limit.

Medusa

Medusa (Cai et al., 2024) freezes the target and adds $K$ heads on $h_t$. The figure is the pipeline: draft → expand → verify. $\mathrm{Anc}(i)$ = ancestors of node $i$ on that tree. For shortlist sizes $s_1,s_2,s_3,\dots$:

• full paths — $s_1\cdot s_2\cdot s_3\cdot \dots$ (what you would verify as separate sequences)
• tree nodes — $s_1+s_1{s}_2+s_1{s}_2{s}_3+\cdots$ (each depth adds one product term)

Heads never see the token just sampled, so accuracy falls with offset (suffix decay). Hydra (Ankner et al., 2024) feeds $x_{t+1}$ to later heads — better drafts, less parallelism.

Training. Supervise head $k$ with the token $k$ steps past Head 0's target, $x_{t+1+k}$:

$$p_t^{(k)}=\mathrm{softmax}({\mathrm{Head}}_k(h_t)),L=\sum _t\sum _{k=0}^K{\lambda }_k\mathrm{CE}(p_t^{(k)},x_{t+1+k}),{\lambda }_k={\lambda }^k,\lambda \approx 0.8$$

• Medusa-1 — freeze the model, train only heads $k\ge 1$. ~2.2x.
• Medusa-2 — also fine-tune the trunk so $h_t$ adapts. ~2.3–3.6x, accepted length ~3.0–3.5.

MTP

Multi-token prediction (Gloeckle et al., 2024) uses the same architecture: shared trunk, $n$ parallel heads on $h_t$, one shared unembedding $f_u$.

$$P_{\theta }(x_{t+i}\mid x_{\le t})=\mathrm{softmax}(f_u(f_{h_i}(h_t))),i=1,\dots ,n$$

Two differences: the heads $f_{h_i}$ are transformer layers (not a single MLP), and they are a pretraining objective — every position is trained on all $n$ next tokens:

$$L_n=-\sum _t\sum _{i=1}^n\log P_{\theta }(x_{t+i}\mid x_{\le t})$$

Aim is sample efficiency and better representations; faster decoding (~3x, extra heads + Medusa-style tree attention) is a by-product.

	Medusa	MTP
Heads	one MLP, added after training	transformer layers, present in pretraining
Trunk	frozen (Medusa-1)	trained under $L_n$
Purpose	speed	representations; speed is a bonus

EAGLE keeps a sequential draft but conditions each step on the token actually sampled — the branch information Medusa-style heads never get.

References

• Cai et al., 2024. Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arXiv 2401.10774
• Gloeckle et al., 2024. Better & Faster Large Language Models via Multi-token Prediction. arXiv 2404.19737
• Ankner et al., 2024. Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding. arXiv 2402.05109