Résumé Technique : Optimisation de Politique Robuste pour Prévenir l'Oubli Catastrophique

Énoncé du Problème

Les grands modèles de langage (LLM) sont généralement entraînés via des pipelines multi-étapes, commençant par l'apprentissage par renforcement à partir de retours humains (RLHF) pour aligner les modèles sur la sécurité et l'utilité, suivi d'un ajustement fin pour des tâches spécifiques en aval (par exemple, raisonnement mathématique, codage). Une vulnérabilité critique de ce paradigme est l'oubli catastrophique : même des mises à jour mineures en aval peuvent dégrader les comportements appris précédemment, tels que les garde-fous de sécurité.

Les recherches antérieures se sont principalement concentrées sur des stratégies d'atténuation au moment de l'ajustement fin en aval (par exemple, la répétition, l'ajustement fin efficace en paramètres, la régularisation ou la fusion de modèles) pour préserver les capacités a posteriori. Cependant, les auteurs soutiennent que ces méthodes sont insuffisantes car elles dépendent de procédures spécifiques en aval. Si un modèle est ultérieurement soumis à un ajustement fin « ordinaire » qui s'écarte de ces protocoles spécifiques, les comportements appris précédemment restent fragiles. Le défi central est de garantir que la politique de base elle-même soit robuste aux adaptations futures, indépendamment de la tâche ou de l'algorithme utilisés en aval.

Méthodologie : Optimisation de Politique Robuste pour l'Ajustement Fin (FRPO)

Les auteurs proposent l'Optimisation de Politique Robuste pour l'Ajustement Fin (FRPO), un cadre conçu pour optimiser la politique de base en vue de la stabilité face aux futurs changements en aval.

1. Formulation du Problème : Robustesse Max-Min

Les auteurs formalisent le problème comme une optimisation max-min. Au lieu de maximiser la récompense espérée uniquement pour la politique actuelle $\pi_\theta$, l'objectif vise à maximiser la récompense minimale réalisable par n'importe quelle politique en aval $Q$ située dans un voisinage borné par la divergence KL de $\pi_\theta$.

Le problème d'optimisation est défini comme suit :
$$\max_{\pi_\theta} \left( \inf_{Q: \mathbb{E}[\text{KL}(Q \parallel \pi_\theta)] \le \rho} \mathbb{E}_{x,y \sim Q}[r(x,y)] - \beta \text{KL}(\pi_\theta \parallel \pi_{\text{ref}}) \right)$$

Ici, $\rho$ représente le rayon du voisinage KL, reflétant les contraintes standard dans l'ajustement fin en aval (par exemple, régions de confiance, faibles taux d'apprentissage ou contraintes LoRA). La contrainte est imposée en espérance sur la distribution des invites, permettant des changements significatifs sur les invites pertinentes pour la tâche tout en maintenant une proximité moyenne avec la politique de base.

2. Dérivation Théorique

En utilisant la théorie de la dualité pour l'optimisation robuste distributionnelle (spécifiquement sous une contrainte de $f$-divergence), les auteurs dérivent une forme duale du problème de minimisation interne.

• Objectif de Risque Entropique : La formulation duale révèle que l'objectif robuste est équivalent à l'optimisation d'un objectif de risque entropique. La minimisation interne sur $Q$ transforme le problème en maximisation de :
  $$-\lambda \log \mathbb{E}_{y \sim \pi_\theta} \left[ e^{-r(x,y)/\lambda} \right]$$
  où $\lambda$ est un hyperparamètre ajustable contrôlant la sensibilité au risque.
• Interprétation : Cet objectif pénalise davantage les trajectoires à faible récompense que la maximisation standard de la récompense espérée. Il cherche efficacement des régions « plates » dans le paysage de récompense au sein de l'espace des politiques, garantissant que de petites perturbations (ajustement fin en aval) ne conduisent pas à des chutes brutales de performance.
• Cas Limite : Lorsque $\lambda \to \infty$, l'objectif retrouve l'objectif standard GRPO (Optimisation de Politique Relative de Groupe), qui maximise la récompense espérée.

3. Implémentation Algorithmique

Les auteurs adaptent l'algorithme GRPO pour optimiser cet nouvel objectif sans introduire de surcharge computationnelle supplémentaire (par exemple, aucun critique appris ou passes avant supplémentaires).

• Échantillonnage d'Importance : L'espérance sous la nouvelle politique est estimée à l'aide d'échantillons provenant d'une politique décalée $\pi_{\text{old}}$ via l'échantillonnage d'importance.
• Recadrage et Lignes de Base : Pour assurer la stabilité, l'algorithme utilise un recadrage de style PPO sur les rapports d'importance. Crucialement, les auteurs dérivent un terme de ligne de base pour annuler la variance élevée dans le gradient de la fonction de partition logarithmique lorsque $\lambda$ est grand.
• Réduction du Biais : L'estimation Monte Carlo de la fonction de partition logarithmique introduit un biais de $O(1/G)$ (où $G$ est la taille du groupe). Les auteurs emploient une technique de jackknife (estimation leave-one-out) pour réduire ce biais à $O(1/G^2)$, assurant une convergence stable même avec de petites tailles de groupe.

Contributions Clés

1. Cadre RLHF Robuste : L'article introduit une formulation max-min qui prend en compte tous les changements en aval dans une boule KL. Il démontre que cette formulation est mathématiquement équivalente à l'optimisation d'un objectif de risque entropique, fournissant une base théorique pour la platitude de la récompense dans l'espace des politiques.
2. Algorithme FRPO : Un remplacement pratique et direct pour GRPO ne nécessitant aucun calcul supplémentaire. Il s'intègre de manière transparente aux pipelines RLHF existants, inclut une ligne de base dérivée pour la stabilité, et retrouve le GRPO standard lorsque $\lambda \to \infty$.
3. Validation Empirique : La méthode est évaluée dans divers scénarios, y compris l'alignement de sécurité et le raisonnement mathématique, démontrant une préservation supérieure des capacités antérieures par rapport à l'ajustement fin standard et à d'autres bases de robustesse.

Résultats Expérimentaux

Les auteurs évaluent FRPO sur deux scénarios principaux : Entraînement à la Sécurité et Raisonnement Mathématique.

1. Entraînement à la Sécurité (Garde-fous de Sécurité)

• Configuration : Des modèles (Mistral-v0.1 et Qwen2.5) ont été entraînés sur un jeu de données de sécurité utilisant FRPO avec des $\lambda$ variables, puis ajustés finement sur des tâches en aval (Alpaca pour le suivi d'instructions, GSM8K pour les mathématiques, et UltraFeedback pour le RL).
• Constats :
  • Les modèles entraînés avec FRPO ($\lambda \approx 0,2$) ont maintenu des taux de refus de sécurité significativement plus élevés et des scores de nocivité plus faibles (StrongREJECT) après l'ajustement fin par rapport au GRPO standard.
  • La méthode a efficacement préservé les garde-fous de sécurité même dans des paramètres d'ajustement fin agressifs (SFT à paramètres complets, divers taux d'apprentissage, LoRA).
  • Analyse des Compromis : Bien qu'un $\lambda$ plus petit ait mieux préservé la sécurité, un $\lambda$ plus grand a légèrement amélioré l'utilité au détriment de la sécurité. Le $\lambda$ optimal s'est avéré dépendre de la tâche mais généralement faible pour la sécurité.
  • Performance en Aval : Crucialement, FRPO n'a pas compromis la performance des tâches en aval (par exemple, précision GSM8K, scores MMLU) par rapport au GRPO, confirmant que la rétention de la sécurité n'était pas due à un échec à apprendre la nouvelle tâche.

2. Apprentissage Continu (Mathématiques vers Code)

• Configuration : Les modèles ont d'abord été entraînés sur le raisonnement mathématique (jeu de données MATH) puis ajustés finement sur la génération de code (OpenCodeInstruct).
• Constats :
  • FRPO avec $\lambda = 2,0$ (optimisé pour la variance plus élevée des récompenses mathématiques) a préservé la précision mathématique de manière significativement meilleure que le GRPO.
  • Plus précisément, FRPO a maintenu une précision 22 % plus élevée sur le benchmark MATH500 après l'ajustement fin pour le code par rapport au GRPO, tout en obtenant des gains comparables sur les benchmarks de codage (MBPP+).

3. Études d'Ablation

• Comparaison avec SAM : Les auteurs ont comparé FRPO avec la Minimisation Sensible à la Netteté (SAM). Bien que SAM (opérant dans l'espace des paramètres) ait offert une certaine amélioration par rapport au GRPO, FRPO (opérant dans l'espace des politiques/KL) s'est révélé supérieur. Les auteurs notent que la boule L2 de SAM dans l'espace des paramètres ne capture pas entièrement les changements de politique en aval, tandis que la boule KL contraint directement les changements de distribution de sortie.
• Rôle de $\lambda$ : Les expériences ont confirmé que $\lambda$ contrôle le compromis entre la récompense moyenne et la variance. Le $\lambda$ optimal est corrélé à l'écart-type des récompenses et au budget KL $\rho$.

Importance et Revendications

L'article revendique que la robustesse à l'ajustement fin en aval devrait être intégrée directement dans la politique de base lors du RLHF, plutôt que de compter sur des interventions ultérieures. En optimisant pour la stabilité de la récompense au sein d'un voisinage borné par la KL, FRPO identifie des politiques résidant dans des régions « plates en récompense », les rendant intrinsèquement résistantes à l'oubli catastrophique.

Les auteurs positionnent FRPO comme une solution polyvalente qui :

• Fonctionne dans des domaines divers (sécurité, mathématiques, codage).
• Ne nécessite aucun coût computationnel supplémentaire par rapport au GRPO standard.
• Fournit une approche principielle et mathématiquement fondée pour l'apprentissage continu dans les LLM, agnostique du protocole spécifique d'ajustement fin en aval.

Ce travail suggère que le principe d'optimisation pour la robustesse à l'adaptation future peut s'étendre au-delà du RLHF vers le pré-entraînement et l'ajustement fin supervisé, offrant une nouvelle direction pour la recherche sur l'alignement et l'apprentissage continu.