CoChat
Sign Up
PreprintCited by 27Natural Language Processing TechniquesTopic ModelingAdvanced Neural Network Applications

The DOI matches its DataCite record; 2 of the 2 references it cites were found in a registry or an index.

Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

Open access
Cai, Tianle, Li, Yuhong, Geng, Zhengyang, Peng, Hongwu, Lee, Jason D., Chen, Deming, Dao, Tri·2024·arXiv
DOI 10.48550/arxiv.2401.10774arXiv 2401.10774
Open this paper in CoChat — free accountOpen PDF

Abstract

Large Language Models (LLMs) employ auto-regressive decoding that requires sequential computation, with each step reliant on the previous one's output. This creates a bottleneck as each step necessitates moving the full model parameters from High-Bandwidth Memory (HBM) to the accelerator's cache. While methods such as speculative decoding have been suggested to address this issue, their implementation is impeded by the challenges associated with acquiring and maintaining a separate draft model. In this paper, we present Medusa, an efficient method that augments LLM inference by adding extra decoding heads to predict multiple subsequent tokens in parallel. Using a tree-based attention mechanism, Medusa constructs multiple candidate continuations and verifies them simultaneously in each decoding step. By leveraging parallel processing, Medusa substantially reduces the number of decoding steps required. We present two levels of fine-tuning procedures for Medusa to meet the needs of different use cases: Medusa-1: Medusa is directly fine-tuned on top of a frozen backbone LLM, enabling lossless inference acceleration. Medusa-2: Medusa is fine-tuned together with the backbone LLM, enabling better prediction accuracy of Medusa heads and higher speedup but needing a special training recipe that preserves the backbone model's capabilities. Moreover, we propose several extensions that improve or expand the utility of Medusa, including a self-distillation to handle situations where no training data is available and a typical acceptance scheme to boost the acceptance rate while maintaining generation quality. We evaluate Medusa on models of various sizes and training procedures. Our experiments demonstrate that Medusa-1 can achieve over 2.2x speedup without compromising generation quality, while Medusa-2 further improves the speedup to 2.3-3.6x.

Overview

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2401.10774

Cites

2 of 2 references verified0 citation errors · 0 not checkable

Statistics

Verification

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

  • Existing model-based approaches need a fine-tuned lightweight draft model that aligns the output of the original model.

    “The modelbased approaches (Li et al. 2024b,a; Cai et al. 2024) need a fine-tuned lightweight draft model that aligns the output of original.”

    From DIAA: A Decoding-Efficient Inference Acceleration Approach for On-Device Large Language Models (2026) · cites this paper as Cai et al. 2024

    The Introduction says MEDUSA has “no requirement for a draft model”: it fine-tunes decoding heads on the original model instead. The paper describes small, separately trained draft models as a feature of other speculative-decoding approaches, so the claimed requirement does not hold for the cited method. (LLM, conf=0.90)

    Checked by a reasoning model — may be wrong.

    Experimental — accuracy not yet measured

  • Static verification often leads to suboptimal performance when adapting to dynamic contexts.

    “Moreover, static verification often leads to suboptimal performance when adapting to dynamic contexts.”

    From DIAA: A Decoding-Efficient Inference Acceleration Approach for On-Device Large Language Models (2026) · cites this paper as Cai et al. 2024

    Section 2.3.1 (page unavailable) says rejection sampling becomes less efficient as sampling temperature rises and proposes an acceptance threshold that responds to the prediction distribution. This supports a narrower point about verification under changing sampling conditions, but the paper does not establish that static verification often performs suboptimally across dynamic contexts generally. (LLM, conf=0.78)

    Experimental — accuracy not yet measured

  • Most existing works do not focus on edge environments, where draft models add memory overhead or retraining costs.

    “However, most existing works do not focus on edge environments, where draft models add memory overhead or retraining costs.”

    From DIAA: A Decoding-Efficient Inference Acceleration Approach for On-Device Large Language Models (2026) · cites this paper as Cai et al. 2024

    Section 2.1.1 says existing speculative-decoding approaches often separately pre-train a smaller draft model, requiring substantial additional computation. The Introduction and Appendix A.1 also describe the complexity of managing a separate model. The paper does not establish that most existing works neglect edge environments or specifically quantify draft-model memory overhead there; its discussion of deployment complexity concerns distributed systems. (LLM, conf=0.85)

    Experimental — accuracy not yet measured

  • Prior works that utilized the attention matrix to compress multiple trajectories adjusted the position ids so each trajectory follows consecutive increasing position ids.

    “For these works, on top of the attention matrix, they adjusted the position ids as well, so each trajectory follows a consecutive increasing position ids.”

    From MEM1: Learning to Synergize Memory and Reasoning for Efficient Long-Horizon Agents (2025) · cites this paper as [8]

  • Prior works that utilized the attention matrix to compress multiple trajectories mainly targeted tree-exploration.

    “Specifically, prior works [28, 10, 8] that utilized the attention matrix to compress multiple trajectories mainly targeted tree-exploration, i.e., generating multiple sequences with the same prefix.”

    From MEM1: Learning to Synergize Memory and Reasoning for Efficient Long-Horizon Agents (2025) · cites this paper as [8]

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Citations

Cites

Flagged first, then by citations

Research works printed but not linked: 36

  1. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv:2305.13245, 2023. · arXiv
  2. Basu, S., Ramachandran, G. S., Keskar, N. S., and Varshney, L. R. {MIROSTAT}: A {neural} {text} {decoding} {algorithm} {that} {directly} {controls} {perplexity}. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum? id=W1G1JZEIy5_.
  3. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
  4. Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. Accelerating large language model decoding with speculative sampling. February 2023. doi: 10.48550/ARXIV.2302.01318. · DOI · Paper page
  5. Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022. · arXiv
  6. Dettmers, T., Lewis, M., Shleifer, S., and Zettlemoyer, L. 8- bit optimizers via block-wise quantization. International Conference on Learning Representations, 2021.
  7. Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339, 2022. · arXiv
  8. Dettmers, T., Pagnoni, A., Holtzman, A., and Zettlemoyer, L. Qlora: Efficient finetuning of quantized llms. arXiv preprint arXiv:2305.14314, 2023. · arXiv
  9. Elfwing, S., Uchibe, E., and Doya, K. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 2017. doi: 10.1016/j.neunet.2017.12.012. · DOI · Paper page
  10. Hewitt, J., Manning, C. D., and Liang, P. Truncation sam- pling as language model desmoothing. October 2022. doi: 10.48550/ARXIV.2210.15191. · DOI
  11. Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022. · arXiv
  12. Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum? id=rygGQyrFvH.
  13. Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., and Chen, W. Lora: Low-rank adaptation of large language models. ICLR, 2021.
  14. Kim, S., Hooper, C., Gholami, A., Dong, Z., Li, X., Shen, S., Mahoney, M. W., and Keutzer, K. Squeezellm: Dense-and-sparse quantization. arXiv preprint arXiv:2306.07629, 2023. · arXiv
  15. Kim, Y. and Rush, A. M. Sequence-level knowledge distil- lation. EMNLP, 2016.
  16. Kumar, A., Raghunathan, A., Jones, R., Ma, T., and Liang, P. Fine-tuning can distort pretrained features and under- perform out-of-distribution. International Conference on Learning Representations, 2022.
  17. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
  18. Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. November 2022. doi: 10.48550/ARXIV.2211.17192. · DOI · Paper page
  19. Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023. · arXiv
  20. Meister, C., Wiher, G., Pimentel, T., and Cotterell, R. On the probability-quality paradox in language generation. March 2022. doi: 10.48550/ARXIV.2203.17217. · DOI
  21. Meister, C., Pimentel, T., Wiher, G., and Cotterell, R. Lo- cally typical sampling. Transactions of the Association for Computational Linguistics, 11:102-121, 2023.
  22. Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Wong, R. Y. Y., Chen, Z., Arfeen, D., Abhyankar, R., and Jia, Z. Specinfer: Accelerating generative llm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023. · arXiv
  23. Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to fol- low instructions with human feedback. arXiv preprint arXiv:2203.02155, 2022. · arXiv
  24. Pillutla, K., Swayamdipta, S., Zellers, R., Thickstun, J., Welleck, S., Choi, Y., and Harchaoui, Z. MAUVE: Mea- suring the gap between neural text and human text using divergence frontiers. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL https: //openreview.net/forum?id=Tqx7nJp7PR.
  25. Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Brad- bury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference. November 2022. doi: 10.48550/ARXIV.2211.05102. · DOI · Paper page
  26. Shazeer, N. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019. · arXiv
  27. Spector, B. and Re, C. Accelerating llm inference with staged speculative decoding. arXiv preprint arXiv:2308.04623, 2023. · arXiv
  28. Stern, M., Shazeer, N. M., and Uszkoreit, J. Blockwise parallel decoding for deep autoregressive models. Neural Information Processing Systems, 2018.
  29. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288, 2023. · arXiv
  30. Xia, H., Ge, T., Chen, S.-Q., Wei, F., and Sui, Z. Speculative decoding: Lossless speedup of autoregressive translation, 2023. URL https://openreview.net/forum? id=H-VlwsYvVi.
  31. Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38099.
  32. PMLR, 2023a.
  33. Xiao, Y., Wu, L., Guo, J., Li, J., Zhang, M., Qin, T., and Liu, T.-y. A survey on non-autoregressive generation for neu- ral machine translation and beyond. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023b.
  34. Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877-28888, 2021.
  35. Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022. · arXiv
  36. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H 2 o: Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306.14048, 2023. · arXiv

Cited by

Sorted by citations

Metadata

Authors
Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., & Dao, T.
Affiliations
  • Princeton University (US)
  • University of Illinois Urbana-Champaign (US)
  • Carnegie Mellon University (US)
  • University of Connecticut (US)
ORCID
Venue
arXiv
Publisher
arXiv
Published
2024
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2401.10774
arXiv
2401.10774
OpenAlex
W4391124766
Cited by
27
References
2
Citations by year
2026: 3 · 2025: 17 · 2024: 7
Type
posted content
Language
en
Open access
Licence
Locations
  • arXiv (Cornell University) (submittedVersion, cc-by)
  • arXiv (Cornell University) (cc-by)
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Natural Language Processing Techniques
  • Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
  • Physical Sciences › Computer Science › Computer Vision and Pattern Recognition › Advanced Neural Network Applications
Keywords
speculative decoding, fine-tuning, self-distillation, Medusa

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review