Translate Policy to Language: Flow Matching Generated Rewards for LLM Explanations

Probleemstelling

Naarmate intelligente agenten (zoals Reinforcement Learning-agenten en Large Language Models) steeds meer geïntegreerd raken in menselijke omgevingen, is het cruciaal dat deze agenten hun beslissingen kunnen uitleggen in natuurlijke taal. Bestaande methoden voor het genereren van uitleg zijn vaak:

1. Taakspecifiek: Ze zijn op maat gemaakt voor specifieke toepassingen (bijv. zelfrijdende auto's of aanbevelingssystemen) en missen generaliseerbaarheid.
2. Afhankelijk van menselijke feedback (RLHF): Het verzamelen van menselijke feedback voor het trainen van uitlegmodellen is duur en tijdrovend.
3. Beperkt door AI-feedback (RLAIF): Bestaande RLAIF-methoden gebruiken vaak "proxy" LLM's om beloningen te genereren. Deze proxy-beloningen zijn echter deterministisch en houden geen rekening met de pluralistische en probabilistische aard van menselijke oordelen over wat een goede uitleg is. Mensen hebben verschillende meningen over dezelfde uitleg, wat leidt tot een complex verdelingsprobleem dat door standaardmethoden niet goed wordt gemodelleerd.

Methodologie

De auteurs stellen een nieuw framework voor dat een Explanation LLM traint via Reinforcement Learning van AI-feedback (RLAIF), waarbij de beloningen worden gegenereerd door een Generative Continuous Normalizing Flow (CNF), specifiek een Rectified Flow (RF) model.

De kerncomponenten van de methode zijn:

1. Doel: Train een LLM om een onderliggend beleid (policy) te herstellen en natuurlijke taal uitleggen te genereren die dit beleid weerspiegelt, zonder dat de daadwerkelijke beslissing van de agent in de prompt wordt verstrekt. De uitleg moet derden in staat stellen de oorspronkelijke beslissing af te leiden.
2. Proxy LLMs: Om menselijke feedback te simuleren, worden meerdere onafhankelijke "Proxy LLM's" gebruikt. Deze krijgen de context en de gegenereerde uitleg als input en voorspellen de waarschijnlijkheid van de juiste beslissing. De output is een verdeling van logit-waarden.
3. Rectified Flow Reward Model: In plaats van de logit-waarden van de proxy LLM direct als beloning te gebruiken, trainen de auteurs een Rectified Flow model ($\phi$).
   • Dit model leert de transformatie van een ruisverdeling (bijv. Gaussisch) naar de complexe verdeling van de proxy-LLM-logits.
   • Theoretisch voordeel: Het paper toont aan dat als de ruis in de proxy-LLM-output dezelfde functionele vorm heeft als de basisverdeling van de flow, het flow-model de afwijking van de ware menselijke beloningsverdeling wiskundig kan begrenzen. Het flow-model fungeert als een "denoiser" die de onderliggende menselijke voorkeursverdeling herstelt uit de noisy proxy-data.
4. Architectuur: Het flow-model is ingebed in een LLM-architectuur. Het gebruikt Cross-Attention om linguïstische aanwijzingen uit de context en de uitleg te integreren bij het genereren van de beloningsverdeling. Dit zorgt ervoor dat het model contextueel bewust is.
5. Trainingscyclus:
   • De Explanation LLM genereert uitleg.
   • Proxy LLM's evalueren deze uitleg.
   • Het Rectified Flow model leert de verdeling van deze evaluaties.
   • De Explanation LLM wordt getraind met PPO (Proximal Policy Optimization) om de door het flow-model gegenereerde beloningen te maximaliseren.
   • Dit proces wordt iteratief herhaald.

Belangrijkste Bijdragen

• Generaliseerbaar Framework: Een taak-agnostische methode voor het genereren van uitleg die werkt voor zowel RL-agenten als LLM-taken.
• Flow Matching voor Beloningen: De eerste toepassing van generatieve flow-modellen (Rectified Flow) als reward-modellen in RLAIF. Dit modelleert expliciet de distributie van beloningen in plaats van een enkel punt, wat beter aansluit bij de pluraliteit van menselijke oordelen.
• Theoretische Garantie: Een wiskundige bewering dat het gebruik van een flow-model de foutmarge ten opzichte van de ware menselijke verdeling beperkt, mits de ruisstructuur overeenkomt met de basisverdeling van de flow.
• Cross-Attention Mechanisme: Een specifieke architectuur die het flow-model in staat stelt om contextuele en linguïstische signalen selectief te integreren bij het toekennen van beloningen.

Resultaten

De methode is geëvalueerd op vier benchmarks:

1. SMAC: Multi-agent strategisch spel (StarCraft II).
2. AI2-THOR: Embodied AI met multi-stap planning.
3. MMLU: Professionele kennisvragen (Geneeskunde, Recht, Boekhouding, Psychologie).
4. MathQA: Wiskundige woordproblemen.

Kernbevindingen:

• Prestatie: De methode overtreft alle baselines (inclusief SFT, RLHF-methoden zoals PPO/DPO/KTO, RLAIF-baselines zoals Skywork, en redeneringsframeworks zoals ReFT) met 1,9% tot 18,2% in nauwkeurigheid (ACC) of success rate.
• Effectiviteit van Flow: Het verwijderen van de Rectified Flow en het direct gebruiken van proxy-LLM beloningen resulteerde in een prestatiedaling van 6,9% tot 12,4%, wat aantoont dat het genereren van distributiebeloningen essentieel is.
• Menselijke Evaluatie: Menselijke beoordelaars beoordeelden de gegenereerde uitleggen als:
  • Meer accuraat in het voorspellen van de agent's beslissing.
  • Logischer en meer handelingsgericht (actionable).
  • Minder cognitieve belasting (cognitive load) voor de lezer.
• Generalisatie naar Negatieve Samples: Het model leert onderscheid te maken tussen "ware negatieven" (slechte uitleg) en "valse negatieven" (goede uitleg die door de proxy LLM verkeerd is ingeschat). Het flow-model corrigeert deze fouten en geeft juiste beloningen, wat leidt tot een verbetering van de Explanation LLM.

Betekenis en Impact

Dit paper biedt een doorbraak in het veld van Explainable AI (XAI) en Reinforcement Learning. Door de probabilistische aard van menselijke oordelen te modelleren via generatieve flows, lost het een fundamenteel probleem op in RLAIF: de beperking van deterministische beloningsmodellen.

De methode maakt het mogelijk om betrouwbare, interpreteerbare uitleggen te genereren voor complexe agenten zonder de hoge kosten van menselijke annotatie. Dit is van groot belang voor het opbouwen van vertrouwen in AI-systemen in kritieke domeinen zoals gezondheidszorg, financiën en autonome systemen, waar transparantie en de mogelijkheid voor menselijke interventie essentieel zijn. De theoretische onderbouwing biedt bovendien een nieuwe richting voor het ontwerpen van robuuste reward-modellen die minder gevoelig zijn voor ruis in proxy-data.