Reformulate LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy

Technische Samenvatting: Herformulering van LLM Reinforcement Learning voor Efficiënte Training onder Black-box Discrepantie

1. Probleemstelling

Reinforcement Learning (RL) is een cruciale post-training paradigm geworden voor het verbeteren van de redeneercapaciteiten van Large Language Models (LLM's). RL post-training lijdt echter vaak aan onvoorspelbare suboptimale prestaties of volledige trainingsinstortingen. Recente onderzoeken schrijven deze mislukkingen toe aan een wijdverspreide train-inference discrepantie (of mismatch).

Deze discrepantie ontstaat omdat industriële eisen actoren er vaak toe dwingen om uiteenlopende onderliggende engines en precisieformaten (bijv. verschillende kernels, parallellisatiestrategieën of numerieke precisies zoals BF16 versus FP8) te gebruiken voor exact dezelfde modelgewichten tijdens training versus inference. Dit creëert een verborgen afwijking tussen de actie-waarschijnlijkheidsverdelingen van het trainingsbeleid ($\pi_{Train}$) en het inference-beleid ($\pi_{Inf}$).

Bestaande mitigatiestrategieën hebben aanzienlijke beperkingen:

• Uniforme Precisie-alignement: Het simpelweg verlagen van de trainingsprecisie om overeen te komen met de inference (bijv. naar FP16) verkleint het representeerbare waardebereik en verergert de risico's op gradiëntexplosie.
• Heuristische Interventies: Methoden zoals het maskeren van gradiënten van tokens met een hoge discrepantie onderdrukken slechts de impact van ruisige data. Ze pakken de kernoorzaak niet aan en negeren het uiteindelijke doel: het optimaliseren van het beleid voor de inference-implementatie, niet alleen voor het trainingsproces zelf.
• Infrastructuur-alignement: Het proberen af te stemmen van low-level operatoren en hardwarestacks is arbeidsintensief en vaak onhaalbaar vanwege het enorme aantal uitzonderingsgevallen.

De kernuitdaging is om het trainingsbeleid in staat te stellen de discrepantie met het inference-beleid dynamisch waar te nemen en actief te neutraliseren, zonder het primaire doel van beloningsmaximalisatie op te offeren.

2. Methodologie: Discrepancy-Constrained MDP (DCMDP)

De auteurs stellen een paradigmaverschuiving voor van standaard Markov Decision Processes (MDP) naar een Discrepancy-Constrained MDP (DCMDP). Dit framework behandelt de train-inference discrepantie als een black-box optimalisatieprobleem, waardoor het systeem naadloos kan omgaan met heterogene training en inference setups.

Kerncomponenten:

• Discrepantie Metriek:
  De auteurs evalueren empirisch drie token-niveau discrepantiesignalen: Probability Difference ($\delta_{diff}$), Probability Ratio ($\delta_{ratio}$) en Log-Probability Difference ($\delta_{log}$). Zij identificeren dat de absolute waarschijnlijkheidsverschil ($\delta_{diff} = |p_{Train} - p_{Inf}|$) het meest effectieve strafmechanisme is. Het biedt een "plastische" regularisatie die deployment-relevante drift corrigeert zonder de beloningsgestuurde exploratie te veel te beperken, in tegen tegenstelling tot ratio-gebaseerde straffen die te conservatief kunnen zijn voor tokens met een lage waarschijnlijkheid.

• De Discrepantie Tolerantiezone:
  Een cruciale empirische bevinding is het bestaan van een discrepantie tolerantiezone. Het agressief bestraffen van kleine, verwaarloosbare discrepanties in een vroeg stadium van de training onderdrukt de beleidsexploratie en vermindert de leer-efficiëntie. Omgekeerd verbetert het verminderen van excessieve discrepanties buiten deze zone de optimalisatieconsistentie.

  • Mechanisme: De DCMDP-formulering past alleen straffen toe wanneer de discrepantie een vooraf gedefinieerde drempel $c$ overschrijdt. Binnen deze tolerantiezone mag het beleid vrij verkennen om beloningen te maximaliseren.

• Dual-Objective Optimalisatie via Lagrangian Relaxatie:
  Het probleem wordt geformuleerd als het maximaliseren van de verwachte beloning onder een beperking op de gemiddelde token-niveau discrepantie:
  $$\max_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum r(o_t) \right]$$
  $$\text{onder voorbehoud van: } \frac{1}{|\tau|} \sum |\pi_{Train}(o_t) - \pi_{Inf}(o_t)| \leq c$$

Om dit op te lossen, maken de auteurs gebruik van Lagrangian Relaxatie. Een leerbare multiplier $\lambda$ past dynamisch het gewicht van de discrepantie-straf aan:
* Als de discrepantie de tolerantiegrens overschrijdt ($>c$), neemt $\lambda$ toe, waardoor de straf wordt versterkt om het beleid terug naar de veilige zone te trekken.
* Als de discrepantie binnen de grens valt, vlakt $\lambda$ af naar nul, waardoor het beleid zich kan concentreren op beloningsmaximalisatie.

Het uiteindelijke doel (Dual-GRPO) injecteert deze adaptieve straf direct in de groep-genormaliseerde advantage:
$$\hat{A}^{DCMDP}_{i,t} = \hat{A}_{i,t} - \lambda_k \cdot \delta_{i,t}$$

3. Belangrijkste Bijdragen

1. Autonome Discrepantie-compensatie: De paper demonstreert dat het integreren van een robuuste, black-box straf (token-niveau absolute waarschijnlijkheidsverschil) het trainingsbeleid de capaciteit geeft om autonoom het risico op discrepantie-geïnduceerde trainingsinstorting te elimineren.
2. Paradigmaverschuiving naar DCMDP: De auteurs herformuleren LLM RL als een Discrepancy-Constrained MDP. Deze aanpak balanceert optimaal de duale doelstellingen van het maximaliseren van de taakbeloning en het minimaliseren van de train-inference discrepantie, waarmee zij verder gaan dan naïeve scalarisatie van doelstellingen.
3. Heterogeen Trainingsparadigma: De methode maakt een trainingssetup mogelijk waarbij modellen worden geoptimaliseerd in hoog-getrouwe, resource-intensieve omgevingen (bijv. BF16 precisie), terwijl ze expliciet worden uitgelijnd voor laag-kosten, resource-beperkte inference deployments (bijv. FP8). Dit transformeert infrastructuur-heterogeniteit van een last naar een controleerbare optimalisatieknop.

4. Experimentele Resultaten

De methode werd geëvalueerd op Qwen3-8B (dense) en Qwen3-30B-A3B (Mixture-of-Experts) modellen over zes wiskundige redeneerbenchmarks (AIME24/25, AMC23, MATH-500, Minerva, OlympiadBench).

• BF16 Training/Inference: DC-GRPO presteerde consistent beter dan de vanilla GRPO baseline over alle benchmarks voor beide modellen, waarbij hogere gemiddelde scores werden behaald (bijv. Qwen3-8B verbeterde van 67.3 naar 71.0). Dit bevestigt dat het afdwingen van discrepantie-constraints de primaire beloningsoptimalisatie versterkt in plaats van hindert.
• Heterogene Training (BF16) + Deployment (FP8): In deze setting leed vanilla GRPO aan significante prestatievermindering of instorting door de versterkte train-inference kloof. DC-GRPO mitigeerde dit succesvol, waardoor de prestaties werden behouden en vaak zelfs de prestaties van de BF16-only baseline overtroffen.
• Ablatie-studies: De methode bleek robuust tegen variaties in de initiële Lagrangian multiplier ($\lambda_0$) en de duale leersnelheid ($\eta_\lambda$). De tolerantiebudget ($c$) bleek kritiek; het te klein instellen veroorzaakte overmatige regularisatie, terwijl het te groot instellen ervoor zorgde dat de methode degradeerde naar vanilla GRPO.

5. Betekenis en Claims

De paper claimt dat DCMDP een verenigd, algoritmisch framework biedt voor het mitigeren van train-inference instabiliteit zonder dat daarvoor uitputtende, op maat gemaakte optimalisaties van de onderliggende hardware- of softwarestacks nodig zijn. Door de discrepantie te behandelen als een black-box feedbacksignaal, maakt de methode het volgende mogelijk:

• Stabiele Training: Het voorkomen van de abrupte instortingen die gebruikelijk zijn bij RL post-training.
• Efficiënte Exploratie: Het beleid de vrijheid geven om te verkennen binnen een veilige discrepantiezone.
• Resource-efficiëntie: Het mogelijk maken van het gebruik van hoog-getrouwe trainingssettings om te simuleren en te adapteren aan beperkte deployment-omgevingen, wat de meer efficiënte real-world model-deployments faciliteert.

De auteurs erkennen beperkingen en merken op dat hun empirische scope momenteel beperkt is tot specifieke BF16 en FP8 precisies en specifieke modelfamilies, waarbij toekomstig werk nodig is om de methode uit te breiden naar andere backend-combinaties en modelarchitecturen.