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

2026-05-28

AI Summary

This work addresses the rapid saturation of acceptance length in existing draft models for speculative decoding, which stems from a mismatch between offline training and online inference strategies. To overcome this limitation, the authors propose Draft-OPD, the first method to enable effective online policy distillation for draft models. Draft-OPD leverages the target model to identify erroneous prediction positions and replays the draft generation process at those locations, allowing the draft model to learn directly from its own inference mistakes. By integrating target-assisted sequence generation, error-position replay, and an online distillation mechanism, Draft-OPD substantially improves token acceptance rates. Experiments demonstrate that the method achieves over 5× lossless speedup across diverse tasks, outperforming EAGLE-3 and DFlash by 23% and 13% in inference efficiency, respectively.

Abstract

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\%.

Problem

• speculative decoding
• on-policy distillation
• draft model
• distribution mismatch
• acceptance length

Innovation

• Speculative Decoding
• On-Policy Distillation
• Draft Model
• Target-Assisted Rollout
• Error-Replay Training