Medusa-Style Multi-Step Decoding Heads

• Medusa-style multi-step heads are lightweight decoding modules attached to large language models that predict several future tokens from a single autoregressive state.
• They leverage a tree-based attention mechanism to concurrently verify candidate continuations, thereby increasing arithmetic intensity and reducing decoding steps.
• Variants like Hydra and Amphista extend the design by introducing sequential dependence and bi-directional interactions to further enhance decoding throughput and efficiency.

Medusa-style multi-step heads are lightweight decoding heads attached to a backbone LLM so that several future tokens can be proposed from a single autoregressive state and then verified in parallel. In the canonical Medusa formulation, the original LM head predicts the next token while auxiliary heads predict later positions, and a tree-based attention mechanism constructs multiple candidate continuations and verifies them simultaneously using the backbone model’s own layers (Cai et al., 2024). The design targets the memory bandwidth bottleneck of large-model decoding: instead of paying the full parameter-transfer cost for one emitted token per step, it increases arithmetic intensity and reduces the number of decoding steps needed for a fixed output length.

1. Systems Motivation and Problem Setting

LLMs employ auto-regressive decoding in which each generation step depends on the previous output. For large transformer LLMs, this creates a sequential, memory-bound regime in which each step requires moving the full model parameters from High-Bandwidth Memory to on-chip cache, while arithmetic throughput is underutilized because only one token is produced per full parameter transfer (Cai et al., 2024). In this setting, reducing wall-clock latency depends on increasing arithmetic intensity and reducing the number of decoding steps.

Earlier acceleration proposals, including speculative decoding, address this bottleneck by introducing a draft mechanism that proposes several future tokens and then verifies them with the base model. Medusa-style heads preserve that general objective but remove the need for a separate draft model. Instead, they add multiple lightweight heads on top of the backbone model’s final hidden state, so the same model instance both drafts and verifies candidate continuations (Cai et al., 2024). This avoids training and serving a second model, avoids duplicate KV-cache management, and simplifies integration relative to two-model speculative systems.

The broader design space that emerged after Medusa treats multi-step drafting as a family of inference-time accelerators. Later work frames Medusa as a baseline in which draft heads are attached directly to the base model’s hidden states, then modifies either the head dependency structure, the interaction among head positions, or the tree used for verification (Ankner et al., 2024).

2. Canonical Architecture

In Medusa, the drafting mechanism attaches to the backbone LLM’s last hidden state $h_t$ at the current position $t$, after the transformer layers compute the representation for the just-generated token. The original LM head produces the distribution for position $t+1$, denoted $p_t^{(0)}$, while auxiliary heads $k=1,\dots,K$ predict positions $t+k+1$ in parallel (Cai et al., 2024). A useful formal characterization given in later work is that standard Medusa-style heads approximate the multi-token draft distribution by an independent factorization,

$$p(y_{t+1:t+K} \mid h_t) \approx \prod_{i=1}^{K} p(y_{t+i} \mid h_t),$$

so later drafted positions are predicted from the same hidden state rather than from previously drafted tokens (Ankner et al., 2024).

The Medusa head parameterization is a single-layer feed-forward mapping with residual connection and SiLU activation:

$$p_t^{(k)} = \mathrm{softmax}\!\left(W_2^{(k)} \cdot \big(\mathrm{SiLU}(W_1^{(k)} \cdot h_t) + h_t\big)\right),$$

with $W_1^{(k)} \in \mathbb{R}^{d \times d}$ and $W_2^{(k)} \in \mathbb{R}^{d \times V}$ (Cai et al., 2024). Initialization aligns head predictions with the backbone LM head by copying and zero-initialization, after which the heads are fine-tuned.

Candidate construction uses the top-$t$0 predictions from each head. If the $t$1-th head provides top-$t$2 tokens, the tree of candidate continuations is formed by the Cartesian product across depths, with a virtual root corresponding to the original LM head prediction. The total number of new candidate tokens materialized in one verification pass is

$t$3

The essential systems point is that all heads share the same backbone representation $t$4 and reuse the same KV-cache as the backbone model; no second model or duplicate cache is required (Cai et al., 2024).

Verification relies on tree-based attention. Medusa flattens all candidate tokens into one “mega sequence” but applies a specialized attention mask so that each candidate token attends only to its own branch’s predecessors. This differs from a standard causal mask by disallowing inter-branch leakage while preserving per-branch autoregression. Positional indices are assigned per branch according to depth, and the backbone model then scores all candidate tokens concurrently in one parallel pass (Cai et al., 2024).

3. Decoding and Acceptance

A Medusa decoding iteration proceeds in five stages. First, the backbone forward pass for the last committed token yields $t$5, the original next-token distribution, and the auxiliary head distributions. Second, the system selects top-$t$6 predictions for each head and forms branches by Cartesian product across depths. Third, it performs a single parallel verification pass using the shared KV-cache and the tree attention mask. Fourth, it applies an acceptance rule to candidate prefixes. Fifth, it commits the longest accepted prefix, appends those tokens to the KV-cache, and repeats (Cai et al., 2024).

The paper defines acceleration rate as the average accepted tokens per decoding step and overhead as the ratio of per-step latency to vanilla decoding, with

$t$7

This makes the trade-off explicit: a larger expected accepted prefix length increases acceleration rate, but larger trees or heavier verification can increase overhead (Cai et al., 2024).

Medusa introduces a typical acceptance scheme to improve acceptance rate while maintaining generation quality. Given context $t$8 and a candidate sequence $t$9, a token is accepted when

$t+1$0

The first token at each step is greedily accepted to guarantee progress; subsequent tokens use the entropy-aware threshold (Cai et al., 2024). At temperature $t+1$1, this reduces to greedy decoding, while at higher temperatures typical acceptance tends to accept longer prefixes, improving speedup.

A recurrent point of confusion is the relation between Medusa verification and exact distribution matching. The framework can support rejection sampling, as in speculative decoding, when exact matching to the original model distribution is required. However, the Medusa paper emphasizes typical acceptance because it prioritizes efficiency and reasonable outputs without strict distribution matching (Cai et al., 2024). This means that “lossless” and “distribution-preserving” are not interchangeable descriptions of every Medusa-style deployment.

4. Training Regimes

The original paper defines two training modes. Medusa-1 fine-tunes only the extra heads on top of a frozen backbone LLM. Its objective is a weighted sum of per-head cross-entropies,

$t+1$2

with depth-dependent weights such as $t+1$3 to reflect increasing uncertainty at later positions (Cai et al., 2024). Because the backbone is frozen, generation quality is preserved, and the paper describes this regime as enabling lossless inference acceleration.

Medusa-2 jointly fine-tunes the heads and the backbone. Its objective preserves the backbone’s next-token capability by combining the standard LM loss with the multi-head objective:

$t+1$4

The training recipe uses differential learning rates, a heads warmup scheme, and gradually increasing $t+1$5 so that early large head gradients do not distort the backbone (Cai et al., 2024). This regime improves head prediction accuracy and acceptance, hence higher speedup, but requires special scheduling to preserve the backbone model’s capabilities.

The paper also describes a self-distillation pipeline for settings without training data. The model generates training conversations from seed prompts, and the backbone without adapters provides teacher distributions. The backbone distillation loss is

$t+1$6

To avoid extra memory, LoRA can be applied to the backbone while the teacher is treated as the base model with the adapter turned off (Cai et al., 2024). For Medusa-1, the backbone can be quantized in a QLoRA-like fashion to reduce memory footprint, and the paper states that one-layer heads can be trained in a few hours on one GPU.

5. Variants and Extensions

Later work modifies different parts of the Medusa design space: the dependency structure among heads, the interaction across drafted positions, and the topology of the verification tree. Hydra replaces sequentially independent draft heads with sequentially-dependent heads that condition each drafted token on previously drafted tokens, describing them as a drop-in replacement for standard draft heads and reporting up to $t+1$7 and $t+1$8 improvement in decoding throughput compared to Medusa decoding and autoregressive decoding respectively (Ankner et al., 2024). Amphista retains the multi-head setting but introduces an Auto-embedding Block with bi-directional attention across drafting heads and Staged Adaptation Layers that bridge the autoregressive teacher and non-autoregressive drafter, reporting up to $t+1$9 speedup over vanilla autoregressive decoding and $p_t^{(0)}$0 over Medusa on Vicuna 33B in wall-clock time (Li et al., 2024). Dynamic tree attention replaces Medusa’s fixed tree with a per-step tree selected from head scores, with reported speed-up improvements on Vicuna-7B from $p_t^{(0)}$1 to $p_t^{(0)}$2 for MEDUSA-1 and from $p_t^{(0)}$3 to $p_t^{(0)}$4 for MEDUSA-2 while maintaining similar MT-Bench quality (Zhang, 9 Feb 2025).

Method	Modification	Reported outcome
Hydra++	Sequentially-dependent draft heads	Up to $p_t^{(0)}$5 vs Medusa decoding; $p_t^{(0)}$6 vs autoregressive decoding
Amphista	Bi-directional head interaction plus staged adaptation	Up to $p_t^{(0)}$7 vs vanilla AR; $p_t^{(0)}$8 vs Medusa on Vicuna 33B
Dynamic tree attention	Per-step dynamic candidate tree	Vicuna-7B: $p_t^{(0)}$9 and $k=1,\dots,K$0 speed up

These variants respond to distinct limitations of the original independent-head formulation. Hydra targets the independence mismatch between later drafted positions and the base model’s chain rule. Amphista targets the absence of information interaction across prediction positions while avoiding the sequential overhead of autoregressive drafting. Dynamic tree attention targets the inefficiency of a fixed tree structure that ignores context-dependent head confidence (Ankner et al., 2024, Li et al., 2024, Zhang, 9 Feb 2025).

A more abstract connection appears in theoretical work on multi-head transformers, which shows that different heads can specialize into distinct roles such as traversal and stage control and coordinate through explicit intermediate sequences. That analysis is not an inference-acceleration paper, but it suggests a mechanistic interpretation of why multi-head, multi-step designs can implement staged procedures within a single autoregressive process (Yang et al., 11 Aug 2025).

6. Empirical Profile, Trade-offs, and Limitations

The Medusa paper reports that Medusa-1 achieves speedup of approximately $k=1,\dots,K$1 on Vicuna-7B and approximately $k=1,\dots,K$2 on Vicuna-13B, with MT-Bench scores comparable to baseline, and summarizes the general result as “over $k=1,\dots,K$3” speedup (Cai et al., 2024). Medusa-2 improves the reported speedup to $k=1,\dots,K$4 for both Vicuna-7B and Vicuna-13B in the reported setup, and across tasks reaches $k=1,\dots,K$5–$k=1,\dots,K$6, with the “Extraction” category at approximately $k=1,\dots,K$7 and “Coding” at approximately $k=1,\dots,K$8 (Cai et al., 2024). Representative MT-Bench examples reported for Medusa-2 include acceleration rate $k=1,\dots,K$9, overhead $t+k+1$0, quality $t+k+1$1 for Vicuna-7B; acceleration rate $t+k+1$2, overhead $t+k+1$3, quality $t+k+1$4 for Zephyr-7B; acceleration rate $t+k+1$5, overhead $t+k+1$6, quality $t+k+1$7 for Vicuna-13B; and, for self-distilled Vicuna-33B, acceleration rate $t+k+1$8, overhead $t+k+1$9, quality $$p(y_{t+1:t+K} \mid h_t) \approx \prod_{i=1}^{K} p(y_{t+i} \mid h_t),$$0 (Cai et al., 2024).

The practical profile is therefore not simply “more heads imply more speed.” The paper states that overhead rises with tree size, that sparse optimized trees can maintain or improve acceleration rate with fewer nodes than dense Cartesian trees, and that speed in tokens per second can decrease as tree size grows because of hardware overhead (Cai et al., 2024). The dynamic-tree paper makes the same systems point from a different angle: although its “speed up” metric improves, the unoptimized implementation reports about $$p(y_{t+1:t+K} \mid h_t) \approx \prod_{i=1}^{K} p(y_{t+i} \mid h_t),$$1 fewer tokens per second than fixed Medusa because of mask and candidate-building overhead (Zhang, 9 Feb 2025).

Several limitations recur across the literature. Very large trees can erode net speed gains. Typical acceptance changes the sampling distribution relative to the original model, so exact distributional equivalence is not guaranteed. Integration into high-throughput serving engines such as batched inference in vLLM requires additional engineering. Self-distillation efficacy depends on seed data and teacher quality, and quantization of the teacher in self-distillation can degrade outcomes (Cai et al., 2024). At the same time, the original paper explicitly positions Medusa as complementary to bandwidth-oriented optimizations such as multi-query or grouped-query attention and quantization, because it reduces the number of decoding steps rather than only the cost of each step (Cai et al., 2024).

Taken together, Medusa-style multi-step heads define a family of inference accelerators in which the central design variable is how much future-token structure can be predicted cheaply enough that parallel verification yields a longer accepted prefix than vanilla decoding. The original Medusa formulation demonstrates that independent auxiliary heads plus tree-based verification can deliver substantial acceleration without a separate draft model; later work shows that sequential dependence, bi-directional head interaction, and dynamic tree construction are the principal axes along which that baseline can be extended.