Pushing Biomolecular Utility-Diversity Frontiers with Supergroup Relative Policy Optimization

Abstract: Biomolecular generators are often adapted with reward feedback to improve task-specific utility, but pushing utility alone can concentrate generation on a narrow family of candidates. Maintaining diversity is difficult because sample diversity is a set-level property. We introduce Supergroup Relative Policy Optimization (SGRPO), a flexible GRPO-style framework that directly constructs rewards from set-level diversity. For each condition, SGRPO samples a supergroup of candidate sets, compares their diversity under the same condition, and redistributes the group diversity reward to individual rollouts through leave-one-out diversity contributions before combining it with rollout-level utility. This design decouples SGRPO from a particular generator, utility reward, or diversity metric, and allows instantiation with different GRPO-style approaches. We evaluate SGRPO on de novo small-molecule design, pocket-based small-molecule design, and de novo protein design, instantiating it with both GRPO and Coupled-GRPO across autoregressive and discrete diffusion generators. Across decoding sweeps, SGRPO expands the utility-diversity Pareto frontier and achieves the best frontier-level metrics relative to pretrained generators, GRPO, and memory-assisted GRPO when applicable. Our analyses further show that direct set-level diversity rewards remain effective with small groups and help preserve broader generation-distribution coverage during post-training. The code is available at https://github.com/IDEA-XL/SGRPO.

• The paper introduces SGRPO, which integrates direct set-level diversity optimization into RL-based biomolecular generators to mitigate mode collapse.
• It employs same-condition supergroups with leave-one-out diversity credit assignment, combining utility and redistributed diversity metrics to guide PPO-based updates.
• Empirical evaluations across small-molecule and protein design domains show SGRPO achieves superior utility-diversity trade-offs, validated by metrics like hypervolume and DIP.

Supergroup Relative Policy Optimization for Utility-Diversity Frontier Expansion in Biomolecular Generation

Motivation and Problem Statement

Biomolecular generation tasks, including small-molecule and protein design, increasingly leverage RL-based post-training to improve candidates against domain-specific utility metrics. However, maximizing utility often leads to mode collapse and reduced diversity, undermining downstream applicability where broad candidate exploration is essential. Diversity is intrinsically a set-level property, yet most RL and reward-shaping approaches treat diversity indirectly, penalizing novelty relative to past samples without directly targeting the diversity among sets generated under a fixed condition. The central problem addressed is whether diversity can be directly optimized—treated as a first-class set-level objective—while still ensuring effective per-candidate reward assignment for policy update.

SGRPO: Methodological Framework

SGRPO introduces a GRPO-style framework that leverages the concept of same-condition supergroups to directly incorporate set-level diversity into policy optimization for biomolecular generators. For a given condition, SGRPO samples $M$ groups of $K$ rollouts each (supergroup), computes group-level diversity scores, and carries out leave-one-out comparisons to derive group-relative diversity advantages. This signal is redistributed within each group according to leave-one-out set diversity contributions, emphasizing candidates genuinely contributing to group diversity. Rollout-level utility and redistributed diversity are combined, and the resulting supergroup-relative advantages are used to drive PPO-style updates with KL regularization towards a reference policy.

Figure 1: SGRPO pipeline—supergroup sampling, rollouts, diversity computation, reward redistribution, and PPO-driven policy update.

SGRPO is decoupled from the underlying generator architecture, utility reward, or diversity metric, and can be instantiated via standard GRPO or coupled-GRPO approaches, making it broadly applicable across autoregressive and discrete diffusion generators.

Experimental Evaluation Across Biomolecular Design Domains

SGRPO is rigorously benchmarked across three domains:

• de novo small-molecule design using GenMol (masked discrete diffusion over SAFE representation)
• pocket-conditioned small-molecule design using GenMol-P (with ESM-IF1 pocket embedding)
• de novo protein design using ProGen2 (autoregressive LM)

In each domain, the evaluation protocol computes the utility-diversity Pareto frontier from model outputs under a sweep of decoding hyperparameters (temperature, randomness), summarizing performance via hypervolume (HV), distance to ideal point (DIP), and R2 indicator.

Figure 2: Utility–diversity operating points across tasks; SGRPO dominates baselines by expanding the frontier.

SGRPO consistently achieves superior frontier-level metrics relative to pretrained generators, vanilla GRPO, and memory-assisted GRPO, particularly in the high-utility regime where mode collapse is prevalent. Notably, in pocket-based small-molecule design, SGRPO retains within-pocket diversity while achieving higher utility, mitigating the severe redundancy observed in utility-only RL approaches.

Mechanistic Insights and Ablation Analysis

SGRPO’s advantage stems not only from group-level diversity reward but also from its leave-one-out credit assignment, which ensures reward alignment at the set level:

• Ablation studies reveal that the diversity reward alone improves the frontier, but full SGRPO further strengthens the trade-off by properly redistributing credit to candidates enhancing group diversity.
• Training-dynamic analysis, using UMAP embeddings on protein sequence outputs, highlights that SGRPO maintains broader cluster coverage through training, whereas GRPO contracts distribution and memory-assisted GRPO drifts to a narrow, distant region.

  Figure 3: SGRPO surpasses GRPO in HV even with small group sizes and across reward weights; frontier expansion peaks at $\lambda$=0.5.

Efficiency analysis demonstrates that SGRPO remains effective with small group sizes and tolerates a range of diversity–utility weightings, making it robust for practical deployment.

Theoretical Properties

SGRPO leverages normalized pairwise diversity metrics whose group-wise estimates are unbiased proxies for full-sample diversity. Concentration inequalities confirm that average diversity over small groups approximates the diversity of the entire same-condition sample as group size increases, justifying the use of finite groups in policy optimization.

Implications and Prospects

Practically, SGRPO enables biomolecular generators to traverse a broader utility-diversity spectrum, enhancing both targeted optimization and exploratory capacity. This lets practitioners achieve improved candidates at fixed diversity or, conversely, greater diversity at fixed utility—critical for hit discovery, lead expansion, and protein engineering.

Theoretically, SGRPO demonstrates that direct optimization of set-level diversity, rather than indirect novelty or entropy objectives, is tractable within RL policy update frameworks. Its modularity supports adaptation to evolving diversity metrics, generator architectures, and reward formulas.

Future directions include:

• Scalable diversity estimation via reusable similarities or differentiable metrics
• Task-specific diversity metrics and adaptive group structures
• Joint optimization of multiple utility axes alongside diversity, enabling more nuanced Pareto operating points

Conclusion

Supergroup Relative Policy Optimization formally addresses the utility-diversity trade-off in biomolecular generation by integrating set-level diversity into RL-based post-training. Across molecular and protein generators, SGRPO robustly expands the attainable utility-diversity frontier, outperforming traditional baselines while preserving scientific applicability through improved candidate exploration. The methodological and empirical evidence supports SGRPO as a principled, flexible foundation for reward-driven biomolecular generation.