Draft-OPD : Distillation on-policy pour modèles de draft spéculatifs

Résumé

Le décodage spéculatif accélère l'inférence des grands modèles de langage en associant un modèle cible à un modèle draft léger dont les tokens proposés sont vérifiés en parallèle. Une méthode courante pour construire des modèles draft, comme EAGLE3 ou DFlash, consiste en un fine-tuning supervisé (SFT) sur des trajectoires générées par la cible. Cependant, nous observons que le SFT atteint rapidement un plateau : la longueur d'acceptation du modèle draft sur les données de test cesse de s'améliorer. La raison en est un décalage entre l'apprentissage hors ligne et l'inférence : lors du SFT, le draft apprend à partir de trajectoires fixes générées par la cible, tandis que pendant le décodage spéculatif, il est évalué sur des blocs proposés selon sa propre politique. Cela motive la distillation on-policy (OPD), où le modèle cible supervise le draft sur des états induits par le draft. Pourtant, l'OPD reste difficile pour les modèles draft, car ils ne peuvent pas produire de manière fiable des séquences complètes de façon autonome, tandis que la génération assistée par la cible fait que les séquences collectées suivent la distribution cible, éliminant ainsi le signal on-policy. Nous proposons donc Draft-OPD, qui utilise un déroulement assisté par la cible pour des continuations stables et rejoue le draft à partir des positions d'erreur exposées à la vérification. Cela permet au draft d'apprendre à partir du retour de la cible sur les propositions acceptées et rejetées, en concentrant l'entraînement sur les erreurs induites par le draft qui limitent l'acceptation spéculative. Les expériences montrent que Draft-OPD atteint une accélération sans perte de plus de 5 fois pour les modèles de réflexion sur diverses tâches, améliorant les performances d'EAGLE-3 et DFlash de 23 % et 13 %.

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