Expandindo as Fronteiras de Utilidade-Diversidade Biomolecular com Otimização de Política Relativa de Supergrupo

Resumo

Geradores biomoleculares são frequentemente adaptados com feedback de recompensa para melhorar a utilidade específica da tarefa, mas maximizar apenas a utilidade pode concentrar a geração em uma família restrita de candidatos. Manter a diversidade é difícil, pois a diversidade amostral é uma propriedade do nível do conjunto. Apresentamos a Otimização Relativa de Políticas de Supergrupo (SGRPO), uma estrutura flexível no estilo GRPO que constrói diretamente recompensas a partir da diversidade no nível do conjunto. Para cada condição, a SGRPO amostra um supergrupo de conjuntos candidatos, compara sua diversidade sob a mesma condição e redistribui a recompensa de diversidade do grupo para trajetórias individuais por meio de contribuições de diversidade leave-one-out, antes de combiná-la com a utilidade do nível da trajetória. Esse design desacopla a SGRPO de um gerador, recompensa de utilidade ou métrica de diversidade específicos, e permite a instanciação com diferentes abordagens no estilo GRPO. Avaliamos a SGRPO no design de pequenas moléculas de novo, design de pequenas moléculas baseado em bolsas e design de proteínas de novo, instanciando-a tanto com GRPO quanto com GRPO Acoplado em geradores autorregressivos e de difusão discreta. Em varreduras de decodificação, a SGRPO expande a fronteira de Pareto utilidade-diversidade e alcança as melhores métricas no nível da fronteira em relação a geradores pré-treinados, GRPO e GRPO assistida por memória, quando aplicável. Nossas análises mostram ainda que recompensas diretas de diversidade no nível do conjunto permanecem eficazes com grupos pequenos e ajudam a preservar uma cobertura mais ampla da distribuição de geração durante o pós-treinamento. O código está disponível em https://github.com/IDEA-XL/SGRPO.

English

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.