Expand and Prune: Maximizing Trajectory Diversity for Effective GRPO in Generative Models

TL;DR

This paper introduces Pro-GRPO, a dynamic framework that enhances trajectory diversity and reduces computational costs in generative model alignment by expanding and pruning trajectories during sampling.

Contribution

We propose Pro-GRPO, a novel dynamic method integrating latent feature-based pruning with an expand-and-prune strategy to improve efficiency and effectiveness in trajectory-based generative model optimization.

Findings

Pro-GRPO reduces computational overhead compared to static methods.

Expanding initial trajectory groups increases diversity and optimization potential.

Pro-GRPO outperforms existing methods on diffusion and flow-based models.

Abstract

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.