Multi-Draft On-Policy Distillation for Diffusion Drafters

Petr Molodyk; Yuchen Zhu; Jaemoo Choi; Molei Tao; Yongxin Chen

cs.CLOct 27, 2026

Multi-Draft On-Policy Distillation for Diffusion Drafters

Authors: Petr Molodyk, Yuchen Zhu, Jaemoo Choi, Molei Tao, Yongxin Chen

Organizations: Georgia Institute of Technology

Abstract

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

\tau

, 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

\alpha=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

\alpha=0.9

and

K=8

for the 4B model and

\alpha=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(\cdot\mid\text{hot})$ , $65\%$	$0.42$	$0.31$	$0.27-\delta$	$\delta$
$p(\cdot\mid\text{cold})$ , $35\%$	$0.02$	$0.01$	$0.10$	$0.87$
Objective	chocolate	cocoa	drink	brew	$A_{0}$	$A_{1}$
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 ( $\alpha=0.9$ )	0.350	0.253	0.278	0.119	1	0.655

Table 3: Numerical two-branch example for Figure 1 , with

\delta=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

A_{T}

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

\tau

. Means are reported to three decimal places.

Figure 4: Reverse-KL controls with

\alpha=0

and

K=1

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 for Diffusion Language Models

Speculative decoding for diffusion language models speeds generation by drafting tokens in parallel and verifying them efficiently.

See all 35 papers

Draft-OPD: On-Policy Distillation for Speculative Draft Models
May 28, 2026Haodi Lei, Yafu Li, Haoran Zhang +8Speculative Decoding LLM Inference Acceleration
Shanghai Jiao Tong University · Shanghai AI Laboratory · The Chinese University of Hong Kong +3
Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds
Oct 7, 2026Yunxiao Zhao, Changxiao CaiSpeculative Decoding Parallel Speculative Decoding
Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, MI, USA.
DStill: Trajectory Consistency Distillation for Accelerated One-Step Drafting in Block Diffusion Speculative Decoding
Oct 19, 2026Shreyas Kulkarni, Uttam Gurram, Madhav Kanda +5Speculative Decoding for Diffusion Language Models LLM Inference Acceleration
AI Foundations, Capital One

May 28, 2026J/K move · Enter open · S save

Draft-OPD: On-Policy Distillation for Speculative Draft Models

Haodi Lei, Yafu Li, Haoran Zhang, Shunkai Zhang+7

Shanghai Jiao Tong University · Shanghai AI Laboratory · The Chinese University of Hong Kong+3

Speculative Decoding LLM Inference Acceleration Language Model Distillation On-Policy Distillation

Open paper AskPDF arXiv

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\times

lossless acceleration for thinking models across diverse tasks, improving over EAGLE-3 and DFlash by 23% and 13%.

Abstract

Speculative Decoding for Diffusion Language Models

Explore similar work

Draft-OPD: On-Policy Distillation for Speculative Draft Models

Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds

DStill: Trajectory Consistency Distillation for Accelerated One-Step Drafting in Block Diffusion Speculative Decoding