Technische Samenvatting: rePIRL: Leer PRM met Inverse RL voor LLM Redeneren

Probleemstelling

In het redeneren van Large Language Models (LLM's) worden Process Reward Models (PRM's) gebruikt om token-niveau tussenliggende beloningen te bieden, met als doel de trainingsefficiëntie te verbeteren, variantie te verminderen en reward hacking te voorkomen. Bestaande methoden voor het leren van effectieve PRM's kampen echter met aanzienlijke beperkingen:

• Offline Methoden: Gesuperviseerde benaderingen (bijv. Uesato et al., 2022; Lightman et al., 2023) en DPO-achtige werken vereisen toegang tot token-niveau annotaties of voorkeursetiketten van expert-trajecten, die kostbaar zijn om te verkrijgen. Op Monte Carlo Tree Search (MCTS) gebaseerde methoden (bijv. Wang et al., 2024) vereisen het draaien van een expert-policy om trajecten te genereren, wat computationeel duur is en onhaalbaar is als de expert-policy niet beschikbaar is.
• Online Methoden: Benaderingen zoals PRIME (Cui et al., 2025a) vertrouwen op sterke aannames met betrekking tot de beloningsfuncties en het te trainen policy-model. Intrinsieke beloningsmethoden (bijv. Zhao et al., 2025b) gebruiken modelentropie of vertrouwen als beloningen, maar lijden aan "entropie-instorting" (entropy collapse), waarbij de prestaties in latere trainingsfasen verslechteren door excessief mode-seeking gedrag.

De kernuitdaging is het leren van een effectieve PRM met minimale aannames: specifiek, zonder dat er token-niveau beloningsannotaties, voorkeursetiketten, toegang tot de expert-policy of sterke aannames over de vorm van de beloningsfunctie nodig zijn.

Methodologie: rePIRL Framework

De auteurs stellen rePIRL voor, een framework geïnspireerd door Inverse Reinforcement Learning (IRL) dat PRM's leert van enkel expert-trajecten. De methode modelleert multi-step LLM-redeneren als een token-niveau Markov Decision Process (MDP) en maakt gebruik van een duaal leerproces dat de policy en de PRM afwisselend bijwerkt.

1. Theoretische Fundering

• Constructie van het Doel: Volgens het klassieke IRL-framework definiëren de auteurs een verborgen variabele $o_t \in \{0, 1\}$ die aangeeft of een actie is gesampled uit de expert-policy. Ze modelleren de waarschijnlijkheid dat een actie door een expert is gesampled met behulp van een energiegebaseerde functie: $p(o_t=1|s_t, a_t) \propto \exp(r_\phi(s_t, a_t))$.
• Likelihood Maximalisatie: De beloningsfunctie $r_\phi$ wordt geleerd door de likelihood van expert-trajecten te maximaliseren. Dit leidt tot een doelstelling die een partitiefunctie $z(\phi)$ bevat, die onhandelbaar is voor LLM's vanwege de enorme state- en actieruimtes.
• Hanteerbaarheid via Sampling: Om de onhandelbaarheid van de partitiefunctie op te lossen, passen de auteurs een gradiënttruc en importance sampling toe. In plaats van te samplen uit de optimale policy (die onbekend is), samplen ze trajecten $\bar{\tau}$ uit de huidige policy $\pi_\theta$. De gradiënt wordt benaderd met behulp van importance weights $w = \exp(r_\phi(\bar{\tau})) / \pi_\theta(\bar{\tau})$.
• Policy Optimalisatie: Om modelinstorting (deterministische policies) te voorkomen, wordt de policy bijgewerkt met een Maximum Entropy RL doelstelling (Haarnoja et al., 2018), die beloningsmaximalisatie balanceert met entropie-regularisatie. Dit wordt opgelost met algoritmen zoals RLOO, PPO of GRPO.

2. Leeralgoritme

De trainingslus wisselt af tussen het bijwerken van de PRM en de policy:

1. Rollout: Genereer policy rollouts en haal expert-demonstraties op.
2. Beloningsberekening: Bereken outcome rewards en gemiddelde process rewards.
3. PRM Update: Update het beloningsmodel met behulp van een loss-functie die expert-trajecten contrasteert met policy rollouts, gewogen door importance sampling scores. Cruciaal is dat de auteurs gemiddelde process rewards gebruiken in plaats van cumulatieve beloningen om reward hacking te mitigeren.
4. Policy Update: Update het policy-model met behulp van de geleerde process rewards gecombineerd met outcome rewards (indien beschikbaar), gebruikmakend van maximum entropy doelstellingen.

3. Unificatie van Bestaande Methoden

Een belangrijke theoretische bijdrage is het aantonen dat rePIRL diverse bestaande PRM-leermethoden (DPO, DQO, Math-Shepherd, PRIME) kan verenigen onder één enkel framework door specifieke aanvullende aannames te introduceren. Dit positioneert rePIRL als de meest algemene methode, die de minste aannames vereist.

Belangrijkste Bijdragen

1. Nieuw Framework: Introductie van rePIRL, een IRL-geïnspireerd framework voor het leren van PRM's uit expert-trajecten zonder token-niveau annotaties, voorkeursetiketten of toegang tot de expert-policy.
2. Verenigde Theorie: Een theoretische analyse die aantoont dat rePirl diverse state-of-the-art offline en online PRM-methoden (inclusief DPO, DQO, PRIME en MCTS-gebaseerde benaderingen) verenigt onder een enkel leerframework, wat hun onderliggende aannames verheldert.
3. Schaalbaar Algoritme: Ontwikkeling van een schaalbaar leeralgoritme afgestemd op LLM's, dat importance sampling gebruikt om onhandelbare partitiefuncties te behandelen en maximum entropy RL om de stabiliteit van de training te waarborgen.
4. Empirische Validatie: Uitgebreide evaluatie op wiskunde- en coding-benchmarks die superieure prestaties laten zien ten opzichte van baselines.

Experimentele Resultaten

De auteurs evalueerden rePIRL met behulp van Qwen2.5-3B-Instruct en Qwen3-4B-Base op zeven standaard benchmarks (AIME-2024, AMC, Math-500, MinervaMath, OlympiadBench, Leetcode, LiveCodeBench).

• Prestaties: rePIRL behaalde state-of-the-art resultaten onder de 3B- en 4B-schaal modellen. Zo bereikte het op Qwen3-4B-Base een gemiddelde nauwkeurigheid van 41,6% op wiskunde-benchmarks en 39,8% op coding-benchmarks, waarmee het sterke baselines zoals PRIME, MCTS en vanilla RLOO overtrof.
• Vergelijking: Het vertoonde absolute verbeteringen van 3,2% (wiskunde) en 4,3% (coding) ten opzichte van vanilla RLOO met Qwen3-4B-Base.
• Toepassingen:
  • Test-time Training: Wanneer getraind met enkel PRM-signalen zonder outcome rewards, evenaarde rePIRL de prestaties van modellen die getraind zijn met verifieerbare outcome rewards.
  • Test-time Scaling: Het gebruik van de PRM als verifieerder voor het selecteren van betere rollouts leverde een monotone prestatieverbetering op, wat beter presteerde dan majority voting.
  • Moeilijke Problemen: De PRM bood informatieve vroege signalen voor moeilijke problemen waar outcome rewards schaars waren, wat leidde tot snellere convergentie vergeleken met outcome-only training.
• Efficiëntie: rePIRL introduceerde slechts een ~1,4x trainingstijd overhead vergeleken met outcome-gebaseerde RL (bijv. RLOO), terwijl het de noodzaak voor aparte value networks vermeed.

Betekenis en Claims

Het artikel beweert het eerste werk te zijn dat Inverse Reinforcement Learning generaliseert voor LLM-redeneren en de eerste die bestaande PRM-leermethoden in één framework verenigt.

• Minimale Aannames: De auteurs benadrukken dat rePIRL opereert onder de minimale noodzakelijke aannames, waardoor het algemeen inzetbaar is dan methoden die token-niveau annotaties of specifieke toegang tot de expert-policy vereisen.
• Stabiliteit: Het framework pakt het probleem van entropie-instorting aan dat veel voorkomt bij intrinsieke beloningsmethoden door gebruik te maken van maximum entropy RL, wat zorgt voor stabiele trainingsdynamiek.
• Nut: De getrainde PRM wordt gepresenteerd als een herbruikbare, op zichzelf staande module die voordelen biedt voor downstream toepassingen zoals test-time training en scaling, vooral in scenario's waar outcome rewards niet beschikbaar of schaars zijn.

De auteurs concluderen dat hoewel rePIRL een verenigd theoretisch perspectief en praktische bruikbaarheid toevoegt, toekomstig werk de mogelijkheid kan verkennen om het te combineren met intrinsieke beloningsmethoden, systematische vergelijkingen van online RL-heuristieken uit te voeren, en op te schalen naar agentic taken en grotere modellen.