The inference process in Large Language Models (LLMs) is often limited due to
the absence of parallelism in the auto-regressive decoding process, resulting
in most operations being restricted by the memory bandwidth of accelerators.
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 introduces only minimal overhead in
terms of single-step latency while substantially reducing 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 推理的效率，并减少所需的解码步骤。通过两个级别的精细调整程序，Medusa 可以满足不同用例的需求。此外，我们还提出了几种扩展方法，包括自我蒸馏和典型接受方案，以提高 Medusa 的效用。我们的实验证明，Medusa-1 可以在不损害生成质量的情况下实现超过 2.2 倍的加速，而 Medusa-2 进一步提高了加速度，可达到 2.3-3.6 倍。