Speculative decoding is used to reduce the inference cost of language models. One-step diffusion drafters propose a block of multiple tokens in parallel, reducing the sequential cost of speculative decoding. On-policy distillation trains these drafters using verifier feedback on their own proposals. However, parallel drafting makes this feedback difficult to use: the drafter predicts each token without observing the preceding tokens sampled within the same block, while the verifier conditions its prediction on those tokens. Different sampled prefixes therefore give different verifier targets for the same drafter prediction. We show that distilling against these targets separately can suppress tokens that are plausible after some prefixes but unlikely after others, harming draft acceptance. We explain the effects of this mismatch and introduce Multi-Draft On-Policy Distillation (MDraft), a new parallel drafter distillation method which samples several draft blocks from the same context and aggregates the verifier's full-vocabulary conditionals across their prefixes. We use all the drafts to collect information from the verifier, allowing support from alternative prefixes to influence the training target. The resulting algorithm improves the acceptance length of a parallel drafter beyond the level of pre-training, and alleviates the effects of the distributional mismatch. MDraft changes only training, retaining the drafter's architecture and inference procedure, allowing the resulting drafter to be used as-is in place of the pretrained checkpoint. Experiments on a suite of mathematical, coding, and conversational tasks show longer accepted drafts than prior distillation methods.
Figures & tables
Figure 1: A parallel drafter predicts the second token without observing its sampled first token. The verifier’s preferred continuation depends on that first token; different training objectives therefore yield different second-token choices. The verifier’s greedy continuation is “hot chocolate.” Bars are schematic; Table 3 gives the numerical distributions.
Math
Code
Chat
Method
GSM8K
MATH500
AIME25
MBPP
HumanEval
MT-Bench
Mean
Temperature T=0
DFlash
6.469
8.023
7.379
6.023
6.581
4.119
6.432
AdaFlash
6.793
8.339
7.709
6.141
6.713
4.329
6.671
Draft-OPD
6.604
7.844
7.113
5.882
6.438
4.198
6.347
RKL
6.688
7.899
7.108
5.940
6.513
4.247
6.399
Table 1: Distilling DFlash with a Qwen3-4B teacher. Each cell reports mean acceptance length τ , including the bonus token (higher is better). Mean is the unweighted average across benchmarks. Bold and underlining mark the best and second-best results within each temperature and column.
Figure 2: Position-wise conditional acceptance gains over DFlash, in percentage points, at T=0 (left) and T=1 (right). Red: MDraft; green: AdaFlash; orange: RKL. The evaluation uses 1,656 prompts from GSM8K, MBPP and MT-Bench. Shaded bands denote 95% confidence intervals. Positions index the 15 drafted tokens. MDraft’s gains extend through the middle and later positions, with a larger advantage over AdaFlash at T=1 .
Figure 3: We fix α=0.9 and train DFlash for 20,000 iterations with varying numbers of parallel drafts K=1,2,4,8 . We evaluate acceptance length on a subset of GSM8K and MBPP.
Method
GSM8K
MATH500
AIME25
MBPP
HumanEval
MT-Bench
Mean
Qwen3-4B
Temperature T=0
DFlash
6.469
8.023
7.379
6.023
6.581
4.119
6.432
AdaFlash
6.793
8.339
7.709
6.141
6.713
4.329
6.671 (+3.7%)
MDraft
6.907
8.322
7.640
6.187
6.740
4.317
6.685 (+3.9%)
Temperature T=0.6
Table 2: Post-training with Qwen3-4B and Qwen3-8B. Each cell reports mean acceptance length; MDraft uses α=0.9 and K=8 for the 4B model and α=1.0 , K=4 for the 8B model. Bold marks the best result in each column within each model and temperature.
Verifier conditional
chocolate
cocoa
drink
brew
p(⋅∣hot) , 65%
0.42
0.31
0.27−δ
δ
p(⋅∣cold) , 35%
0.02
0.01
0.10
0.87
Objective
chocolate
cocoa
drink
brew
A0
A1
Per-prefix RKL
0.338
0.217
0.445
0.000
0
0.582
Per-prefix FKL
0.280
0.205
0.210
0.305
0
0.604
MDraft ( α=0.9 )
0.350
0.253
0.278
0.119
1
0.655
Table 3: Numerical two-branch example for Figure 1 , with δ=10−9 and fixed branch probabilities 0.65/0.35 . Top: normalized verifier conditionals. Bottom: optimal second-position draft probabilities, rounded to three decimals, and second-token acceptance AT at evaluation temperature T . Bold probabilities mark each drafter’s greedy choice; bold acceptance values mark the best of the three objectives.
Method
Weights
T=0
T=0.6
T=1
AdaFlash
Uniform
6.671
6.168
5.546
AdaFlash
Survival
6.691
6.219
5.568
MDraft
Uniform
6.568
6.260
5.836
MDraft
Survival
6.685
6.321
5.903
Table 4: Uniform versus survival weights: unweighted six-benchmark mean acceptance length τ . Means are reported to three decimal places.
Figure 4: Reverse-KL controls with α=0 and K=1 or K=8 , evaluated on a GSM8K subset at T=0 (left) and T=1 (right). Curves include the initial checkpoint and evaluations every 2,000 training steps. Dashed horizontal lines mark MDraft’s acceptance length with K=8 at step 14,000; they are endpoint references, not MDraft training trajectories.
Figure 5: Post-training acceptance length on a GSM8K subset at T=0 (left) and T=1 (right) over 34,000 training steps. Starting from the same pretrained checkpoint, FKL initially reduces acceptance length and only partially recovers, remaining below its initial value. RKL improves over initialization, while MDraft achieves higher acceptance length throughout the subsequent evaluations.
Domain
Available source pairs
Prepared training pairs
Training share
Math
410,464
410,448
61.52%
Code
135,536
128,372
19.24%
Chat
2,182,748
128,372
19.24%
Total
2,728,748
667,192
100%
Table 5: Composition of the Qwen3-8B training mixture. Counts are prompt–completion pairs.
Speculative decoding accelerates large language model inference by pairing a target model with a lightweight draft model whose proposed tokens are verified in parallel. A common way to build draft models, like EAGLE3 or DFlash is supervised fine-tuning (SFT) on target-generated trajectories. However, we observe that SFT quickly plateaus: the draft model's acceptance length on test data stops improving. The reason is an offline-to-inference mismatch: In SFT, the drafter learns from fixed target-generated trajectories, whereas during speculative decoding it is evaluated on blocks proposed under its own policy. This motivates on-policy distillation (OPD), where the target model supervises the drafter on draft-induced states. Yet OPD remains difficult for draft models, as they cannot reliably roll out complete sequences independently, whereas target-assisted generation makes the collected sequences follow the target distribution and thus eliminates the on-policy signal. We therefore propose Draft-OPD, which uses target-assisted rollout for stable continuations and replays drafting from the verification-exposed error positions. This allows the drafter to learn from target feedback on both accepted and rejected proposals, focusing training on the draft-induced errors that limit speculative acceptance. Experiments show that Draft-OPD achieves over 5× lossless acceleration for thinking models across diverse tasks, improving over EAGLE-3 and DFlash by 23% and 13%.
Haodi Lei, Yafu Li, Haoran Zhang +8
Shanghai Jiao Tong University · Shanghai AI Laboratory · The Chinese University of Hong Kong +3
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.
Yunxiao Zhao, Changxiao Cai
Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, MI, USA.
Block-diffusion speculative decoding accelerates large language model inference by drafting multiple tokens in parallel. However, one-step drafting in discrete diffusion is limited by token factorization error: each token is predicted from the same masked block, without conditioning on the tokens selected at other positions in that step. This can produce inconsistent drafts and reduce acceptance length. Iterative refinement mitigates this limitation by conditioning later predictions on previously committed tokens, though additional passes increase drafting cost. To address this trade-off between draft fidelity and inference latency, we introduce DStill, a two-stage framework that distills a target-conditioned drafter's own multi-step refinement into one-step drafting. We first train the drafter to complete partially revealed blocks, then use its refinement trajectories to train a student to predict the final completion in one pass. Using a curriculum that gradually transitions from partial to fully masked inputs, we successfully transfer the multi-step refinement capabilities into one step. Across reasoning, math, code benchmarks on three model families DStill improves acceptance length by up to 28% and speedup by up to 25% over one-step DFlash, achieving up to 6.45x speedup over auto-regressive decoding.