Prefix Grouper：通过共享前缀前向提高GRPO训练效率
机器学习
2025-06-09 v1 人工智能
摘要
Group Relative Policy Optimization（GRPO）通过计算来自共享相同输入前缀的候选输出之间相对比较的梯度来增强策略学习。尽管其有效性突出，但在处理长共享前缀时，GRPO会引入巨大的计算开销，这些前缀必须为每个组成员冗余地编码。这种效率低下在长上下文学习场景中成为主要的可扩展性瓶颈。我们提出了Prefix Grouper，一种高效的GRPO训练算法，通过Shared-Prefix Forward策略消除冗余的前缀计算。具体而言，通过对自注意力机制进行重构，将其分为两个部分，我们的方法使共享前缀仅编码一次，同时保持完全可微和与端到端训练的兼容性。我们提供了理论和实证证据表明，Prefix Grouper与标准GRPO在训练等价性：它产生相同的前向输出和反向梯度，确保优化动力学和最终策略性能保持不变。实验结果表明，Prefix Grouper在显著降低训练计算成本方面实现了一致的结果，尤其是在长前缀场景中。该方法完全即插即用：它与现有的GRPO-based架构兼容，可作为 drop-in 替代品无缝集成到当前训练管道中，无需结构修改，仅需对输入构建和注意力计算进行最小更改。Prefix Grouper 使 larger group sizes 在相同的计算预算下成为可能，从而提高了GRPO针对更复杂任务和更大模型的可扩展性。代码现已公开available at https://github.com/johncaged/PrefixGrouper
引用
@article{arxiv.2506.05433,
title = {Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward},
author = {Zikang Liu and Tongtian Yue and Yepeng Tang and Longteng Guo and Junxian Cai and Qingbin Liu and Xi Chen and Jing Liu},
journal= {arXiv preprint arXiv:2506.05433},
year = {2025}
}
备注
10 pages, technical report