FlashMap: drafter, który generuje cały blok tokenów naraz i pilnuje zależności między nimi

W dekodowaniu spekulatywnym mały model, tzw. drafter, proponuje blok tokenów, a duży model sprawdza je jednym równoległym przebiegiem. Wynik jest taki sam, jaki dałby sam duży model, a zysk zależy od tego, ile propozycji przejdzie weryfikację i jak tanio je wygenerowano. W preprincie opublikowanym 17 października Jeff Shen, Zhejian Zhou, Ryan Lee, Jonathan May i Shirley Ho opisują FlashMap, drafter, który próbuje poprawić oba te czynniki naraz.

Problem z dyfuzyjnymi drafterami

Drafterzy autoregresyjni, jak EAGLE-3, produkują blok token po tokenie, więc koszt rośnie z jego długością. Drafterzy dyfuzji blokowej, np. DFlash, generują cały blok w jednym przebiegu. Według dokumentacji vLLM DFlash używa do tego małego modelu dyfuzyjnego. Autorzy FlashMap zwracają jednak uwagę na słabość tego podejścia: każdy token w bloku przewiduje się niezależnie, więc zależności między nimi są ignorowane. To obniża liczbę propozycji akceptowanych przez model docelowy.

FlashMap zastępuje dyskretną dyfuzję ciągłą mapą przepływu (flow map). Model warunkowany jest ukrytymi stanami modelu docelowego i z szumu gaussowskiego tworzy cały blok w jednym przebiegu. Jedna próbka szumu wyznacza spójnie zawartość wszystkich pozycji, dlatego tokeny w bloku pasują do siebie.

Wyniki i konfiguracja

Według autorów na Qwen3-8B FlashMap jest szybszy od porównywalnego DFlash na każdym teście, ze średnio o 7% wyższym przyspieszeniem i do 6,1x bez utraty jakości. Według autorów preprintu szczytowy wynik to 6,14x na GSM8K przy temperaturze 0 (średnia długość akceptowanego odcinka: 7,69 tokena), a najsłabszy to 2,28x na Alpaca przy temperaturze 1.

• Drafter ma ok. 0,5 mld parametrów (500 mln dla Qwen3-8B, 475 mln dla Qwen3-4B) i generuje bloki po 16 tokenów.
• To 6-blokowy transformer dyfuzyjny, który korzysta ze stanów ukrytych z warstw 1, 9, 17, 25 i 33 modelu docelowego.
• Trening przebiega dwuetapowo (nauczyciel przepływowy, potem uczeń z mapą przepływu) na 500 tys. próbek z Open-PerfectBlend.
• Autorzy sami przetrenowali EAGLE-3 i DFlash na tych samych danych i ustawieniach, a porównanie objęło 9 testów z matematyki, kodu i czatu.

Wariant jednokrokowy przewyższa DFlash pod względem długości akceptowanego odcinka w 15 z 18 ustawień. Przy temperaturze 0 przewaga wynosi 5,5% w matematyce i 11,1% w zadaniach czatowych. Tryb dwukrokowy wygrywa we wszystkich 18 ustawieniach, średnio o 10%.

Wyniki mają ograniczenia. Pomiary przyspieszenia wykonano na jednej karcie NVIDIA H200 NVL, w bf16, przy współbieżności 1, a porównanie opiera się na bazowych modelach przetrenowanych przez samych autorów. Dla mniejszego Qwen3-4B przy temperaturze 0 FlashMap przegrywa z DFlash w 4 z 9 zadań. Wyniki przedstawiono w preprincie.

Wagi drafterów dla Qwen3-8B i Qwen3-4B są na Hugging Face na licencji Apache-2.0, kod na MIT. Dostępne są też wtyczki do vLLM i SGLang. Nad projektem pracowali badacze z Polymathic AI, Princeton, USC ISI, NYU i Flatiron Institute.