Technische Samenvatting: d-TreeRPO

Probleemstelling

Diffusie-taalmodellen (dLLMs) bieden een parallelle decoderingsalternatief voor autoregressieve (AR) modellen, waardoor snellere inferentie mogelijk wordt. Het toepassen van Versterkingsleer (RL) om de redeneercapaciteiten van dLLMs te verbeteren, stuit echter op twee kritieke betrouwbaarheidsknelpunten die bestaande methoden niet aanpakken:

1. Reward-Sparsiteit en Gebrek aan Verifieerbaarheid: Bestaande methoden voor beleidsoptimalisatie bij dLLMs vertrouwen vaak op grove uitkomstrewards (bijvoorbeeld binaire correctheid van het uiteindelijke antwoord) of ongeverifieerde processignalen. Deze sparsiteit belemmert een nauwkeurige berekening van het voordeel, terwijl ongeverifieerde signalen kunnen leiden tot reward-hacking.
2. Onbetrouwbare Kansschatting: In tegenstelling tot AR-modellen, die sequenties via de kettingregel ontleden voor exacte kansberekening in één forward pass, genereren dLLMs tekst via een "any-order" denoising-proces. De ware log-kans van een actie vereist een verwachting over alle mogelijke decoderingsordes, wat onberekenbaar is. Bestaande benaderingen (bijvoorbeeld ELBO-gebaseerde ondergrenzen of single-step forward passes) lijden onder niet-geanalyseerde benaderingsfouten, wat de betrouwbaarheid van beleidsgradiënt-updates ondermijnt.

Methodologie: d-TreeRPO

De auteurs stellen d-TreeRPO voor, een betrouwbaar RL-kader dat specifiek is ontworpen voor dLLMs. De methode adresseert de bovengenoemde knelpunten via drie kerncomponenten:

1. Boomgestructureerde Rollouts voor Granulaire Rewards

Om reward-sparsiteit te overwinnen, structureert d-TreeRPO het rollout-proces als een boom in plaats van een vlakke reeks onafhankelijke steekproeven.

• Structuur: De $N$ denoising-stappen worden gegroepeerd in $H$ boomstappen. Bij elke niet-leaf-node worden $B$ onafhankelijke steekproeven uitgevoerd, waardoor een boom ontstaat met $B^H$ leaf-nodes die terminale toestanden vertegenwoordigen.
• Bottom-Up Berekening van Voordeel: Leaf-nodes krijgen verifieerbare uitkomstrewards (bijvoorbeeld correctheid van de puzzeloplossing). State-waarden voor parent-nodes worden bottom-up berekend als het gemiddelde van de waarden van hun kinderen.
• Fijnmazige Voordelen: Het voordeel voor een overgang van een parent-node $p$ naar een child-node $c$ wordt gedefinieerd als $A_p^c = V_c - V_p$. Dit biedt stap-voor-stap, verifieerbare reward-signalen die granulerer zijn dan uniforme uitkomstrewards.

2. Theoretische Analyse van Kansschatting

Het artikel adresseert de onberekenbaarheid van onbevooroordeelde kansschatting in dLLMs.

• Theoretisch Bewijs: De auteurs bewijzen (Stelling 1) dat de kloof tussen de kansschatting via een single-time forward pass en de onbevooroordeelde verwachte kans afneemt naarmate de voorspellingszekerheid van het model toeneemt. Specifiek wordt de schattingsfoutgrens strakker naarmate het model deterministischer wordt.
• Implicatie: Hoewel hoge zekerheid de schattingsnauwkeurigheid verbetert, creëert dit een afweging tussen exploratie en exploitatie, aangezien te deterministische beleidslijnen in een vroeg stadium van het training het onderzoek kunnen belemmeren.

3. Tijdgeplande Self-Distillation-Loss

Om de afweging tussen exploratie (vroege training) en nauwkeurige kansschatting (late training) in evenwicht te brengen, introduceert d-TreeRPO een tijdgeplande self-distillation-doelstelling.

• Mechanisme: De loss moedigt het beleid aan om zich aan te passen aan een doelverdeling die is opgebouwd uit child-nodes met positieve voordelen.
• Planning:
  • Temperatuur ($\tau(t)$): Start hoog om diverse steekproeven toe te staan en neemt af in de tijd om de doelverdeling scherper te maken.
  • Gewicht ($\lambda(t)$): Start dicht bij nul om de loss te minimaliseren tijdens de exploratiefase en neemt progressief toe om determinisme af te dwingen en schattingsfouten in latere stadia te verminderen.
• Doel: Deze loss wordt toegevoegd aan de standaard GRPO-stijl beleidsgradiënt-doelstelling, waardoor modeldeterminisme alleen wordt versterkt wanneer het veilig is, en zo de nauwkeurigheid van de single-pass kansschattingen die worden gebruikt in de RL-update wordt verbeterd.

Belangrijkste Bijdragen

1. d-TreeRPO Kader: Een nieuw algoritme voor beleidsoptimalisatie voor dLLMs dat boomgestructureerde rollouts gebruikt om fijnmazige, verifieerbare stap-voor-stap rewards en voordelen te bieden.
2. Theoretisch Inzicht & Self-Distillation: Een theoretisch bewijs dat voorspellingszekerheid koppelt aan vermindering van schattingsfouten, wat leidt tot het ontwerp van een tijdgeplande self-distillation-loss die dynamisch de afweging tussen exploratie en exploitatie beheert.
3. Empirische Validatie: Uitgebreide experimenten die aantonen dat d-TreeRPO bestaande baselines (inclusief Diffu-GRPO, VRPO en wd1) overtreft op meerdere redeneerbenchmarks.

Experimentele Resultaten

De auteurs hebben d-TreeRPO geëvalueerd op LLaDA-8B-Instruct en LLaDA-MoE-7BA1B-Instruct over vier redeneerbenchmarks: Sudoku, Countdown, GSM8K en Math500.

• Prestatiewinst: In vergelijking met het basismodel behaalde d-TreeRPO aanzienlijke verbeteringen:
  • Sudoku: +86,2% (LLaDA-8B)
  • Countdown: +51,6% (LLaDA-8B)
  • GSM8K: +4,5% (LLaDA-8B)
  • Math500: +5,3% (LLaDA-8B)
• Ablatiestudies:
  • Sparse Rewards: Varianten met sparse uitkomstrewards (het uitzenden van finale rewards naar alle tokens) presteerden aanzienlijk slechter, wat de noodzaak bevestigt van fijnmazige credit-toewijzing via boomstructuren.
  • Self-Distillation: Het verwijderen van de self-distillation-loss leidde tot een daling in prestaties en hogere fouten in kansschatting. Omgekeerd resulteerde een diversiteitsbevorderende loss (bedoeld om hoge entropie te behouden) in de slechtste prestaties en onstabiele training.
  • Planning: Een omgekeerde planning (sterke distillatie vroeg, zwak later) veroorzaakte prestatiedalingen na initiële winsten, wat het voorwaarts-geplande ontwerp valideert.
• Efficiëntie: Hoewel d-TreeRPO door boom-rollouts een hogere rekenkosten heeft dan sommige baselines (bijvoorbeeld Diffu-GRPO), convergeert het sneller in termen van reward en bereikt het superieure finale prestaties, waardoor het een praktische keuze is.

Betekenis en Claims

Het artikel stelt dat d-TreeRPO een betrouwbaarder fundament legt voor RL in dLLMs door de dubbele problemen van reward-sparsiteit en bias in kansschatting op te lossen. Door gebruik te maken van boomgestructureerde rollouts, zorgt de methode ervoor dat voordeelsignalen zowel verifieerbaar als granulier zijn. Bovendien navigeert het kader, door de noodzaak van zekerheid voor kansschatting theoretisch te onderbouwen en een tijdgepland self-distillation-mechanisme te implementeren, effectief door de voor dLLMs unieke afweging tussen exploratie en exploitatie. De resultaten suggereren dat betrouwbare beleidsoptimalisatie voor diffusiemodellen vereist dat men verder gaat dan eenvoudige uitkomstrewards en single-pass-benaderingen, en in plaats daarvan gestructureerd zoeken en dynamische zekerheidsregulatie omarmt.