Technische Samenvatting: Privileged Value Functions voor LLM Reinforcement Learning

Probleemstelling

Reinforcement Learning (RL) voor Large Language Models (LLMs) wordt momenteel gedomineerd door critic-vrije, groep-relatieve methoden zoals GRPO (Group Relative Policy Optimization). Hoewel deze methoden de variantie van de gradiënt verminderen door meerdere reacties per prompt te samplen en deze met elkaar te vergelijken, lijden ze aan twee primaire beperkingen:

1. Sequentie-niveau Credit Assignment: Ze bieden voordelen alleen op sequentieniveau, wat faalt in het bieden van granulariteit op token-niveau, wat theoretisch superieur is voor temporele credit assignment.
2. Straggler-effecten en Off-Policyness: De training wordt geblokkeerd door de traagste rollout in een groep, wat de doorvoer vermindert. In asynchrone settings verhoogt dit de off-policyness doordat het beleid (policy) afdrijft terwijl het wacht op trage reacties.

Learned value functions pakken theoretisch beide problemen aan door token-niveau voordelen te bieden en te opereren zonder grote groepen. Echter, hun adoptie is gehinderd door:

• Infrastructurele Complexiteit: Het onderhouden van een apart waardemodel voegt engineering-overhead toe.
• Gevoeligheid voor Fitting: Als de critic slecht is aangepast (fitted), kunnen value baselines ruis of bias introduceren, wat vaak slechter presteert dan eenvoudige groepsgemiddelden.
• Gebrek aan Rechtvaardiging: Het praktische succes van critic-vrije methoden heeft het moeilijk gemaakt om de extra kosten van value functions te rechtvaardigen.

Methodologie

Het paper stelt twee complementaire strategieën voor om value-function-gebaseerde RL te verbeteren, met als doel hun inclusie te rechtvaardigen door prestaties en robuustheid te verhogen.

1. Privileged Value Functions (PVF)

De kerninzicht is dat value functions kunnen conditioneren op geprivilegieerde informatie—data die beschikbaar is tijdens de training maar verborgen blijft voor het beleid tijdens de inferentie—zonder de onbevooroordeeldheid van de policy gradient te schenden.

• Mechanisme: Een Privileged Value Function $V_\phi(h_{i,t}, z_{i,t})$ conditioneert de critic op de standaard policy historie $h_{i,t}$ en aanvullende context $z_{i,t}$.
• Toelaatbaarheid: Om te garanderen dat de policy gradient onbevooroordeeld blijft, moet de geprivilegieerde informatie $z_{i,t}$ conditioneel onafhankelijk zijn van de huidige token $y_{i,t}$ gegeven de historie $h_{i,t}$. Het mag geen toekomstige tokens of beloningen bevatten die gegenereerd worden door het huidige traject.
• Vormen van Geprivilegieerde Informatie:
  • Referentie-oplossingen: Oracle antwoorden, bewijs-schetsen of gouden patches (bijv. in Sudoku of code repair).
  • Leave-One-Out (LOO) Groepsinformatie: Andere gesamplede reacties en hun beloningen uit dezelfde prompt-groep. Dit behandelt de groep als een in-context learning prompt voor de critic.
  • Verifier Rubrics: Gedetailleerde taakspecificaties die niet beschikbaar zijn voor het beleid.
• Theoretisch Voordeel: Conditioneren op rijkere informatie vermindert de variantie van de waarde-voorspelling ($E[(R - E[R|H, Z])^2] \le E[(R - E[R|H])^2]$), waardoor de variantie van de policy gradient wordt verminderd.

2. TETHER (Adaptive Group-Value Baseline)

Om het risico van slecht aangepaste critics aan te pakken, introduceren de auteurs TETHER, een baseline die adaptief interpoleert tussen groep-relatieve en value-gebaseerde baselines.

• Mechanisme: TETHER berekent een gemengde baseline $b^{TETHER}_{i,t} = (1-\rho)b^{LOO}_i + \rho V_{i,t}$, waarbij $b^{LOO}_i$ het leave-one-out groepsgemiddelde is en $V_{i,t}$ de geleerde token waarde.
• Adaptieve Coëfficiënt ($\rho$): De mengverhouding $\rho$ is geen vaste hyperparameter. Deze wordt via kleinste kwadraten (least squares) aangepast om de voorspellingsfout van de Monte Carlo returns ($R_i$) op voltooide groepen te minimaliseren.
  • $\rho = 0$: Herstelt de groep-baseline (betrouwbaar, geen token-niveau credit).
  • $\rho = 1$: Herstelt de pure value baseline (token-niveau credit, gevoelig voor de fit van de critic).
  • Tussenliggende waarden: Ruilen soepel tussen de twee uit.
• Trainingsdynamiek: De coëfficiënt $\rho$ wordt bijgewerkt met behulp van een Exponential Moving Average (EMA) op batches na de policy update voor die batch, waardoor de baseline voor een batch niet afhankelijk is van zijn eigen returns (behoud van onbevooroordeeldheid).

Belangrijkste Bijdragen

1. Privileged Value Functions (PVF): Een formeel mechanisme om taak-relevante, token-niveau signalen in de value function te injecteren met behulp van informatie die verborgen is voor het beleid. Dit verbetert de waarde-schatting en vermindert de gradiënt-variantie zonder het policy objectief te wijzigen.
2. TETHER: Een adaptieve baseline die dynamisch balanceert tussen de betrouwbaarheid van groepsgemiddelden en de granulariteit van value functions op basis van de huidige nauwkeurigheid van de critic, waardoor de noodzaak voor handmatige hyperparameter tuning om de transitie te beheren wordt weggenomen.
3. Empirische Validatie: Demonstratie dat beide strategieën consistent beter presteren dan standaard value-gebaseerde baselines en competitief of superieur zijn aan gemiddelde-baseline GRPO over meerdere reasoning-zware taken.

Experimentele Resultaten

De auteurs evalueerden hun methoden op vier taken: Reasoning Gym (procedurele redenering), CodeIO (programmaredenering), Sudoku (multi-turn grid filling), en MiniF2F (formele wiskunde).

• PVF Prestaties:

  • PVF presteerde consistent beter dan zowel de standaard Value Function (VF) als de Group Mean (MEAN) baselines in alle settings.
  • Sudoku: Toonde de grootste winst. Toegang tot de opgeloste grid (geprivilegieerde info) vereenvoudigde de taak van de critic om tussenliggende zetten te valideren, wat anders een impliciete oplossing van de puzzel vereist.
  • CodeIO: PVF maakte gebruik van LOO groep-reacties en beloningen als geprivilegieerde info, en presteerde beter dan zowel VF als MEAN, wat aantoont dat groepsinformatie effectief kan worden benut zelfs zonder oracle antwoorden.
  • Explained Variance: PVF verklaarde consistent meer return variantie dan VF, wat correleert met hogere uiteindelijke beloningen.

• TETHER Prestaties:

  • TETHER presteerde consistent beter dan de standaard VF baseline over alle taken.
  • Het mitigeerde succesvol de degradatie die bij VF optrad wanneer de critic vroeg in de training slecht was aangepast.
  • In Reasoning Gym en MiniF2F presteerde TETHER beter dan de MEAN baseline. In Sudoku verkleinde het de kloof met MEAN, en in CodeIO kwam het overeen met MEAN.
  • Coëfficiënt Dynamiek: De adaptieve $\rho$ begon nabij 0 (vertrouwend op groepsgemiddelden) en nam toe naarmate de critic verbeterde, wat de capaciteit van de methode bevestigt om zelf te reguleren op basis van de nauwkeurigheid van de critic.

Betekenis en Claims

Het paper betoogt dat value functions heroverwogen moeten worden als een praktisch onderdeel van LLM RL-pipelines. De auteurs claimen:

• Rechtvaardiging van Kosten: Door geprivilegieerde informatie te gebruiken (PVF) en adaptieve mixing (TETHER), kunnen de prestatiewinsten van value functions de infrastructuurkosten overtreffen, zelfs in scenario's waar critic-vrije methoden momenteel dominant zijn.
• Robuustheid: TETHER biedt een betrouwbaar pad om value functions in bestaande GRPO-pipelines te introduceren zonder het risico op trainingsinstabiliteit door een slecht aangepaste critic.
• Toekomstige Richtingen: Het werk suggereert dat uitgebreide value pretraining op diverse policies en systematische afstemming van de $\lambda$ parameter (voor de bias-variance trade-off) de schaalbaarheid en bruikbaarheid van value-gebaseerde RL voor LLM's verder kunnen verbeteren.

De auteurs benadrukken dat hun aanpak de theoretische garanties van onbevooroordeelde policy gradients handhaaft, terwijl het praktische verbeteringen biedt in trainingsefficiëntie en uiteindelijke prestaties.