Agrupador de Prefijos: Entrenamiento Eficiente de GRPO mediante Propagación hacia Adelante de Prefijos Compartidos

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

Resumen

La Optimización de Políticas Relativas por Grupos (GRPO, por sus siglas en inglés) mejora el aprendizaje de políticas al calcular gradientes a partir de comparaciones relativas entre salidas candidatas que comparten un prefijo de entrada común. A pesar de su efectividad, GRPO introduce un sobrecosto computacional significativo al procesar prefijos compartidos largos, los cuales deben codificarse de manera redundante para cada miembro del grupo. Esta ineficiencia se convierte en un cuello de botella importante para la escalabilidad en escenarios de aprendizaje con contextos largos. Proponemos Prefix Grouper, un algoritmo de entrenamiento eficiente para GRPO que elimina el cómputo redundante del prefijo mediante una estrategia de Avance de Prefijo Compartido. En particular, al reestructurar la autoatención en dos partes, nuestro método permite que el prefijo compartido se codifique una sola vez, manteniendo al mismo tiempo la diferenciabilidad completa y la compatibilidad con el entrenamiento de extremo a extremo. Proporcionamos evidencia tanto teórica como empírica de que Prefix Grouper es equivalente en entrenamiento al GRPO estándar: produce salidas directas y gradientes inversos idénticos, lo que garantiza que la dinámica de optimización y el rendimiento final de la política permanezcan inalterados. Empíricamente, nuestros experimentos confirman que Prefix Grouper logra resultados consistentes mientras reduce significativamente el costo computacional del entrenamiento, especialmente en escenarios con prefijos largos. El método propuesto es completamente plug-and-play: es compatible con arquitecturas basadas en GRPO existentes y puede integrarse sin problemas en las tuberías de entrenamiento actuales como un reemplazo directo, sin requerir modificaciones estructurales y solo cambios mínimos en la construcción de entradas y el cómputo de atención. Prefix Grouper permite el uso de tamaños de grupo más grandes bajo el mismo presupuesto computacional, mejorando así la escalabilidad de GRPO para tareas más complejas y modelos más grandes. El código está disponible en 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