Résumé Technique : Sur l'espace de conception de l'adaptation en ligne de la diffusion discrète pour l'optimisation moléculaire

Énoncé du Problème

L'optimisation moléculaire commence généralement par un modèle génératif pré-entraîné (par exemple, un modèle de diffusion discrète) qui capture un large a priori sur les structures chimiques valides. Cependant, l'objectif pratique au moment du test n'est pas d'échantillonner à partir de cet a priori, mais d'identifier des molécules qui maximisent une fonction de récompense spécifique à une tâche $r(x)$ (par exemple, l'affinité de liaison ou la fitness protéique) en utilisant un budget d'oracle limité. L'oracle, qui évalue la récompense, est souvent coûteux (par exemple, des essais en laboratoire humide ou des pipelines de docking coûteux).

Le défi central est l'adaptation de modèle en ligne : orienter un générateur pré-entraîné $p_{\theta_0}$ vers des régions de haute récompense de l'espace chimique tout en maintenant la validité et la diversité. Cela implique une boucle étroitement couplée où le système doit décider :

1. Acquisition : Quels candidats évaluer compte tenu du budget de l'oracle.
2. Adaptation : Comment convertir les récompenses observées en mises à jour du modèle.
3. Réutilisation du Feedback : Quelles données historiques conserver pour l'entraînement.
4. Débiaisage de l'A Priori : À quel point éloigner le générateur de ses modes pré-entraînés sans provoquer d'effondrement de la validité.

Bien que ces composants aient été étudiés isolément, il reste incertain s'ils se complètent, interfèrent ou deviennent redondants lorsqu'ils sont intégrés dans une boucle complète d'adaptation en ligne.

Méthodologie

Les auteurs proposent une étude contrôlée d'un « harnais » modulaire pour l'adaptation en ligne des modèles de diffusion discrète. La boucle (Figure 1) alterne entre l'échantillonnage de candidats, la sélection d'un lot pour l'évaluation par l'oracle et l'ajustement fin (fine-tuning) du générateur. L'étude évalue cinq choix de conception spécifiques de « plug-in » (Tableau 1) à travers neuf tâches : six tâches d'affinité de liaison de petites molécules et trois tâches de fitness protéique.

Composants de Conception Clés

1. Pilotage (Acquisition) :

   • Échantillonnage de Thompson : Au lieu d'une sélection gourmande basée sur la prédiction de la moyenne d'un modèle de substitution, les auteurs utilisent l'échantillonnage de Thompson pour sélectionner les candidats. Cela équilibre l'exploration et l'exploitation en échantillonnant des estimations de récompense $\hat{r}_i = \mu_i + \epsilon_i \sigma_i$, où $\sigma_i$ représente le désaccord de l'ensemble (incertitude). Cela encourage l'évaluation de candidats ayant une incertitude élevée même si leur prédiction moyenne est plus faible.

2. Adaptation (Mise en forme de la Récompense et Débiaisage) :

   • Mise en forme de la Récompense par CVaR : Pour se concentrer sur la queue de haute récompense plutôt que sur la moyenne, la récompense d'entraînement est mise en forme à l'aide de la Valeur à Risque Conditionnelle (CVaR). La récompense est transformée en $\tilde{r}(x) = (r(x) - Q_\tau)_+ / (1-\tau)$, où $Q_\tau$ est le quantile $\tau$ (fixé à 0,8). Cela concentre les mises à jour de gradient sur les échantillons les plus performants.
   • Régularisation de l'Entropie de Densité (Débiaisage du Modèle) : Pour empêcher le modèle de s'effondrer dans les modes de haute densité déjà existants de l'a priori, un régularisateur est ajouté : $\log \tilde{r}_{debias}(x) = \log r(x) - \gamma \log p_{\theta_t}(x)$. Comme le calcul de la vraisemblance marginale exacte $\log p_{\theta_t}(x)$ pour la diffusion discrète est intraitable, les auteurs utilisent un estimateur de Monte Carlo efficace en termes de calcul, basé sur l'ELBO de la diffusion (une approximation à passage unique), pour approximer la log-vraisemblance négative.

3. Corrections de l'Inadéquation d'Objectif (Objective-Mismatch) :

   • Tampon de Rejeu (Replay Buffer) : Pour empêcher le modèle d'oublier les découvertes de haute récompense faites lors des premiers tours à mesure que la distribution d'échantillonnage change, un tampon FIFO à capacité fixe avec éviction prioritaire (supprimant d'abord les échantillons de plus faible récompense) est utilisé. Les mini-lots sont tirés via un échantillonnage stratifié pour s'assurer que les valeurs aberrantes de haute récompense ne sont pas diluées.
   • Pénalité de Sortie Invalide : Pour maintenir la validité (par exemple, des chaînes SMILES valides), les candidats qui échouent au parsing (par exemple, par RDKit) se voient attribuer une pénalité négative fixe $r_{inv}$. Cela ramène le générateur vers le manifold moléculaire valide lorsque l'exploration produit des structures invalides.

Configuration Expérimentale

• Tâches : Optimisation de petites molécules pour la liaison à six cibles protéiques (en utilisant GenMol comme générateur et FlashAffinity/Boltz-2 comme oracles) et optimisation de la fitness protéique pour trois familles (GB1, CreiLOV, TrpB).
• Baselines : La boucle complète en ligne est comparée à :
  • L'ajustement fin hors ligne (entraînement sur un jeu de données statique).
  • La recherche au moment de l'inférence (Beam Search, MCTS, Best-of-N, etc.) avec un générateur fixe.
  • Des versions ablatives de la boucle en ligne (supprimant un composant à la fois).
• Métriques : La performance est mesurée par la meilleure récompense trouvée, les moyennes top-k et la récompense normalisée, évaluées par rapport aux appels d'oracle cumulés et aux heures GPU. Les métriques secondaires incluent la validité, la diversité (Tanimoto), le QED et la l'accessibilité synthétique (SA).

Résultats Clés

1. Complémentarité des Composants

L'étude conclut que les composants proposés apportent des bénéfices complémentaires, particulièrement pour les tâches de petites molécules où les candidats de haute récompense se situent loin de l'a priori pré-entraîné.

• Petites Molécules : L'acquisition (échantillonnage de Thompson), la mise en forme CVaR et la régularisation de l'entropie de densité produisent les gains de récompense les plus clairs. L'échantillonnage de Thompson améliore l'allocation des appels d'oracle ; la mise en forme CVaR concentre les mises à jour sur la queue de la récompense ; et le débiaisage pousse l'exploration loin des modes de l'a priori suroccupés.
• Fitness Protéique : Les effets d'ablation sont généralement plus faibles. Les a prisri spécifiques aux familles pour les protéines couvrent déjà les régions de haute fitness, rendant les changements de distribution importants moins critiques. Cependant, l'échantillonnage de Thompson apporte toujours des gains constants en améliorant l'allocation de l'oracle.

2. Mécanismes de Stabilité

• Les Tampons de Rejeu (Replay Buffers) et les Pénalités de Validité agissent comme des stabilisateurs plutôt que comme des moteurs primaires de récompense. Le tampon de rejeu aligne les gradients avec la distribution moyenne des données évaluées, empênant la perte des premières découvertes de haute récompense. La pénalité de validité empêche le « détournement de récompense » (reward hacking) où le modèle génère des chaînes chimiquement invalides mais hautement scorées, un mode de défaillance qui survient rapidement sans la pénalité.

3. Adaptation en Ligne vs Recherche au Moment de l'Inférence

La boucle complète en ligne surpasse significativement les baselines de recherche au moment de l'inférence (ex: Beam Search, MCTS) sous des budgets d'appels d'oracle et d'heures GPU équivalents.

• Efficacité : Sur les tâches de petites molécules, la boucle en ligne atteint le score final de la meilleure baseline de recherche avec environ 58 fois moins d'appels d'oracle. Sur les tâches de protéines, l'accélération est de 16 fois.
• Raisonnement : Les méthodes de recherche dépensent le budget pour générer et scorer des candidats à partir d'un générateur fixe. En revanche, l'adaptation en ligne réutilise le feedback de l'oracle pour mettre à jour le générateur lui-même, multipliant la valeur de chaque appel d'oracle. Les méthodes de recherche consomment également un calcul important pour la génération, laissant peu de temps pour les étapes de gradient de l'ajustement fin dans un budget de temps réel fixe.

4. Analyse du Changement de Distribution

Les gains de performance sont les plus prononcés lorsque la solution optimale nécessite un changement important par rapport à l'a priori pré-entraîné.

• La Régularisation de l'Entropie de Densité est le principal moteur du changement de distribution, poussant la masse du générateur loin des régions de haute densité de l'a priori vers des zones de faible densité et de haute récompense.
• Le CVaR et l'Échantillonnage de Thompson facilitent ce changement de manière conditionnelle, selon l'emplacement des candidats de haute récompense par rapport à l'a priori.
• Les tâches de protéines montrent des changements de distribution plus faibles car les a prisri spécifiques aux familles sont déjà concentrés près des régions de haute fitness.

Signification et Revendications

Le papier prétend fournir une recette pratique pour une optimisation moléculaire efficace en termes de feedback : l'ajustement fin en ligne avec acquisition, mise en forme de la récompense, débiaisage, rejeu et contrôle de la validité.

• Nouveauté : Ce travail intègre et évalue systématiquement ces composants au sein d'un cadre de diffusion discrète unique, démontrant qu'ils ne sont pas redondants mais mutuellement renforçants.
• Efficacité : La méthode proposée surpasse à la fois l'ajustement fin hors ligne et les baselines de recherche au moment de l'inférence, particulièrement dans les scénarios où l'oracle est coûteux et les molécules cibles sont hors distribution par rapport au modèle pré-entraîné.
• Robustesse : La recette est montrée comme robuste à travers différents objectifs d'ajustement fin (DDPP-LB et VIDD) et différents domaines (petites molécules vs protéines), bien que l'ampleur des gains varie selon la distance entre l'a priori et la distribution cible.

Les auteurs concluent que, bien que les composants soient plus critiques pour la conception de petites molécules (où l'a priori est large et la cible est distante), le cadre offre une approche généralisable pour orienter les modèles génératifs discrets dans des tâches de découverte scientifique avec un feedback limité.