💬 NLP

LogitSpec: Accelerating Retrieval-based Speculative Decoding via Next Next Token Speculation

LogitSpec is a training-free, plug-and-play retrieval-based speculative decoding method that accelerates LLM inference by leveraging the last token's logits to speculate the "next-next" token, thereby expanding the retrieval range to achieve up to 2.61× speedup and higher token acceptance rates.

Original authors: Tianyu Liu, Qitan Lv, Hao Li, Xing Gao, Xiao Sun, Xiaoyan Sun

Published 2026-04-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tianyu Liu, Qitan Lv, Hao Li, Xing Gao, Xiao Sun, Xiaoyan Sun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef (the Large Language Model, or LLM) trying to write a complex recipe. The old way of doing this is very slow: you think of one word, write it down, wait for the computer to check if it's correct, then think of the next word, and so on. It's like writing a letter one letter at a time, waiting for a stamp of approval after every single letter.

Speculative Decoding is a new way to speed this up. Instead of writing one word at a time, you hire a sous-chef (a draft model) to guess the next few words for you. You then quickly check if the sous-chef was right. If they were, you accept all those words at once. If they were wrong, you just throw away the wrong guesses and try again. This saves a lot of time.

However, hiring a sous-chef has problems:

You have to train them specifically for your kitchen (expensive and hard).
They take up extra space in your kitchen (memory).
If you change your main recipe, you have to retrain the sous-chef.

The Problem with "Retrieval" (The Library Approach)

Some researchers tried to solve this by firing the sous-chef and using a library instead. Instead of a person guessing, the computer looks at what you've already written and searches a massive database of past recipes to find a matching phrase.

The Flaw: This is like trying to find the next word in a sentence by looking at the immediately previous word.

Your sentence: "What is the area of the..."
The Library's guess: It sees "the" and thinks, "Oh, 'the' usually goes with 'triangle'!" (Because it found a past sentence: "What is the area of the triangle?").
The Reality: You actually wanted to say "What is the area of the circle?" or "What is the area of the field?"
Result: The library gets stuck on the wrong word because it's only looking at the immediate context, not the whole idea.

The Solution: LogitSpec (The "Crystal Ball" Chef)

The authors of this paper, LogitSpec, realized that the main chef (the LLM) actually has a hidden superpower it doesn't use often enough.

When the chef predicts the next word, they also have a "feeling" (mathematically called logits) about what comes after that next word. Even though the chef is only supposed to say the next word, their brain is already whispering, "And after that, it's probably 'circle'."

LogitSpec uses this whisper to fix the library search.

Here is how it works, step-by-step, using a creative analogy:

1. The "Crystal Ball" Step

Instead of just looking at the last word ("the"), LogitSpec asks the chef: "If you had to guess the word two steps ahead, what would it be?"

The chef looks at its internal "feelings" and says, "I bet the word after 'the' is 'circle'."
This is the Next-Next Token Speculation.

2. The "Super-Search" Step

Now, instead of searching the library for just "the", LogitSpec searches for the phrase "the circle".

Old Way (Vanilla Retrieval): Searches for "the". Finds "the triangle", "the sky", "the dog". (Wrong!)
LogitSpec Way: Searches for "the circle". Finds "the circle of friends", "the circle of life", "the circle of the field". (Much more accurate!)

3. The Verification

The computer takes these better guesses and checks them against the master chef. Because the guesses are now much more accurate, the master chef accepts them more often.

Why is this a big deal?

No Extra Training: You don't need to hire or train a new sous-chef. You just use the main chef's existing "feelings" (logits) that were already there.
Plug-and-Play: It works with any existing language model without changing its code or weights.
Speed: In their tests, this method made the computer 2.6 times faster at writing text. It also meant the computer could accept, on average, 3.28 words at a time instead of just one.

The Bottom Line

Think of LogitSpec as a detective who doesn't just look at the crime scene (the last word) to guess what happened next. Instead, the detective uses a psychic hunch (the logit) to guess the next crime scene, and then uses that hunch to find the perfect evidence in the library.

By looking two steps ahead, the system stops getting confused by similar-looking words and finds the right words much faster, making AI write text significantly quicker without needing any extra training or expensive hardware.

1. Problem Statement

Large Language Models (LLMs) suffer from high inference latency due to auto-regressive decoding, where only one token is generated per forward pass. Speculative Decoding (SD) addresses this by using a small "draft" model to propose multiple tokens, which the target model validates in parallel.

However, existing SD approaches face two main limitations:

Draft-Model-Based SD: Requires training and deploying an extra draft model (e.g., Medusa, Eagle). This introduces significant memory overhead, complex deployment challenges, and training costs.
Draft-Model-Free (Retrieval-Based) SD: Methods like REST or PLD retrieve draft tokens from the context or an external database without a draft model. However, they rely on a matching paradigm that often fails. If the current context does not have an exact n-gram match in the reference, the retrieval fails, leading to low acceptance rates. As shown in the paper, standard retrieval methods fail to find matched tokens in over 30% of decoding steps in some scenarios.

Core Challenge: How to improve the accuracy of retrieval-based speculative decoding without training an extra model, specifically when exact context matches are scarce.

2. Methodology: LogitSpec

The authors propose LogitSpec, a training-free, plug-and-play framework that enhances retrieval-based SD by leveraging the logits of the last token to speculate the next-next token ( $x_{i+2}$ ).

Key Insight

The paper empirically observes that the logit distribution of the current token ( $x_i$ ) contains sufficient information to predict the next-next token ( $x_{i+2}$ ) with high accuracy, even without fine-tuning. In over 50% of decoding steps, the ground-truth next-next token appears within the top-60 entries of the last token's logit distribution.

The LogitSpec Workflow

LogitSpec operates in two main phases per decoding step:

Speculation & Retrieval Expansion:
- Step 1 (Next Token): The model predicts the next token ( $x_{i+1}$ ) normally.
- Step 2 (Next-Next Speculation): Instead of stopping, LogitSpec takes the top- $k$ entries from the last logit (excluding the predicted $x_{i+1}$ ) as candidate speculations for the next-next token ( $\hat{x}_{i+2}$ ).
- Step 3 (Dual Retrieval): The system retrieves relevant references for both the next token ( $x_{i+1}$ $x_{i + 1}$ ) and the speculated next-next tokens ( $\hat{x}_{i+2}$ $\overset{x}{^}_{i + 2}$ ).
  - Query 1: $g^{(1)} = (x_{i-m+2}, \dots, x_i, x_{i+1})$
  - Query 2: $g^{(2)} = (x_{i-m+3}, \dots, x_i, x_{i+1}, \hat{x}_{i+2})$
- This effectively expands the retrieval search space. Even if the exact sequence for $x_{i+1}$ isn't found, the model might find a match for the sequence $x_{i+1} + \hat{x}_{i+2}$ , significantly increasing the probability of finding valid draft tokens.
Draft Tree Construction & Verification:
- Retrieved tokens are organized into a draft tree where $x_{i+1}$ is the root, and branches correspond to different speculated next-next tokens.
- A tree attention mask is applied to allow parallel verification of all branches by the target model.
- Pruning Strategy: To control computational cost, a heuristic pruning strategy is applied based on the rank of the speculated next-next token. Higher-confidence speculations (lower rank) retain more retrieved tokens, while lower-confidence ones retain fewer.

3. Key Contributions

Novel Observation: The authors demonstrate that the last token's logit can robustly predict the next-next token across various models (Vicuna, Llama, Qwen) and tasks, without fine-tuning.
LogitSpec Framework: A new retrieval-based SD method that uses next-next token speculation as a guidance anchor to expand the retrieval range. It is training-free and plug-and-play.
Performance Gains: Extensive experiments show LogitSpec achieves up to 2.61× speedup and a mean of 3.28 accepted tokens per decoding step, significantly outperforming existing retrieval-based baselines (REST, PLD, Lookahead).
Robustness: Unlike other retrieval methods that struggle in low-repetition contexts (e.g., translation or complex reasoning), LogitSpec maintains high performance because the speculation mechanism compensates for the lack of exact context matches.

4. Experimental Results

The authors evaluated LogitSpec on diverse benchmarks including Spec-Bench (MT, Trans, Sum, QA, Math, RAG), HumanEval (Code), GSM8K (Math), MATH, AIME, and LongBench.

Speedup:
- Vicuna-7B: Achieved 2.61× speedup on CNN/DM and 2.47× on HumanEval.
- Average: Across all benchmarks, LogitSpec achieved a 2.26× speedup on average (Vicuna-7B), compared to ~1.33× for REST and ~1.78× for PLD.
Acceptance Rate:
- Achieved a Mean Accepted Tokens (MAT) of 3.28 on CNN/DM, compared to 1.64 for REST.
- On LongBench (long-context), it achieved 2.01× speedup, demonstrating effectiveness even with massive context windows.
Retrieval Success:
- Standard retrieval methods (PLD) failed to find matches in >30% of steps.
- LogitSpec achieved a 97.76% successful retrieval rate on Spec-Bench, a massive improvement driven by the next-next token guidance.
Overhead:
- The retrieval process adds negligible overhead (~1.17% of total wall-clock time) due to efficient hash-based implementation.
- The method is compatible with existing frameworks (e.g., Hugging Face Transformers) and requires no extra GPU memory for a draft model.

5. Significance and Impact

Eliminates Training Costs: LogitSpec removes the need to train or maintain a separate draft model, making speculative decoding accessible for any pre-trained LLM immediately.
Solves the "Cold Start" Retrieval Problem: By using the model's own predictive capability (logits) to guide retrieval, it overcomes the fundamental limitation of retrieval-based methods that fail when context repetition is low.
Scalability: The method scales well across different model sizes (0.5B to 70B) and architectures, proving that the "next-next token" signal is a fundamental property of modern LLMs.
Practical Deployment: Being plug-and-play, it can be easily integrated into production inference engines (like vLLM) to accelerate LLM serving without architectural changes or retraining.

In conclusion, LogitSpec represents a significant step forward in efficient LLM inference by cleverly repurposing the model's internal logits to enhance retrieval, achieving near-draft-model performance without the associated costs.