Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds
cs.LG, cs.AI, cs.CL, stat.ML
Submitted: 2026-10-07
Updated: 2026-10-07
Code: https://github.com/y-x-zhao/AngelSpec-EDR
License: http://creativecommons.org/licenses/by/4.0/
The gist: Speculative decoding accelerates large language model inference by using a low-cost draft model to propose tokens that the full-size target model verifies in parallel.
Terminology
Abstract
Speculative decoding accelerates large language model inference by using a low-cost draft model to propose tokens that the full-size target model verifies in parallel. Parallel and semi-autoregressive (semi- AR) drafters improve drafting efficiency by proposing an entire block in a single forward pass, but training them raises a new difficulty: the draft distribution for a given position depends on where the decoding round starts, and where rounds start depends on how many tokens earlier rounds accepted. Existing training objectives typically rely on block-local surrogates that ignore this cross-round coupling, and therefore do not directly optimize the global decoding efficiency. In this work, we develop a theoretical framework for training and evaluating these drafters by representing speculative decoding as a Markov reward process. This formulation yields the Expected Decoding Rounds (EDR) objective, which weights local rejection costs by state occupancies and exactly equals the expected number of decoding rounds. Unlike prior surrogate objectives, EDR introduces no auxiliary hyperparameters. We then derive an exact temporal-difference gradient that supports unbiased stochastic optimization from target-model rollouts. The same framework also yields an exact offline evaluator for round counts, enabling paired drafter comparisons on shared target rollouts without running speculative decoding. Finetuning two state-of-the- art drafters, DSpark and DFly, with EDR consistently improves mean accepted length and outperforms existing training objectives across nine benchmarks spanning math reasoning, code generation, and chat.
Sources
- Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding
- Program Synthesis with Large Language Models
- Confidence-Based Decoding is Provably Efficient for Diffusion Language Models
- Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
- Accelerating Large Language Model Decoding with Speculative Sampling
- DFlash: Block Diffusion for Flash Speculative Decoding
- Evaluating Large Language Models Trained on Code
- Optimal Inference Schedules for Masked Diffusion Models
- Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
- LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling
- DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation
- Training Verifiers to Solve Math Word Problems
- Efficient Sampling with Discrete Diffusion Models: Sharp and Adaptive Guarantees
- Verification-Aware Training for Speculative Decoding
- Measuring Mathematical Problem Solving With the MATH Dataset
- Classifier-Free Diffusion Guidance
- ECHO: Elastic Speculative Decoding with Sparse Gating for High-Concurrency Scenarios
- SpecDec++: Boosting Speculative Decoding via Adaptive Candidate Lengths
- Mercury: Ultra-Fast Language Models Based on Diffusion
- Draft-OPD: On-Policy Distillation for Speculative Draft Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks