Pushing Biomolecular Utility-Diversity Frontiers with Supergroup Relative Policy Optimization

Xinwu Ye; He Cao; Hao Li; Bin Feng; Zijing Liu; Xiangru Tang; Yu Li; Shenghua Gao

Source-linked AI summary

Pushing Biomolecular Utility-Diversity Frontiers with Supergroup Relative Policy Optimization

Xinwu Ye, He Cao, Hao Li, Bin Feng, Zijing Liu, Xiangru Tang, Yu Li, Shenghua Gao

arXiv:2605.08659v3May 9, 2026cs.CEq-bio.BM

TL;DR

Biomolecular reinforcement learning must improve utility without collapsing generation diversity, but diversity is defined over sets rather than individual samples. SGRPO compares same-condition candidate groups, redistributes leave-one-out diversity contributions to individual rollouts, and combines them with utility. Across small-molecule and protein design settings, it expands the utility-diversity Pareto frontier over pretrained generators and reinforcement-learning baselines.

Problem
Utility optimization can concentrate biomolecular generation on narrow candidate families, while diversity is a set-level property that individual-rollout updates do not directly represent.
Method
SGRPO samples same-condition candidate sets, compares their diversity, redistributes group rewards through leave-one-out contributions, and combines them with rollout-level utility in GRPO-style optimization.
Results
SGRPO improves the utility-diversity Pareto frontier across de novo small-molecule, pocket-based small-molecule, and de novo protein design relative to pretrained generators and relevant RL baselines.
Takeaways & Limitations
Directly rewarding diverse generated sets provides a practical way to expand the operating points available to biomolecular generation models.
Takeaways & Limitations
Memory-assisted coupled-GRPO is excluded from the de novo small-molecule comparison because global memory mixes unrelated pockets, while pocket-specific memories require separate optimization and extra compute.

Abstract

from arXiv · show

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.

1 Introduction

Biomolecular post-training must balance task-specific utility with diversity because utility optimization can concentrate outputs in narrow candidate families. SGRPO addresses this set-level challenge through same-condition group comparisons and redistributed diversity rewards.

Utility-only reward optimization can concentrate biomolecular generation on narrow candidate families, while highly diverse generators may produce too few high-utility samples.
SGRPO is evaluated across different generator families, conditioning settings, utility functions, and diversity metrics as a broadly applicable post-training principle.
Set-level diversity is difficult to optimize because policy updates act on individual rollouts rather than collections of candidates.
SGRPO samples multiple candidate sets under each condition, scores their diversity, compares same-condition sets, and redistributes diversity rewards through leave-one-out contributions.
Across three biomolecular generation settings, SGRPO consistently improves the attainable utility-diversity Pareto frontier over pretrained generators and relevant GRPO baselines.It remains effective with small group sizes and better preserves generation-distribution coverage during post-training.

2 Related Work

Prior diversity-aware reinforcement-learning methods either depend closely on particular generators or use indirect novelty and memory-based reward shaping. These approaches can help empirically but do not directly optimize the diversity of the current candidate set.

Biomolecular objective optimization includes reinforcement learning, preference optimization, iterative retraining, and latent-space Bayesian or evolutionary methods.
Some diversity-promoting methods are tailored to specific generators or design tasks, including joint SMILES generation, fragment construction, and protein sequence-design regularization.
Broader diversity-aware approaches use indirect mechanisms such as memory penalties, novelty rewards, entropy regularization, visitation bonuses, or diverse mini-batch selection.

3 Problem Setup: Utility–Diversity Frontier in Biomolecular Generation

Biomolecular generation evaluates both candidate-level utility and set-level diversity under conditioning and decoding choices. The resulting utility-diversity frontier captures the operating points available to a generator and motivates methods that push it outward.

A conditional generator assigns each candidate an individual utility score based on the conditioning input, which may be empty, a task specification, or a target environment.
Diversity is a set-level score over candidates generated under the same condition, measuring properties such as internal, scaffold, sequence, or cluster diversity.
Decoding strategies jointly determine expected individual utility and expected set-level diversity for the generator.
The attainable frontier P(θ) consists of utility-diversity pairs induced by decoding strategies, with Pareto-optimal points lacking another strategy that improves both quantities.
The goal is to push the frontier outward, enabling better utility at fixed diversity, better diversity at fixed utility, or both.

4 Supergroup Relative Policy Optimization

SGRPO optimizes utility and set-level diversity by comparing candidate groups sampled under the same condition, redistributing group diversity signals to individual rollouts, and applying a PPO-style update. Its construction uses local same-condition comparisons and combines utility with diversity through a tunable trade-off.

4 Supergroup Relative Policy Optimization: SGRPO samples several same-condition candidate groups, scores group diversity, redistributes group signals by within-group contributions, and updates the policy with a relative PPO-style objective.
4 Supergroup Relative Policy Optimization: For each condition, the supergroup contains M groups of K rollouts, with M controlling alternative-group comparisons and K controlling group size.
4 Supergroup Relative Policy Optimization: Same-condition comparisons avoid confounding policy quality with differences in intrinsic condition difficulty.
4.1 Supergroup-relative group diversity: Utility is assigned to individual candidates, whereas diversity is evaluated over entire groups.
4.2 Utility and group-level diversity: Group diversity is centered within the supergroup through leave-one-out comparisons, so positive or negative values indicate whether a group is more or less diverse than its alternatives.
4.2 Utility and group-level diversity: For normalized pairwise diversity, average diversity over groups of size K is an unbiased proxy for diversity of a larger same-condition sample.
4.3 Diversity reward redistribution: Leave-one-out set contributions assign more diversity signal to candidates that matter more for their group’s diversity.
4.3 Diversity reward redistribution: Sign-aware redistribution concentrates positive signal on high contributors in more-diverse groups and negative signal on low contributors in less-diverse groups while preserving the original group reward on average.

5 Experiments

The evaluation measures utility–diversity Pareto frontiers across three biomolecular generation tasks and compares SGRPO with pretrained, GRPO, and memory-assisted baselines where applicable. Across decoding sweeps, SGRPO generally preserves or expands the frontier, with especially strong gains in pocket-based design and protein design.

Evaluation protocol: The evaluation decodes each model across operating points and summarizes each point by utility and set-level diversity.The resulting operating-point set is evaluated through its non-dominated subset, or Pareto frontier.
Evaluation protocol: Hypervolume measures the area dominated by the non-dominated frontier, while DIP and R2 quantify distance or shortfall relative to the ideal point.HV is higher-is-better; DIP and R2 are lower-is-better, and the experiment-specific reference point makes HV a within-experiment metric.
De novo small-molecule design: Coupled-SGRPO achieves the best de novo small-molecule frontier, with HV 0.0670, DIP 0.2542, and R2 0.0977.Its main advantage is slower diversity loss at the high-utility end rather than uniform improvement across all decoding settings.
Pocket-based small-molecule design: Pocket-based design shows SGRPO’s strongest benefit, retaining markedly higher within-pocket diversity at comparable utility and achieving the best HV, DIP, and R2.The advantage is most pronounced in the high-utility regime, where docking optimization otherwise concentrates generation on a few chemotypes.
De novo protein design: In de novo protein design, SGRPO reaches similarly high utility while preserving diversity closer to the pretrained model, achieving the best HV, DIP, and R2.GRPO and Memory-assisted GRPO improve utility while collapsing diversity, with Memory-assisted GRPO showing the most severe concentration.

6 Analysis

The analyses show that both components of SGRPO contribute to utility-diversity improvements, while its training dynamics preserve broader generation coverage than comparison methods. SGRPO remains effective with small groups and across nonzero diversity-reward weights.

6.1 Ablation Study: Removing the diversity reward produces the innermost utility-diversity curve, while adding group-level diversity pressure moves the curve outward.The ablation isolates the supergroup diversity reward from leave-one-out credit assignment.
6.1 Ablation Study: Leave-one-out credit assignment further expands the utility-diversity curve by rewarding rollouts that contribute to set-level diversity.Full SGRPO dominates the variant that uses group diversity without leave-one-out credit.
6.2 Training Dynamics of Generated Distributions: After 20 steps, SGRPO spreads across multiple sequence clusters, whereas GRPO and Memory-assisted GRPO move into relatively concentrated regions.The comparison uses shared two-dimensional UMAP embeddings from sequences sampled at the original, intermediate, and final checkpoints.
6.2 Training Dynamics of Generated Distributions: From 20 to 100 steps, SGRPO retains multiple clusters while GRPO contracts and Memory-assisted GRPO drifts toward a distant narrow region.The authors attribute the memory-assisted drift to discouraging revisits to high-density regions rather than directly optimizing set diversity.
6.3 Robustness to Diversity-Estimator Efficiency and Reward Weighting: SGRPO outperforms GRPO with small groups under a fixed total supergroup size and improves HV over GRPO for every tested nonzero λ.The strongest frontier expansion occurs at λ = 0.5.

7 Conclusion

SGRPO directly combines rollout-level utility with set-level diversity for biomolecular post-training. Across three generation settings, it improves the attainable utility-diversity Pareto frontier and suggests that diverse-set rewards can expand available operating points.

7 Conclusion: SGRPO combines rollout-level utility with set-level sample diversity in biomolecular post-training.It compares candidate sets generated under the same condition and redistributes set rewards through leave-one-out diversity contributions.
7 Conclusion: Across de novo small-molecule, pocket-based small-molecule, and de novo protein design, SGRPO improves the utility-diversity Pareto frontier over pretrained generators and RL baselines.The conclusion presents direct rewards for diverse generated sets as a practical way to expand available operating points.

A Full Training Procedure of SGRPO

The full SGRPO procedure samples same-condition supergroups, computes utility and diversity signals, redistributes diversity through leave-one-out contributions, and performs a PPO-style reference-regularized update. Its small-group analysis establishes partition consistency and concentration properties for the normalized pairwise diversity signal.

A Full Training Procedure of SGRPO: For each condition, SGRPO samples a same-condition supergroup, computes rollout utility and group diversity, and forms supergroup-relative advantages.The training procedure uses M groups of K rollouts per condition.
A Full Training Procedure of SGRPO: Leave-one-out contributions redistribute each group’s diversity reward to individual rollouts according to their effect on set diversity.The contribution for a rollout compares the group’s diversity with the diversity after removing that rollout.
A Full Training Procedure of SGRPO: The composed rollout reward weights utility by 1 − λ and redistributed group diversity by λ before advantage computation.The resulting supergroup-relative advantages are used in a PPO-style objective with KL regularization toward a reference policy.
A Full Training Procedure of SGRPO: All relative comparisons are centered within same-condition supergroups, separating policy-quality comparisons from variation in condition difficulty.The diversity signals are centered across M groups, and advantages across the corresponding MK rollouts.
B Properties of Small-Group Pairwise Diversity Rewards: Normalized pairwise diversity has partition consistency and concentration properties that support using small-group diversity as a training signal.Average diversity over randomly partitioned groups is an unbiased proxy for full-sample diversity, and the proxy becomes more stable as group size increases.
B Properties of Small-Group Pairwise Diversity Rewards: The normalized pairwise diversity is based on average pairwise dissimilarity under a biomolecular similarity function.For a set A, the diversity score is defined over pairwise relationships among its members.

B.1 Partition Consistency

For normalized pairwise diversity, uniformly random balanced partitions preserve the full-set diversity in expectation. Thus, averaging diversity across groups provides an unbiased same-condition proxy for the larger sample.

B.1 Partition Consistency: Proposition B.1 establishes partition consistency for normalized pairwise diversity under uniformly random balanced partitions.The set has size N = MK and is partitioned into M groups of size K.
B.1 Partition Consistency: The average diversity of random groups is exactly aligned in expectation with the diversity of the full same-condition sample.This shows that SGRPO’s small-group objective targets the intended full-sample diversity rather than an unrelated objective.

B.2 Concentration Around Full-Sample Diversity

The average diversity of small groups not only matches full-sample diversity in expectation but also concentrates around it as group size increases. This clarifies why larger groups provide a more stable diversity signal.

B.2 Concentration Around Full-Sample Diversity: Proposition B.2 shows that average small-group diversity concentrates around full-sample diversity for sufficiently large groups.The analysis assumes i.i.d. samples from the policy under a fixed condition and uniformly random balanced partitions.
B.2 Concentration Around Full-Sample Diversity: Both full-sample and averaged group diversity estimate the same population quantity under the stated sampling and partitioning assumptions.The proof treats full-sample diversity as a bounded order-two U-statistic and group averages as averages of independent bounded U-statistics.
B.2 Concentration Around Full-Sample Diversity: Increasing group size K improves the stability of the small-group diversity signal as a proxy for full-sample diversity.Partition consistency removes objective-level bias, while concentration governs proxy stability.
B.2 Concentration Around Full-Sample Diversity: The concentration result provides a theoretical role for K beyond unbiasedness: larger groups make the proxy less variable around full-sample diversity.The stated guarantee is a sufficient condition, not a claim that every group size performs equally well.

C.2 Training

GenMol-P is trained as a pocket-conditioned masked discrete diffusion model, then used as the initialization for coupled-GRPO and coupled-SGRPO post-training.

C.2 Training: GenMol-P is supervised-tuned on CrossDocked2020 pocket-ligand pairs using a pocket structural prefix and SAFE-tokenized ligands.Pocket backbones are converted into residue-level structural context before molecular diffusion training.
C.2 Training: The supervised objective trains recovery of clean SAFE tokens from corrupted ligand sequences conditioned on the pocket prefix.Training samples a diffusion time t and uses the masked discrete diffusion forward process.
C.2 Training: The reported GenMol-P checkpoint uses 8 H200 GPUs, bf16 precision, global batch size 3072, and AdamW with learning rate 3 × 10^-4.These are implementation settings for the supervised checkpoint.
C.2 Training: Coupled-GRPO and coupled-SGRPO are initialized from the verified GenMol-P checkpoint after 5,500 supervised optimization steps.The checkpoint is followed by reinforcement-learning post-training runs.

D Experimental Implementation Details

The experiments evaluate diffusion-based molecular generators with coupled policy optimization, explicit QED–SA utility, and set-level diversity across decoding sweeps and pocket-conditioned settings.

D. Experimental Implementation Details: GenMol generates SAFE molecular strings with masked discrete diffusion, iteratively denoising masked positions while retaining confident predictions.Post-training methods preserve this native diffusion sampler through coupled-GRPO.
D. Experimental Implementation Details: The de novo small-molecule rollout utility combines QED and transformed synthetic accessibility as u(x) = 0.6 QED(x) + 0.4 sSA(x).QED and sSA are high-is-better scores normalized to [0, 1].
D. Experimental Implementation Details: SGRPO adds internal set diversity based on one minus mean pairwise Tanimoto similarity between Morgan fingerprints, unlike coupled-GRPO’s rollout-only utility reward.This directly distinguishes the set-level reward from the utility-only baseline.
D. Experimental Implementation Details: Coupled-GRPO evaluates completed molecules using complementary diffusion masks, preserving the denoising interface while reducing policy-gradient variance relative to independent masks.Each token contributes once across the coupled pair under a partially observed context.
D. Experimental Implementation Details: The main Pareto evaluation uses six paired decoding settings and generates 1,000 molecules per model at each setting.Utility is averaged over valid molecules, while diversity is computed over the same valid set.
D. Experimental Implementation Details: For pocket-conditioned design, GenMol-P prepends continuous pocket embeddings to SAFE-token embeddings, and utility weights QED, sSA, and docking as 0.3, 0.2, and 0.5.The pocket comparison uses original GenMol-P, coupled-GRPO, and SGRPO without a memory-assisted baseline.
D. Experimental Implementation Details: Protein-design utility combines four normalized scores targeting distinct failure modes of unconstrained protein generation.The supplied implementation passage identifies naturalness and developability components, including solubility and liability.

E Discussion and Future Work

SGRPO provides a direct set-level diversity signal but requires additional rollout and diversity-computation overhead. Its scope treats utility and diversity as user-specified axes, leaving efficiency and task-specialized extensions for future work.

Computational Cost: SGRPO requires same-condition candidate groups, adding rollout overhead compared with objectives that score candidates independently.The analysis reports effectiveness with near-minimal group sizes, but notes that larger-scale efficiency remains an open direction.
Computational Cost: O(MK^2Csim) work computes within-group pairwise similarities for M groups of K rollouts, with the same matrix reused for group rewards and leave-one-out contributions.This reuse avoids a second pass over pairwise similarities for leave-one-out credit assignment.
Computational Cost: Memory-Assisted GRPO can incur greater similarity-lookup cost when its memory contains group-scale or larger numbers of stored index–bucket pairs.In one de novo small-molecule run, SGRPO uses K = 64 and M = 8, while the memory exceeds 300 index–bucket pairs after the first training step.
Computational Cost: Diversity-computation overhead remains a practical constraint, especially for non-pairwise metrics that may require batching, learned proxies, or approximate set summaries.For pairwise metrics, caching repeated similarities or reusable fingerprints, embeddings, and neighborhood structures could reduce cost.
Scope and Future Work: The experiments evaluate trade-offs against task-specific scalar utilities, while future work could expose individual utility components as additional controllable axes.SGRPO is intentionally decoupled from a particular generator, task, utility objective, or diversity metric, enabling specialized diversity notions and adaptive group construction.