Accelerating AI Inference with Medusa and EAGLE

This article provides a comprehensive overview of how the Medusa and EAGLE frameworks accelerate Large Language Model (LLM) inference through speculative decoding. Standard AI text generation processes words one at a time, creating a major memory bottleneck. To solve this, both Medusa and EAGLE eliminate the need for a separate draft model by attaching lightweight decoding heads directly to the base model. While Medusa uses multiple independent heads to guess several future tokens in parallel, EAGLE introduces a more advanced single-head autoregression technique that leverages hidden feature states for higher accuracy. Both methods ultimately verify these predicted tokens simultaneously, drastically increasing generation speed without sacrificing output quality.

The Bottleneck of Standard LLM Inference

Traditional autoregressive Large Language Models generate text sequentially. To predict the fifth word in a sentence, the model must first generate the fourth word, feed it back into the system, and load all the model's massive parameters from memory all over again. This makes LLM inference heavily memory-bandwidth bound rather than compute-bound. Speculative decoding solves this by guessing multiple future tokens at once and using a single forward pass of the main model to verify them in parallel. Medusa and EAGLE are two prominent architectures that achieve this by adding lightweight decoding heads directly to the original model.

How Medusa Works: Parallel Decoding Heads

Medusa simplifies speculative decoding by removing the need for a separate, smaller "draft" model. Instead, it attaches multiple lightweight decoding heads directly onto the final hidden state of the frozen base LLM.

  • Architecture: While the original language model head predicts the immediate next token (at position \(t+1\)), Medusa adds \(K\) additional heads. Head 1 predicts the token at \(t+2\), Head 2 predicts \(t+3\), and so on. Each head is a simple feed-forward neural network with a residual connection.
  • Training: During training, the base LLM remains frozen (in the standard Medusa-1 setup), and only the new heads are trained using cross-entropy loss to match the base model's predictions at their respective future steps.
  • Tree Attention Verification: Because the heads predict future tokens independently, they can easily make conflicting guesses. To maximize the chance of a correct sequence, each head generates its top-k most likely tokens. Medusa combines these predictions into a candidate tree. Using a customized "tree attention" mask, the base LLM processes all branches of this tree simultaneously in a single forward pass, accepting the longest correct sequence of tokens.

How EAGLE Works: Feature-Level Autoregression

EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency) identifies a major flaw in Medusa's approach: predicting the 3rd or 4th future token without knowing the 2nd token is inherently inaccurate. EAGLE solves this by making the lightweight drafting process autoregressive, while keeping it blazingly fast.

  • Architecture: Instead of multiple independent heads, EAGLE attaches a single lightweight autoregressive head (typically consists of a single Transformer decoder layer) to the base model.
  • Feature Extrapolation: Rather than predicting discrete tokens directly, EAGLE operates on the continuous "feature" level—specifically, the second-to-top hidden states of the LLM. Hidden states contain much richer contextual information than a single chosen word.
  • Addressing Uncertainty: Predicting future features has a catch: the hidden state of the next step depends on the exact token sampled in the current step. EAGLE resolves this by concatenating the current hidden state with the embedding of the just-sampled token, feeding both into its lightweight head to predict the next hidden state.
  • Sequential Drafting, Parallel Verification: EAGLE runs its lightweight head sequentially a few times to rapidly draft a tree of future tokens. Once this fast drafting phase is complete, the base LLM verifies the entire tree in a single parallel forward pass, just like Medusa.

Key Differences in Simultaneous Prediction

While both methods aim to verify multiple tokens in one step of the base model, their approach to generating the guesses differs fundamentally:

  1. Drafting Mechanism: Medusa predicts future tokens in parallel using separate heads for each future position. EAGLE predicts future states sequentially using a single recurrent head, which is still exceptionally fast because the head is minuscule compared to the base model.
  2. Context Awareness: Medusa's heads are blind to what the other heads are doing; the head predicting the third word does not know what the first head picked for the second word. EAGLE passes the sampled token embedding from one step to the next, ensuring that future predictions are contextually aligned with earlier guesses.
  3. Acceptance Rate: Because EAGLE resolves token uncertainty step-by-step at the feature level, its draft trees are significantly more accurate than Medusa's. This leads to more tokens being accepted per forward pass, resulting in a higher overall speedup (often 3x faster than standard decoding, compared to Medusa’s roughly 2x speedup).