Prefix Grouper : Entraînement efficace de GRPO via le partage de préfixes en propagation avant

Résumé

L'Optimisation de Politique Relative par Groupe (GRPO) améliore l'apprentissage de politiques en calculant des gradients à partir de comparaisons relatives entre des sorties candidates partageant un préfixe d'entrée commun. Malgré son efficacité, GRPO introduit une surcharge computationnelle importante lors du traitement de longs préfixes partagés, qui doivent être encodés de manière redondante pour chaque membre du groupe. Cette inefficacité devient un goulot d'étranglement majeur pour la scalabilité dans les scénarios d'apprentissage à contexte long. Nous proposons Prefix Grouper, un algorithme d'entraînement GRPO efficace qui élimine le calcul redondant des préfixes via une stratégie de Forward à Préfixe Partagé. En particulier, en restructurant l'auto-attention en deux parties, notre méthode permet d'encoder le préfixe partagé une seule fois, tout en préservant la différentiabilité complète et la compatibilité avec l'entraînement de bout en bout. Nous fournissons des preuves théoriques et empiriques que Prefix Grouper est équivalent en entraînement au GRPO standard : il produit des sorties avant et des gradients arrière identiques, garantissant que la dynamique d'optimisation et les performances finales de la politique restent inchangées. Empiriquement, nos expériences confirment que Prefix Grouper obtient des résultats cohérents tout en réduisant significativement le coût computationnel de l'entraînement, en particulier dans les scénarios à long préfixe. La méthode proposée est entièrement plug-and-play : elle est compatible avec les architectures existantes basées sur GRPO et peut être intégrée de manière transparente dans les pipelines d'entraînement actuels en tant que remplacement direct, sans nécessiter de modifications structurelles et seulement des changements minimaux dans la construction des entrées et le calcul de l'attention. Prefix Grouper permet l'utilisation de tailles de groupe plus importantes sous le même budget computationnel, améliorant ainsi la scalabilité de GRPO pour des tâches plus complexes et des modèles plus grands. Le code est désormais disponible à l'adresse https://github.com/johncaged/PrefixGrouper.

English

Group Relative Policy Optimization (GRPO) enhances policy learning by computing gradients from relative comparisons among candidate outputs that share a common input prefix. Despite its effectiveness, GRPO introduces substantial computational overhead when processing long shared prefixes, which must be redundantly encoded for each group member. This inefficiency becomes a major scalability bottleneck in long-context learning scenarios. We propose Prefix Grouper, an efficient GRPO training algorithm that eliminates redundant prefix computation via a Shared-Prefix Forward strategy. In particular, by restructuring self-attention into two parts, our method enables the shared prefix to be encoded only once, while preserving full differentiability and compatibility with end-to-end training. We provide both theoretical and empirical evidence that Prefix Grouper is training-equivalent to standard GRPO: it yields identical forward outputs and backward gradients, ensuring that the optimization dynamics and final policy performance remain unchanged. Empirically, our experiments confirm that Prefix Grouper achieves consistent results while significantly reducing the computational cost of training, particularly in long-prefix scenarios. The proposed method is fully plug-and-play: it is compatible with existing GRPO-based architectures and can be seamlessly integrated into current training pipelines as a drop-in replacement, requiring no structural modifications and only minimal changes to input construction and attention computation. Prefix Grouper enables the use of larger group sizes under the same computational budget, thereby improving the scalability of GRPO to more complex tasks and larger models. Code is now available at https://github.com/johncaged/PrefixGrouper

Prefix Grouper : Entraînement efficace de GRPO via le partage de préfixes en propagation avant

Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward

Résumé

Support