Prefix Grouper: Treinamento Eficiente de GRPO por meio de Forward de Prefixo Compartilhado

Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward

June 5, 2025

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

cs.AI

Resumo

O Group Relative Policy Optimization (GRPO) aprimora o aprendizado de políticas ao calcular gradientes a partir de comparações relativas entre saídas candidatas que compartilham um prefixo de entrada comum. Apesar de sua eficácia, o GRPO introduz uma sobrecarga computacional significativa ao processar prefixos compartilhados longos, que precisam ser codificados de forma redundante para cada membro do grupo. Essa ineficiência se torna um grande gargalo de escalabilidade em cenários de aprendizado com contexto longo. Propomos o Prefix Grouper, um algoritmo de treinamento eficiente para GRPO que elimina a computação redundante de prefixos por meio de uma estratégia de Forward de Prefixo Compartilhado. Especificamente, ao reestruturar a auto-atenção em duas partes, nosso método permite que o prefixo compartilhado seja codificado apenas uma vez, mantendo a diferenciabilidade completa e a compatibilidade com o treinamento de ponta a ponta. Fornecemos evidências teóricas e empíricas de que o Prefix Grouper é equivalente ao GRPO padrão em termos de treinamento: ele produz saídas diretas e gradientes retropropagados idênticos, garantindo que a dinâmica de otimização e o desempenho final da política permaneçam inalterados. Empiricamente, nossos experimentos confirmam que o Prefix Grouper alcança resultados consistentes enquanto reduz significativamente o custo computacional do treinamento, especialmente em cenários com prefixos longos. O método proposto é totalmente plug-and-play: é compatível com arquiteturas baseadas em GRPO existentes e pode ser integrado de forma transparente aos pipelines de treinamento atuais como uma substituição direta, sem exigir modificações estruturais e apenas mudanças mínimas na construção de entradas e no cálculo de atenção. O Prefix Grouper permite o uso de tamanhos de grupo maiores sob o mesmo orçamento computacional, melhorando assim a escalabilidade do GRPO para tarefas mais complexas e modelos maiores. O código está disponível em 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