Future-KL Regularized GRPO: Process-Level Credit Assignment from $f$-Divergence Regularization

Résumé Technique : GRPO Régularisé par le KL du Futur

Énoncé du Problème
L'Optimisation de Politique Relative par Groupe (GRPO) s'est imposée comme un paradigme dominant pour l'apprentissage par renforcement (RL) sans critique dans le post-entraînement des grands modèles de langage (LLM), en particulier pour le raisonnement mathématique. GRPO élimine le besoin d'une fonction de valeur apprise en normalisant les récompenses de résultat au sein d'un groupe de réponses échantillonnées. Cependant, les implémentations standards de la régularisation KL dans GRPO traitent généralement la pénalité comme une pénalité locale de token côté perte (par exemple, en utilisant l'estimateur $k_3$).

Les auteurs identifient un défaut fondamental dans cette approche : elle ignore l'attribution de crédit au niveau du processus induite par la génération autoregressive. Dans une trajectoire autoregressive, un token précoce non seulement engendre son propre coût de régularisation local, mais modifie également la distribution des futurs préfixes, affectant ainsi les coûts KL de tous les tokens subséquents. Les termes de perte locaux standards ne différencient que la pénalité du token échantillonné, omettant le « retour à venir de régularisation future » causal.

De plus, l'article met en évidence une incompatibilité structurelle entre la normalisation par groupe de GRPO et les formulations standard de KL dans la récompense. Contrairement aux objectifs de récompense linéaires (par exemple, RL MaxEnt) où la récompense et le régularisateur peuvent être fusionnés, la normalisation par groupe de GRPO induit une utilité non linéaire au niveau du prompt (spécifiquement, une transformation arcsin pour les récompenses binaires). Cette non-linéarité signifie que l'ajout de pénalités KL aux récompenses avant la normalisation déforme l'objectif implicite, effondrant l'information KL au niveau du token en un scalaire au niveau de la trajectoire et échouant à préserver le signal de gradient correct.

Méthodologie
L'article dérive le gradient de politique on-policy pour des objectifs de style GRPO sous une régularisation $f$-divergence générale au niveau des tokens. L'insight théorique clé est que le gradient se décompose en deux termes distincts avec des propriétés d'échelle différentes :

1. Le Terme de Récompense : Hérité de l'utilité non linéaire $\ell(p)$, ce terme porte un gain dépendant du prompt $\ell'(p)$. Pour les récompenses binaires, cela retrouve l'avantage GRPO standardisé.
2. Le Terme de Régularisateur : Ce terme est indépendant de $\ell'$ et se décompose en une pénalité locale de token plus un retour à venir de régularisation future causal.

Sur la base de cette dérivation, les auteurs proposent l'Optimisation de Politique Régularisée par le KL du Futur (FRPO). L'algorithme modifie la mise à jour GRPO standard de manière découplée :

1. Calculer l'avantage de résultat normalisé par groupe ($\hat{A}^{GRPO}$) strictement à partir des récompenses de vérificateur originales, préservant l'utilité non linéaire implicite.
2. Construire un terme de correction KL du futur au niveau du token en calculant la somme cumulée inverse des rapports log par token ($\delta_{i,j} = \log \frac{\pi_\theta}{\pi_{ref}}$) du token actuel jusqu'à la fin de la trajectoire.
3. Ajouter ce terme KL du futur à l'avantage après la construction de l'avantage, plutôt que de façonner la récompense avant la normalisation.

La règle de mise à jour pour un token $t$ dans la réponse $i$ est :
$$\hat{A}^{FRPO}_{i,t} = \text{sg}\left( \hat{A}^{GRPO}_i - \beta \sum_{j=t}^{T_i} \delta_{i,j} \right)$$
où $\text{sg}$ désigne l'arrêt du gradient. Cette approche ne nécessite ni critique, ni passages de modèle supplémentaires, et maintient la même surcharge computationnelle que le GRPO standard.

Contributions Clés

1. Identification de l'Utilité Non Linéaire : L'article démontre que la normalisation par groupe de GRPO induit une utilité non linéaire ($\ell(p) = 2 \arcsin \sqrt{p}$ pour les récompenses binaires), rendant la récompense et le KL non fusionnables. L'ajout de KL aux récompenses avant la normalisation déforme structurellement l'objectif.
2. Dérivation du Gradient On-Policy : Les auteurs dérivent le gradient on-policy correct pour la régularisation $f$-divergence au niveau des tokens, montrant qu'il nécessite un terme de retour à venir de régularisation future absent dans les implémentations locales côté perte.
3. Algorithme FRPO : Ils proposent FRPO, une mise à jour simple et découplée qui préserve l'avantage de résultat de GRPO tout en injectant l'attribution de crédit KL du futur nécessaire. Cette méthode sépare le classement au niveau de la trajectoire de la régularisation KL au niveau du processus.
4. Interprétation d'Optimalité : La structure KL du futur est montrée comme étant alignée avec les conditions d'optimalité du RL régularisé par l'entropie, où la politique optimale satisfait une forme de récompense de processus impliquant les coûts KL futurs.

Résultats Expérimentaux
Les auteurs évaluent FRPO sur des tâches de raisonnement mathématique en utilisant des modèles de base Qwen3 (1.7B, 4B et 30B-A3B) entraînés sur le jeu de données DAPO-Math-17k.

• Performance : Sur le modèle 30B-A3B, FRPO améliore le pass@16 d'un gain absolu d'environ 5 % par rapport à la meilleure base de référence (GRPO standard avec KL côté perte) sur le benchmark MATH500. Il surpasse également les bases de référence sur AIME24, AIME25 et d'autres benchmarks de niveau Olympiade.
• Stabilité et Dérive : FRPO maintient une entropie finale plus élevée et une dérive de politique plus faible (mesurée par PPO-KL) par rapport aux bases de référence KL côté perte. Cela suggère que la méthode fait émerger des patterns de raisonnement par Chaîne de Pensée (CoT) longs sans déviation excessive par rapport à la politique de référence.
• Ablation : Les expériences confirment que le KL inverse (recherche de mode) fonctionne mieux que le KL direct (couverture de mode) pour cette tâche. L'article note également que l'utilisation de l'estimateur $k_1$ (rapport log) est préférée à $k_3$ pour le terme KL du futur, car la non-négativité de $k_3$ crée un signal de pénalité unilatéral qui entrave l'attribution de crédit.
• Modes d'Échec : L'article observe que les approches « KL dans la récompense » (ajout de KL aux récompenses avant la normalisation) conduisent à un effondrement du modèle ou à une dégradation des performances, validant l'affirmation théorique selon laquelle une telle fusion est structurellement incorrecte sous l'utilité non linéaire de GRPO.

Signification
L'article soutient que la signification de ce travail réside dans la correction d'une erreur d'implémentation subtile mais critique dans un cadre RL largement utilisé. En déplaçant la perspective de « pénalité de perte locale » vers « attribution de crédit au niveau du processus », les auteurs montrent que GRPO peut être amélioré sans introduire de fonctions de valeur complexes ni de coûts computationnels supplémentaires. La méthode FRPO proposée fournit une modification légère et théoriquement fondée qui améliore les capacités de raisonnement tout en préservant les propriétés de stabilité et de généralisation du modèle de base. Ce travail souligne que, dans les contextes autoregressifs, la régularisation est intrinsèquement un signal au niveau du processus, et ignorer le coût futur à venir conduit à une attribution de crédit sous-optimale.