Technische Samenvatting: Hergebruik je FLOPs: Schalen van RL op Moeilijke Problemen door te Conditioneren op Zeer Off-Policy Prefixes

1. Probleemstelling

Reinforcement Learning (RL) is de de facto methode geworden voor het verbeteren van de redeneercapaciteiten van Large Language Models (LLM's), met name in wiskunde en programmeren. Echter, standaard on-policy RL-methoden worden geconfronteerd met een kritieke bottleneck wanneer ze worden toegepast op moeilijke problemen waarbij de pass rate van het basismodel bijna nul is (bijv. pass@2k ≈ 0).

In dit regime samplet het RL-beleid (policy) zelgen een correct redeneertraject. Hierdoor ontstaan de volgende problemen:

• Verdwijnende leersignalen: De policy gradient-schattingen worden nul of bijna nul omdat alle gesamplede trajecten een beloning van nul ontvangen.
• Verspilling van rekenkracht: Het systeem besteedt enorme hoeveelheden FLOPs aan sampling zonder dat er een leersignaal wordt ontvangen, wat ervoor zorgt dat de training stagneert en de beloningen een plateau bereiken.
• Inefficiëntie van bestaande oplossingen:
  • Supervised Fine-Tuning (SFT) op off-policy trajecten: Hoewel dit het model kan bootstrappen, leidt dit vaak tot entropie-instorting (entropy collapse) en memorisatie, wat verdere exploratie belemmert.
  • Directe Off-Policy RL: Het direct gebruiken van off-policy trajecten via importance weighting leidt tot gradiëntschattingen met een hoge variantie en instabiliteit.

De kernuitdaging is hoe de "verspilde" rekenkracht uit eerdere inferentie- of RL-runs (die mogelijk zeldzame correcte trajecten bevatten) kan worden hergebruikt om het online RL-beleid te sturen, zonder de instabiliteiten te introduceren die gepaard gaan met directe off-policy supervisie.

2. Methodologie: PrefixRL

De auteurs stellen PrefixRL voor, een framework dat het on-policy RL-proces conditioneert op prefixes van succesvolle off-policy trajecten, in plaats van die trajecten als directe supervisiedoelen te gebruiken.

Kernmechanisme

1. Dataverzameling: Een dataset van off-policy trajecten ($D_{off}$) wordt verzameld, doorgaans via grootschalige rejection sampling op het basismodel totdat er voor elk probleem een correct traject is gevonden.
2. Probleemaugmentatie: Voor elk origineel probleem $x$ extraheert het systeem een prefix (de eerste $h$ tokens) uit een correct off-policy traject $y_x$. Deze prefix wordt aan de originele probleemprompt toegevoegd om een geprefixeerd probleem te creëren: $\text{concat}(x, y_{x, 1:h})$.
3. Trainingsdoelstelling: Het model wordt getraind met on-policy RL op een mix van:
   • Originele (zonder prefix) problemen.
   • Geprefixeerde problemen.
   • Cruciaal: Gradiënten worden gemaskeerd op de off-policy prefix-tokens. Het model wordt alleen bijgewerkt op de tokens die het genereert na de prefix.

Theoretische Garanties

Het artikel biedt theoretische bewijzen met betrekking tot de PrefixRL-doelstelling:

• Consistentie: Als de off-policy prefixes correct zijn en realiseerbaar door een policy binnen de modelklasse, dan is elke maximizer van de PrefixRL-doelstelling ook een maximizer van de standaard RL-doelstelling. De methode introduceert geen bias in de optimale oplossing.
• Sample-efficiëntie: PrefixRL vermindert de suboptimaliteitskloof vergeleken met standaard RL. Door het beleid in toestanden met hogere beloningen te plaatsen (toestanden waar een correcte oplossing al gedeeltelijk gevormd is), wordt de exploratiebelasting verminderd. Theoretische grenzen suggeren een reductie in het aantal benodigde on-policy samples, wat gunstig schaalt met de contextlengte.

Belangrijk Fenomeen: Back-Generalization

Empirisch ontdekken de auteurs back-generalization: het trainen alleen op geprefixeerde problemen verhoogt de prestaties op de originele, ongeprefixeerde problemen aanzienlijk, zelfs wanneer het model tijdens die fase nooit expliciet op de ongeprefixeerde versies is getraind.

• Mechanisme: Het model imiteert niet simpelweg de prefix. In plaats daarvan leert het strategieën die generaliseren naar de ongeprefixeerde staat.
• Strategie-ontdekking: Het model kan suboptimale strategieën die aanwezig zijn in de prefix "ontleren" en nieuwe, betere strategieën ontdekken. Het wijzigt de next-token distributie op ongeziene toestanden, waardoor het effectief "terugwaarts" beweegt (backtracking) vanuit de hoog-belonende toestanden die door de prefix worden geboden om betere paden in de ongeprefixeerde ruimte te vinden.
• Gerelateerdheid: Back-generalization is het meest effectief wanneer de prefix (of in-context probleem) semantisch gerelateerd is aan het doelprobleem.

3. Belangrijkste Bijdragen

1. PrefixRL Framework: Een nieuwe methode om off-policy FLOPs te recyclen door te conditioneren op correcte prefixes, waarbij de instabiliteit van directe off-policy RL of SFT-geïnduceerde entropie-instorting wordt vermeden.
2. Theoretische Analyse: Bewijs dat PrefixRL consistent is met standaard RL-doelstellingen en voorzien is van bewijsbare winst in sample-efficiëntie door de exploratiebelasting te verminderen.
3. Ontdekking van Back-Generalization: De identificatie van een fenomeen waarbij conditionering op prefixes kan leiden tot superieure prestaties op ongeprefixeerde taken, wat vaak beter is dan standaard generalisatie over gerelateerde taken.
4. Zelfverbeteringslus: Demonstratie van een praktische pipeline waarbij off-policy trajecten worden gegenereerd via rejection sampling van het basismodel zelf, wat een zelfverbeterende cyclus creëert zonder externe docentmodellen.

4. Experimentele Resultaten

De auteurs evalueerden PrefixRL op moeilijke wiskundige redeneerbenchmarks (DAPO, OMNI-MATH, AIME '25, HMMT '25) met behulp van Llama-3.1-8B en Qwen-3-4B modellen.

• Reken-efficiëntie: PrefixRL bereikt dezelfde trainingsbeloning 2× sneller dan de sterkste baseline (SFT op off-policy data gevolgd door RL), zelfs na rekening te houden met de rekenkosten van de initiële rejection sampling.
• Nauwkeurigheidswinst:
  • Op moeilijke trainingsproblemen verhoogt PrefixRL de uiteindelijke nauwkeurigheid met >45% (meer dan 3× relatief) vergeleken met de baseline.
  • Op held-out benchmarks zoals AIME '25 verbetert PrefixRL de pass@1 met 12% ten opzichte van de sterkste baseline in een compute-matched vergelijking.
• Robuustheid van de Bron: De methode blijft effectief, zelfs wanneer off-policy prefixes afkomstig zijn van een andere modelfamilie (bijv. het gebruik van Qwen-prefixes om een Llama-policy te trainen), wat de flexibiliteit in praktische settings aantoont.
• Trainingsdynamiek:
  • Stabiliteit: In tegen tegenstelling tot SFT+RL, behoudt PrefixRL de token-niveau entropie, wat voorkomt dat exploratie instort.
  • Gradiëntkwaliteit: PrefixRL vertoont hogere gradiëntnormen en een lagere variantie vergeleken met standaard off-policy RL (dat lijdt onder hoog-variante importance weights).
  • Efficiëntie: Het model genereert gemiddeld kortere antwoorden, wat aangeeft dat het eerder beslissende stappen bereikt en "onproductief dwalen" vermindert.

5. Betekenis en Claims

Het artikel claimt dat PrefixRL de aanpak van het schalen van RL op moeilijke problemen fundamenteel verandert. In plaats van off-policy data te zien als een doelwit voor imitatie (wat leidt tot instabiliteit of instorting) of het weg te gooien als ruis, behandelt PrefixRL het als een conditioneringsmechanisme om exploratie te sturen.

De betekenis ligt in:

• Hergebruik van FLOPs: Het transformeert voorheen verspilde rekenkracht (mislukte pogingen die uiteindelijk een correct traject opleverden via rejection sampling) in een direct signaal voor leren.
• Oplossen van het "Moeilijk Probleem" Regime: Het maakt RL-training mogelijk op problemen waarbij on-policy rollouts bijna nooit slagen, een regime waarin standaard RL stagneert.
• Back-Generalization: De ontdekking dat conditionering op prefixes kan leiden tot superieure prestaties op ongeprefixeerde taken suggereert een krachtig mechanisme om eerdere rekenkracht te benutten dat verder gaat dan eenvoudige generalisatie, wat potentieel de manier waarop toekomstige RL-algoritmen historische data incorporeren, kan hervormen.

De auteurs concluderen dat het benutten van back-generalization een cruciale ingrediënt is voor het bootstrappen van eerdere rekenkracht en het mogelijk maken van continue zelfverbetering in redenerende LLM's.