Prefix Grouper: Эффективное обучение GRPO с использованием разделяемого префикса в прямом проходе

Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward

June 5, 2025

Авторы: Zikang Liu, Tongtian Yue, Yepeng Tang, Longteng Guo, Junxian Cai, Qingbin Liu, Xi Chen, Jing Liu

cs.AI

Аннотация

Групповая Оптимизация Относительной Политики (Group Relative Policy Optimization, GRPO) улучшает обучение политик за счет вычисления градиентов на основе относительных сравнений между кандидатными выходами, которые имеют общий входной префикс. Несмотря на свою эффективность, GRPO вносит значительные вычислительные издержки при обработке длинных общих префиксов, которые должны избыточно кодироваться для каждого члена группы. Эта неэффективность становится основным узким местом масштабируемости в сценариях обучения с длинным контекстом. Мы предлагаем Prefix Grouper, эффективный алгоритм обучения GRPO, который устраняет избыточные вычисления префиксов с помощью стратегии Shared-Prefix Forward. В частности, путем реструктуризации самовнимания на две части наш метод позволяет кодировать общий префикс только один раз, сохраняя при этом полную дифференцируемость и совместимость с сквозным обучением. Мы предоставляем как теоретические, так и эмпирические доказательства того, что 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