Research map / Large language models / Language, safety and interpretability
Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Medusa adds multiple lightweight decoding heads to a frozen or jointly trained LLM, verifies tree-structured candidate continuations in one step with tree attention, and uses typical acceptance to speed up decoding by about 2.2-2.8x.
Problem
Autoregressive LLM decoding is memory-bandwidth-bound and sequential, so each token requires a full forward pass; speculative decoding needs a separately trained draft model that is hard to obtain and serve.
Approach
Attach K extra single-layer decoding heads to the last hidden state to predict future tokens in parallel, build candidate continuations as a Cartesian-product tree verified with a tree attention mask in one forward pass, and accept them via a typical acceptance criterion based on the original model's entropy-scaled threshold instead of rejection sampling.
Key results
- Medusa-1 on Vicuna-7B achieves 2.18x speedup with no quality loss (MT-Bench 6.23 vs 6.17 baseline); Medusa-2 reaches 2.83x.
- Medusa-2 on Vicuna-13B achieves 2.83x speedup and on Vicuna-33B 2.35x with self-distillation.
- Speculative decoding with open draft models gives about 1.47-1.60x speedup versus Medusa-2's 2.35-2.83x in the Table comparison.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add tree attention to Medusa heads (ablation) | Medusa-1 on Vicuna-7B · MT-Bench writing and roleplay | speedup: $\sim$1.5x -> $\sim$1.9x | Helped |
| Use an optimized tree configuration (ablation) | Medusa on Vicuna-7B · MT-Bench writing and roleplay | speedup: $\sim$1.9x -> $\sim$2.2x | Helped |
| Train heads jointly with the backbone using Medusa-2 (ablation) | Medusa on Vicuna-7B · MT-Bench writing and roleplay | speedup: $\sim$2.2x -> $\sim$2.8x | Helped |
| Use two-stage Medusa-2 fine-tuning instead of direct fine-tuning (ablation) | Vicuna-7B · MT-Bench | quality: 5.925 -> 6.18 | Helped |
| Use two-stage Medusa-2 fine-tuning instead of Medusa-1 (ablation) | Vicuna-7B · MT-Bench | speedup: 2.18 -> 2.83 | Helped |
| Increase the typical-acceptance threshold ε (ablation) | Medusa-2 on Vicuna-7B · MT-Bench writing and roleplay | quality | Helped |
| Increase the typical-acceptance threshold ε (ablation) | Medusa-2 on Vicuna-7B · MT-Bench writing and roleplay | speedup | Hurt |
| Use typical acceptance rather than rejection sampling (ablation) | Medusa-2 · MT-Bench | speedup | Helped |
| Use optimized sparse tree with 64 nodes rather than dense tree with 256 nodes (ablation) | Medusa-2 on Vicuna-7B · MT-Bench writing and roleplay | acceleration rate | Helped |
| Increase the number of candidate tokens in the tree (ablation) | Medusa-2 on Vicuna-7B · MT-Bench writing and roleplay | speed | Hurt |
| Use self-distillation for Medusa-2 training (ablation) | Vicuna-33B · MT-Bench | speedup: 2.35 | Mixed |
Surprising findings
- Directly fine-tuning the backbone together with Medusa heads without the two-stage recipe degrades MT-Bench quality (5.925 vs 6.17 baseline).
- Rejection sampling yields diminishing efficiency as temperature increases, motivating typical acceptance.
Limitations
- Experiments focus on batch size 1; larger batch behavior is only discussed qualitatively.
- Self-distillation on Vicuna-33B and Zephyr-7B yields lower acceleration rates (2.35x and 2.66x) due to the quality-speed trade-off.
- Typical acceptance trades off quality and speed via thresholds, and decoding slows as tree size grows (compute-bound).
Experimental setup
Datasets: ShareGPT UltraChat Alpaca-eval
Benchmarks: MT-Bench AlpacaEval
Metrics: speedup acceleration rate overhead MT-Bench quality score tokens per second
Baselines: Speculative decoding (Llama-68M, Llama-160M, TinyLlama draft models) Huggingface default decoding
Builds on: speculative decoding Blockwise parallel decoding LoRA QLoRA truncation sampling (typical decoding)
Closest papers by meaning
- Optimized Multi-Token Joint Decoding With Auxiliary Model for LLM Inference (ICLR 2025 · similarity 0.91)
- Progressive Mixed-Precision Decoding for Efficient LLM Inference (ICLR 2025 · similarity 0.91)
- Break the Sequential Dependency of LLM Inference Using Lookahead Decoding (ICML 2024 · similarity 0.90)
- PARD: Accelerating LLM Inference with Low‑Cost PARallel Draft Model Adaptation (ICLR 2026 · similarity 0.90)
- Tandem Transformers for Inference Efficient LLMs (ICML 2024 · similarity 0.90)
- Accelerating LLM Inference with Lossless Speculative Decoding Algorithms for Heterogeneous Vocabularies (ICML 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.