Prefix Grouper: Addestramento Efficiente di GRPO tramite Forward a Prefisso Condiviso

Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward

June 5, 2025

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

cs.AI

Abstract

Group Relative Policy Optimization (GRPO) migliora l'apprendimento delle politiche calcolando i gradienti da confronti relativi tra output candidati che condividono un prefisso di input comune. Nonostante la sua efficacia, GRPO introduce un sovraccarico computazionale significativo quando elabora prefissi condivisi lunghi, che devono essere codificati in modo ridondante per ogni membro del gruppo. Questa inefficienza diventa un collo di bottiglia importante per la scalabilità negli scenari di apprendimento con contesti lunghi. Proponiamo Prefix Grouper, un algoritmo di addestramento GRPO efficiente che elimina il calcolo ridondante del prefisso attraverso una strategia di Shared-Prefix Forward. In particolare, ristrutturando l'auto-attenzione in due parti, il nostro metodo consente di codificare il prefisso condiviso una sola volta, preservando al contempo la completa differenziabilità e la compatibilità con l'addestramento end-to-end. Forniamo sia prove teoriche che empiriche che Prefix Grouper è equivalente in addestramento al GRPO standard: produce output in avanti e gradienti all'indietro identici, garantendo che le dinamiche di ottimizzazione e le prestazioni finali della politica rimangano invariate. Empiricamente, i nostri esperimenti confermano che Prefix Grouper ottiene risultati consistenti riducendo significativamente il costo computazionale dell'addestramento, specialmente negli scenari con prefissi lunghi. Il metodo proposto è completamente plug-and-play: è compatibile con le architetture basate su GRPO esistenti e può essere integrato senza soluzione di continuità nelle pipeline di addestramento attuali come sostituzione diretta, senza richiedere modifiche strutturali e solo minimi cambiamenti alla costruzione dell'input e al calcolo dell'attenzione. Prefix Grouper consente l'uso di dimensioni di gruppo maggiori a parità di budget computazionale, migliorando così la scalabilità di GRPO per compiti più complessi e modelli più grandi. Il codice è ora disponibile all'indirizzo 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