Ключевой идеей работы является использование представлений более мощной таргет - модели для драфта. А именно, берутся активации с разных слоев таргет-модели, конкатенируются, прогоняются через некий адаптер и подаются через kv-кэш (причем во все блоки).
Все токены (в пределах расшумляемого блока) предсказываются за один проход сети.
Процедура обучения выглядит следующим образом:
📌 Сэмплируются anchor токены начала блока
📌 Первый токен в последовательности - чистый, остальные заменяются на
[MASK]. Это отличается от типичного обучения диффузии, где зашумляются токены на случайных позициях.Кроме того, утверждается что более ранние токены более важны для генерации, поэтому лосс убывает по экспоненте от начала расшумляемого блока.
Голова и эмбеддинги берутся от таргет-модели.
🧪 Эксперименты
Учат драфт-модели поверх Llama-3.1 Instruct (8B) , Qwen-3 (4B, 8B, Coder-30B-A3B-Instruct). Драфт-модель - диффузионка (той же архитектуры что и базовая модель?) с 5 блоками, 8 для Qwen-3-Coder. Учат драфт на смеси примеров из Nemotron Post-Training Dataset V2 и Code Alpaca.
Замеряют ускорение (качество не надо, ибо lossless) на разных бенчах на математику/код и MTBench. Оно дает ускорение до 6 раз против ~2 раз у EAGLE-3. Acceptance length тоже заметно больше.
На бенчах на ризонинг тоже все хорошо.
Есть и SGLang интеграция, дающая ускорение до 4х раз на батче 1, и до 2 с чем-то раз на батче 32.
В ablation показывают, что выбранный размер драфта более-менее оптимален по скорости / acceptance. Больше фичей от таргета улучшает качество, но повышает расход памяти. Выбранный размер блока расшумления 16 тоже около оптимален.
💡 Выводы
Если все на самом деле так хорошо, то выглядит просто как имба. Надо учить диффузионный драфтер, но это все равно дешево по сравнению с обучением базовой авторегрессии. Команды, выкладывающие новые модели, могли бы выкладывать их вместе с таким драфтером. Оригинальный репозиторий уже содержит большое количество моделей разного размера - вплоть до Kimi K2.6.