Bridging Draft Policy Misalignment: Group Tree Optimization for Speculative Decoding

Dit paper introduceert Group Tree Optimization (GTO), een methode die de prestaties van speculatieve decoding verbetert door het trainingsdoel van het draft-model af te stemmen op de boom-structuur van de decoding, wat leidt tot een hogere acceptatielengte en een extra snelheidswinst van 7,7% ten opzichte van de huidige state-of-the-art.

De kern

Het Grote Probleem: De Verkeerde Voorspeller

Stel je voor dat je een zeer slimme, maar trage robot hebt (de Target Model) die zinnen voor je moet schrijven. Omdat deze robot traag is, heb je een snelle, slimmerik-assistent (het Draft Model) die voor hem een paar woorden voorspelt. De robot kijkt dan even of die woorden kloppen. Als ze kloppen, schrijft hij ze direct; zo niet, dan moet hij het zelf doen. Dit heet Speculative Decoding.

Het probleem in de huidige systemen is als volgt:
De assistent (het draft model) wordt getraind alsof hij één enkele lijn moet voorspellen. Hij leert: "Kies altijd het woord dat de meeste kans van slagen heeft." Hij denkt: "Ik ga het woord 'de' kiezen, want dat is het meest waarschijnlijk."

Maar tijdens het echte werk (de decoding) werkt de robot anders. Hij laat de assistent niet één lijn maken, maar een boom van opties. Hij zegt: "Maak een boom met 'de', 'een', 'het' en 'die'." Vervolgens kijkt de robot naar alle takken van die boom en kiest hij de tak die het beste past bij de context.

De mislukking: De assistent is getraind om de "beste" enkele tak te kiezen, maar in de echte boom wordt die tak soms weggegooid omdat een andere tak (die de assistent minder belangrijk vond) beter bleek te zijn. Het is alsof je een voetballer traint om alleen te scoren op het middenveld, maar in de wedstrijd moet hij een hele strategie spelen met veel spelers. Hij is getraind voor de verkeerde spelregels.

De Oplossing: GTO (Groep Boom Optimalisatie)

De auteurs van dit paper, Shijing Hu en collega's, hebben een nieuwe methode bedacht genaamd GTO (Group Tree Optimization) om dit probleem op te lossen. Ze zeggen: "Laten we de assistent trainen alsof hij al een boom maakt, niet alsof hij één lijn trekt."

Hier zijn de twee belangrijkste onderdelen van hun oplossing, vertaald naar alledaagse taal:

1. De "Boom-beloning" (Draft Tree Reward)

In plaats van de assistent te belonen voor het goed voorspellen van één woord, belonen we hem op basis van de hele boom.

• De Analogie: Stel je voor dat je een gids bent in een doolhof. De oude manier was: "Als je de juiste afslag kiest, krijg je een punt." De nieuwe manier (GTO) is: "Kijk naar de hele kaart die je tekent. Hoeveel wegen leiden er naar de uitgang? Hoe langer de reeks van goede wegen die je kunt vinden, hoe meer punten je krijgt."
• Wat het doet: De assistent leert nu niet alleen het "beste" woord te kiezen, maar leert een boom van woorden te maken waarvan de hele verzameling de meeste kans heeft om door de robot te worden goedgekeurd. Hij leert dat het soms beter is om een minder waarschijnlijk woord te kiezen als dat zorgt voor een sterkere tak in de boom.

2. De "Groep-Vergelijking" (Group-Based Optimization)

Het trainen van zo'n boom is lastig omdat de resultaten soms willekeurig lijken. Soms is een zin gewoon moeilijker dan een andere. GTO lost dit op door slim te vergelijken.

• De Analogie: Stel je voor dat je een groep leerlingen hebt die een moeilijke test maken. Als je ze allemaal apart bekijkt, is het lastig om te zien wie echt slim is en wie gewoon een makkelijke of moeilijke vraag kreeg.
  • GTO pakt een kleine groep leerlingen (bijvoorbeeld 8) die net naast elkaar zitten in de klas (dezelfde context).
  • Hij vergelijkt hun antwoorden met elkaar. Als de assistent in deze groep een boom maakt die veel beter werkt dan de "oude versie" van de assistent, krijgt hij een grote beloning.
  • Door ze in groepjes te vergelijken, verdwijnt de ruis. Je ziet niet of de vraag moeilijk was, maar of de strategie van de assistent beter was dan die van zijn concurrenten in dezelfde groep.

Waarom is dit zo goed?

De auteurs hebben dit getest op verschillende taken: chatten, programmeren en wiskunde.

• Het resultaat: De assistent maakt nu veel minder fouten in zijn strategie. Hij leert de "boom" te bouwen die de robot echt nodig heeft.
• De snelheid: Doordat de robot minder hoeft na te denken en vaker de voorspellingen van de assistent kan accepteren, gaat het hele proces 7,7% sneller dan de huidige beste methoden (zoals EAGLE-3).
• De kwaliteit: De antwoorden worden niet slechter; ze zijn net zo slim, maar komen veel sneller binnen.

Samenvatting in één zin

GTO is als het trainen van een voorspeller niet voor een solovogel, maar voor een heel team: hij leert een boom van opties te maken die samenwerken, zodat de trage robot minder werk heeft en alles veel sneller kan doen.

Kortom: Ze hebben de training van de AI aangepast zodat deze precies doet wat hij in de praktijk moet doen, in plaats van wat we dachten dat hij moest doen. Het resultaat is een snellere en slimmere AI.

Technische samenvatting

Titel: Bridging Draft Policy Misalignment: Group Tree Optimization for Speculative Decoding

Publicatie: ICLR 2026 (voorgesteld)
Auteurs: Shijing Hu, Jingyang Li, Zhihui Lu, Pan Zhou

---

1. Het Probleem: Misalignement tussen Training en Decoding

Speculatieve decoding (Speculative Decoding) versnelt de inferentie van Large Language Models (LLM's) door een lichtgewicht "draft model" meerdere tokens te laten voorstellen die parallel worden geverifieerd door het grotere "target model".

Het paper identificeert een fundamenteel probleem: Draft Policy Misalignment.

• Training: Bestaande methoden (zoals EAGLE-3, HASS, GRIFFIN) optimaliseren het draft model voor een greedy pad. Het doel is om de kans te maximaliseren dat het draft model precies hetzelfde token genereert als het target model, stap voor stap, in een enkele lijn.
• Decoding: Tijdens het daadwerkelijke gebruik (decoding) gebruikt het systeem echter een boom-strategie (tree policy). Het draft model bouwt een boom met meerdere takken (sequenties). Deze takken worden opnieuw gerangschikt (re-ranked) op basis van hun waarschijnlijkheid, en de beste takken worden geselecteerd voor verificatie door het target model.

De gevolgen van deze misalignement:

1. Pruning van het optimale pad: Het pad dat tijdens de training als "beste" (greedy) werd gezien, wordt tijdens de decoding vaak weggegooid (gepruned) omdat een broertak (sibling branch) een hogere totale betrouwbaarheid heeft.
2. Verificatie-mismatch: Zelfs als het greedy pad overleeft, accepteert het target model vaak een andere tak dan het pad waar het model voor is getraind.
3. Resultaat: De trainingsinspanning wordt verspild aan een strategie die niet overeenkomt met de realiteit van de decoding, wat de potentiële snelheidswinst beperkt.

2. Methodologie: Group Tree Optimization (GTO)

Om dit probleem op te lossen, stellen de auteurs Group Tree Optimization (GTO) voor. Dit is een trainingsframework dat de trainingsdoelstelling direct afstemt op de boom-gebaseerde decoding-strategie. GTO bestaat uit twee kerncomponenten:

A. Draft Tree Reward (Debeloning voor de Boom)

In plaats van het optimaliseren van token-accuraatheid op een enkel pad, definieert GTO een beloning die de verwachte acceptatielengte van de volledige boom meet.

• Tijdens de training wordt, net als bij decoding, een boom van kandidaat-sequenties gegenereerd.
• De beloning ($r_t$) is de verwachte lengte van de sequenties die door het target model worden geaccepteerd.
• Dit wordt berekend met een smooth maximum (log-sum-exp) over alle takken in de boom. Dit zorgt ervoor dat de beloning differentiabel blijft, maar toch prioriteit geeft aan de sterkste takken in plaats van simpelweg een gemiddelde te nemen.
• Theoretisch bewijs: Het paper bewijst dat het maximaliseren van deze reward theoretisch leidt tot een toename van de verwachte acceptatielengte en dus de snelheidswinst.

B. Groepsgewijze Optimalisatie van het Draft Beleid

Het direct maximaliseren van deze boom-reward is uitdagend vanwege de hoge variantie en de sparsiteit van de beloning. GTO lost dit op met een stabiel, groepsgewijs trainingsalgoritme (geïnspireerd door GRPO):

1. Groepering: Trainingsvoorbeelden worden opgedeeld in groepen van aangrenzende posities in een sequentie. Omdat deze posities een vergelijkbare context hebben, kunnen hun boom-rewards eerlijk met elkaar worden vergeleken.
2. Debiasing: Om systematische moeilijkheidsverschillen tussen verschillende contexten te elimineren, wordt een referentie-draft model (een bevroren model, vaak een goed getraind EAGLE-3 model) gebruikt. De beloning van het huidige model wordt gecorrigeerd door de beloning van het referentiemodel af te trekken.
3. Standaardisatie: Binnen elke groep worden de beloningen genormaliseerd (gemiddelde en standaarddeviatie) om de variantie te verlagen en de "credit assignment" (toewijzing van succes aan specifieke takken) te verbeteren.
4. PPO-achtige Surrogaat: Het model wordt getraind met een afgeknipte (clipped) likelihood-ratio doelstelling, vergelijkbaar met PPO (Proximal Policy Optimization), maar toegepast op de langste geaccepteerde sequentie binnen de boom.

Het trainingsproces verloopt in twee fasen:

• Fase 1 (Warmup): Een referentiemodel wordt getraind met standaard token-level objectives.
• Fase 2 (GTO): Het huidige model wordt gefinetuned met de groepsgewijze boom-reward, gebruikmakend van het referentiemodel voor debiasing.

3. Belangrijkste Bijdragen

1. Identificatie van de Misalignement: Het paper kwantificeert en demonstreert empirisch dat bestaande methoden falen omdat ze trainen voor een greedy pad maar decoderen met een boom-strategie (tot 34% van de greedy pads wordt weggepruned).
2. Novel Training Objective: Introductie van de Draft Tree Reward, een sampling-vrije beloning die direct de prestaties tijdens decoding meet (verwachte acceptatielengte).
3. Stabiel Optimalisatiealgoritme: Ontwikkeling van een groepsgewijze trainingsmethode met debiasing en standaardisatie, die zorgt voor stabiele convergentie ondanks de hoge variantie van boom-rewards.
4. Theoretische Garantie: Een wiskundig bewijs dat het maximaliseren van de voorgestelde reward leidt tot een strikte verbetering van de verwachte acceptatielengte.

4. Resultaten

GTO is getest op diverse LLM's (LLaMA-3.1-8B, LLaMA-3.3-70B, Vicuna-13B, DeepSeek-R1, Qwen3) en benchmarks (MT-Bench, HumanEval, GSM8K).

• Prestaties: GTO overtreft de huidige state-of-the-art (EAGLE-3) consistent.
  • Acceptatielengte: Een stijging van 7,4% in het aantal geaccepteerde tokens per cyclus.
  • Snelheidswinst (Speedup): Een extra 7,7% versnelling in wall-clock tijd ten opzichte van EAGLE-3 (bij temperatuur 0).
• Robuustheid: De verbeteringen zijn consistent over verschillende taken (chat, codering, wiskunde) en temperaturen (0 en 1).
• Generalisatie: GTO werkt ook als een finetuning-laag bovenop andere draft-modellen (zoals GRIFFIN en HASS), wat aantoont dat het een algemene oplossing is voor het misalignement-probleem.
• Ablatie Studies:
  • De keuze voor Log-Sum-Exp (LSE) als aggregatie-operator bleek superieur aan het gebruik van een simpel maximum of gemiddelde.
  • Een groepsgrootte van 4-8 tokens bleek optimaal voor het balanceren van variantie en context-afstemming.
  • Debiasing was cruciaal; zonder debiasing daalde de prestatie aanzienlijk door ruis.

5. Betekenis en Conclusie

Dit paper biedt een praktische en algemene oplossing voor een fundamenteel knelpunt in speculatieve decoding. Door de kloof tussen hoe modellen worden getraind en hoe ze daadwerkelijk worden gebruikt te dichten, haalt GTO de volledige potentieel van speculatieve decoding naar boven.

• Efficiëntie: Het biedt een significante snelheidswinst zonder de kwaliteit van de output te verliezen (lossless).
• Toepasbaarheid: Het is model-agnostisch en kan worden toegepast op bestaande draft-modellen zonder architecturale wijzigingen aan de verificatiestack.
• Kosten: Hoewel de training iets meer rekenkracht vereist (door het bouwen van bomen), weegt dit niet op tegen de aanzienlijke winst in inferentie-efficiëntie, wat GTO zeer aantrekkelijk maakt voor productieomgevingen met hoge latentie-eisen.

Kortom, GTO verschuift de focus van "het voorspellen van het juiste volgende token" naar "het genereren van een boom van tokens die het target model het langst accepteert", wat leidt tot een nieuwe staat van de kunst in snelle LLM-inferentie.