Technische Samenvatting: Outcome-Supervised Process Reward Modeling via Learnable Credit Assignment (LCA)

1. Probleemstelling

Process Reward Models (PRM's) verbeteren de redeneervaardigheden van Large Language Models (LLM's) door fijne, op stappen gebaseerde feedback te geven. Het trainen van PRM's vereist echter doorgaans dure menselijke annotaties voor elke redeneerstap. Om dit aan te pakken, proberen Outcome-Supervised PRM's te leren van enkel de correctheid van het uiteindelijke antwoord, waardoor stap-specifieke labels worden vermeden.

De kernuitdaging in deze setting is credit assignment (toeschrijving van krediet): het toeschrijven van de uiteindelijke uitkomst (correct/incorrect) aan de specifieke redeneerstappen die de uitkomst hebben veroorzaakt. Bestaande benaderingen lijden aan fundamentele tekortkomingen:

• Uniform Assignment: Gaat ervan uit dat alle stappen evenveel bijdragen aan de uiteindelijke uitkomst (bijv. via additieve aggregatie). Dit onderscheidt geen correcte voorstappen van de specifieke stap waar een fout optreedt.
• Causal Assignment: Propageert invloed achterwaarts vanuit het uiteindelijke antwoord met behulp van Monte Carlo-rollouts (bijv. het schatten van de waarschijnlijkheid om een correct antwoord te bereiken vanaf een bepaalde stap). Dit vervaagt de logische correctheid van een stap met de verwachte succesratio van downstream completions, wat het moeilijk maakt om de werkelijke eerste fout te lokaliseren.

Beide methoden falen in het verankeren van krediet in de feitelijke correctheid van individuele stappen, wat het vermogen van het model belemmert om procesfouten nauwkeurig te identificeren.

2. Methodologie: LCA-framework

De auteurs stellen Outcome-Supervised Process Reward Modeling via Learnable Credit Assignment (LCA) voor, een framework dat zowel credit assignment als reward modeling gezamenlijk leert zonder stap-niveau supervisie.

2.1 Theoretische Formulering: Weakest Link Assignment

LCA is gebaseerd op het Weakest Link Assignment principe: een redeneerketen is zo sterk als zijn zwakste schakel. Een traject is correct als en slechts als elke stap correct is; het is incorrect als elke stap fout is.

• Dit principe transformeert het outcome-supervised PRM probleem naar een Multiple Instance Learning (MIL) probleem.
• Bag: Het volledige redeneetraject.
• Instances: De redeneerprefixes (stappen) binnen het traject.
• Label: De traject-niveau correctheid (positief als het traject een fout bevat, negatief anders).
• Aanname: De standaard Multi-Instance Assumption houdt stand: een bag is positief als en alleen als deze ten minste één positieve instance bevat (een foutieve stap).

2.2 De Uitdaging van Afhankelijkheid

Standaard MIL-methoden gaan vaak uit van onafhankelijke of zwak afhankelijke instances. Echter, redeneertrajecten vertonen autoregressieve, geneste en monotone eigenschappen. In een traject bepaalt de laatste prefix alleen de bag-label, waardoor standaard pooling-technieken (zoals Max Pooling of Attention Pooling) degenereren. Het model heeft de neiging om tussenliggende stappen te negeren, waardoor het effectief een Outcome Reward Model (ORM) wordt dat alleen de eindtoestand voorspelt.

2.3 Softmax-Weighted-Sum (SWS) Pooling

Om de wederzijdse afhankelijkheid tussen credit assignment en reward modeling op te lossen, introduceert LCA Softmax-Weighted-Sum (SWS) pooling.

• Mechanisme: Gegeven instance-niveau foutkansen $p = (p_1, \dots, p_T)$, is de bag-niveau voorspelling een gewogen som:
  $$\hat{p} = \sum_{t=1}^T w_t p_t, \quad \text{waarbij } w_t = \frac{e^{p_t/\alpha}}{\sum_{j=1}^T e^{p_j/\alpha}}$$
• Rol van Temperatuur ($\alpha$): De parameter $\alpha$ controleert de entropie van de gewichtsverdeling.
  • Wanneer $\alpha \to 0$, nadert SWS Max Pooling (strikte Weakest Link).
  • Wanneer $\alpha \to \infty$, nadert SWS Average Pooling (Uniform Assignment).
• Voordelen:
  1. Niet-verwaarloosbaar Krediet: In tegen tegenstelling tot Max Pooling zorgt SWS ervoor dat elke instance een strikt positief gewicht ontvangt, wat voorkomt dat het model tussenliggende stappen negeert.
  2. Structurele Prior: De softmax-weging codeert op natuurlijke wijze de Standard Multi-Instance Assumption, waarbij instances met hogere voorspelde foutkansen (positieve instances) de bag-voorspelling domineren.
  3. Bayesiaanse Consistentie: De auteurs bewijzen dat onder milde aannames, het minimaliseren van de verwachte loss met SWS pooling leidt tot een Bayes-optimale instance classifier, zelfs zonder instance-niveau labels.

2.4 Trainingsdoelstelling

De PRM wordt end-to-end getraind met een traject-niveau cross-entropy loss. Het model geeft per-prefix foutkansen output, die via SWS pooling worden geaggregeerd om overeen te komen met de traject-niveau grondwaarheid (eindantwoord correctheid). Een stop-gradient operatie wordt toegepast op de softmax-gewichten om de training te stabiliseren.

3. Belangrijkste Bijdragen

1. Inzichtelijke Analyse: Het artikel identificeert dat bestaande outcome-supervised PRM's (Uniform en Causal) falen in het verankeren van krediet in stap-correctheid. Het stelt Weakest Link Assignment voor als een principieel alternatief dat aansluit bij de "first-error" criterium die wordt gebruikt bij de evaluatie van PRM's.
2. Technische Innovatie (LCA):
   • Formaliseert outcome-supervised PRM als een MIL-probleem onder het Weakest Link principe.
   • Introduceert SWS pooling, een MIL-techniek die is afgestemd op de sterke afhankelijkheid en redundantie van redeneertoestanden, waarbij een gesmoothde Weakest Link assignment als een structurele prior wordt ingebed.
   • Levert een theoretisch bewijs van Bayesiaanse consistentie voor de leerdoelstelling op het instance-niveau.
3. Empirische Prestaties: Toont aan dat LCA consistent beter presteert dan state-of-the-art outcome-supervised PRM's over meerdere taken (foutidentificatie, best-of-N scaling) en backbone modellen (Llama, Qwen).

4. Experimentele Resultaten

De auteurs evalueerden LCA op wiskundige redeneertaken met de Math-Shepherd dataset voor training en ProcessBench voor evaluatie.

• Foutidentificatie (RQ1): Op ProcessBench behaalde LCA een nieuwe state-of-the-art.
  • Met een 7B backbone overtrof LCA de op één na beste methode met 5,5% in F1-score.
  • Met een 3B backbone presteerde het aanzienlijk beter dan alle concurrerende 3B methoden.
  • Opvallend genoeg behaalde LCA deze resultaten zonder de zware computationele kosten van Monte Carlo sampling die vereist zijn door baselines zoals MathShepherd, SCAN of OmegaPRM.
• Test-Time Scaling (RQ2): In Best-of-N sampling experimenten op MATH-500 behaalde LCA consistent de hoogste gemiddelde scores over verschillende generator modellen (Mistral, Llama, Qwen) en steekproefgroottes ($N=8, 32, 64$).
• Ablatie-studies:
  • Temperatuur ($\alpha$): Prestaties variëren met $\alpha$. Een lagere $\alpha$ (dichter bij Max Pooling) is gunstig voor eenvoudige datasets (GSM8K), terwijl een hogere $\alpha$ helpt bij complexe datasets met veel procesfouten (OlympiadBench), wat een trade-off laat zien tussen theoretische consistentie en empirische effectiviteit.
  • Pooling-technieken: SWS pooling met $\alpha=0.2$ presteerde aanzienlijk beter dan andere MIL pooling-technieken (Instance, Embedding, Attention, Conjunctive) en benaderde de prestaties van volledig gesuperviseerde PRM's die getraind zijn op menselijke stap-niveau labels.

5. Betekenis en Claims

Het artikel claimt dat LCA een geprinciepeerd alternatief biedt voor fijnmazige proces-supervisie in complexe redeneertaken. Door de "kip-en-ei" paradox tussen credit assignment en reward modeling op te lossen via het MIL-framework en SWS pooling, produceert LCA PRM's die accuraat krediet toeschrijven aan de stappen die de uitkomst bepalen.

De auteurs benadrukken dat hun methode:

• De trainings-test kloof verkleint: Door de trainingsdoelstelling (Weakest Link) af te stemmen op de evaluatiemetriek (First Error Detection).
• Computationele overhead vermindert: Het vermijdt de dure Monte Carlo-rollouts die vereist zijn door andere outcome-supervised methoden.
• Generalisatie mogelijk maakt: De geleerde PRM's vertonen sterke generalisatie over verschillende backbone modellen en redeneerniveaus.

Het artikel concludeert bescheiden door te erkennen dat hoewel LCA beter presteert dan andere outcome-supervised methoden, het nog steeds achterblijft bij PRM's die getraind zijn met menselijke stap-niveau labels. Zij identificeren leren met ruisende labels (waarbij correcte antwoorden worden bereikt via gebrekkige redenering) en het integreren van PRM's in Reinforcement Learning (RL) als belangrijke richtingen voor toekomstig werk.