Pushing Biomolecular Utility-Diversity Frontiers with Supergroup Relative Policy Optimization

2026-05-08

AI Summary

This work addresses the challenge in biomolecular generation where optimizing task utility alone often yields solutions lacking diversity, and set-level diversity is difficult to incorporate into existing reinforcement learning frameworks. To overcome this, the authors propose Super-Group Relative Policy Optimization (SGRPO), which for the first time formulates set diversity as a differentiable reward signal. SGRPO constructs candidate sets and evaluates each molecule’s contribution to diversity via leave-one-out estimation, jointly optimizing this diversity reward alongside utility. The framework is agnostic to the underlying generative model, utility function, and diversity metric, and supports both autoregressive and discrete diffusion generators. Evaluated on de novo small-molecule design, pocket-conditioned design, and de novo protein design, SGRPO and its coupled variant consistently expand the utility–diversity Pareto frontier, outperforming pretrained models, GRPO, and memory-augmented GRPO across multiple decoding settings while maintaining strong distributional coverage even with small candidate sets.

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.

Problem

• biomolecular generation
• utility-diversity trade-off
• set-level diversity
• molecular design
• generative models

Innovation

• Supergroup Relative Policy Optimization
• set-level diversity
• utility-diversity trade-off
• molecular generation
• Pareto frontier