Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds
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.
Related
- Divergence-Minimizing Distribution Matching for Parallel Channels and Channels with Memory
- Minimizing the Arithmetic and Communication Complexity of Jacobi's Method for Eigenvalues and Singular Values: Part Two -- Parallel Algorithms
- GpFlash: Maximizing Speed and Minimizing Memory for Gaussian Process Training and Prediction
- Parallel Optimized Particle Swarm Optimization: An Efficient Parallel Approach for Neural Network Training
- The Limits of Speculation: Bounding Speculative Decoding in Mixture-of-Experts
BibTeX
@misc{zhao2026training,
title = {{Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds}},
author = {Yunxiao Zhao and Changxiao Cai},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.10411},
eprint = {2610.10411},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.10411},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds
AU - Yunxiao Zhao
AU - Changxiao Cai
PY - 2026
DA - 2026/10/07
PB - arXiv
DO - 10.48550/arxiv.2610.10411
UR - https://arxiv.org/abs/2610.10411
AB - 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.
ER -