Medusa

핵심 아이디어

Speculative decoding은 효과적이지만 별도의 draft 모델을 유지해야 한다. Medusa는 target 모델 자체에 여러 개의 decoding head를 추가해서, 하나의 모델이 동시에 여러 토큰을 예측하게 만든다.

각 헤드는 “다음 토큰”, “그 다음 토큰”, “그 다음다음 토큰”을 독립적으로 예측한다. Tree attention으로 가능한 후보 경로들을 효율적으로 검증한다.

메커니즘

1. Base LLM에 여러 개의 Medusa head를 추가
2. 각 헤드는 미래의 특정 위치 토큰을 예측
3. 추론 시 여러 후보 경로를 tree 형태로 구성
4. Tree attention으로 한 번에 모든 경로를 평가
5. 검증된 경로를 선택

별도 모델 없이 하나의 forward pass로 여러 토큰을 생성할 수 있다.

장점

• 별도 draft 모델 불필요
• 메모리 오버헤드 작음 (head 파라미터는 base 대비 미미)
• 추가 학습 비용도 base 학습 대비 훨씬 작음

참고

• 논문: Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
• arXiv: 2401.10774