SPS: Steering Probability Squeezing for Better Exploration in Reinforcement Learning for Large Language Models

Probleemstelling: Het "Squeezing Effect" in RL

Hoewel Reinforcement Learning (RL) een veelbelovende paradigma is voor het trainen van redenerende Large Language Models (LLMs) via regelgebaseerde beloningssignalen, vertoont het een fundamentele beperking. Traditionele RL-methoden (zoals GRPO) verbeteren vaak de prestaties voor een enkele steekproef (Pass@1), maar falen in het effectief verkennen van diverse redeneertrajecten, wat essentieel is voor multi-sample prestaties (Pass@k).

De auteurs identificeren een fundamenteel "Squeezing Effect" (klem-effect) als de oorzaak:

• Mechanisme: Tijdens het RL-training worden negatieve gradiënten toegepast op low-probability tokens (foutieve antwoorden). Door de normalisatie-eigenschappen van de softmax-functie, wordt de waarschijnlijkheidsmassa die van deze zeldzame tokens wordt verwijderd, niet gelijkmatig herverdeeld. In plaats daarvan wordt deze massa "samengeperst" naar de al dominante, hoogstwaarschijnlijke tokens (de "greedy" antwoorden).
• Gevolg: De output-verdeling wordt extreem scherp (peaked), wat leidt tot een verlies aan entropie en diversiteit. Het model blijft vastzitten in een smalle subset van hoge-beloning trajecten en kan geen nieuwe, potentieel betere redeneerpaden ontdekken. Dit resulteert in een empirische bovengrens voor Pass@k, ongeacht hoe lang het training doorgaat.

Methodologie: Steering Probability Squeezing (SPS)

Om dit klem-effect te mitigeren en de exploratie te verbeteren, stellen de auteurs Steering Probability Squeezing (SPS) voor. Dit is een trainingsparadigma dat conventioneel RL alternerend combineert met Inverse Reinforcement Learning (IRL).

Het SPS-framework werkt in iteratieve lussen met twee fasen:

1. RL-fase (Exploratie):

   • Er wordt standaard RL (bijv. GRPO) uitgevoerd om rollouts te genereren en de policy te updaten op basis van beloningen.
   • Hierbij worden rollouts verzameld, waarbij een subset wordt geselecteerd die nadruk legt op trajecten met een relatief lage waarschijnlijkheid (Low-Likelihood Trajectory Emphasis - L2TE), maar die toch correct kunnen zijn.

2. IRL-fase (Distributieherstructurering):

   • In plaats van externe expertdata te gebruiken, behandelt SPS de gegenereerde on-policy rollouts als "demonstraties".
   • Een Forward-KL objective wordt gebruikt om de policy-distributie opnieuw te vormen. De loss-functie ($L_{IRL}$) minimaliseert de divergentie tussen de huidige policy en de empirische verdeling van de geselecteerde rollouts.
   • Doel: Dit herverdeelt de waarschijnlijkheidsmassa weg van de overmatig dominante (greedy) antwoorden naar onderbelichte, maar waardevolle alternatieven. Het "stuurt" het klem-effect dus in een richting die exploratie bevordert in plaats van concentratie.

Deze twee fasen worden herhaaldelijk afgewisseld, waardoor het model continu de grenzen van oplosbare problemen kan uitbreiden zonder voorbarig te convergeren naar een smalle set van gedragingen.

Belangrijkste Bijdragen

1. Analyse van RL-dynamiek: De auteurs bieden een theoretische en empirische analyse die het "squeezing effect" als een intrinsieke beperking van RL-basistraining identificeert. Ze formaliseren hoe probability mass redistribution leidt tot progressieve concentratie en stellen een empirische bovengrens voor Pass@k vast.
2. Het SPS-framework: De introductie van een nieuwe trainingsstrategie die IRL gebruikt om de door RL veroorzaakte distributie te herschikken, zonder externe supervisie of prior kennis.
3. Iteratieve Trainingsstrategie: Een cyclus van RL en IRL-updates die zorgt voor een progressieve herverdeling van waarschijnlijkheidsmassa, waardoor het model in staat blijft om diverse redeneertrajecten te ontdekken.

Resultaten

De methode is geëvalueerd op vijf olympiade-niveau wiskundige benchmarks (AIME, BRUMO, HMMT) met modellen uit de Qwen2.5-Math-serie (1.5B en 7B parameters).

• Prestatieverbetering: SPS toont consistente en aanzienlijke verbeteringen in Pass@128 (multi-sample prestatie) vergeleken met state-of-the-art baselines zoals GRPO, DAPO en GSPO.
• Specifiek Voorbeeld: Op het BRUMO-benchmark bereikte het Qwen2.5-Math-1.5B-model met SPS een Pass@128-score van 63.33. Dit is een verbetering van +10.00 punten ten opzichte van de standaard GRPO.
• Diversiteit: Analyse toont aan dat SPS de traject-similariteit verlaagt (meer diversiteit) en de log-probabiliteit dichter bij het basismodel houdt, wat aantoont dat het klem-effect effectief wordt tegengegaan.
• Efficiëntie: De IRL-fase voegt slechts een minimaal computatie-overhead toe (ongeveer 3% van de totale trainingstijd per iteratie), terwijl de RL-rollout de meeste tijd kost.

Betekenis en Conclusie

Dit paper biedt een cruciale inzichten in de beperkingen van huidige RL-methoden voor redenerende LLM's. Het toont aan dat het simpelweg verhogen van de sample-efficiëntie (Pass@1) niet voldoende is voor complexe redeneertaken; echte exploratie vereist het actief beheersen van de kansverdeling.

Door het "squeezing effect" te sturen via een combinatie van RL en IRL, biedt SPS een effectief pad om de exploratiecapaciteit van redenerende modellen te vergroten. Dit leidt tot modellen die niet alleen vaker het juiste antwoord vinden bij één poging, maar ook een breder scala aan correcte redeneerpaden kunnen genereren, wat essentieel is voor de volgende generatie AI-systemen die complexe problemen moeten oplossen. De bevindingen suggereren dat het afwisselen van RL en IRL een fundamentele verbetering is ten opzichte van puur RL-gebaseerde training.