Medusa Heads

Medusa heads are lightweight decoding heads added on top of a large language model's last hidden layer, each trained to predict a token further ahead than the model's own output layer. In one forward pass the model produces its usual next-token prediction plus guesses for the next few positions, and a tree-attention step checks those guesses in the following pass and keeps the longest run the model agrees with. Medusa is a form of speculative decoding that needs no separate draft model, so there is no second network to train, host, or keep in sync. Tianle Cai, Tri Dao, and colleagues introduced it in Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads in January 2024.

How the Heads Predict Ahead#

A standard text generation step feeds the last hidden state at position t into the language-model head, which scores the token at t + 1. Medusa head k reads the same hidden state and scores the token at t + k + 1. Each head is a single feed-forward layer with a SiLU activation and a residual connection, followed by a projection to the vocabulary. The projection starts as a copy of the language-model head and the feed-forward layer starts at zero, so before training every head reproduces the base model's next-token prediction. The paper finds five heads sufficient. This PyTorch sketch builds three heads and compares their outputs with the base model's head on one hidden state:

import torch
import torch.nn.functional as F

hidden_size, vocab_size, num_heads = 256, 1000, 3
lm_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)  # the frozen model's output layer

# Each Medusa head: one residual SiLU layer (zero-initialized) and a copy of the LM head
w1 = [torch.nn.Linear(hidden_size, hidden_size) for _ in range(num_heads)]
w2 = [torch.nn.Linear(hidden_size, vocab_size, bias=False) for _ in range(num_heads)]
for layer, out in zip(w1, w2):
    torch.nn.init.zeros_(layer.weight)
    torch.nn.init.zeros_(layer.bias)
    out.weight.data.copy_(lm_head.weight.data)

# All heads read the same last hidden state; head k targets position t + k + 1
h = torch.randn(1, hidden_size)
head_logits = [out(h + F.silu(layer(h))) for layer, out in zip(w1, w2)]
print([torch.equal(logits, lm_head(h)) for logits in head_logits])

The output [True, True, True] shows that the zero-initialized residual layer makes each new head start as an exact copy of the base model's prediction. Training then shifts head k toward its own future position, while lm_head and the backbone stay unchanged.

Tree Attention and Acceptance#

The heads' top guesses combine into candidate continuations. With the top two tokens from head 1 and the top three from head 2, the Cartesian product gives six candidates that share prefixes and form a tree:

On the next forward pass, every candidate token is appended to the input at once, and an attention mask lets each token see only its own ancestors in the tree, so the model scores all branches in parallel. The longest branch the model accepts is kept, and its final hidden state feeds the heads for the next round. Because the heads all read the same hidden state, head 2 proposes the same tokens under every head-1 branch. The authors found that a sparse tree of 64 nodes, built from the most accurate combinations, gave a higher acceleration rate than a dense 256-node tree.

With greedy decoding, a candidate token is accepted when it matches the model's top choice. For sampling, Medusa offers typical acceptance, which keeps any candidate whose probability under the base model clears a threshold tied to the distribution's entropy, accepting more tokens at the cost of exact distribution matching.

Medusa-1 and Medusa-2 Training#

Medusa-1 trains only the new heads on a frozen backbone, so the model's outputs keep their quality, and the paper reports more than 2.2x faster generation. Training Vicuna-7B heads took about five hours on a single A100 GPU with 60,000 ShareGPT conversations. This is the same split that freeze=23 sets up when fine-tuning Ultralytics YOLO26: layers 0-22 stay fixed and only the layer-23 detection head learns from the new data.

Medusa-2 fine-tunes the heads and the backbone together, which makes the heads more accurate and raises the speedup to 2.3-3.6x, using a training recipe that protects the backbone's original abilities. When the original training data is unavailable, self-distillation trains the heads on text the model generates itself.

Medusa sits among several ways to draft tokens from the target model itself. EAGLE drafts from the model's internal features with a small autoregressive layer, and multi-token prediction, used in DeepSeek-V3, trains extra prediction modules during pretraining that serving engines later reuse as drafters. All three cut inference latency most at the small batch sizes of interactive chatbots, where each forward pass leaves the GPU mostly waiting on memory.