OR Else: A Differentiable Trust Region for Policy Optimization

Technische Samenvatting: OR ELSE: EEN DIFFERENTIEERBARE VERTROUWENSREGIO VOOR BELEIDSOOPTIMALISATIE

Probleemstelling
Reinforcement Learning from Human Feedback (RLHF) voor Large Language Models (LLMs) vertrouwt doorgaans op Proximal Policy Optimization (PPO) en de variant daarvan, Group Relative Policy Optimization (GRPO). Beide methoden maken gebruik van een afgeknipte (clipped) surrogaat-objectief om beleidsupdates te beperken. Hoewel het clipping-mechanisme een eenzijdige nul-gradiënt regio creëert om excessieve beleidsdrift in de "gunstige" richting te voorkomen, is de overgang naar deze regio geometrisch abrupt: de afgeleide van de objectieffunctie met betrekking tot de waarschijnlijkheidsratio verandert discontinu bij de clipping-grens. De auteurs onderzoeken of Output Reset (OR), een gladde, eenzijdige verzadigingsregel die oorspronkelijk is ontworpen voor classifier-training, een superieur alternatief biedt voor LLM post-training. Specifiek stelt de paper de vraag of het vervangen van de afgeknipte surrogaat door een differentieerbare OR-gebaseerde loss de optimalisatiestabiliteit en de uiteindelijke beloningsscores kan verbeteren zonder nieuwe pathologieën te introduceren.

Methodologie
De paper introduceert PPO-OR en GRPO-OR, twee varianten die de standaard afgeknipte beleidsterm vervangen door een OR-afgeleide gekwadrateerde marge-loss in de ruimte van token log-ratio's.

• Kernmechanisme: In plaats van de waarschijnlijkheidsratio $r_t(\theta)$ te gebruiken, opereren de methoden op de per-token log-waarschijnlijkheidsratio $\rho_t = \log \pi_\theta(a_t|s_t) - \log \pi_{\theta_{old}}(a_t|s_t)$.
• De OR Loss: Voor een token met een niet-nul geschatte advantage $\hat{A}_t$, wordt een doel-marge $\sigma_t \alpha_{OR}$ ingesteld (waarbij $\sigma_t = \text{sign}(\hat{A}_t)$). De loss is een gekwadrateerde eenzijdige marge: $\ell_{OR,t} = [\max(0, \alpha_{OR} - \sigma_t \rho_t)]^2$.
• Output Reset Gedrag: Zodra $\rho_t$ de marge kruist in de richting die door het teken van de advantage wordt bevoordeeld, wordt het doel "gereset" naar de huidige $\rho_t$. Dit resulteert in een nul-residue en een nul directe gradiëntbijdrage voor dat token. De afgeleide nadert continu nul bij de grens, waardoor de loss $C^1$ continu is (in tegen tegenstelling tot de $C^0$ continue maar niet-differentieerbare PPO clip).
• Advantage Estimatie:
  • PPO-OR: Behoudt Generalized Advantage Estimation (GAE) met een geleerde value head.
  • GRPO-OR: Behoudt groep-relatieve advantages (berekend uit $G=2$ responsen per prompt) zonder een critic.
• Experimentele Opzet: De studie gebruikt Llama-3.2-1B-Instruct op de Anthropic hh-rlhf dataset. Alle methoden delen een enkele bevroren reward model en worden geëvalueerd over drie seeds. De experimenten zijn ontworpen als twee gematchte binnen-familie vergelijkingen: PPO-clip vs. PPO-OR (onder GAE) en GRPO vs. GRPO-OR (onder groep-relatieve advantages). Cross-familie vergelijkingen worden genoteerd als ongecontroleerd vanwege verschillen in advantage-estimatie, critic-gebruik en responslengtes.

Belangrijkste Bijdragen

1. Algoritmisch Voorstel: De introductie van PPO-OR en GRPO-OR, die de afgeknipte surrogaat vervangen door een continu differentieerbare, eenzijdige gekwadrateerde marge-loss in de log-ratio ruimte.
2. Theoretische Karakterisering: Formeel bewijs dat de OR-objectief $C^1$ continu is bij de verzadigingsgrens en dat de sample-niveau update-richting strikt wordt bepaald door het teken van de advantage, en continu verdwijnt bij het kruisen van de marge.
3. Empirische Evaluatie: Een rigoureuze vergelijking van clipping versus OR binnen twee verschillende advantage-estimatie families, waarbij wordt benadrukt dat de voordelen van OR niet uniform zijn maar afhangen van de omliggende optimalisatie-pipeline.
4. Diagnostisch Inzicht: De identificatie dat het vervangen van clipping door OR niet inherent het probleem van grote cumulatieve beleidsdrift oplost die wordt waargenomen in groep-relatieve settings, wat onderscheid maakt tussen lokale gradiëntverzadiging en globale beleidscontrole.

Resultaten
De paper rapporteert training-tijd reward-model scores en optimalisatiediagnostics over drie seeds:

• GAE Familie (PPO-clip vs. PPO-OR):

  • Beloningsscore: PPO-OR behaalde een hogere gemiddelde eindscore (+0.500) vergeleken met PPO-clip (+0.195), een verschil van +0.305.
    • Stabiliteit: PPO-OR vertoonde een grotere over de seeds verspreide standaarddeviatie (0.158) vergeleken met PPO-clip (0.110).
  • Drift: Beide methoden behielden vergelijkbare, lage rollout-naar-huidige log-ratio verplaatsingen (ongeveer 0.3–0.5).

• Groep-Relatieve Familie (GRPO vs. GRPO-OR):

  • Beloningsscore: GRPO-OR verbeterde de gemiddelde score niet; het was iets lager (+0.457) dan GRPO (+0.488).
  • Stabiliteit: GRPO-OR vertoonde een aanzienlijk kleinere over de seeds verspreide standaarddeviatie (0.027) vergeleken met GRPO (0.080).
  • Diagnostiek: GRPO-OR bereikte een bijna-nul terminaal OR-residue en een dalend overshoot-fractie, terwijl de gematchte GRPO-trace groot en zeer variabel bleef.
  • Drift: Cruciaal, GRPO-OR verminderde de grote rollout-naar-huidige verplaatsing waargenomen in GRPO niet. Beide groep-relatieve methoden vertoonden substantieel hogere verplaatsingen (ongeveer 5–13) vergeleken met de GAE methoden.

Betekenis en Claims
De paper concludeert bescheiden en benadrukt dat de geobserveerde effecten van OR context-afhankelijk zijn:

• Geen Universele Vertrouwensregio: De auteurs stellen expliciet dat de gladde verzadiging van de OR loss niet constitueert als een formele vertrouwensregio. Het onderdrukt de directe gradiëntbijdrage van individuele tokens nadat zij een marge hebben gekruist, maar beperkt de cumulatieve beleidsdrift veroorzaakt door gedeelde parameters, hulp-losses of veranderende advantage-tekens over iteraties heen.
• Pipeline Afhankelijkheid: De effectiviteit van OR is gekoppeld aan de advantage-estimatiemethode. In de GAE-setting correleerde OR met hogere beloningsscores maar verhoogde het de variantie. In de groep-relatieve setting verminderde OR de variantie en stabiliseerde het de loss-trace, maar verbeterde het de gemiddelde beloningsscore of de beleidsdrift niet.
• Beperkingen: De studie erkent dat de resultaten specifief zijn voor een lichtgewicht 1B model, een enkele bevroren reward model, en een beperkte groepsgrootte ($G=2$). De bevindingen karakteriseren deze specifieke experimentele setting eerder dan het vaststellen van algemene superioriteit over PPO of GRPO. De paper stelt dat de grote beleidsdrift in groep-relatieve methoden een onopgelost probleem blijft dat OR alleen niet kan oplossen, wat suggereert dat toekomstig werk de cumulatieve verplaatsing moet adresseren via sterkere referentie-regularisatie of expliciete verplaatsingsstraffen.

Samenvattend demonstreert de paper dat Output Reset een differentieerbaar alternatief biedt voor clipping dat de optimalisatiegedrag verandert en in specifieke contexten de beloningsscores kan verbeteren, maar dat het niet inherent de uitdagingen van beleidsstabiliteit of drift in groep-relatieve RLHF oplost.