Medusa: Simple LLM Inference Acceleration Framework with Multiple
  Decoding Heads
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.
65add289ac588f2a1c2b22c2
Tianle Cai
admin_assigned
2024-01-22T11:22:03.769Z
64459357c525660aa20be337
https://cdn-avatars.huggingface.co/v1/production/uploads/noauth/1fryBXRuSvAgMIDkpWBpI.jpeg
Tianle Cai
user
tianlecai
65add289ac588f2a1c2b22c3
Yuhong Li
65add289ac588f2a1c2b22c4
Zhengyang Geng
admin_assigned
2024-01-22T11:22:37.019Z
626039656089b62d9f332854
https://cdn-avatars.huggingface.co/v1/production/uploads/1650473291256-noauth.jpeg
Zhengyang Geng
user
Gsunshine
65add289ac588f2a1c2b22c5
Hongwu Peng
admin_assigned
2024-01-22T11:22:55.102Z
655a9f12bf604a7732716115
/avatars/2f5094af95bf3eb236e39dacb909a5d7.svg
Hongwu Peng
user
harveyp123
65add289ac588f2a1c2b22c6
Jason D. Lee
65add289ac588f2a1c2b22c7
Deming Chen
65add289ac588f2a1c2b22c8
Tri Dao
admin_assigned
2024-01-22T11:23:52.871Z
64b8a6b5cf14c2fabe98159b
/avatars/dbc009451865435bf290791beadc4723.svg
Tri Dao
user
tridao
Medusa: Simple LLM Inference Acceleration Framework with Multiple
  Decoding Heads
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 enhances Large Language Model inference by adding parallel decoding heads and using tree-based attention to predict multiple tokens simultaneously, achieving significant speedup with minimal latency.
Large Language Models LLMs auto-regressive decoding parallelism memory bandwidth speculative decoding decoding heads tree-based attention inference acceleration fine-tuning self-distillation acceptance scheme generation quality speedup