Inference Optimization2024~108 minarXiv:2401.10774

Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

score 10Tianle Cai, Yuhong Li, Zhengyang Geng +4↑ 59★ 2,715Jan 19, 2024cs.LGhttps://github.com/fasterdecoding/medusa

🎯 Pitch

MEDUSA replaces the separate draft model of speculative decoding with multiple lightweight decoding heads attached directly to the frozen LLM, predicting several future tokens in parallel and verifying them via tree attention in a single forward pass. This achieves over 2.2× lossless speedup without the infrastructure headache of maintaining a draft model, and joint training further pushes it to 2.8×.

1. Executive Summary

MEDUSA introduces a method that accelerates LLM inference by attaching multiple lightweight decoding heads to a frozen or jointly fine-tuned backbone model, enabling the parallel prediction of several subsequent tokens without a separate draft model. Evaluated on MT-Bench with Vicuna-7B, 13B, 33B, and Zephyr-7B, the approach uses MEDUSA heads (single-layer feed-forward networks appended to the last hidden states—each head predicting the token at a specific future position) and tree attention (a structured attention mask that verifies multiple candidate continuations simultaneously within a single forward pass). MEDUSA-1 achieves over 2.2× wall-time speedup losslessly, while MEDUSA-2 jointly trains the heads with the backbone to reach 2.3–2.8× speedup, establishing that multi-head parallel prediction can substitute for a draft model in speculative decoding only when the training recipe explicitly preserves the backbone's next-token capability through combined loss, differential learning rates, and head warmup.

2. Context and Motivation

The Core Problem: Auto-Regressive Decoding Is Memory-Bandwidth-Bound

The fundamental problem MEDUSA addresses is deceptively simple: generating text with large language models is slow, and the reason has less to do with raw computation than with how data moves between memory and processors. The paper identifies this as a memory-bandwidth-bound phenomenon (Section 1, Appendix A.1). When an LLM generates tokens auto-regressively—producing one token at a time, each step dependent on the previous one—the forward pass loads the entire model weight matrix from High-Bandwidth Memory (HBM, which is fast but limited in capacity) into the accelerator's cache (SRAM, which is much faster but tiny). For each token generated, all billions of parameters must be transferred across this memory bus. Modern accelerators can execute matrix multiplications far faster than they can move data from HBM, meaning the arithmetic units spend most of their time idle, waiting for parameters to arrive. The paper captures this precise bottleneck:

"This bottleneck is inherent to the sequential nature of auto-regressive decoding, where each forward pass requires transferring the complete model parameters from High-Bandwidth Memory (HBM) to the accelerator's cache. This process, which generates only a single token, underutilizes the arithmetic computation potential of modern accelerators, leading to inefficiency." (Section 1)

This matters for several practical reasons the authors highlight explicitly (Section 1, Impact Statement). First, latency-sensitive applications—chatbots, coding assistants, interactive writing tools—require near-instantaneous token generation to feel responsive. Second, local deployment of models on consumer hardware (e.g., on-device LLMs) is hampered by memory constraints, making inference speed a barrier to democratization. Third, operational cost in production systems scales directly with the number of forward passes; reducing decoding steps translates to lower cloud-compute bills. Fourth, the growing disparity between model size and memory bandwidth: models have scaled to hundreds of billions of parameters while HBM bandwidth has improved much more slowly, making the bottleneck increasingly severe over time.

The Intuition: Increase Arithmetic Intensity

The paper frames the solution space through a specific systems concept: arithmetic intensity—the ratio of total floating-point operations (FLOPs) to total data movement (bytes transferred). Auto-regressive decoding has low arithmetic intensity because each forward pass moves a large weight matrix but performs computation only for a single token's embedding. The fix is conceptually straightforward: generate multiple tokens per forward pass, so the expensive parameter loads are amortized over more predicted tokens, increasing arithmetic intensity and reducing the total number of data transfers.

However, the fundamental constraint of auto-regression is that the model needs the previous token to predict the next one. How do you predict multiple tokens without violating this dependency? This is the architectural challenge that all acceleration methods must solve.

Prior Approaches and Their Limitations

Quantization and KV-Cache Reduction

The paper acknowledges two categories of existing approaches (Appendix A.1). The first category focuses on reducing memory consumption rather than reducing decoding steps: quantization (Xiao et al., 2023a; Dettmers et al., 2022; Frantar et al., 2022; Lin et al., 2023; Kim et al., 2023) compresses model weights into lower bit-widths, and KV-cache reduction techniques (Shazeer, 2019; Ainslie et al., 2023; Zhang et al., 2023) shrink the key-value states stored during generation. These methods are orthogonal to MEDUSA—they reduce the per-step cost but do not reduce the number of steps. The paper does not position against these; rather, it treats them as complementary (they can be applied alongside MEDUSA, and MEDUSA-1 is explicitly compatible with quantized backbones via QLoRA).

Speculative Decoding: The Direct Competitor

The second category—and MEDUSA's direct intellectual competition—is speculative decoding (Leviathan et al., 2022; Chen et al., 2023). The idea: use a small, fast "draft model" to speculatively generate a sequence of candidate tokens (say, 4–8 tokens), then run the large model once to verify all candidates in parallel. Where the draft model's prediction matches the large model's distribution, tokens are accepted; where they diverge, tokens are rejected and re-sampled. The large model scores all candidates in a single forward pass, so if the acceptance rate is high, the effective throughput multiplies by the average accepted length.

Why this approach is elegant. Speculative decoding provides lossless acceleration: the rejection sampling scheme (a specific token-level accept/reject procedure based on comparing the draft model's probability to the original model's probability) guarantees that the output distribution exactly matches what the original model would have produced. This means you can drop it into any deployment without affecting generation quality—a strong practical property.

Where speculative decoding falls short. The paper identifies three concrete pain points with the draft model approach:

Acquiring an appropriate draft model is difficult and expensive. The draft model must be small (fast inference) yet capable of generating continuations the large model will accept. Existing approaches often require separate pre-training of the draft model, which Miao et al. (2023) reports consumed "275 NVIDIA A100 GPU hours." This is not a trivial upfront cost—it's a substantial engineering investment that must be repeated for every new large model.
Distribution shift between draft and target models. A separately trained draft model—even if small and efficient—produces tokens from a different distribution than the target model. This reduces the acceptance rate, directly degrading speedup. The rejection sampling mechanism guarantees correctness but cannot recover the lost efficiency: if the draft model generates tokens the large model would never have produced, those tokens are simply rejected, and the compute spent on them is wasted.
Serving complexity in distributed systems. Chen et al. (2023) highlighted that deploying two separate models in a distributed serving environment introduces operational challenges: managing memory for both models, routing between them, handling different throughput characteristics, and coordinating parallel strategies. This is not just a one-time integration cost—it's an ongoing operational burden.

Additional practical friction includes: (a) the pre-training data used for the draft model may not match the target model's training distribution, (b) the draft model must be continuously updated if the target model evolves, and (c) hyperparameters like draft length ( $\gamma$ ) need careful tuning per model size—the paper's own speculative decoding results in Appendix D (Figure 7) show that the optimal $\gamma$ differs across Vicuna-7B ( $\gamma=4$ with Llama-68M), Vicuna-13B ( $\gamma=3$ with Llama-68M), and Vicuna-33B ( $\gamma=3$ with Tiny-Vicuna). There is no one-size-fits-all.

Ancient History: Blockwise Parallel Decoding (Stern et al., 2018)

MEDUSA is not the first attempt to avoid a separate draft model by attaching prediction heads to the backbone. Stern et al. (2018) proposed Blockwise Parallel Decoding for machine translation and image super-resolution, where multiple decoding heads appended to the model's hidden states predict several future tokens in parallel. This directly inspired MEDUSA's architecture.

Why didn't this solve the problem for LLMs? Stern et al. (2018) was developed in a pre-LLM era (2018) for sequence-to-sequence models on tasks like translation. The techniques weren't studied in the context of large-scale auto-regressive language models where (a) the memory-bandwidth bottleneck is acute, (b) generation tasks span diverse and open-ended domains, and (c) quality preservation is non-negotiable. The method also lacked several critical innovations that MEDUSA introduces to make it effective: generating and verifying multiple candidate continuations via tree attention (Stern et al. only produced a single candidate sequence per step), a systematic training methodology that preserves the backbone's capability, and a principled acceptance scheme that balances speed and quality.

The MEDUSA paper "revisits and refines" this concept (Section 1), acknowledging Stern et al. as the intellectual ancestor while substantially re-engineering the approach for modern LLMs.

How MEDUSA Positions Itself

MEDUSA occupies a specific, well-motivated niche in the inference acceleration landscape:

It is a zero-draft-model alternative to speculative decoding. The central design philosophy is: use the original model's own knowledge for prediction, not a separate model. Instead of training an external draft model, MEDUSA adds lightweight decoding heads directly on top of the frozen (or jointly trained) backbone model. These heads leverage the rich hidden representations that the backbone already produces—the last hidden states encode substantial information about what tokens are likely to follow. By training only single-layer feed-forward networks (not a full transformer), MEDUSA achieves parameter efficiency and eliminates distribution shift (the heads learn from the same representations the backbone itself produces).

It is an enabling tool for democratized LLM deployment. The paper emphasizes accessibility: MEDUSA-1 can be trained on a single consumer GPU using a quantized backbone (via QLoRA integration), taking approximately 5 hours on an A100 PCIe for Vicuna-7B (Section 2.2.1). This contrasts sharply with speculative decoding's multi-hundred-GPU-hour pre-training requirements. The paper's self-distillation pipeline (Section 2.3.2) further reduces barriers by generating training data from the model itself when the original training dataset is unavailable—for example, with RLHF-tuned models like Zephyr-7B where the output distribution differs from the pre-training data.

It treats generation quality as a first-class constraint, not an afterthought. The authors explicitly acknowledge that naive approaches fail. Directly fine-tuning the backbone with MEDUSA heads without the proper recipe degrades generation quality (Section 2.2.2, validated in Table 2: "Direct Fine-tuning" scores 5.925 vs. 6.17 for the baseline on MT-Bench). MEDUSA-2's training recipe (combined loss that includes the backbone's cross-entropy, differential learning rates, head warmup via two-stage training) is designed specifically to prevent this degradation. The typical acceptance scheme (Section 2.3.1) further demonstrates awareness of the quality-speed tradeoff, providing a tunable threshold that allows practitioners to choose their operating point.

It scales with model size without proportional overhead. The MEDUSA heads themselves are negligible in parameter count (single-layer FFNs with residual connections, $K \cdot 2 \cdot d^2$ parameters for $K$ heads and hidden dimension $d$ , compared to the backbone's billions of parameters). The tree attention mechanism processes multiple candidates in a single forward pass, and the paper's study of hardware constraints (Appendix G) shows that the overhead from additional candidate tokens is modest because the linear layers remain memory-bandwidth-bound for typical batch sizes.

The Missing Piece in Prior Work: A Systematic Training Methodology

Prior work on multi-head parallel decoding (Stern et al., 2018) demonstrated the architectural possibility but didn't address the practical challenge of training the heads without destroying the model's capabilities. Speculative decoding solved this by keeping the draft model separate—the backbone is never modified. MEDUSA takes the harder path of integrating prediction heads into the backbone, which requires solving a specific sub-problem: how do you add new parameters and train them (potentially alongside the backbone) while guaranteeing that the model's next-token prediction quality remains intact?

This is non-trivial because fine-tuning language models on auxiliary tasks can cause catastrophic forgetting or distort the learned representations (Kumar et al., 2022, cited by the authors). The paper's two-level training strategy (MEDUSA-1 frozen, MEDUSA-2 joint with careful regularization) is therefore a core contribution that makes the architectural idea practically viable.

Summary of the Gap and Position

Aspect	Speculative Decoding	MEDUSA
Draft mechanism	Separate small model	Multi-head layers on backbone
Pre-training cost	High (hundreds of GPU hours)	Low (hours on single GPU, MEDUSA-1)
Distribution shift	Present (separate model)	Eliminated (shared representations)
Serving complexity	Two models to manage	Single model system
Quality guarantee	Rejection sampling (lossless)	Typical acceptance or rejection sampling
Training data dependency	Requires draft model data	Can use self-distillation
Integration with fine-tuning	Separate from backbone training	Native via MEDUSA-2 recipe

MEDUSA directly addresses speculative decoding's friction points (draft model acquisition, distribution shift, serving complexity) while building on the older parallel-decoding idea from Stern et al. (2018). The key insight that enables this: the backbone's last hidden states already contain rich predictive information about future tokens, and lightweight heads can exploit this information without the heavyweight process of training a separate transformer model.

3. Technical Approach

3.1 Reader orientation (approachable technical breakdown)

MEDUSA is a set of extra "prediction heads" that sit on top of an existing LLM and guess several future tokens at once, plus a tree-shaped attention mask that lets the model check all those guesses in a single forward pass. The technique solves the memory-bandwidth bottleneck of auto-regressive decoding—where each token requires a full traversal of model weights from slow memory to fast cache—by amortizing that expensive weight transfer over multiple candidate tokens predicted and verified together, rather than generating tokens one at a time.

3.2 Big-picture architecture (diagram in words)

The MEDUSA inference pipeline has five major components operating within a single extended decoding step:

The backbone LLM (e.g., Vicuna-7B)—a standard transformer that processes the input context through all its layers to produce the last hidden state $\mathbf{h}_t$ at position $t$ . This model performs the expensive weight-loading operation once per MEDUSA decoding step, not once per token.
MEDUSA heads ( $K$ additional single-layer feed-forward networks appended to $\mathbf{h}_t$ )—the $k$ -th head predicts a probability distribution $\mathbf{p}^{(k)}_t$ over the vocabulary for the token at position $t + k + 1$ . The original language model (LM) head predicts the immediate next token $\mathbf{p}^{(0)}_t$ at position $t + 1$ . All heads compute in parallel from the same hidden state.
Tree-structured candidate construction—for each head, the $s_k$ most probable tokens are selected. The Cartesian product across heads builds a tree of candidate continuations rooted at the current context, with the $k$ -th level having $\prod_{i=1}^k s_i$ nodes. A virtual root (level 0) corresponds to the original LM head's prediction.
Tree attention mask and positional encoding—a custom attention pattern that restricts each candidate token to attend only to its direct ancestors in the tree (not to tokens in other branches). The full model processes the entire batch of candidate tokens in one forward pass, computing logits for all positions simultaneously.
Verification and acceptance—either rejection sampling (guaranteeing distributional equivalence with the original model) or typical acceptance (a threshold-based scheme that accepts plausible candidates based on the original model's probability), which selects the longest valid prefix from the candidate tree and commits those tokens to the output sequence. The next decoding step begins from the final accepted position.

Information flow per decoding step: (1) backbone computes $\mathbf{h}_t$ from the current prefix → (2) all $K+1$ heads predict distributions over the vocabulary → (3) top- $s_k$ tokens from each head form a candidate tree → (4) tree attention processes all candidates through the backbone in one forward pass, producing verification logits → (5) acceptance scheme selects a prefix of candidates → (6) the selected prefix extends the output sequence, and the process repeats.

3.3 Roadmap for the deep dive

First, the MEDUSA heads architecture: the mathematical form of each head, the parameter count, the initialization strategy, and why a single FFN layer suffices. This is the core prediction mechanism that replaces the draft model.
Second, tree attention: the candidate construction procedure, the tree-structured attention mask, the positional encoding adjustments, and the computational cost analysis. This is the mechanism that enables parallel verification of many candidate sequences.
Third, the two training strategies (MEDUSA-1 and MEDUSA-2): the loss functions, the differential learning rates, the head warmup procedure, and the combined loss formulation. This covers how the heads are fitted to a model without destroying its capabilities.
Fourth, typical acceptance: the mathematical criterion for accepting candidate tokens, the entropy-dependent threshold, the relationship to truncation sampling, and how it differs from rejection sampling in both mechanism and speed-quality tradeoff.
Fifth, self-distillation: the data generation pipeline for models without accessible training data, the KL-divergence distillation loss for the backbone, and the LoRA-based trick that avoids maintaining two models in memory.
Sixth, optimized tree construction: the calibration-based greedy algorithm that selects which nodes to include in the tree to maximize expected acceptance length under a fixed node budget.

3.4 Detailed, sentence-based technical breakdown

This is a systems architecture paper with an empirical validation methodology, not a theoretical contribution. The core idea is that the last hidden state of a transformer LLM encodes enough predictive information to guess multiple future tokens, and lightweight heads trained on top of that hidden state can replace a separate draft model for speculative decoding—but only if the training recipe explicitly preserves the backbone model's next-token capability.

MEDUSA Heads: Architecture and Mechanism

Each MEDUSA head is a single feed-forward layer with a residual connection applied to the backbone's last hidden state. For a given position $t$ in the sequence, the backbone LLM computes a hidden state vector $\mathbf{h}_t \in \mathbb{R}^d$ (where $d$ is the model's hidden dimension, e.g., 4096 for Llama-7B). The MEDUSA heads all consume this same $\mathbf{h}_t$ —they do not see intermediate hidden states from earlier transformer layers, only the final output.

The $k$ -th MEDUSA head (for $k = 1, 2, \ldots, K$ ) produces a probability distribution over the vocabulary $V$ for the token at position $t + k + 1$ :

$\mathbf{p}^{(k)}_t = \text{softmax}\left( \mathbf{W}^{(k)}_2 \cdot \left( \text{SiLU}(\mathbf{W}^{(k)}_1 \cdot \mathbf{h}_t) + \mathbf{h}_t \right) \right)$

where $\mathbf{W}^{(k)}_1 \in \mathbb{R}^{d \times d}$ is the first weight matrix (a dense projection that transforms the hidden state), $\mathbf{W}^{(k)}_2 \in \mathbb{R}^{d \times V}$ is the second weight matrix (mapping back to vocabulary-sized logits), $d$ is the backbone's hidden dimension, and $V$ is the vocabulary size (e.g., 32,000 for Llama). The SiLU (Sigmoid Linear Unit) activation is the same non-linearity used in the Llama model family (Touvron et al., 2023), ensuring consistency with the backbone's own activations.

What it computes, operationally: For each candidate future position $t+k+1$ , the head takes the backbone's hidden state $\mathbf{h}_t$ , applies a learned linear transformation $\mathbf{W}^{(k)}_1$ , passes it through the SiLU gating function, adds back the original hidden state $\mathbf{h}_t$ via the residual connection, and projects to vocabulary size via $\mathbf{W}^{(k)}_2$ . The softmax converts the resulting logits into a proper probability distribution. Each head $k$ maps $\mathbf{h}_t \rightarrow$ a $V$ -dimensional probability vector over tokens that could appear $k+1$ positions in the future.

Why this form: The residual connection ( $+ \mathbf{h}_t$ ) serves a critical initialization purpose. The authors initialize $\mathbf{W}^{(k)}_2$ to be identical to the original language model head's weight matrix and $\mathbf{W}^{(k)}_1$ to zero. This means that at initialization, before any MEDUSA-specific training, $\mathbf{W}^{(k)}_1 \cdot \mathbf{h}_t = \mathbf{0}$ , the SiLU of zero is zero, and the residual connection passes $\mathbf{h}_t$ directly to $\mathbf{W}^{(k)}_2$ , making the head's initial prediction exactly equal to the original LM head's prediction:

"This aligns the initial prediction of MEDUSA heads with that of the original model." (Section 2.1.1)

Without this initialization, the heads would start from a random distribution, requiring much more training to converge and potentially causing large initial gradients that could destabilize training. The residual connection also provides a gradient highway that allows information from the frozen (or slowly-updating) backbone to flow directly to the output, preventing the head from needing to re-learn the entire mapping from scratch.

The authors empirically determined that five heads are sufficient for most scenarios (Section 2.2.3): "Empirically, we found that five heads are sufficient at most." This means $K \leq 5$ , and the heads predict tokens at positions $t+2, t+3, \ldots, t+6$ (since the original LM head handles $t+1$ ). The parameter count is modest: each head contributes $d \times d + d \times V$ parameters (for the two weight matrices), plus negligible bias terms. For Llama-7B with $d = 4096$ and $V = 32000$ , each head has approximately $4096 \times 4096 + 4096 \times 32000 \approx 16.8\text{M} + 131.1\text{M} = 147.9\text{M}$ parameters. Five heads total roughly 740M parameters—substantial in absolute terms but small relative to the 7B backbone (about 10.5% overhead), and these additional parameters are only used during the candidate-generation phase, not stored in the KV cache.

Design choice: why not deeper networks? The paper explicitly chose a single-layer design, following Stern et al. (2018). The rationale is pragmatic: deeper prediction heads would (a) add more parameters and latency, (b) require more training data to converge, and (c) be unnecessary because the backbone's last hidden state already encodes rich contextual information—the head primarily needs to project this information to future positions, not to re-process it. The authors state: "We find that this simple design is sufficient to achieve satisfactory performance" (Section 2.1.1), indicating that more complex architectures were considered but did not justify their additional cost.

Tree Attention: Parallel Verification of Multiple Candidates

With $K$ MEDUSA heads plus the original LM head, MEDUSA obtains $K+1$ probability distributions: $\mathbf{p}^{(0)}_t$ (immediate next token from the LM head) and $\mathbf{p}^{(1)}_t, \ldots, \mathbf{p}^{(K)}_t$ (future tokens from the MEDUSA heads). The naive approach would be to sample a single continuation from these distributions—take the most likely token from $\mathbf{p}^{(0)}_t$ , then from $\mathbf{p}^{(1)}_t$ , and so on. However, this would produce exactly one candidate sequence per decoding step, underutilizing the available parallel compute.

The key insight of tree attention is to sample multiple tokens from each head and verify many candidate continuations in parallel. The procedure works as follows:

Step 1: Top- $s_k$ selection. For each head $k$ (where $k = 0$ is the LM head, $k = 1, \ldots, K$ are MEDUSA heads), the system selects the $s_k$ most probable tokens from the head's output distribution. The value $s_k$ is a hyperparameter per head—for example, in Figure 2, $s_1 = 2$ (two candidates from the first MEDUSA head) and $s_2 = 3$ (three candidates from the second MEDUSA head). These are not required to be equal across heads; they can be tuned based on each head's prediction accuracy.

Step 2: Cartesian product construction. The candidates form a tree. The root (level 0) is the current context prefix. Level 1 contains the $s_0$ candidates from the LM head. For each level-1 candidate, level 2 contains $s_1$ candidates from the first MEDUSA head, giving $s_0 \times s_1$ nodes at level 2. In general, level $k+1$ contains $\prod_{i=0}^{k} s_i$ nodes. This is explicitly described as:

"These candidates are established by determining the Cartesian product of the top- $s_k$ predictions from each head. For instance, in Figure 2, with $s_1 = 2$ and $s_2 = 3$ , each first head prediction can be succeeded by any prediction from the second head." (Section 2.1.2)

The total number of new tokens introduced (candidate nodes in the tree, excluding the root) is:

$\sum_{k=1}^{K} \prod_{i=1}^{k} s_i$

where $K$ is the number of MEDUSA heads (the LM head at level 1 is treated separately; the paper's notation counts from $k=1$ for the first MEDUSA head, with the LM head providing the $s_0$ root-level branching).

Step 3: Attention mask with tree topology. Normally, transformer decoders use a causal attention mask—each token attends to all preceding tokens. With a tree of candidates, this would be incorrect: a candidate in branch A should not attend to a candidate in branch B at the same depth, because they represent mutually exclusive continuations. The tree attention mask therefore restricts each token to attend only to its ancestors in the tree—the tokens along the path from the root to that node.

"Within this framework, only tokens from the same continuation are regarded as historical data. Drawing inspiration from the concept of embedding graph structures into attention as proposed in the graph neural network domain (Ying et al., 2021), we incorporate the tree structure into our attention mask, visualized in Figure 2." (Section 2.1.2)

Concretely, the attention mask is a binary matrix $\mathbf{M} \in \{0, 1\}^{N_{\text{total}} \times N_{\text{total}}}$ where $N_{\text{total}}$ is the number of tokens in the extended sequence (original prefix plus all candidate nodes). For a candidate token at node $i$ and a potential attended token at node $j$ , $\mathbf{M}_{ij} = 1$ if node $j$ is an ancestor of node $i$ in the tree, and $\mathbf{M}_{ij} = 0$ otherwise. All original prefix tokens are ancestors of all candidate nodes (since all candidates extend from the prefix).

Step 4: Positional encoding adjustment. Transformers use positional encodings (rotary position embeddings in Llama) that assign each token a position index. In the tree structure, multiple candidate tokens at the same depth should have the same positional index—they all represent the $t + \ell$ -th token in their respective branches. The paper adjusts positional indices "in line with this structure" (Section 2.1.2): each candidate at depth $\ell$ from the root receives the position index $t + \ell$ , regardless of which branch it belongs to. Original prefix tokens retain their original positions $1, 2, \ldots, t$ .

Step 5: Single forward pass through the backbone. The entire extended sequence (prefix + all $N_{\text{total}}$ candidate tokens) is processed through the backbone transformer in one forward pass. The tree mask ensures that the model computes each candidate's hidden state conditioned only on its valid ancestral context. The forward pass produces logits for every candidate position simultaneously.

Step 6: Verification. For each candidate token (say, a token predicted by the first MEDUSA head to be at position $t+3$ ), the backbone's forward pass produces a logit vector at that position. These logits represent the backbone's actual prediction of what token should appear there, conditioned on the specific branch context (the ancestral tokens in that candidate's path). The verification step compares the MEDUSA head's proposed token against the backbone's distribution at that position, using either rejection sampling or typical acceptance (detailed in the acceptance scheme section below).

What tree attention achieves, operationally: It transforms $\prod_{i=0}^{K} s_i$ independent candidate sequences into a single batched computation. Without tree attention, verifying $C$ candidate sequences would require either $C$ separate forward passes (defeating the purpose of acceleration) or padding all sequences to equal length and processing them with a batch dimension (wasting computation on pad tokens and preventing candidates from sharing the prefix computation). Tree attention amortizes the prefix computation across all candidates (each prefix token is computed once and attended to by all candidate branches) and processes all candidates in a single forward pass by exploiting the sparse connectivity of the tree.

Why this specific tree shape (Cartesian product): The authors acknowledge that the Cartesian product is "the most simple and regular way to construct the tree structure" (Section 2.1.2), and they immediately note that more sophisticated tree structures are possible. The Cartesian product is easy to implement, predictable in memory footprint, and naturally decomposes the prediction problem—each head specializes in a specific future offset, and the product combines their independent predictions combinatorially. The optimized tree construction (Section 2.3.3) refines this by pruning low-accuracy nodes.

Computational cost analysis. The forward pass processes $N_{\text{prefix}} + N_{\text{candidates}}$ tokens, where $N_{\text{candidates}} = \sum_{k=1}^{K} \prod_{i=1}^{k} s_i$ . The linear layers ( $\mathbf{X}\mathbf{W}_Q, \mathbf{X}\mathbf{W}_K, \mathbf{X}\mathbf{W}_V$ , and the feed-forward layers $\mathbf{X}\mathbf{W}_u, \mathbf{X}\mathbf{W}_g, \mathbf{X}\mathbf{W}_d$ ) scale with the number of tokens being processed. However, the key insight from the roofline analysis (Appendix G.2) is that for small batch sizes (the paper focuses on batch size 1), these linear layers are memory-bandwidth-bound—their arithmetic units are underutilized. Adding candidate tokens increases the FLOPs per weight-load, moving the operational intensity (FLOPs per byte transferred) higher and actually improving hardware utilization. Figure 20 and Table 8 show that for a batch size of 1, processing 64 candidate tokens increases the linear layer throughput from 1.26 TFLOP/s (operational intensity 1.0) to 76.57 TFLOP/s (operational intensity 63.17), a 61× improvement in effective FLOP/s because the same weight transfer now feeds more computation. This is the dual benefit: more tokens decoded per step and better hardware utilization per step.

However, there is a tradeoff. The attention matrix multiplication ( $\mathbf{Q}\mathbf{K}^T$ ) grows quadratically with the number of candidate tokens (since each candidate must attend to all its ancestors, not just a single prefix). Table 6 shows that for sequence length 1024 and batch size 16, going from 1 candidate to 112 candidates increases attention latency from 1.24 TFLOP/s (operational intensity 0.99) to 75.45 TFLOP/s (operational intensity 56.44). While throughput improves, the absolute latency still increases. The paper's ablation in Figure 4b demonstrates this: speed initially increases with tree size as the acceleration rate (average tokens accepted per step) grows faster than the overhead, but beyond approximately 64 candidate nodes, the overhead from additional matrix multiplications overtakes the marginal acceleration gains, and speed declines. The optimal tree size is therefore a hardware-dependent quantity that balances acceptance rate improvements against computational overhead.

Training Strategies: MEDUSA-1 and MEDUSA-2

The paper proposes two levels of training depending on whether the backbone model is frozen or jointly optimized. The critical challenge is that directly fine-tuning the backbone alongside MEDUSA heads without proper constraints degrades the model's core next-token prediction quality—validated empirically in Table 2, where "Direct Fine-tuning" on Vicuna-7B drops MT-Bench score from 6.17 to 5.925.

Training data assumption. Both strategies assume access to a dataset matching the target model's output distribution. For models fine-tuned on public data (e.g., Vicuna-7B on ShareGPT), the original training dataset is used directly. For models with private training data or RLHF post-training (e.g., Zephyr-7B), the self-distillation pipeline (Section 2.3.2) generates synthetic training data.

MEDUSA-1: Frozen Backbone

In MEDUSA-1, the backbone model parameters are completely frozen. Only the MEDUSA heads' weight matrices ( $\mathbf{W}^{(k)}_1$ and $\mathbf{W}^{(k)}_2$ for each $k$ ) are updated. The backbone is used purely as a feature extractor—it computes $\mathbf{h}_t$ for each training position, and the heads learn to predict future tokens from those hidden states.

Loss function. For each training position $t$ , the $k$ -th MEDUSA head produces a probability distribution $\mathbf{p}^{(k)}_t$ over the vocabulary. The ground truth token at position $t + k + 1$ is $y_{t+k+1}$ . The loss for head $k$ is the negative log-likelihood:

$\mathcal{L}_k = -\log p^{(k)}_t(y_{t+k+1})$

where $p^{(k)}_t(y)$ denotes the probability assigned to token $y$ by the $k$ -th head at position $t$ .

The total MEDUSA-1 loss is a weighted sum across heads:

$\mathcal{L}_{\text{MEDUSA-1}} = \sum_{k=1}^{K} -\lambda_k \log p^{(k)}_t(y_{t+k+1})$

where $\lambda_k$ is a per-head weight and $K$ is the number of MEDUSA heads (the LM head is not included—it is frozen and not trained).

What it computes, operationally: For each position in the training sequence, the backbone produces a hidden state $\mathbf{h}_t$ . Each MEDUSA head $k$ is a function that maps $\mathbf{h}_t \rightarrow \mathbf{p}^{(k)}_t$ . The loss measures how well this predicted distribution concentrates probability on the actual token that appears $k+1$ positions later in the training data. The weighted sum combines losses across all future offsets, with larger offsets ( $k$ larger) typically receiving smaller weights because predictions become more uncertain further into the future.

Why the weighting $\lambda_k$ : The authors observe that $\mathcal{L}_k$ is "larger when $k$ is larger, which is reasonable since the prediction of the $k$ -th head is more uncertain when $k$ is larger" (Section 2.2.1). The weights balance the gradients so that heads predicting further into the future don't dominate the training signal. The paper sets $\lambda_k = 0.8^k$ (the $k$ -th power of a constant), giving a decaying geometric progression: $\lambda_1 = 0.8$ , $\lambda_2 = 0.64$ , $\lambda_3 = 0.512$ , $\lambda_4 = 0.4096$ , $\lambda_5 = 0.32768$ . Without this weighting, the distant heads would contribute larger loss terms (since they're more uncertain) and could receive disproportionately large gradient updates, potentially destabilizing training or causing the model to sacrifice near-future prediction accuracy to marginally improve far-future guesses.

Memory efficiency via quantization. Since the backbone is frozen and only used to compute hidden states (no gradient flows through it), MEDUSA-1 is compatible with quantized backbones. The authors explicitly draw the parallel to QLoRA (Dettmers et al., 2023): "we can use a quantized version of the backbone model to reduce the memory consumption. This introduces a more democratized way to accelerate LLM inference, as with the quantization, MEDUSA can be trained for a large model on a single consumer GPU similar to QLoRA" (Section 2.2.1). The paper reports that MEDUSA-1 training on Vicuna-7B with a 4-bit quantized backbone takes "5 hours for MEDUSA-1 on Vicuna 7B model with a single NVIDIA A100 PCIE GPU to train on 60k ShareGPT samples" (Section 2.2.1).

Training hyperparameters (Appendix B.3): Global batch size 64, peak learning rate $5 \times 10^{-4}$ for the backbone (which is frozen—this likely refers to the optimizer state or is a typo; the heads are the only trainable parameters) and $2 \times 10^{-3}$ for MEDUSA heads, cosine learning rate schedule with warmup for 40 steps, 8-bit AdamW optimizer. Training runs for 2 epochs on the ShareGPT dataset.

Design choice: why freeze the backbone? Freezing guarantees zero degradation of the original model's next-token capability—there is no risk of catastrophic forgetting or distribution shift because the original model's weights are untouched. This makes MEDUSA-1 a safe, lossless acceleration method: if the MEDUSA heads perform poorly, the model can fall back to standard auto-regressive decoding with exactly the same quality as before training. The tradeoff is that frozen backbones don't adapt their representations to help the MEDUSA heads, limiting head prediction accuracy and therefore the achievable speedup (the heads must learn to extract future-token information from hidden states that were optimized solely for next-token prediction).

MEDUSA-2: Joint Training

MEDUSA-2 trains the MEDUSA heads together with the backbone model, allowing the backbone to adapt its hidden representations to better support multi-token prediction. However, this introduces the risk that the backbone's primary capability—predicting the immediate next token accurately—could degrade. The paper proposes three specific mechanisms to prevent this degradation.

Combined loss function. The total MEDUSA-2 loss adds the backbone's standard next-token prediction loss to the weighted MEDUSA head losses:

$\mathcal{L}_{\text{MEDUSA-2}} = \mathcal{L}_{\text{LM}} + \lambda_0 \mathcal{L}_{\text{MEDUSA-1}}$

where:

$\mathcal{L}_{\text{LM}} = -\log p^{(0)}_t(y_{t+1})$

is the standard cross-entropy loss of the original language model head predicting the immediate next token $y_{t+1}$ from hidden state $\mathbf{h}_t$ , and $\lambda_0$ is a hyperparameter controlling the relative weight of the MEDUSA auxiliary loss.

What it computes, operationally: At each training step, the backbone computes $\mathbf{h}_t$ for every position. The LM head produces $\mathbf{p}^{(0)}_t$ and incurs loss $\mathcal{L}_{\text{LM}}$ for predicting $y_{t+1}$ . Meanwhile, each MEDUSA head $k$ produces $\mathbf{p}^{(k)}_t$ and incurs loss $\mathcal{L}_k$ for predicting $y_{t+k+1}$ . The total loss is the backbone's loss plus $\lambda_0$ times the sum of weighted MEDUSA head losses. Gradients flow through all parameters: the backbone is updated to simultaneously optimize for next-token prediction and for providing features that help the MEDUSA heads predict future tokens.

Why this form: The $\mathcal{L}_{\text{LM}}$ term acts as a regularizer—it explicitly constrains the backbone to maintain its core capability even as it adapts to the auxiliary task. Without it, the backbone could drift toward representations that are excellent for multi-token prediction but poor for next-token prediction, which would defeat the purpose (since the verification step in inference relies on the backbone's accurate next-token probabilities to accept or reject candidates). The weight $\lambda_0$ controls the strength of the auxiliary task: a small $\lambda_0$ (e.g., 0.2 as used for Vicuna experiments, Appendix B.3) ensures the auxiliary loss doesn't overwhelm the primary objective.

Differential learning rates. The backbone model is already well-trained; the MEDUSA heads start from the initialization described earlier (identity-mapped to LM head predictions). To allow the heads to converge faster without destabilizing the backbone, the paper uses separate learning rates:

"Since the backbone model is already well-trained and the MEDUSA heads need more training, we can use separate learning rates for them to enable faster convergence of MEDUSA heads while preserving the backbone model's capability." (Section 2.2.2)

The paper sets the MEDUSA head learning rate to be 4 times larger than the backbone learning rate (Appendix B.2). Concretely, for the Vicuna experiments (Appendix B.3), the backbone peak learning rate is $5 \times 10^{-4}$ and the MEDUSA head peak learning rate is $2 \times 10^{-3}$ , a $4\times$ ratio. This differential allows the heads to rapidly adapt to the multi-token prediction task while the backbone evolves more cautiously.

Head warmup (two-stage training). At the start of joint training, the MEDUSA heads have high loss (they haven't yet learned to predict future tokens accurately). Large gradients from this loss could distort the backbone's parameters before the heads converge. The paper introduces a warmup procedure inspired by Kumar et al. (2022), who showed that fine-tuning can distort pre-trained features:

Stage 1 (MEDUSA-1 training): Train only the MEDUSA heads with the backbone frozen, using $\mathcal{L}_{\text{MEDUSA-1}}$ . This gets the heads to a reasonable accuracy without touching the backbone. The stage 1 trained heads serve as initialization for stage 2.
Stage 2 (Joint training with warmup): Unfreeze the backbone and train jointly. The paper describes two variants of warmup within stage 2:
- Simple strategy: "first train the backbone model for a few epochs, then train the MEDUSA heads together with the backbone model" (Section 2.2.2). This means the backbone initially trains alone (with $\mathcal{L}_{\text{LM}}$ only), giving it time to adjust to the joint training regime before the heads join.
- Sophisticated strategy: "gradually increasing the weight $\lambda_0$ of the backbone model's loss" (Section 2.2.2). Start with a small $\lambda_0$ and increase it to the target value over training, preventing the MEDUSA loss from dominating the early gradient signal. For the self-distillation experiments (Appendix B.4), the authors use a sine schedule for $\lambda_0$ to "gradually increase the value to its peak at the end of the training."

The authors state that "both strategies work well in practice" (Section 2.2.2).

Design choice: why this three-part recipe? Each component addresses a distinct failure mode. The combined loss prevents catastrophic forgetting of next-token prediction (the backbone's gradients always include a term pushing toward accurate immediate token prediction). Differential learning rates handle the asymmetry that the backbone is near-optimal while the heads are randomly initialized (or stage-1 pre-trained). Head warmup prevents early-stage gradient interference—the heads' initially large, poorly-conditioned gradients from distorting the backbone's carefully learned representations before the heads have converged enough to provide useful signal. Together, these mechanisms enable joint training without the quality degradation observed with "Direct Fine-tuning" (Table 2).

Parameter-efficient fine-tuning via LoRA. For MEDUSA-2, the paper uses either LoRA (Hu et al., 2021) or QLoRA (Dettmers et al., 2023) for the backbone updates, rather than full fine-tuning. LoRA approximates weight updates with low-rank matrices: for a weight matrix $\mathbf{W} \in \mathbb{R}^{d \times d}$ , the update is $\Delta\mathbf{W} = \mathbf{B}\mathbf{A}$ where $\mathbf{B} \in \mathbb{R}^{d \times r}$ , $\mathbf{A} \in \mathbb{R}^{r \times d}$ , and $r \ll d$ is the rank. This dramatically reduces the number of trainable parameters (the backbone's full weights remain frozen; only the low-rank adapters are updated). The paper applies LoRA to "all the linear layers of the backbone model, including the language model head" (Appendix B.2), with rank $r = 32$ and scaling factor $\alpha = 16$ . A dropout of 0.05 is added to the LoRA adapter.

Why LoRA? (1) Memory efficiency: only the adapter parameters and MEDUSA head parameters require optimizer states, allowing MEDUSA-2 training on hardware that cannot fit the full model. (2) Implicit regularization: the low-rank constraint prevents the backbone from deviating too far from its pre-trained state, acting as an additional safeguard against catastrophic forgetting. (3) Self-distillation compatibility: as described in Section 2.3.2, the LoRA adapter can be toggled off to recover the original model for teacher predictions, enabling a clean self-distillation setup without maintaining two separate model instances in memory.

Typical Acceptance Scheme

In speculative decoding, the standard verification mechanism is rejection sampling. For each candidate token $x$ proposed by the draft model with probability $q(x)$ , the original model computes its own probability $p(x)$ . The token is accepted with probability $\min(1, p(x)/q(x))$ ; if rejected, a new token is sampled from a modified distribution $\text{norm}(\max(0, p(x) - q(x)))$ . This procedure guarantees that the output sequence is sampled exactly from the original model's distribution $p$ —it is a lossless acceleration method.

However, rejection sampling has a practical problem: as the sampling temperature increases, the acceptance rate drops. Intuitively, at temperature 0 (greedy decoding), both models always output the highest-probability token, so every candidate is accepted (assuming the draft model is well-aligned). At higher temperatures, the distributions spread out, and the draft model's sample is more likely to land in a region where $p(x) < q(x)$ , triggering rejection. The paper notes:

"subsequent implementations (Joao Gante, 2023; Spector & Re, 2023) reveal that this sampling strategy results in diminished efficiency as the sampling temperature increases. Intuitively, this can be comprehended in the extreme instance where the draft model is the same as the original one: Using greedy decoding, all output of the draft model will be accepted, therefore maximizing the efficiency. Conversely, rejection sampling introduces extra overhead, as the draft model and the original model are sampled independently. Even if their distributions align perfectly, the output of the draft model may still be rejected." (Section 2.3.1)

The typical acceptance criterion. Rather than requiring distributional equivalence, MEDUSA proposes a simpler acceptance rule: accept a candidate token if it is "typical"—not implausible under the original model's distribution. The criterion is adapted from Hewitt et al. (2022)'s truncation sampling framework:

A candidate token $x_{n+k}$ at position $n+k$ (with context $x_1, x_2, \ldots, x_{n+k-1}$ ) is accepted if:

$p_{\text{original}}(x_{n+k} \mid x_1, \ldots, x_{n+k-1}) > \min\left( \epsilon, \; \delta \cdot \exp\left( -H(p_{\text{original}}(\cdot \mid x_1, \ldots, x_{n+k-1})) \right) \right)$

where $p_{\text{original}}(\cdot \mid \text{context})$ is the original backbone model's predicted probability distribution at that position, $H(\cdot)$ is the entropy function (measured in nats, since natural log is standard), $\epsilon$ is a hard (absolute) threshold, and $\delta$ is an entropy-dependent threshold multiplier.

What it computes, operationally: For each candidate token proposed by the MEDUSA heads, the backbone's forward pass (the same forward pass used for tree attention verification) produces a probability $p_{\text{original}}(x \mid \text{context})$ —how likely the backbone thinks this token is at this position. Separately, the entropy $H$ of the backbone's distribution at that position is computed: $H = -\sum_{v \in V} p(v) \log p(v)$ , measuring how uncertain the backbone is about what token should appear. The threshold is $\epsilon$ (a floor value) or $\delta \cdot e^{-H}$ (an entropy-adaptive threshold), whichever is smaller. If the candidate token's probability exceeds this threshold, it is accepted.

Why this form: The two-parameter threshold handles two regimes:

Low entropy (confident predictions): When the backbone is very certain ( $H$ is small), $\delta \cdot e^{-H}$ is large (since $e^{-H} \approx 1$ when $H \approx 0$ ). The $\min$ operation selects the hard threshold $\epsilon$ , preventing the acceptance criterion from becoming so lax that low-probability tokens are accepted. The $\epsilon$ parameter provides a floor: no token with probability below $\epsilon$ is ever accepted, regardless of entropy.
High entropy (uncertain predictions): When the backbone is uncertain ( $H$ is large), $\delta \cdot e^{-H}$ becomes small (since $e^{-H}$ decays as entropy increases). The $\min$ operation selects this smaller value, making the acceptance criterion more lenient—the logic being that when there are many reasonable continuations (high entropy), a wider range of tokens should be considered acceptable.

The specific form $\delta \cdot e^{-H}$ is adapted from Hewitt et al. (2022), where it emerges from an information-theoretic analysis of truncation sampling. The exponential-of-negative-entropy term represents the "typical set" probability—tokens in the typical set have log-probability close to $-H$ , meaning their probability is approximately $e^{-H}$ . The $\delta$ parameter scales this threshold.

Whole-candidate and first-token handling. The acceptance check is applied token-by-token along each candidate path. A candidate prefix is accepted up to the first token that fails the criterion. To guarantee at least one token is generated per step, the first token (the LM head's prediction) is unconditionally accepted using greedy decoding: "we apply greedy decoding for the first token and unconditionally accept it while employing typical acceptance for subsequent tokens" (Section 2.3.1). The final accepted sequence for the step is the longest accepted prefix among all candidate branches.

Relationship to temperature. The paper notes that typical acceptance naturally handles sampling temperature: "when the temperature is set to 0, it reverts to greedy decoding, as only the most probable token possesses non-zero probability. As the temperature surpasses 0, the outcome of greedy decoding will consistently be accepted with appropriate $\epsilon, \delta$ , since those tokens have the maximum probability, yielding maximal speedup. Likewise, in general scenarios, an increased temperature will correspondingly result in longer accepted sequences" (Section 2.3.1). This addresses the problem that rejection sampling's efficiency degrades with temperature—typical acceptance's efficiency improves with temperature because the threshold $\delta \cdot e^{-H}$ relaxes as entropy increases.

Empirical behavior (Figure 5). The paper sweeps $\epsilon$ from 0.01 to 0.25 in steps of 0.01 on the writing and roleplay categories of MT-Bench with Vicuna-7B MEDUSA-2. The results show a quality-speed tradeoff: as $\epsilon$ increases, quality scores rise (from ~7.0 to ~7.6 on a GPT-4-evaluated scale) while the acceleration rate decreases (from ~3.5 to ~3.0). The default random sampling (rejection sampling equivalent at temperature 0.7) achieves ~7.5 quality at ~3.5 acceleration rate, while typical sampling with $\epsilon$ around 0.1 achieves similar quality (~7.5) at slightly lower acceleration (~3.3). The tradeoff is controllable: practitioners can tune $\epsilon$ and $\delta$ to match their quality requirements.

Why this scheme over rejection sampling: Rejection sampling guarantees distributional equivalence but at the cost of efficiency—especially at higher temperatures where creative, diverse outputs are desired. The authors argue that "it is typically unnecessary to match the distribution of the original model" (Section 2.3.1) in practice. Users care about output quality, not about whether the sampling procedure is theoretically unbiased. Typical acceptance trades the strict distributional guarantee for higher acceptance rates, particularly at elevated temperatures where the entropy-dependent threshold recognizes that many continuations are reasonable. This aligns with the practical observation that sampling from LLMs in real applications uses temperature as a "creativity" knob, not as a mechanism requiring exact distribution matching.

Self-Distillation: Training Data Generation Without Access to the Original Dataset

When the model's training dataset is unavailable (private data, or models that underwent RLHF post-training), MEDUSA cannot use the standard supervised fine-tuning approach. The self-distillation pipeline (Section 2.3.2) addresses this by using the model itself to generate synthetic training data.

Data generation process. The procedure has two variants:

Multi-turn from seed prompts: Take prompts from a public seed dataset (e.g., ShareGPT or UltraChat) that is "from a domain similar to the target model." Feed each prompt to the model and collect the model's responses. For multi-turn conversations, "sequentially feed the prompts from the seed dataset to the model" (Section 2.3.2)—each prompt from a multi-turn seed conversation is fed in sequence, with the model generating the assistant's response at each turn.
Self-talk (for bidirectional models): For models like Zephyr-7B that are "trained on both roles of the conversation, they have the ability to self-talk, and we can simply feed the first prompt and let the model generate multiple rounds of conversation" (Section 2.3.2). The model generates both user and assistant turns automatically.

The paper collects "a dataset at about 100k samples" for both Vicuna-33B and Zephyr-7B using ShareGPT and UltraChat as seed datasets (Section 3.2).

The distillation loss for MEDUSA-2. For MEDUSA-1, the synthetically generated dataset can be used directly as supervised training data for the MEDUSA heads (the backbone provides hidden states, the generated tokens serve as targets). For MEDUSA-2, however, the backbone itself would be trained on its own outputs. Naively training the backbone on self-generated data typically degrades quality (the paper states: "even without training MEDUSA heads, training the backbone model with this dataset will lead to performance degradation," Section 2.3.2). This is a form of model collapse—the backbone drifts from its original distribution by fitting to its own samples.

To prevent this, the paper uses a knowledge distillation loss:

$\mathcal{L}_{\text{LM-distill}} = \text{KL}\left( \mathbf{p}^{(0)}_{\text{original}, t} \; \| \; \mathbf{p}^{(0)}_t \right)$

where $\mathbf{p}^{(0)}_{\text{original}, t}$ is the original model's predicted probability distribution at position $t$ (the teacher), $\mathbf{p}^{(0)}_t$ is the training model's predicted distribution (the student), and $\text{KL}(\cdot \| \cdot)$ denotes the Kullback-Leibler divergence.

What it computes, operationally: For each position in the self-generated sequence, the original (frozen) model computes its full probability distribution $\mathbf{p}^{(0)}_{\text{original}, t}$ over the vocabulary. The training model (with LoRA adapters and MEDUSA heads) computes its own distribution $\mathbf{p}^{(0)}_t$ . The KL divergence measures how much information is lost when using the training model's distribution to approximate the original's: $\text{KL}(P \| Q) = \sum_v P(v) \log(P(v)/Q(v))$ . Minimizing this divergence encourages the training model to match the original model's distribution at every position, not just at the sampled token. This preserves the full distributional shape, preventing the model from overfitting to the particular tokens that were sampled in the self-generated data.

Why KL divergence rather than cross-entropy with the sampled token: Standard supervised training uses $-\log p^{(0)}_t(y_{t+1})$ where $y_{t+1}$ is the single sampled token. This discards information about the rest of the distribution—the model only learns that $y_{t+1}$ was the correct token, not how probable other tokens should be. Over many self-generated samples, this can cause the model to drift toward a sharper, less well-calibrated distribution (overconfidence on sampled tokens, underconfidence on alternatives). The KL divergence uses the full distribution $\mathbf{p}^{(0)}_{\text{original}, t}$ as a soft target, preserving the original model's uncertainty structure.

Memory-efficient teacher via LoRA. Maintaining two full model instances (the frozen teacher and the training student) would double memory consumption. The paper's solution: "use a parameter-efficient adapter like LoRA for fine-tuning the backbone model. In this way, the original model is simply the model with the adapter turned off. Therefore, the distillation does not require additional memory consumption" (Section 2.3.2). The student model = base model + LoRA adapter (turned on); the teacher model = base model + LoRA adapter (turned off) = original model. Switching between them is a single flag toggle, not a separate model load.

Quantization caveat. The paper warns: "one tip about using self-distillation is that it is preferable to use LoRA without quantization in this case, otherwise, the teacher model will be the quantized model, which may lead to a lower generation quality" (Section 2.3.2). If the backbone is 4-bit quantized, toggling LoRA off yields a 4-bit quantized teacher, whose distributions are degraded by quantization error. The student would then learn to match this degraded distribution, propagating the quality loss. For self-distillation, full-precision LoRA (no quantization) is preferred.

Design choice: why not just use greedy targets from the seed dataset? The self-distillation pipeline is necessary for models where the output distribution has been shaped by processes beyond supervised fine-tuning—specifically RLHF, which optimizes a reward model and fundamentally changes the distribution. The seed dataset (e.g., ShareGPT) contains human-written or GPT-generated responses, which may differ significantly from what the RLHF-tuned model would produce. Training on those mismatched targets could undo the RLHF alignment. Generating data from the model itself preserves the alignment.

Optimized Tree Construction: Greedy Node Selection by Expected Acceptance

The Cartesian product tree construction (Section 2.1.2) is simple but suboptimal: it builds a full, regular tree where every branch at level $k$ has exactly $s_k$ children at level $k+1$ . However, not all branches are equally useful. Some MEDUSA heads are more accurate than others; some top predictions succeed more often than others. An optimized tree would allocate its fixed node budget to branches with the highest probability of correctly predicting future tokens.

The paper proposes a calibration-based greedy algorithm to construct a tree that maximizes expected acceptance length (Section 2.3.3):

Step 1: Calibrate head accuracies. On a calibration dataset (the paper uses Alpaca-eval for this), for each head $k$ and each rank $i$ (where $i = 1$ is the top prediction, $i = 2$ is the second-top, etc.), measure $a^{(i)}_k$ : the empirical probability that the $i$ -th ranked token from head $k$ is the correct token at that future position. The paper defines this as:

"this accuracy is equal to top-i accuracy minus top-(i-1) accuracy" (footnote, Section 2.3.3)

This "marginal" accuracy measures the additional contribution of the $i$ -th prediction beyond the more confident $i-1$ predictions. For example, if the top-1 accuracy is 0.40 and top-2 accuracy is 0.55, then $a^{(1)}_k = 0.40$ (the top-1 token alone is correct 40% of the time) and $a^{(2)}_k = 0.55 - 0.40 = 0.15$ (the second-top token is correct in the remaining 15% of cases).

Step 2: Independence assumption for path accuracy. For a candidate path selecting the $i_j$ -th top prediction from head $j$ (for $j = 1, 2, \ldots, \ell$ ), the estimated probability that the entire path is correct is:

$\prod_{j=1}^{\ell} a^{(i_j)}_j$

This assumes that the outcomes of different heads are independent—a simplifying assumption that allows tractable construction. In reality, head accuracies are likely correlated (a difficult context may cause all heads to fail), but the independence assumption works sufficiently well as a heuristic.

Step 3: Expected acceptance length. For a given tree (a set of nodes, each associated with a specific head's rank prediction), the expected acceptance length is the sum over all leaf nodes of the product of accuracies along the path to that leaf, times the number of tokens represented by that leaf:

$\sum_{[i_1, i_2, \ldots, i_\ell] \in \mathcal{I}} \ell \cdot \prod_{j=1}^{\ell} a^{(i_j)}_j$

where $\mathcal{I}$ is the set of all paths in the tree. More generally, the marginal contribution of adding a new node (extending an existing path by one more head prediction) is exactly the accuracy $a^{(i)}_k$ associated with that node—because the new node adds one more accepted token to the expected length, and this occurs only when all ancestors on the path and the new token are correct.

Step 4: Greedy tree construction. Starting from a tree containing only the root (current context), at each step, consider all possible nodes that could be added—each corresponds to a specific head $k$ and rank $i$ that extends an existing path in the tree. The node with the highest accuracy $a^{(i)}_k$ is added to the tree. This process repeats until the tree reaches the desired total number of nodes (e.g., 64), which is a hyperparameter chosen based on the hardware tradeoff analysis (Figure 4b—speed peaks around 64 nodes for the evaluated hardware).

What it computes, operationally: The calibration accuracies $a^{(i)}_k$ tell us which head-rank combinations are most reliable. The greedy algorithm builds a tree by always adding the most reliable available prediction, regardless of its depth or which branch it extends. The resulting tree is "sparse" and irregular—some branches go deeper (extending predictions further into the future) while others are shallow. Figure 6 in the paper visualizes this: a tree for MEDUSA-2 Vicuna-7B with 64 nodes, depth 4, that "leans towards the left," indicating the algorithm's preference for nodes with higher accuracy.

Why greedy rather than exhaustive search: The number of possible trees with $N$ nodes is combinatorially explosive. Exhaustive optimization over tree structures is intractable. The greedy algorithm is justified by the submodular-like property that the marginal contribution of a new node is simply its accuracy (independent of which other nodes are in the tree, under the independence assumption). This makes greedy selection a natural heuristic that, while not guaranteed to be globally optimal, is computationally efficient and empirically effective.

Design choice: Cartesian product vs. optimized tree. The Cartesian product tree is a special case of the greedy construction where all top- $s_k$ predictions at each level are included (a fully balanced tree). The optimized tree allows unbalanced allocation: if head 3's top-1 prediction is very reliable ( $a^{(1)}_3 = 0.45$ ) but head 2's top-2 prediction is unreliable ( $a^{(2)}_2 = 0.03$ ), the greedy algorithm will include the former and skip the latter, allocating the node budget to high-accuracy predictions regardless of their position in the tree hierarchy. Figure 4a demonstrates the benefit: an optimized sparse tree with 64 nodes achieves a better acceleration rate than a dense Cartesian product tree with 256 nodes, because the 64 nodes are concentrated on the most reliable predictions while the 256-node tree wastes capacity on low-accuracy branches.

Summary of Design Choices and Their Justifications

Single FFN layer with residual connection for MEDUSA heads: Sufficient capacity because the backbone's hidden states already encode rich context; the residual with zero-initialized $\mathbf{W}^{(k)}_1$ ensures heads start at LM head equivalence, avoiding large initial gradients.
Cartesian product tree with top- $s_k$ selection: Simplest structure that combinatorially multiplies candidates; later refined by accuracy-calibrated greedy pruning for optimal node allocation.
Tree attention mask with ancestry-only connections: Enables parallel verification without cross-branch contamination; avoids the $O(N^2)$ cost of full causal attention by exploiting tree sparsity.
MEDUSA-1 frozen backbone: Guarantees no quality degradation (safe deployment); compatible with quantization (democratized training); serves as initialization for MEDUSA-2.
MEDUSA-2 combined loss + differential LR + head warmup: Each component addresses a distinct failure mode: combined loss prevents catastrophic forgetting, differential LR handles convergence rate asymmetry, warmup prevents early gradient interference from untrained heads.
Typical acceptance with entropy-dependent threshold: Addresses rejection sampling's temperature-dependent efficiency drop; accepts tokens that are "reasonable" rather than distributionally guaranteed; provides tunable quality-speed tradeoff via $\epsilon$ and $\delta$ .
Self-distillation with KL divergence loss and LoRA teacher: Generates training data from the model itself when original data is unavailable; KL divergence preserves full distributional shape (not just sampled tokens); LoRA toggle trick avoids memory duplication.
Greedy node selection by calibrated accuracy: Allocates fixed node budget to highest-expected-value predictions; empirically outperforms balanced Cartesian product at same node count; computationally efficient (linear in nodes, not exponential in tree depth).

4. Key Insights and Innovations

Innovation 1: Diagnosing the Draft Model as an Architectural Bottleneck, Not a Necessary Component

MEDUSA's most fundamental intellectual move is reframing what speculative decoding requires. The dominant framing before this paper was that accelerating auto-regressive decoding demands two separate models: a large target model and a smaller, independently trained draft model (Leviathan et al., 2022; Chen et al., 2023). This was treated almost as a logical necessity—you need something fast to generate candidates and something powerful to verify them, and these roles map naturally onto separate parameter sets. Spector & Re (2023) and Miao et al. (2023) accepted this framing and worked within it, proposing staged hierarchies of draft models or ensembles of draft models, each introducing additional complexity.

MEDUSA challenges this framing at the architectural level. The paper's core diagnostic move is asking: should the candidate generator and the verifier actually be separate, or can they share representations? The answer hinges on a specific observation about transformer hidden states: the last hidden state at position t already encodes substantial predictive information about multiple future tokens, not just the immediate next one. If this is true—and MEDUSA's empirical results demonstrate that it is—then a separate draft model is not just unnecessary; it is wasteful. It discards representation-sharing opportunities, introduces distribution shift, and adds serving complexity, all for a function that the original model already has the internal capacity to perform.

This is not merely an engineering simplification. It is a conceptual reframing of the acceleration problem from "coordinate two models" to "extract more information from the one model you already have." The MEDUSA heads are not a draft model replacement in the sense of providing the same function through a different architecture. They are a fundamentally different category of solution: representation exploitation rather than model coordination. The heads consume representations the backbone already computes; they don't introduce a separate inference pipeline. This shift in framing has downstream consequences for training methodology (the two-level strategy exists because representations are shared, not despite it), for deployment simplicity (no second model to manage), and for democratization (training becomes radically cheaper).

The significance extends beyond speculative decoding. The paper implicitly demonstrates that transformer hidden states contain more predictive structure than is utilized by the standard next-token prediction objective. The fact that single-layer FFN heads—trained with modest data—can predict tokens 2–6 positions into the future from a single hidden state h_t suggests that these representations are not position-myopic; they encode a local trajectory through token space, not just a point estimate. This has implications for representation learning research beyond inference acceleration: if hidden states already contain multi-step predictive information, training objectives that exploit this (rather than only optimizing next-token log-likelihood) might produce richer representations. MEDUSA-2's joint training, where the backbone adapts to support multi-token prediction, is a concrete step in this direction.

The evidence for this reframing's validity is the paper's direct comparison to speculative decoding (Table 1): MEDUSA-2 achieves speedups of 2.35–2.83× across model sizes, compared to 1.47–1.60× for speculative decoding with open-source draft models, while maintaining quality. The speedup gap is not small—it is roughly 1.5–1.8× larger—and it comes from the representation-sharing architecture, not from incremental tuning of the draft model approach.

Innovation 2: Identifying and Solving the Quality-Preservation Problem in Integrated Acceleration

Prior work on multi-head parallel decoding (Stern et al., 2018) demonstrated the architectural possibility but sidestepped a critical practical obstacle: how do you add prediction heads to a pre-trained model without degrading its core capability? Stern et al. trained models from scratch or fine-tuned them for specific tasks (machine translation, image super-resolution) where the auxiliary heads' predictions were the primary output. In contrast, MEDUSA operates on general-purpose LLMs that have been expensively pre-trained and often RLHF-aligned, where next-token prediction quality is non-negotiable—users expect the accelerated model to produce outputs indistinguishable from the original.

The paper's diagnostic insight is that this degradation is not hypothetical or minor. Direct fine-tuning of the backbone with MEDUSA heads (without the MEDUSA-2 training recipe) causes a measurable quality drop: Table 2 shows MT-Bench scores falling from 6.17 (baseline) to 5.925 (direct fine-tuning) for Vicuna-7B. This is not a trivial difference—on MT-Bench's 0–10 scale, it's a meaningful regression that users would notice. The degradation mechanism is recognizable from the fine-tuning literature (Kumar et al., 2022, cited by the paper): auxiliary task gradients distort pre-trained features, causing the model to lose capabilities on the original task.

What makes MEDUSA's contribution distinctive is not that it observes this problem—catastrophic forgetting during fine-tuning is well-known—but that it provides a reproducible, three-component recipe that solves it. The combined loss (backbone cross-entropy plus weighted MEDUSA head losses), differential learning rates, and two-stage head warmup form an integrated protocol where each component addresses a specific failure mode. The conceptual contribution is recognizing that these components must work together; applying any subset is insufficient because the failure modes interact. For example, differential learning rates alone cannot prevent early-stage gradient interference from untrained heads; head warmup alone cannot prevent slow drift during extended joint training; combined loss alone cannot handle the convergence-rate asymmetry between the near-optimal backbone and randomly initialized heads.

This is a methodological innovation rather than an architectural one—it is about how to train, not what to train. But it is no less significant for being methodological. The field's default assumption (implicit in speculative decoding's popularity) was that modifying the backbone for acceleration was too risky—better to keep the acceleration mechanism entirely external. MEDUSA-2 demonstrates that integrated acceleration is viable if the training recipe is designed with explicit quality-preservation mechanisms. This opens a design space that was previously considered off-limits: future work can explore tighter integration between prediction heads and backbone, confident that the quality-preservation problem has a known solution strategy.

The evidence is in Table 2: MEDUSA-2 with the full training recipe achieves 6.18 on MT-Bench versus 6.17 for the baseline, a negligible difference, while providing 2.83× speedup. The preservation is not perfect (the paper reports -0.07 and -0.14 quality deltas for Zephyr-7B and Vicuna-13B in Table 1), but it is dramatically better than the 5.925 from direct fine-tuning and adequate for practical deployment.

Innovation 3: Typical Acceptance as a Practical Reframing of the Quality-Speed Tradeoff

Rejection sampling in speculative decoding guarantees distributional equivalence with the original model—a theoretically elegant property that makes the acceleration "lossless." However, the paper identifies a practical crack in this guarantee: as sampling temperature increases, rejection sampling's efficiency drops because the draft model's samples increasingly diverge from the target model's distribution, triggering more rejections (Section 2.3.1). This creates a tension: creative, diverse generation (which users want) is exactly the regime where speculative decoding provides the least speedup.

MEDUSA's typical acceptance scheme reframes the problem. Rather than asking "does this token come from the exact same distribution as the original model?" it asks "is this token plausible under the original model?" This is a shift from distributional matching to output acceptability—a pragmatic move that recognizes most users don't care about theoretical unbiasedness; they care about getting good outputs quickly. The insight connects to the truncation sampling literature (Hewitt et al., 2022; Meister et al., 2023), which established that generation quality often improves when you restrict sampling to a "typical set" of reasonable tokens rather than sampling from the full distribution. MEDUSA inverts this logic: instead of using typicality to improve quality (by filtering low-probability tokens during generation), it uses typicality to improve speed (by accepting more candidate tokens during verification).

The entropy-dependent threshold (δ·exp(-H)) is the conceptual key. It encodes an intuition that is absent from rejection sampling: when the model is uncertain (high entropy), many continuations are reasonable, so the acceptance criterion should relax. When the model is confident (low entropy), only high-probability continuations should be accepted. This threshold adapts automatically to context—it doesn't require per-task or per-temperature tuning of acceptance parameters. The ε hard threshold provides a floor that prevents acceptance of vanishingly improbable tokens even in high-entropy contexts.

The innovation is not the threshold formula itself (which is adapted from Hewitt et al., 2022) but the recognition that this formula can replace rejection sampling in speculative decoding pipelines, trading the lossless guarantee for higher throughput at the temperatures where throughput matters most. The paper's ablation (Figure 5) demonstrates the controllability: by sweeping ε, practitioners can choose their operating point on the quality-speed curve, from near-greedy (low ε, high speed, slightly reduced creativity) to near-rejection-sampling (high ε, lower speed, distributional fidelity). This tunability makes typical acceptance a practical tool, not just an interesting alternative.

The broader significance is in challenging the assumption that "lossless" is the right goal for inference acceleration. In many deployment scenarios, users want good enough outputs fast, not exactly-the-distribution outputs slowly. The paper doesn't argue that lossless guarantees are worthless—rejection sampling remains available as an option within MEDUSA—but it demonstrates that relaxing the guarantee can yield substantial practical gains without meaningful quality degradation. This is a pragmatic contribution that may influence how the field evaluates acceleration methods: acceptance rate and wall-clock speed under realistic generation settings (with temperature > 0) are more informative than theoretical distributional equivalence.

Innovation 4: Self-Distillation with LoRA Toggle as a Deployment Enabler

Speculative decoding assumes access to a suitable draft model, which in practice means either a pre-trained small model from the same family (if one exists) or the resources to pre-train one. MEDUSA's training strategies assume access to a dataset matching the target model's output distribution—an assumption that breaks for many real-world models. RLHF-tuned models (like Zephyr-7B), models trained on private data (like Vicuna-33B's experimental training set), and many proprietary API models have output distributions that differ from any publicly available training corpus. Training MEDUSA heads on mismatched data would produce heads that predict tokens the target model wouldn't generate, reducing acceptance rates and speedup.

The paper's self-distillation pipeline (Section 2.3.2) solves this by generating training data from the model itself. This is not novel as a general technique—model-generated data for knowledge distillation is well-established (Kim & Rush, 2016). What's distinctive is the specific memory-efficient teacher-student architecture via LoRA toggle. The insight: if you fine-tune the backbone with a LoRA adapter, the "teacher" (original model) is just the student with the adapter disabled. This means self-distillation with KL divergence loss requires zero additional GPU memory beyond what's already allocated for training—no second model instance, no separate inference pipeline, no offloading.

The KL divergence loss (rather than cross-entropy with the sampled token) is also a deliberate choice that addresses a subtle failure mode. Training on self-generated data with standard token-level cross-entropy causes the model to overfit to the particular tokens it happened to sample, losing calibration on the full distribution. KL divergence with the original model's full distribution preserves the shape of the original model's uncertainty, preventing the distributional collapse that the paper observes with naive self-training ("even without training MEDUSA heads, training the backbone model with this dataset will lead to performance degradation").

This is a practical deployment innovation, not a theoretical one. It doesn't change the acceleration mechanism; it removes a barrier to adoption. The paper demonstrates its effectiveness on Vicuna-33B (private training data) and Zephyr-7B (RLHF-tuned), achieving 2.35× and 2.66× speedups respectively (Table 1, Figure 8), with quality deltas of +0.05 and -0.07 on MT-Bench. These results matter because they show MEDUSA is viable for models where the original training recipe is unavailable—which is the common case for practitioners using open-weight models or API-accessed models. This bridges the gap between the paper's controlled experiments (Vicuna-7B/13B with known training data) and the messy reality of deployed model ecosystems.

5. Experimental Analysis

Evaluation Methodology

Dataset. All experiments use MT-Bench (Zheng et al., 2023), a multi-turn, conversational-format benchmark with 80 questions across 8 categories (Writing, Roleplay, Reasoning, Math, Coding, Extraction, STEM, Humanities). The paper also reports supplementary results on AlpacaEval (Li et al., 2023) in Appendix F (Table 4). The evaluation judges model responses using GPT-4 as an evaluator, assigning scores from 0–10, with quality comparisons reported as score deltas relative to the baseline model.
Base model(s). The primary models are Vicuna-7B, Vicuna-13B, and Vicuna-33B (Chiang et al., 2023)—chat models fine-tuned from Llama (Touvron et al., 2023) on the ShareGPT dataset, with the 33B variant using a private dataset—plus Zephyr-7B (Tunstall et al., 2023), an RLHF-aligned model distilled from Llama-2. These span 7B–33B parameters and cover different training recipes (public SFT, private SFT, RLHF alignment), allowing the paper to test MEDUSA's generality across model scales and training paradigms.
Metrics. The paper reports three interrelated metrics (defined in Appendix B.1): (1) Acceleration rate—the average number of tokens decoded per decoding step (baseline auto-regressive = 1.0); (2) Overhead—the ratio of average per-step latency of MEDUSA to vanilla decoding, measuring how much slower each MEDUSA step is; (3) Speedup—the wall-clock acceleration, computed as Speedup = Acceleration rate / Overhead. Generation quality is measured via GPT-4-evaluated scores on MT-Bench (0–10 scale), with quality differences (+0.01, -0.07, etc.) indicating how MEDUSA models compare to their original counterparts. Wall-time speed (tokens per second) is also reported directly.
Baselines. The paper compares against: (1) Vanilla auto-regressive decoding—the default HuggingFace implementation with no acceleration; (2) Speculative decoding (Leviathan et al., 2022; Chen et al., 2023)—implemented using open-source draft models including Llama-68M, Llama-160M (Miao et al., 2023), Tiny-Llama (Zhang et al., 2024), and Tiny-Vicuna (Pan, 2023), evaluated across different draft lengths γ (detailed in Appendix D, Figure 7); (3) Direct fine-tuning—training MEDUSA heads alongside the backbone without the MEDUSA-2 training recipe (combined loss, differential learning rates, warmup), used as an ablation to isolate the recipe's contribution (Table 2).
Generation budget / compute accounting. The primary unit of comparison is wall-clock time (tokens per second), not FLOPs or parameter counts. This is appropriate for a memory-bandwidth-bound regime where arithmetic utilization is not the bottleneck. For speculative decoding comparisons, draft length γ is swept across {1, 2, ..., 14} to find the optimal configuration per model size (Figure 7). For tree attention, the number of candidate tokens (tree nodes) is varied from 1 to ~256 in ablation studies (Figure 4), with the optimal configuration at 64 nodes for the main experiments. Training compute is measured in GPU hours: MEDUSA-1 on Vicuna-7B takes "5 hours... with a single NVIDIA A100 PCIE GPU to train on 60k ShareGPT samples" (Section 2.2.1).
Cross-validation / statistical protocol. The paper does not employ formal cross-validation. Strategy selection (e.g., optimal tree configuration, typical acceptance thresholds) is performed on a per-category subset of MT-Bench (specifically the writing and roleplay categories for ablation studies in Sections 3.3.1 and 3.3.2). The main speedup results are then reported across all 8 MT-Bench categories. Training uses standard practices: cosine learning rate schedule with warmup, 8-bit AdamW optimizer, and 2-epoch training on the ShareGPT dataset for Vicuna models (Appendix B.3).

Main Quantitative Results

MEDUSA-1 vs. MEDUSA-2 on Vicuna-7B and 13B (Section 3.1, Figure 3)

The headline result for MEDUSA-1: 2.18× speedup on Vicuna-7B and 2.33× on Vicuna-13B over the vanilla HuggingFace baseline, with zero quality degradation (MEDUSA-1 keeps the backbone frozen, so quality is trivially preserved). For MEDUSA-2: 2.83× speedup on both Vicuna-7B and 13B, representing an additional ~30% improvement over MEDUSA-1 from joint training.

Figure 3a presents the raw throughput: Vicuna-7B baseline achieves approximately 39 tokens/second; MEDUSA-1 raises this to approximately 85 tokens/second; MEDUSA-2 reaches approximately 110 tokens/second. For Vicuna-13B: baseline approximately 28 tokens/second, MEDUSA-1 approximately 65 tokens/second, MEDUSA-2 approximately 79 tokens/second. The absolute throughput decreases with model size (larger models = more parameters to move per forward pass), but the relative speedup is consistent or slightly better for the larger model, confirming that MEDUSA's benefit scales with the memory-bandwidth bottleneck severity.

The per-category breakdown for MEDUSA-2 Vicuna-7B (Figure 3b) reveals substantial variation: Coding achieves 3.29× speedup, Extraction achieves 3.62×, while Reasoning (2.58×) and Humanities (2.58×) are lower. The paper interprets this variation optimistically—"this points to a significant potential for optimizing coding LLMs"—but does not deeply analyze why these categories differ. Likely explanations (not stated in the paper): coding tasks have more predictable syntactic structure (easier for MEDUSA heads to guess future tokens), while reasoning tasks require more context-dependent token choices.

Table 1 (Section 3.2) provides the component metrics behind these speedups. For Vicuna-7B MEDUSA-2: acceleration rate = 3.47 (average tokens decoded per step), overhead = 1.22 (each step is 22% slower than vanilla), giving speedup = 3.47 / 1.22 = 2.84× (rounding accounts for the 2.83× reported). For Vicuna-13B: acceleration rate = 3.51, overhead = 1.23, speedup = 2.85×. The acceleration rates of ~3.5 tokens per step mean that on average, 2.5 of the 5 MEDUSA heads are producing correct predictions—a substantial fraction of the theoretical maximum.

Self-Distillation Results: Vicuna-33B and Zephyr-7B (Section 3.2, Table 1, Figure 8)

For models where the original training data is unavailable, self-distillation enables MEDUSA-2 training with quality preservation.

Vicuna-33B: MEDUSA-2 achieves 2.35× speedup with acceleration rate = 3.01 and overhead = 1.27 (Table 1). Quality on MT-Bench is 7.18 (+0.05 vs. original), indicating self-distillation successfully preserved the model's output distribution despite the private training data mismatch. The speedup is lower than for 7B/13B models (2.35× vs. 2.83×), which the authors hypothesize is "due to a mismatch between the hidden training dataset and the dataset we used for self-distillation" (Section 3.2). The self-distillation dataset (100k samples from ShareGPT and UltraChat seeds) may not fully capture the 33B model's private training distribution, making the MEDUSA heads less accurate.

Zephyr-7B: MEDUSA-2 achieves 2.66× speedup with acceleration rate = 3.14 and overhead = 1.18 (Table 1). Quality on MT-Bench is 7.25 (-0.07 vs. original), a small but measurable degradation. This is notable because Zephyr-7B is RLHF-aligned—its output distribution is shaped by a reward model, not just supervised fine-tuning. The self-distillation pipeline successfully generates training data that preserves this alignment (the -0.07 drop is small), but the quality is not perfectly maintained.

Figure 8 (bar chart) visualizes these results alongside Vicuna-7B and 13B: the self-distillation-trained models (Zephyr-7B, Vicuna-33B) show consistently lower speedup than the directly-trained models (Vicuna-7B, 13B), with the gap being ~0.5× (2.83× → 2.35–2.66×). The paper is transparent about this tradeoff: "models trained with self-distillation... have weaker speedup due to the trade-off between preserving quality and boosting speed" (Figure 8 caption). The mechanism: to prevent quality degradation during self-distillation, the training uses a small λ₀ = 0.01 (Appendix B.4), which limits how much the backbone adapts to help the MEDUSA heads, capping the achievable acceleration rate.

Comparison with Speculative Decoding (Table 1, Figure 7 in Appendix D)

Table 1's bottom section directly compares speedup: speculative decoding achieves 1.47× (Vicuna-7B), 1.56× (Vicuna-13B), and 1.60× (Vicuna-33B), while MEDUSA-2 achieves 2.83×, 2.83×, and 2.35× respectively. The gap is substantial: MEDUSA provides 1.5–1.8× more speedup than speculative decoding on the same models.

However, this comparison requires nuance. The speculative decoding results in Figure 7 (Appendix D) show that optimal draft model selection matters significantly. For Vicuna-7B, Llama-68M with γ = 4 achieves ~55 tokens/second vs. ~42 tokens/second baseline—a 1.31× speedup, not the 1.47× reported. The reported 1.47× likely comes from the best configuration across all draft models and γ values, but even this best-case underperforms MEDUSA-2's 2.83× by a factor of 1.9×. For Vicuna-33B, Tiny-Vicuna with γ = 3 achieves ~27 tokens/second vs. ~18 tokens/second baseline (1.50×), while MEDUSA-2 achieves 2.35×. The pattern is consistent: speculative decoding is fundamentally limited by the quality of available draft models, while MEDUSA's integrated heads benefit from shared representations.

The paper also notes (Appendix D) that speculative decoding's optimal γ varies with model size (γ = 4 for 7B, γ = 3 for 13B and 33B), requiring per-model hyperparameter tuning that MEDUSA avoids through its optimized tree construction procedure.

Speedup on AlpacaEval (Appendix F, Table 4)

To validate generalizability beyond MT-Bench, the paper reports AlpacaEval results. MEDUSA-2 achieves: Vicuna-7B: 2.88× speedup (acc. rate 3.23, base speed 37.07 → MEDUSA speed 106.76 tokens/s); Vicuna-13B: 3.16× speedup (acc. rate 3.28, 29.01 → 91.54 tokens/s); Vicuna-33B: 2.26× speedup; Zephyr-7B: 2.91× speedup. These are broadly consistent with MT-Bench results, confirming that the speedup is not benchmark-specific. The Vicuna-33B result (2.26×) is notably lower than on MT-Bench (2.35×), suggesting some variance across evaluation datasets, but the overall pattern holds.

Ablation Studies and Robustness Checks

Tree attention configuration (Section 3.3.1, Figure 4): Sparse optimized trees significantly outperform dense Cartesian product trees at equivalent node counts. A sparse tree with 64 nodes achieves a better acceleration rate than a dense tree with 256 nodes (Figure 4a, comparing red stars to blue dots at those x-axis positions). However, speed (tokens/second) peaks around 64 candidate nodes and declines thereafter (Figure 4b)—the trend line shows acceleration rate growing logarithmically while computational overhead grows superlinearly. The paper attributes the speed decline to "increased overhead introduced by the compute-bound" (Section 3.3.1), specifically the attention matrix multiplication QKᵀ which scales quadratically with candidate count. This non-monotonic relationship between tree size and speed is a critical practical finding: larger trees increase the expected acceptance length but past ~64 nodes, the additional matrix multiplication cost outweighs the marginal acceptance gains.

Typical acceptance thresholds (Section 3.3.2, Figure 5): Sweeping ε from 0.01 to 0.25 (in steps of 0.01) on writing and roleplay categories with temperature 0.7 reveals a smooth quality-speed tradeoff. At ε = 0.01 (strict threshold), the quality score is ~7.0 with acceleration rate ~3.5. At ε = 0.25 (lenient threshold), quality rises to ~7.6 while acceleration rate drops to ~3.0. The default rejection sampling (RS) achieves ~7.5 quality at ~3.5 acceleration rate—higher quality with better speed than the comparable ε setting. However, the typical acceptance curve crosses near the RS operating point: at ε ≈ 0.1, quality is ~7.5 with acceleration ~3.3, close to RS's 7.5 at 3.5. The paper frames this as "comparable with random sampling when ε increases" (Section 3.3.2), though the data shows a consistent small speed penalty for typical acceptance at matched quality levels. Greedy decoding achieves ~7.0–7.1 quality at ~3.5 acceleration rate, confirming that for creative tasks (writing, roleplay), greedy outputs are lower quality despite the speed advantage.

Two-stage fine-tuning effectiveness (Section 3.3.3, Table 2): This ablation is the paper's most important methodological validation. Four conditions are compared on Vicuna-7B:

Condition	Quality (MT-Bench)	Speedup
Baseline (original Vicuna-7B)	6.17	N/A
Direct fine-tuning (backbone + heads, no special recipe)	5.925	N/A
MEDUSA-1 (frozen backbone, heads only)	6.23	2.18×
MEDUSA-2 (full recipe)	6.18	2.83×

Direct fine-tuning causes a quality drop of 0.245 points (6.17 → 5.925), confirming that naive joint training degrades the model's core capability. MEDUSA-1 not only preserves quality but slightly improves it (6.17 → 6.23)—likely because the frozen backbone provides stable features, and the head training doesn't affect the backbone at all. MEDUSA-2 matches baseline quality (6.18 vs. 6.17) while providing 30% more speedup than MEDUSA-1. This ablation isolates the contribution of the training recipe: without it, integrated acceleration fails; with it, the backbone's capability is maintained while the heads achieve higher accuracy through representation co-adaptation.

Impact of individual techniques on speedup (Table 3): The paper provides a cumulative breakdown starting from a MEDUSA-1 baseline without tree attention (~1.5× speedup). Adding tree attention raises speedup to ~1.9×, using optimized tree configuration (sparse tree with calibrated node selection) raises it to ~2.2×, and training heads with MEDUSA-2 (joint backbone-head training) achieves the final ~2.8×. Each technique contributes additively: tree attention adds ~0.4×, optimization adds ~0.3×, joint training adds ~0.6×. The largest single contribution comes from MEDUSA-2's joint training, underscoring that head accuracy (not just candidate quantity) is the dominant factor.

Hardware scaling and roofline analysis (Appendix G, Figures 9–23): The paper provides extensive hardware profiling, though primarily as analysis rather than controlled ablation. Key findings: (1) at batch size 1, linear layers are memory-bandwidth-bound with operational intensity ~1–2 FLOP/byte, while attention matrix multiplications are even more severely bound (~0.5–1 FLOP/byte); (2) MEDUSA's additional candidate tokens shift linear layers toward compute-bound regime—processing 64 candidate tokens at batch size 1 increases up/gate/down linear layer throughput from 1.26 TFLOP/s to 76.57 TFLOP/s, a 61× improvement in FLOP utilization (Table 8); (3) the simulated speedup curves (Figures 21–23) predict optimal candidate counts of ~64 for batch size 1, with speedup declining at larger batch sizes (>32) as linear layers transition fully to compute-bound and overhead dominates. The simulation in Figure 22 shows speedup dropping below 1.0× at batch size 64 with 112 candidates—MEDUSA would actually slow down inference in high-throughput batched settings.

Speculative decoding with different draft models (Appendix D, Figure 7): While not technically a MEDUSA ablation, this analysis characterizes the competing approach. The optimal γ varies across model sizes (γ = 4 for Vicuna-7B with Llama-68M, γ = 3 for Vicuna-13B and 33B), and the best draft model also varies (Llama-68M for 7B and 13B, Tiny-Vicuna for 33B). Even at optimal configurations, speculative decoding's speedup (1.47–1.60×) is substantially below MEDUSA's. The draft model's throughput is itself a bottleneck: Tiny-Vicuna-1B generates tokens at ~40 tokens/second (inferred from Figure 7c where the baseline is ~18 tokens/s and best speculative is ~27 tokens/s), meaning the draft model's own inference time limits how many candidates can be generated per target model forward pass.

Critical Assessment

Claim 1: MEDUSA-1 achieves over 2.2× speedup without compromising generation quality.

Supported with an important boundary condition. The 2.18× speedup for Vicuna-7B and 2.33× for Vicuna-13B (Figure 3a) are directly measured wall-clock improvements. Quality preservation for MEDUSA-1 is definitional—the backbone is frozen, so the model's next-token predictions are identical to the original. However, this "lossless" guarantee applies only to the model outputs, not to the generation process. MEDUSA-1 with typical acceptance may produce different output sequences than the original model (even with identical model weights) because typical acceptance does not guarantee distributional equivalence. The paper acknowledges this implicitly by providing typical acceptance as an option, but the "without compromising generation quality" claim should be qualified: quality is preserved in the sense that MT-Bench scores don't degrade (Table 2 shows 6.23 for MEDUSA-1 vs. 6.17 baseline—a slight improvement, likely within noise), but the output distribution may differ. Users who require exact distributional matching must use rejection sampling, which may reduce speedup relative to the reported 2.18× (the paper does not report MEDUSA-1 speedup specifically with rejection sampling).

Claim 2: MEDUSA-2 jointly trains the heads with the backbone to reach 2.3–2.8× speedup.

Strongly supported, with quality caveats. The 2.83× speedup on Vicuna-7B/13B (Figure 3a) is the strongest result. Table 1 provides the full range: 2.35× (Vicuna-33B) to 2.83× (Vicuna-7B/13B), with Zephyr-7B at 2.66×. Quality preservation is good but not perfect: Vicuna-7B shows +0.01 (effectively unchanged), Zephyr-7B shows -0.07 (small regression), Vicuna-13B shows -0.14 (modest regression), Vicuna-33B shows +0.05 (noise-level improvement). The -0.14 drop for Vicuna-13B is the largest quality degradation and warrants scrutiny—it represents approximately 2.2% of the 6.43 baseline score on MT-Bench's 0–10 scale. Whether this is "acceptable" depends on the deployment context; the paper does not discuss it as a concern. The quality deltas are small enough to be within GPT-4 evaluator variance (the paper does not report evaluator confidence intervals), but the consistent negative deltas for two of four models suggest a systematic, if minor, quality cost to MEDUSA-2.

Claim 3: MEDUSA establishes that multi-head parallel prediction can substitute for a draft model in speculative decoding only when the training recipe explicitly preserves the backbone's next-token capability through combined loss, differential learning rates, and head warmup.

Strongly supported. Table 2 provides the direct evidence: direct fine-tuning degrades quality (6.17 → 5.925), while MEDUSA-2 preserves it (6.17 → 6.18). The ablation isolates the training recipe as the causal factor—all conditions use the same architecture (5 MEDUSA heads, single FFN layer each); only the training procedure differs. The three-component recipe (combined loss, differential LR, warmup) is justified individually through ablation-like reasoning in Section 2.2.2, though the paper does not ablate which of the three components is most critical. It's possible that, for example, combined loss alone suffices and the other components provide marginal improvements; the paper cannot distinguish this because it never tests partial recipe configurations. This is a methodological gap: the claim that all three components are necessary is asserted but not experimentally proven.

Genuine weaknesses and missing experiments:

Single evaluation dataset. All quality evaluations are on MT-Bench (80 questions, 8 categories). AlpacaEval results (Appendix F) report only speed, not quality deltas versus baseline. This is a thin evaluation for a method that claims to preserve generation quality—a broader set of benchmarks (MMLU, HellaSwag, HumanEval for coding-specific quality, etc.) would strengthen confidence that no domain-specific degradation occurs.
MT-Bench evaluator variance is unaddressed. GPT-4-as-judge produces scores with non-trivial variance. The quality deltas of ±0.01–0.14 reported in Table 1 may be within the noise floor of the evaluator. Without confidence intervals or multiple evaluation runs, it's impossible to determine whether Vicuna-13B's -0.14 quality drop is a real degradation or sampling noise. This is particularly problematic because the paper's central narrative depends on "quality preservation"—small degradations that go undetected due to evaluator variance would undermine the claim.
No ablation of individual MEDUSA-2 components. The three-part training recipe (combined loss, differential LR, warmup) is presented as an integrated whole. The paper never tests whether combined loss alone (without warmup), or warmup alone (without differential LR), achieves similar quality preservation. This makes it impossible to attribute the recipe's success to specific mechanisms or to simplify the training procedure. The claim that "each component addresses a distinct failure mode" is architectural reasoning, not experimental fact.
Self-distillation quality is tested on only two models. The self-distillation pipeline is validated on Vicuna-33B (private SFT data) and Zephyr-7B (RLHF). These cover two failure modes (private data, alignment shift), but the sample size is small. The -0.07 quality drop for Zephyr-7B and -0.14 for Vicuna-13B suggest that quality preservation is harder when training data is synthetic, but the paper doesn't systematically explore this (e.g., by comparing self-distillation vs. original-data training on the same model).
Batch size 1 only. All experiments use batch size 1. The paper's own hardware analysis (Appendix G, Figure 22) shows that simulated speedup declines with batch size and becomes negative at batch size 64. This means MEDUSA's benefits may be limited to interactive, single-query deployments and may not transfer to high-throughput batched inference. The paper acknowledges this in the Discussion ("we want to emphasize that the ideas presented in our paper can be generalized to larger batch-size settings, which are now supported by libraries like TensorRT and Huggingface TGI following our paper"), but provides no experimental evidence for batched settings.
Speculative decoding baseline may be suboptimal. The paper's speculative decoding results use open-source draft models (Llama-68M, Tiny-Llama, etc.) that were not specifically designed as draft models for Vicuna. Speculative decoding papers (Leviathan et al., 2022; Chen et al., 2023) typically use draft models from the same model family, fine-tuned on similar data. It's possible that a purpose-built draft model (e.g., a Vicuna-1B distilled from Vicuna-7B) would achieve higher acceptance rates and narrow the speedup gap. The paper's speculative decoding comparison is therefore against a practical baseline (what an open-source user could deploy) rather than an optimal baseline (what speculative decoding can achieve with bespoke draft model training). The paper could have strengthened this comparison by training a Vicuna-specific draft model, controlling for the training data investment.
No latency analysis for the first token. MEDUSA's tree attention processes all candidates in a single forward pass, but the forward pass latency increases with candidate count. For interactive applications, time-to-first-token (TTFT) is a critical metric that is not reported. The paper reports average tokens-per-second over full generations, which could mask higher latency on initial steps where the prefix is short and the tree attention overhead is proportionally larger.
Limited exploration of tree size vs. sequence length interaction. The hardware analysis (Appendix G) shows that sequence length affects attention matrix multiplication overhead non-trivially (Table 6: at seq_len 128, 112 candidates achieves 36.57 TFLOP/s; at seq_len 8192, the same 112 candidates achieves 84.5 TFLOP/s—more than double). This means the optimal tree size likely depends on current context length, which the paper does not address. A fixed tree size of 64 nodes may be suboptimal for very short or very long contexts.

6. Limitations and Trade-offs

Batch Size 1 Only: MEDUSA's Benefits May Not Transfer to High-Throughput Inference

The assumption or constraint. The paper's entire experimental evaluation operates at batch size 1, the regime "representative of the use case where LLMs are locally hosted for personal use" (Section 1). The authors explicitly acknowledge this scope restriction in the Discussion: "In the paper, we focus on the setting with batch size 1 for simplicity. Yet, we want to emphasize that the ideas presented in our paper can be generalized to larger batch-size settings" (Section 4). However, no experiments support this generalization claim.

The consequence. The mechanism that produces MEDUSA's speedup—amortizing weight transfers over multiple candidate tokens—relies on the linear layers being memory-bandwidth-bound. At larger batch sizes, these same layers shift toward the compute-bound regime because more tokens per batch mean more FLOPs per weight load. The paper's own hardware simulation (Appendix G, Figure 22) demonstrates this failure mode concretely: at sequence length 1024 with batch size 64, simulated speedup actually drops below 1.0× when processing 112 candidate tokens—meaning MEDUSA would slow down inference relative to vanilla auto-regressive decoding. Even at batch size 32, the speedup curve shows substantial degradation relative to batch size 1. This matters for production deployments where batching is the primary mechanism for achieving high throughput: serving systems handling many concurrent users typically aggregate queries into batches precisely to escape the memory-bandwidth bottleneck that MEDUSA targets. In that regime, MEDUSA's additional computation (processing candidate tokens that may be rejected) becomes pure overhead rather than an amortization benefit. The consequence is that MEDUSA's headline 2.3–2.8× speedup is likely confined to interactive, single-query settings (chatbots, coding assistants, local LLM usage) and does not apply to the high-throughput batched inference that dominates cloud API deployments.

What evidence exists in the paper. The roofline analysis in Appendix G.2 (Figures 18–20, Tables 6–8) shows that linear layers shift from operational intensity ~1 FLOP/byte (memory-bandwidth-bound) at batch size 1 to operational intensity ~63 FLOP/byte at batch size 16 and 64 candidates (Table 8). The simulated speedup curves in Figure 22 directly model this: at batch size 1 with 112 candidates, speedup is ~2.7×; at batch size 32, it drops to ~1.3×; at batch size 64, it falls to ~0.9× (below baseline). Figure 23 additionally shows that sequence length compounds this effect—longer sequences increase attention matrix multiplication overhead, further reducing speedup at any given batch size. The paper does not report any end-to-end batched inference experiments to validate or refute the simulation predictions.

Mitigation status. The authors acknowledge the limitation implicitly through the batch-size-1 scope and the simulation results, and they gesture toward future generalization ("can be generalized to larger batch-size settings, which are now supported by libraries like TensorRT and Huggingface TGI following our paper," Section 4). However, this is an aspiration, not a demonstrated capability. The simulation results suggest that MEDUSA's core mechanism is fundamentally at odds with large-batch inference, and no mitigation strategy (e.g., dynamic candidate count reduction at larger batch sizes, selective head deactivation) is proposed or evaluated. A practitioner deploying MEDUSA in a production serving system with batching would need to implement their own throttling logic to disable MEDUSA at high batch sizes, which is not described in the paper.

Self-Distillation Introduces a Quality-Speed Tradeoff That Limits MEDUSA-2's Applicability to Proprietary or RLHF-Tuned Models

The assumption or constraint. MEDUSA-1 and MEDUSA-2 both assume "the availability of a training dataset that aligns with the target model's output distribution" (Section 2.2). For models where this assumption breaks—those trained on private data (Vicuna-33B) or with RLHF post-training (Zephyr-7B)—the self-distillation pipeline (Section 2.3.2) generates synthetic training data by having the model respond to seed prompts. The paper is transparent that this introduces a tradeoff: "models trained with self-distillation... have weaker speedup due to the trade-off between preserving quality and boosting speed" (Figure 8 caption).

The consequence. Self-distillation does not simply work less well—it introduces a specific coupling between quality preservation and speedup that constrains MEDUSA-2's utility for the very models that most need it (those without public training data). The mechanism: to prevent the backbone from degrading when trained on self-generated data, the paper uses a KL-divergence distillation loss (matching the original model's full distribution) with a very small MEDUSA auxiliary loss weight (\lambda_0 = 0.01 in Appendix B.4, compared to \lambda_0 = 0.2 for directly-trained models in Appendix B.3). This small \lambda_0 limits how much the backbone adapts its representations to help the MEDUSA heads, capping the achievable acceleration rate. The consequence in practice: self-distillation-trained MEDUSA-2 models show speedups of 2.35–2.66× (Table 1) versus 2.83× for directly-trained models—a 6–17% speedup penalty. More subtly, the quality is not perfectly preserved: Zephyr-7B shows a -0.07 MT-Bench score drop and Vicuna-13B (which, though trained with public data, uses the self-distillation variant in its MEDUSA-2 training per Appendix B.3's description) shows a -0.14 drop. These quality deltas are small, but they matter: for RLHF-aligned models where users specifically value the alignment properties (helpfulness, harmlessness, instruction-following), even modest regression could be unacceptable. The paper does not evaluate whether the specific capabilities that RLHF imparts (refusal of harmful requests, adherence to stylistic preferences) are preserved under self-distillation.

A further complication: the self-distillation pipeline generates 100k samples using seed prompts from ShareGPT and UltraChat (Section 3.2). If the target model's training distribution differs significantly from these seed domains, the generated data may not cover the full output distribution, creating blind spots where MEDUSA heads are poorly calibrated. The paper hypothesizes that Vicuna-33B's lower speedup (2.35×) is "due to a mismatch between the hidden training dataset and the dataset we used for self-distillation" (Section 3.2), acknowledging this coverage issue without resolving it.

What evidence exists in the paper. Table 1 provides the direct comparison: directly-trained Vicuna-7B achieves 2.83× speedup with +0.01 quality delta, while self-distillation-trained Vicuna-33B achieves 2.35× with +0.05 quality delta and Zephyr-7B achieves 2.66× with -0.07 quality delta. Figure 8 visualizes the speedup gap between self-distillation and direct-training models. The paper's own ablation of the self-distillation approach reveals the quality degradation risk: "even without training MEDUSA heads, training the backbone model with this dataset will lead to performance degradation" (Section 2.3.2). The KL divergence loss and small \lambda_0 are presented as mitigation, but their effectiveness is only demonstrated on two models (one with private data, one with RLHF), neither of which has publicly available ground-truth training data for controlled comparison.

Mitigation status. The paper acknowledges the tradeoff explicitly (Figure 8 caption) and provides the KL-divergence distillation loss as the mechanism for managing it. However, this is a partial mitigation: it prevents catastrophic quality collapse but does not eliminate the speedup penalty or the small residual quality regression. The paper does not explore alternative approaches that could improve self-distillation, such as: using a larger or more diverse seed dataset to improve coverage; iterative self-distillation where the MEDUSA model generates new training data for itself; or mixing self-distillation data with publicly available data when partial overlap with the training distribution is known. The recommendation to "use LoRA without quantization" for self-distillation (Section 2.3.2) improves teacher quality but does not address the fundamental coverage-vs-adaptation tension.

Quality Evaluation Is Insufficient to Validate the "Lossless" and "Quality-Preserving" Claims

The assumption or constraint. The paper's central value proposition is that MEDUSA accelerates inference "without compromising generation quality" (Section 1, reinforced in Table 2 and throughout). Quality is evaluated exclusively through MT-Bench (80 multi-turn questions across 8 categories, scored by GPT-4 on a 0–10 scale), with supplementary speed-only results on AlpacaEval (Appendix F, Table 4, which reports no quality deltas). The paper states quality differences as point estimates (e.g., "+0.01", "-0.07", "-0.14") without confidence intervals, statistical tests, or evaluator variance analysis.

The consequence. This leaves a fundamental question unanswered: are the reported quality deltas real degradations or sampling noise? MT-Bench GPT-4 evaluation has known variance. Zheng et al. (2023), the paper that introduced MT-Bench, reports that GPT-4 evaluations have a "moderate agreement" with human judgments and that re-prompting can produce different scores. Without multiple evaluation runs or variance estimates, a -0.14 quality delta for Vicuna-13B could mean anything from "a meaningful capability regression" to "statistically indistinguishable from the baseline." This matters enormously for deployment decisions: if a practitioner observes a -0.14 MT-Bench drop, should they reject MEDUSA-2 for their use case, or is this within normal variance?

More broadly, MT-Bench alone is insufficient to validate quality preservation for a general-purpose acceleration method. The benchmark tests multi-turn conversation ability across 8 categories—it does not test: factual accuracy (no knowledge-intensive QA like MMLU or TriviaQA), reasoning capability (MT-Bench includes a "Reasoning" category with only ~10 questions), code generation correctness (no pass@k evaluation on HumanEval or MBPP), mathematical problem-solving (no MATH or GSM8K), or safety/alignment properties (no harmfulness or refusal evaluation). MEDUSA-2's joint training modifies the backbone's weights—even with the careful recipe, this could introduce subtle distributional shifts that degrade performance on tasks not represented in MT-Bench. The self-distillation variant (which trains on model-generated data) is particularly susceptible to this: the KL divergence loss preserves the shape of the output distribution on in-distribution prompts but provides no guarantee about out-of-distribution generalization. A model that performs identically to the baseline on conversational MT-Bench prompts might still exhibit degraded factual recall or reasoning if the backbone's representations were perturbed during joint training.

What evidence exists in the paper. Table 1 reports quality deltas of +0.01 (Vicuna-7B), -0.07 (Zephyr-7B), -0.14 (Vicuna-13B), and +0.05 (Vicuna-33B) on MT-Bench. Table 2 compares baseline (6.17), direct fine-tuning (5.925), MEDUSA-1 (6.23), and MEDUSA-2 (6.18). The paper does not report standard deviations, confidence intervals, inter-evaluator agreement, or test-retest reliability for these scores. AlpacaEval results (Table 4) report only speed, not quality. No other quality benchmarks are included.

Mitigation status. The paper does not address this limitation. It treats the single-run MT-Bench scores as definitive evidence of quality preservation and does not acknowledge the evaluator variance problem or the narrow scope of the quality evaluation. The authors could have addressed this by: reporting multiple evaluation runs with variance estimates; evaluating on at least one additional benchmark from a different capability domain (e.g., MMLU for knowledge, HumanEval for code); or, for MEDUSA-1, noting that quality preservation is definitional for the model weights (since the backbone is frozen) but not for the generation process when typical acceptance is used (since typical acceptance changes the sampling distribution). The absence of these checks makes the "lossless" and "quality-preserving" claims less robust than they appear.

The MEDUSA-2 Training Recipe's Components Are Not Ablated, Making Attribution and Simplification Impossible

The assumption or constraint. MEDUSA-2's training recipe combines three mechanisms: (1) a combined loss that adds the backbone's next-token cross-entropy to the MEDUSA head losses, (2) differential learning rates where the MEDUSA heads learn 4× faster than the backbone, and (3) a two-stage head warmup where heads are first trained with the backbone frozen before joint training begins. The paper states that this recipe is necessary because directly fine-tuning the backbone with MEDUSA heads degrades quality (Table 2: 6.17 → 5.925). However, the paper never isolates which components are individually responsible for quality preservation versus speedup.

The consequence. The paper cannot answer several practically important questions: Is the combined loss alone sufficient to prevent quality degradation, with differential learning rates and warmup providing marginal benefits? Could a simpler recipe (e.g., combined loss with equal learning rates) achieve comparable results? Is the two-stage warmup necessary, or could careful learning rate scheduling achieve the same effect? Without ablation, practitioners cannot simplify the training procedure for their own use cases, and researchers cannot build on the work with confidence about which components are essential. The current presentation treats the recipe as a monolithic package, which limits reproducibility: if a practitioner's training setup differs (different model architecture, different dataset, different optimizer), they cannot diagnose which component to adjust because the paper provides no signal about component-level sensitivity.

This also weakens the paper's conceptual contribution. The claim that "each component addresses a distinct failure mode" (Section 2.2.2) is architectural reasoning, not experimental fact. The paper describes why each component might help (combined loss prevents catastrophic forgetting, differential LR handles convergence rate asymmetry, warmup prevents early gradient interference), but it never tests whether removing any component actually causes the predicted failure mode. For example, does training without warmup actually produce distorted backbone features (as Kumar et al., 2022 would predict), or is the quality degradation from direct fine-tuning (Table 2) caused entirely by the missing combined loss term? These are empirically distinguishable hypotheses that the paper does not test.

What evidence exists in the paper. Table 2 provides the only relevant comparison: baseline (6.17), direct fine-tuning without any of the three components (5.925), and MEDUSA-2 with all three components (6.18). There is no condition testing combined loss alone, combined loss plus differential LR, warmup alone, or any other partial configuration. Table 3 attributes speedup contributions to individual techniques (tree attention, optimized tree configuration, MEDUSA-2 training), but these are additive components, not ablations of the MEDUSA-2 recipe itself. The paper's discussion of the three components (Section 2.2.2) provides theoretical motivation but no component-level experimental validation.

Mitigation status. There is no mitigation. The paper presents the recipe as a validated whole and does not acknowledge the absence of component-level ablations as a limitation. A full set of ablations (even just the 2³ = 8 combinations of the three binary components) would be a substantial experimental undertaking, but even a targeted ablation of the most critical suspected component (combined loss, since its absence would most directly explain the direct fine-tuning degradation) would substantially strengthen the claims. The paper's statement that "both strategies work well in practice" (Section 2.2.2, referring to the two warmup variants) is the only signal about component-level robustness, and it concerns only the warmup phase.

Memory Overhead of MEDUSA Heads Is Non-Trivial for Deployment, and Training Compute Costs Are Partially Externalized

The assumption or constraint. The paper presents MEDUSA as lightweight: the heads are "just a single layer akin to the original language model head" (Section 2.1.1) and "MEDUSA does not add complexity to the serving system design and is friendly to distributed settings" (Section 2.1.1). The training cost is described as accessible: "5 hours for MEDUSA-1 on Vicuna 7B model with a single NVIDIA A100 PCIE GPU" (Section 2.2.1). MEDUSA-2 additionally uses parameter-efficient fine-tuning (LoRA/QLoRA) to reduce memory.

The consequence. The parameter count tells a more nuanced story. Each MEDUSA head contains d × d + d × V parameters where d is the hidden dimension and V is the vocabulary size. For Llama-7B (d = 4096, V = 32000), each head contributes approximately 4096² + 4096 × 32000 ≈ 16.8M + 131.1M = 147.9M parameters. Five heads total roughly 740M parameters. Relative to the 7B backbone, this is a ~10.6% parameter increase—not negligible for memory-constrained deployments. These parameters must be loaded from HBM during every forward pass (since the heads operate on the last hidden state), adding to the memory-bandwidth burden that MEDUSA is designed to alleviate. The paper's roofline analysis (Appendix G) focuses on the backbone's linear layers and attention, not on the MEDUSA heads themselves, so the heads' contribution to memory traffic is unaccounted for in the overhead measurements.

The training cost also deserves scrutiny. The 5-hour MEDUSA-1 figure is for a single A100 (80GB) PCIe GPU—a high-end datacenter accelerator that costs thousands of dollars. Training on a "single consumer GPU" (as the paper suggests is possible via quantization, Section 2.2.1) would take proportionally longer (an RTX 3090 has roughly one-third the memory bandwidth of an A100 PCIe). The paper does not report consumer GPU training times. For MEDUSA-2, the two-stage training procedure (MEDUSA-1 stage followed by joint training with warmup) requires additional epochs beyond the 2 epochs reported for MEDUSA-1. The total training compute is not systematically reported across all experiments, making cost-benefit analysis difficult: a practitioner evaluating whether to adopt MEDUSA needs to weigh the training cost (GPU hours, engineering time for implementation and hyperparameter tuning) against the inference speedup, amortized over their expected inference volume. The paper provides only point estimates for training time without discussing how these scale with model size, dataset size, or hardware tier.

What evidence exists in the paper. The MEDUSA head architecture is specified in Section 2.1.1 (single FFN layer with residual, W₁ ∈ ℝ^{d×d}, W₂ ∈ ℝ^{d×V}). The training time is mentioned in Section 2.2.1 (5 hours, A100, 60k samples). Appendix B.3 specifies 2-epoch training for Vicuna-7B/13B. The overhead measurements (Table 1: 1.18–1.27× per-step latency increase) include the MEDUSA heads' computation implicitly (since they are part of the forward pass), but the parameter count and its contribution to memory pressure are not discussed. The paper does not report MEDUSA-2's total training time or compare it to MEDUSA-1's.

Mitigation status. The paper partially addresses the training cost barrier through its emphasis on quantization (MEDUSA-1 with 4-bit quantized backbone, analogous to QLoRA) and parameter-efficient fine-tuning (LoRA for MEDUSA-2). The self-distillation pipeline (Section 2.3.2) reduces the data acquisition cost for models without public training data. However, the parameter overhead at inference time is not acknowledged or mitigated—the 740M extra parameters are always loaded regardless of how many candidate tokens are ultimately accepted. The paper's claim that MEDUSA "does not add complexity to the serving system" (Section 2.1.1) is true in the sense that no separate model process is required, but the memory and compute overhead of the heads themselves is non-zero and should be factored into deployment planning. The paper's optimized tree construction (Section 2.3.3) provides a mechanism to use fewer than the full set of 5 heads (since the greedy selection may prune entire heads that have low accuracy), but this optimization is presented as a speedup technique, not as a memory reduction strategy, and its impact on parameter count is not quantified.

Typical Acceptance Sacrifices the Distributional Guarantee That Makes Speculative Decoding Safe for Arbitrary Deployment

The assumption or constraint. The paper introduces typical acceptance (Section 2.3.1) as an alternative to rejection sampling that can "accelerate the decoding speed further while maintaining a similar generation quality" (Section 1). The key tradeoff is stated explicitly: "it is typically unnecessary to match the distribution of the original model" (Section 2.3.1). The paper frames this as a practical choice, not a limitation. However, the consequence of abandoning distributional equivalence is that MEDUSA with typical acceptance produces outputs from an unknown, modified distribution whose relationship to the original model's distribution is governed by the acceptance thresholds ε and δ and the MEDUSA heads' prediction accuracy.

The consequence. This creates a deployment safety concern that does not arise with rejection-sampling-based speculative decoding. With rejection sampling, the output distribution is provably identical to the original model's—if the original model was safe, calibrated, and aligned, the accelerated model inherits those properties exactly. With typical acceptance, the output distribution is a filtered version of the MEDUSA head proposals, where the filter is the entropy-dependent threshold. This can produce subtle quality changes that are difficult to anticipate:

Mode-seeking behavior: The typical acceptance criterion accepts tokens with probability above a threshold, which disproportionately accepts high-probability tokens and rejects low-probability ones. In high-entropy contexts (where many tokens are reasonable), this may not matter. In low-entropy contexts, it should approximate greedy decoding (since only the top token exceeds the threshold). But in intermediate regimes, it may systematically suppress the tail of the distribution, reducing output diversity in ways that are hard to detect without distribution-level evaluation.
Interaction with alignment: For RLHF-tuned models, the output distribution was shaped by a reward model to exhibit specific properties (helpfulness, harmlessness, honesty). Typical acceptance modifies this distribution without any feedback from the reward model. A token that would have been sampled by the original model (because it had moderate probability and satisfied alignment constraints) might be rejected by typical acceptance because it falls below the threshold, while a different token with higher probability but worse alignment properties is accepted. The paper provides no safety or alignment evaluation to detect such shifts.
Threshold sensitivity: Figure 5 shows that quality and acceleration rate are sensitive to ε. At ε = 0.01, quality is ~7.0 with acceleration ~3.5; at ε = 0.25, quality rises to ~7.6 while acceleration drops to ~3.0. A practitioner must tune ε to match their quality requirements, but there is no principled guidance for doing so beyond the empirical sweep on a specific benchmark. The optimal ε likely depends on the task, the temperature, and the model, requiring per-deployment tuning that the paper does not systematize.

What evidence exists in the paper. Figure 5 provides the primary evidence: quality and acceleration rate vary systematically with ε, and typical acceptance with appropriate ε can approximate rejection sampling's quality-speed operating point (at ε ≈ 0.1, typical acceptance achieves ~7.5 quality at ~3.3 acceleration rate vs. rejection sampling's ~7.5 at ~3.5). The paper is transparent that typical acceptance does not match the original distribution: "we diverge because we do not insist on an exact correspondence between the output and language model distribution" (Appendix A.2). However, the paper does not characterize how the modified distribution differs—there is no measurement of output diversity (e.g., entropy of generated sequences, distinct n-gram counts, or distributional distance metrics like KL divergence from the original model's output distribution). The quality evaluation uses GPT-4 MT-Bench scores, which assess output quality in an absolute sense but cannot detect distributional shifts: a model could produce high-quality outputs that are systematically less diverse or differently styled without affecting the average MT-Bench score.

Mitigation status. The paper partially mitigates this by making typical acceptance optional and configurable. Rejection sampling remains available for deployments that require distributional equivalence (the paper mentions it but does not report speedup numbers for MEDUSA with rejection sampling). The ε and δ parameters provide a tuning knob for the quality-speed tradeoff. The entropy-dependent threshold design ensures that in low-entropy contexts (where the model is confident), the acceptance criterion is strict, while in high-entropy contexts (where diversity matters), it relaxes—this is a theoretically motivated adaptive behavior that partially addresses the diversity concern. However, the paper does not provide a method for setting ε and δ that guarantees a specific quality property (e.g., "outputs are within ε-KL divergence of the original model's distribution"), which would make the tradeoff more principled. The absence of alignment-specific evaluation (e.g., harmfulness rates, refusal rates, or instruction-following accuracy) leaves open the possibility that typical acceptance interacts poorly with safety-tuned models in ways that MT-Bench quality scores do not capture. A practitioner deploying MEDUSA in a safety-critical application would need to conduct their own alignment evaluation, which the paper does not facilitate.

7. Implications and Future Directions

How This Work Changes the Landscape

MEDUSA changes the conversation around LLM inference acceleration by demonstrating that the draft model—long treated as a necessary architectural component of speculative decoding—is actually a removable abstraction. Before this paper, the field's default framing was that accelerating auto-regressive decoding requires two separate models: one to propose candidates and one to verify them (Leviathan et al., 2022; Chen et al., 2023). Even work that pushed the boundaries of speculative decoding, such as staged hierarchies (Spector & Re, 2023) or ensembles of draft models (Miao et al., 2023), accepted this two-model premise and worked to make it more sophisticated. MEDUSA does not iterate on draft model design—it eliminates the draft model entirely by asking a different question: can the original model's own hidden representations serve as the candidate generation mechanism?

This reframing matters because it shifts the problem from "how do we coordinate two models efficiently" to "how do we extract more predictive signal from representations we already compute." The evidence that this shift is genuine rather than cosmetic comes from the speedup gap: 2.3–2.8× for MEDUSA versus 1.47–1.60× for speculative decoding with comparable draft models (Table 1). The factor-of-two difference is not explained by incremental engineering improvements—it comes from the architectural choice to share representations between candidate generation and verification, which eliminates distribution shift and amortizes the cost of computing hidden states across both functions.

The paper also resolves a previously muddy question: why did blockwise parallel decoding (Stern et al., 2018) not become the standard for LLM inference? The answer, as MEDUSA demonstrates, is that the architectural idea alone is insufficient. Stern et al. showed that you can attach prediction heads to hidden states, but they did not solve the quality-preservation problem that arises when you train those heads on a pre-trained general-purpose model. MEDUSA-2's training recipe—combined loss, differential learning rates, head warmup—is the missing piece that makes the architecture practically viable. The fact that "Direct Fine-tuning" degrades MT-Bench quality from 6.17 to 5.925 (Table 2) while the full MEDUSA-2 recipe preserves it at 6.18 validates that the training methodology is a first-class contribution, not an implementation detail.

This shifts research attention in several directions. More attractive: work on representation exploitation—what other predictive signals are latent in transformer hidden states? If single-layer FFN heads can predict tokens 2–6 positions ahead from h_t, what about predicting syntactic structure, semantic roles, or retrieval needs? MEDUSA opens a design space where auxiliary prediction heads serve as lightweight extensions of existing models rather than separate systems. Less attractive: research that assumes the draft model is an unavoidable cost of speculative decoding. The paper demonstrates that draft model pre-training (275 A100 GPU hours in Miao et al., 2023) is not a necessary investment—MEDUSA-1 trains in 5 hours on a single A100. Future work on draft-model-free acceleration now has a validated baseline to improve upon, making draft-model-centric approaches harder to justify on cost-effectiveness grounds alone.

The paper also forces a re-examination of what "lossless acceleration" means in practice. Speculative decoding's rejection sampling provides a distributional equivalence guarantee that is mathematically elegant but practically limiting—efficiency drops at higher temperatures where diverse outputs are desired. MEDUSA's typical acceptance scheme demonstrates that relaxing this guarantee yields substantial practical gains (Figure 5) without measurable quality degradation on MT-Bench. This challenges the field's default assumption that distributional matching is the right goal: if users cannot distinguish outputs from the accelerated model and the original, does the mathematical guarantee matter? The paper does not answer this definitively, but it makes the question empirically tractable by providing a controllable alternative (tunable ε and δ thresholds) that spans the spectrum from near-greedy to near-rejection-sampling behavior.

Follow-Up Research This Work Enables

Ablation of the MEDUSA-2 training recipe to identify which components are individually necessary. The paper presents combined loss, differential learning rates, and head warmup as an integrated package and never tests partial configurations. This leaves open the question: is the combined loss alone sufficient to prevent quality degradation, with the other components providing marginal benefits? A follow-up study would train MEDUSA-2 on Vicuna-7B with the 2³ = 8 combinations of these three binary components enabled/disabled, measuring both MT-Bench quality and acceleration rate for each configuration. The specific hypothesis to test: combined loss alone (backbone cross-entropy plus weighted MEDUSA head losses, equal learning rates, no warmup) achieves most of the quality preservation, while differential LR and warmup primarily affect convergence speed and final head accuracy. If confirmed, this would simplify the recipe and make it more accessible. If disconfirmed—if quality degrades whenever any single component is removed—the paper's "each component addresses a distinct failure mode" claim would be experimentally validated and future work could focus on optimizing each component independently rather than treating the recipe as monolithic.

Batched inference with dynamic candidate count throttling. MEDUSA's speedup is validated only at batch size 1, and the paper's own hardware simulation (Appendix G, Figure 22) predicts speedup drops below 1.0× at batch size 64. A critical follow-up would implement a production MEDUSA serving system with dynamic batch-size-dependent behavior: at batch sizes below some threshold (e.g., 4), enable the full tree attention with optimized candidate count; at moderate batch sizes (4–16), reduce the tree to a small number of candidates (e.g., top-1 per head only); at high batch sizes (>16), disable MEDUSA entirely and fall back to standard auto-regressive decoding. The evaluation would measure throughput (tokens/second) and latency (time-to-first-token, time-per-output-token) across a sweep of batch sizes, sequence lengths, and model sizes, comparing against both vanilla decoding and speculative decoding at equivalent batch configurations. This would determine whether MEDUSA's practical deployment envelope extends beyond single-query interactive use to the batched throughput regime that dominates cloud API serving.

Cross-architecture generalization of MEDUSA heads. All experiments use Llama-family models with the same architecture (SiLU activations, rotary position embeddings, specific hidden dimension ratios). Do MEDUSA heads transfer to other architectures? A systematic study would test MEDUSA-1 on models with different design choices: non-Llama architectures (e.g., Mistral, Gemma, Falcon), models with different activation functions (GeLU in OPT, GELU in BLOOM), models with grouped-query attention (GQA) versus multi-head attention (MHA), and encoder-decoder architectures (T5, BART). The key question: is the single-FFN-layer head design universal, or does it exploit specific properties of Llama's hidden representations? A negative result—MEDUSA heads underperform on certain architectures—would reveal what representation properties enable multi-token prediction and guide architecture-specific head design (e.g., deeper heads for models with less predictive hidden states). A positive result (consistent 2×+ speedup across architectures) would establish MEDUSA as a universal acceleration primitive.

Combining MEDUSA with speculative decoding for draft-model-enhanced candidate generation. MEDUSA eliminates the draft model but does not preclude using one. A natural hybrid: use MEDUSA heads on a draft model (small, fast) to generate candidates, then verify with a large target model using tree attention. The draft model's MEDUSA heads would predict multiple future tokens per draft forward pass, and the target model would verify them in batches. This combines the speed of a small draft model with the multi-token parallelism of MEDUSA. A concrete experiment: train MEDUSA heads on TinyLlama-1.1B, use them to generate candidate trees, verify with Vicuna-7B. The expected speedup is the product of (draft model throughput advantage) × (MEDUSA multi-token advantage) × (acceptance rate). The risk: distribution shift between draft and target models may cause the target to reject many MEDUSA-generated candidates, negating the multi-token benefit. The experiment would measure whether the acceptance rate on MEDUSA-head-generated draft tokens is high enough for the combination to beat either approach alone. The results would clarify whether representation sharing (MEDUSA's core advantage) is essential or whether multi-token prediction is beneficial even across a model boundary.

Distributional characterization of typical acceptance to make the quality-speed tradeoff principled rather than empirical. The paper tunes ε and δ by sweeping on a benchmark (Figure 5), yielding a quality-speed curve without theoretical characterization of what distribution the accepted outputs come from. A follow-up would measure the KL divergence between the original model's output distribution and the typical-acceptance-modified distribution as a function of ε, δ, temperature, and task type. The specific experiment: for a fixed set of prompts, sample 10,000 completions from the original Vicuna-7B and 10,000 completions from MEDUSA-2 Vicuna-7B with typical acceptance at various (ε, δ) settings. Measure: (1) KL divergence between empirical token distributions, (2) diversity metrics (distinct n-grams, entropy of generated sequences, self-BLEU), (3) MAUVE score (Pillutla et al., 2021) comparing the two output distributions, and (4) task-specific quality metrics (MT-Bench, AlpacaEval). This would provide a principled way to set ε and δ: instead of "sweep on your benchmark and pick," a practitioner could specify a maximum acceptable KL divergence and derive the corresponding thresholds. It would also clarify whether the quality preservation observed on MT-Bench masks distributional changes that could matter for downstream tasks (e.g., reduced diversity in creative writing, mode collapse in few-shot settings).

Safety and alignment evaluation of typical acceptance on RLHF-tuned models. The paper validates quality preservation on Zephyr-7B (RLHF-aligned) using MT-Bench scores (-0.07 delta), but does not test whether specific alignment properties (refusal of harmful requests, instruction-following accuracy, honesty in sensitive contexts) are preserved under typical acceptance. A safety-focused follow-up would evaluate MEDUSA-2 Zephyr-7B with typical acceptance on: (1) a harmfulness benchmark (e.g., Anthropic's harmlessness evaluation, or a set of adversarial prompts designed to elicit unsafe outputs), measuring refusal rates and harmful completion rates; (2) instruction-following accuracy on a diverse set of constrained generation tasks; and (3) truthfulness on a QA benchmark (TruthfulQA). The comparison would be against both the original Zephyr-7B (with standard sampling) and Zephyr-7B with MEDUSA-2 using rejection sampling (to isolate the effect of typical acceptance from the effect of joint training). If typical acceptance does not degrade safety metrics, this provides strong evidence for deployment in safety-sensitive applications. If it does degrade safety (even while maintaining MT-Bench scores), this would reveal that MT-Bench's coarse quality evaluation is insufficient for alignment-critical deployments and motivate the development of safety-aware acceptance schemes.

Practical Applications and Downstream Use Cases

Interactive coding assistants on consumer hardware. MEDUSA-1 can be trained on a quantized 7B code model (e.g., CodeLlama-7B) in hours on a single GPU, then deployed on a consumer laptop with a GPU. The 2.2× speedup (MEDUSA-1) reduces per-token latency from ~25ms to ~11ms, making interactive code completion feel substantially more responsive. The Coding category on MT-Bench achieves 3.29× speedup with MEDUSA-2 (Figure 3b), suggesting that code-specific fine-tuning of MEDUSA heads could push speedup even higher for programming tasks where syntax is highly predictable. The self-distillation pipeline enables this without access to the original fine-tuning dataset: generate code completions from the target model using open-source coding prompts (e.g., from The Stack or GitHub), train MEDUSA heads on these completions, and deploy. The memory overhead (~740M parameters for 5 heads on a 7B model, ~10.6% increase) is manageable on consumer GPUs with 8–16 GB VRAM, especially when combined with 4-bit quantization of the backbone. This application directly realizes the paper's democratization argument: high-quality code LLM inference becomes feasible on hardware that millions of developers already own.

Low-latency chatbot deployments where GPU resources are shared or limited. Consider a small company running a customer-support chatbot using Vicuna-13B on a single A100 GPU. At baseline throughput (~28 tokens/second, Figure 3a), a 100-token response takes ~3.6 seconds, which is at the boundary of acceptable latency for chat. MEDUSA-1's 2.33× speedup brings this to ~1.5 seconds—a perceptible improvement that makes the interaction feel instantaneous. Critically, MEDUSA-1 achieves this without any risk to response quality (frozen backbone) and without the operational complexity of deploying a separate draft model (no additional model server, no inter-model routing, no memory partitioning between two models). The training cost (5 GPU-hours on an A100) is negligible compared to the ongoing inference cost savings: if the chatbot serves 10,000 queries/day, reducing per-query GPU time by 2.33× saves ~15 A100-hours/day, paying back the training cost within hours of deployment. This use case is where MEDUSA's batch-size-1 design assumption (interactive queries arrive one at a time, not batched) matches reality, and where the simplicity of single-model deployment (no draft model to maintain) provides the strongest operational advantage over speculative decoding.

Self-improvement data generation pipelines. When using LLMs to generate training data for themselves (e.g., STaR, ReST^EM, or rejection sampling fine-tuning), inference cost often dominates the total compute budget—the model generates thousands or millions of completions, and only a fraction are retained as training examples. MEDUSA-2's 2.83× speedup on Vicuna-7B directly translates to 2.83× more generated training examples per GPU-hour, or equivalently, 65% lower cost to generate the same volume of data. The self-distillation pipeline (Section 2.3.2) is specifically designed for this scenario: the model's own outputs are the training data, and the KL-divergence loss prevents quality degradation when training MEDUSA heads on that data. A concrete pipeline: (1) train MEDUSA-2 on the target model using self-distillation with 100k seed prompts; (2) use the accelerated model to generate 1M training examples for a downstream task; (3) filter these examples using a reward model or verifier; (4) fine-tune the original (non-MEDUSA) model on the filtered examples. The MEDUSA heads are a transient acceleration tool—used during data generation, discarded afterward—so the inference-time parameter overhead is irrelevant. This application exploits MEDUSA's fastest training path (self-distillation with KL divergence, which the paper shows preserves quality) and its strongest speedup regime (batch size 1, interactive-length generations), making it perhaps the most immediately practical use case in the paper.

Evaluation and benchmarking infrastructure. Organizations that maintain LLM evaluation pipelines (e.g., HELM, LM Evaluation Harness, AlpacaEval) run thousands of model inferences across hundreds of prompts and tasks. These pipelines operate at batch size 1 (each prompt is processed independently) and are latency-sensitive (total evaluation time determines iteration speed). MEDUSA-1 can be applied to any evaluated model without modifying its weights, providing ~2× speedup on evaluation runs with zero risk of changing model outputs (the backbone is untouched, and evaluation typically uses greedy decoding where acceptance is deterministic). The 5-hour training cost per model is amortized over many evaluation runs. For a benchmark like AlpacaEval (805 prompts), reducing inference time from ~30 minutes to ~15 minutes per model enables faster experimentation cycles. The key advantage over speculative decoding in this setting: no draft model compatibility issues (different model architectures, tokenizers, and training distributions break speculative decoding's draft-target alignment), and no need to maintain draft model infrastructure across the diverse set of models being evaluated.

Dependency tree

Papers this one builds on, walking backwards through the same domain and shared topics.

SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference2023· inference-optimization← LLM-judged similar
- Read More2023· inference-optimization← in-domain predecessor
- Fast KVzip: Efficient and Accurate LLM Inference with Gated KV Eviction2026· inference-optimization← similar topic
- Locality-aware Parallel Decoding for Efficient Autoregressive Image Generation2025· inference-optimization← similar topic
- Set Block Decoding is a Language Model Inference Accelerator2025· inference-optimization← similar topic
- Break the Sequential Dependency of LLM Inference Using LOOKAHEAD DECODING2024· inference-optimization← similar topic
- LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding2024· inference-optimization← similar topic
- LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference2024· inference-optimization← similar topic
- LLM in a Flash: Efficient Large Language Model Inference with Limited Memory2023· inference-optimization← similar topic
- Fast Inference from Transformers via Speculative Decoding2022· inference-optimization← similar topic
EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test2025· inference-optimization← similar topic
- LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention2025· inference-optimization← in-domain predecessor
- Reward-Guided Speculative Decoding for Efficient LLM Reasoning2025· reasoning← topic thread
An Efficient Inference Framework for Early-exit Large Language Models2024· inference-optimization← similar topic
DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads2024· inference-optimization← similar topic
- PrefixQuant: Eliminating Outliers by Prefixed Tokens for Large Language Models Quantization2024· inference-optimization← in-domain predecessor
- A Controlled Study on Long Context Extension and Generalization in LLMs2024· training-methods← topic thread

What it influenced

Later papers in the same chain. Direction reversed.

Speculative Streaming: Fast LLM Inference without Auxiliary Models2024· inference-optimization→ LLM-judged similar
- Ouroboros: Generating Longer Drafts Phrase by Phrase for Faster Speculative Decoding2024· inference-optimization→ in-domain predecessor
- Fast Inference from Transformers via Speculative Decoding2022· inference-optimization→ similar topic
- Accelerating LLM Inference with Staged Speculative Decoding2023· inference-optimization→ similar topic
- Break the Sequential Dependency of LLM Inference Using LOOKAHEAD DECODING2024· inference-optimization→ similar topic
Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting2024· inference-optimization→ LLM-judged similar
- Clover: Regressive Lightweight Speculative Decoding with Sequential Knowledge2024· inference-optimization→ in-domain predecessor
- LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding2024· inference-optimization→ similar topic
Set Block Decoding is a Language Model Inference Accelerator2025· inference-optimization→ LLM-judged similar
- Bridging the Gap Between Promise and Performance for Microscaling FP4 Quantization2025· inference-optimization→ in-domain predecessor
- SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference2023· inference-optimization→ similar topic
EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty2024· inference-optimization→ in-domain predecessor
- EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test2025· inference-optimization→ LLM-judged similar
- TAPS: Task Aware Proposal Distributions for Speculative Sampling2026· inference-optimization→ LLM-judged similar