Polymathic AI та дослідники з Princeton представили FlashMap — модель-чернетку Flow Map для декодування LLM у 6,14 раза швидше

Дослідники з Polymathic AI, Принстонського університету, USC ISI, NYU та Flatiron Institute випустили FlashMap. Спекулятивне декодування FlashMap замінює дискретний дифузійний чернетковий модуль на безперервне відображення потоку. Він створює повний блок із 16 токенів за 1 прямий прохід. Токени в цьому блоці залишаються узгодженими між собою. Для Qwen3-8B дослідницька команда повідомляє про до 6.14x прискорення без втрати точності на GSM8K. Ваги чернеткового модуля для Qwen3-8B і Qwen3-4B доступні на Hugging Face, разом із плагінами для vLLM і SGLang.

Коротко

Розмір: чернетковий модуль із 0.5B параметрів (500M для Qwen3-8B, 475M для Qwen3-4B); створює блоки по K = 16 токенів.
Працює на: тому самому GPU, що й цільова модель. У статті прискорення вимірювалися на 1 NVIDIA H200 NVL за паралелізму 1, у bf16.
Продуктивність: найшвидший чернетковий модуль, протестований на Qwen3-8B, у середньому приблизно на 7% швидший за перенавчений DFlash.
Найкращий результат: 6.14x прискорення на GSM8K за температури 0 (довжина прийняття — 7.69 токена).
Найгірший результат: 2.28x прискорення на Alpaca за температури 1. На Qwen3-4B (температура 0) він поступається DFlash у 4 із 9 завдань.

Що таке FlashMap?

FlashMap — це чернетковий модуль для спекулятивного декодування, який перетворює гаусівський шум на блок токенів за 1 крок. Невеликий чернетковий модуль пропонує K токенів. Потім цільова модель перевіряє всі K токенів за 1 паралельний прохід і залишає найдовший префікс, що збігається. Результат залишається точно таким, яким його створила б цільова модель.

Швидкість залежить від 2 речей: скільки створених токенів буде прийнято та наскільки дешево їх було створено. EAGLE-3 створює токени авторегресивно, тому вартість зростає разом із довжиною блоку. DFlash створює повний блок за 1 прохід із використанням блокової дифузії.

Чому дифузійні чернеткові модулі втрачають прийняття?

DFlash прогнозує кожну позицію в блоці незалежно. Дослідницька команда називає це факторизованим наближенням. Уявімо продовження, яким із однаковою ймовірністю можуть бути “new york” або “san diego”. Незалежний чернетковий модуль у половині випадків запропонує “new diego” або “san york”. Цільова модель відхилить їх на 2-й позиції.

FlashMap розвиває ідеї Flow Map Language Models. Кожен зразок шуму відображається на 1 повний блок. Шум спільно визначає модус для кожної позиції.

Як працює FlashMap?

Чернетковий модуль — це дифузійний трансформер із 6 блоків (DiT), розмірністю ембедингів 2048 і 16 головами уваги. Він зчитує приховані стани з шарів цільової моделі 1, 9, 17, 25 і 33. Матриця ембедингів і LM-голова копіюються з цільової моделі та залишаються замороженими.

Навчання проходить у 2 етапи:

Вчитель потоку: безперервний потік над one-hot блоками токенів, навчений за допомогою перехресної ентропії. Перепараметризація часу враховує словник Qwen3 із 151 936 токенів.
Учень відображення потоку: дистиляція напівгрупи перетворює потік на денойзер із 2 часовими точками. Він переходить від шуму до даних за 1 крок.

Обидва чернеткові модулі навчаються на 500K зразках із Open-PerfectBlend. Загальні обчислення для всіх моделей і абляцій становили близько 4 600 GPU-годин на H200.

Як FlashMap показує себе на бенчмарках?

Дослідницька команда перенавчила EAGLE-3 і DFlash на тих самих даних і з такими самими налаштуваннями. Тести охоплюють 9 бенчмарків із математики, програмування та чатів за температур 0 і 1.

Швидкість: однокроковий FlashMap — найшвидший чернетковий модуль у кожному завданні Qwen3-8B за обох температур.
Точність: однокроковий FlashMap перевершує DFlash за довжиною прийняття у 15 із 18 налаштувань. Він має приблизно вдвічі менше параметрів.
Найбільший приріст у чатах: за температури 0 перевага над DFlash становить 5.5% у математиці та 11.1% у чатах.
Двокроковий режим: довжина прийняття перевищує показник DFlash у 18 із 18 налаштувань, у середньому на 10%.
EAGLE-3: за температури 1 він має на 9% більшу довжину прийняття, але працює на 35% повільніше.

Тест із перемішуванням перевіряє, чи використовують чернетки реальну структуру між токенами. За 2 кроки FlashMap відновлює більшу частину переваги спільної структури 16-крокового вчителя. За 1 крок він зберігає лише невелику її частину. Створення 4 зразків шуму у вигляді дерева підвищило прискорення в середньому на 11%.

Порівняння: FlashMap та інші чернеткові модулі для спекулятивного декодування

Характеристика	FlashMap	DFlash	EAGLE-3	DSpark
Стиль створення чернетки	Безперервне відображення потоку, 1 або 2 проходи	Блокова дифузія, 1 прохід	Авторегресивний, дерево	Напівавторегресивний
Залежності між токенами	Так, спільні для всього блоку	Ні, незалежні для кожної позиції	Так, послідовні	Так, легкий послідовний модуль
Розмір чернеткового модуля (Qwen3-8B)	500M	1B	302M (перенавчений, згідно зі статтею FlashMap)	Не розкрито
Головна заява (власне джерело)	До 6.1x	До 6.17x на Qwen3-8B	До 6.5x	На 60–85% швидше на користувача порівняно з MTP-1
Середнє прискорення, Qwen3-8B, T=0 (зіставне, стаття FlashMap)	4.36x	4.07x	2.55x	Не тестувалося
Прискорення на GSM8K, T=0 (зіставне)	6.14x	5.77x	3.20x	Не тестувалося
Ліцензія на ваги	Apache-2.0	MIT	Apache-2.0 (код)	Не розкрито
Обслуговування	Transformers, vLLM, SGLang	Transformers, vLLM, SGLang	SGLang та інші	Обслуговування DeepSeek-V4

Середні показники прискорення обчислені Marktechpost за рисунком 1(c) зі статті про FlashMap. У зіставних рядках використано перенавчені базові моделі авторів.

Як розгорнути FlashMap?

Пакету flashmap потрібні лише torch і transformers>=4.56. Встановіть його за допомогою uv pip install flashmap. Виклик spec_generate запускає чернеткову та цільову моделі разом на 1 GPU.

Для обслуговування сторінка проєкту пропонує 2 способи:

vLLM 0.31: передайте чернетковий модуль у --speculative-config, встановивши num_speculative_tokens на 16.
SGLang 0.5.20: використовуйте --speculative-algorithm DSPARK зі шляхом до чернетки FlashMap.

Встановіть num_steps=2, щоб обміняти частину швидкості на вище прийняття. Для семплювання за температури понад 0 у vLLM потрібно вказати "draft_sample_method": "probabilistic".

Головні висновки

FlashMap спільно створює 16 токенів за 1 прохід із використанням безперервного відображення потоку.
До 6.14x прискорення без втрати точності на Qwen3-8B із чернетковим модулем на 0.5B параметрів.
Швидший за перенавчений DFlash у всіх 9 бенчмарках Qwen3-8B.
Двокроковий режим зберігає стабільну довжину прийняття в усьому блоці.
Ваги ліцензовані за Apache-2.0, код — за MIT, доступні плагіни для vLLM і SGLang.

Ознайомтеся зі статтею, сторінкою проєкту, репозиторієм GitHub і вагами моделей. Також підписуйтеся на нас у Twitter і не забудьте приєднатися до нашого субреддіту про машинне навчання з понад 150 тис. учасників та підписатися на нашу розсилку. Стривайте! Ви є в Telegram? тепер ви також можете приєднатися до нас у Telegram.

Потрібна наша партнерська підтримка для просування вашого репозиторію GitHub, сторінки на Hugging Face, релізу продукту, вебінару тощо? Зв’яжіться з нами

Перекладено автоматично з англійської. Оригінал статті — за посиланням нижче.