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

1. Problématique

L'article aborde le défi de la transparence et de l'interprétabilité des agents intelligents (RL et LLM) qui interagissent de plus en plus avec les humains. L'objectif est de former un LLM d'explication capable de générer des justifications en langage naturel qui reflètent fidèlement la politique sous-jacente d'un agent, sans avoir accès à la décision finale de l'agent lors de la génération (le LLM doit inférer la décision uniquement à partir du contexte).

Les défis majeurs identifiés sont :

• Pluralité et Probabilité des jugements humains : Les explications sont subjectives ; différents humains peuvent juger une même explication différemment. Les méthodes actuelles (RLHF, RLAIF) traitent souvent les récompenses comme des valeurs déterministes, ignorant cette distribution complexe.
• Bruit des récompenses par proxy : L'utilisation de LLMs comme "proxy" pour générer des récompenses (RLAIF) introduit du bruit et des biais par rapport aux véritables préférences humaines.
• Généralisation : Il existe un manque de méthodes agnostiques à la tâche capables de fonctionner aussi bien sur des tâches de RL (comme StarCraft) que sur des tâches de raisonnement LLM (comme MMLU).

2. Méthodologie

Les auteurs proposent un cadre général utilisant l'apprentissage par renforcement à partir de feedback d'IA (RLAIF), mais amélioré par des Flux Normaux Continus (CNF) et spécifiquement des Flux Rectifiés (Rectified Flow - RF).

A. Architecture Globale

Le système repose sur deux composants principaux entraînés de manière itérative :

1. LLM d'Explication (EXPLANATION LLM) : Génère des explications basées uniquement sur le contexte (trajectoires d'états/actions pour le RL, ou questions pour les LLM), sans connaître la décision réelle.
2. Modèle de Récompense Génératif (Rectified Flow) : Un modèle de flux qui apprend à générer des distributions de récompenses probabilistes plutôt que des scores scalaires fixes.

B. Le Modèle de Flux Rectifié (Rectified Flow Reward Model)

Au lieu d'utiliser un modèle de récompense déterministe classique, les auteurs utilisent un Flux Rectifié pour mapper un bruit initial (distribution de base) vers une distribution complexe de récompenses.

• Entraînement du Flux : Le modèle est entraîné sur des échantillons générés par plusieurs PROXY LLMs. Ces LLMs reçoivent le contexte et l'explication, puis tentent de prédire la décision. Les logits (probabilités) de ces prédictions servent d'échantillons de récompense.
• Mécanisme d'Attention Croisée (Cross-Attention) : L'architecture du flux intègre un mécanisme d'attention croisée qui permet au modèle de récompense de "lire" les indices linguistiques du contexte et de l'explication générée. Cela permet au modèle d'ajuster la distribution de récompense en fonction de la qualité de l'explication.
• Gestion du Bruit (Théorie) : L'article démontre théoriquement que si le bruit introduit par les LLMs proxy suit la même forme fonctionnelle que la distribution de base du flux (ex: Gaussienne), le Flux Rectifié peut bornner mathématiquement l'écart entre la distribution de récompense apprise et la véritable distribution de récompense humaine. Le flux agit essentiellement comme un débruiteur inversant le processus de corruption du proxy.

C. Procédure d'Entraînement

1. Génération de Données : Un LLM proxy (o1-mini) génère des paires d'explications "préférées" (permettant d'inférer la bonne décision) et "rejetées".
2. Entraînement du Flux : Le modèle de flux est entraîné pour reproduire la distribution des logits des PROXY LLMs sur les explications préférées.
3. Optimisation du LLM d'Explication : Le LLM d'explication est entraîné via PPO (Proximal Policy Optimization). La récompense utilisée n'est pas un score unique, mais la probabilité (logit) que le modèle de flux attribue à la décision réelle, calculée sentence par sentence (récompenses denses).
4. Itération : Le processus alterne entre l'amélioration du modèle de flux et celle du LLM d'explication.

3. Contributions Clés

• Modélisation Distributionnelle des Récompenses : Introduction de Flux Rectifiés pour générer des récompenses probabilistes, capturant la nature plurielle des jugements humains, contrairement aux approches déterministes classiques.
• Garantie Théorique de Robustesse : Démonstration que le modèle de flux peut corriger les erreurs systématiques des LLMs proxy, à condition que le bruit soit modélisable, offrant une borne d'erreur sur la récupération de la distribution humaine.
• Architecture Hybride : Intégration réussie d'un modèle génératif (Flux) dans un cadre LLM via des mécanismes d'attention croisée, permettant une adaptation contextuelle fine des récompenses.
• Approche Agnostique : Une méthode unique applicable aussi bien aux agents de RL (SMAC, AI2-THOR) qu'aux tâches de compréhension de langage (MMLU, MathQA).

4. Résultats Expérimentaux

Les expériences ont été menées sur quatre benchmarks : SMAC (Stratégie multi-agent), AI2-THOR (Robotique/Planification), MMLU (Connaissances professionnelles) et MathQA (Problèmes mathématiques).

• Performance Supérieure : La méthode proposée surpasse systématiquement les bases de référence (SFT, RLHF avec PPO/DPO/KTO, RLAIF standard, et des modèles de raisonnement comme ReFT).
  • Gain de 6,9% à 12,4% par rapport à l'utilisation directe de récompenses de LLM proxy.
  • Gain de 1,9% à 18,2% par rapport aux meilleures méthodes de l'état de l'art (SOTA) sur tous les benchmarks.
• Évaluation Humaine : Les explications générées par la méthode obtiennent les scores les plus élevés en termes de :
  • Précision de prédiction : Les humains infèrent plus correctement la décision de l'agent.
  • Cohérence logique : Meilleure structure argumentative.
  • Actionnabilité : Plus faciles à utiliser pour prendre des décisions.
  • Charge cognitive : Moins d'effort mental requis pour les évaluateurs.
• Robustesse aux Échantillons Négatifs : Le modèle apprend à distinguer les "vrais négatifs" (mauvaises explications) des "faux négatifs" (bonnes explications mal classées par le proxy), attribuant des récompenses appropriées même aux cas où le proxy échoue.

5. Signification et Impact

Cet article représente une avancée significative dans le domaine de l'IA explicable (XAI) et de l'alignement des modèles de langage.

• Au-delà du déterminisme : Il déplace le paradigme de l'alignement vers une modélisation distributionnelle, reconnaissant que la "vérité" humaine est souvent une distribution de préférences et non un point unique.
• Efficacité et Coût : En utilisant des LLMs proxy corrigés par un flux génératif, la méthode réduit la dépendance coûteuse à l'annotation humaine massive tout en maintenant, voire en améliorant, la qualité de l'alignement.
• Confiance et Sécurité : En produisant des explications plus fiables et logiques, cette méthode renforce la confiance des humains dans les systèmes d'IA autonomes, un élément crucial pour leur déploiement dans des domaines sensibles (santé, finance, robotique).

En résumé, cette recherche propose une solution élégante et théoriquement fondée pour transformer des politiques d'agents opaques en explications naturelles fiables, en utilisant la puissance des modèles génératifs de flux pour maîtriser l'incertitude inhérente aux jugements humains.