Résumé Technique : Pousser les Frontières Utilité-Diversité Biomoléculaires avec l'Optimisation Relative de Politique de Supergroupe

1. Énoncé du Problème

La génération biomoléculaire vise à produire des candidats (petites molécules ou protéines) satisfaisant des objectifs chimiques ou biologiques spécifiques. Bien que l'apprentissage par renforcement (RL) constitue un cadre standard pour l'entraînement postérieur de générateurs préentraînés afin de maximiser l'utilité (par exemple, similarité avec des médicaments, scores d'amarrage, stabilité), l'optimisation de l'utilité seule conduit souvent à un effondrement de mode, où le générateur concentre la masse de probabilité sur une famille étroite de candidats à haut score. À l'inverse, le maintien de la diversité est difficile car la diversité est une propriété au niveau de l'ensemble (définie sur des collections d'échantillons), alors que l'optimisation standard de la politique met à jour des trajectoires individuelles.

Les méthodes existantes de RL favorisant la diversité reposent souvent sur des proxies indirects, tels que des pénalités de nouveauté par rapport à des échantillons historiques ou des filtres basés sur la mémoire. Ces approches peuvent pénaliser excessivement des modes utiles à haute densité ou induire une dérive distributionnelle. Le défi central abordé par cet article est le suivant : Comment pouvons-nous optimiser directement la diversité d'un ensemble d'échantillons comme objectif de première classe tout en attribuant toujours un crédit utile aux candidats générés individuellement ?

2. Méthodologie : Optimisation Relative de Politique de Supergroupe (SGRPO)

Les auteurs proposent l'Optimisation Relative de Politique de Supergroupe (SGRPO), un cadre flexible inspiré de l'Optimisation Relative de Politique de Groupe (GRPO) qui optimise directement le compromis entre l'utilité au niveau de la trajectoire et la diversité au niveau de l'ensemble.

Mécanisme Central

SGRPO fonctionne en échantillonnant plusieurs ensembles de candidats (groupes) sous la même condition et en les comparant pour redistribuer les récompenses de diversité. Le processus implique quatre étapes clés :

1. Échantillonnage de Supergroupe à Même Condition :
   Pour une condition donnée $C$ (par exemple, une poche protéique spécifique ou aucune condition), la politique échantillonne $M$ groupes, chacun contenant $K$ trajectoires générées indépendamment. Cet ensemble $S(C) = \{G_1, \dots, G_M\}$ forme un « supergroupe » de $N=MK$ candidats. Les comparaisons sont restreintes à ce supergroupe pour éviter de confondre la qualité de la politique avec la difficulté intrinsèque de différentes conditions.

2. Évaluation de la Diversité au Niveau du Groupe :
   Chaque groupe $G_m$ se voit attribuer un score de diversité $R_m = D(G_m)$ basé sur une métrique spécifiée par l'utilisateur (par exemple, distance Tanimoto interne pour les molécules, similarité Levenshtein pour les protéines). Le signal de diversité relative du groupe est calculé via une comparaison leave-one-out (un élément retiré) par rapport aux autres groupes du supergroupe :
   $$A_m^{grp} = R_m - \frac{1}{M-1}\sum_{h \neq m} R_h$$
   Un $A_m^{grp}$ positif indique que le groupe est plus diversifié que ses alternatives.

3. Redistribution Sensible à l'Ensemble (Contribution Leave-One-Out) :
   Pour combler l'écart entre les signaux au niveau de l'ensemble et les mises à jour de politique individuelles, SGRPO redistribue la récompense de diversité du groupe vers les trajectoires individuelles en fonction de leur contribution à la diversité du groupe.

   • Pour chaque trajectoire $x_{m,i}$, une contribution leave-one-out est calculée : $c_{m,i} = D(G_m) - D(G_m \setminus \{x_{m,i}\})$.
   • Ces contributions sont standardisées et utilisées pour former des poids de softmax sensibles au signe.
   • Si un groupe est très diversifié ($A_m^{grp} > 0$), le signal positif est concentré sur les candidats ayant des contributions de diversité élevées. Si un groupe est moins diversifié ($A_m^{grp} < 0$), le signal négatif est concentré sur les candidats ayant de faibles contributions.
   • Cela garantit que la récompense moyenne du groupe est préservée tout en attribuant un crédit aux molécules spécifiques qui favorisent la diversité.

4. Mise à Jour de Politique Relative au Supergroupe :
   La récompense finale pour une trajectoire combine son utilité individuelle $r_{m,i}$ et la récompense de diversité redistribuée $\tilde{R}_{m,i}$ :
   $$\bar{r}_{m,i} = (1-\lambda)r_{m,i} + \lambda \tilde{R}_{m,i}$$
   où $\lambda$ contrôle le compromis utilité-diversité. La fonction d'avantage est calculée comme une référence leave-one-out sur toutes les trajectoires du supergroupe, et la politique est mise à jour à l'aide d'un objectif de type PPO tronqué avec régularisation KL.

3. Contributions Clés

• Optimisation Directe au Niveau de l'Ensemble : SGRPO est le premier cadre à optimiser directement la diversité au niveau de l'ensemble comme objectif principal dans la génération biomoléculaire, plutôt que de s'appuyer sur des proxies de nouveauté indirects ou des pénalités dépendantes de l'historique.
• Conception Découplée : La méthode est découplée des architectures de générateurs spécifiques (autoregressives ou de diffusion), des récompenses d'utilité ou des métriques de diversité, permettant une instanciation avec divers optimiseurs de style GRPO (par exemple, GRPO standard pour les modèles autoregressifs, Coupled-GRPO pour la diffusion discrète).
• Mécanisme d'Attribution de Crédit : L'introduction d'un mécanisme de contribution de diversité leave-one-out permet au système d'attribuer un crédit aux candidats individuels qui soutiennent véritablement la diversité de l'ensemble, résolvant le problème de l'optimisation d'une propriété d'ensemble avec des mises à jour individuelles.
• Robustesse à la Taille du Groupe : La méthode reste efficace même avec de petits groupes, la rendant computationnellement réalisable pour la conception biomoléculaire à grande échelle.

4. Résultats Expérimentaux

Les auteurs ont évalué SGRPO dans trois contextes distincts de génération biomoléculaire :

1. Conception de Nouvelles Petites Molécules : Utilisation de GenMol (diffusion discrète).
2. Conception de Petites Molécules Basée sur des Poches : Utilisation de GenMol-P (diffusion conditionnée par la poche).
3. Conception de Nouvelles Protéines : Utilisation de ProGen2 (autoregressif).

Métriques de Performance :
L'évaluation s'est concentrée sur la Frontière de Pareto Utilité-Diversité, mesurée par :

• Hypervolume (HV) : La zone dominée par les points de fonctionnement non dominés.
• Distance au Point Idéal (DIP) : Distance vers la meilleure utilité et diversité théoriques.
• Indicateur R2 : Moyenne pondérée du pire déficit.

Constats :

• Expansion de la Frontière : Dans toutes les tâches et balayages de décodage, SGRPO a constamment élargi la frontière de Pareto utilité-diversité accessible par rapport aux générateurs préentraînés, au GRPO standard et au GRPO assisté par mémoire.
• Préservation de la Diversité : Dans les régimes à haute utilité où les lignes de base (en particulier GRPO) souffraient d'un effondrement sévère de la diversité, SGRPO a maintenu une diversité significativement plus élevée. Par exemple, dans la conception basée sur des poches, SGRPO a déplacé la trajectoire du point de fonctionnement vers l'extérieur, conservant une diversité nettement plus élevée à des niveaux d'utilité comparables.
• Études d'Ablation : La suppression de la récompense de diversité a ramené les performances aux niveaux du GRPO standard. La suppression de l'attribution de crédit leave-one-out a affaibli les performances, confirmant l'importance de redistribuer les signaux de diversité vers les contributeurs spécifiques.
• Dynamiques Distributionnelles : Les visualisations des distributions de séquences ont montré que tandis que GRPO et GRPO assisté par mémoire se contractaient dans des régions étroites ou dérivaient, SGRPO explorait et préservait plusieurs clusters, indiquant une couverture plus large de l'espace des séquences.

5. Signification et Revendications

L'article revendique que SGRPO offre un principe d'entraînement postérieur largement applicable pour élargir les points de fonctionnement disponibles pour les modèles de génération biomoléculaire. En récompensant directement des ensembles générés diversifiés, la méthode repousse la frontière utilité-diversité vers l'extérieur, permettant aux modèles d'atteindre une meilleure utilité à des niveaux de diversité fixes (et vice versa) sans sacrifier la couverture.

Les auteurs soulignent que cette approche adresse une limitation fondamentale dans la conception biomoléculaire basée sur le RL actuelle : l'incapacité à traiter la diversité comme un objectif de première classe lors de l'optimisation de la politique. Les résultats suggèrent que l'optimisation directe de la diversité au niveau de l'ensemble est une stratégie pratique et robuste pour atténuer l'effondrement de mode dans les tâches de génération de petites molécules et de protéines. Le code est mis à disposition pour faciliter de futures recherches et instanciations dans différents domaines.