Prefix Grouper: 공유 접두어 순전파를 통한 효율적인 GRPO 학습

초록

그룹 상대 정책 최적화(Group Relative Policy Optimization, GRPO)는 공통 입력 접두사를 공유하는 후보 출력들 간의 상대적 비교를 통해 그래디언트를 계산함으로써 정책 학습을 강화한다. 그러나 GRPO는 긴 공유 접두사를 처리할 때 각 그룹 멤버에 대해 중복적으로 인코딩해야 하므로 상당한 계산 오버헤드를 유발한다. 이러한 비효율성은 장문맥 학습 시나리오에서 주요 확장성 병목 현상으로 작용한다. 본 연구에서는 공유 접두사 계산의 중복을 제거하는 효율적인 GRPO 학습 알고리즘인 Prefix Grouper를 제안한다. 특히, 셀프 어텐션을 두 부분으로 재구성함으로써 공유 접두사를 한 번만 인코딩하면서도 완전한 미분 가능성과 엔드투엔드 학습과의 호환성을 유지한다. Prefix Grouper가 표준 GRPO와 학습적으로 동등함을 이론적 및 실험적으로 입증한다: 동일한 순방향 출력과 역방향 그래디언트를 생성하여 최적화 동역학과 최종 정책 성능이 변하지 않음을 보장한다. 실험 결과, Prefix Grouper는 특히 긴 접두사 시나리오에서 학습의 계산 비용을 크게 줄이면서도 일관된 결과를 달성함을 확인하였다. 제안된 방법은 완전히 플러그앤플레이 방식으로, 기존 GRPO 기반 아키텍처와 호환되며 현재의 학습 파이프라인에 드롭인 대체로 원활하게 통합될 수 있다. 구조적 수정 없이 입력 구성과 어텐션 계산에 최소한의 변경만 필요하다. Prefix Grouper는 동일한 계산 예산 내에서 더 큰 그룹 크기를 사용할 수 있게 하여 GRPO의 확장성을 더 복잡한 작업과 더 큰 모델로 향상시킨다. 코드는 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

Prefix Grouper: 공유 접두어 순전파를 통한 효율적인 GRPO 학습

Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward

초록

Support