Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

1. Problématique : Le Décalage de Distribution et l'Oubli Catastrophique

L'article aborde un défi majeur dans l'adaptation des Grands Modèles de Langage (LLM) : l'oubli catastrophique lors du Supervised Fine-Tuning (SFT) sur des données de domaine spécifique.

• Le cœur du problème : Lorsque les données de downstream (domaine cible) présentent un décalage de distribution important par rapport à la distribution d'entraînement initiale du modèle, le SFT standard agit comme un signal d'apprentissage hors politique (off-policy).
• Conséquences :
  1. Instabilité de l'optimisation : Les démonstrations experts peuvent avoir une faible vraisemblance sous la politique actuelle du modèle, entraînant des mises à jour de gradient à forte variance et instables.
  2. Oubli Catastrophique : Le modèle se spécialise trop rapidement sur la distribution de downstream, au détriment de ses capacités générales acquises précédemment.
• Limites des approches existantes : Les méthodes de réécriture de données (data rewriting) précédentes tentent de combler cet écart en réécrivant les données avant le SFT. Cependant, elles souffrent de deux défauts majeurs :
  • Mauvaise alignement distributionnel : Elles génèrent des réécritures conditionnées par un prompt de réécriture spécifique ($\pi_0(\cdot | x, y^*, x_{prompt})$), ce qui ne correspond pas à la distribution naturelle de génération de type "Question-Réponse" du modèle ($\pi_0(\cdot | x)$).
  • Effondrement de la diversité : L'utilisation de templates fixes conduit à des sorties homogènes et artificielles, réduisant la richesse des données d'entraînement.

2. Méthodologie : L'Agent de Réécriture basé sur le RL

Les auteurs proposent un cadre novateur qui traite la réécriture de données non pas comme une tâche de promptage statique, mais comme un problème d'apprentissage de politique résolu par Reinforcement Learning (RL).

Architecture Globale

Le système se déroule en deux étapes principales :

1. Entraînement de l'Agent de Réécriture (Stage I) : Un agent $R_\phi$ apprend à réécrire les données.
2. Construction du Dataset et SFT (Stage II) : Utilisation des données réécrites pour affiner le modèle cible.

Détails Techniques

• Politique de Réécriture Efficace en Paramètres :

  • Le modèle de base $\pi_0$ (froid) est gelé.
  • L'agent de réécriture $R_\phi$ est paramétré via des adaptateurs LoRA (Low-Rank Adaptation) appliqués sur $\pi_0$. Cela permet un ajustement local et contrôlé sans dériver la politique globale du modèle.

• Optimisation par GRPO (Group Relative Policy Optimization) :

  • L'agent est entraîné avec GRPO, une variante de PPO qui estime la ligne de base à partir de groupes d'échantillons, éliminant le besoin d'un réseau de valeur séparé.
  • Pour chaque entrée $(x, y^*)$, l'agent génère un groupe de $K$ réécritures candidates.

• Objectif Unifié et Fonction de Récompense :
  L'optimisation vise trois objectifs simultanés, évalués automatiquement :

  1. Cohérence de la Tâche (Hard Gate) : Une vérification stricte (correctitude de la réponse + validité du raisonnement) agit comme un filtre binaire. Seules les réécritures valides reçoivent des récompenses supplémentaires.
  2. Alignement Distributionnel ($r_{dist}$) : La récompense favorise les réécritures qui sont plus probables sous la distribution de génération QA standard du modèle de base ($\pi_0(\cdot | x)$), sans le prompt de réécriture. Cela assure que les données réécrites sont "naturelles" pour le modèle.
  3. Diversité ($r_{div}$) : Une récompense basée sur la distance sémantique entre les candidats valides d'un même groupe encourage la variété des solutions, évitant l'effondrement de mode.

• Pipeline de Construction de Données :
  Un processus GENERATE-VERIFY-FALLBACK est utilisé. Si une réécriture échoue à la vérification de cohérence, le système revient à la démonstration experte originale ($y^*$) pour ne pas perdre de données d'entraînement.

3. Contributions Clés

1. Insight sur le Décalage de Distribution : Identification que les méthodes de réécriture actuelles créent un nouveau décalage (conditionné par le prompt) plutôt que de s'aligner sur la distribution naturelle de génération du modèle.
2. Cadre d'Optimisation Unifié : Introduction d'un objectif RL qui combine la cohérence de la tâche (contrainte dure), l'alignement distributionnel et la diversité, permettant de générer des données de haute qualité.
3. Agent de Réécriture Apprenant : Remplacement des heuristiques de prompting fixe par une politique apprise (LoRA + GRPO) capable de s'adapter dynamiquement pour maximiser la qualité des données.

4. Résultats Expérimentaux

Les expériences ont été menées sur plusieurs modèles (Llama-3.2-1B/3B, Mistral-7B) avec des tâches de raisonnement mathématique comme domaine de downstream.

• Performance en Downstream (Math) : La méthode proposée atteint des gains de performance comparables, voire supérieurs, au SFT standard (Vanilla SFT) sur les benchmarks mathématiques (ex: +55.23% sur Mistral-7B).
• Réduction de l'Oubli Catastrophique : C'est le résultat le plus significatif. Contrairement au SFT standard qui dégrade fortement les performances sur les benchmarks généraux (MMLU, AGIEval), la méthode proposée réduit la chute de performance de 12,34 % en moyenne.
  • Exemple : Sur Mistral-7B, la chute de généralisation passe de 19,17 % (SFT standard) à seulement 7,35 % avec l'agent de réécriture.
• Stabilité de l'Entraînement : Les courbes de perte montrent que le SFT sur les données réécrites par l'agent converge plus rapidement et vers un plateau plus bas et stable, indiquant une meilleure facilité d'ajustement (fitting) des données.
• Analyse d'Ablation :
  • Supprimer la récompense d'alignement ($r_{dist}$) dégrade fortement la rétention des capacités générales.
  • Supprimer la diversité ($r_{div}$) affecte les performances en domaine mathématique.
  • L'utilisation d'un "Hard Gate" (filtrage strict) est cruciale pour éviter le bruit dans l'apprentissage.

5. Signification et Impact

Cet article propose un changement de paradigme dans la préparation des données pour le SFT. Au lieu de simplement collecter ou filtrer des données, il apprend à réécrire activement les données pour qu'elles soient intrinsèquement compatibles avec la distribution du modèle cible.

• Stabilité vs Plasticité : La méthode résout la tension classique entre la nécessité de s'adapter à une nouvelle tâche (plasticité) et celle de conserver les connaissances antérieures (stabilité), en assurant que l'adaptation se fait via des signaux d'apprentissage "in-distribution".
• Généralité : Bien que testée sur les mathématiques, le cadre est conçu pour être applicable à d'autres domaines, offrant une voie prometteuse pour l'adaptation robuste des LLM sans sacrifier leurs capacités générales.
• Efficacité : L'approche LoRA rend la méthode légère et économiquement viable par rapport à l'entraînement complet de modèles de réécriture.

En résumé, cette méthode "patche" le décalage de distribution à la source, transformant le SFT hors politique en un processus plus stable et moins destructeur pour les capacités générales des modèles.