Technical Summary: Draft-OPD

Problem Statement

Speculative decoding (SD) accelerates Large Language Model (LLM) inference by using a lightweight draft model to propose token blocks, which are then verified in parallel by a larger target model. While current mainstream methods (e.g., EAGLE-3, DFlash) train these draft models via Supervised Fine-Tuning (SFT) on fixed trajectories generated by the target model, this approach hits a performance ceiling.

The authors identify a fundamental offline-to-inference mismatch:

1. Training Phase: The draft model learns from static trajectories where every prefix is generated by the target model.
2. Inference Phase: The draft model is evaluated on blocks it proposes itself. The acceptance length (the number of tokens accepted per verification step) is determined by the draft model's own policy-induced states, not the static states seen during SFT.

Consequently, continued SFT quickly plateaus, and the draft model fails to improve its acceptance length on test data. Furthermore, standard On-Policy Distillation (OPD) is difficult to apply directly to draft models because:

• Draft-only rollouts (where the draft model generates sequences independently) are unstable and often degenerate into repetitive outputs.
• Target-assisted rollouts (where the target model corrects the draft) produce high-quality sequences but eliminate the "on-policy" signal. The target model discards the draft's errors, meaning the draft model only learns from the target's distribution, failing to address the specific errors that limit speculative acceptance.

Methodology: Draft-OPD

To bridge this gap, the authors propose Draft-OPD, an on-policy distillation framework designed specifically for training-based draft models. The method enables the draft model to learn from its own inference-time states, including both successful proposals and the specific errors that cause verification failures.

The framework consists of three coordinated components:

1. Rollout with Error-Position Collection

Instead of forcing the draft model to generate full sequences alone, Draft-OPD uses target-assisted rollouts to ensure stable, high-quality continuations.

• During speculative decoding, the draft model proposes a block of tokens, and the target model verifies them.
• The system records the starting position of each drafted block as an anchor.
• This process generates a valid rollout (following the target distribution) while preserving the specific states where the draft model took action.

2. Replay for Log-Probability Computation

To recover the "on-policy" signal lost during the target-assisted rollout, the system replays the drafting process from the recorded anchors.

• Starting from an anchor state (the prefix verified by the target), the system re-simulates the draft model's proposal of the subsequent block.
• It computes log-probabilities for both the draft model (student) and the target model (teacher) on these draft-generated prefixes.
• Crucially, this allows the system to evaluate the target model's distribution on tokens that were rejected during the initial verification, capturing the specific errors that the draft model made.

3. Acceptance-Aware Distillation Objective

The training objective treats accepted and rejected tokens differently to optimize for speculative acceptance:

• Accepted Tokens: For tokens that pass verification, the method uses Forward KL Divergence ($D_{KL}(P_{target} || P_{draft})$). This encourages the draft distribution to cover the target distribution where they already agree, reinforcing reliable behavior.
• Rejected Tokens: For tokens that fail verification, the method uses Reverse KL Divergence ($D_{KL}(P_{draft} || P_{target})$). This penalizes the draft model for assigning high probability to tokens the target model rejects, effectively teaching the draft model to avoid its own high-confidence errors.
• Weighting: Rejected tokens are weighted exponentially based on their position within the block ($w_k = \gamma^{k-1}$). Errors occurring earlier in a block are weighted more heavily because they invalidate the entire subsequent suffix, having a larger impact on the overall acceptance length.

Key Contributions

1. Identification of Limitation: The paper demonstrates that offline SFT for draft models plateaus quickly due to the mismatch between training states (target-generated) and inference states (draft-generated).
2. Analysis of OPD Challenges: It explains why standard OPD fails for draft models: independent rollouts are unstable, while naive target-assisted rollouts erase the on-policy error signal.
3. Draft-OPD Framework: It introduces a novel framework combining target-assisted rollouts with error-position replay, enabling efficient post-training on verification-time errors.
4. Acceptance-Aware Objective: It proposes a hybrid KL-divergence objective that differentiates between accepted and rejected tokens to specifically target the errors limiting speculative speedup.

Experimental Results

The authors evaluated Draft-OPD on the Qwen3 family (4B, 8B, and 30B-A3B-Thinking) across mathematical reasoning (GSM8K, MATH-500, AIME), code generation (MBPP, HumanEval, SWE-bench Lite), and general chat (MT-Bench).

• Performance Gains: Under a matched training FLOPs budget, Draft-OPD consistently outperformed EAGLE-3 and DFlash.
  • It achieved over 5× lossless acceleration for thinking models.
  • It improved average acceptance length and speedup by 23% over EAGLE-3 and 13% over DFlash.
• Throughput: In SGLang deployment benchmarks, Draft-OPD showed consistent throughput gains across concurrency levels (1 to 32), with average acceptance length improvements of ~11.2% over DFlash.
• Ablation Studies:
  • Simply applying SFT on OPD-generated data (without the replay mechanism) failed to improve performance, confirming the necessity of the on-policy error signal.
  • Using a single KL direction (all-forward or all-reverse) was inferior to the mixed, acceptance-aware objective.
  • Randomly selecting replay anchors was less effective than selecting anchors based on error positions, highlighting the importance of focusing training on failure states.

Significance and Claims

The paper claims that Draft-OPD offers a significant step forward in optimizing training-based draft models. By shifting from static, offline supervision to dynamic, on-policy distillation that explicitly targets verification-time errors, the method breaks the performance plateau of traditional SFT.

The authors emphasize that their approach preserves the lossless nature of speculative decoding (maintaining the target model's output distribution) while substantially improving efficiency. They position this work as a practical solution for accelerating "thinking" models and complex reasoning tasks, where standard draft models often struggle due to the high uncertainty of long reasoning traces. The paper concludes that post-training draft models on their own inference-time errors is a critical direction for improving speculative decoding, offering a path to higher acceptance lengths and faster inference without compromising generation quality.