Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs

Probleemstelling

Speculatieve decodering is een veelgebruikte techniek om de inferentie van grote taalmodellen (LLM's) te versnellen. Hierbij gebruikt een kleiner "draft model" (of drafter) om een reeks tokens te voorspellen die door het grotere "target model" parallel worden geverifieerd. Als de voorspellingen correct zijn, worden meerdere tokens in één keer gegenereerd, wat de latentie verlaagt.

Het huidige probleem is echter dat er vaak meerdere kandidaat-drafters beschikbaar zijn, elk gespecialiseerd in een bepaald domein (bijv. wiskunde, code, biomedische tekst). Een enkele drafter presteert goed binnen zijn specialisatie maar faalt vaak buiten die context. Bestaande methoden, zoals BanditSpec, behandelen de selectie van de beste drafter als een Multi-Armed Bandit (MAB) probleem. Dit vereist een afweging tussen exploratie (het testen van verschillende drafters) en exploitatie (het gebruiken van de momenteel beste drafter). Het nadeel van deze aanpak is dat ze slechts feedback ontvangen van de gekozen drafter, wat leidt tot trage convergentie en suboptimale prestaties, vooral als het aantal drafters groeit.

Methodologie: HedgeSpec

De auteurs introduceren HedgeSpec, een framework dat het probleem transformeert van een bandit-probleem (gedeeltelijke informatie) naar een full-information online learning probleem.

1. Het "Niet-Bandit" Inzicht:
Het kernidee is dat tijdens de verificatiefase van speculatieve decodering, het target model een traject van tokens produceert dat als "counterfactual evidence" kan dienen voor alle drafters in de pool, niet alleen de gekozen ene.

• Na het verifiëren van een chunk van tokens door het target model, kunnen de auteurs de acceptatiekansen voor elke andere drafter in de pool berekenen door hun voorspellingen te vergelijken met het geverifieerde traject.
• Dit gebeurt zonder extra calls naar het dure target model, omdat de verificatie al plaatsvindt. Hierdoor krijgen ze "panoramische feedback" (full-information feedback).

2. Schatting van Acceptatielengte (Theorem 3):
Om de prestaties van niet-gekozen drafters te evalueren, gebruiken de auteurs een schatter voor de verwachte acceptatielengte (Expected Acceptance Length - EAL).

• Ze definiëren $\gamma_k$ als de kans dat het $k$-de token wordt geaccepteerd, gegeven dat alle voorgaande tokens zijn geaccepteerd.
• Zelfs als een drafter niet is gekozen, kunnen ze een onbevooroordeelde schatter berekenen voor hoe lang de acceptatiereeks zou zijn geweest als die drafter wel was gekozen. Dit lost het probleem op van "censoring" (waarbij je geen informatie hebt over hoe lang een niet-gekozen drafter zou hebben gepresteerd).

3. Online Learning Algoritme:
In plaats van bandit-algoritmen (zoals UCB of EXP3) te gebruiken, past HedgeSpec algoritmen toe voor full-information learning, zoals Hedge of NormalHedge.

• Verliesfunctie: Het verlies wordt gebaseerd op het verschil tussen de maximale mogelijke acceptatielengte ($K+1$) en de geschatte acceptatielengte van de drafter.
• Vertraging (Delayed Feedback): Omdat tokens in chunks worden verwerkt, is er een vertraging in de feedback. De auteurs modelleren dit als een online learning probleem met vertraagde feedback en gebruiken een reduction naar standaard algoritmen om de regret-bounds te garanderen.

Belangrijkste Bijdragen

1. Full-Information Framework: Het eerste bewijs dat het mogelijk is om alle drafters nauwkeurig te evalueren zonder extra kosten voor het target model, waardoor het bandit-probleem overbodig wordt.
2. Theoretische Garantie: Het paper levert strikte "no-regret" garanties. De prestaties van HedgeSpec convergeren exponentieel sneller naar de beste drafter in hindsight dan bandit-methoden naarmate het aantal drafters ($N$) toeneemt. De regret schaling is $O(\sqrt{K \log N / T})$ in plaats van de slechtere schaling bij bandits.
3. Systeem-efficiëntie: De auteurs tonen aan dat de overhead van het evalueren van alle drafters verwaarloosbaar is in vergelijking met de winst in acceptatiekansen. De evaluatie kan parallel worden uitgevoerd.
4. Robuustheid: Het systeem werkt met diverse speculatieve decodering methoden (single draft, multi-drafts, draft-trees zoals EAGLE) en is bestand tegen distributieveranderingen (Out-of-Distribution queries).

Resultaten

De auteurs hebben uitgebreide experimenten uitgevoerd met open-source LLM's (Llama-3.1-8B, Qwen-3-8B/32B) en 21 gespecialiseerde drafters over zeven domeinen (Python, Wiskunde, Biologie, Chemie, Medische QA, CNN DM, SQL).

• Prestatieverbetering: HedgeSpec overtreft consistent de state-of-the-art methoden EAGLE (een sterke generalist) en BanditSpec (UCB en EXP3 varianten).
  • In enkele domeinen werd een toename van 83.7% in tokens per seconde behaald ten opzichte van EAGLE.
  • Gemiddeld over alle domeinen was er een winst van 46.1% in tokens per seconde.
  • De Mean Accepted Tokens (MAT) steeg met tot 49% ten opzichte van bandit-baselines.
• Schaalbaarheid: Terwijl de prestaties van bandit-methoden sterk verslechteren naarmate het aantal drafters toeneemt (door de hoge kosten van exploratie), schaalt HedgeSpec naadloos en blijft stabiel presteren.
• Robuustheid bij OOD: In tegenstelling tot een offline router (een statische classifier die trainingsdata nodig heeft), presteert HedgeSpec uitstekend bij queries die buiten de trainingsdistributie vallen (bijv. promptvariaties die de intentie veranderen), omdat het dynamisch leert tijdens de runtime.

Betekenis en Impact

Dit paper biedt een paradigmaverschuiving in het veld van speculatieve decodering. Het toont aan dat de beperkingen van bandit-algoritmen (trage exploratie) kunnen worden omzeild door de structuur van het decodingsproces slim te benutten.

• Efficiëntie: Het maakt het mogelijk om grote pools van gespecialiseerde "expert" drafters te combineren zonder de overhead van traag leren. Dit is cruciaal voor productieomgevingen waar queries uit zeer diverse domeinen komen.
• Theoretische Onderbouwing: Het biedt wiskundige zekerheid dat het systeem optimaal leert, wat belangrijk is voor betrouwbare AI-servingsystemen.
• Praktische Toepasbaarheid: De methode is plug-and-play met bestaande frameworks zoals EAGLE en vereist geen hertraining van de target modellen of toegang tot hun trainingsdata.

Kortom, HedgeSpec is een krachtige, theoretisch onderbouwde oplossing die de snelheid en kwaliteit van LLM-inferentie aanzienlijk verbetert door slimme, datagedreven selectie van de beste drafter in real-time.