LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

Deze paper introduceert LK-losses, een trainingsdoel dat de acceptatiekans bij speculatieve decoding direct optimaliseert in plaats van de gebruikelijke KL-divergentie, wat leidt tot aanzienlijke snelheidswinsten voor diverse grote taalmodellen zonder extra rekentijd.

De kern

Stel je voor dat je een zeer slimme, maar trage robot hebt die verhalen schrijft (een groot taalmodel of LLM). Deze robot is geweldig, maar hij werkt heel langzaam omdat hij elke zin woord voor woord moet bedenken, wachtend op het vorige woord voordat hij het volgende kan schrijven.

Om dit te versnellen, hebben onderzoekers een truc bedacht: Speculatieve Decoding.

Het Probleem: De Snelle Assistent en de Strikte Chef

Stel je voor dat de trage robot de Chef is. Hij wil een recept schrijven.
Om tijd te winnen, heeft hij een Snelle Assistent (het draft model) in dienst.

1. De Assistent probeert snel 5 woorden achter elkaar te raden die de Chef waarschijnlijk zou gebruiken.
2. De Chef kijkt dan in één keer naar die 5 woorden.
   • Als de Chef denkt: "Ja, dat zou ik ook hebben gezegd!", dan worden alle 5 woorden direct geaccepteerd. Groot voordeel! De Chef hoeft niet zelf na te denken over die 5 woorden.
   • Als de Chef denkt: "Nee, dat klinkt raar", dan gooit hij het eerste verkeerde woord weg en stopt hij met kijken naar de rest. De Assistent heeft zijn tijd dan verloren.

De snelheid van het hele systeem hangt af van hoe vaak de Chef "Ja" zegt. Dit noemen ze de Acceptatiegraad. Hoe meer woorden de Assistent goed raadt, hoe sneller de Chef kan werken.

Het Oude Probleem: De Verkeerde Oefening

Tot nu toe trainden onderzoekers de Assistent door te vragen: "Hoe dichtbij jouw antwoord komt het antwoord van de Chef?"
Ze gebruikten een wiskundige maatstaf (KL-divergentie) die kijkt of de verdeling van de antwoorden hetzelfde is.

Het probleem: De Assistent is klein en niet zo slim als de Chef. Hij kan de Chef nooit 100% perfect nabootsen.
Als je de Assistent alleen laat oefenen op "dichtbij komen", kan hij vastlopen in een valkuil. Hij leert misschien een antwoord dat wiskundig "dichtbij" is, maar dat de Chef toch afkeurt. Het is alsof je een student laat oefenen op het gevoel van een examen, in plaats van op het daadwerkelijke cijfer dat hij nodig heeft om te slagen.

De Oplossing: LK Losses (De Nieuwe Trainingsmethode)

De auteurs van dit paper (van Nebius) zeggen: "Waarom oefenen we niet direct op het doel?"
In plaats van te vragen: "Is je antwoord vergelijkbaar met dat van de Chef?", vragen ze direct: "Zal de Chef dit woord accepteren?"

Ze hebben twee nieuwe trainingsmethoden bedacht, die ze LK Losses noemen (een knipoog naar de oude methode, maar dan gericht op Likelihood of Kwaliteit van acceptatie).

Analogie 1: De Twee Trainingsmethoden

1. De Hybride Methode (De Slimme Coach):

   • Aan het begin: De Assistent is nog heel onzeker en maakt veel fouten. De Coach zegt: "Probeer eerst maar gewoon in de buurt te komen van wat de Chef zegt" (de oude methode). Dit helpt de Assistent om niet in de war te raken.
   • Naarmate hij beter wordt: Zodra de Assistent wat zekerder is, schakelt de Coach over. Hij zegt: "Nee, wacht even! Kijk niet meer naar wat de Chef zou zeggen, maar kijk puur naar welke woorden de Chef zeker zal accepteren."
   • Dit is als een leerling die eerst de theorie leert, en later direct begint met oefenvragen die echt tellen voor het examen.

2. De Waarschijnlijkheids-Methode (De Directe Benadering):

   • Deze methode kijkt puur naar de kans dat een woord wordt geaccepteerd en probeert die kans direct te maximaliseren. Het is alsof je een speler laat oefenen op het scoren van doelpunten, in plaats van op het "mooi spelen" van voetbal.

Wat Leverde Dit Op?

De onderzoekers hebben dit getest met verschillende "Chefs" (grote modellen van 8 miljard tot 685 miljard parameters) en verschillende "Assistenten".

• Het resultaat: De Assistenten die met de nieuwe LK-methode werden getraind, raakten 8% tot 10% meer woorden goed dan degenen die met de oude methode werden getraind.
• Bijna gratis: Dit kostte geen extra rekenkracht om te trainen. Het is gewoon een andere manier om de "vraag" te stellen tijdens het leren.
• Werkt overal: Het werkte beter bij kleine, simpele Assistenten (die het hardst nodig hebben) en ook bij de grootste, meest complexe modellen.

Samenvatting in Eén Zin

In plaats van een snelle assistent te trainen om "zo veel mogelijk op de chef te lijken", trainen ze hem nu direct om woorden te raden die de chef zeker zal goedkeuren, waardoor de hele schrijfsessie veel sneller verloopt.

Het is alsof je van een student die probeert de leraar na te bootsen, een student maakt die gewoon weet welke antwoorden de leraar als "goed" zal beoordelen.

Technische samenvatting

Probleemstelling

Speculatieve decoding (speculative decoding) is een techniek om de inferentie van autoregressieve Large Language Models (LLM's) te versnellen. Hierbij gebruikt een lichtgewicht "draft model" kandidaat-tokens die parallel worden geverifieerd door het zwaardere "target model". De snelheidswinst wordt primair bepaald door de acceptatiekans (acceptance rate): de waarschijnlijkheid dat een door het draft model voorgestelde token door het target model wordt geaccepteerd.

Het standaardtrainingsparadigma voor draft modellen is het minimaliseren van de Kullback-Leibler (KL) divergentie tussen de verdeling van het target model en het draft model. Hoewel KL-divergentie en acceptatiekans dezelfde globale optimum hebben (waarbij beide perfect zijn), is dit een proxy-objective.

• Het kernprobleem: Draft modellen hebben vaak een zeer beperkte capaciteit (vaak slechts 1-5% van de parameters van het target model). Hierdoor convergeren ze naar suboptimale oplossingen. Op deze suboptimale punten garandeert het minimaliseren van KL-divergentie niet dat de acceptatiekans wordt gemaximaliseerd. De optimale verdeling voor KL is niet noodzakelijk dezelfde als die voor maximale acceptatie.

Methodologie: LK Losses

De auteurs introduceren LK Losses (Losses for Acceptance Rate), een nieuwe reeks trainingsdoelstellingen die de acceptatiekans direct maximaliseren in plaats van een proxy te gebruiken. Ze analyseren eerst de gradienten van verschillende divergenties en stellen twee varianten voor:

1. Gradientanalyse en Motivatie:

   • Forward KL: Biedt gladde, goed geconditioneerde gradienten, maar optimaliseert een proxy.
   • Total Variation (TV) Distance: Is wiskundig equivalent aan het maximaliseren van de acceptatiekans ($\alpha = 1 - TV$). Echter, bij willekeurig geïnitieerde modellen (waarbij de draft-verdeling ver van de target-verdeling ligt) leiden TV-gradiënten tot verdwijnende gradienten (vanishing gradients) en niet-gladde optimalisatieoppervlakken, wat training onpraktisch maakt.

2. Variant 1: Hybride Doelstelling met Adaptieve Blending ($L^\lambda_{LK}$)

   • Dit combineert KL-divergentie en TV-distance in een hybride verliesfunctie:
     $$L^\lambda_{LK} = \lambda \cdot KL(p\|q) + (1 - \lambda) \cdot TV(p, q)$$
   • Adaptieve planning: De parameter $\lambda$ wordt dynamisch aangepast op basis van de huidige acceptatiekans ($\alpha$).
     • In het begin van de training (lage $\alpha$) is $\lambda$ hoog, waardoor de stabiele KL-gradiënten de optimalisatie sturen.
     • Naarmate de acceptatiekans verbetert, daalt $\lambda$, waardoor de focus verschuift naar het direct maximaliseren van de acceptatiekans via TV.
   • Dit werkt analoog aan "trust-region" methoden in beleidsoptimalisatie.

3. Variant 2: Kansengebaseerde Benadering ($L^\alpha_{LK}$)

   • Deze methode minimaliseert de negatieve log-marginal likelihood van de acceptatiekans:
     $$L^\alpha_{LK} = -\log \sum_{x \in V} \min(p(x), q(x))$$
   • Dit is in feite een TV-optimalisatie met adaptieve gradiënt-schaling (de factor $1/\alpha$). Wanneer de acceptatiekans laag is, worden de gradienten automatisch versterkt, wat het probleem van verdwijnende gradienten oplost.

4. Vocabulaire Truncatie:

   • De methode is robuust tegen vocabulaire-truncatie (waarbij het draft model alleen hoge-frequentie tokens voorspelt). In tegenstelling tot KL, die oneindig wordt als $q(x)=0$ maar $p(x)>0$, hebben tokens buiten het draft-vocabulaire geen invloed op de acceptatiekans en vereisen LK-losses geen aanpassing van de target-verdeling.

Belangrijkste Bijdragen

• Directe Optimalisatie: Het introduceren van trainingsdoelstellingen die de acceptatiekans direct maximaliseren, in plaats van te vertrouwen op KL-divergentie als proxy.
• Architectuur- en Modelonafhankelijkheid: De methode is getest op vier verschillende draft-architecturen (EAGLE-3, MEDUSA, MLP-speculator, MTP) en zes target-modellen van verschillende groottes (8B tot 685B parameters), inclusief MoE-architecturen.
• Open Source: De auteurs hebben de trainingsdatasets en de weights van de draft-modellen vrijgegeven om reproduceerbaarheid te waarborgen.
• Implementatiegemak: De losses zijn eenvoudig te implementeren, voegen geen rekenkundige overhead toe tijdens training en kunnen direct worden geïntegreerd in bestaande frameworks.

Resultaten

De auteurs evalueren de methode op drie domeinen: conversatie (MT-bench), codering (HumanEval) en wiskunde (GSM8K).

• Consistente Verbetering: LK-losses presteren consistent beter dan de standaard KL-baseline over alle configuraties.
• Acceptatielengte ($\tau$): Er wordt een toename gemeld in de gemiddelde acceptatielengte (het aantal tokens dat per ronde wordt gegenereerd).
  • Voor low-capacity modellen (zoals MEDUSA en MLP) zijn de verbeteringen het grootst (tot 8-10%).
  • Voor EAGLE-3 op grote target-modellen (zoals Qwen3-235B en DeepSeek-V3) worden verbeteringen van 3,8% tot 8,2% behaald.
  • Bij DeepSeek-V3 (waarbij het native MTP-module werd fine-ge-tuned) leverde LK-losses een extra 5,6% winst op ten opzichte van KL-finetuning.
• Hybride vs. Pure TV: Pure TV-training presteert slecht vanwege optimisatieproblemen bij willekeurige initialisatie. De hybride aanpak met adaptieve planning lost dit op en levert de beste resultaten op.

Betekenis en Conclusie

Dit artikel biedt een fundamentele verbetering voor de efficiëntie van speculatieve decoding. Het toont aan dat het vertrouwen op KL-divergentie als trainingsdoel voor beperkte draft-modellen suboptimaal is. Door de acceptatiekans direct te targeten via LK-losses, kunnen systemen aanzienlijk sneller infereren zonder extra rekenkracht tijdens de trainingsfase.

De methode is vooral waardevol bij het gebruik van zeer kleine draft-modellen voor enorme target-modellen (bijv. een 8B model als draft voor een 685B model), waar de capaciteitskloof het grootst is en de voordelen van directe optimalisatie het meest tot hun recht komen. Dit maakt LK-losses een krachtig alternatief voor de huidige industriestandaard.