Learning Reasoning Reward Models from Expert Demonstration via Inverse Reinforcement Learning

Titel: Learning Reasoning Reward Models from Expert Demonstration via Inverse Reinforcement Learning

Auteurs: Claudio Fanconi, Nicolás Astorga, Mihaela van der Schaar (Universiteit van Cambridge)
Publicatie: ICLR 2026 Workshop on Logical Reasoning of Large Language Models

---

1. Probleemstelling

Huidige methoden om redeneervermogen in Large Language Models (LLMs) te verbeteren, vertrouwen voornamelijk op twee benaderingen, die beide beperkingen hebben:

1. Supervised Fine-Tuning (SFT): Dit is fundamenteel imitatief. Het model leert om de getrainde "expert traces" na te bootsen. Hoewel dit effectief is zolang het model binnen het bereik van de getrainde data blijft, biedt het geen duidelijke richtlijnen wanneer het model zich in een staat bevindt die afwijkt van de getoonde voorbeelden (off-demonstration states). Het leert geen expliciete stap-voor-stap doelstelling om alternatieve redeneerpaden te scoren.
2. Reinforcement Learning (RL) met uitkomst-rewards: Deze methode vereist een goed gespecificeerde verificateur (bijv. een code-uitvoerder of een menselijke beoordelaar) die alleen een beloning geeft aan het einde van de reeks (outcome-level). Dit is vaak lastig te schalen voor complexe taken en biedt geen feedback over waar in het redeneerproces een fout is gemaakt.

Er is een behoefte aan een methode die proces-rewards (stap-voor-stap feedback) leert direct uit expert-demonstraties, zonder afhankelijk te zijn van handmatig ontworpen verificateurs, en die zowel tijdens het trainen als tijdens het infereren nuttig is.

2. Methodologie: Adversarial Inverse Reinforcement Learning (AIRL)

De auteurs stellen een Adversarial Inverse Reinforcement Learning (AIRL) framework voor om een redeneer-beloningssysteem (reward model) te leren puur uit expert-demonstraties.

Kernconcepten:

• Formulering als IRL: Het probleem wordt geformuleerd als een min-max spel. Een discriminator ($D_\phi$) probeert expert-traces te onderscheiden van traces gegenereerd door het beleid ($\pi_\theta$), terwijl het beleid probeert de discriminator te misleiden door traces te genereren die lijken op expert-traces.
• Discriminator Training:
  • In plaats van alle expert-traces als positief en alle beleid-traces als negatief te labelen (wat kan leiden tot het leren van stijlartefacten in plaats van redeneerkwaliteit), gebruiken de auteurs antwoord-overeenkomst als label.
  • Positieve set ($D_{pos}$): Expert-traces + beleid-traces die hetzelfde eindantwoord geven als de expert.
  • Negatieve set ($D_{neg}$): Beleid-traces met een ander antwoord + synthetisch verontreinigde expert-traces (om robuustheid te vergroten).
  • De discriminator leert zo om de kwaliteit van het redeneerproces te beoordelen, niet alleen de bron of het eindantwoord.
• Granulariteit van Beloningen: De auteurs evalueren drie niveaus van beloningsresolutie via een masker ($m_t$):
  1. Sparse: Alleen beloning bij het laatste token (uitkomst).
  2. Interval: Beloning op vaste intervallen.
  3. Dense: Beloning op elk token (proces-niveau).
  • Voor sparse/interval wordt de beloning "backfilled" (uitgebreid) naar de tussenliggende tokens.
• Beleidsoptimalisatie: Het taalmodel wordt geoptimaliseerd met Group Relative Policy Optimization (GRPO), aangepast voor dichte proces-beloningen. Om instabiliteit te voorkomen, worden rewards geclipt en worden replay-buffers gebruikt om de discriminator stabiel te houden.

3. Belangrijkste Bijdragen

1. Adaptatie van AIRL voor LLMs: Het toepassen van adversariele IRL om dichte redeneer-beloningen te leren uit expert-demonstraties op meerdere granulariteiten.
2. Universele Reward-Model: Het aantonen dat één enkel geleerd reward-model herbruikbaar is voor zowel training (als signaal voor beleidsoptimalisatie) als inferentie (als reranker voor gegenereerde traces).
3. Foutlokalisatie: Het vermogen van het dichte reward-model om precies aan te geven waar een redeneertraject voor het eerst afwijkt van het juiste pad, wat diagnose van fouten mogelijk maakt.

4. Resultaten

De methode is geëvalueerd op drie benchmarks: GSM8K (wiskunde), MEDREASON (medisch redeneren) en MMLU-PRO (wetenschappelijke vragen).

• D1: Trainingssignaal (Verbetering t.o.v. SFT):

  • De geleerde rewards presteren over het algemeen beter dan Supervised Fine-Tuning (SFT). In 8 van de 9 configuraties (combinaties van model en dataset) overtrof de beste AIRL-variant de SFT-baseline.
  • Sparse rewards bleken het meest stabiel en vaak het sterkst. Dense rewards bieden fijnere feedback maar zijn gevoeliger voor instabiliteit tijdens het trainen (mode collapse).
  • Voorbeeld: Op GSM8K verbeterde een Qwen2.5-7B model met sparse rewards van 70.1% (SFT) naar 85.8%.

• D2: Inferentie-tijd Assistentie (Reranking):

  • Het geleerde reward-model fungeert als een krachtige reranker. Bij een vaste sampling-begroting (Best-of-16) levert het aanzienlijke winsten op.
  • De grootste winst was +17.4 percentagepunten (van 38.4% naar 55.7%) voor Qwen2.5-7B op GSM8K met een dichte reward.
  • Het combineren van reward-reranking met meerderheidsvoting (majority voting) gaf de beste resultaten bij grotere sampling-budgetten.

• D3: Generalisatie en Foutlokalisatie:

  • Transfer: De rewards generaliseren goed over verschillende taken en modelbackbones (bijv. een op MEDREASON getrainde reward werkt ook goed op GSM8K), hoewel niet volledig invariant.
  • Lokalisatie: Het dichte reward-model kan de exacte stap identificeren waarop een fout optreedt. In visualisaties zagen de auteurs dat de reward positief bleef voor correcte tussenstappen, maar scherp daalde op het moment dat de berekening of logica foutief werd, zelfs als het eindantwoord later nog steeds probeerde correct te zijn.

5. Betekenis en Conclusie

Dit werk sluit de kloof tussen pure imitatie (SFT) en beloningsgestuurde optimalisatie (RL) voor LLM-redenering.

• Efficiëntie: Het elimineert de noodzaak van handmatig ontworpen, taak-specifieke verificateurs door beloningen direct uit expert-data te leren.
• Interpreteerbaarheid: Het biedt inzicht in het redeneerproces door fouten op token-niveau te lokaliseren, wat helpt bij het diagnosticeren van waarom een model faalt.
• Dualiteit: Hetzelfde model dient als leraar tijdens het trainen en als bewaker tijdens het infereren.

De auteurs concluderen dat hoewel er een afweging bestaat tussen de dichtheid van de beloning en de stabiliteit van de optimalisatie, het AIRL-framework een robuust en herbruikbaar middel biedt om redeneerredenen te distilleren en te optimaliseren.