Draft-OPD: On-Policy дистилляция для спекулятивных драфт-моделей

Аннотация

Спекулятивное декодирование ускоряет вывод больших языковых моделей, объединяя целевую модель с легковесной моделью-черновиком, чьи предлагаемые токены верифицируются параллельно. Распространённым способом построения моделей-черновиков, таких как EAGLE3 или DFlash, является контролируемая тонкая настройка (SFT) на траекториях, порождённых целевой моделью. Однако мы обнаружили, что SFT быстро выходит на плато: длина принятия модели-черновика на тестовых данных перестаёт улучшаться. Причина заключается в несоответствии между обучением на фиксированных данных и инференсом: в SFT модель-черновик обучается на фиксированных траекториях целевой модели, тогда как при спекулятивном декодировании она оценивается на блоках, предложенных в соответствии с её собственной политикой. Это мотивирует применение дистилляции на политике (OPD), при которой целевая модель контролирует черновик на состояниях, порождённых им самим. Однако OPD остаётся сложной для моделей-черновиков, поскольку они не могут надёжно самостоятельно разворачивать полные последовательности, в то время как генерация с помощью целевой модели заставляет собранные последовательности следовать целевому распределению и тем самым устраняет сигнал на политике. Поэтому мы предлагаем Draft-OPD, который использует развёртывание с помощью целевой модели для устойчивых продолжений и повторяет генерацию черновика с позиций ошибок, выявленных при верификации. Это позволяет черновику учиться на обратной связи от целевой модели как по принятым, так и по отклонённым предложениям, сосредотачивая обучение на ошибках, вызванных черновиком, которые ограничивают спекулятивное принятие. Эксперименты показывают, что Draft-OPD достигает более чем пятикратного ускорения без потери качества для моделей рассуждения на различных задачах, превосходя EAGLE-3 и DFlash на 23% и 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\%.