Technische Samenvatting: Dropout-GRPO voor Continue Latente Redenering

1. Probleemstelling

Het artikel behandelt een fundamentele structurele incompatibiliteit tussen Group Relative Policy Optimization (GRPO) en continue latente redeneermodellen (specifiek de COCONUT-architectuur).

• De GRPO-vereiste: GRPO vertrouwt op de diversiteit van $K$ rollouts binnen een groep om een niet-nul groepsgemiddelde advantage te berekenen ($A^{(k)} = r^{(k)} - \mu_r$). Als alle rollouts identieke trajecten produceren, stort de advantage in tot nul, wat optimalisatie doet stagneren.
• De uitdaging van Latente Redenering: In tegenstelling tot standaard Chain-of-Thought (CoT) modellen die diversiteit afleiden uit stochastische token-sampling, opereren continue latente redeneermodellen (zoals COCONUT) via deterministische verborgen staat-recurrentie. Gegeven een prompt en vaste parameters, leveren meerdere rollouts identieke latente trajecten en antwoorden op.
• Het Gevolg: Het toepassen van standaard GRPO op deze modellen resulteert in nul variantie ($\sigma_r = 0$), waardoor de policy gradient verdwijnt en er geen leersignaal wordt gegenereerd.

2. Methodologie: Dropout-GRPO

De auteurs stellen Dropout-GRPO voor, een methode die noodzakelijke stochastiek injecteert in het latente redeneerproces via gestructureerde dropout, terwijl de theoretische garanties van variationele inferentie behouden blijven.

Kernmechanisme: Shared-Mask Variational Dropout

In plaats van standaard token-niveau sampling, introduceert de methode stochastiek op parameterniveau:

1. Mask Generatie: Voor elke rollout $k$ wordt een enkele Bernoulli-masker $\xi^{(k)}$ getrokken.
2. Consistentie: Dit masker wordt constant gehouden over alle $T$ latente recurrentie-stappen binnen die specifieke rollout.
3. Perturbatie: Het masker induceert een gestructureerde perturbatie $\tilde{\theta}(\xi)$ op de netwerkgewichten. De latente recurrentie wordt $h_t = f_{\tilde{\theta}(\xi)}(h_{t-1})$.
4. Mask Replay: Cruciaal is dat de RNG-seed voor het masker wordt opgeslagen tijdens de rollout-fase. Tijdens de policy update wordt hetzelfde masker geregenereerd en toegepast. Dit zorgt ervoor dat het traject dat geëvalueerd wordt voor de beloning bit-identiek is aan het traject dat gebruikt wordt voor de gradiëntberekening wanneer $\theta = \theta_{old}$.

Theoretische Interpretatie

• Bayesiaanse Model-gemiddelde: Volgens Gal en Ghahramani behandelt de gedeelde mask de elke rollout als een enkele posterior sample uit een gestructureerde distributie over parameters. De marginale policy $\bar{\pi}_\theta$ wordt geïnterpreteerd als een Bayesiaans model-gemiddelde.
• Gradiënt Onbevoordeeldheid: De auteurs bewijzen dat de surrogaatgradiënt met de gemiddelde-alleen advantage ($A^{(k)} = r^{(k)} - \mu_r$) een onbevoordeelde schatter is van de gradiënt van de verwachte beloning van de marginale policy, geschaald door $(1 - 1/K)$.
• Variantiereductie: Door hetzelfde masker te hergebruiken (Common Random Numbers) tijdens de update als dat gebruikt werd tijdens de rollout, vermindert de methode de variantie vergeleken met het opnieuw samplen van een nieuwe mask tijdens de update.

Implementatieverfijningen

Om training te stabiliseren, worden vier specifieke verfijningen geïntroduceerd:

1. Huberized $k_3$ KL: Vervangt de standaard clamping in de KL-divergentie estimator met een Huber-stijl extensie om te voorkomen dat gradiënten verdwijnen voor tokens met grote log-ratio's, waardoor "ontsnapte" tokens worden vermeden die zonder regularisatie afdrijven.
2. Trust-Region Annealing: De KL-coëfficiënt $\beta$ wordt proportioneel geanneald met de learning rate $\eta_t$ om te voorkomen dat de policy vast komt te zitten in een evenwicht waar de KL-trek de policy gradient wegtrekt.
3. Groepsniveau Nauwkeurigheidsfilter: Groepen met een empirische nauwkeurigheid $\mu_r$ buiten een geconfigureerd venster (bijv. $[0.20, 0.95]$) worden gedropt. Dit fungeert als een emergent curriculum, waarbij rekenkracht wordt gefocust op prompts met een niet-triviale variantie.
4. DDP-Coördineerde Skip: Zorgt voor gedistribueerde trainingsconsistentie bij het droppen van groepen.

3. Belangrijkste Bijdragen

1. Nieuw Paradigma voor Latente RL: Het artikel identificeert de deterministische aard van latente recurrentie als de oorzaak van GRPO-falen en stelt shared-mask variational dropout voor als een minimale, theoretisch zuivere oplossing. Het is de eerste methode die groep-relatieve RL succesvol toepast op continue latente redeneermodellen.
2. Theoretische Rechtvaardiging:
   • Bewijst de goed gedefinieerde aard van de latente Jacobiaan $\partial h_T / \partial \theta$ onder standaard Transformer-regulariteit.
   • Bewijst de onbevoordeeldheid van de surrogaatgradiënt voor de marginale policy.
   • Demonstreert dat mask replay zorgt voor variantiereductie via Common Random Numbers.
3. Empirische Validatie: Past GRPO succesvol toe op het COCONUT-model op de GSM8K benchmark, een regime waarin deterministische GRPO volledig faalt.
4. Referentie-implementatie: Publiceert code die bit-identieke mask replay realiseert over de gehele latente recurrentie.

4. Experimentele Resultaten

Experimenten werden uitgevoerd met Qwen2.5-1.5B als basismodel met een latente diepte van $T=6$.

• Prestatieverbetering: Dropout-GRPO verbeterde de GSM8K pass@1 nauwkeurigheid van een COCONUT-baseline van 27,29% naar 29,01%.
• Herstel van Capaciteiten: De methode herstelde een prestatie-regressie van 1,33 punt op de SVAMP-dataset veroorzaakt door curriculum training, waarbij de prestatie werd teruggebracht naar 44,00% (overeenkomend met de variantie van het basismodel).
• Ablatie Studies:
  • Deterministische GRPO: Faalde volledig, wat een vlak trainingsverloop opleverde dat identiek is aan de SFT-initialisatie door nul advantage.
  • REINFORCE + Dropout (EMA): Faalde door hoge variantie; de single-rollout gradiënt was te ruisig om de SFT-initialisatie te stabiliseren.
  • Dropout-GRPO: De enige configuratie die een levensvatbaar leersignaal produceerde, wat bevestigt dat groep-relatieve baselines essentieel zijn voor variantiereductie in deze setting.

5. Betekenis en Claims

Het artikel positioneert dit werk als een praktische en theoretisch onderbouwde aanpak voor post-training latente redeneerende Large Language Models (LLMs).

• Haalbaarheid: Het demonstreert dat RL post-training haalbaar is voor modellen die direct in verborgen toestanden redeneren, een capaciteit die voorheen geblokkeerd werd door het gebrek aan rollout-diversiteit.
• Mechanisme: Het herformuleert dropout niet als een regularisator (die vaak wordt uitgeschakeld in moderne LLM's), maar als een bron van gestructureerde exogene stochastiek die vereist is voor variationele policy optimalisatie.
• Beperkingen: De auteurs zijn bescheiden over de omvang van de verbetering (een winst van 1,72 punt op GSM8K), wat zij toeschrijven aan de kleine schaal van de backbone (1,5B) en de schaarste aan geschikte beknopte stap-redeneerdatasets. Ze merken op dat de ingezette policy (inferentie zonder dropout) verschilt van de geoptimaliseerde marginale policy, wat een bias introduceert die afhankelijk is van Lipschitz-constanten.

Samenvattend overbrugt Dropout-GRPO de kloof tussen groep-relatieve reinforcement learning en continue latente redenering, waardoor de optimalisatie mogelijk wordt gemaakt van modellen die voorheen niet in staat waren GRPO te gebruiken vanwege hun deterministische interne dynamiek.