Search papers, labs, and topics across Lattice.
This paper introduces Draft-OPD, an innovative approach to on-policy distillation for speculative decoding in large language models, addressing the limitations of traditional supervised fine-tuning (SFT) which plateaus in performance. By employing target-assisted rollout and focusing on draft-induced errors, Draft-OPD enables the draft model to learn effectively from both accepted and rejected proposals, enhancing its ability to propose tokens in parallel with the target model. Experimental results demonstrate that Draft-OPD achieves over five times lossless acceleration in inference speed while outperforming existing models like EAGLE-3 and DFlash by 23% and 13%, respectively.
Draft-OPD accelerates inference by over 5x while improving speculative decoding accuracy, transforming how draft models learn from target feedback.
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 5times lossless acceleration for thinking models across diverse tasks, improving over EAGLE-3 and DFlash by 23\% and 13\%.
Tsinghua AI