Self-Speculative Decoding

Medusa heads are lightweight, parallel prediction heads attached to the intermediate layers of a transformer model. They enable self-speculative decoding by proposing multiple future tokens in a single forward pass. Each head predicts a token at a different future position (e.g., head 1 predicts the next token, head 2 predicts the token two steps ahead).

• Architecture: Typically implemented as simple linear layers added after a chosen transformer block.
• Training: The heads are fine-tuned on the target model's outputs using a standard next-token prediction loss, ensuring they learn to mimic the model's own future distributions.
• Purpose: To generate a candidate sequence (or tree) of tokens for parallel verification without requiring a separate draft model's forward pass.

Tree attention is a modified attention mechanism that allows a transformer to process a tree of candidate token sequences in a single, batched forward pass. This is the core mechanism for efficient parallel verification in self-speculative decoding.

• Mechanism: Instead of a single causal attention mask for a linear sequence, it uses a mask that allows candidate tokens to attend to all preceding tokens in their branch of the tree.
• Efficiency: It enables the target model (the model itself) to score all speculative candidates simultaneously, calculating their true probabilities.
• Contrast with Standard Verification: In classic speculative decoding with a separate draft model, verification uses a simple batch. Tree attention handles a more complex, branching candidate structure generated by multiple Medusa heads.

The candidate sequence is the set of future tokens proposed by the Medusa heads. The model's own forward pass, using tree attention, then verifies these tokens in parallel.

• Generation: For a given context, all Medusa heads fire simultaneously, producing a set of candidate next tokens. This often forms a tree where the first head's output influences the candidates for later positions.
• Verification Forward Pass: The model runs one forward pass on this tree of candidates. It computes the true probability for each candidate token given the actual previous tokens (not the draft predictions).
• Acceptance Rule: A candidate token is accepted if its probability from the verified model is greater than the probability of the draft token proposed by the Medusa head. This is a form of token-level acceptance.

This component manages the outcome of verification, deciding which tokens to commit to the output sequence.

• Acceptance Rate: The percentage of proposed tokens accepted. High acceptance is critical for speedup, as each rejection incurs a verification cost without adding a new token to the output.
• Early Stopping: Verification proceeds sequentially through the candidate tree. If a token is rejected, the process for that branch halts, avoiding wasted computation on subsequent candidates in that branch.
• Rollback Mechanism: When a token is rejected, generation reverts to the last accepted position. The model then generates the correct token autoregressively (using its standard next-token prediction) before resuming speculative drafting. This ensures output distribution matches the original model.

The primary goal is to increase the speedup factor—the ratio of generation time for standard autoregressive decoding versus self-speculative decoding—while preserving the model's original output distribution.

• Theoretical Speedup: Governed by the formula: Speedup ≈ (1 + γ * acceptance_rate) / (1 + γ), where γ (speculative factor) is the number of candidate tokens proposed per step.
• Latency-Accuracy Tradeoff: Increasing γ (proposing more tokens) raises the potential speedup but may lower the acceptance rate, as predicting further into the future is less accurate. The optimal γ balances verification cost with acceptance gains.
• Throughput Improvement: In batched serving scenarios, the parallel verification of candidate trees across multiple requests can lead to significant throughput improvement by better utilizing GPU parallel compute units.

Self-speculative decoding differs fundamentally from the classic small-big model pair approach.

• Single Model vs. Two Models: Eliminates the need to host, load, and manage a separate draft model. This simplifies deployment and avoids draft model distillation overhead.
• Architectural Integration vs. External Drafting: Drafting is performed by integrated Medusa heads, not an external model. Verification uses tree attention instead of simple batch verification.
• Overhead: The computational overhead is the added cost of the Medusa heads and tree attention, which is typically far less than running a separate draft model's forward pass. This makes it particularly efficient for reducing inference latency in resource-constrained or high-throughput environments.