Medusa Heads

A Medusa head is a lightweight, fully-connected layer that predicts a token at a specific future position. Multiple heads are attached in parallel to the base model's final layer output, each trained to predict a token at a different offset (e.g., 1, 2, 3, ... steps ahead). This allows the model to propose a candidate sequence of N future tokens in a single, batched forward pass, eliminating the need for a separate draft model.

To verify the drafted tokens efficiently, Medusa employs a tree attention mechanism. The candidate tokens from all heads are arranged into a tree structure. The base (target) model then performs a single, modified forward pass that computes attention over this tree, scoring all candidate branches in parallel. This is the verification forward pass, and it determines which drafted tokens are accepted.

The verification pass compares the model's own probability distribution for each position against the Medusa head's draft token. A token is accepted if its probability meets a threshold. The system uses an early stopping strategy: verification halts at the first rejected token. A rollback mechanism then reverts to the last accepted position, and the next token is generated autoregressively by the base model before drafting resumes. This ensures output distribution remains identical to standard autoregressive generation.

Medusa heads are not random; they are fine-tuned on the base model's own outputs. Using a next-token prediction objective on diverse text, each head learns to approximate the base model's behavior at its specific lookahead position. This training is parameter-efficient, as only the new head layers are updated, keeping the original model weights frozen. This process is a form of self-distillation.

Medusa is integrated into the model serving stack. Key components include:

• Speculative KV Cache: Manages the key-value pairs for the tree of candidate tokens to avoid redundant computation.
• Modified Sampler: Orchestrates the drafting, verification, and rollback steps within the generation loop.
• Lightweight Overhead: The heads add minimal parameters (e.g., <1% of base model size) and computational cost to the forward pass, preserving the core speedup benefit.

The effectiveness of Medusa is measured by:

• Acceptance Rate: The percentage of drafted tokens accepted. Higher rates (e.g., >0.8) lead to greater speedups.
• Speedup Factor: Typically achieves 1.5x to 3x faster token generation compared to standard autoregressive decoding on the same hardware.
• Latency-Accuracy Tradeoff: Introduces no accuracy loss, as rejected tokens are corrected by the base model. The tradeoff is between speedup and the added memory/compute for the parallel heads and tree attention.

Medusa heads are lightweight neural network modules appended to the final transformer layer of a base model. They operate in parallel, each head predicting a token at a different future position (e.g., head 1 predicts the next token, head 2 predicts the token two steps ahead). This architecture enables a single forward pass to generate a candidate sequence of tokens, which is then verified by the model's original output head. The heads are typically trained via knowledge distillation from the base model's own autoregressive outputs to maximize prediction accuracy.

The Medusa decoding process follows a distinct, optimized loop:

• Drafting Phase: Given the current context, the base model and all Medusa heads execute in a single forward pass. The model's primary head and the Medusa heads collectively output a tree of candidate tokens.
• Verification Phase: The base model's original output head verifies this candidate tree via a tree attention mechanism, scoring all proposed tokens in parallel.
• Acceptance & Rollback: Tokens are accepted based on the verification scores. If a token is rejected, the process rolls back to the last accepted token and continues, either autoregressively or with another Medusa-augmented forward pass.

Medusa heads offer specific improvements by eliminating the need for a separate draft model:

• Reduced System Complexity: No need to load, manage, or synchronize two distinct models.
• Lower Memory Overhead: Medusa heads add minimal parameters compared to a full auxiliary model.
• Improved Draft Quality: Since heads are distilled from the target model itself, the acceptance rate for proposed tokens is typically higher than with a small, independently trained draft model.
• Eliminated Communication Latency: All computations occur within a single model graph, avoiding data transfer costs between separate models.

Medusa heads are not random; they are carefully trained to mimic the base model's future token distributions. The standard process involves:

• Dataset Creation: Generating sequences from the base model and extracting target tokens for each future position.
• Distillation Loss: Training the heads using a cross-entropy loss against the base model's soft labels (probability distributions) for tokens at offsets t+1, t+2, etc.
• Freezing Base Weights: Typically, only the Medusa head parameters are updated, while the base model's pretrained weights remain frozen to preserve its core capabilities and prevent catastrophic forgetting.

A core innovation enabling Medusa is tree attention, a modified attention mechanism that allows efficient parallel verification of a tree of candidate tokens. Unlike standard causal attention which processes a linear sequence, tree attention constructs a computational graph where:

• Each candidate branch shares the key-value (KV) cache for the common prefix.
• The attention operation computes scores for all leaf nodes (candidate tokens) in the tree simultaneously.
• This allows the model's verification forward pass to score a speculative factor of k tokens with near-constant overhead, rather than k sequential operations.