Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs

1. Problem Statement

Context: Speculative decoding accelerates Large Language Model (LLM) inference by using a smaller "drafter" model to predict tokens, which are then verified in parallel by a larger "target" model.
The Challenge: A single drafter often performs well only in specific domains (e.g., code, math) and poorly in others. While domain-specific drafters exist, dynamically selecting the best one for each incoming query is difficult.
Existing Approach: Previous works (e.g., MetaSD, BanditSpec) frame this as a Multi-Armed Bandit (MAB) problem. They balance exploration (trying different drafters) and exploitation (using the best known drafter).
Limitation of MAB: Bandit algorithms suffer from high regret because they only receive feedback on the chosen drafter (partial information). They must waste time exploring suboptimal drafters to learn their performance, leading to slower convergence and lower efficiency, especially as the number of candidate drafters ($N$) increases.

Core Question: Can we design an algorithm that competes with the best drafter in hindsight for every query without the high cost of exploration required by bandit methods?

2. Methodology: HedgeSpec

The authors propose HedgeSpec, a framework that transforms the drafter selection problem from a Bandit setting (partial information) to a Full-Information Online Learning setting.

Key Insight: Counterfactual Evaluation

The paper's central observation is that exploration is unnecessary because the target model's verification process provides feedback for all drafters, not just the one selected.

• Mechanism: After a chunk of tokens is generated by a chosen drafter and verified by the target model, the system "prefills" the verified trajectory into all other candidate drafters in the pool.
• Result: This allows the system to compute the hypothetical acceptance probability and expected acceptance length for every drafter in the pool based on a single target model pass. This is done without incurring additional calls to the expensive target model.

Algorithmic Framework

1. Full-Information Feedback: Instead of observing a loss only for the selected action $i_t$, the algorithm observes a loss vector $\mathbf{f}_t \in [0,1]^N$ for all $N$ drafters.
2. Loss Functions:
   • Acceptance Probability: Loss = $1 - \gamma_t[i]$ (where $\gamma$ is the acceptance probability).
   • Expected Acceptance Length (EAL): The paper derives an unbiased estimator for the expected number of accepted tokens for any drafter given a single verified trajectory (Theorem 3).
3. Handling Delayed Feedback (Censoring):
   • In speculative decoding, feedback for a specific token is only available after the chunk is verified. Furthermore, if a chosen drafter stops early, the system cannot immediately know how many tokens a different drafter would have accepted (censoring).
   • Solution: The authors model this as an online learning problem with delayed feedback. They utilize algorithms (specifically NormalHedge or Hedge) combined with a reduction technique (Joulani et al., 2013) to handle the delay. The regret bound increases only mildly with the delay.
4. Theoretical Guarantee: The method achieves no-regret guarantees. The average performance converges to that of the best fixed drafter in hindsight at a rate of $O(\sqrt{\frac{K \log N}{T}})$, where $K$ is the draft depth and $T$ is the number of steps. This is exponentially faster in $N$ than bandit-based approaches.

3. Key Contributions

1. Full-Information Framework: First to demonstrate that drafter selection can be treated as a full-information problem by leveraging the structure of speculative decoding to evaluate all drafters simultaneously via counterfactual analysis.
2. Theoretical Guarantees: Proved no-regret bounds for both acceptance probability and expected acceptance length objectives, even under delayed feedback and censoring conditions.
3. System Efficiency: Designed a practical implementation where the overhead of evaluating all drafters is negligible (roughly 1/25th of a target forward pass) because drafters are lightweight and evaluations can be parallelized.
4. Robustness to Distribution Shift: Unlike static routers (classifiers trained offline), HedgeSpec adapts online to Out-of-Distribution (OOD) prompts without prior knowledge.

4. Experimental Results

The authors evaluated HedgeSpec on LLaMA-3.1-8B-IT, Qwen-3-8B, and Qwen-3-32B with pools of 7 domain-specialized drafters (Python, Math, Biology, Chemistry, MedQA, CNN DM, SQL).

• Performance vs. Baselines:
  • vs. EAGLE (State-of-the-art single drafter): HedgeSpec achieved up to 83.7% higher token throughput (Token/s) and 46.1% average improvement across mixed domains.
  • vs. BanditSpec (UCB/EXP3): HedgeSpec significantly outperformed bandit methods, achieving up to 49% higher Mean Accepted Tokens (MAT) and 41% higher Token/s.
• Scalability: As the number of drafters increases, bandit methods degrade rapidly due to exploration costs. HedgeSpec scales gracefully, maintaining near-zero regret and high efficiency.
• OOD Robustness: In tests with prompts that shifted the distribution (e.g., requesting longer reasoning chains), static offline routers failed (misrouting 90%+ of queries), while HedgeSpec adapted dynamically, achieving up to 2.34x gains over the offline router.
• Overhead: The evaluation overhead is minimal. The gain in acceptance rate from selecting the optimal drafter easily offsets the cost of evaluating the entire pool.

5. Significance

• Paradigm Shift: Moves the field from "exploration-heavy" bandit approaches to "full-information" learning for speculative decoding, fundamentally changing how multi-drafter systems are optimized.
• Practical Impact: Enables the use of diverse, specialized "expert" drafters in production LLM serving without the penalty of slow adaptation. This is crucial for domains requiring long reasoning chains or specific expertise (e.g., coding, medical QA).
• Generality: The approach is agnostic to the underlying speculative decoding method (works with single-draft, multi-draft, and draft-trees like EAGLE) and can be applied to any set of candidate drafters.

In summary, HedgeSpec proves that by leveraging the inherent structure of speculative decoding, one can obtain "panoramic feedback" on all candidate models, allowing for provably optimal, low-latency selection that outperforms both static experts and traditional bandit-based adaptive methods.