Medusa: Faster LLM text generation

Introduction

Autoregressive language models generate one next token at a time. Medusa adds several lightweight prediction heads to a base model so that one forward pass can propose a short continuation. The base model then verifies the proposals with tree attention. The paper studies this approach as an inference tradeoff: speed depends on how many proposed tokens are accepted and on the cost of the extra heads.

The base language-model head and Medusa heads read the same hidden state. Their candidate tree is checked by the base model, and the accepted prefix becomes the next context.

Open figure at full size

The Medusa paper describes this decoding pipeline.

Training

A Medusa head predicts a token at a fixed offset from the current position. The first head predicts the token after the base model’s next-token position, the second predicts the following token, and so on. For a prefix such as “2025 would be”, the base model predicts the next token while the first head predicts the token after that position. The heads use the hidden state already computed by the base model.

The paper describes two training variants:

Medusa-1 freezes the base model and trains only the new heads.
Medusa-2 trains the base model and heads together with a joint objective. Its training recipe includes a self-distillation signal so the base model can continue to provide useful next-token supervision.

The loss gives lower weight to farther-ahead predictions. That weighting is part of the paper’s setup, rather than a guarantee that every head will be equally accurate.

Inference

Proposals and a tree

Suppose the current prefix is “The future of artificial intelligence will”. The base model and the heads produce candidates at several offsets. Keeping the top two candidates at each step creates branches such as:

“transform” → “the” → “world”
“change” → “our” → “industry”

A real decode can have a different number of branches and a different accepted prefix. The four-token example below is idealized and assumes that all proposed tokens are accepted:

The base model proposes “transform”.
The first head proposes “the”.
The second head proposes “world”.
The third head proposes “in”.

If every candidate passes verification, the next context includes all four tokens. If verification rejects a candidate, decoding keeps only the accepted prefix and continues from the rejection point.

Medusa organizes candidate continuations as branches. Tree attention lets each candidate read the prompt and its own branch predecessors, so branches can be verified in one batched pass.

Open figure at full size

The paper describes the candidate tree and its attention mask in the tree-attention section.

Tree attention

The target model evaluates the candidate tree with one batched forward pass. The attention mask allows a candidate to see the original context and its same-branch predecessors. For example, a candidate after “transform” can attend to the prefix and “transform”; a candidate after “transform the” can also attend to “the”. It cannot use tokens from a different branch as context.

This shares computation between branches while keeping each verification path causally valid. The tree reduces repeated work, but it also adds bookkeeping and extra attention positions.

Acceptance rules

Verification can use different acceptance rules. Typical acceptance accepts a proposal when its probability under the verifying model is high enough under the rule’s threshold. It is an efficiency and quality tradeoff, so the accepted sequence can depend on the configuration.

An exact speculative-decoding rule instead uses rejection sampling. At the first rejected token, it samples a replacement from the normalized positive residual max(0, p − q), where p is the target distribution and q is the proposal distribution. With the required conditions and implementation, this preserves the target model’s distribution. The usual Medusa acceptance configurations are designed for useful speed and quality, but should not be described as universally distribution-preserving.

Results

The paper’s shared setup trains five Medusa heads and evaluates a configured candidate tree at batch size 1. The speedup entries below use the paper’s tree-based acceptance procedure for each comparison; the candidate budget and acceptance details are experimental settings, not universal defaults. The paper separately studies typical-acceptance thresholds at temperature 0.7 on selected MT-Bench categories.

Model and variant	Workload and setting	Acceptance context	Reported speedup
Vicuna-7B, Medusa-1	MT-Bench, batch size 1, sequence length 4096	Five-head Medusa tree configuration	2.18×
Vicuna-7B, Medusa-2	MT-Bench, batch size 1, sequence length 4096	Five-head Medusa tree configuration	2.83×
Vicuna-13B, Medusa-1	MT-Bench, batch size 1, sequence length 4096	Five-head Medusa tree configuration	2.33×
Vicuna-13B, Medusa-2	MT-Bench, batch size 1, sequence length 4096	Five-head Medusa tree configuration	2.83×
Vicuna-7B, Medusa-2	MT-Bench coding category	Five-head Medusa tree configuration	3.29×
Vicuna-7B, Medusa-2	MT-Bench extraction category	Five-head Medusa tree configuration	3.62×

For Vicuna-7B on MT-Bench, the paper reports GPT-4 judge scores of 6.17 for the base model, 6.23 for Medusa-1, and 6.18 for Medusa-2. Those scores are close within that evaluation, so the speedup should be read alongside the benchmark quality result.

The paper’s category panel reports these speedups for Vicuna-7B with Medusa-2 on eight MT-Bench categories. The common experimental setting uses batch size 1; the paper does not specify a separate acceptance threshold for each category bar, so these values retain the paper’s configured decoding setup.

MT-Bench category	Model and setting	Reported speedup
Humanities	Vicuna-7B, Medusa-2, batch size 1	2.58×
Reasoning	Vicuna-7B, Medusa-2, batch size 1	2.58×
Roleplay	Vicuna-7B, Medusa-2, batch size 1	2.70×
Writing	Vicuna-7B, Medusa-2, batch size 1	2.72×
STEM	Vicuna-7B, Medusa-2, batch size 1	2.77×
Math	Vicuna-7B, Medusa-2, batch size 1	3.01×
Coding	Vicuna-7B, Medusa-2, batch size 1	3.29×
Extraction	Vicuna-7B, Medusa-2, batch size 1	3.62×

The values are from Figure 3 of the Medusa paper; they are measurements from that model, workload, and decoding setup, not universal speedups.

The extra heads avoid deploying a separate draft model, but they still consume parameters and computation. More candidates can increase verification overhead, and low agreement between proposals and the base model reduces the number of accepted tokens. Medusa is therefore best understood as a decoding design whose benefit depends on the model, workload, hardware, and acceptance behavior.