TAPS: Dlaczego dane treningowe modelu-draftu mają większe znaczenie niż jego architektura

Speculative decoding to jedna z najbardziej eleganckich sztuczek w inferencji LLM: mały, szybki model-draft model-draft Lekki model językowy, który szybko proponuje kandydujące tokeny. Większy model ‘weryfikator’ sprawdza te propozycje równolegle, akceptując poprawne i odrzucając błędne - przyspieszając generowanie bez zmiany jakości wyjścia. proponuje tokeny, a duży weryfikator weryfikator Pełnowymiarowy docelowy model językowy, który sprawdza propozycje draftu. Przetwarza wszystkich kandydatów w jednym przebiegu, akceptując te zgodne z własną dystrybucją - gwarantując identyczną jakość jak standardowe dekodowanie autoregresyjne. zatwierdza lub odrzuca je równolegle. Ta sama dystrybucja wyjściowa, mniej kosztownych przebiegów.

Ale jest pytanie, którego nikt nie zadawał: co się dzieje, gdy model-draft był trenowany na niewłaściwych danych?

Praca “TAPS: Task Aware Proposal Distributions for Speculative Sampling” (Zbib et al., KAUST i AUB, marzec 2026) pokazuje, że dane treningowe draftu mają takie samo znaczenie jak architektura. Draft trenowany na matematyce świetnie radzi sobie z matematyką, ale potyka się na czacie. Draft trenowany na czacie - odwrotnie. A sposób, w jaki łączy się specjalistów w czasie inferencji - nie w przestrzeni wag - decyduje o tym, czy uzyskamy to, co najlepsze z obu światów, czy to, co najgorsze.

Problem: jeden draft nie pasuje do wszystkiego

Większość dotychczasowych prac nad speculative decoding koncentruje się na ulepszeniach architektury - lepsze draftowanie na poziomie cech (EAGLE, EAGLE-2, EAGLE-3), dynamiczne drzewa, kaskadowe draftery, metody self-speculative. Ale prawie wszystkie modele-drafty trenuje się na tych samych szerokich korpusach, takich jak ShareGPT.

To tworzy martwy punkt. Jeśli draft jest trenowany na danych konwersacyjnych, a Twoje zadanie to rozumowanie matematyczne, dystrybucja propozycji dystrybucja propozycji Rozkład prawdopodobieństwa następnych tokenów produkowany przez model-draft. Gdy ta dystrybucja jest dobrze dopasowana do zachowania weryfikatora na danym zadaniu, więcej tokenów jest akceptowanych i inferencja jest szybsza. draftu będzie słabo dopasowana do zachowania weryfikatora. Algorytm speculative decoding wciąż jest bezstratny - dostajesz poprawną odpowiedź - ale długość akceptacji spada, a razem z nią przyspieszenie.

TAPS stawia pięć pytań badawczych:

1. Czy trening zadaniowy poprawia akceptację na dopasowanych zadaniach?
2. Czy mieszany trening odzyskuje odporność międzydomenową?
3. Jak łączyć wielu wyspecjalizowanych drafterów?
4. Jakie sygnały (confidence vs. entropia) są przydatne do routingu?
5. Jak głębokość spekulatywna wpływa na specjalizację zadaniową?

Setup eksperymentalny

Kontrolowany design to jedna z mocnych stron pracy. Wszystko jest ustalone poza dwoma zmiennymi: danymi treningowymi i strategią kompozycji.

Elementy stałe

• Weryfikator: Meta-Llama-3-8B-Instruct
• Architektura draftu: lekki dekoder w stylu LLaMA, jedna warstwa transformera, hidden size 4096, ~0.8B parametrów
• Tokenizer: współdzielony między draftem a weryfikatorem
• Trening: 20 epok, learning rate $3 \times 10^{-5}$, batch size 8

Zmienna: dane treningowe

• MathInstruct (70k przykładów) - rozumowanie matematyczne
• ShareGPT (70k przykładów) - dane konwersacyjne
• Mixed 35k+35k - zbalansowana mieszanka, ta sama łączna liczba
• Mixed 70k+70k - podwojona łączna liczba, obie domeny w pełni reprezentowane

Zmienna: backbone spekulatywny

Dwie najnowocześniejsze architektury testowane niezależnie:

• EAGLE-2 EAGLE-2 Metoda speculative decoding, która przewiduje przyszłe ukryte cechy (nie tokeny bezpośrednio) i używa kontekstowo-zależnego dynamicznego drzewa do organizacji sekwencji kandydujących. - draftowanie na poziomie cech z dynamiczną konstrukcją drzewa
• HASS HASS Harmonized Speculative Sampling - poprawia dopasowanie draft-target przez destylację Top-K i zharmonizowane wyrównanie kontekstu, trenując na niedoskonałych cechach draftu zamiast tylko na czystych cechach docelowych. - zharmonizowane reprezentacje z destylacją Top-K

Benchmarki

Główna metryka: długość akceptacji - średnia liczba tokenów draftu zaakceptowanych na jedno wywołanie weryfikatora. Wyższa = lepsze dopasowanie = szybsza inferencja.

RQ1: Specjalizacja jest realna

Wyniki są jednoznaczne. Pod HASS przy temperaturze 0:

ShareGPT jest o 37% lepszy na MT-Bench. MathInstruct jest o 23% lepszy na GSM8K i 34% lepszy na MATH-500. Ten sam wzorzec utrzymuje się dla EAGLE-2.

To nie jest subtelna różnica. Draft trenowany na niewłaściwej domenie nie traci kilku punktów procentowych - on znacząco degraduje przewagę speculative decoding.

Wniosek: Jakość draftu zależy nie tylko od architektury, ale od dopasowania dystrybucji treningowej do docelowego obciążenia.

RQ2: Mieszany trening pomaga, ale nie dominuje

Czy można naprawić specjalizację przez mieszanie danych treningowych? Częściowo.

Pod HASS przy temperaturze 0, Mixed 70k+70k osiąga najwyższą średnią długość akceptacji (5.18 na wszystkich benchmarkach). Ale przy temperaturze 1 spada do 3.69 - poniżej 4.29 osiąganego przez Mixed 35k+35k.

Wzorzec powtarza się dla EAGLE-2: Mixed 70k+70k jest najsilniejszy przy temperaturze 0 (średnia 4.48), ale Mixed 35k+35k jest stabilniejszy przy temperaturze 1 (3.81 vs. 3.26).

Wniosek: Mieszany trening poszerza pokrycie, ale nie poprawia jednolicie we wszystkich temperaturach dekodowania. Nie eliminuje potrzeby dostrojenia mieszanki do reżimu dekodowania.

RQ3: Komponuj w czasie inferencji, nie w przestrzeni wag

To najbardziej praktycznie użyteczne odkrycie pracy. Gdy masz dwóch wyspecjalizowanych drafterów, jak ich połączyć?

Trzy strategie

1. Uśrednianie checkpointów - liniowa interpolacja w przestrzeni parametrów:

$$\theta_{merge} = \lambda \theta_{math} + (1 - \lambda) \theta_{chat}$$

Proste, ale destrukcyjne. Średnie długości akceptacji 2.59 (HASS) i 2.62 (EAGLE-2) - najgorsze wyniki w całej tabeli, konsekwentnie poniżej każdego jednodomendowego specjalisty.

2. Routing oparty na confidence - uruchom oba draftery, wybierz ten z wyższą średnią pewnością:

$$\mathcal{T}^* = \arg\max_{\mathcal{T} \in \{\mathcal{T}_{math}, \mathcal{T}_{chat}\}} \text{Score}(\mathcal{T}), \quad \text{Score}(\mathcal{T}) = \frac{1}{|\mathcal{T}|} \sum_{v \in \mathcal{T}} c(v)$$

Znaczący skok: 4.80 (HASS) i 4.63 (EAGLE-2) średniej długości akceptacji przy temperaturze 0.

3. Merged-tree verification - upakuj oba drzewa draftu pod wspólnym korzeniem z oddzielnymi maskami uwagi, zweryfikuj wszystkich kandydatów w jednym przebiegu:

Najlepszy wynik ogólnie: 5.11 (HASS) i 5.02 (EAGLE-2) przy temperaturze 0. To bije każdy jednodomendowy checkpoint i każdy wariant z mieszanymi danymi.

Tabela wyników (Temperatura 0)

Uśrednianie wag niszczy wiedzę specjalistów. Kompozycja w czasie inferencji ją zachowuje.

Wniosek: Trzymaj specjalistów osobno. Łącz ich propozycje, nie parametry.

RQ4: Confidence bije entropię w routingu

Zarówno confidence, jak i entropia to potencjalne sygnały do decydowania, któremu specjaliście zaufać. TAPS testuje oba na poziomie benchmarków używając EAGLE-2.

Routing na bazie confidence

Routing na bazie confidence daje niemal idealną separację domen: specjalista matematyczny jest wybierany dla >90% przykładów matematycznych, specjalista konwersacyjny dla >80% konwersacji.

Routing na bazie entropii

Routing entropijny jest niewiele lepszy od losowego. Wykrywa, że odrzucone tokeny mają wyższą entropię (przydatne jako diagnostyka), ale produkuje niemal zbalansowane podziały, które nie oddzielają domen.

Wniosek: Confidence to sygnał routingu. Entropia to sygnał diagnostyczny. Nie myl jednego z drugim.

RQ5: Głębokość ujawnia kompromis eksploracja-eksploatacja

Eleganckie odkrycie: jak specjalizacja wchodzi w interakcję z głębokością drzewa draftu.

Na płytkich głębokościach (1–2 tokeny) drafty z mieszanych danych często wypadają najlepiej. Szersze pokrycie oznacza więcej szans na wyprodukowanie akceptowalnego pierwszego tokenu - przewaga eksploracyjna.

Na głębszych poziomach (3–5 tokenów) specjalista dopasowany do zadania coraz bardziej dominuje, szczególnie na benchmarkach rozumowania. Utrzymanie zgodności z weryfikatorem wymaga bliskiego dopasowania dystrybucji - przewaga eksploatacyjna.

To tłumaczy, dlaczego merged-tree verification działa tak dobrze: zapewnia eksplorację (różnorodność od dwóch specjalistów) na płytkich poziomach, jednocześnie pozwalając lepiej dopasowanemu specjaliście napędzać głębszą akceptację.

Algorytm Merged-Tree

Najbardziej nowatorski wkład techniczny to procedura merged-tree verification:

1. Wygeneruj drzewa draftu $\mathcal{T}_{math}$ i $\mathcal{T}_{chat}$ z tego samego tokenu korzeniowego
2. Scal pod wspólnym korzeniem przez konkatenację węzłów i przemapowanie indeksów
3. Zbuduj maski uwagi zachowujące przodków - węzły z jednego poddrzewa nie mogą obserwować węzłów z drugiego
4. Przypisz identyfikatory pozycji według głębokości drzewa
5. Zweryfikuj całe scalone drzewo w jednym przebiegu weryfikatora
6. Wyodrębnij ścieżki kandydujące i zastosuj standardową akceptację spekulatywną
7. Zatwierdź zaakceptowany prefiks

Izolacja masek uwagi jest kluczowa: każde poddrzewo utrzymuje własną wewnętrzną hierarchię przodków, więc weryfikator ocenia propozycje każdego specjalisty na ich własnych warunkach. Różnorodność pochodzi z posiadania obu zestawów kandydatów, nie z krzyżowania ich obliczeń.

Praktyczne implikacje

Dla inżynierów inferencji

Jeśli Twoje obciążenie jest domenowo specyficzne (kodowanie, matematyka, prawo, medycyna), wytrenuj lub dostroij model-draft na dopasowanych danych domenowych. Backbone architektoniczny (EAGLE-2 vs. HASS) ma mniejsze znaczenie niż dopasowanie dystrybucji treningowej.

Dla wdrożeń wielodomenowych

Nie uśredniaj checkpointów. Zamiast tego:

• Użyj confidence routing (niższy narzut, wymaga uruchomienia obu drafterów na prefiksie)
• Użyj merged-tree verification (wyższa akceptacja, wymaga większego batcha weryfikatora ze scalonego drzewa)

Dla badaczy

Sekcja ograniczeń pracy jest odświeżająco uczciwa: jeden model docelowy, dwie domeny źródłowe, dwa backbony spekulatywne, cztery benchmarki. Polityka routingu jest celowo prosta (oparta na confidence, nie uczona). Ulepszenia latencji end-to-end z merged-tree verification nie są mierzone - długość akceptacji jest proxy. To wszystko jasne kierunki dalszych prac.

Ograniczenia

• Jeden weryfikator (Llama-3-8B-Instruct) - niejasne, jak wyniki przenoszą się na większe lub inne architektury
• Tylko dwie domeny - realne wdrożenia obejmują wiele więcej typów zadań
• Brak latencji end-to-end - merged-tree verification zwiększa batch size weryfikatora; czy zysk w długości akceptacji kompensuje ten koszt, pozostaje otwartym pytaniem
• Prosta polityka routingu - oparta na confidence, nie uczona; wyuczony router mógłby działać lepiej
• Długość akceptacji ≠ przyspieszenie wall-clock - narzut uruchamiania wielu drafterów lub większych batchy weryfikacji nie jest uwzględniony

Podsumowanie

TAPS dostarcza prosty, ale ważny komunikat: speculative decoding to problem projektowania systemów, nie tylko problem architektoniczny. Model-draft nie jest stałym komponentem pomocniczym - to wybór projektowy, który powinien być dopasowany do docelowego obciążenia. Gdy istnieje wielu specjalistów, komponowanie ich w czasie inferencji (przez routing lub scaloną weryfikację) znacząco przewyższa łączenie w przestrzeni wag.

Praktyczny przepis: trenuj domenowo specyficzne drafty, trzymaj je osobno, łącz ich propozycje w czasie inferencji i używaj confidence (nie entropii) do routingu między nimi.

📎 Linki

• Na podstawie publikacji 📄 2603.27027