💬 NLP

Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space

Dit artikel introduceert Order-Token Search, een nieuwe decodemethode voor Diffusion Language Models die gezamenlijk de generatievolgorde en de tokenruimte verkent om bestaande baselines te overtreffen op benchmarks voor wiskundige redenering en coderen.

Oorspronkelijke auteurs: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

Gepubliceerd 2026-02-06

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een complexe puzzel probeert op te lossen, zoals een wiskundig probleem of een programmeeruitdaging, maar je hebt een magische assistent (het Diffusion Language Model) die niet simpelweg van links naar rechts schrijft zoals een mens die een letter typt. In plaats daarvan begint deze assistent met een blanco pagina vol "MASK"-placeholders en probeert hij alle gaten in één keer in te vullen door te raden welke woorden er thuishoren.

Het probleem is dat de assistent een beetje chaotisch is. Hij kan de gaten in elke gewenste volgorde invullen. Soms raadt hij het juiste woord voor het eerste gat, maar loopt hij vast omdat hij het verkeerde woord voor het tweede gat heeft geraden. Andere keren raadt hij de juiste woorden, maar in de verkeerde volgorde, wat leidt tot een doodlopend pad.

De Oude Manier: Raden en Snoeien

Voorheen probeerden mensen dit op twee belangrijke manieren op te lossen, die beide gebreken vertoonden:

De "Zelfverzekerde" Aanpak: De assistent kijkt naar de gaten waar hij zich het meest zeker van voelt en vult die als eerste in. Het is als een wandelaar die alleen over het pad loopt dat er het meest solide uitziet.
- Het Goede: Het krijgt meestal snel een goed antwoord.
- Het Slechte: Als het "solide pad" naar een klif leidt (een fout antwoord), zit de wandelaar vast. Hij onderzoekt nooit andere paden die misschien beter waren geweest.
De "Willekeurige" Aanpak: De assistent kiest de gaten om volledig willekeurig in te vullen.
- Het Goede: Het verkent een enorme variëteit aan paden, dus het is zeer waarschijnlijk dat het uiteindelijk de juiste oplossing vindt.
- Het Slechte: Het is zo versnipperd dat het zelden het beste pad bij de eerste poging vindt. Het is als een wandelaar die cirkels loopt; ze vinden de schat misschien wel, maar ze zullen ook veel lege gaten graven voordat ze er zijn.

De Nieuwe Oplossing: Order-Token Search

De auteurs van dit artikel introduceerden een nieuwe methaan genaamd Order-Token Search. Denk aan dit als een Team van Ontdekkingsreizigers die samenwerken.

In plaats van één wandelaar of een chaotische menigte uit te sturen, stuurt de methode een klein team (een "beam") van ontdekkingsreizigers uit. Zo werken zij:

Divergerende Paden (De Zoektocht): Op regelmatige intervallen splitst het team zich op. Elke ontdekkingsreiziger probeert een andere strategie:
- Ontdekkingsreiziger A besluit het eerste ontbrekende woord in te vullen.
- Ontdekkingsreiziger B besluit het laatste ontbrekende woord in te vullen.
- Ontdekkingsreiziger C probeert een ander woord voor het middelste gat.
- Analogie: Ze verkennen zowel waar er als volgende geschreven wordt (de volgorde) als wat er geschreven wordt (het token).
Het Scorebord (De Likelihood Estimator): Dit is het magische deel. Het team heeft een speciale rechter (de likelihood estimator) die niet alleen naar het uiteindelijke antwoord kijkt. In plaats daarvan kijkt de rechter naar elke stap die de ontdekkingsreizigers hebben genomen.
- Maakte de ontdekkingsreiziger een logische zet?
  *, Maakt deze gedeeltelijke zin zin gegeven wat er eerder is geschreven?
- Analogie: Stel je een coach voor die een estafette bekijkt. Als een hardloper vroeg in de race struikelt, wacht de coach niet tot hij de race heeft voltooid om hem af te snijden; de coach stopt hem onmiddellijk omdat de stap die hij zette fout was.
Doodlopende Wegen Afsnijden (Pruning): De rechter scoort de voortgang van elke ontdekkingsreiziger. Als een ontdekkingsreiziger een pad bewandelt dat onwaarschijnlijk succesvol lijkt (zelfs als hij nog niet klaar is), snijdt het team dat pad af en concentreert de middelen zich op de ontdekkingsreizigers die op de beste trajecten zitten.

Waarom Dit Er Toe Doet

Het paper testte dit op moeilijke wiskundige problemen (zoals de GSM8K en MATH500 datasets) en programmeertaken (HumanEval).

Het Resultaat: De "Team van Ontdekkingsreizigers" (Order-Token Search) vond consequent vaker de juiste antwoorden dan de oude "Zelfverzekerde" of "Willekeurige" methoden.
De Vergelijking: Het presteerde zo goed dat het methoden evenaarde of zelfs versloeg die vereisen dat de AI maandenlang dure her-training ondergaat (zoals diffu-GRPO). Dit betekent dat je een veel slimmere AI kunt krijgen door simpelweg te veranderen hoe hij denkt tijdens de test, zonder de hersenen zelf opnieuw te hoeven trainen.

Een Bijzondere Opmerking over Sudoku

Het paper probeerde dit ook op Sudoku-puzzels. Opvallend genoeg werkte het daar niet goed. De auteurs leggen uit dat Sudoku strikte, globale regels vereist (zoals "geen herhalende cijfers in een rij") die de interne "scorekaart" van de AI simpelweg niet lijkt te begrijpen. Het is als het geven van een kaart aan een team van wandelaars die de kliffen niet laat zien; ongeacht hoe goed ze verkennen, ze kunnen de afgrond niet vermijden als de kaart fout is. Dit suggereert dat voor sommige taken de AI zelf anders getraind moet worden, niet alleen de zoekmethode.

In een Notendop

Het paper laat zien dat door een AI te laten exploreren naar meerdere verschillende manieren van schrijven (volgorde) en meerdere verschillende woorden (tokens) tegelijkertijd, en vervolgens een slim scoresysteem te gebruiken om slechte ideeën vroegtijdig af te snijden, we veel betere resultaten kunnen behalen van Diffusion Language Models zonder ze opnieuw te hoeven trainen. Het verandert een chaotisch gokspel in een gestructureerde, efficiënte zoektocht naar de waarheid.

Technische Samenvatting: Verbeteren van de Decodering van Diffusie-Taalmodellen door Gezamenlijke Zoektocht in Generatievolgorde en Tokenruimte

Probleemstelling

Diffusie-Taalmodellen (DLM's), in het bijzonder Masked Diffusion Models (MDM's), bieden een duidelijk voordeel ten opzichte van autoregressieve (AR) modellen: orde-agnostische generatie. In tegenstelling tot AR-modellen die tokens strikt van links naar rechts genereren, kunnen MDM's meerdere posities parallel de-noisen, wat adaptieve generatievolgordes mogelijk maakt. De huidige decoderingsmethoden slagen er echter niet in om deze flexibiliteit volledig te benutten.

Bestaande decoderingsstrategieën committeren zich doorgaans aan een enkele trajectorie per run, gestuurd door heuristische remasking-regels:

Low-Confidence Remasking: Fixeert greedily tokens met een hoge betrouwbaarheid in een vroeg stadium. Hoewel dit vaak sterke single-sample accuratesse (pass@1) oplevert, beperkt het de zoektocht tot een smalle set trajectorieën, wat leidt tot een trage groei in multi-sample dekking (pass@k) en het risico op vastlopen in onjuiste lokale optima.
Random Remasking: Samplet remasking-patronen uniform om de diversiteit van trajectorieën te maximaliseren. Hoewel dit de pass@k verbetert door een breder oplossingsveld te verkennen, lijdt het aan een zwakke pass@1 omdat het geen principieel mechanisme heeft om de juiste trajectorie te selecteren tussen concurrerende kandidaten.

Dit creëert een exploratie–exploitatie-trade-off: huidige methoden of exploiteren te veel lokale betrouwbaarheid (missen geldige oplossingen) of exploreren blind (falen in het selecteren van de beste oplossing). Het artikel betoogt dat het potentieel van DLM's om diverse generatievolgordes en tokenwaarden te verkennen onbenut blijft omdat standaard decodering geen expliciete zoektocht uitvoert in deze gezamenlijke ruimte.

Methodologie: Order-Token Search

Om dit aan te pakken, stellen de auteurs Order-Token Search voor, een decoderingsalgoritme dat beam search uitvoert in de gezamenlijke ruimte van generatievolgorde en tokenwaarden. Het algoritme wisselt af tussen het uitbreiden van kandidaat-trajectorieën en het wegfilteren (pruning) van onbelovende trajectorieën.

1. Zoekproces (Expansie)

Order-Token Search houdt een beam bij van $K$ gedeeltelijk gedenoisede kandidaten. In tegenstelling tot standaard sampling, dat een vaste sequentie van beslissingen volgt, vertakt het algoritme expliciet over twee bronnen van onzekerheid op door de gebruiker gespecificeerde intervallen:

Generatievolgorde: Welke posities (tokens) als volgende ontmaskerd worden.
Tokenwaarden: Welke tokens aan die posities worden toegewezen.

Om de computationele complexiteit te beheersen, maakt de zoektocht gebruik van block diffusion. In plaats van bij elke individuele denoising-stap te vertakken (wat een prohibitieve overhead zou veroorzaken), breidt het algoritme kandidaten uit op de grenzen van aaneengesloten tokenblokken. Binnen elk blok vindt onafhankelijke denoising plaats. Dit reduceert de beslismomenten van de totale sequentielengte naar het aantal blokken, waardoor een substantiële diversiteit in zowel volgorde als tokenassignaties mogelijk wordt gemaakt terwijl de inferentiekosten beheersbaar blijven.

2. Pruning Proces (Likelihood Schatting)

De kerninnovatie ligt in de likelihood estimator die wordt gebruikt om partiële trajectorieën te scoren en te filteren. Een standaard MDM-doelstelling (de waarschijnlijkheid van het reconstrueren van de volledige schone sequentie vanuit een ruisige staat) is instabiel voor intermediaire staten met veel resterende maskers, omdat het het model dwingt moeilijke 'infilling'-problemen op te lossen waar het niet op is getraind tijdens de training.

De auteurs introduceren een incrementele action-level scoring functie:

In plaats van de gehele staat te scoren, scoort het algoritme de specifieke denoising acties die zijn ondernomen om van een meer gecorrumpeerde staat $x_t$ naar een minder gecorrumpeerde staat $x_s$ te gaan.
De score $s(x_t; x_s)$ meet hoe zelfverzekerd het model de nieuw onthulde tokens in $x_s$ kan her-voorspellen, geconditioneerd op de volledige-sequentie voorspelling $x_0$ (die als context dient).
Deze aanpak evalueert of de eerdere beslissingen van een traject (volgorde en tokenkeuzes) worden ondersteund door het model, wat een signaal met een lagere variantie en een grotere discriminatiekracht biedt voor pruning.

De cumulatieve score voor een kandidaat is de som van deze incrementele scores over de zoekintervallen. Het algoritme behoudt de top- $K$ kandidaten na elke expansie, waardoor onjuiste of suboptimale paden effectief worden weggefilterd voordat de decodering wordt voortgezet.

Belangrijkste Bijdragen

Identificatie van de Trade-off: Het artikel toont empirisch aan dat standaard MDM-decoderingmethoden de tegenovergestelde uiteinden van het exploratie–exploitatie-spectrum bezetten. Low-confidence remasking over-exploiteert (hoge pass@1, lage pass@k groei), terwijl random remasking blind exploreert (hoge pass@k, lage pass@1).
Order-Token Search Algoritme: De voorstel van een nieuw decoderingsalgoritme dat zoekt in de gezamenlijke ruimte van generatievolgorde en tokenwaarden. Het maakt gebruik van een stabiele, incrementele likelihood estimator om partiële trajectorieën te filteren, waardoor trajectdiversiteit wordt omgezet in verbeterde single-sample accuratesse.
Empirische Validatie: De methode presteert consistent beter dan baselines op wiskundige redenerings- en coding-benchmarks, waarbij het de winsten evenaart of zelfs overtreft die typisch worden bereikt door dure post-training methoden zoals diffu-GRPO.

Experimentele Resultaten

De auteurs evalueerden Order-Token Search (OTS) op GSM8K, MATH500, Countdown, en HumanEval met behulp van LLaDA-8B-Instruct en LLaDA-1.5 (een RL post-getrainde variant).

Prestatiewinsten: OTS behaalde consequent de hoogste gemiddelde accuratesse over de benchmarks.
- Op LLaDA-8B verbeterde OTS de pass@1 met +3,1% op GSM8K, +3,8% op MATH500, +7,9% op Countdown, en +6,8% op HumanEval vergeleken met de low-confidence remasking baseline.
- Op LLaDA-1.5 behaalde OTS vergelijkbare relatieve verbeteringen, waarbij een gemiddelde accuratesse van 40,4% werd bereikt vergeleken met 37,2% voor de beste baseline (Low-conf + MV).
Vergelijking met Post-Training: OTS kwam in de buurt van of overtrof de prestaties van modellen die onderhevig zijn aan reinforcement learning post-training (diffu-GRPO), wat suggereert dat zoektochten tijdens de inferentie even effectief kunnen zijn als het hertrainen van het model voor het verbeteren van redeneercapaciteiten.
Ablatie-studies:
- Likelihood Estimator: De specifieke incrementele scoring functie (Eq. 2) was cruciaal. Varianten die de joint likelihood van alle onthulde blokken of AR-stijl likelihoods gebruikten, presteerden aanzienlijk slechter, wat bevestigt dat de op maat gemaakte estimator noodzakelijk is voor stabiele pruning in MDM's.
- Zoekruimte: OTS presteerde beter dan "Order Search" (zoeken alleen naar volgorde) en "Token Search" (zoeken alleen naar tokens) die AR-stijl likelihoods gebruikten, wat de noodzaak van gezamenlijke zoektocht en de mismatch van AR-heuristieken met MDM's benadrukt.
Efficiëntie: OTS schaalt effectiever met test-time compute (NFE) dan majority voting. Bij een gelijkwaardig compute-budget behaalde OTS een hogere accuratesse (bijv. 29,3% vs. 19,9% op Countdown) door gestructureerde zoektocht te gebruiken in plaats van ongestructureerde sampling.

Betekenis en Claims

Het artikel claimt dat gezamenlijke order–token zoektocht een sleutelcomponent is voor het bevorderen van DLM-decodering. Door expliciet de ruimte van generatievolgordes en tokenwaarden te verkennen en te filteren op basis van een stabiele likelihood estimator, zet de methode de inherente multi-trajectorie potentie van diffusiemodellen om in betrouwbare single-sample prestaties.

De auteurs positioneren dit werk als een oplossing voor de "exploratie–exploitatie" bottleneck in de huidige MDM-decodering. Zij stellen dat hun aanpak DLM's in staat stelt om te ontsnappen aan de lokale optima die greedy methoden gevangen houden en de juiste oplossingen te selecteren die random sampling mist, en dat dit alles kan zonder aanvullende modeltraining. Het werk vestigt dat inferentie-tijd zoekstrategieën, specifelijk afgestemd op de parallelle, orde-agnostische aard van diffusiemodellen, systematische verbeteringen kunnen opleveren voor redeneer- en coding-taken.