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

AI summary

Overview

Research area: Machine learning systems for large language model (LLM) inference acceleration, specifically online model selection within speculative decoding, sitting at the intersection of online learning theory and practical LLM serving.

Technical level: Advanced. The paper combines regret analysis (no-regret online learning, delayed feedback reductions, counterfactual estimators) with systems-level latency measurements on multi-GPU hardware, and assumes familiarity with speculative decoding mechanics (acceptance probability, total variation distance, draft trees).

Scope: The paper proposes HedgeSpec, a full-information online learning method for selecting among multiple draft models during speculative decoding, with no-regret guarantees and end-to-end speedup measurements on three target LLMs and 21 curated drafters.

What This Paper Is About

Speculative decoding speeds up LLM inference by having a small "draft" model guess several tokens that a larger target model then verifies in parallel. No single drafter works well on every task: a code drafter excels on code but collapses on medical questions, causing long-tail latency. The paper asks how to automatically pick the best drafter for each incoming query, and shows that unlike prior work (which treated this as a multi-armed bandit where only the chosen drafter's feedback is observed), it is possible to evaluate all drafters from a single verified trajectory without extra target-model calls, turning the problem into full-information online learning that converges exponentially faster in the number of drafters N.

Key Contributions

1. Full-information framework for multi-drafter decoding: The authors design an evaluation phase that accurately estimates how every drafter in a pool would have performed on the realized target trajectory, not just the drafter that was chosen, at very low overhead and without additional queries to the target model.
2. Theoretical guarantees: They formulate drafter selection under two objectives — token acceptance probability and expected acceptance length — and prove no-regret guarantees via an off-policy (counterfactual) estimator plus a reduction for delayed feedback.
3. A counterfactual acceptance-length estimator (Theorem 3): The "one-step counterfactual" estimator is unbiased in expectation for the number of accepted tokens, with variance bounded by K²/4, compared to O(NK²) for the BanditSpec EXP3-Spec estimator, which grows with N.
4. System-efficient implementation and empirical validation: They build practical versions of the online learner and show consistent improvements over EAGLE and BanditSpec across LLaMA-3.1-8B-IT and Qwen-3-8B/32B with seven drafters per target model, reporting up to 83.7% token/s gain in a single domain and up to 46.1% on average over EAGLE, and up to 49% MAT gain over bandit baselines.

Main Findings

• Panoramic feedback beats bandit feedback: Because HedgeSpec evaluates all drafters from a single verified trajectory, it adapts far faster than bandit methods. Cumulative regret on LLaMA-3.1 with a Python workload shows bandit methods accumulating regret quickly, while HedgeSpec converges to near-zero regret within a handful of steps.
• End-to-end gains on LLaMA-3.1-8B-IT: HedgeSpec reached an average MAT of 7.15 and 90.41 Token/s, versus EAGLE at 5.69 and 74.34, UCBSpec at 5.09 and 68.89, and EXP3Spec at 4.86 and 65.22.
• End-to-end gains on Qwen-3-8B: HedgeSpec reached an average MAT of 6.37 and 69.44 Token/s, versus EAGLE at 4.23 and 47.53, UCBSpec at 4.58 and 52.72, and EXP3Spec at 4.37 and 49.75.
• End-to-end gains on Qwen-3-32B: HedgeSpec reached an average MAT of 6.21 and 40.41 Token/s, versus EAGLE at 2.88 and 20.76, UCBSpec at 4.50 and 33.24, and EXP3Spec at 4.33 and 31.55.
• Largest single-domain improvement: On SQL requests with Qwen, HedgeSpec improved MAT from 4.2 to 7.52 (a 79% gain) and Token/s from 44.6 to 81.94 (an 83.7% gain).
• Average improvement: Across all mixed queries, HedgeSpec achieved a 46.1% average improvement over EAGLE.
• Specialists alone are weaker than the generalist: The 7 curated drafters each perform strongly in-domain (the diagonal of Table 1) but degrade outside their specialization, and on average perform worse than the vanilla EAGLE model.
• Longer reasoning helps adaptation: Qwen-3 series reasoning models tend to produce longer outputs than LLaMA-3.1-IT — 1.64× length in the Math workload — giving learners more time to converge. Bandit methods underperform EAGLE on LLaMA-3.1-8B-IT but outperform it on Qwen-3-8B for this reason, and longer reasoning chains amplify HedgeSpec's advantage.
• Scalability with pool size: As the number of candidate drafters increases, bandit methods deteriorate sharply while HedgeSpec scales gracefully, remaining nearly unaffected because global information lets it rapidly identify the best drafter.
• Robustness under distribution shift: Compared to a static offline router, HedgeSpec held up on out-of-distribution workloads. On MedQA with LLaMA-3.1-8B-IT, the static router dropped to MAT 3.03 / 40.83 Token/s while HedgeSpec reached 5.56 / 70.38; on MedQA with Qwen-3-8B, the router fell to 2.49 / 28.18 while HedgeSpec reached 6.04 / 65.99.
• Low evaluation overhead: Per-component timings for LLaMA-3.1-8B-IT were 75.709 ms for the Llama forward pass, 2.497 ms for the EAGLE forward pass, and 0.413 ms for the hedge update. Evaluating a drafter costs roughly 1/25 of a target forward, so securing just one additional MAT offsets evaluating up to 25 drafters in sequence; in practice the cost is lower because drafter evaluations run in parallel.

Methodology in Plain English

The setup: a pool of N drafters, and at each step the system must choose one to propose tokens while a large target model verifies them. Prior work framed this as a bandit: you only learn how the drafter you picked did, so you must spend effort exploring alternatives. The key observation here is that speculative decoding is lossless — the verified output is always distributed according to the target model — so the trajectory that was actually produced can be reused as counterfactual evidence for every other drafter. Each other drafter is "prefilled" with the verified tokens, and its probability of accepting each token is computed from quantities already available, such as 1 − TV(p, q) for standard speculative decoding or the total probability of the draft tree's child nodes for EAGLE. Because this requires no extra target-model calls, the problem becomes full-information online learning rather than a bandit.

The authors then define a loss for each drafter — either 1 − γ (for acceptance probability) or a chunk-length-based loss built from a weighted sum of per-step acceptance probabilities (for acceptance length) — and feed these losses to a hedging algorithm, specifically NormalHedge, with Hedge as an alternative base learner. A subtlety the paper highlights is a censoring problem: tokens are not observed one at a time but in verified chunks, and alternative drafters might accept more tokens than the selected one, so the needed feedback arrives late. The authors handle this by casting it as delayed feedback and using a black-box reduction (Joulani et al., 2013) that tolerates bounded, non-constant delay, with a maximum delay of 2K. In practice they found the "stochastic setting" variant of that reduction worked best, maintaining a queue of feedback and applying the base algorithm's next available action. The core theoretical result is that as long as a single target-rolled-out trajectory exists, an unbiased estimator of acceptance length can be computed for any alternative drafter, avoiding the combinatorial number of rollouts that direct computation would require.

For evaluation, the authors fine-tuned 21 drafters of the EAGLE-3 family (seven per target model) on seven open-source datasets spanning Python, Math, Biology, Chemistry, MedicalQA, CNN_DM, and SQL, using SpecForge as the training pipeline. All experiments used FP16 precision, batch size 1, and nodes with 8 NVIDIA A100 GPUs connected via NVLink, reporting Mean Number of Accepted Tokens (MAT) alongside wall-clock token throughput.

Why This Matters

Impact on research: This work reframes a problem the community had modeled as a multi-armed bandit and shows the bandit formulation is unnecessarily pessimistic — full information is available for free given the structure of speculative decoding. It connects inference-systems engineering to classical online learning with delayed feedback, and it opens a line of work on how the "not a bandit" insight generalizes to other adaptive inference decisions.

Real-world applications (derived from the domains and drafters studied):

• Serving mixed workloads where code, math, scientific, medical, and summarization queries arrive through the same endpoint, so a single generalist drafter leaves substantial latency on the table.
• Code assistants and SQL generation tools, where the largest single-domain gain in the paper (SQL with Qwen: 79% MAT gain, 83.7% token/s gain) directly targets interactive latency.
• Retrieval-augmented and summarization workloads, where the paper notes retrieval-based drafters work well only when outputs closely match the input.
• Long-chain reasoning services, where Qwen-3-style models generate longer outputs (1.64× length in the Math workload) and HedgeSpec's advantage grows because there is more time to converge.

Industry relevance: The method is orthogonal to how drafters are built, so it can be layered on top of existing EAGLE-3 deployments without changing the drafter training pipeline. The overhead analysis — a drafter forward pass costing roughly 1/25 of a target forward pass — makes the economics concrete for latency-sensitive serving: reduced target calls can more than pay for evaluating a whole pool of drafters. It also offers a practical alternative to maintaining a static router, which the paper shows mis-routes requests under distribution shift.

Future Directions

• Integration with broader speculative decoding frameworks: The paper states that additional experiments and discussion on integrating HedgeSpec with a wider range of speculative decoding frameworks are deferred to Appendix D.9, and the truncated content does not report those results.
• Jointly trained drafters: The paper defers discussion of jointly trained drafters to Appendix D.5; how selection interacts with a shared training process is left open in the content provided.
• Hedging algorithm variants: Additional hedging variants and their trade-offs are deferred to Appendix D.6, suggesting the choice of base learner is not fully settled.
• Game formulation choices: The paper notes in Appendix B.3 that each round can be a token or a chunk, and the loss can target acceptance length or acceptance probability, with different pros and cons — the most natural setting was chosen, but the alternatives remain live design questions. Heuristics for practical implementation are also deferred to Appendix B.5.

Target Audience

Researchers and practitioners working on LLM inference efficiency and speculative decoding will benefit most, particularly those who need to serve heterogeneous query mixes and want principled rather than heuristic drafter selection. It also suits theoretically inclined readers interested in applied online learning, delayed feedback, and off-policy estimation, since the paper's central technical move is showing that an existing bandit formulation can be upgraded to full-information learning with provable guarantees. Engineers deploying EAGLE-style speculative decoding will find the overhead breakdown and the A100-based end-to-end tables directly actionable; readers without a background in speculative decoding or regret analysis should expect a steep curve, as the paper assumes fluency in both.

Authors’ abstract

Speculative decoding is widely used in accelerating large language model (LLM) inference. In this work, we focus on the online draft model selection problem in speculative decoding. We design an algorithm that provably competes with the best draft model in hindsight for each query in terms of either the token acceptance probability or expected acceptance length. In particular, we show that we can accurately evaluate all draft models, instead of only the chosen model without incurring additional queries to the target model, which allows us to improve exponentially over the existing bandit-based approach as the number of draft models increases. Our approach is generically applicable with any speculative decoding methods (single draft, multi-drafts and draft-trees). Moreover, we design system-efficient versions of online learners and demonstrate that the overhead in computation and latency can be substantially reduced. We conduct extensive experiments on open-source LLMs and diverse datasets, demonstrating that our methods substantially outperform the state-of-the-art EAGLE3 and the BanditSpec baseline in a variety of domains where specialized domain-expert drafters are available, especially when long reasoning chains are required.