Stabilizing Reinforcement Learning for Diffusion Language Models

Each language version is independently generated for its own context, not a direct translation.

Stel je voor dat je een zeer slimme, creatieve robot (een Diffusion Large Language Model of dLLM) traint om moeilijke wiskundepuzzels op te lossen. Je wilt dat deze robot beter wordt door te oefenen met een systeem van beloningen: als hij een goed antwoord geeft, krijgt hij een puntje; als hij fouten maakt, krijgt hij een straf.

In de wereld van kunstmatige intelligentie heet dit Versterkend Leren (Reinforcement Learning). Voor de traditionele robots (die zogenoemde "autoregressieve" modellen zijn) werkt een bepaalde trainingsmethode, genaamd GRPO, fantastisch. Het is als een strenge maar eerlijke coach die de robot helpt om stap voor stap te verbeteren.

Maar hier zit de twist: als je dezezelfde coach (GRPO) probeert te gebruiken op de nieuwe, moderne robots (de Diffusion-modellen), gaat het volledig mis. De robot begint te trillen, raakt in paniek en stopt met leren. Dit fenomeen noemen de auteurs van dit paper een "beloningsinstorting" (reward collapse).

Waarom gaat het mis? (De "Gestoorde Scorekaart")

Het probleem zit hem in hoe de coach de prestaties meet.

Bij de oude robots: De coach kan precies zien hoeveel waarschijnlijker het goede antwoord is dan het oude antwoord. Het is als een perfecte scorekaart.
Bij de nieuwe robots: De coach kan die exacte score niet berekenen. Hij moet het schatten. En dat schatten is erg onnauwkeurig en ruisig, alsof je probeert de temperatuur te meten met een gebroken thermometer.

Deze "ruis" (fouten in de schatting) zorgt voor twee grote problemen:

Het "Paniek-Effect" (Gradient Spikes): Soms denkt de coach door de ruis dat de robot een enorme verbetering heeft gemaakt, terwijl hij eigenlijk niets gedaan heeft. Omdat de coach dit gelooft, geeft hij de robot een gigantische duw in de goede richting. Maar omdat het een leugen was, duwt hij de robot de verkeerde kant op. De robot raakt in paniek en maakt nog grotere fouten.
De "Zelfversterkende Lussen": De robot maakt nu grotere fouten, waardoor de coach de volgende keer nog slechtere schattingen maakt. Die slechte schattingen leiden weer tot nog grotere duwen. Het is een vicieuze cirkel: Fouten leiden tot paniek, paniek leidt tot meer fouten.

De Oplossing: StableDRL (De "Stabiele Coach")

De auteurs van dit paper hebben een nieuwe methode bedacht, genaamd StableDRL. Ze hebben de coach (GRPO) op twee cruciale manieren aangepast om deze paniek te stoppen:

1. De "Onvoorwaardelijke Rem" (Unconditional Clipping)

In het oude systeem mocht de coach alleen remmen als de robot al te ver weg was. Maar omdat de schattingen zo ruisig waren, kon de coach soms "vergeten" te remmen als de robot een enorme, nep-duw kreeg.
De oplossing: De nieuwe coach heeft een harde rem die altijd werkt, ongeacht wat er gebeurt. Als de schatting van de verbetering te groot wordt (bijvoorbeeld: "Hij is 1000% beter geworden!"), zegt de coach: "Nee, dat is onmogelijk door ruis. We gaan uit van maximaal 10%." Dit voorkomt die gigantische, destabiliserende duwen.

2. De "Eigen Groep-Norm" (Self-Normalization)

Stel je voor dat de coach een groep van 10 robots tegelijk bekijkt. In het oude systeem deelde hij de beloning door het vaste aantal robots (10). Maar als de schattingen van die 10 robots heel wisselend zijn (sommige zeggen "super!", andere "slecht!"), wordt de gemiddelde duw chaotisch.
De oplossing: De nieuwe coach kijkt niet naar het vaste getal 10, maar naar de totale kracht van de groep. Hij past de duw aan op basis van wat de groep echt heeft gedaan. Hierdoor blijft de duw altijd binnen een veilige, voorspelbare zone. Het is alsof je de kracht van een team niet deelt door het aantal mensen, maar door de totale energie die ze samen hebben.

Het Resultaat: Een Rustige Leraar

Met deze twee aanpassingen (de harde rem en de slimme groepsberekening) kunnen de nieuwe robots nu eindelijk stabiel leren.

Ze trillen niet meer.
Ze raken niet in paniek.
Ze kunnen duizenden trainingsstappen maken zonder te crashen.

In de praktijk betekent dit dat deze robots nu veel beter worden in logisch redeneren, wiskunde en puzzels oplossen dan ooit tevoren. Ze halen zelfs resultaten die beter zijn dan de beste traditionele robots.

Samenvattend in één zin:

De auteurs hebben een nieuwe trainingsmethode bedacht die de "ruis" in de schattingen van moderne AI-robots filtert, waardoor ze niet meer in paniek raken en eindelijk hun volle potentieel kunnen bereiken.

Each language version is independently generated for its own context, not a direct translation.

Titel: Stabilisatie van Reinforcement Learning voor Diffusie Taalmodellen

Auteurs: Jianyuan Zhong, Kaibo Wang, Ding Ding, et al. (Huawei, CUHK, HKUST)

1. Het Probleem: Instabiliteit bij GRPO voor Diffusiemodellen

Hoewel Group Relative Policy Optimization (GRPO) zeer effectief is voor het post-trainen van autoregressieve (AR) taalmodellen, faalt de directe toepassing ervan op Discrete Diffusion Large Language Models (dLLMs) vaak door ernstige instabiliteit, wat leidt tot een abrupte ineenstorting van de beloning (reward collapse).

De auteurs identificeren twee fundamentele oorzaken voor deze incompatibiliteit:

Onberekenbare Belangverhoudingen (Importance Ratios): In AR-modellen zijn de waarschijnlijkheidsverhoudingen tussen strategieën exact berekenbaar. In dLLMs is de log-waarschijnlijkheid echter onberekenbaar (intractable) en moet deze worden geschat via schatters zoals de Evidence Lower Bound (ELBO) of mean-field benaderingen. Deze schattingen zijn inherent ruisig en vertonen een "long-tail" verdeling met extreme uitbijters.
Ontwerpgebreken in standaard GRPO:
- Conditionele Clipping: Standaard GRPO gebruikt conditionele clipping (alleen clippen bij bepaalde voordelen). In dLLMs kan ruis in de geschatte verhoudingen ervoor zorgen dat deze conditionele drempel onbedoeld wordt omzeild, wat leidt tot gradiëntspikes (plotselinge, enorme updates).
- Vaste Normalisatie: GRPO normaliseert updates op basis van een vaste groepsgrootte. Bij hoge variatie in de geschatte verhoudingen (door ruis) veroorzaakt dit grote schommelingen in de grootte van de gradiënten.

Deze factoren creëren een zelfversterkende instabiliteitslus: Ruis veroorzaakt gradiëntspikes $\rightarrow$ dit duwt het beleid (policy) weg van de oorspronkelijke verdeling (policy drift) $\rightarrow$ deze drift verergert de variatie in toekomstige verhoudingsschattingen $\rightarrow$ wat leidt tot nog grotere spikes en uiteindelijk tot falen.

2. Methodologie: StableDRL

Om deze instabiliteitslus te doorbreken, stellen de auteurs StableDRL voor, een herschreven versie van GRPO die specifiek is ontworpen voor dLLMs. De methode bestaat uit twee kerncomponenten:

A. Unconditional Clipping (Onvoorwaardelijke Clipping)

In plaats van de conditionele clipping van GRPO, legt StableDRL een strikte, onvoorwaardelijke grens op aan de geschatte belangverhoudingen ( $\hat{\rho}$ ).

De verhouding wordt altijd begrensd binnen het interval $[1-\epsilon, 1+\epsilon]$ , ongeacht het teken van het voordeel (advantage).
Doel: Dit voorkomt dat uitbijters (outliers) door de ruis onbeperkte gradiëntspikes genereren, zelfs wanneer het voordeel negatief is.

B. Self-Normalization (Zelf-normalisatie)

Standaard GRPO deelt door de vaste groepsgrootte ( $G$ ). StableDRL vervangt dit door het delen door de som van de geknipte verhoudingen ( $\sum \text{clip}(\hat{\rho}_i)$ ).

Doel: Dit beperkt de update binnen de convexe hull (het convexe omhulsel) van de per-staal gradiënten. Hierdoor wordt de updategrootte niet langer beïnvloed door de fluctuaties in de totale groepsgrootte van de gewichten, wat de variatie in de updategrootte drastisch vermindert.

C. Uitbreiding naar Block Diffusion: Trappenhuis-Attention

Voor block-wise diffusiemodellen (waarbij tekst in blokken wordt gegenereerd) introduceert de auteurs een Staircase Attention-mechanisme.

Dit maakt het mogelijk om de ELBO-schatting in één doorloop ( $O(1)$ ) te berekenen zonder dat er informatielekken optreden (waarbij tokens "cheaten" door naar hun eigen grondtruth te kijken).
Het gebruikt een dubbele stream (schone context + korrupte target) met een specifiek masker dat strikte conditionele onafhankelijkheid garandeert.

3. Belangrijkste Bijdragen

Theoretische Diagnose: De auteurs hebben de mechanismen achter de reward collapse in dLLMs wiskundig en empirisch in kaart gebracht, specifiek de rol van ruis in belangverhoudingen en de daaruit voortvloeiende zelfversterkende lus.
StableDRL Framework: Introductie van een nieuw RL-framework dat onvoorwaardelijke clipping en zelf-normalisatie combineert om stabiel full-parameter training mogelijk te maken.
Architecturale Generalisatie: De methode werkt zowel voor full-attention diffusiemodellen (zoals LLaDA) als voor block-wise diffusiemodellen (zoals SDAR), dankzij de Staircase Attention.
State-of-the-Art Resultaten: Het bereiken van nieuwe state-of-the-art prestaties in redeneertaken voor dLLMs, wat aantoont dat stabilisatie essentieel is om het redeneervermogen van deze modellen volledig te ontsluiten.

4. Resultaten

De auteurs testen StableDRL op twee modellen: LLaDA-8B (full-attention) en SDAR-8B (block diffusion).

Stabiliteit: Waar eerdere methoden (zoals ESPO en SPG) vaak faalden na ~300 stappen door reward collapse, slaagt StableDRL erin om >1000 stappen stabiel te trainen zonder instorting.
Prestaties:
- Op GSM8K en MATH500 (wiskundige redenering) behaalt StableDRL de hoogste gemiddelde nauwkeurigheid, met name een verbetering van +5.2% op MATH500 ten opzichte van de beste concurrent (SPG).
- Op Countdown (planning) en Sudoku toont het superieure consistentie over verschillende generatielengtes.
- Op de strenge AIME'24 benchmark (wiskunde) overtreft StableDRL zelfs sterke autoregressieve baselines (Qwen3-8B) met een score van 16.7% vs 10.0%.
Robuustheidstest: In een "Exploding Weight Stress Test" (waarbij ruis kunstmatig wordt opgeblazen) blijft StableDRL stabiel, terwijl concurrenten direct instorten. Dit bevestigt dat de methode effectief is in het onderdrukken van ruisgevoeligheid.

5. Betekenis en Impact

Dit werk is van cruciaal belang voor de evolutie van Large Language Models (LLMs):

Paradigmaverschuiving: Het toont aan dat diffusiemodellen een levensvatbaar alternatief zijn voor autoregressieve modellen, mits de trainingsstabiliteit wordt opgelost. Diffusiemodellen bieden voordelen zoals parallelle decoding en bidirectionele context.
Oplossing voor een fundamenteel probleem: Het biedt een algemene oplossing voor het probleem van ruis in belangverhoudingen bij on-line RL, wat eerder een grote barrière was voor het toepassen van GRPO op niet-autoregressieve modellen.
Schalbaarheid: Door het mogelijk maken van stabiel full-parameter fine-tuning, kunnen dLLMs hun volledige redeneercapaciteit benutten, wat leidt tot significante prestatieverbeteringen in complexe taken zoals wiskunde en logisch redeneren.

Samenvattend introduceert StableDRL de eerste methode die stabiel, full-parameter reinforcement learning mogelijk maakt voor zowel full-attention als block-wise diffusion taalmodellen, waardoor een nieuwe staat van de kunst wordt bereikt in redeneertaken.