Исследователи из Polymathic AI, Принстонского университета, USC ISI, NYU и Института Флатирон опубликовали FlashMap — «драфтер» для спекулятивного декодирования языковых моделей. По словам команды, он может ускорить генерацию на Qwen3-8B до 6,14 раза, не изменяя получаемый результат. Веса драфтера для Qwen3-8B и Qwen3-4B доступны на Hugging Face.
Спекулятивное декодирование заключается в том, что небольшой вспомогательный модель предлагает несколько токенов (единиц текста, обрабатываемых моделью), которые затем проверяются основной моделью за один проход. FlashMap предлагает целый блок из 16 токенов сразу, и, по словам авторов, токены в этом блоке остаются согласованными между собой. Драфтер содержит 0,5 миллиарда параметров (500 миллионов для Qwen3-8B и 475 миллионов для Qwen3-4B) и работает на том же GPU, что и целевая модель.
Почему драфтеры при диффузии теряют точность?
Драфтер DFlash предсказывает каждую позицию в блоке независимо — это исследователи называют факторизованным приближением. В источнике приводится пример последовательности слов, которая с равной вероятностью может быть «new york» или «san diego». Независимый драфтер тогда в половине случаев предложит «new diego» или «san york», и целевая модель отвергнет эти последовательности на второй позиции.
FlashMap основан на Flow Map Language Models. Каждый шумовой сэмпл соответствует целому блоку, и этот шум совместно определяет выбранную версию для каждой позиции. Драфтер представляет собой диффузионный трансформер (DiT) из 6 блоков с размерностью эмбеддинга 2048 и 16 головами внимания. Он считывает скрытые состояния слоев 1, 9, 17, 25 и 33 целевой модели.
Заявленные результаты по девяти бенчмаркам
Исследователи переобучили EAGLE-3 и DFlash на тех же данных и в тех же условиях, а затем протестировали драфтеры на девяти бенчмарках, охватывающих математику, код и чат, при температурах 0 и 1. По данным источника, одноэтапный FlashMap является самым быстрым драфтером в каждой протестированной задаче на Qwen3-8B при обеих температурах. Он превосходит DFlash по длине принятия в 15 конфигурациях из 18, имея примерно вдвое меньше параметров.
В двухэтапном режиме длина принятия превышает показатель DFlash в 18 конфигурациях из 18, в среднем на 10% лучше. Авторы также указывают, что на Qwen3-4B при температуре 0 FlashMap уступает DFlash в 4 из 9 задач. Наименьший заявленный прирост составляет 2,28 раза на Alpaca при температуре 1. Модель обучалась на 500 000 сэмплов из Open-PerfectBlend, что в сумме заняло около 4600 часов вычислений на H200.
Для развертывания пакет flashmap требует только torch и transformers версии 4.56 или выше, а проекты указывают пути использования через vLLM 0.31 и SGLang 0.5.20. Веса опубликованы под лицензией Apache-2.0, а код — под лицензией MIT.
