Draft-OPD: On-Policy Distillation für spekulative Draft-Modelle

Zusammenfassung

Spekulative Dekodierung beschleunigt die Inferenz großer Sprachmodelle, indem ein Zielmodell mit einem leichten Entwurfsmodell kombiniert wird, dessen vorgeschlagene Token parallel verifiziert werden. Eine gängige Methode zur Erstellung von Entwurfsmodellen wie EAGLE-3 oder DFlash ist das überwachte Feintuning (SFT) auf zielgenerierten Trajektorien. Wir beobachten jedoch, dass SFT schnell ein Plateau erreicht: Die Akzeptanzlänge des Entwurfsmodells bei Testdaten verbessert sich nicht mehr. Der Grund ist eine Offline-zu-Inferenz-Diskrepanz: Beim SFT lernt das Entwurfsmodell aus festen, zielgenerierten Trajektorien, während es während der spekulativen Dekodierung auf Blöcken evaluiert wird, die unter seiner eigenen Policy vorgeschlagen wurden. Dies motiviert die On-Policy-Destillation (OPD), bei der das Zielmodell das Entwurfsmodell auf entwurfsinduzierten Zuständen überwacht. Dennoch bleibt OPD für Entwurfsmodelle schwierig, da sie keine vollständigen Sequenzen zuverlässig eigenständig durchführen können, während die zielunterstützte Generierung dazu führt, dass die gesammelten Sequenzen der Zielverteilung folgen und somit das On-Policy-Signal eliminieren. Daher schlagen wir Draft-OPD vor, das eine zielunterstützte Durchführung für stabile Fortsetzungen nutzt und das Entwerfen von den verifikations-exponierten Fehlerpositionen wiederholt. Dies ermöglicht es dem Entwurfsmodell, aus Zielrückmeldungen sowohl über akzeptierte als auch über abgelehnte Vorschläge zu lernen, wobei das Training auf die entwurfsinduzierten Fehler fokussiert wird, die die spekulative Akzeptanz einschränken. Experimente zeigen, dass Draft-OPD für Denkmodelle bei verschiedenen Aufgaben eine über 5-fache verlustfreie Beschleunigung erreicht und damit eine Verbesserung von 23 % bzw. 13 % gegenüber EAGLE-3 und DFlash erzielt.

English

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