💬 NLP

Beyond the Target: From Imitation to Collaboration in Speculative Decoding

This paper introduces Collaborative Speculative Decoding (CoSpec), a novel framework that replaces the traditional imitation-based approach with a reinforcement learning-driven arbitration policy, allowing a smaller draft model to selectively contribute tokens even when they mismatch the larger target model, thereby achieving both significant inference speedups and superior final accuracy.

Original authors: Jinze Li, Yixing Xu, Guanchen Li, Jinfeng Xu, Shuo Yang, Yang Zhang, Xuanwu Yin, Dong Li, Edith C. H. Ngai, Emad Barsoum

Published 2026-05-26

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Jinze Li, Yixing Xu, Guanchen Li, Jinfeng Xu, Shuo Yang, Yang Zhang, Xuanwu Yin, Dong Li, Edith C. H. Ngai, Emad Barsoum

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 trying to solve a very difficult puzzle, like a complex math problem or a tricky coding challenge. You have two people helping you:

The Expert (Target Model): A brilliant, highly trained professor who knows almost everything but is very slow. They take their time to think about every single step before writing it down.
The Intern (Draft Model): A fast, energetic, but less experienced assistant. They can think and write very quickly, but they make mistakes more often.

The Old Way: "The Strict Supervisor"

In the traditional method (called Speculative Decoding), the Intern tries to guess the next few steps of the solution and writes them down quickly. Then, the Expert checks them.

The Rule: If the Intern's guess matches the Expert's thought exactly, the Expert says, "Great, keep going!" and they move forward together.
The Problem: If the Intern guesses even one word differently, the Expert immediately says, "No, that's wrong," crosses it out, and writes their own version. The Intern's work is thrown away, even if the Intern was actually right about that specific part!

The paper argues that this is wasteful. Sometimes, the Intern might have a brilliant idea that the Expert missed because the Expert is stuck in a certain way of thinking. By blindly rejecting the Intern's different ideas, the system misses out on good solutions.

The New Way: "The Collaborative Team" (CoSpec)

The authors introduce a new system called Collaborative Speculative Decoding (CoSpec). Instead of the Expert being the only boss, they add a Smart Mediator (an arbitration policy).

Here is how it works:

The Setup: The Intern still writes down a block of ideas quickly. The Expert still checks them in parallel (at the same time).
The Match: If they agree, great! They move forward fast.
The Disagreement: If the Intern and Expert disagree on a step, the Mediator steps in.
- The Mediator looks at the context, the Intern's idea, and the Expert's idea.
- It asks: "Which of these two paths is more likely to lead to the correct final answer?"
- If the Mediator thinks the Intern's different idea is actually the better choice, it says, "Keep the Intern's idea!"
- If the Intern is clearly wrong, it says, "Go with the Expert."

Why This Matters

Think of it like a sports team. In the old system, the coach (Expert) would bench the star player (Intern) the moment they tried a different play than the coach's playbook, even if that different play was the winning move.

In the new system, the coach listens to the assistant coach. If the assistant suggests a different play that might win the game, the team tries it.

The Results:

Faster: Because the team accepts more of the Intern's good guesses, they don't have to stop and rewrite as often. They finish the puzzle much faster.
Smarter: Because they sometimes keep the Intern's "different" ideas, they actually solve more puzzles correctly than the Expert could have solved alone.

The Bottom Line

This paper shows that by treating the fast, smaller model as a partner rather than just a copycat, we can make AI both faster and more accurate. The "mistakes" or differences between the two models aren't always errors; sometimes, they are opportunities to find a better solution that the big model would have missed on its own.

Technical Summary: Beyond the Target: From Imitation to Collaboration in Speculative Decoding

1. Problem Statement

Speculative Decoding (SPD) is a standard technique for accelerating Large Language Model (LLM) inference. It employs a smaller "draft" model to propose a sequence of future tokens, which are then verified in parallel by a larger "target" model. The dominant paradigm in SPD relies on imitation: the draft model is trained to approximate the target model's distribution, and a draft token is accepted only if it exactly matches the target's prediction.

The authors argue that this paradigm is suboptimal because it implicitly assumes the target model is the sole authority at every token position. In reality, while the target model is superior overall, it is not uniformly better at every local decision. In tasks like reasoning and code generation, a mismatch between the draft and target does not necessarily indicate a draft error; it may represent a useful alternative path that avoids a target model mistake. Current SPD methods discard these "non-nested" strengths of the draft model, limiting the system to the performance ceiling of the target model alone.

2. Methodology: Collaborative Speculative Decoding (CoSpec)

The paper introduces Collaborative Speculative Decoding (CoSpec), a framework that shifts the relationship between draft and target models from pure imitation to active collaboration. Instead of treating mismatches as verification failures, CoSpec treats them as opportunities to select the branch (draft or target) that maximizes final task utility.

Core Architecture

CoSpec retains the standard computational structure of SPD:

Drafting: The draft model proposes a block of $K$ tokens.
Verification: The target model verifies the block in a single parallel pass.
Arbitration: At positions where the draft and target disagree (mismatches), a learned arbitration policy decides whether to accept the draft token or fall back to the target token.

The Arbitration Policy

The arbitration policy ( $\pi_\theta$ ) is a lightweight model (using LoRA adapters on the draft backbone) that observes:

The input prompt and verified prefix.
The draft block.
The target verification tokens.

It utilizes a hybrid attention mask: the context attends causally, while the draft-target region uses bidirectional attention to compare candidates. The policy outputs a probability of accepting the draft token at each mismatch.

Training via Reinforcement Learning

The policy is trained using a two-stage pipeline:

Supervised Warm-up: Initialized using judge-labeled preferences on vanilla SPD rollouts.
Reinforcement Learning (RL): Optimized using Group Relative Policy Optimization (GRPO) on full CoSpec rollouts.

Reward Shaping:
The reward function is designed to maximize final task utility (e.g., answer correctness) rather than distribution preservation.

Correct Rollouts: Rewards are proportional to the length of the accepted span beyond the first mismatch ( $\rho_r$ ), encouraging the policy to accept useful draft continuations.
Incorrect Rollouts: Penalties are applied for:
- Accepting draft tokens that are highly improbable under the target distribution (risk regularization).
- Rejecting a mismatch and falling back to a target branch that ultimately leads to an error (preventing the policy from treating fallback as a "safe" default).
Normalization: Advantages are normalized within prompt groups, with specific scaling for mixed groups of correct and incorrect rollouts to balance the magnitude of positive and negative signals.

3. Key Contributions

Paradigm Shift: The paper formulates a generalized SPD paradigm that moves from "target-bounded imitation" to "collaborative decoding," allowing the draft model to contribute when its continuation is more useful for the final task outcome.
CoSpec Framework: The introduction of a learned arbitration policy that selectively accepts draft tokens at mismatches, leveraging the complementary strengths of both models.
Theoretical Foundation: Proposition 3.1 establishes a "local complementarity gap," proving that an oracle arbitrator can strictly outperform a target-faithful verifier whenever the draft branch has higher expected utility on a non-zero subset of mismatch states.
Empirical Validation: Extensive experiments demonstrating that CoSpec achieves state-of-the-art performance in both speed and task accuracy across multiple benchmarks.

4. Experimental Results

The authors evaluated CoSpec using LLaMA-3.3-70B/8B and Qwen3-32B/4B pairs on math reasoning (GSM8K, MATH-500), code generation (HumanEval, MBPP), and general knowledge tasks.

Performance vs. Target-Only: CoSpec consistently surpasses the target model's standalone performance.
- LLaMA-70B (Greedy): Mean score improved from 90.05 (Target-only) to 91.73 with a 3.96× speedup.
- Qwen3-32B (Greedy): Mean score improved from 83.84 to 86.46 with a 2.05× speedup.
Comparison with Baselines: CoSpec outperforms target-preserving methods (SpS, REST, EAGLE-3) in accuracy and loosely verified methods (FLy) in both speed and accuracy. It also outperforms utility-improving baselines like AASD.
Complementarity Analysis: Experiments show that while the target model solves more examples overall, the "oracle union" (examples solved by either model) is consistently higher than the target alone. CoSpec recovers approximately 50–65% of these non-nested successes.
Generalization: The trained arbitrator generalizes to unseen domains (dialogue, factual QA, multilingual reasoning) without task-specific retraining, suggesting the arbitration policy learns interaction patterns rather than task-specific content.

5. Significance and Claims

The paper claims that CoSpec offers a new perspective on speculative decoding by redefining the role of the draft model. Rather than serving merely as an accelerator that approximates the target, the draft model acts as a complementary proposer.

The authors emphasize that this approach is particularly relevant for tasks where outputs are judged by answer correctness (e.g., math, code, factual QA) rather than strict adherence to the target distribution. By shifting the objective from "distribution preservation" to "final task utility," CoSpec demonstrates that it is possible to break the performance ceiling of the target model while maintaining substantial inference speedups. The work suggests that future speculative decoding systems should focus on exploiting non-nested model strengths through learned arbitration rather than strict imitation.