Medusa: Single-Model Speculative Decoding

Autoregressive language model generation is constrained by memory bandwidth. Generating text token by token requires moving multi-billion parameter matrices from High-Bandwidth Memory (HBM) to local GPU registers for every single generated token. Because each forward pass computes logits for a single query token, the arithmetic intensity remains low, leaving high-performance Tensor Cores underutilized while waiting for parameter transfers.

Standard speculative decoding addresses this bottleneck by deploying a smaller secondary draft model to propose a sequence of tokens. The primary model then verifies these draft tokens in a single parallel forward pass. While effective at increasing token throughput, hosting two separate models introduces deployment complexity. 

Medusa eliminates the secondary draft model requirement by attaching lightweight auxiliary decoding heads directly to the base model's final hidden state representation.

The Architecture of Auxiliary Medusa Heads

Rather than routing draft generation to an external neural network, Medusa modifies the output architecture of the base model. The primary model consists of a transformer backbone that produces a final hidden state vector at each decoding step. Standard generation passes this vector through a single language modeling head to produce the probability distribution over the vocabulary for the immediate next token.

Medusa attaches multiple additional linear or single-layer Feed-Forward Network (FFN) heads directly to that same hidden state representation. Each auxiliary head trains specifically to predict a token at a distinct positional offset further down the sequence. At any given decoding step, the base model generates logits for the next token via its standard language modeling head while the first auxiliary head predicts the second token, the second auxiliary head predicts the third token, etcetera.

Because all auxiliary heads read from the same underlying hidden representation, token proposals generate concurrently in a single forward pass without executing sequential autoregressive steps for drafting.

Candidate Token Trees and Tree-Attention

In speculative decoding, draft proposals form a linear chain of tokens. If the primary model rejects the second token in the sequence, all subsequent candidate tokens are discarded regardless of their individual quality.

Medusa increases token acceptance rates by constructing a candidate token tree from top probability predictions across its auxiliary heads. For each head, the system samples the highest-probability tokens. Combining these candidate predictions creates a branching tree of potential sequences.

The system constructs a dynamic tree-attention mask to evaluate all branches of the candidate tree simultaneously inside the base model. Tree-attention allows tokens in different branches of the tree to attend to their direct ancestral path while remaining isolated from competing sibling branches.

By passing the candidate token tree and the tree-attention mask into the base model, Tensor Cores evaluate dozens of candidate paths concurrently in one forward pass. The system processes the accepted path matching the base model's probability distribution, recovering multiple tokens per forward pass while preserving decoding accuracy.

Parameter-Efficient Training Strategies

Training auxiliary Medusa heads requires minimal compute compared to pre-training or full fine-tuning. The base model parameters remain frozen during training.

Engineering teams train the auxiliary parameters using a multi-head cross-entropy loss function. A decay hyperparameter penalizes error less heavily further down the sequence, prioritizing accuracy for tokens closer to the immediate prediction point.

Because the base model weights are locked, training strictly updates the lightweight projection layers. This parameter-efficient setup allows infrastructure developers to fine-tune Medusa heads on domain datasets in a matter of hours on standard GPUs.

Systems Efficiency

Attaching auxiliary heads directly to the base transformer backbone provides several advantages for enterprise serving platforms:

By combining lightweight auxiliary prediction heads with tree-attention verification, Medusa provides a parameter-efficient blueprint for reducing inference latency across open-source transformers.