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

1. Problemstellung

Immer mehr autonome Agenten (basierend auf Reinforcement Learning, RL, oder Large Language Models, LLMs) interagieren mit Menschen. Um das Vertrauen und die Sicherheit in diesen Systemen zu gewährleisten, ist es entscheidend, dass diese Agenten ihre Entscheidungen in natürlicher Sprache erklären können.

Das Hauptproblem besteht darin, dass menschliches Feedback zu Erklärungen pluralistisch (verschiedene Menschen haben unterschiedliche Meinungen) und probabilistisch (nicht deterministisch) ist. Herkömmliche Methoden wie Reinforcement Learning from Human Feedback (RLHF) oder Reinforcement Learning from AI Feedback (RLAIF) nutzen oft statische oder deterministische Belohnungssignale (Rewards).

• RLHF: Ist teuer und skaliert schlecht aufgrund der Notwendigkeit menschlicher Annotationen.
• RLAIF: Nutzt Proxy-LLMs, um Belohnungen zu generieren, führt aber zu Fehlern, da diese Modelle die komplexe Verteilung menschlicher Urteile nicht exakt abbilden können. Bestehende Ansätze zur Modellierung von Verteilungs-Belohnungen (Distributional Reward Models) erfordern oft eine Diskretisierung des Feedbacks, was Informationsverluste verursacht.

Das Ziel ist es, ein Framework zu entwickeln, das LLMs trainiert, um generalisierbare, kontextabhängige Erklärungen zu generieren, die die zugrunde liegende Entscheidungsstrategie (Policy) eines Agenten korrekt widerspiegeln, unter Berücksichtigung der probabilistischen Natur menschlicher Bewertungen.

2. Methodik

Die Autoren schlagen ein Framework vor, das Reinforcement Learning from AI Feedback (RLAIF) mit generativen Continuous Normalizing Flows (CNFs), speziell Rectified Flows (RF), kombiniert, um hochwertige, verteilungsbasierte Belohnungen zu erzeugen.

Kernkomponenten:

1. Erklärungs-LLM (EXPLANATION LLM):

   • Ein LLM (z. B. Llama-3.1-8B), das trainiert wird, um basierend auf einem Kontext (z. B. Zustandsverlauf in RL oder eine Frage in QA) eine Erklärung zu generieren.
   • Wichtig: Die tatsächliche Entscheidung (der "Ground Truth") wird dem Prompt nicht mitgeteilt, um zu erzwingen, dass das Modell die Logik aus dem Kontext ableitet und nicht nur die Antwort wiederholt.

2. Rectified Flow Reward Model (RF):

   • Anstelle eines einfachen Scores nutzt das System einen generativen Fluss, der eine Basisverteilung (z. B. Gaußsches Rauschen) in eine komplexe Verteilung von Belohnungen transformiert.
   • Proxy-LLMs: Mehrere unabhängige Proxy-LLMs generieren zunächst Logits (Wahrscheinlichkeiten) für die korrekte Entscheidung basierend auf der vom Erklärungs-LLM generierten Erklärung.
   • Training des RF: Der Rectified Flow wird trainiert, um die Verteilung dieser Proxy-Logits zu approximieren. Theoretisch wird gezeigt, dass der RF in der Lage ist, das Rauschen der Proxy-LLMs zu "denoisen" und sich der wahren menschlichen Belohnungsverteilung anzunähern, sofern die Rauschverteilung der Proxy-LLMs der Startverteilung des Flows entspricht.
   • Architektur: Der RF ist in ein LLM eingebettet. Er nutzt Cross-Attention, um linguistische Hinweise aus dem Kontext und der Erklärung zu verarbeiten und diese mit dem latenten Raum des Flows zu verbinden. Dies ermöglicht es dem Modell, kontextsensitive Belohnungen zu generieren.

3. Trainingsprozess:

   • Iteratives Training: Der Erklärungs-LLM und das RF-Modell werden abwechselnd trainiert.
   • Belohnungssignal: Die Belohnung für den Erklärungs-LLM ist die Wahrscheinlichkeit (Logit), dass der Proxy-LLM (bzw. der denoisierte Flow) die korrekte Entscheidung aus der Erklärung ableitet.
   • Dichte Belohnungen: Das System berechnet Belohnungen auf Satzebene (Änderung der Wahrscheinlichkeit der korrekten Entscheidung nach jedem neuen Satz), was einen dichteren und informativeren Lernsignal liefert als eine einzelne Belohnung am Ende der Erklärung.
   • Optimierung: Der Erklärungs-LLM wird mittels PPO (Proximal Policy Optimization) mit LoRA (Low-Rank Adaptation) feinabgestimmt, um die generierten Belohnungen zu maximieren.

3. Schlüsselbeiträge

• Theoretische Garantie: Die Autoren beweisen, dass ein Rectified Flow, der auf noisy Proxy-Rewards trainiert wird, die Abweichung von der wahren menschlichen Belohnungsverteilung nach oben beschränken kann (Error Bound). Der Flow fungiert als Denoiser, der die inhärenten Fehler der Proxy-LLMs korrigiert.
• Generative Reward Modeling: Einführung von Continuous Normalizing Flows als Reward-Modelle, die kontinuierliche, probabilistische Feedback-Verteilungen modellieren können, ohne Diskretisierung.
• Kontextsensitive Architektur: Entwicklung einer speziellen RF-Architektur mit Cross-Attention, die linguistische Nuancen in Erklärungen und Kontexten berücksichtigt, um generalisierbare Belohnungen zu erzeugen.
• Generalisierung auf Negative Samples: Das Modell lernt, zwischen "wahren Negativen" (schlechte Erklärungen) und "falschen Negativen" (gute Erklärungen, die von Proxy-LLMs fälschlich abgelehnt wurden) zu unterscheiden und diese differenziert zu belohnen.

4. Ergebnisse

Die Methode wurde auf vier Benchmarks evaluiert: SMAC (Multi-Agent RL), AI2-THOR (Embodied AI), MMLU (Wissensfragen) und MathQA (Mathe-Aufgaben).

• Überlegenheit gegenüber Baselines: Die Methode übertrifft State-of-the-Art-Baselines (SFT, RLHF mit PPO/DPO/KTO, RLAIF mit Skywork, Reasoning-Frameworks wie ReFT) konsistent.
  • Auf SMAC, MMLU und MathQA wurden Verbesserungen der Genauigkeit (ACC) von 6,9 % bis 12,4 % gegenüber dem direkten Proxy-LLM-Ansatz und 1,9 % bis 18,2 % gegenüber anderen Baselines erzielt.
  • Auf dem schwierigen AI2-THOR-Benchmark wurde die Erfolgsrate um 2,5 % bis 7,4 % gesteigert.
• Human Evaluation: Menschliche Evaluatoren bewerteten die generierten Erklärungen als:
  • Logisch fundierter (höchste Scores für Logik).
  • Handlungsorientierter (Actionability).
  • Kognitiv weniger belastend (niedrigere Cognitive Load).
  • Die Genauigkeit, mit der Dritte die wahre Entscheidung aus der Erklärung ableiten konnten, war signifikant höher als bei allen Vergleichsmethoden.
• Robustheit: Die Methode zeigte sich robust gegenüber verschiedenen Backbone-Modellen (Llama, Qwen, Gemma) und auch unter Bedingungen mit verrauschten Belohnungssignalen.

5. Bedeutung und Ausblick

Dieses Paper stellt einen bedeutenden Fortschritt im Bereich der Explainable AI (XAI) und des RLHF/RLAIF dar.

• Paradigmenwechsel: Es bewegt sich weg von deterministischen Belohnungssignalen hin zu probabilistischen, generativen Modellen, die die Komplexität menschlicher Urteile besser abbilden.
• Skalierbarkeit: Durch die Nutzung von AI-Feedback und die theoretische Korrektur von Proxy-Fehlern wird der Bedarf an teuren menschlichen Annotationen reduziert, ohne an Qualität zu verlieren.
• Vertrauen und Sicherheit: Die Fähigkeit, Agentenentscheidungen durch logisch fundierte und verständliche Erklärungen zu untermauern, ist essenziell für den Einsatz von KI in sicherheitskritischen Bereichen wie Medizin, Recht und autonomen Systemen.

Zusammenfassend demonstriert das Paper, dass Flow-Matching-Modelle als Reward-Generatoren nicht nur die Trainingsleistung von Erklärungs-LLMs steigern, sondern auch theoretisch fundierte Garantien für die Qualität der generierten Erklärungen bieten.