# Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

**Authors:** Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, Tri Dao  
**arXiv:** [2401.10774](https://arxiv.org/abs/2401.10774)  
**Date:** 2024-01-19  
**Summary page:** https://www.summarizepaper.com/en/paper/2401.10774v1/

## Key points

- Large Language Models (LLMs) face limitations in inference due to lack of parallelism in auto-regressive decoding
- Medusa introduces multiple decoding heads to predict multiple tokens simultaneously, improving efficiency
- Tree-based attention mechanism used to generate and verify candidate continuations concurrently, reducing decoding steps
- Two levels of fine-tuning procedures: Medusa-1 for lossless acceleration, Medusa-2 for enhanced accuracy and speedup
- Extensions include self-distillation and acceptance scheme to improve utility
- Experimental results show Medusa can achieve 2.2x speedup with Medusa-2 reaching 2.3-3.6x speedup while maintaining generation quality
- Code implementation available at https://github.com/FasterDecoding/Medusa

## Summary

The paper "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads" by Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao addresses the limitations of Large Language Models (LLMs) in the inference process due to the lack of parallelism in auto-regressive decoding. This limitation results in operations being constrained by accelerator memory bandwidth. To tackle this problem, the authors introduce Medusa, an efficient method that enhances LLM inference by incorporating additional decoding heads to predict multiple subsequent tokens simultaneously. By utilizing a tree-based attention mechanism, Medusa generates multiple candidate continuations and verifies them concurrently in each decoding step. This approach minimizes overhead in terms of single-step latency while significantly reducing the number of decoding steps required. The paper presents two levels of fine-tuning procedures for Medusa tailored to different use cases: Medusa-1 involves direct fine-tuning on top of a frozen backbone LLM for lossless inference acceleration, while Medusa-2 entails joint fine-tuning with the backbone LLM to enhance prediction accuracy and achieve higher speedup. However, Medusa-2 requires a specialized training recipe to preserve the capabilities of the backbone model. Furthermore, several extensions are proposed to enhance Medusa's utility including self-distillation for scenarios where training data is unavailable and a typical acceptance scheme to boost acceptance rates while maintaining generation quality. The authors evaluate Medusa across models of varying sizes and training procedures. Experimental results demonstrate that Medusa-1 can achieve over 2.2x speedup without compromising generation quality, with Medusa-2 further improving speedup performance to 2.3-3.6x. Overall  offers a promising solution for accelerating LLM inference by introducing parallel processing through multiple decoding heads while maintaining high-quality generation outcomes. The code implementation for this framework is available at https://github.com/FasterDecoding/Medusa.