OR Else: A Differentiable Trust Region for Policy Optimization

Résumé technique : OR ELSE : UNE RÉGION DE CONFIANCE DIFFÉRENCIABLE POUR L'OPTIMISATION DE POLITIQUE

Énoncé du problème
L'apprentissage par renforcement à partir de la rétroaction humaine (RLHF) pour les grands modèles de langage (LLM) repose couramment sur l'optimisation de politique proximale (PPO) et sa variante, l'optimisation de politique relative au groupe (GRPO). Les deux méthodes utilisent un objectif de substitut tronqué (clipped surrogate objective) pour contraindre les mises à jour de la politique. Bien que le mécanisme de troncature crée une région de gradient nul unilatérale pour empêcher une dérive excessive de la politique dans la direction « favorable », la transition vers cette région est géométriquement abrupte : la dérivée de la fonction d'objectif par rapport au ratio de probabilité change de manière discontinue au niveau de la limite de troncature. Les auteurs examinent si l'Output Reset (OR - Réinitialisation de Sortie), une règle de saturation unilatérale lisse initialement conçue pour l'entraînement de classificateurs, offre une alternative supérieure pour le post-entraînement des LLM. Plus précisément, l'article demande si le remplacement du substitut tronqué par une perte basée sur l'OR, plus différentiable, peut améliorer la stabilité de l'optimisation et les scores de récompense finaux sans introduire de nouvelles pathologies.

Méthodologie
L'article introduit PPO-OR et GRPO-OR, deux variantes qui remplacent le terme de politique tronqué standard par une perte de marge au carré dérivée de l'OR dans l'espace des log-ratios de jetons (tokens).

• Mécanisme central : Au lieu d'utiliser le ratio de probabilité $r_t(\theta)$, les méthodes opèrent sur le log-ratio de probabilité par jeton $\rho_t = \log \pi_\theta(a_t|s_t) - \log \pi_{\theta_{old}}(a_t|s_t)$.
• La perte OR : Pour un jeton ayant un avantage estimé non nul $\hat{A}_t$, une marge cible $\sigma_t \alpha_{OR}$ est fixée (où $\sigma_t = \text{sign}(\hat{A}_t)$). La perte est une marge au carré unilatérale : $\ell_{OR,t} = [\max(0, \alpha_{OR} - \sigma_t \rho_t)]^2$.
• Comportement de l'Output Reset : Une fois que $\rho_t$ franchit la marge dans la direction favorisée par le signe de l'avantage, la cible est « réinitialisée » à la valeur actuelle de $\rho_t$. Cela résulte en un résidu nul et une contribution de gradient direct nulle pour ce jeton. La dérivée tend vers zéro de manière continue à la frontière, rendant la perte $C^1$ continue (contrairement au clip de PPO qui est $C^0$ continu mais non différentiable).
• Estimation de l'avantage :
  • PPO-OR : Conserve l'estimation d'avantage généralisé (GAE) avec une tête de valeur apprise.
  • GRPO-OR : Conserve les avantages relatifs au groupe (calculés à partir de $G=2$ réponses par prompt) sans critique.
• Configuration expérimentale : L'étude utilise Llama-3.2-1B-Instruct sur le jeu de données Anthropic hh-rlhf. Toutes les méthodes partagent un modèle de récompense unique gelé et sont évaluées sur trois graines (seeds). Les expériences sont conçues comme des comparaisons appariées au sein d'une même famille : PPO-clip vs PPO-OR (sous GAE) et GRPO vs GRPO-OR (sous avantages relatifs au groupe). Les comparaisons inter-familles sont notées comme non contrôlées en raison des différences d'estimation d'avantage, d'utilisation de critique et de longueurs de réponse.

Contributions clés

1. Proposition algorithmique : L'introduction de PPO-OR et GRPO-OR, qui substituent le substitut tronqué par une perte de marge au carré continûment différentiable dans l'espace des log-ratios.
2. Caractérisation théorique : Preuve formelle que l'objectif OR est $C^1$ continu à la frontière de saturation et que la direction de mise à jour au niveau de l'échantillon est strictement déterminée par le signe de l'avantage, s'annulant continûment lors du franchissement de la marge.
3. Évaluation empirique : Une comparaison rigoureuse entre la troncature et l'OR au sein de deux familles distinctes d'estimation d'avantage, soulignant que les bénéfices de l'OR ne sont pas uniformes mais dépendent du pipeline d'optimisation environnant.
4. Aperçu diagnostique : L'identification du fait que remplacer la troncature par l'OR ne résout pas intrinsèquement le problème de la grande dérive de politique cumulative observée dans les configurations relatives au groupe, distinguant la saturation du gradient local du contrôle global de la politique.

Résultats
L'article rapporte les scores du modèle de récompense pendant l'entraînement et les diagnostics d'optimisation sur trois graines :

• Famille GAE (PPO-clip vs PPO-OR) :

  • Score de récompense : PPO-OR a obtenu un score moyen final plus élevé (+0,500) que PPO-clip (+0,195), soit une différence de +0,305.
  • Stabilité : PPO-OR a présenté un écart-type entre les graines plus important (0,158) que PPO-clip (0,110).
  • Dérive : Les deux méthodes ont maintenu un déplacement du log-ratio entre le rollout et l'actuel similaire et faible (environ 0,3–0,5).

• Famille Relative au Groupe (GRPO vs GRPO-OR) :

  • Score de récompense : GRPO-OR n'a pas amélioré le score moyen ; il était légèrement inférieur (+0,457) à celui de GRPO (+0,488).
  • Stabilité : GRPO-OR a montré un écart-type entre les graines nettement plus faible (0,027) que GRPO (0,080).
  • Diagnostics : GRPO-OR a atteint un résidu OR terminal proche de zéro et une fraction de dépassement décroissante, tandis que la trace correspondante de GRPO est restée élevée et très variable.
  • Dérive : Crucialement, GRPO-OR n'a pas réduit le grand déplacement entre le rollout et l'actuel observé dans GRPO. Les deux méthodes relatives au groupe ont présenté un déplacement substantiellement plus élevé (environ 5–13) par rapport aux méthodes GAE.

Signification et Revendications
L'article conclut avec modestie, soulignant que les effets observés de l'OR sont dépendants du contexte :

• Pas une Région de Confiance Universelle : Les auteurs affirment explicitement que la saturation lisse de la perte OR ne constitue pas une région de confiance formelle. Elle supprime la contribution directe du gradient des jetons individuels après qu'ils ont franchi une marge, mais ne contraint pas la dérive de politique cumulative causée par les paramètres partagés, les pertes auxiliaires ou le changement de signes d'avantage au cours des itérations.
• Dépendance au Pipeline : L'efficacité de l'OR est liée à la méthode d'estimation de l'avantage. Dans le cadre GAE, l'OR est corrélé à des scores de récompense plus élevés mais augmente la variance. Dans le cadre relatif au groupe, l'OR réduit la variance et stabilise la trace de la perte, mais n'améliore pas le score de récompense moyen ni ne réduit la dérive de la politique.
• Limites : L'étude reconnaît que les résultats sont spécifiques à un modèle léger de 1B, un modèle de récompense gelé unique et une taille de groupe contrainte ($G=2$). Les conclusions caractérisent ce cadre expérimental spécifique plutôt que d'établir une supériorité générale sur PPO ou GRPO. L'article pose que la grande dérive de politique dans les méthodes relatives au groupe reste un problème non résolu que l'OR seul ne peut pas corriger, suggérant que les travaux futurs doivent traiter le déplacement cumulatif par une régularisation de référence plus forte ou des pénalités de déplacement explicites.

En résumé, l'article démontre que l'Output Reset offre une alternative différentiable à la troncature qui modifie le comportement d'optimisation et peut améliorer les scores de récompense dans des contextes spécifiques, mais qu'il ne résout pas intrinsèquement les défis de la stabilité de la politique ou de la dérive dans le cadre du RLHF relatif au groupe.