Technische Samenvatting: Over de schattingen van voordelen voor Max@K Policy Gradients

Probleemstelling

Reinforcement learning met verifieerbare beloningen (RLVR) is een dominant paradigma voor het post-trainen van grote taalmodellen (LLM's) op redeneertaken. Deze instellingen vertrouwen echter doorgaans op ijle, op uitkomst gebaseerde beloningen (bijv. een binaire correct/incorrect signaal), wat exploratie in grote combinatorische ruimtes moeilijk maakt. Hoewel methoden zoals entropie-regularisatie en count-based bonussen bestaan, introduceren zij hulpdoelen die zorgvuldige afstemming vereisen.

Een complementaire aanpak is het direct optimaliseren van inferentie-doelen zoals pass@K (de waarschijnlijkheid dat ten minste één van $K$ gesamplede voltooiingen correct is) en zijn continue generalisatie, max@K (de verwachte maximale beloning onder $K$ samples). Verschillende recente werken hebben onbevooroordeelde policy-gradient (PG) schatters voorgesteld voor deze doelen. De relaties tussen deze schatters blijven echter onduidelijk. Specifiek maken bestaande schatters vaak gebruik van verschillende signalen, baselines en normalisaties. Een cruciaal probleem dat in dit werk wordt geïdentificeerd, is dat vooraanstaande schatters (bijv. gebaseerd op Expected Improvement) onbevooroordeeld zijn als PG-schatters, maar niet-gecentreerde voordelen opleveren. Een niet-gecentreerd voordeel (één met een niet-nul gemiddelde) komt overeen met een onderschatte baseline, wat de optimalisatie naar suboptimale acties kan sturen en de gradiëntvariantie kan verhogen.

Methodologie

Het paper stelt MaxPO (Max@K Policy Optimization) voor, een methode die een Leave-Two-Out (L2O) baseline introduceert om de schatting van het voordeel voor max@k-doelen te corrigeren.

1. Analyse van bestaande schatters

De auteurs herzien de Expected Improvement (EI) schatter die in eerder werk wordt gebruikt (bijv. Walder en Karkhanis [54]). Ze laten zien dat hoewel de EI-gebaseerde schatter een onbevooroordeelde PG-schatter is, het resulterende voordeelssignaal $s(a)$ strikt niet-negatief is (vanwege de ReLU-operatie in de definitie van EI). Bijgevolg is $E[s(a)] > 0$, wat betekent dat het voordeel niet gecentreerd is. Dit gebrek aan centrering laat ruimte voor variantiereductie.

2. De Leave-Two-Out (L2O) Baseline

Om een exact gecentreerd voordeel te bereiken terwijl de PG-onbevooroordeeldheid behouden blijft, introduceren de auteurs de L2O baseline.

• De Uitdaging: Een standaard Leave-One-Out (L1O) baseline faalt in de max@K-setting omdat de beloning van een sample $i$ de berekening van het voordeel van andere samples $j$ impliciet kan beïnvloeden (omdat $i$ fungeert als een potentiële comparator in de max-operatie). Deze afhankelijkheid schendt de onafhankelijkheidsconditie die vereist is voor onbevooroordeelde PG-schatting.
• De Oplossing: De L2O baseline $b_{-i}^{L2O}$ wordt berekend door het gemiddelde te nemen van de Expected Improvement-scores van alle andere samples $j$, waarbij elke score wordt berekend over een subset van comparators die zowel de doel-sample $i$ als de sample $j$ zelf uitsluit.
• Eigenschappen:
  • Onafhankelijkheid: Door $i$ uit de comparator-set voor alle $j$ te sluiten, is de baseline $b_{-i}^{L2O}$ onafhankelijk van de beloning $r_i$, waardoor de PG-schatter onbevooroordeeld blijft.
  • Centrering: De auteurs bewijzen dat het gerealiseerde batch-gemiddelde voordeel exact nul is voor elke gesamplede batch: $\frac{1}{B}\sum_{i=1}^B (s_i - b_{-i}^{L2O}) = 0$.
  • Efficiëntie: Hoewel een naïeve berekening $O(B^3)$ zou zijn, leiden de auteurs een vectoriseerde implementatie af met behulp van ReLU-gebaseerde formuleringen die de L2O baseline in $O(B^2)$ tijd berekent, wat het haalbaar maakt voor GPU-versnelde LLM-training.

3. Canonieke Finite-Batch Advantage

Het paper leidt een canonieke finite-batch advantage af voor max@K, gedefinieerd als $u_i - v_i$:

• $u_i$: De conditionele verwachte max@K waarde van een groep die respons $i$ bevat.
• $v_i$: De leave-one-out verwachte max@K waarde berekend uit groepen die $i$ uitsluiten.
  De auteurs tonen aan dat de MaxPO-schatter (EI + L2O) exact dit canonieke signaal herstelt ($s_i - b_{-i}^{L2O} = u_i - v_i$). Dit biedt een verenigd beeld van bestaande schatters en classificeert ze als "niet-gecentreerd" (bijv. EI-alleen/PKPO) en "(genormaliseerd) canoniek" (bijv. MaxPO, en bepaalde analytische pass@K varianten).

4. Integratie in LLM Training

MaxPO wordt geïntegreerd in groep-gebaseerde RL-algoritmen (zoals GRPO). Voor elke prompt wordt een groep van $G$ outputs gesampled. De L2O baseline wordt binnen elke groep berekend om sequentie-niveau voordelen te genereren, die vervolgens worden doorgegeven aan token-niveau losses.

Belangrijkste Bijdragen

1. Identificatie van Niet-Gecentreerde Voordelen: Het paper identificeert dat vooraanstaande max@K PG-schatters (specifiek EI-gebaseerde) onbevooroordeeld zijn als gradiëntschatters, maar falen in het produceren van gecentreerde voordelen, wat leidt tot suboptimale variantie-eigenschappen.
2. Leave-Two-Out (L2O) Baseline: Een nieuw baseline-ontwerp dat de PG-onbevooroordeeldheid behoudt in de niet-lineaire max@K-setting, terwijl het zorgt dat het gerealiseerde batch-voordeel een gemiddelde van nul heeft.
3. Canonieke Unificatie: De afleiding van de canonieke finite-batch advantage ($u_i - v_i$) en het bewijs dat MaxPO dit signaal herstelt, biedt een verenigd kader om bestaande pass@K/max@K schatters te vergelijken en te categoriseren.
4. Efficiënte Implementatie: Een $O(B^2)$ vectoriseerbaar algoritme voor het berekenen van de L2O baseline, wat praktische toepassing in LLM post-training mogelijk maakt.

Experimentele Resultaten

De auteurs valideren de methode via gecontroleerde bandit/maze experimenten en LLM-redeneringstaken.

• Variantiereductie: In bandit-omgevingen vermindert de L2O baseline de gradiëntvariantie aanzienlijk vergeleken met ruwe EI-schatters, waarbij het effect groter wordt naarmate de grootte van de actieruimte toeneemt. In LLM-training (Llama-3.2-3B-Instruct) verminderde de methode de gradiëntvariantie-proxy met 77,4% vergeleken met PKPO (de EI-alleen baseline).
• Prestatiewinst: Op vijf wiskundige redeneringsbenchmarks (AIME24, AIME25, AMC23, MATH500, Minerva) presteerde MaxPO beter dan sterke baselines inclusief PKPO, GRPO en Entropy-Adv.
  • Op Qwen2.5-Math-7B verbeterde MaxPO de taakgemiddelde pass@256 met 5,2% ten opzichte van PKPO.
  • Op Llama-3.2-3B-Instruct verbeterde het pass@256 met 2,4% ten opzichte van PKPO.
• Inference Compute Scaling: De prestatiewinsten van MaxPO waren het meest uitgesproken bij hogere niveaus van inference compute (grotere $k$ in pass@k), wat suggereert dat de methode het beleid helpt om extra test-time samples beter te benutten.

Betekenis en Claims

Het paper claimt dat MaxPO een principieel fundament biedt voor het ontwerpen en vergelijken van policy-gradient schatters voor inferentie-doelen zoals max@K. Door de canonieke vorm van het voordeel te identificeren en de bias in bestaande schatters te corrigeren, biedt het werk een methode om de efficiëntie en stabiliteit van post-training redeneermodellen te verbeteren.

De auteurs merken bescheiden op dat hun analyse zich richt op finite-batch, on-policy schatting onder i.i.d. sampling. Zij erkennen beperkingen met betrekking tot off-policy training, adaptieve sampling en settings met gecorreleerde generaties als richtingen voor toekomstig werk. Bovendien benadrukken zij dat hoewel de methode de optimalisatie verbetert, een betere optimalisatie van herhaalde-sampling redeneerdoelen bredere impacten heeft, wat potentieel modellen kan versterken die misbruikt kunnen worden voor misleidende assistentie of geautomatiseerde probleemoplossing.