Draft-OPD: Destilación On-Policy para Modelos de Draft Especulativos

Resumen

La decodificación especulativa acelera la inferencia en modelos de lenguaje grandes al emparejar un modelo objetivo con un modelo borrador ligero cuyas propuestas de tokens se verifican en paralelo. Un método común para construir modelos borrador, como EAGLE3 o DFlash, es el ajuste fino supervisado (SFT) sobre trayectorias generadas por el modelo objetivo. Sin embargo, observamos que el SFT se estanca rápidamente: la longitud de aceptación del modelo borrador en datos de prueba deja de mejorar. La razón es un desajuste entre el entrenamiento offline y la inferencia: en el SFT, el borrador aprende a partir de trayectorias fijas generadas por el modelo objetivo, mientras que durante la decodificación especulativa se evalúa sobre bloques propuestos bajo su propia política. Esto motiva la destilación on-policy (OPD), donde el modelo objetivo supervisa al borrador en estados inducidos por este último. Sin embargo, la OPD sigue siendo difícil para los modelos borrador, ya que no pueden desplegar secuencias completas de forma fiable e independiente, mientras que la generación asistida por el modelo objetivo hace que las secuencias recogidas sigan la distribución de este, eliminando así la señal on-policy. Por lo tanto, proponemos Draft-OPD, que utiliza un despliegue asistido por el modelo objetivo para continuaciones estables y reproduce el borrador desde las posiciones de error expuestas durante la verificación. Esto permite que el borrador aprenda de la retroalimentación del modelo objetivo tanto en propuestas aceptadas como rechazadas, centrando el entrenamiento en los errores inducidos por el borrador que limitan la aceptación especulativa. Los experimentos muestran que Draft-OPD logra una aceleración sin pérdidas de más de 5 veces para modelos de razonamiento en diversas tareas, mejorando en un 23% y un 13% respecto a EAGLE-3 y DFlash, respectivamente.

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