Prefix Grouper: Efficiënte GRPO-training via gedeelde voorvoegsels (Shared-Prefix Forward)

Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward

June 5, 2025

Auteurs: Zikang Liu, Tongtian Yue, Yepeng Tang, Longteng Guo, Junxian Cai, Qingbin Liu, Xi Chen, Jing Liu

cs.AI

Samenvatting

Group Relative Policy Optimization (GRPO) verbetert het leren van beleid door gradiënten te berekenen uit relatieve vergelijkingen tussen kandidaat-uitvoeren die een gemeenschappelijk invoervoorvoegsel delen. Ondanks de effectiviteit introduceert GRPO aanzienlijke rekenkundige overhead bij het verwerken van lange gedeelde voorvoegsels, die voor elk groepslid redundant gecodeerd moeten worden. Deze inefficiëntie wordt een grote schaalbaarheidsbelemmering in scenario's met lange contexten. Wij stellen Prefix Grouper voor, een efficiënt GRPO-trainingsalgoritme dat redundante voorvoegselberekeningen elimineert via een Shared-Prefix Forward-strategie. In het bijzonder, door self-attention op te splitsen in twee delen, maakt onze methode het mogelijk dat het gedeelde voorvoegsel slechts één keer wordt gecodeerd, terwijl volledige differentieerbaarheid en compatibiliteit met end-to-end training behouden blijven. Wij leveren zowel theoretisch als empirisch bewijs dat Prefix Grouper trainingsequivalent is aan standaard GRPO: het levert identieke voorwaartse uitvoeren en achterwaartse gradiënten op, waardoor de optimalisatiedynamiek en de uiteindelijke beleidsprestaties ongewijzigd blijven. Empirisch bevestigen onze experimenten dat Prefix Grouper consistente resultaten behaalt terwijl de rekenkundige kosten van het trainen aanzienlijk worden verlaagd, vooral in scenario's met lange voorvoegsels. De voorgestelde methode is volledig plug-and-play: het is compatibel met bestaande GRPO-gebaseerde architecturen en kan naadloos worden geïntegreerd in huidige trainingspijplijnen als een drop-in vervanging, zonder structurele aanpassingen en met slechts minimale wijzigingen aan invoerconstructie en aandachtberekening. Prefix Grouper maakt het gebruik van grotere groepsgroottes mogelijk binnen hetzelfde rekenkundige budget, waardoor de schaalbaarheid van GRPO naar complexere taken en grotere modellen wordt verbeterd. Code is nu beschikbaar op 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