Résumé Technique : PS-PPO : Prefix-Sampling PPO pour le RLHF sans critiqueur

1. Énoncé du Problème

L'apprentissage par renforcement à partir de la rétroaction humaine (RLHF) est la méthode standard pour aligner les grands modèles de langage (LLM) avec les préférences humaines. Cependant, l'entraînement des LLM avec le RLHF fait face à des goulots d'étranglement computationnels importants, particulièrement dans les approches sans critiqueur (ex: GRPO, RLOO) qui évitent l'entraînement d'un réseau de valeur séparé.

Dans ces méthodes sans critiqueur, une récompense scalaire unique est assignée à une trajectoire complète (complétion) après la génération. Cette récompense est ensuite diffusée uniformément à tous les jetons (tokens) de la trajectoire pour calculer les avantages (advantages). Par conséquent, les mises à jour de la politique nécessitent une rétropropagation sur la trajectoire complète pour chaque déroulement (rollout). Les auteurs soutiennent que cela est inefficace car, dans de nombreuses tâches de raisonnement (telles que le calcul étape par étape), les préfixes intermédiaires contiennent souvent suffisamment d'informations pour déterminer le résultat final. L'analyse empirique des auteurs (Figure 1) montre que le taux de réussite d'une trajectoire se stabilise souvent bien avant la fin de la complétion, impliquant que les jetons de la suffixe transportent des signaux d'apprentissage redondants. Les mises à jour de la trajectoire complète gaspillent ainsi un calcul et une mémoire GPU substantiels sur ces suffixes redondants.

2. Méthodologie : PS-PPO

Les auteurs proposent le Prefix-Sampling Proximal Policy Optimization (PS-PPO), une méthode sans critiqueur et efficace en termes de calcul, qui exploite la redondance temporelle en tronquant la passe arrière tout en maintenant un estimateur de gradient non biaisé.

Mécanisme Central

Au lieu de rétropropager à travers toute la complétion $o = [o_1, \dots, o_T]$, PS-PPO échantillonne un instant de coupure $H$ pour chaque trajectoire. La mise à jour de la politique est calculée uniquement sur le préfixe $o_{\le H} = [o_1, \dots, o_H]$.

Pour garantir que cette troncature n'introduise pas de biais, PS-PPO emploie un schéma de repondération :

1. Distribution de Coupure : Une distribution conditionnée par le prompt $\xi_{1:T}$ est définie, où $\xi_t = \Pr(H \ge t | x)$ représente la probabilité que le $t$-ème jeton soit conservé. Cette distribution est non croissante ($\xi_1 \ge \xi_2 \ge \dots \ge \xi_T$).
2. Estimateur Non Biaisé : Le gradient pour une trajectoire tronquée est repondéré par l'inverse de la probabilité d'inclusion $1/\xi_t$. L'estimateur résultant $\hat{G}(\theta)$ est :
   $$\hat{G}(\theta) := \frac{1}{K} \sum_{k=1}^K \sum_{t=1}^{H^{(k)}} \frac{1}{\xi_t} g_t^{(k)}(\theta)$$
   où $g_t^{(k)}$ est le gradient par instant de temps. L'espérance de cet estimateur sur l'échantillonnage de $H$ permet de retrouver la mise à jour de la trajectoire complète.

Optimisation de la Distribution de Coupure

Les auteurs formulent la sélection de $\xi_{1:T}$ comme un problème d'optimisation convexe. L'objectif est de minimiser la variance de l'estimateur de gradient tronqué sous un budget de calcul $B$ (le nombre attendu de jetons rétropropagés).

• Surrogate de Variance : Les auteurs dérivent un substitut traitable de la variance induite par l'échantillonnage d'importance, l'approximant par $\sum w_t(x) (1/\xi_t - 1)$, où $w_t(x)$ représente l'importance de la mise à jour à l'instant $t$.
• Proxy d'Importance : Puisque le calcul de la norme exacte du score $\|\nabla_\theta \log \pi_\theta\|^2$ nécessite une passe arrière complète (ce qui contredirait l'objectif), les auteurs utilisent un proxy de type "forward-only" basé sur les activations de la tête de sortie et l'incertitude de la récompense.
  • Incertitude de la Récompense ($u_t(x)$) : Estimée via la variance des récompenses données l'état du préfixe, approximée par la distance entre les distributions de jetons suivants des déroulements réussis et infructueux.
  • Proxy de Norme de Gradient ($\bar{\gamma}_t$) : Dérivé de la norme du gradient de la tête de sortie.
• Optimisation : Le problème minimise $\sum \frac{\bar{\gamma}_t(x,t) u_t(x)}{\xi_t}$ sous la contrainte de budget et de monotonie. La solution est trouvée en utilisant l'algorithme Pool Adjacent Violators (PAV) pour imposer la contrainte de monotonie.

Flux de l'Algorithme

1. Générer $K$ complétions pour un prompt.
2. Calculer les récompenses terminales et diffuser les avantages.
3. Estimer l'incertitude par instant et les proxies de gradient à partir des déroulements.
4. Résoudre l'optimisation convexe pour obtenir les probabilités de coupure monotones optimales $\xi^*_{1:T}$.
5. Échantillonner un instant de coupure $H^{(k)}$ pour chaque trajectoire à partir de la distribution définie par $\xi^*$.
6. Effectuer la rétropropagation uniquement sur les jetons $t \le H^{(k)}$, en repondérant la perte par $1/\xi^*_t$.

3. Contributions Clés

1. Cadre PS-PPO : Une nouvelle méthode de RLHF sans critiqueur qui introduit une troncature stochastique de la passe arrière conditionnée par le prompt, réduisant considérablement les coûts de calcul et de mémoire.
2. Estimateur de Troncature Non Biaisé : Une dérivation théorique d'un estimateur repondéré par la probabilité d'inclusion qui retrouve la mise à jour complète par diffusion en espérance, sans nécessiter de déroulements supplémentaires ou de modèles de valeur auxiliaires.
3. Stratégie de Coupure Optimisée : Une approche fondée sur les principes pour déterminer où tronquer en résolvant un problème d'optimisation convexe qui équilibre la réduction de la variance et un budget de calcul, en utilisant des proxies "forward-only" pour l'importance du gradient et l'incertitude de la récompense.
4. Validation Empirique : Démonstration que PS-PPO atteint une précision comparable aux bases fortes sans critiqueur (GRPO, DAPO, RLOO) tout en réduisant substantiellement le temps d'entraînement et la mémoire GPU de pointe.

4. Résultats Expérimentaux

Les auteurs ont évalué PS-PPO sur des benchmarks de raisonnement mathématique (MATH500, AMC 2023, Minerva Math, AIME 2024/2025) en utilisant Llama-3.1-8B-Instruct et Qwen2.5-Math-7B.

• Performance : PS-PPO (Optimisé) a atteint une précision Pass@1 comparable ou légèrement supérieure aux bases fortes comme DAPO et S-GRPO sur tous les benchmarks. Par exemple, sur MATH500 avec Llama-3.1-8B, PS-PPO a atteint 47,6 % contre 46,8 % pour DAPO.
• Efficacité :
  • Temps d'entraînement : PS-PPO a réduit le temps d'entraînement par étape de 33 % à 45 % par rapport aux bases. Cela est principalement dû à la réduction des coûts de passage avant/arrière, compensant largement le surcoût du calcul de la distribution de coupure.
  • Mémoire : L'utilisation de la mémoire GPU de pointe a diminué de 15 % à 17 % grâce à la réduction du stockage des activations pour les séquences tronquées.
  • Scalabilité : Les gains d'efficacité sont devenus plus prononcés à mesure que la longueur maximale de la complétion ($T_{max}$) augmentait. À $T_{max}=4096$, PS-PPO était 2,8x plus rapide que S-GRPO et 3,3x plus rapide que DAPO lors de l'étape de mise à jour.
• Ablations :
  • Budget ($B$) : Un budget de $B=128$ jetons offrait le meilleur compromis entre précision et vitesse. Des budgets plus petits augmentaient la variance ; des budgets plus élevés offraient des rendements de précision décroissants avec des coûts de calcul plus élevés.
  • Stratégie de Coupure : La stratégie "Optimisée" (résolution du problème convexe) a nettement surperformé les stratégies de troncature uniforme, de priorité temporelle et heuristique, confirmant que la sélection adaptative de préfixes informatifs est cruciale.
  • Déroulements ($K$) : Augmenter $K$ améliorait la stabilité et la performance mais augmentait le calcul par étape. $K=8$ a été identifié comme un compromis pratique.

5. Signification et Revendications

L'article affirme que PS-PPO offre une voie scalable pour aligner les grands modèles de langage avec des exigences de ressources nettement réduites. La signification centrale réside dans l'observation que les préfixes intermédiaires dans les traces de raisonnement contiennent souvent suffisamment d'informations pour déterminer le résultat final, rendant les mises à jour de la trajectoire complète computationnellement redondantes.

En découplant le signal d'apprentissage de la longueur totale de la séquence via l'échantillonnage de préfixe stochastique et la repondération par importance, PS-PPO permet :

• Réduction des Barrières Matérielles : La baisse de la mémoire de pointe et des coûts de calcul rend le RLHF plus accessible aux groupes académiques et aux structures aux ressources limitées.
• Entraînement Efficace pour les Contextes Longs : La méthode est particulièrement efficace pour les longues traces de raisonnement où le coût d'une rétropropagation complète est prohibitif.

Les auteurs maintiennent un ton modeste, notant que PS-PPO est une méthode d'optimisation qui améliore l'efficacité mais ne détermine pas intrinsèquement la sécurité ou le comportement du modèle ; ceux-ci restent dépendants des spécifications de récompense, des données d'entraînement et des garde-fous de déploiement. La méthode préserve les propriétés théoriques de l'objectif de séquence complète en espérance, garantissant que les gains d'efficacité ne se font pas au détriment d'un apprentissage biaisé.