Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

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.

한 줄 요약

Medusa는 병렬 토큰 예측을 통해 LLM 추론 속도를 2.3~3.6× 가속하는 프레임워크이다.

핵심 기여도

Medusa-1: 기존 모델을 동결하고 추가 디코딩 헤드만 미세조정하여 2.2× 속도 향상.
Medusa-2: 디코딩 헤드와 백본 모델을 함께 미세조정하여 2.3~3.6× 속도 향상.
Tree-based attention 메커니즘을 통해 병렬 토큰 후보 생성 및 검증.
Self-distillation과 typical acceptance scheme으로 데이터 부족 및 수용률 저하 문제 해결.

핵심 아이디어

Medusa는 기존의 순차적 auto-regressive decoding 방식의 한계를 극복하기 위해 병렬 토큰 예측을 도입한 새로운 추론 가속 기법이다. 기존의 speculative decoding이 별도의 draft 모델을 필요로 하는 반면, Medusa는 기존 백본 모델에 추가적인 디코딩 헤드를 붙여 복수의 토큰을 동시에 예측한다. 이는 tree-based attention 메커니즘을 통해 여러 후보 토큰을 생성하고 동시에 검증함으로써, 단계 수를 줄이고 추론 속도를 향상시킨다.

Medusa의 핵심 통찰은 기존 모델의 파라미터를 변경하지 않고도 병렬성을 도입할 수 있다는 점이다. Medusa-1은 백본 모델을 동결하고 디코딩 헤드만 미세조정하여 손실 없는 속도 향상을 달성한다. Medusa-2는 디코딩 헤드와 백본 모델을 함께 학습하여 예측 정확도와 속도 향상을 동시에 달성하지만, 별도의 학습 레시피가 필요하다.

또한, self-distillation을 통해 학습 데이터가 없는 상황에서도 Medusa 헤드를 학습할 수 있으며, typical acceptance scheme은 rejection sampling의 복잡성을 줄이면서도 생성 품질을 유지한다.

기술적 접근법

**Tree-based attention**: 병렬 토큰 후보 생성을 위한 attention 마스크 조정.
**Medusa-1**: 동결된 백본 모델 위에 디코딩 헤드만 미세조정 (parameter-efficient).
**Medusa-2**: 디코딩 헤드와 백본 모델을 함께 미세조정 (더 높은 정확도, 속도 향상).
**Self-distillation**: 학습 데이터가 없을 때, 기존 모델의 출력을 사용한 학습.
**Typical acceptance scheme**: temperature 기반 후보 선택으로 rejection sampling 대체.
**하이퍼파라미터**: QLoRA와 유사한 양자화 기법 적용 가능.

주요 결과

**Vicuna-7B**: Medusa-1 2.18×, Medusa-2 2.83× 속도 향상.
**Vicuna-13B**: Medusa-1 2.33×, Medusa-2 2.83× 속도 향상.
**Vicuna-33B & Zephyr-7B**: Medusa-2 2.3~3.6× 속도 향상.
**Task별 성능**: "Coding" 3.29×, "Extraction" 3.62× 최대 속도 향상.
**MT-Bench**: Medusa-2는 속도 향상과 생성 품질 유지의 균형 달성.

의의 및 한계

Medusa는 별도의 draft 모델 없이도 LLM 추론을 병렬화할 수 있는 간단하고 효율적인 방법을 제시한다. 기존 시스템에 쉽게 통합 가능하며, 다양한 모델 크기와 학습 방식에서 일관된 성능 향상을 보인다. 특히, self-distillation과 typical acceptance scheme은 데이터 부족 및 수용률 저하 문제를 해결하는 실용적 접근법이다.

그러나, Medusa-2는 별도의 학습 레시피가 필요하며, 학습 데이터와 모델 크기의 불일치가 성능에 영향을 줄 수 있다. 또한, 배치 추론 환경에서의 통합은 추가 엔지니어링이 필요하다는 한계가 있다.

실용적 활용

Medusa는 소프트웨어 개발, 정보 추출, 코드 생성 등 병렬 처리가 유리한 분야에서 유용하게 활용될 수 있다. 특히, 개인용 LLM 호스팅 환경에서 실시간 추론 성능 향상이 필요한 경우에 적합하며, 기존 추론 엔진(vLLM 등)에 통합될 경우 대규모 배치 처리에서도 활용 가능하다.