Repousser les frontières d'utilité-diversité biomoléculaire avec l'optimisation relative de politique par supergroupe

Résumé

Les générateurs biomoléculaires sont souvent adaptés avec une rétroaction par récompense pour améliorer leur utilité spécifique à une tâche, mais maximiser uniquement l'utilité peut concentrer la génération sur une famille restreinte de candidats. Maintenir la diversité est difficile car la diversité d'un échantillon est une propriété au niveau de l'ensemble. Nous présentons l'Optimisation Relative de Politique par Supergroupe (SGRPO), un cadre flexible de type GRPO qui construit directement les récompenses à partir de la diversité au niveau des ensembles. Pour chaque condition, SGRPO échantillonne un supergroupe d'ensembles candidats, compare leur diversité sous la même condition, et redistribue la récompense de diversité du groupe aux roulages individuels via des contributions de diversité par leave-one-out, avant de la combiner avec l'utilité au niveau du roulage. Cette conception découple SGRPO d'un générateur, d'une récompense d'utilité ou d'une métrique de diversité particuliers, et permet son instanciation avec différentes approches de type GRPO. Nous évaluons SGRPO sur la conception de petites molécules de novo, la conception de petites molécules basée sur une poche, et la conception de protéines de novo, en l'instanciant avec GRPO et Coupled-GRPO sur des générateurs autorégressifs et à diffusion discrète. Sur les balayages de décodage, SGRPO élargit la frontière de Pareto utilité-diversité et atteint les meilleures métriques au niveau de la frontière par rapport aux générateurs pré-entraînés, à GRPO et à GRPO assisté par mémoire, lorsque applicable. Nos analyses montrent en outre que les récompenses directes de diversité au niveau des ensembles restent efficaces avec de petits groupes et aident à préserver une couverture plus large de la distribution de génération pendant le post-entraînement. Le code est disponible à 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.