Bootstrapped Mixed Rewards for RL Post-Training: Injecting Canonical Action Order

Dit artikel toont aan dat het combineren van een taakbeloning met een beloning voor de canonieke volgorde van acties tijdens RL-post-training de prestaties van een Transformer op zebra-puzzels verbetert, zelfs wanneer het model is gefinetuned op gerandomiseerde oplossingssequenties.

De kern

Stel je voor dat je een slimme robot leert een complex raadsel op te lossen, zoals een "Zebra-puzzel" (een logische puzzel waar je moet raden wie welke hond heeft, in welk huis woont, etc.).

Deze paper beschrijft een slimme manier om die robot nog slimmer te maken, zonder dat je hem opnieuw moet leren hoe de puzzel werkt. Hier is de uitleg in gewone taal, met een paar leuke vergelijkingen:

1. Het Probleem: De Robot die "Willekeurig" Loopt

Stel je voor dat je de robot eerst leert de oplossing te schrijven, maar je laat hem de stappen in een willekeurige volgorde doen.

• Vergelijking: Het is alsof je iemand leert een huis te bouwen, maar je laat hem eerst het dak leggen, dan de muren, dan de fundering, en dan weer het dak. De robot leert wel wat er in het huis moet staan, maar hij leert niet hoe je het logisch bouwt.

Wanneer de robot later zelfstandig moet werken, maakt hij vaak fouten omdat hij geen duidelijk plan heeft. Hij probeert gewoon het antwoord te gissen.

2. De Oplossing: Een "Fluisterend" Hintsysteem

De onderzoekers wilden weten: Kunnen we de robot helpen om de stappen in de goede volgorde te doen, zonder dat we hem opnieuw moeten leren trainen?

Ze gebruikten een techniek genaamd Reinforcement Learning (beloningstherapie). Hierbij krijgt de robot een puntje als hij het raadsel oplost. Maar ze voegden iets nieuws toe: een extra beloning voor de volgorde.

• De Analogie: Stel je voor dat de robot een speler is in een computerspel.
  • De oude manier: Je krijgt alleen punten als je de eindbaas verslaat.
  • De nieuwe manier: Je krijgt punten als je de eindbaas verslaat, PLUS een klein extra puntje als je de levels in de juiste volgorde hebt doorlopen (eerst level 1, dan level 2, etc.).

Zelfs als de robot die extra punten maar heel weinig krijgt (bijvoorbeeld 1% van de totale score), helpt het hem enorm om een beter "plan" te volgen.

3. Hoe werkt het precies? (De "Bootstrapped" Mix)

De onderzoekers hadden een slimme truc nodig om de twee soorten beloningen eerlijk te vergelijken.

• Soms is het heel moeilijk om de puzzel helemaal op te lossen (dat puntje is zeldzaam).
• Soms is het makkelijk om de volgorde een beetje te volgen (dat puntje krijg je vaker).

Als je deze twee gewoon optelt, zou het ene puntje het andere kunnen "overstemmen". Daarom gebruikten ze een automatische schaalvergroting (de "bootstrapped scaling").

• Vergelijking: Het is alsof je twee geluiden mengt in een stereo-installatie. Als het ene geluid (het oplossen van de puzzel) heel zacht is en het andere (de volgorde) heel hard, zou je het harde geluid moeten dempen. Ze stelden de volume-knoppen zo in dat beide geluiden even hard klinken aan het begin, zodat de robot beide signalen even goed kan horen.

4. Het Resultaat: Een Klein Duwtje in de Rug

Het resultaat was verrassend goed:

• De robot die alleen leerde op "oplossen" (zonder volgorde-hints), loste ongeveer 28% van de puzzels op.
• De robot die een heel klein beetje extra beloning kreeg voor de juiste volgorde (zelfs maar 1% van de beloning), kon plotseling 36% van de puzzels oplossen.

De les: Je hoeft de robot niet te vertellen hoe hij moet denken. Je hoeft alleen maar een heel klein "flauw" hintje te geven over de volgorde, en de robot begint vanzelf de logische stappen te volgen alsof hij een eigen wereldmodel heeft.

Samenvatting in één zin

Door een slimme robot een heel klein extra puntje te geven voor het volgen van de juiste volgorde (in plaats van alleen voor het eindantwoord), wordt hij veel beter in het oplossen van logische puzzels, zonder dat je hem opnieuw hoeft te leren hoe de puzzel werkt.

Het is alsof je iemand die een raadsel probeert op te lossen niet vertelt wat het antwoord is, maar alleen fluistert: "Oh, en probeer maar eerst de makkelijkste hints te gebruiken." Dat kleine hintje maakt het verschil tussen vastlopen en slagen.

Technische samenvatting

Probleemstelling

Reinforcement Learning (RL) voor post-training van taalmodellen optimaliseert vaak een enkelvoudig scalair doel (zoals taaksucces), maar negeert de structurele aspecten van hoe oplossingen worden gegenereerd, zoals de volgorde van tussenstappen. Bestaande methoden negeren vaak de volgorde waarin een model acties uitvoert, zelfs als er een "canonieke" (logische) volgorde bestaat die de oplossing efficiënter maakt.

De auteurs stellen de vraag: Kan een scalair hint over een canonieke oplossingsvolgorde, dat uitsluitend tijdens de RL-post-training wordt gebruikt, de prestaties verbeteren, zelfs als het model is gefinetuned op gerandomiseerde oplossingssequenties?

Het onderzoek focust op "Zebra-puzzels" (logische raadsels), waarbij een canonieke solver een deterministische volgorde van stappen volgt om het raadsel op te lossen, terwijl een willekeurige volgorde geen logische structuur heeft.

Methodologie

1. Data en Opzet

• Dataset: De auteurs gebruiken de Zebra-puzzel-dataset (Shah et al., 2024). Elke puzzel wordt opgelost met precies 9 acties (het invullen van een grid).
• Varianten: Er worden twee sequenties onderscheiden:
  • Solver-order: De chronologische volgorde waarin een canonieke solver de cellen invult (logische redenering).
  • Random-order: Dezelfde acties, maar willekeurig geschud.
• Model: Een GPT-2-achtige Transformer (4 lagen, 4 attention heads, hidden size 256), getraind vanaf nul.

2. Trainingsfasen

• Standaard Finetuning: Het model wordt eerst gefinetuned op de dataset met randomized solution orders onder een causaal taalmodelleringsdoel. Het model leert dus de puzzels op te lossen, maar zonder de logische volgorde van de stappen te zien.
• RL Post-training (GRPO): Vervolgens wordt het model gepost-getraind met Group Relative Policy Optimization (GRPO). Hierbij worden twee soorten beloningen (rewards) gebruikt:
  1. Solved Reward ($R_{solve}$): Een schaarse beloning (1 of 0). Het is 1 alleen als het model een volledig correcte oplossing produceert.
  2. Ordering Reward ($R_{order}$): Een vorm van reward shaping. Deze beloning meet hoe goed de emissievolgorde van het model overeenkomt met de canonieke solver-volgorde, onafhankelijk van of de waarden correct zijn. De beloning is hoger naarmate de indices van de gegenereerde cellen dichter bij de canonieke indices liggen.

3. Beloningsmix en Bootstrapped Scaling
Om de effecten van deze signalen te isoleren, worden de beloningen gecombineerd via een vaste gewogen som:
$$R_{total} = \alpha \cdot R_{solve} + (1 - \alpha) \cdot R_{order}$$

Een cruciale innovatie is de bootstrapped reward scaling:

• Omdat de absolute grootte van $R_{solve}$ en $R_{order}$ sterk kan verschillen, zou een simpele mix leiden tot overheersing door één component.
• De auteurs evalueren het gefinetunde model op een validatieset om de gemiddelde beloningen ($\bar{R}_{solve}$ en $\bar{R}_{order}$) te berekenen.
• Vervolgens worden globale schalingsfactoren ingesteld zodat bij initialisatie de bijdrage van elke component exact overeenkomt met de gewenste mix $\alpha$. Deze factoren blijven constant tijdens de hele post-training.

Belangrijkste Bijdragen

1. Scalair Hint-methode: Een methode om solver-volgorde in GRPO te injecteren zonder de supervised trainingsdata of de modelarchitectuur aan te passen. Het model ziet nooit de canonieke volgorde tijdens finetuning; het leert deze alleen via de RL-beloning.
2. Bootstrapped Scaling Procedure: Een procedure om heterogene beloningsmagnitudes te normaliseren, waardoor gecontroleerde mix-studies mogelijk zijn zonder dat de schaal van de beloning de resultaten verstoort.
3. Empirisch Bewijs: Het aantonen dat ruwe, grove volgorde-signalen, wanneer gemixt met correctheid, de nauwkeurigheid van RL-post-training verbeteren.

Resultaten

De prestaties werden gemeten aan de hand van puzzelnauwkeurigheid (fraction van volledig opgeloste puzzels) op een testset.

• Baseline: Het model gefinetuned op random volgorde behaalde 27,9% nauwkeurigheid.
• Alleen Taak Reward (1 : 0): Post-training met alleen de "solved" reward leverde 32,6% op.
• Gemixte Rewards: Het toevoegen van zelfs een zeer kleine component van de ordering reward leidde tot significante verbeteringen:
  • 0.99 : 0.01 (Solve : Order): Bereikte de beste prestatie van 36,3%.
  • Andere mixen (bijv. 0.95 : 0.05 en 0.9 : 0.1) presteerden ook aanzienlijk beter dan de taak-only baseline (rond de 35-35,5%).

Conclusie van de resultaten: Zelfs een minimale "ordering hint" stuurt het beleid effectief naar canonieke trajecten, wat resulteert in betere oplossingen, zonder dat het model expliciet op geordende data is getraind.

Betekenis en Conclusie

Dit werk toont aan dat structurele bias (in dit geval de volgorde van redenering) effectief kan worden ingebracht in RL-post-training via een simpel scalair signaal, zonder de noodzaak van nieuwe, duur te cureren datasets of complexe architectuurwijzigingen.

• Efficiëntie: Het biedt een goedkope en modulaire "knop" om de redeneerstructuur van een model te verbeteren.
• Generalisatie: Hoewel het experiment beperkt is tot Zebra-puzzels en één modelgrootte, suggereert het dat het begrijpen en modelleren van de "wereld" (in dit geval de dynamiek van de puzzeloplossing) essentieel is voor succesvolle RL.
• Toekomst: De auteurs merken op dat de vaste schalingsfactoren een beperking kunnen zijn als de beloningscomponenten tijdens het trainen verschillende snelheden van verbetering vertonen. Dynamische aanpassing van deze schaling is een logische volgende stap.

Samenvattend bewijst dit paper dat het "leren van de wereld" (de volgorde van acties) via beloningsvorming een krachtige methode is om de redeneercapaciteiten van taalmodellen te verbeteren, zelfs wanneer ze oorspronkelijk op ongeordende data zijn getraind.