Technische Samenvatting: Mask-Aware Policy Gradients voor Diffusion Language Models

Probleemstelling

Reinforcement learning (RL) heeft de redeneervaardigheden van autoregressieve Large Language Models (LLM's) succesvol verbeterd door policy gradients te optimaliseren op basis van tractabele log-likelihoods. Het uitbreiden van deze methoden naar Masked Diffusion Language Models (MDLMs) blijft echter uitdagend. In tegenstelling tot autoregressieve modellen, die tekst sequentieel van links naar rechts genereren, genereren MDLMs tekst door iteratief tokens te ontmaskeren vanuit een volledig gemaskeerde sequentie.

De kern van de moeilijkheid ligt in de ontractabele schatting van de log-likelihood voor MDLMs. Bestaande benaderingen benaderen deze likelihood met behulp van Evidence Lower Bound (ELBO)-varianten of door alleen het component van de tokenvoorspelling van het generatiepad te modelleren. Deze methoden negeren een cruciaal aspect van MDLM-generatie: de volgorde waarin posities worden ontmaskerd. In MDLMs bestaat het generatiepad uit twee afzonderlijke beslissingen bij elke stap: (1) het voorspellen van welke tokens op de gemaskeerde posities moeten worden geplaatst, en (2) het beslissen welke van die posities ongemaskeerd blijven (en welke opnieuw gemaskeerd worden). Vorige methoden die de tweede beslissing negeren, maken geen gebruik van het volledige signaal dat beschikbaar is in het generatieproces voor policy-optimalisatie.

Methodologie

De auteurs formaliseren MDLM-generatie als een tweestaps actie Markov-beslissingsproces (MDP). Bij elke denoising-stap $t$ voert het model uit:

1. Tokenvoorspelling: Het voorspellen van tokens voor alle momenteel gemaskeerde posities.
2. Positiekeuze (Remaskering): Het beslissen welke subset van voorspelde posities wordt onthuld en welke terugkeert naar de [MASK] staat.

Om policy gradient-optimalisatie over dit proces mogelijk te maken, introduceren de auteurs Probabilistische Remaskering. Standaard MDLM-inferentie gebruikt vaak een deterministische, greedige top-$K$-strategie om posities te selecteren om te unmasken op basis van betrouwbaarheidsscores. Deze greedische selectie is niet-differentieerbaar, wat gradiëntgebaseerde optimalisatie van de unmasking-volgorde verhindert. De voorgestelde methode vervangt dit door een probabilistische variant:

• In plaats van deterministisch de top-$K$ posities te selecteren, samplet het model een subset van posities $U_t$ uit een Plackett–Luce distributie.
• De kansen voor deze distributie worden direct afgeleid van de eigen token log-likelihoods (logits) van het model, geschaald door een temperatuurparameter $\tau$.
• Deze formulering creëert een differentieerbare distributie over subsets van posities, waardoor de log-waarschijnlijkheid van de unmasking-beslissing kan worden opgenomen in de traject-likelihood.

Policy Gradient Decompositie

Door de volledige traject-likelihood $\pi_\theta(\hat{z}|c)$ te modelleren als een product van tokenvoorspellingskansen en unmasking-kansen, deelt de policy gradient zich natuurlijk op in twee termen:
$$\nabla_\theta J(\theta) = \mathbb{E}_{\hat{z}} \left[ R(c, z_T) \sum_{t=1}^T \left( \underbrace{\nabla_\theta \log \pi_\theta(\hat{z}_t | c, z_{t-1})}_{\text{Token Gradiënt}} + \underbrace{\nabla_\theta \log p_{\text{unmask}}(U_t | \hat{z}_t, z_{t-1}, c)}_{\text{Masking Gradiënt}} \right) \right]$$

• Token Term: Versterkt de correctheid van de voorspelde tokens.
• Masking Term: Versterkt de selectie van de unmasking-volgorde.
• Efficiëntie: De masking term wordt berekend vanuit de bestaande logits van het model die worden gebruikt voor de tokenvoorspelling, waardoor geen extra parameters, architecturale wijzigingen of extra forward passes nodig zijn.

De auteurs implementeren dit met het Group Sequence Policy Optimization (GSPO) algoritme, waarbij de importance ratios voor zowel de token- als de positietermen afzonderlijk worden afgeklemd (clipping). Om de computationele kosten van het evalueren van volledige trajecten te beheersen, gebruiken ze de StepMerge benadering, waarbij denoising-stappen worden gegroepeerd in segmenten om het aantal forward passes te verminderen terwijl de trajectstructuur behouden blijft.

Belangrijkste Bijdragen

1. Formalisering van MDLM als een tweestaps MDP: Het artikel biedt een theoretisch kader dat aantoont dat MDLM-generatie zowel tokenselectie als positiekeuze omvat, en dat de policy gradient dienovereenkomstig uiteenvalt.
2. Probabilistische Remaskering: De afleiding van een differentieerbare unmasking-distributie (Plackett–Luce) gebaseerd op de eigen logits van het model, wat de optimalisatie van de unmasking-volgorde mogelijk maakt zonder architecturale wijzigingen.
3. Theoretisch Inzicht: Demonstratie dat het negeren van de positie-component van de policy gradient kan leiden tot het missen van richtingen die de verwachte return verbeteren, zelfs wanneer de token-waarschijnlijkheden ongewijzigd blijven.
4. State-of-the-Art Prestaties: Empirische validatie die aantoont dat het gezamenlijk optimaliseren van beide termen consistente verbeteringen oplevert ten opzichte van bestaande traject-gebaseerde en ELBO-gebaseerde methoden.

Experimentele Resultaten

De methode werd geëvalueerd op benchmarks voor wiskundig redeneren (GSM8K, MATH500) en code-generatie (HumanEval, MBPP), waarbij LLaDA-8B-Instruct als basismodel werd gebruikt.

• Prestatiewinst: De voorgestelde methode behaalde nieuwe state-of-the-art resultaten over alle benchmarks.
  • GSM8K: 87,1% nauwkeurigheid (bij generatielengte 128), een verbetering van +2,5% ten opzsten de vorige beste (StepMerge).
  • MATH500: 53,4% nauwkeurigheid, een verbetering van +4,0%.
  • HumanEval: 47,1% nauwkeurigheid, een verbetering van +2,2%.
  • MBPP: 53,4% nauwkeurigheid, een verbetering van +2,5%.
• Ablatie-studies:
  • Blokgrootte: Het prestatieverschil tussen de voorgestelde methode en de baselines werd groter naarmate de inference blokgrootte toenam (van 32 naar de volledige sequentie), wat suggereert dat het voordeel van het optimaliseren van de positiekeuze groeit met de complexiteit van de unmasking-beslissingsruimte.
  • Efficiëntie: De methode convergeerde sneller in wall-clock tijd vergeleken met de sterke ELBO-gebaseerde baseline (SPG), waarbij de uiteindelijke nauwkeurigheid van SPG werd bereikt in ~15 uur versus ~18 uur, en een hogere uiteindelijke nauwkeurigheid behaalde (81,0% vs 78,5% op GSM8K).
  • Generalisatie: De aanpak generaliseerde naar een ander basismodel (Dream-7B) en plannings-taken (Sudoku, Countdown), waarbij het consequent de baselines overtrof.
• Vergelijking met DCoLT: Wanneer vergeleken met LLaDOU (dat een aparte head gebruikt voor positiekeuze), behaalde de voorgestelde methode een vergelijkbare of betere nauwkeurigheid met aanzienlijk lagere trainingskosten (~160 GPU-uren versus ~800), omdat het geen extra modules vereist.

Betekenis en Claims

Het artikel claimt dat de flexibiliteit van MDLMs gepaard gaat met een rijkere trajectstructuur dan autoregressieve modellen, specifelijk de beslissing van wanneer posities te unmasken. Door deze beslissing te behandelen als onderdeel van de policy en deze te optimaliseren via een differentieerbaar probabilistisch framework, demonstreren de auteurs dat de gezamenlijke optimalisatie van tokenvoorspellingen en positiekeuze cruciaal is voor het maximaliseren van de prestaties van diffusion language models in RL-omgevingen.

De auteurs benadrukken dat hun aanpak praktisch efficiënt is, aangezien het gebruik maakt van bestaande model-logits zonder dat er nieuwe parameters of architecturale wijzigingen nodig zijn, wat het compatibel maakt met standaard MDLM-trainingsinfrastructuren zoals StepMerge. De resultaten suggereren dat de "masking gradient" een betekenisvol trainingssignaal levert dat eerdere methoden, die zich uitsluitend op token-likelihoods richtten, niet konden vangen.