通过共享前缀计算,大幅降低GRPO训练的冗余开销。
Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward
- 将自注意力拆分为两部分,仅编码一次共享前缀
- 在长前缀场景下计算成本显著下降,性能与原版一致
- 无需修改结构,可直接替换现有GRPO训练流程
组相对策略优化(GRPO)通过比较共享相同输入前缀的候选输出来提升策略学习效果。然而,在处理长共享前缀时,传统方法需对每个组内成员重复编码前缀,造成显著计算开销,成为长上下文学习中的主要扩展瓶颈。本文提出Prefix Grouper,采用共享前缀前向传播策略,通过重构自注意力机制,使共享前缀仅需编码一次,同时保持端到端训练的可微性。理论与实验证明,Prefix Grouper与标准GRPO完全等价:前向输出和反向梯度一致,优化动态与最终策略性能不变。实验表明,该方法在长前缀场景中显著降低计算成本,且支持更大组规模,在相同算力下提升GRPO的可扩展性。该方法为即插即用设计,兼容现有GRPO架构,无需结构调整,仅需轻微输入构造与注意力修改即可集成。代码已开源。
原文摘要 · Abstract (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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。