Stabilizing Policy Optimization via Logits Convexity

Dit artikel introduceert Logits Convex Optimization (LCO), een nieuw kader voor beleidsoptimalisatie dat de stabiliteit van het trainen van grote taalmodellen verbetert door de convexiteit van de SFT-verliesfunctie te benutten, wat ontbreekt in traditionele methoden zoals PPO.

De kern

Stel je voor dat je een zeer getalenteerde, maar nogal onrustige kunstenaar hebt: een Grote Taalmodel (LLM). Deze kunstenaar kan prachtige verhalen schrijven, wiskundige problemen oplossen en vragen beantwoorden. Maar om hem echt goed te laten presteren op specifieke taken, moet je hem "trainen".

Er zijn twee manieren om dit te doen:

1. SFT (Supervised Fine-Tuning): Je geeft de kunstenaar een boek met de perfecte antwoorden en zegt: "Kijk, dit is hoe het moet." Dit gaat heel rustig en stabiel.
2. RL (Reinforcement Learning): Je geeft de kunstenaar een opdracht en zegt: "Probeer het zelf, en als het goed is krijg je een sterretje, als het fout is een rode vlag." Dit is de methode die momenteel populair is voor complexe taken, maar het is berucht om zijn instabiliteit. Soms wordt de kunstenaar zo enthousiast of zo verward door de feedback dat hij helemaal de controle verliest en stopt met leren.

De auteurs van dit paper hebben gekeken naar waarom die tweede methode (RL) zo vaak uit de hand loopt, terwijl de eerste (SFT) zo rustig verloopt.

Het Geheim: De "Logits" en de Bolle Weg

Stel je voor dat het leren van de kunstenaar een wandeling is over een landschap.

• Bij SFT is dit landschap als een perfect ronde kom. Als je een balletje in de kom legt, rolt het vanzelf rustig naar het diepste punt (het beste antwoord). De weg is glad en voorspelbaar. In de wiskundige taal van de paper noemen ze dit convexiteit op het niveau van de "logits" (de ruwe, onbewerkte voorkeuren van het model). Omdat de weg zo glad is, weet het model precies welke kant op te gaan.
• Bij RL (met een populaire methode genaamd PPO) is het landschap echter als een ruwe, hobbelige berg met valkuilen. Soms denkt het model dat het naar beneden gaat, maar dan stuitert het plotseling omhoog. Dit komt door de manier waarop de feedback wordt verwerkt. De "weg" is niet convex; hij is vol gaten en scherpe randen. Hierdoor kan het model plotseling enorme, chaotische sprongen maken in plaats van kleine, gestage stapjes.

De Oplossing: LCO (Logits Convex Optimization)

De auteurs hebben een nieuwe methode bedacht, genaamd LCO. Ze zeggen eigenlijk: "Waarom proberen we de kunstenaar te laten klimmen over die hobbelige berg, als we hem gewoon een kaart kunnen geven van de perfecte bestemming?"

In plaats van te wachten op de chaotische feedback van RL, gebruiken LCO de wiskundige inzichten om het model direct te sturen naar het ideale doel. Ze maken de "hobbelige berg" om tot een gladde kom, net als bij SFT.

Hoe werkt het in de praktijk?
Stel je voor dat je een GPS hebt die niet zegt: "Ga links, oh wacht, rechts, oh nee, links!" (zoals PPO doet). In plaats daarvan zegt LCO: "Het perfecte punt ligt precies daar, en de weg erheen is een rechte, gladde lijn."

• Ze gebruiken een nieuwe formule die zorgt dat elke stap die het model zet, logisch en veilig is.
• Hierdoor worden de "schokken" en "explosies" in de training verdwenen. Het model leert sneller en stabieler.

Wat is het resultaat?

De paper toont aan dat deze nieuwe methode (LCO) beter werkt dan de oude, onstabiele methoden:

• Minder crashen: De training stopt niet halverwege omdat het model "op hol slaat".
• Beter resultaat: Of het nu gaat om wiskundige problemen oplossen of het beantwoorden van vragen, de modellen die met LCO zijn getraind, presteren vaak beter dan die met de traditionele methoden.
• Sneller: Omdat de weg glad is, komt het model sneller bij het doel.

Samenvattend

Dit paper lost een groot probleem op in de wereld van AI. Het zegt: "We hebben te lang geprobeerd om een onstabiel systeem (RL) te fixeren met plakband en lijm (zoals 'clipping' en 'strafregels'). In plaats daarvan hebben we de basis van het systeem veranderd zodat het van nature stabiel is, door de 'weg' die het model moet lopen, glad te maken."

Het is alsof je van een ruwe, gevaarlijke bergpad overstapt naar een snelweg: je komt sneller, veiliger en met minder stress op je bestemming aan.

Technische samenvatting

Probleemstelling: Instabiliteit in RL voor LLM's

Hoewel Reinforcement Learning (RL) een hoeksteen is geworden voor het uitlijnen van Large Language Models (LLM's) met menselijke voorkeuren en het verbeteren van redeneercapaciteiten, staat de training van RL bekend om zijn inherente instabiliteit. In tegenstelling tot Supervised Fine-Tuning (SFT), dat een stabiel trainingspad volgt, lijden RL-algoritmen zoals Proximal Policy Optimization (PPO) vaak aan:

• Volatile gradiënten: PPO gebruikt een "geclipte surrogate objective" die, ondanks stabiliserende technieken zoals clipping en KL-straf, nog steeds kan leiden tot explosieve gradiënten.
• Training collapse: Deze gradiëntfluctuaties kunnen leiden tot te grote parameterupdates, wat de training irreversibel kan laten instorten.
• Gebrek aan theoretisch inzicht: De onderliggende oorzaken van dit stabiliteitsverschil tussen SFT en RL zijn nog niet volledig begrepen.

De auteurs identificeren dat het ontbreken van een specifieke eigenschap in het optimalisatielandschap van RL de boosdoener is, terwijl SFT deze wel bezit.

Methodologie: Logits Convexity en LCO

De kern van het paper is een theoretische analyse van de geometrie van het optimalisatielandschap, specifiek gericht op de convexiteit van logits (de ruwe uitkomsten van het model voorafgaand aan de softmax-functie).

1. Theoretische Analyse:

   • De auteurs tonen aan dat de SFT-loss functie logits-convex is. Dit betekent dat de Hessiaan-matrix ten opzichte van de logits positief semi-definiet is.
   • Deze convexiteit garandeert dat de gradiënt in de parameter-ruimte richting de optimale parameters wijst (gunstige gradiënt-directionaliteit), waardoor het model niet vastloopt in valse stationaire punten.
   • In tegenstelling hieraan is de PPO-loss niet logits-convex. Dit leidt tot een gebrek aan directionele consistentie, wat de oorzaak is van de waargenomen gradiëntspikes en instabiliteit, vooral bij acties met negatieve voordelen (advantages).

2. De Oplossing: Logits Convex Optimization (LCO):
   Gebaseerd op deze inzichten stellen de auteurs LCO voor, een nieuw optimalisatiekader dat het RL-probleem herschrijft als een supervisie-alignment-probleem naar een optimale doelstelling.

   • Doelstelling: In plaats van direct de surrogate objective van PPO te maximaliseren, traint LCO het beleid om te convergeren naar een optimale doellogits ($z^*$) of een optimale verdeling ($\pi^*$), afgeleid van de originele RL-doelstelling (via de advantage-functie).
   • Implementaties: De auteurs introduceren drie varianten van LCO die allemaal de eigenschap van logits-convexity behouden:
     • LCO-MSE: Minimaliseert de Mean Squared Error tussen de huidige logits en de optimale doel-logits.
     • LCO-LCH: Gebruikt een log-cosh-verliesfunctie, wat robuuster is tegen uitbijters (outliers) in de logits-ruimte.
     • LCO-KLD: Minimaliseert de forward KL-divergentie tussen de huidige beleidsverdeling en de optimale verdeling.
   • Voordelen: Door convexiteit te behouden, garanderen deze methoden dat de gradiëntnormen monotoon afnemen naarmate het model convergeert, wat plotselinge spikes voorkomt.

Belangrijkste Bijdragen

• Theoretisch Inzicht: Het identificeren van "logits convexity" als de fundamentele reden voor de stabiliteit van SFT versus de instabiliteit van PPO.
• Nieuw Kader: De ontwikkeling van LCO, een framework dat de stabiliteit van SFT combineert met de optimalisatiekracht van RL.
• Analyse van Gradiëntdynamiek: Het aantonen dat de instabiliteit in PPO voornamelijk wordt veroorzaakt door negatieve advantage-samples in niet-convexe gebieden, wat leidt tot onnodig grote updates.
• Bewijs van Convergentie: Wiskundige bewijzen dat LCO-objectieven lineaire convergentie garanderen en dat de gradiëntnormen begrensd zijn door de resterende fout, wat een zelf-stabiliserend mechanisme creëert.

Resultaten

De auteurs hebben LCO getest op diverse model-families (Qwen, Llama, Mistral) en taken, waaronder wiskundig redeneren, machine reading comprehension en instructie-opvolging.

• Wiskundig Redeneren (MATH500, AMC23, MinervaMath):
  • LCO-varianten presteerden consistent beter dan standaard RL-baselines (PPO, GRPO, DAPO, GSPO) en zelfs beter dan de Reward Model (RM) zelf in sommige gevallen.
  • Bijvoorbeeld, op de MATH500 benchmark behaalde LCO-KLD met een Qwen-3-4B backbone een Pass@1 van 73.20%, wat een state-of-the-art resultaat is en significant hoger ligt dan PPO (67.80%).
• Machine Reading Comprehension (QA-Feedback):
  • LCO-KLD behaalde de hoogste gemiddelde beloningen voor relevantie, feitelijke juistheid en volledigheid, en overtrof PPO en andere distillatiemethoden aanzienlijk.
• Instructie-opvolging (AlpacaEval 2.0):
  • LCO-varianten behaalden de hoogste win-rates (bijv. 32.93% LC WR voor LCO-KLD op Qwen-3-4B), wat aantoont dat de methode robuust is voor diverse taaktypes.
• Stabiliteit en Efficiëntie:
  • Grafieken tonen aan dat LCO-KLD stabiele gradiëntnormen behoudt tijdens de hele training, terwijl PPO na ongeveer 6.000 stappen begint te oscilleren en instort.
  • LCO is ook sample-efficiënter: LCO-KLD bereikte vergelijkbare prestaties als PPO met bijna de helft van de trainingsdata, en LCO-LCH was zelfs drie keer zo efficiënt dankzij de sterke convexiteit.

Betekenis en Impact

Dit paper biedt een fundamentele theoretische verklaring voor een van de grootste uitdagingen in het trainen van LLM's: de instabiliteit van RL. Door de focus te verleggen van heuristische beperkingen (zoals clipping in PPO) naar het behoud van wiskundige eigenschappen (convexiteit) in het optimalisatielandschap, biedt LCO een meer robuuste en betrouwbare aanpak.

De resultaten suggereren dat het mogelijk is om RL-training voor LLM's te stabiliseren zonder in te boeten aan prestaties, wat cruciaal is voor het schalen van RL-toepassingen in complexe domeinen zoals redeneren en creatief schrijven. De methode is breed toepasbaar en vereist geen ingewikkelde architectuurveranderingen, maar herformuleert de leerdoelstellingen op een wiskundig onderbouwde manier.