Résumé Technique : Apprentissage par Renforcement par Feedback Humain Distributionnellement Robuste

Énoncé du Problème

L'apprentissage par renforcement à partir de feedback humain (RLHF) est devenu une technique standard pour aligner les grands modèles de langage (LLM) avec les intentions humaines. Cependant, les pipelines RLHF existants souffrent d'un manque de robustesse. Lorsque la distribution des prompts rencontrés lors de l'inférence (tâches en aval) diffère significativement de la distribution de l'ensemble de données de préférences utilisé lors de l'ajustement fin (fine-tuning), les performances du modèle se détériorent souvent. Ce problème est exacerbé par des phénomènes tels que la sur-optimisation, où les politiques améliorent les récompenses proxies au détriment des fonctions de récompense réelles, et par les décalages de distribution causés par des changements d'environnements ou l'évolution des préférences humaines. Les approches antérieures de la robustesse se sont concentrées sur des modèles de corruption adverses spécifiques, des décalages de modèles de préférence, ou des méthodes heuristiques comme la fusion de modèles, qui échouent souvent face à de légers décalages distributionnels.

Méthodologie

Les auteurs proposent un cadre d'Optimisation Distributionnellement Robuste (DRO) pour le RLHF. L'objectif central est de garantir que les modèles ajustés (tant les modèles de récompense que les politiques) maintiennent leurs performances même lorsque la distribution des prompts change à l'intérieur d'une distance bornée de la distribution d'entraînement. Les auteurs formulent cela en introduisant un problème d'optimisation minimax où l'apprenant minimise la perte face à la pire distribution dans une boule de $\phi$-divergence (spécifiquement la distance de Variation Totale) autour de la distribution empirique d'entraînement.

Le papier développe des variantes robustes pour deux pipelines RLHF principaux :

1. Estimation de Récompense Robuste :

   • Formulation : Le modèle de récompense $r$ est entraîné pour minimiser l'espérance de la pire perte sur des distributions $D$ telles que $d_{TV}(D, D_{src}) \leq \rho$.
   • Algorithme : Un algorithme de descente de gradient par mini-lots (Algorithm 1) est proposé. À chaque itération, l'algorithme échantillonne un mini-lot et calcule les poids optimaux $q^*_t$ correspondant au décalage de distribution le plus défavorable dans une distance $\rho$. Le gradient est ensuite repondéré par ces poids optimaux avant la mise à jour des paramètres de la récompense.
   • Convergence : Sous l'hypothèse d'une classe de fonctions de récompense linéaires, les auteurs prouvent que l'algorithme converge vers une solution $\epsilon$-optimale avec $T = O(1/\epsilon^2)$ itérations et $n = \tilde{O}(1/\epsilon^2)$ échantillons par itération.

2. Optimisation de Politique Robuste :

   • Formulation : La politique $\pi$ est optimisée pour maximiser la fonction de valeur KL-régularisée la plus défavorable.
   • Algorithme : Les auteurs proposent une méthode de Gradient de Politique Naturelle (NPG) robuste (Algorithm 2). Similairement à l'estimation de la récompense, elle calcule les poids les plus défavorables pour le mini-lot. Crucialement, elle met à jour la politique en utilisant une matrice d'information de Fisher pondérée dérivée de ces poids les plus défavorables, plutôt que de la matrice de Fisher empirique standard.
   • Convergence : En supposant une classe de politiques log-linéaires et des conditions de concentrabilité standards, la méthode atteint un taux de convergence linéaire ($T = O(\log(1/\epsilon))$) mais nécessite une complexité d'échantillonnage plus élevée de $n = \tilde{O}(1/\epsilon^4)$ par itération pour assurer la convergence du dernier itérant.

3. Optimisation Directe de Préférence (DPO) Robuste :

   • Formulation : Une version distributionnellement robuste de DPO est formulée, optimisant directement la politique sans un modèle de récompense explicite.
   • Algorithme : L'algorithme 3 applique la descente de gradient stochastique par mini-lots avec des gradients repondérés basés sur le décalage de distribution le plus défavorable.
   • Convergence : L'analyse montre une complexité d'itération de $T = O(1/\epsilon^2)$ et une complexité d'échantillonnage de $n = \tilde{O}(1/\epsilon^2)$, similaire à la phase d'estimation de la récompense robuste.

Contributions Clés

• Cadre Théorique : Le papier établit les premières formulations de robustesse distributionnelle pour les deux phases d'estimation de la récompense et d'optimisation de la politique du RLHF, ainsi que pour le DPO.
• Algorithmes : Il introduit des algorithmes pratiques, basés sur des mini-lots, pour l'Estimation de Récompense Robuste, le NPG Robuste et le DPO Robuste. Ces algorithmes sont conçus pour être minimaux dans leur intrusion aux pipelines existants, ne nécessitant que le calcul des facteurs de repondération (résolu via l'optimisation convexe) et la repondération des gradients.
• Garanties de Convergence : Les auteurs fournissent des preuves théoriques de la convergence de ces algorithmes. Ils établissent des bornes sur la complexité d'échantillonnage et d'itération pour les modèles de récompense linéaires et les politiques log-linéaires, démontant que la robustesse peut être obtenue avec des garanties prouvables.
• Validation Empirique : Les méthodes sont évaluées sur des tâches hors-distribution (OOD) en utilisant le jeu de données Unified-Feedback pour l'entraînement et RewardBench, HHH-Alignment, et MT-Bench pour l'évaluation.

Résultats Expérimentaux

Les auteurs ont évalué leurs méthodes en utilisant les modèles de base Gemma-2B-it et Mistral-7B-Instruct-v0.2.

• Modèles de Récompense : Sur le benchmark RewardBench, les modèles de récompense distributionnellement robustes (DR RM) ont montré une précision moyenne améliorée par rapport aux modèles Bradley-Terry (BTL) standards, particulièrement lorsque le paramètre de robustesse $\rho$ était fixé à 0,1 ou 0,2. Notamment, la performance sur le sous-ensemble Reasoning (Raisonnement) s'est nettement améliorée (ex: +23,4 % sur un sous-ensemble avec $\rho=0,1$). Cependant, la performance a décliné lorsque $\rho$ était trop élevé (ex: $\rho \geq 0,6$), indiquant un compromis entre robustesse et ajustement à la distribution d'entraînement.
• Optimisation de Politique (PPO) : Lors de l'utilisation des modèles de récompense robustes pour entraîner des politiques via PPO, les variantes robustes (DR PPO) ont généralement égalé ou dépassé les performances du PPO standard sur les jeux de données en-distribution (Unified-Feedback) et hors-distribution (HHH-Alignment, MT-Bench). Une amélioration marquée a été observée dans les tâches de raisonnement, avec une augmentation de +10,0 % de la précision pour $\rho=0,1$.
• Optimisation Directe de Préférence (DPO) : Le DPO robuste (DR DPO) a systématiquement amélioré les performances sur les jeux de données OOD par rapport au DPO standard. Pour HHH-Alignment, la meilleure performance a été atteinte à $\rho=0,8$, et pour MT-Bench à $\rho=0,5$. Le DR DPO a généralement surpassé le DR PPO sur les tâches OOD.

Signification et Revendications

Le papier affirme que l'optimisation distributionnellement robuste fournit une approche fondée sur des principes pour améliorer la généralisation des LLM aux tâches inédites et aux décalages de distribution. Les auteurs soulignent que leurs méthodes nécessitent peu de modifications aux pipelines RLHF efficaces existants (par exemple, ceux utilisant les bibliothèques Hugging Face) tout en offrant des garanties de convergence théoriques.

Une découverte clé mise en évidence par les auteurs est l'amélioration inattendue des capacités de raisonnement sur les jeux de données OOD lorsque la robustesse distributionnelle est imposée. Les auteurs notent qu'il n'est pas immédiatement clair pourquoi le DRO devrait affecter positivement le raisonnement, suggérant que ce phénomène mérite une investigation plus approfondie sur les liens entre raisonnement et robustesse distributionnelle.

Ce travail se positionne comme une étape vers un alignement véritable des LLM avec les intentions humaines dans des environnements réels dynamiques où les distributions de données ne sont pas statiques, allant au-delà des limites du RLHF standard qui échoue souvent à généraliser au-delà des exemples de préférences spécifiques vus durant l'entraînement.