d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation

Het artikel introduceert d2, een nieuw reinforcement learning-framework voor masked diffusion taalmodellen dat gespecialiseerde traject-likelihoodschatters (d2-AnyOrder en d2-StepMerge) gebruikt om de redeneercapaciteiten op logische en wiskundige benchmarks aanzienlijk te verbeteren, waarbij nieuwe state-of-the-art prestaties worden behaald.

De kern

Stel je voor dat je een zeer getalenteerde kunstenaar hebt die een meesterwerk kan schilderen door te beginnen met een leeg canvas en geleidelijk details toe te voegen totdat de afbeelding compleet is. Dit is hoe Diffusion Language Models (DLMs) werken. In tegenstelling tot traditionele AI-schrijvers die zinnen woord voor woord opbouwen (zoals het stapelen van bakstenen), beginnen deze modellen met een rommelige massa van "ruis" (zoals statische ruis op een oude tv) en "ontruisen" deze stapsgewijs om een samenhangend verhaal of een oplossing te onthullen.

Het paper introduceert een nieuwe trainingsmethode genaamd d2 om deze kunstenaars te leren hoe ze veel beter moeilijke puzzels (zoals wiskundeproblemen of logische spellen) kunnen oplossen. Hier is de uitleg over hoe het werkt, met behulp van eenvoudige analogieën.

Het Probleem: De "Blinde" Criticus

Om een AI te leren redeneren, gebruiken onderzoekers Reinforcement Learning (RL). Denk hierbij aan een spel waarbij de AI een puzzel probeert op te lossen, en een "Critic" (een beloningssysteem) een score geeft. Als de AI een goede score haalt, leert hij dit opnieuw te doen.

Er is echter een addertje onder het gras bij Diffusion-modellen. Om effectief te leren, moet de AI weten hoe waarschijnlijk het was dat hij precies die specifieke reeks woorden genereerde.

• Voor traditionele AI (Autoregressief): Het is als het schrijven van een brief. Je weet precies wat je eerst, tweede en derde hebt geschreven. Het berekenen van de "waarschijnlijkheid" is eenvoudig en snel.
• Voor Diffusion AI: Het is als het kijken naar een beeldhouwwerk dat uit een blok steen tevoorschijn komt. De AI maakt duizenden kleine aanpassingen. Om de exacte waarschijnlijkheid van het uiteindelijke beeldhouwwerk te weten, zou je theoretisch de hele video moeten terugspoelen en elke enkele frame van het beeldhouwproces opnieuw moeten afspelen. Dit doen voor elke trainingsstap is computationeel onmogelijk (te traag en te duur).

Omdat ze deze "score" niet nauwkeurig konden berekenen, waren eerdere methoden aan het gokken, wat leidde tot zwakke redeneervaardigheden.

De Oplossing: d2 (De Slimme Scorekeeper)

De auteurs creëerden d2, een framework dat fungeert als een super-efficiënte scorekeeper. Het bepaalt de "waarschijnlijkheid" van de reis van de AI zonder de hele film telkens opnieuw te hoeven afspelen. Ze bouwden twee verschillende tools voor twee verschillende soorten kunstenaars:

Tool 1: d2-AnyOrder (De "Magische Spiegel")

Sommige diffusion-modellen zijn speciaal; ze kunnen de uiteindelijke afbeelding in elke willekeurige volgorde onthullen. Je zou de ogen eerst kunnen schilderen, dan de neus, en dan de achtergrond, of andersom, en het werkt nog steeds.

• De Analogie: Stel je voor dat je een magische spiegel hebt. In plaats van het schilderproces stap voor stap te bekijken, kun je naar het voltooide schilderij kijken en direct zien wat de kunstenaar in elke stap zou hebben gezegd, allemaal tegelijkertijd.
• Het Resultaat: Deze tool, genaamd d2-AnyOrder, berekent de perfecte score in één enkele blik (één computerpass). Het is ongelooflijk nauwkeurig en zorgt ervoor dat de AI zeer snel leert.

Tool 2: d2-StepMerge (De "Chunking Strategie")

De meeste populaire diffusion-modellen (zoals het beroemde LLaDA) kunnen niet in elke willekeurige volgorde schilderen; ze moeten een specifieke sequentie volgen. Ze hebben geen "magische spiegel".

• De Analogie: Omdat we de hele film niet in één keer kunnen zien, breken we de film op in stukjes (chunks). In plaats van elke enkele frame te bekijken, kijken we de film in segmenten van 10 seconden. We schatten de score voor elk segment in en tellen deze bij elkaar op.
• Het Resultaat: Deze tool, genaamd d2-StepMerge, is een benadering. Het is niet perfect, maar de auteurs hebben wiskundig bewezen dat de fout klein en controleerbaar is. Het ruilt een klein beetje nauwkeurigheid in voor een enorme winst in snelheid, waardoor het praktisch bruikbaar is voor standaardmodellen.

De Resultaten: Van Novice naar Grandmaster

De onderzoekers testten deze tools op enkele van de moeilijkste puzzels voor AI:

1. Logische Spellen: Sudoku en Countdown (een wiskundespel).
2. Wiskundeproblemen: GSM8K en MATH500 (standaard benchmarks voor wiskundig redeneren).

De Uitkomst:

• Zonder extra "spiekbriefjes" (Supervised Fine-Tuning): Het d2-framework leerde de modellen beter te redeneren dan welke eerdere methode dan ook.
• De Competitie Verslaan: Op de Sudoku- en Countdown-tests behaalden de met d2 getrainde modellen scores van respectievelijk bijna 92% en 56%, waarmee ze de vorige beste methoden, die rond de 22% en 42% bleven hangen, verpletterden.
• Wiskundige Meesterschap: Op complexe wiskundige benchmarks vestigde d2 een nieuw "State-of-the-Art" record, waarmee bewezen werd dat diffusion-modellen net zo goed kunnen redeneren als traditionele tekstgenererende modellen, maar met de toegevoegde voordelen van snelheid en controleerbaarheid.

Samenvatting

Het paper zegt: "We hebben een manier gevonden om Diffusion AI-modellen helder te leren denken."

• Ze realiseerden zich dat de oude manier om het werk van de AI te beoordelen te traag en onnauwkeurig was.
• Ze hebben d2 uitgevonden, een nieuw beoordelingssysteem dat ofwel perfect nauwkeurig is (voor speciale modellen) of slim efficiënt (voor standaardmodellen).
• Met behulp van dit nieuwe systeem leerden de AI-modellen logische en wiskundige puzzels aanzienlijk beter op te lossen dan voorheen, zonder dat ze vooraf de antwoorden hoefden te leren.

Kortom, d2 geeft diffusion-modellen het vermogen om nauwkeurig terug te kijken op hun eigen denkproces, waardoor ze uit hun fouten kunnen leren en echte redeneerexperts kunnen worden.

Technische samenvatting

Technische Samenvatting: d2: Verbeteren van Redeneren in Diffusie Taalmodellen via Traject-waarschijnlijkheidsschatting

Probleemstelling

Diffusie Taalmodellen (DLMs), in het bijzonder gemaskeerde DLMs, zijn opgekomen als competitieve alternatieven voor autoregressieve (AR) modellen vanwege hun controleerbaarheid en parallelle generatiecapaciteiten. Het verbeteren van hun redeneervermogen via Reinforcement Learning (RL) blijft echter uitdagend. Hoewel RL de standaard is geworden voor het induceren van redeneren in AR-modellen, is het toepassen ervan op DLMs niet triviaal. De kern van de moeilijkheid ligt in de policy gradient-formulering: in tegen tegenstelling tot AR-modellen waarbij waarschijnlijkheden netjes factoriseren over tokenposities, is de exacte waarschijnlijkheid van een diffusietraject computationeel onhandelbaar. Naieve berekening van deze waarschijnlijkheden vereist $T$ forward passes (waarbij $T$ het aantal diffusiestappen is), wat standaard policy gradient-methoden zoals Group Relative Policy Optimization (GRPO) computationeel prohibitief maakt. Bestaande benaderingen introduceren vaak significante bias, wat leidt tot suboptimale policy-updates.

Methodologie

De auteurs introduceren d2, een redeneerkader dat is afgestemd op gemaskeerde DLMs. Het kader centreert zich rond een nieuw policy gradient-algoritme afgeleid van de GRPO-objectief, die expliciet een nauwkeurige schatting van de sampling-traject-waarschijnlijkheden vereist. Om de computationele bottleneck aan te pakken, stelt het paper een familie van estimators voor die zijn afgestemd op verschillende modelklassen:

1. Theoretische Fundering

De auteurs leiden een GRPO-stijl objectief af voor gemaskeerde DLMs door de waarschijnlijkheid te marginaliseren over tijd-latents en gebruik te maken van importance sampling. De resulterende gradiëntdecompositie (Theorem 3.1) is afhankelijk van de ratio van traject-waarschijnlijkheden tussen de huidige policy $\pi_\theta$ en een verouderde policy $\pi_{old}$. De belangrijkste uitdaging is het efficiënt evalueren van $\pi(x_{0:T})$.

2. De Estimators

Het d2-kader biedt twee specifieke estimators gebaseerd op de architectuur van het model:

• d2-AnyOrder (Exacte Estimator):

  • Doel: DLMs die any-order decoding (AO-dLLMs) ondersteunen. Deze modellen laten toe dat tokens in elke willekeurige volgorde worden gedecodeerd, mits de attention-mask de specifieke causaliteitsbeperkingen respecteert (onafhankelijke masks en volgorde-causaliteit).
  • Mechanisme: Het construeert een sequentie van lengte $2L$ (het samenvoegen van de schone tokens en mask-tokens) met specifieke positionele encodings en attention masks. Dit stelt het model in staat om de exacte conditionele waarschijnlijkheden van alle tokens in een enkele forward pass te berekenen.
  • Beperking: Standaard gemaskeerde DLMs (zoals LLaDA) ondersteunen van nature geen any-order decoding; hun trainingsdoelstellingen garanderen niet de noodzakelijke attention-eigenschappen voor deze estimator om onbevooroordeeld te zijn.

• d2-StepMerge (Benaderende Estimator):

  • Doel: Standaard gemaskeerde diffusiemodellen (MDMs) die geen any-order decoding ondersteunen.
  • Mechanisme: Geïnspireerd door block composite likelihood, benadert deze estimator de traject-waarschijnlijkheid door het $T$-staps traject te verdelen in $N$ opeenvolgende tijdssegmenten. In plaats van elke stap te evalueren, evalueert het de waarschijnlijkheid van tokens die binnen elk segment niet gemaskeerd zijn, met behulp van een enkele forward pass per segment.
  • Trade-off: Het vermindert het aantal forward passes van $T$ naar $N$. Het paper biedt een theoretische bovengrens (Theorem 4.1) op de benaderingsfout (KL-divergentie), waarbij wordt aangetoond dat de fout monotoon afneemt naarmate $N$ toeneemt, wat de compute-bias trade-off kwantificeert.

Belangrijkste Bijdragen

1. d2 Framework: Een principieel RL post-training kader voor gemaskeerde DLMs dat een GRPO-objectief afleidt dat afhankelijk is van nauwkeurige traject-waarschijnlijkheidsschatting.
2. d2-AnyOrder: Een exacte, single-pass likelihood estimator voor AO-dLLMs, die aantoont dat onbevooroordeelde likelihood-estimatie bestaande baselines in RL-settings aanzienlijk overtreft.
3. d2-StepMerge: Een deterministische, multi-pass estimator voor standaard MDMs met analytisch begrensde benaderingsfouten, die een superieure efficiëntie-accuratesse trade-off biedt vergeleken met eerdere methoden zoals diffu-GRPO.
4. Empirische Validatie: Uitgebreide experimenten die aantonen dat d2 state-of-the-art prestaties bereikt op redeneerbenchmarks zonder te vertroukken op supervised chain-of-thought (CoT) fine-tuning.

Experimentele Resultaten

De auteurs evalueerden d2 op logische redeneertaken (Countdown, Sudoku) en wiskundige redeneertaken (GSM8K, MATH500), evenals programmeertaken en toxiciteit-sturing.

• Prestaties op AO-dLLMs: Wanneer toegepast op modellen die any-order decoding ondersteunen (bijv. Eso-LM, Set Diffusion, en een aangepaste Any-Order Causal LLaDA), presteerde d2-AnyOrder aanzienlijk beter dan baselines zoals DDPO en diffu-GRPO. Bijvoorbeeld, op GSM8K met een Qwen3-1.7B base, behaalde d2-AnyOrder een nauwkeurigheid van 67% vergeleken met 63% voor diffu-GRPO.
• Prestaties op Standaard MDMs: Toegepast op LLaDA-8B-Instruct (een standaard gemaskeerd DLM), zette d2-StepMerge (met $N=16$) nieuwe state-of-the-art resultaten neer:
  • Sudoku: 91.9% (vs. 22.1% voor d1 met SFT).
  • Countdown: 56.6% (vs. 42.2% voor d1).
  • GSM8K: 85.0% (vs. 82.1% voor d1).
  • MATH500: 41.6% (vs. 39.0% voor wd1).
  • Opvallend genoeg werden deze resultaten behaald zonder supervised fine-tuning op extra CoT-data, terwijl sommige baselines daar wel afhankelijk van waren.
• Generalisatie: De methode generaliseerde naar andere architecturen (Dream 7B) en taken (HumanEval, MBPP), waarbij een verbeterde lengte-generalisatie werd getoond vergeleken met diffu-GRPO.
• Ablatie: Experimenten bevestigden dat exacte likelihoods (d2-AnyOrder) hogere-getrouwheidssignalen leveren dan benaderingen, en dat d2-StepMerge met een optimale $N$ (bijv. 16) de balans tussen compute-kosten en accuratesse beter beheert dan extreme waarden.

Betekenis en Claims

Het paper claimt dat nauwkeurige traject-waarschijnlijkheidsschatting centraal staat voor het succes van RL voor redeneren in DLMs. De auteurs stellen dat vorige methoden niet optimaal presteerden omdat ze vertrouwden op bevooroordeelde of computationeel dure likelihood-schattingen.

• Theoretische Bijdrage: Het werk biedt de eerste rigoureuze afleiding van een GRPO-objectief voor gemaskeerde DLMs en stelt theoretische grenzen vast aan de benaderingsfout van step-merging strategieën.
• Praktische Impact: d2 demonstreert dat diffusiemodellen redeneervermogen kunnen bereiken dat concurrerend is met, of superieur aan, autoregressieve modellen wanneer ze correct worden getraind met RL, zelfs zonder uitgebreide supervised CoT-data.
• Toekomstperspectief: De auteurs merken op dat hoewel d2-AnyOrder superieur is, de toepasbaarheid momenteel beperkt is tot modellen die specifiek zijn getraind voor any-order decoding. Zij beschouwen de ontwikkeling van sterkere algemene purpose any-order dLLMs als een natuurlijke volgende stap om de exacte estimator volledig te benutten.

Het paper concludeert dat d2 een nieuwe state-of-the-art instelt voor DLMs op logische en wiskundige redeneertaken, waarmee de effectiviteit van hun likelihood-estimatiestrategieën bij het ontsluiten van het redeneerpotentieel van diffusie-gebaseerde taalmodellen wordt gevalideerd.