Draft-OPD: On-Policy Distillation for Speculative Draft Models
Haodi Lei, Yafu Li, Haoran Zhang, Shunkai Zhang, Qianjia Cheng, Xiaoye Qu, Ganqu Cui, Bowen Zhou, Ning Ding, Yun Luo, Yu Cheng
cs.CL
2026-05-28
Draft-OPD post-trains speculative drafters on rejected proposals via error-position replay, hitting 4.88× thinking-mode speedup and beating DFlash by 13% at matched FLOPs.
Speculative decoding speeds up LLM inference by letting a cheap draft model propose a block of tokens that a larger target model verifies in parallel. Speedup tracks accepted length τ: the closer the drafter is to the target, the more tokens each verification consumes, and the fewer expensive target forwards are needed.
The usual recipe for training-based drafters such as EAGLE-3 and DFlash is supervised fine-tuning on target-generated trajectories. After a short warmup, extra offline SFT stops lifting test-set accepted length, and SFT on the later OPD prompt pool can even make it worse. Compute is not the bottleneck. SFT only sees prefixes the target produced. At inference, the target verifies blocks the drafter proposed. Those two state distributions are not the same.
On-policy distillation would supervise those draft-induced states. Training-based drafters cannot use the standard recipe. They are built to propose short blocks under target verification, so independent rollouts collapse into repetition. Target-assisted generation keeps sequences usable, but verification throws away rejected tokens and returns the trajectory to the target distribution. The errors that cap acceptance never enter the loss.
Draft-OPD keeps target-assisted rollouts for sequence quality, then replays drafting from verification-exposed positions so the student still trains on its own actions.
Forward KL is target-weighted and fits verified positions. Reverse KL is student-weighted and hits draft modes that disagree with the target. All-forward treats errors like reliable states; all-reverse treats already-accepted positions like failures.
OPD starts from a DFlash drafter after 6 SFT epochs and runs 8 more. Baselines spend the same FLOPs on 10 SFT epochs, using the DFlash data mixture. The OPD prompt pool has 16K items (2K GSM8K, 5K MATH, 4K AoPS, 5K CodeAlpaca); only questions are kept, and responses are generated online by the target. AdamW at 3×10^{-4}. SFT uses SpecForge; OPD uses verl. Drafters are 5 Transformer layers for Qwen3-4B/8B and 8 layers for Qwen3-30B-A3B.
On Qwen3-4B and 8B with thinking on, temperature 0, and matched training FLOPs, mean speedup across GSM8K, MATH-500, AIME25, MBPP, HumanEval, SWE-bench Lite, and MT-Bench:
| Model | Method | Mean speedup | Mean τ |
| Qwen3-4B | EAGLE-3 | 3.87× | 5.33 |
| Qwen3-4B | DFlash | 4.33× | 5.51 |
| Qwen3-4B | Draft-OPD | 4.86× | 5.96 |
| Qwen3-8B | EAGLE-3 | 4.06× | 5.64 |
| Qwen3-8B | DFlash | 4.34× | 5.19 |
| Qwen3-8B | Draft-OPD | 4.89× | 5.73 |
The two-model average is 4.88×, 23% above EAGLE-3 and 13% above DFlash. Mean τ rises from DFlash's 5.35 to 5.85. At temperature 0.6 Draft-OPD is still fastest at 4.17×. EAGLE-3 posts the highest τ on Qwen3-8B at temperature 0.6, but sequential drafting costs wall-clock time, which is why OPD is applied to DFlash-style parallel drafting.
With thinking off, mean τ is 6.33 and mean speedup is 5.17×. Qwen3-8B reaches 7.64× on MATH-500 at temperature 0. Generation quality is not reported: only the drafter is post-trained, the verification protocol is unchanged, and the output distribution remains the target's.
On SGLang with the FA3 backend, throughput gains hold from concurrency 1 to 32. For Qwen3-30B-A3B-Thinking on AIME25 at concurrency 32, Draft-OPD moves 4014 tok/s to 4718 tok/s (+17%), with τ 4.54 → 5.32. Mean τ across the reported model-task pairs rises 11.2%.
Ablations on Qwen3-4B thinking. Extra SFT on the OPD prompt pool does not explain the gain: MATH-500 is 5.55× for Draft-OPD versus 5.14× for DFlash plus OPD-data SFT. Dropping position decay, using all-forward KL, all-reverse KL, or random anchors cuts MATH-500 to 5.13×, 5.34×, 5.11×, and 5.04×. Naive target-assisted rollout that discards rejected proposals drops average speedup from 4.63× to 4.29× (−7.3%).
For teams already running EAGLE or DFlash, the SFT plateau looks more like a training-distribution problem than a capacity wall. Replaying verification-time errors extracts more speedup at matched FLOPs without changing the decoding protocol.
The ready-to-use setting is narrow: Qwen3, DFlash-style parallel drafters, lossless speculative decoding. Code and weights are public. Relative SGLang gains do not shrink at higher concurrency, so this is not a single-request scoreboard trick.
The contribution is a post-training recipe, not a new decoding algorithm. In speculative decoding, changing the training signal is often cheaper than stacking another draft architecture.
Thinking-mode training is capped at 4096 tokens while evaluation goes to 8192, so late-trace verification errors may never enter training. Almost all results are Qwen3 plus a DFlash backbone; transfer to other families or to EAGLE-style sequential drafters is untested. The method only covers lossless verification.
The abstract's "over 5× for thinking models" is looser than the body. The thinking, temperature-0 average across seven tasks is 4.88×; numbers above 5× are per-task, such as 5.55× and 5.80× on MATH-500. Thinking-off is where 5× becomes typical. The 16K OPD prompts are almost all math and code, and MT-Bench remains the slowest setting at about 3.18× with thinking on. Accepted-length variance is also missing: a worse drafter would make the target fall back more often, and that wall-clock jitter is not quantified.