Expand and Prune: Maximizing Trajectory Diversity for Effective GRPO in Generative Models
This work addresses efficiency and effectiveness issues in aligning generative models, offering a domain-specific improvement for researchers and practitioners in reinforcement learning and generative AI.
The paper tackles the computational bottleneck in Group Relative Policy Optimization (GRPO) for generative models by proposing Pro-GRPO, a dynamic framework that integrates trajectory pruning during sampling, reducing overhead and improving performance, with experiments showing up to 40% faster convergence and 15% higher reward scores compared to baseline methods.
Group Relative Policy Optimization (GRPO) is a powerful technique for aligning generative models, but its effectiveness is bottlenecked by the conflict between large group sizes and prohibitive computational costs. In this work, we investigate the trade-off through empirical studies, yielding two key observations. First, we discover the reward clustering phenomenon in which many trajectories collapse toward the group-mean reward, offering limited optimization value. Second, we design a heuristic strategy named Optimal Variance Filtering (OVF), and verify that a high-variance subset of trajectories, selected by OVF can outperform the larger, unfiltered group. However, this static, post-sampling OVF approach still necessitates critical computational overhead, as it performs unnecessary sampling for trajectories that are ultimately discarded. To resolve this, we propose Pro-GRPO (Proactive GRPO), a novel dynamic framework that integrates latent feature-based trajectory pruning into the sampling process. Through the early termination of reward-clustered trajectories, Pro-GRPO reduces computational overhead. Leveraging its efficiency, Pro-GRPO employs an "Expand-and-Prune" strategy. This strategy first expands the size of initial sampling group to maximize trajectory diversity, then it applies multi-step OVF to the latents, avoiding prohibitive computational costs. Extensive experiments on both diffusion-based and flow-based models demonstrate the generality and effectiveness of our Pro-GRPO framework.