Source-linked AI summary
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
TL;DR
Draft-model SFT plateaus because fixed target-generated trajectories do not match the draft-induced states that determine speculative acceptance. Draft-OPD replays drafting from verification-exposed error positions, achieving over 5× lossless acceleration for thinking models and improving over EAGLE-3 and DFlash by 23% and 13%.
Problem
Draft-model SFT plateaus because fixed target-generated trajectories mismatch the draft-induced states that determine speculative acceptance.
Method
Draft-OPD combines target-assisted rollouts with error-position replay and acceptance-aware distillation on accepted and rejected draft proposals.
Results
Over 5× lossless acceleration is achieved for thinking models, with improvements over EAGLE-3 and DFlash of 23% and 13% under matched FLOPs.
Takeaways & Limitations
Draft-OPD improves acceptance length and translates those gains into practical serving throughput across evaluated models, tasks, and concurrency levels.
Takeaways & Limitations
Experiments focus on Qwen3 models and the DFlash-style architecture, so transfer to other model families, draft architectures, and inference backends remains untested.
Abstract
from arXiv · showhide
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\%.
1 Introduction
Speculative decoding reduces inference cost by verifying lightweight draft-model proposals with a larger target model, but offline SFT plateaus because training trajectories differ from draft-induced inference states. Draft-OPD addresses this mismatch through target-assisted rollout and error-position replay, achieving substantial lossless acceleration.
- Speculative decoding uses a lightweight draft model to propose tokens that a larger target model verifies in parallel, preserving the target output distribution.
- Offline SFT quickly plateaus because drafters train on fixed target trajectories but are evaluated on blocks induced by their own policy.Continuing SFT can leave test acceptance length fluctuating around a fixed plateau, while SFT on OPD data can reduce it.
- Standard OPD is difficult for draft models because draft-only rollouts are unstable, whereas target-assisted rollouts remove the on-policy signal.
- Draft-OPD combines target-assisted rollout with error-position replay to train on verification-time states where the drafter acted.It includes both accepted and rejected proposals by replaying drafting from error positions and scoring the resulting prefixes with the target model.
- Over 5× lossless acceleration is achieved by Draft-OPD for thinking models, improving over EAGLE-3 and DFlash by 23% and 13% under matched FLOPs.
2 Related Work
Speculative decoding uses a lightweight draft model to propose tokens that a larger target model verifies in parallel while preserving the target distribution. Related work improves draft modeling through target features, block-level architectures, and distillation focused on matching draft and target behavior.
- Speculative decoding: Speculative decoding preserves the target distribution by verifying multiple tokens proposed by a lightweight draft model with a larger target model in parallel.This framework accelerates autoregressive inference.
- Draft-model improvements: EAGLE leverages feature-level context from the frozen target model to improve speculative decoding.EAGLE is cited as a representative feature-level approach.
- Draft-model improvements: DFlash improves speculative decoding through a stronger block-level draft architecture.The passage identifies DFlash as a block-level architectural approach.
- Distillation: Distillation methods train draft models to better match the target distribution, with DistillSpec emphasizing on-policy data and task-specific divergence choices.Temperature-centric distillation further indicates that matching training and decoding configurations can improve speculative decoding under challenging sampling settings.
3 Preliminary
Speculative decoding pairs a target model with a lightweight drafter whose token blocks are verified in parallel, preserving the target distribution while reducing expensive decoding steps. On-policy distillation addresses the mismatch between fixed target-generated training trajectories and draft-induced states encountered during speculative decoding.
- Speculative Decoding: Speculative decoding pairs target model pθ with lightweight draft model qϕ, which proposes K-token blocks for parallel target verification.The verifier accepts the longest valid prefix of each drafted block.
- Speculative Decoding: Parallel verification preserves the target model’s output distribution while reducing the number of expensive target-model decoding steps.Generation continues from the verified prefix.
- Speculative Decoding: Accepted length τ measures the number of draft tokens accepted per verification round and reflects draft–target alignment and decoding efficiency.Higher τ means each target-model verification generates more tokens.
- On-Policy Distillation: On-policy distillation trains students on states induced by their own policies, directly targeting exposure mismatch between offline supervised training and inference-time generation.The teacher supplies next-token distributions at student-induced states.
- On-Policy Distillation: For speculative decoding, SFT uses fixed target-generated trajectories, whereas OPD supervises draft-induced states encountered during draft verification.This targets prefixes on which the drafter itself produces blocks.
4 Method
Draft-OPD addresses the instability and missing draft-policy signal that make direct OPD unsuitable for speculative draft models. It combines target-assisted rollouts with error-position replay and an acceptance-aware objective to train on both accepted and rejected draft proposals.
- Motivation: Direct OPD is unsuitable because draft modules produce repetitive or degenerate full-sequence rollouts, while strict target-assisted verification removes the draft-policy signal.Target-assisted sequences follow the target distribution, retain accepted tokens, and discard rejected proposals that expose draft errors.
- Rollout with error-position collection: Draft-OPD uses speculative decoding to collect stable target-quality rollouts while recording each drafted block’s start position as an anchor.Anchors preserve local draft actions without requiring standalone full-sequence generation and reuse target hidden states from the same rollout.
- Replay for Log-Probability Computation: Replay from every anchor computes student and teacher token-level log-probabilities on the draft-generated prefixes, including proposals rejected during verification.This exposes the target model to draft-induced states rather than only final rollout tokens.
- Acceptance-Aware Distillation Objective: Draft-OPD uses forward KL for accepted tokens and reverse KL for rejected tokens, reflecting agreement states versus target-disagreed draft modes.Verification partitions each block into accepted and rejected token sets.
- Acceptance-Aware Distillation Objective: Rejected-token weights decay exponentially across block positions because earlier errors prevent the verifier from using the remaining suffix and affect acceptance length more.The final objective averages accepted-token and rejected-token losses, with λ_acc = λ_rej = 1 in all experiments.
5 Experiment
Experiments across Qwen3 models and reasoning, coding, and software-engineering benchmarks show that Draft-OPD improves acceptance length and decoding speed over EAGLE-3 and DFlash under matched training FLOPs. Ablations and serving tests indicate that its gains depend on preserving draft-induced errors and replaying informative failure states.
- Models and tasks: Experiments use Qwen3-4B, Qwen3-8B, and Qwen3-30B-A3B-Thinking-2507 across mathematical reasoning, code generation, and software-engineering benchmarks.Benchmarks include GSM8K, MATH-500, AIME, MBPP, HumanEval, and SWE-bench Lite.
- Baselines: Draft-OPD is compared with EAGLE-3 and DFlash under approximately matched total training FLOPs.Both baselines use the DFlash data mixture, while their SFT budgets match Draft-OPD’s combined SFT and OPD budget.
- Thinking Mode Enabled: At temperature 0 with thinking enabled, Draft-OPD raises average τ from 5.35 for DFlash to 5.85 and achieves 4.88× average speedup.This improves over EAGLE-3 and DFlash by 23% and 13%, respectively.
- Decoding results: With thinking disabled, Draft-OPD maintains an average acceptance length of 6.33 and achieves a 5.17× average speedup across temperatures and model sizes.At temperature 0.6 with thinking enabled, it remains fastest with 4.17× average speedup.
- Performance on SGLang: On SGLang, Draft-OPD improves acceptance length by 11.2% on average and achieves up to a 17% speedup on Qwen3-30B-A3B-Thinking.Throughput gains persist under concurrency up to 32, with the average relative gain higher at concurrency 32 than at concurrency 1.
- Ablations: Ablations show that acceptance-aware KL directions, error-position replay, and rejected-token position decay each outperform their alternatives, while naive target-assisted rollout reduces average speedup from 4.63× to 4.29×.The relative drop is 7.3%, indicating that stable target-assisted rollouts alone are insufficient without preserving draft-induced errors.
6 Conclusion
Draft-OPD is an on-policy distillation framework for training-based draft models that combines target-assisted rollouts with error-position replay. It also uses an acceptance-aware objective to learn from both accepted and rejected draft tokens while preserving the draft-policy errors that determine speculative acceptance.
- Framework: Draft-OPD trains draft models with target-assisted rollouts and an error-position replay mechanism.This keeps training samples stable while preserving draft-policy errors relevant to speculative acceptance.
- Framework: The framework preserves the draft-policy errors that determine speculative acceptance.Its replay mechanism focuses training on error positions exposed during verification.
- Distillation objective: An acceptance-aware distillation objective treats accepted and rejected draft tokens differently.This enables learning from both reliable proposals and informative failures.
Limitations
Draft-OPD’s limitations concern training coverage, evaluation scope, and its restriction to lossless speculative decoding. Longer rollouts, broader model and system evaluations, and approximate verification remain future directions.
- Training Length: Training rollouts are capped at 4096 tokens, whereas evaluation uses 8192 tokens, potentially leaving late-stage states in long generations underrepresented.Scaling OPD training to longer rollouts could expose the drafter to more verification-time errors and further improve draft-target alignment.
- Evaluation Scope: Experiments primarily use Qwen3 models, the DFlash-style parallel draft architecture, and SGLang, so transferability to other families, architectures, and backends remains uncertain.The work includes comparisons with EAGLE-3, but further evaluation is needed across broader model and deployment settings.
- Lossless Decoding: Draft-OPD targets efficiency under lossless verification, leaving extension of on-policy draft-model training to approximate or lossy verification as future work.Lossless speculative decoding preserves the target model distribution, so Draft-OPD is designed to improve decoding efficiency rather than generation quality.
A Training Details
Draft-OPD initializes its OPD stage from a six-epoch SFT checkpoint, then trains for eight epochs with specified sequence lengths, optimization settings, and equal accepted/rejected objective weights. EAGLE-3 and DFlash baselines instead use ten epochs of supervised training and select their best-evaluation checkpoints.
- SFT stage: Draft-OPD initializes OPD from the draft-model checkpoint after 6 SFT epochs, while EAGLE-3 and DFlash use 10 supervised-training epochs and report their best-evaluation checkpoints.The SFT configuration largely follows DFlash and uses SpecForge.
- OPD stage: Draft-OPD uses rejected-token position weights with γ = 0.8 and trains on the Section 5 OPD mixture for 8 epochs.The OPD stage is implemented with verl.
- OPD stage: Draft-OPD uses maximum response lengths of 4096 tokens for thinking-enabled and 2048 tokens for thinking-disabled training.Optimization uses AdamW with a learning rate of 3 × 10^-4, cosine scheduling, and a 0.05 warmup ratio.
- OPD stage: The final Draft-OPD objective sets λ_acc = λ_rej = 1 in all experiments.The accepted and rejected objective weights are equal across experiments.
B Loss Design for Accepted and Rejected Draft Tokens
Draft-OPD uses forward KL on accepted replay positions and reverse KL on rejected positions because the former provide target-weighted supervision at verified states, while the latter expose draft-induced errors. This acceptance-aware loss distinguishes reliable target states from draft-proposed modes that disagree with the target.
- Objective justification: For a fixed replay state and unconstrained draft distribution, forward KL, reverse KL, and the mixed objective share the optimum q = p.Their difference is how finite-capacity optimization weights target-supported versus draft-proposed tokens.
- Accepted positions: Accepted replay positions use target-weighted cross-entropy because the draft proposal has passed target verification and the state is reliable.The accepted-position objective is equivalent, up to an additive constant, to minimizing forward KL.
- Rejected positions: Rejected replay positions use draft-weighted disagreement because failed or invalidated draft suffixes provide draft-induced error signals.Reverse KL directly penalizes high-probability draft modes that disagree with the target distribution.
- Acceptance-aware loss: The acceptance-aware Draft-OPD loss aggregates forward KL for accepted positions with reverse KL for rejected positions.This decomposition applies each KL direction to the replay positions whose supervision matches its weighting scheme.
C Thinking-Mode Drafting Gap
Thinking-mode responses exhibit higher token-level next-token NLL than non-thinking responses across evaluated datasets, revealing a drafting gap. Draft-OPD partially improves acceptance by training on verification-time errors, while reasoning-specific draft-model methods remain future work.
- C Thinking-Mode Drafting Gap: Qwen3-4B was evaluated in thinking and non-thinking modes on ShareGPT, AoPS, and math and code splits of Nemotron-Post-Training-Dataset-v2.The study computes next-token NLL at each generated-response token position.
- C Thinking-Mode Drafting Gap: Higher next-token NLL occurs in thinking-mode responses than non-thinking responses across the evaluated datasets.Figure 5 reports token-level negative log-likelihood under both modes.
- C Thinking-Mode Drafting Gap: Draft-OPD partially mitigates the challenge by post-training draft models on verification-time errors, improving acceptance for reasoning-oriented decoding.Because Draft-OPD is designed as a general post-training framework across decoding settings, reasoning-specific draft-model approaches remain an important future direction.