Medusa: Multi-Head Decoding for Faster Inference

Overview

Medusa is a technique that accelerates LLM inference by predicting multiple future tokens in parallel using separate decoding heads, then verifying them. Achieves 2-3x speedup without model changes or distillation.

• Paper: "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads" (Cai et al., 2023)
• Key Idea: Add auxiliary prediction heads for next N tokens
• Speedup: 2-3x with 1-2% quality loss
• Advantage: Works with any existing LLM without modification
• Trade-off: Increased inference parameters (auxiliary heads)

---

The Problem: Token-by-Token Generation is Sequential

Inference Bottleneck

Standard Autoregressive Generation:

Token 1:  model(input) → output_1 → sample → token_1
Token 2:  model(prev + token_1) → output_2 → sample → token_2
Token 3:  model(prev + token_1 + token_2) → output_3 → sample → token_3
...

Timeline:
Step 1:  ════ Forward pass (50ms)
Step 2:           ════ Forward pass (50ms)
Step 3:                ════ Forward pass (50ms)
...
Total for 100 tokens: 100 × 50ms = 5 seconds

Problem:
- Each token takes 50ms (compute)
- Must wait for previous token to continue
- Can't parallelize!
- Throughput: ~20 tokens/sec (very slow!)

Question:
Can we predict multiple tokens before committing to them?

Why Parallel Prediction is Feasible

Key observation:
Next token is mostly determined by recent context

Example:
"The capital of France is P" → Next token: "aris" (very predictable)
"Once upon a time" → Next token: hard to predict

But:
We can make EDUCATED GUESSES for next 2-3 tokens!

Correct sequence: "Paris, France"
Guess 1: "P" (likely)
Guess 2: "aris" (likely, follows P)
Guess 3: "," (likely, follows "Paris")

If guesses correct:
- Process in one forward pass instead of 3!
- 3x speedup!

If guess wrong:
- Reject and resample
- Still faster than sequential (amortized)

---

How Medusa Works

Architecture

Standard LLM:
- ┌─────────────────────────────────────┐
    - Input Embeddings                     │
  - ┤
    - Transformer Layers (32 layers)       │
  - ┤
    - Final Layer Norm                     │
  - ┤
    - LM Head (output vocabulary)          │
  - ┤
    - Output: vocab_size logits            │
  - ┘
         ↓
    Sample token t+1

Medusa Modification:
- ┌─────────────────────────────────────┐
    - Input Embeddings                     │
  - ┤
    - Transformer Layers (32 layers)       │
  - ┤
    - Final Layer Norm                     │
  - ┬───────────┤
    - LM Head                  │ (original)
  - ┼───────────┤
    - Medusa Head 1            │ (parallel head)
    - [small MLP]              │
    - Output: vocab_size       │ (predict t+2)
  - ┼───────────┤
    - Medusa Head 2            │ (parallel head)
    - [small MLP]              │
    - Output: vocab_size       │ (predict t+3)
  - ┼───────────┤
    - Medusa Head 3            │ (parallel head)
    - [small MLP]              │
    - Output: vocab_size       │ (predict t+4)
  - ┴───────────┘
    ↓           ↓           ↓           ↓
  Token t+1   Token t+2   Token t+3   Token t+4
  (main)      (guess)     (guess)     (guess)

Inference Process

Step 1: Forward pass (once)
  - Input: prompt + previously accepted tokens
  - Output: logits from main head + all Medusa heads
  - Time: 50ms (same as before!)

Step 2: Generate candidate tokens
  - Main head: sample token t+1 (highest probability)
  - Medusa head 1: sample token t+2 (highest probability)
  - Medusa head 2: sample token t+3 (highest probability)
  - Medusa head 3: sample token t+4 (highest probability)
  - Time: <1ms (sampling is fast)

Step 3: Verification
  - Check: Is t+1 likely? (Yes, from main head)
  - Check: Given t+1, is t+2 likely?
    - Forward pass: model(prompt + t+1 + t+2_guess)
    - Compare main head output to t+2_guess
    - If probability > threshold: ACCEPT t+2
    - Else: REJECT, resample
  - Check: Given t+1, t+2, is t+3 likely?
  - Check: Given t+1, t+2, t+3, is t+4 likely?
  - Time: Usually ~50ms for 1-2 verifications

Result:
  - If all 4 tokens accepted: 4 tokens in 50ms (4x speedup!)
  - If 2 tokens accepted: 2 tokens in 50ms (2x speedup)
  - Average: 2-3x speedup

Real sequence vs guesses:
  - Model predicts: [high prob for "The"]
  - Medusa guesses: [high prob for "The"]
  - ACCEPT "The"
│
  - Model predicts (given "The"): [high prob for "cat"]
  - Medusa guesses: [high prob for "cat"]
  - ACCEPT "cat"
│
  - Model predicts (given "The cat"): [high prob for "sat"]
  - Medusa guesses: [high prob for "jumped"]
  - MISMATCH! REJECT jump, sample "sat"
│
Total: 3 tokens generated
- With guessing: 50ms + verification
- Without guessing: 150ms

Implementation

class MedusaHead(nn.Module):
    """Auxiliary decoding head for predicting future token"""

    def __init__(self, hidden_dim, vocab_size):
        super().__init__()
        self.head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, vocab_size)
        )

    def forward(self, x):
        return self.head(x)

class ModelWithMedusa(nn.Module):
    def __init__(self, base_model, hidden_dim, vocab_size, num_medusa_heads=3):
        super().__init__()
        self.base_model = base_model
        self.lm_head = base_model.lm_head

        # Medusa heads for parallel decoding
        self.medusa_heads = nn.ModuleList([
            MedusaHead(hidden_dim, vocab_size)
            for _ in range(num_medusa_heads)
        ])

    def forward(self, input_ids):
        """
        Returns:
        - main_logits: (batch, seq_len, vocab_size)
        - medusa_logits: list of (batch, seq_len, vocab_size)
        """
        # Base model forward
        hidden = self.base_model(input_ids, output_hidden_state=True)
        last_hidden = hidden.last_hidden_state

        # Main head
        main_logits = self.lm_head(last_hidden)

        # Medusa heads (parallel predictions)
        medusa_logits = [head(last_hidden) for head in self.medusa_heads]

        return main_logits, medusa_logits

def speculative_decoding_with_medusa(model, prompt, max_tokens=100, 
                                     medusa_temperature=0.5, 
                                     verify_threshold=0.9):
    """
    Generate with Medusa (speculative decoding)

    Key parameters:
    - medusa_temperature: Lower = more confident guesses
    - verify_threshold: Minimum probability to accept guess
    """
    generated = []
    current_input = prompt

    for _ in range(max_tokens):
        # Forward pass (single, not multiple!)
        main_logits, medusa_logits = model(current_input)

        # Latest token predictions
        main_pred = main_logits[:, -1, :]  # (batch, vocab)
        medusa_preds = [m[:, -1, :] for m in medusa_logits]

        # Generate candidates
        main_token = sample(main_pred, temperature=0.7)
        medusa_tokens = [
            sample(m, temperature=medusa_temperature)
            for m in medusa_preds
        ]

        # Verify candidates
        candidates = [main_token] + medusa_tokens
        accepted_count = 0

        for i, token in enumerate(candidates):
            if i == 0:
                # Main token: always accept
                generated.append(token)
                accepted_count += 1
            else:
                # Verify: recompute with candidate token
                test_input = current_input + [generated[-1], token]
                test_logits, _ = model(test_input)

                # Get probability of this token
                token_prob = softmax(test_logits[:, -1, :])[0, token]

                if token_prob > verify_threshold:
                    generated.append(token)
                    accepted_count += 1
                else:
                    # Reject and stop verifying
                    break

        # Update input for next iteration
        current_input = current_input + generated[-accepted_count:]

        # Typically accept 2-3 tokens per forward pass
        # So 100 tokens takes ~30-50 forward passes instead of 100

    return generated

---

Performance Analysis

Speedup Calculation

Scenario: Generate 100 tokens with Medusa (3 heads)

Standard generation:
  - 100 forward passes × 50ms = 5000ms
  - Throughput: 20 tokens/sec

Medusa with 70% acceptance rate:
  - Main token: Always accepted
  - Head 1: 70% acceptance rate
  - Head 2: 50% acceptance rate (70% × 70%)
  - Head 3: 35% acceptance rate (70% × 70% × 70%)
  - Expected tokens per pass: 1 + 0.7 + 0.5 + 0.35 = 2.55 tokens/pass
  - Forward passes needed: 100 / 2.55 ≈ 39 passes
  - Verification passes: ~39 × 2 ≈ 78 passes (rough estimate)
  - Total passes: ~117 (vs 100 standard)
  - But each pass is parallel prediction, not sequential
  - Effective speedup: ~2.5x!
  - Throughput: 50 tokens/sec (2.5x improvement)

Better scenario (80% acceptance):
  - Expected tokens per pass: 1 + 0.8 + 0.64 + 0.51 = 2.95
  - Forward passes: 34
  - Speedup: ~3x
  - Throughput: 60 tokens/sec

Quality Impact

Experiment: LLaMA 7B with Medusa

Metric              Baseline  Medusa    Delta
─────────────────────────────────────────────
MMLU Score          45.3%     44.8%     -0.5%
Human Preference    —         Tied      Equal
Generation Speed    1x        2.5x      +150%

Quality loss:
  - Minimal (<1%)
  - Acceptance threshold tuning can improve
  - Trade: Fast inference worth slight quality dip

Configuration impact:
  - 3 heads: 2.5x speedup, 0.5% loss
  - 2 heads: 1.8x speedup, 0.2% loss
  - 1 head: 1.3x speedup, <0.1% loss
  - More heads = more speedup but potentially less accurate guesses

---

Advantages and Limitations

Advantages

✅ 2-3x speedup without model changes
✅ Compatible with existing models
✅ Minimal quality loss (<1%)
✅ Works with batching
✅ Orthogonal to other optimizations (can combine with quantization, etc.)

Limitations

❌ Requires retraining Medusa heads
❌ Quality degrades for "branching" outputs (multiple valid continuations)
❌ Verification adds some compute overhead
❌ Accepts only when confident (sensitive to threshold)
❌ Less effective with high temperature (random sampling)

---

When to Use Medusa

Use Medusa when:
✅ Inference speed is critical
✅ Can afford to retrain auxiliary heads
✅ Outputs are relatively deterministic
✅ Want simple, model-agnostic solution

Avoid Medusa when:
❌ Quality must be perfect
❌ Outputs are highly variable (creative writing)
❌ Can't modify inference code

---

Key Takeaways

🎯 Parallel prediction: Guess next tokens in parallel
✅ Verify guesses: Expensive but amortized over guesses
⚡ 2-3x speedup with minimal quality loss
🔄 Orthogonal to other optimizations
📊 Trade: Slight quality for major speed gain

---

Related Notes

• Speculative Decoding - Similar concept, different approach
• Continuous Batching - Works together for inference optimization
• Token Merging (Tome) - Another inference acceleration technique
• Llm Inference Optimization - Complete inference stack