PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding

Stel je voor dat je probeert het volgende woord in een verhaal te voorspellen. Lange tijd deden de slimste computers (Large Language Models) dit woord voor woord, zoals iemand een boek hardop voorleest, waarbij ze na elk enkel woord stoppen om na te denken over wat er volgt. Dit is accuraat, maar traag. Recentelijk hebben wetenschappers een nieuwe manier ontdekt om deze verhalen te schrijven met behulp van "diffusie"-modellen. Denk hierbij aan een beeldhouwer die begint met een blok marmer en in één keer het hele beeld uithouwt, in plaats van telkens een klein stukje weg te hakken. Dit stelt de computer in staat om veel woorden tegelijkertijd te raden, wat ongelooflijk snel is.

Er is echter een addertje onder het gras. Wanneer je veel woorden tegelijkertijd raadt, kun je er een paar fout doen. Om dit te herstellen, is er een slimme truc genaamd "speculative decoding". Het is alsoals het hebben van een snelle, junior assistent die de volgende paar woorden raadt, en dan een trage, superintelligente baas die controleert of die gokken juist zijn. Als de baas ermee instemt, accepteert hij de hele batch woorden onmiddellijk, wat een enorme tijd bespaart. Het probleem is dat de junior assistent (het diffusiemodel) geweldig is in het raden van individuele woorden, maar niet altijd weet hoe die woorden in een specifieke volgorde bij elkaar passen. Het is alsof de assistent heel goed is in het kiezen van individuele ingrediënten voor een taart, maar niet altijd weet welke combinatie van ingrediënten uiteindelijk echt lekker zal smaken.

Dit is waar een nieuw paper over gaat. De onderzoekers, onder leiding van Zheng Wang en collega's, realiseerden zich dat de huidige manier van het gebruiken van deze snelle assistenten veel snelheid onbenut laat. Ze ontdekten dat hoewel de assistent een enorme variëteit aan mogelijke woordcombinaties kan genereren, de huidige methode slechts één enkel pad controleert, alsof je door één enkele gang loopt en hoopt dat de deur aan het einde openstaat. De auteurs stellen een nieuw systeem voor genaamd PRESTO (Prefix-Aligned Tree Drafting). In plaats van slechts één gang af te lopen, bouwt PRESTO een boom van mogelijkheden, waarbij meerdere paden tegelijkert === worden verkend. Maar hier is de magie: het lost een fundamenteel gebrek op in de manier waarop het vertrouwen van de assistent wordt gemeten. De oorspronkelijke betrouwbaarheid van de assistent is "prefix-blind", wat betekent dat het niet geeft om welke woorden eraan voorafgingen. PRESTO voegt een "prefix-aligned" score toe, die fungeert als een kompas, om ervoor te zorgen dat de paden die gekozen worden voor exploratie, de paden zijn die het meest waarschijnlijk geaccepteerd zullen worden door de baas.

Het resultaat is een systeem dat aanzienlijk sneller is. In hun tests hielp PRESTO de computer om meer woorden te accepteren in elke ronde van het raden. Op sommige van de beste bestaande opstellingen maakte het het hele proces 1,5 keer sneller. Op andere opstellingen zorgde het voor een versnelling van 1,12 keer. Het paper suggereert dat door het raden van woorden te behandelen als een boom-vertakkend avontuur in plaats van een rechte lijn, en door ervoor te zorgen dat de takken worden gekozen op basis van hoe goed ze bij het verhaal passen, we het beste van beide werelden krijgen: de snelheid van diffusie en de nauwkeurigheid van zorgvuldige controle.

Het Problek: De "Eén-Pad" Valstrik

Om te begrijpen waarom PRESTO nodig is, stel je voor dat je een spelletje "Mad Libs" speelt met een vriend die probeert de ontbrekende woorden te raden. Je vriend is een diffusiemodel. Ze zijn geweldig in het bekijken van een lege ruimte en zeggen: "Ik wed dat het woord hier 'kat' is!" of "Misschien is het 'hond'?" of "Of misschien wel 'raket'?" Ze kunnen al deze opties op exact hetzelfde moment uitroepen.

De huidige methode van het gebruiken van deze vriend is echter erg rigide. Het neemt de beste gok van de vriend, schrijft deze op, en vraagt dan aan de "baas" (het doelmodel) of die gok juist is. Als de baas zegt "Nee", wordt het hele ding weggegooid en moet je opnieuw beginnen. Als de baas "Ja" zegt, ga je naar het volgende woord en herhaal je het proces. Dit wordt lineaire drafting genoemd. Het is alsof je door een bos loopt en alleen naar het pad kijkt dat recht voor je ligt.

De auteurs merkten op dat deze aanpak inefficiënt is. Omdat het diffusiemodel veel opties tegelijkertijd genereert, is er een enorme "combinatorische ruimte" aan mogelijkheden. Het is alsof je een kaart hebt met een duizend verschillende paden, maar je mag er slechts op één lopen. Het paper laat zien dat door vast te houden aan slechts één pad, het systeem veel geldige routes mist. In feite, op wiskundige problemen zoals GSM8K, accepteerde de huidige methode gemiddeld ongeveer 6,5 woorden, maar de onderzoekers berekenden dat als ze alle beste paden hadden kunnen controleren, ze bijna 10 woorden hadden kunnen accepteren. Dat is een enorm gat!

De Mismatch: Het "Blinde" Kompas

De onderzoekers gingen dieper in en ontdekten een specifieke reden waarom het simpelweg controleren van meer paden (het bouwen van een boom) niet perfect werkte met de oude methoden. Ze identificeerden een "fundamentele mismatch".

In de wereld van standaard AI (autoregressieve modellen) hangt de betrouwbaarheidsscore van een woord sterk af van de woorden die eraan voorafgingen. Als de zin is "De kat zat op de...", weet het model dat "mat" een zeer waarschijnlijke volgende stap is, maar "pizza" niet. Dit is prefix-aligned.

Maar diffusiemodellen werken anders. Ze genereren een "marginale" waarschijnlijkheid voor elke positie onafhankelijk van elkaar. Het is alsof het model zegt: "Op positie 5 is 'kat' voor 80% waarschijnlijk," zonder te geven of positie 4 "De" of "De snelle bruine" was. Dit is prefix-blind.

Wanneer je een boom van gokjes probeert te bouwen met deze blinde scores, krijg je een rangschikkingsprobleof. Je kiest misschien een pad dat geweldig lijkt voor het eerste woord, maar verschrikkelijk voor het tweede woord, omdat het model niet besefte dat het eerste woord de context veranderde. Het is als een GPS die je de weg wijst op basis van alleen de huidige straat, zonder rekening te houden met het feit dat je net links bent afgeslagen en nu op een eenrichtingsweg rijdt. Het paper betoogt dat het gebruik van deze blinde scores om een boom te bouwen leidt tot "onbetrouwbare pad-rangschikking", wat betekent dat het systeem de verkeerde takken verkent en tijd verspilt.

De Oplossing: PRESTO

PRESTO (Prefix-Aligned Scoring and priority-based Tree search for diffusion Speculative decOding) lost dit op door een "correctie" toe te voegen aan de scores van het diffusiemodel.

1. Prefix-Aligned Scoring: De auteurs realiseerden zich dat ze het sterke "marginale" signaal van het diffusiemodel (hoe waarschijnlijk een woord op zichzelf is) moesten combineren met een "prefix-conditioned" signaal (hoe waarschijnlijk het is gegeven de voorgaande woorden). Ze creëerden een nieuwe scoringsformule die de diffusiewaarheid vermenigvuldigt met een correctiefactor afgeleid van een eenvoudig n-gram model (een lichtgewicht hulpmiddel dat woordcombinaties bekijkt). Dit creëert een score die de flow van het verhaal respecteert.
2. Priority-Based Tree Search: In plaats van alleen het bovenste pad te kiezen, bouwt PRESTO een boom. Het gebruikt de nieuwe, gecorrigeerde scores om te beslissen welke takken gegroeid moeten worden. Het geeft prioriteit aan paden die de grootste potentie hebben om geaccepteerd te worden door de baas. Het is als een wandelaar die, in plaats van gewoon rechtdoor te lopen, een kaart bekijkt en het pad kiest dat het meest waarschijnlijk naar de top leidt, zelfs als dat pad niet het meest voor de hand liggend is aan het begin.

Het paper testte twee manieren om deze boom te laten groeien: Beam Search (het behouden van een vast aantal top-paden op elke stap) en Best-First Search (altijd het enkelvoudige beste pad dat tot nu toe is gevonden uitbreiden). Ze ontdekten dat voor hun specifieke opstelling Beam Search net zo goed werkte als de complexere Best-First Search, dus kozen ze voor de simpelere, efficiëntere optie.

De Resultaten: Sneller en Slimmer

De auteurs testten PRESTO op diverse taken, waaronder wiskundige problemen (GSM8K, Math500), programmeeruitdagingen (HumanEval, LiveCodeBench) en chatgesprekken. Ze gebruikten twee verschillende soorten systemen:

• Dedicated Diffusion Drafters: Een klein, snel diffusiemodel dat gokt voor een groter autoregressief model (zoals dFlash).
• Self-Speculative Diffusion LLMs: Een enkel diffusiemodel dat zichzelf gokt en controleert (zoals Nemotron-Labs-Diffusion).

De resultaten waren consistent over de hele linie. PRESTO verhoogde consequent de Average Acceptance Length, wat het aantal woorden is dat de baas in één keer accepteert.

• Op het dFlash-systeem (gebruikmakend van Qwen3-8B) steeg de gemiddelde acceptance length van ongeveer 6,6 woorden naar 9,6 woorden. Dit vertaalde zich in een end-to-end versnelling van 1,5 keer.
• Op het Nemotron-Labs-Diffusion-systeem steeg de acceptance length van 8,8 naar 9,9 woorden, wat een versnelling van 1,12 keer opleverde.

Misschien wel het meest indrukwekkende is dat het paper liet zien dat PRESTO werkt zelfs wanneer het systeem "stochastisch" (gerandomiseerd) is, wat meestal moeilijker te voorspellen is. In deze gevallen was de versnelling nog prominenter aanwezig, waarbij sommige benchmarks bijna een verdubbeling van de throughput lieten zien.

De auteurs controleerden ook de "kosten" van deze nieuwe methode. Ze vonden dat de extra arbeid die nodig was om de boom te bouwen en de nieuwe scores te berekenen minimaal was — minder dan 4% van de totale tijd. Het overgrote deel van de tijd (meer dan 90%) werd nog steeds besteed aan de eigenlijke verificatie door het baasmodel. Dit betekent dat PRESTO een zeer efficiënte upgrade is die het systeem niet vertraagt met overhead.

Wat PRESTO NIET is

Het is belangrijk om te vermelden wat het paper niet beweert. De auteurs stellen expliciet dat het simpelweg toepassen van een naïeve boomstructuur (zonder hun prefix-aligned scoring) suboptimaal is. Als je simpelweg de ruwe scores van het diffusiemodel neemt en een boom bouwt, haal je niet het volledige voordeel omdat van de "blinde" aard van de scores. PRESTO is specifiek bedoeld om die scoring-mismatch op te lossen.

Bovendien claimt het paper niet dat het het probleem van diffusiemodellen volledig heeft opgelost. Ze erkennen dat hun methode vertrouwt op een "tractable prefix-aligned signaal" (zoals het n-gram model dat ze gebruikten) om het zware werk te doen. Ze suggereren dat toekomstig werk zelfs rijkere signalen zou kunnen verkennen, maar voor nu is hun eenvoudige correctie voldoende om enorme winsten te boeken.

Waarom dit ertoe doet

In de race om AI sneller en efficiënter te maken, telt elk beetje snelheid. Speculative decoding is een populair onderwerp geworden omdat het ons in staat stelt de beste, meest accurate modellen te gebruiken zonder de volledige tijdstraf te betalen. Echter, de huidige methoden werden beperkt door het feit dat ze diffusiemodellen behandelden als lineaire machines, waarbij ze hun unieke vermogen om tegelijkertijd veel opties te genereren negeerden.

PRESTO verandert de regels door diffusiemodellen te behandelen als de multi-path verkenners die ze zijn. Door de scoring af te stemmen op de manier waarop de baas het werk controleert, ontsluit het het volledige potentieel van de parallelle generatie van het diffusiemodel. Het resultaat is een systeem dat niet alleen iets sneller is, maar aanzienlijk efficiënter, waardoor we tekst kunnen genereren, wiskundige problemen kunnen oplossen en code kunnen schrijven met snelheden die voorheen als onmogelijk werden beschouwd voor dit type modellen. Zoals de auteurs zeggen: ze hebben een "één-pad" wandeling veranderd in een "boom-gebaseerde" expeditie, waardoor ze ervoor zorgen dat elke stap een stap richting het juiste antwoord is.