Draft-OPD: On-Policy Distillatie voor Speculatieve Draftmodellen

Samenvatting

Speculatieve decodering versnelt de inferentie van grote taalmodellen door een doelmodel te combineren met een lichtgewicht conceptmodel waarvan de voorgestelde tokens parallel worden geverifieerd. Een gebruikelijke manier om conceptmodellen zoals EAGLE3 of DFlash te bouwen, is supervised fine-tuning (SFT) op door het doelmodel gegenereerde trajecten. We merken echter dat SFT snel een plateau bereikt: de acceptatielengte van het conceptmodel op testdata stopt met verbeteren. De reden is een mismatch tussen offline en inferentie: bij SFT leert het conceptmodel van vaste, door het doelmodel gegenereerde trajecten, terwijl het tijdens speculatieve decodering wordt geëvalueerd op blokken die onder zijn eigen beleid worden voorgesteld. Dit motiveert on-policy distillatie (OPD), waarbij het doelmodel het conceptmodel begeleidt op door het concept geïnduceerde toestanden. Toch blijft OPD moeilijk voor conceptmodellen, omdat ze niet betrouwbaar volledige reeksen zelfstandig kunnen uitrollen, terwijl doelondersteunde generatie ervoor zorgt dat de verzamelde reeksen de doelverdeling volgen en daarmee het on-policy signaal elimineert. Daarom stellen we Draft-OPD voor, dat gebruikmaakt van doelondersteunde uitrol voor stabiele voortzettingen en het concept opnieuw afspeelt vanuit de verificatie-blootgestelde foutposities. Hierdoor kan het conceptmodel leren van feedback van het doelmodel op zowel geaccepteerde als afgewezen voorstellen, waarbij de training zich richt op de door het concept veroorzaakte fouten die de speculatieve acceptatie beperken. Experimenten tonen aan dat Draft-OPD een meer dan vijfvoudige verliesvrije versnelling realiseert voor denkmodellen bij diverse taken, met een verbetering van 23% ten opzichte van EAGLE-3 en 13% ten opzichte van DFlash.

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