BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

Technische Samenvatting: BudgetDraft – Acceptance-Aware Multi-View Training voor Sparse-KV Speculative Decoding

1. Probleemstelling

Het artikel behandelt een kritieke bottleneck bij het implementeren van Speculative Decoding (SD) voor inferentie met middel- tot lange context (4K–16K tokens). Hoewel SD autoregressieve (AR) generatie versnelt door een klein "drafter"-model te gebruiken om tokens voor te stellen die door een groter "verifier"-model gevalideerd worden, wordt de standaardimplementatie geconfronteerd met een sparse/full mismatch.

• De beperking: Om het piekgeheugen van de GPU (VRAM) en de latentie te beheersen, wordt de drafter vaak gedwongen om met een sparse KV-cache te werken onder een vast budget, terwijl de verifier een full KV-cache behoudt om de outputkwaliteit te waarborgen.
• De faalmodus: Naarmate de contextlengte toeneemt van 4K naar 16K, verslechtert de uitlijning tussen de sparse-drafter en de full-verifier snel. Het artikel identificeert een "16K collapse", waarbij de acceptatiegraden (acceptance rates) snel dalen naar bijna nul, waardoor SD ineffectief wordt (speedup $\le$ 1$\times$).
• Het gevoeligheidsprobleem: Naïve sparse drafting is zeer gevoelig voor het specifieke KV-budget. Een drafter die is afgestemd op één enkel budget, presteert vaak slecht wanneer deze wordt ingezet onder verschillende geheugenbeperkingen, en de relatie tussen de budgetgrootte en de acceptatiegraad is niet-monotoon (kleinere budgetten leveren soms hogere acceptatie op dan grotere budgetten).

2. Methodologie: BudgetDraft

De auteurs stellen BudgetDraft voor, een oplossing tijdens de trainingsfase die een enkele, budget-robuuste drafter creëert die in staat is om stabiele acceptatiegraden te handhaven over variërende KV-budgetten zonder dat er extra componenten tijdens de inferentie nodig zijn.

Kernarchitectuur:
De methode maakt gebruik van een Multi-View Sparse Training strategie waarbij de drafter (student) wordt getraind om zich gelijktijdig met meerdere sparse condities uit te lijnen met een bevroren verifier (teacher).

• Teacher Targets: De verifier genereert greedy teacher targets ($x^*_t$) met behulp van zijn full KV-cache.
• Dual-Branch Loss Function: De trainingsdoelstelling combineert twee complementaire losses:
  1. Acceptance-Aware Loss ($L_A$): Een full-cache branch waarbij de drafter de greedy tokens van de teacher voorspelt met een full prefix cache. Dit zorgt voor directe uitlijning met het acceptatiecriterium van de verifier (top-1 match).
  2. Multi-View Sparse Loss ($L_C$): Een sparse-cache branch waarbij de drafter dezelfde sequentie verwerkt maar met een willekeurig gesampled KV-budget ($B$) uit een set (bijv. $\{256, 512, 1024, 2048\}$). De sparse cache wordt geconstrueerd door de prefix te segmenteren en de top-chunks te behouden op basis van cumulatieve attention weights.
• Trainingsdoelstelling: De totale loss is $L = L_A + \lambda L_C$. Door de drafter tijdens de training bloot te stellen aan diverse sparse views van dezelfde prefix, leert deze te generaliseren over verschillende sparsity-niveaus in plaats van te overfitten op een enkel budget.
• Inferentie: Bij implementatie gebruikt de getrainde drafter een sparse KV-cache met een gekozen budget $B$. De verifier valideert deze met zijn full cache. Er worden geen structurele wijzigingen of tussenliggende modellen toegevoegd aan de inferentie-pipeline.

3. Belangrijkste Bijdragen

1. Karakterisering van Acceptance Collapse: Het artikel stelt empirisch vast dat naïeve sparse speculative decoding lijdt onder een ernstige acceptatiegraad-instorting naarmate de contextlengte toeneemt van 4K naar 16K, waarbij 16K wordt geïdentificeerd als een praktische faalgrens. Het onthult ook de niet-monotone relatie tussen de grootte van het KV-budget en de acceptatiegraad.
2. BudgetDraft Framework: De voorstel van een multi-view trainingsmethode die acceptance-aware alignment combineert met multi-budget sparse training. Dit produceert een budget-invariante drafter die de acceptatiegraden over alle sparsity-niveaus herstelt zonder extra overhead tijdens de inferentie.
3. Empirische Validatie: Demonstratie dat BudgetDraft significante end-to-end versnellingen bereikt ten opzichte van standaard AR-decoding terwijl de geheugenefficiëntie behouden blijft, specifiek in het uitdagende 8K–16K regime.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op drie datasets (PG-19, LongBench, LWM) met een 68M parameter drafter en een 7B parameter verifier op een enkele NVIDIA A100 GPU.

• Speedup Prestaties:
  • 4K Context: Bereikte tot 6.55$\times$ versnelling ten opzichte van AR.
  • 8K Context: Bereikte tot 4.46$\times$ versnelling ten opzichte van AR.
  • 16K Context: Bereikte tot 2.10$\times$ versnelling ten opzichte van AR (waar naïeve SD vaak volledig faalt).
• Robuustheid: In tegen tegenover de baseline SD (sparse/full), die een scherpe prestatiedaling of bijna nul acceptatie vertoont naarmate de contextlengte toeneemt of budgetten veranderen, behoudt BudgetDraft stabiele acceptatiegraden over verschillende KV-budgetten ($B \in \{256, 512, 1024, 2048\}$).
• Vergelijkingen: BudgetDraft presteerde beter dan structurele mitigatiemethoden zoals TriForce en geavanceerde draftingmethoden zoals EAGLE-3 in het 8K–16K regime. Zo behaalde BudgetDraft op 16K op LWM een versnelling van ~1.94$\times$, vergeleken met 1.19$\times$ voor TriForce en 1.36$\times$ voor EAGLE-3.
• Ablatie: Het verwijderen van de multi-view sparse loss ($L_C$) resulteerde in instabiele prestaties onder budgetverschuivingen, wat de noodzaak van multi-budget sampling tijdens de training bevestigt.

5. Betekenis en Claims

Het artikel claimt dat BudgetDraft het "sparse/full mismatch" probleem oplost dat momenteel de praktische inzetbaarheid van speculative decoding in long-context scenario's beperkt. Door de complexiteit naar de trainingsfase te verschuiven, maakt het een eenvoudige, geheugenvriendelijke inferentie-pipeline mogelijk die robuust is tegen variërende hardwarebeperkingen (beschikbaarheid van VRAM) en workload-concurrency.

De auteurs benadrukken dat hun aanpak geen complexe structurele wijzigingen (zoals retrieval caches) of meerdere modellen tijdens de inferentie vereist. In plaats daarvan biedt het een enkele, lichtgewicht drafter die kan adapteren aan verschillende deployment-budgetten, waardoor speculative decoding levensvatbaar wordt voor mid-to-long context toepassingen (4K–16K) waar het voorheen ineffectief was.

Beperkingen erkend door het artikel:

• Positional Extrapolation: De 68M drafter heeft een natuurlijke positie-embedding limiet (2048), wat extrapolatie vereist voor 8K–16K contexten, wat de acceptatiegraden nog steeds kan beperken.
• Verifier Specificiteit: De drafter wordt getraind tegen een specifieke verifier; hertraining is vereist als het verifier-model verandert.
• Greedy Decoding: De resultaten zijn gebaseerd op greedy decoding; het uitbreiden naar sampling-gebaseerde decoding (bijv. temperature sampling) wordt overgelaten aan toekomstig werk.