Суть прорыва: от дискретной диффузии к непрерывному потоку
Команда из Polymathic AI, Принстонского университета, USC ISI, NYU и Института Флэтингтона выпустила FlashMap — новый драфтер для спекулятивного декодирования больших языковых моделей (LLM). В отличие от предыдущих решений, использующих дискретную диффузию (например, DFlash), FlashMap применяет непрерывную карту потока (Flow Map). Это позволяет малой модели-драфтеру (0.5B параметров) за один прямой проход (forward pass) сгенерировать согласованный блок из 16 токенов, которые затем проверяются целевой моделью параллельно.
Ключевая проблема, которую решает FlashMap, — это «факторизованное приближение» в диффузионных драфтерах, где каждый токен предсказывается независимо. Это часто приводило к несоответствиям (например, генерация «new diego» вместо «new york»). FlashMap использует архитектуру DiT (Diffusion Transformer) с 6 блоками, которая учитывает совместные зависимости между токенами, маппируя гауссовский шум в полный блок токенов за один шаг.
Технические детали и архитектура
Драфтер FlashMap для модели Qwen3-8B имеет размер всего 500 миллионов параметров (для Qwen3-4B — 475M). Он считывает скрытые состояния с слоев 1, 9, 17, 25 и 33 целевой модели. Обучение проходило в два этапа:
- Flow teacher: непрерывный поток по одномерным токенам, обученный с перекрестной энтропией.
- Flow map student: дистилляция полугруппы, превращающая поток в 2-временной денуйзер, способный перейти от шума к данным за 1 шаг.
Общий объем вычислений для обучения всех моделей и абляций составил около 4,600 GPU-часов на NVIDIA H200. Датасет для обучения — 500K выборок из Open-PerfectBlend.
Результаты бенчмарков: скорость против точности
Исследователи протестировали FlashMap на 9 бенчмарках (математика, код, чат) при температурах 0 и 1. На модели Qwen3-8B FlashMap показал лучшие результаты среди всех протестированных драфтеров. При температуре 0 (детерминированный режим) на задаче GSM8K достигнуто ускорение 6.14x при длине принятия токенов 7.69. Это означает, что целевая модель обрабатывает 7.69 токенов за каждый шаг драфтера.
| Метрика / Модель | FlashMap (Qwen3-8B) | DFlash (Qwen3-8B) | EAGLE-3 (Qwen3-8B) |
|---|---|---|---|
| Скорость (GSM8K, T=0) | 6.14x | 5.77x | 3.20x |
| Среднее ускорение (9 задач, T=0) | 4.36x | 4.07x | 2.55x |
| Длина принятия (Chat, T=0) | На 11.1% выше DFlash | База | Информация недостаточна |
| Размер драфтера | 500M | 1B | 302M (переобучен) |
| Стиль драфтинга | Непрерывный поток (1 шаг) | Блочная диффузия (1 шаг) | Авторегрессионное дерево |
Практическое применение и развертывание
FlashMap доступен под лицензией Apache-2.0 (модели) и MIT (код). Для использования достаточно установить пакет через uv pip install flashmap. Поддерживается интеграция с популярными движками:
- vLLM 0.31: настройка через
--speculative-configс параметромnum_speculative_tokens=16. - SGLang 0.5.20: использование алгоритма
DSPARKс путем драфтера FlashMap.
Исследователи отмечают, что использование режима 2-шагового драфтинга позволяет еще больше выровнять длину принятия токенов по всему блоку, хотя и требует дополнительных вычислений. Для температур выше 0 в vLLM рекомендуется метод сэмплирования "draft_sample_method": "probabilistic".
Источник: MarkTechPost ↗