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

Abstract: Speculative decoding accelerates LLM 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\%.

• The paper introduces a novel on-policy distillation framework, Draft-OPD, that systematically addresses the offline-to-inference policy mismatch in speculative decoding.
• It leverages error-aware post-training on verification-time state distributions to optimize draft block acceptance lengths and decoding speed.
• Empirical results demonstrate up to 23% improvement over EAGLE-3 and 13% over DFlash, enhancing both throughput and model alignment in practical LLM deployments.

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

Overview

This paper introduces Draft-OPD, a principled on-policy distillation (OPD) framework tailored for the training of speculative draft models for large-scale LLM inference. The central challenge addressed is the suboptimal plateau encountered when training draft models for speculative decoding using conventional supervised fine-tuning (SFT) on target-generated trajectories. Draft-OPD systematically resolves the offline-to-inference policy mismatch inherent in this paradigm by enabling error-focused post-training on verification-time state distributions. Empirical results demonstrate substantial improvements in speedup and draft-target alignment over strong baselines, EAGLE-3 and DFlash, with up to 23% and 13% relative gains, respectively, under matched computational budgets.

Motivation and Background

Speculative decoding has become a prominent technique for improving the inference efficiency of autoregressive LLMs. The deployed system architecture pairs a large target model with a lightweight draft model; the draft model proposes candidate spans, which the target model verifies in parallel. The core determinant of speculative speedup is the draft model's acceptance length $\tau$—the count of tokens in a draft block accepted in each verification round. Existing approaches, such as EAGLE-3 and DFlash, use SFT on target-model generations to train such drafters. However, SFT on static datasets fails to expose the draft model to the full inference-time state distribution, leading to an acceptance-length plateau even as training continues.

The exposure bias—where the drafter is trained on prefixes sampled from the teacher and evaluated on its self-induced prefixes—is especially acute for block-level, training-based draft architectures like DFlash or EAGLE. These drafters are structurally incapable of rolling out complete, high-quality autoregressive sequences independently, causing degenerate off-policy rollouts if standard OPD is naively applied.

Draft-OPD: Methodology

Draft-OPD introduces an efficient on-policy distillation algorithm leveraging speculative decoding rollouts to collect informative error signals (Figure 1).

Figure 1: The Draft-OPD algorithm collects anchor positions from speculative rollouts, replays drafted blocks from these anchors, and computes student and teacher log-probabilities for both accepted and rejected tokens.

The method is composed of three essential components:

1. Rollout with Error-Position Collection: During speculative decoding, each draft block's start index is recorded as an anchor. The sequence of anchors localizes draft-induced deviations that are critical for improving acceptance rates.
2. Replay for Log-Probability Computation: From each anchor, the algorithm replays the drafting process, evaluating both the student (draft model) and teacher (target model) next-token distributions for all positions, including those rejected during verification. This exposes both successful and unsuccessful policy choices.
3. Acceptance-Aware Distillation Objective: Draft-OPD employs asymmetric divergence minimization: a forward KL ($D_{KL}(p \| q)$) for accepted tokens, encouraging coverage of the target, and a reverse KL ($D_{KL}(q \| p)$) for rejected tokens, penalizing over-confident errors. Moreover, losses for later rejected tokens in a block are down-weighted exponentially, focusing optimization on early, more impactful mistakes.

This workflow circumnavigates the instability of draft-only rollouts and the over-filtered signals of naive target-assisted rollouts.

Experimental Evaluation

Draft-OPD was evaluated on Qwen3 series models, with benchmarks spanning mathematical reasoning, code generation, and general-purpose dialogue (e.g., GSM8K, MATH-500, AIME, MBPP, HumanEval, SWE-Lite, MT-Bench). Key metrics reported include speedup ratio over standard autoregressive decoding and the average draft block acceptance length $\tau$.

• With thinking mode enabled, Draft-OPD increased mean $\tau$ from 5.35 (DFlash baseline) to 5.85, yielding a 23% improvement over EAGLE-3 and a 13% improvement over DFlash across seven tasks under greedy decoding.
• In SGLang deployment, Draft-OPD delivered throughput improvements of 8–17% at batch concurrency levels up to 32, showing highly favorable scaling for practical serving scenarios.
• Ablation studies confirmed that the asymmetric KL loss and error-position anchor selection are crucial; replacing the loss with all-forward or all-reverse KL, or randomly choosing anchors, degraded performance.

Theoretical and Practical Implications

Draft-OPD presents a formal resolution to the off-policy mismatch in training-based speculative draft models, where standard OPD is rendered ineffective by the draft model’s inability to generate high-quality, long sequences unaided. By combining target-assisted rollouts with error-focused replay and an acceptance-aware divergence objective, Draft-OPD establishes an efficient and robust training protocol.

Practically, this technique results in lossless $5\times$ decoding speedup for "thinking" tasks, where draft-target alignment is intrinsically more difficult due to higher response uncertainty and longer traces. The method's architecture-agnostic concepts should translate to more advanced speculative draft models and complex generation settings.

Limitations and Future Directions

Current experiments are largely limited to Qwen3 model variants and the DFlash-style parallel draft architecture, with training rollouts capped at 4096 tokens due to compute constraints. Extension to other LLMs, alternative speculative decoding backends, and even lossy verification (for approximate decoding settings) present important avenues for investigation. Draft-OPD is designed as a general framework, but further tailoring for highly uncertain, multi-path reasoning traces seen in chain-of-thought tasks may unlock further gains.

Conclusion

Draft-OPD advances speculative decoding by shifting draft-model training from offline SFT or naive OPD to a fine-grained, error-aware, on-policy framework that leverages real inference-time errors. Its methodologically principled, empirically validated improvements in both acceptance length and end-to-end speedup profile it as a robust foundation for next-generation accelerated LLM inference architectures (2605.29343).