Résumé technique : Génération de molécules physiquement cohérentes avec des modèles basés sur l'énergie

Énoncé du problème

Les états moléculaires à l'équilibre thermique sont régis par la distribution de Boltzmann, $p(x) \propto \exp(-E(x)/(k_BT))$, où la probabilité d'une configuration est déterminée par son énergie potentielle. Bien que les jeux de données standard pour la génération moléculaire (par exemple, QM9, GEOM-Drugs) contiennent des structures optimisées géométriquement, relâchées vers des minima d'énergie locaux, ils ne sont pas tirés de la distribution de Boltzmann ; leurs fréquences relatives reflètent des stratégies d'énumération plutôt que des poids thermodynamiques. Par conséquent, un paysage énergétique scalaire où les minima coïncident avec des configurations stables constitue la cible de modélisation naturelle.

Cependant, l'apprentissage et l'échantillonnage à partir de tels paysages restent des défis ouverts. Les modèles basés sur l'énergie (EBM) traditionnels nécessitent des données distribuées selon Boltzmann pour l'entraînement ou s'appuient sur des méthodes coûteuses de Monte Carlo par chaîne de Markov (MCMC) qui peinent à traverser les régions accidentées et non physiques entre le bruit et les structures valides. À l'inverse, les modèles de diffusion et d'appariement de flux (flow-matching) les plus avancés contournent ces problèmes en apprenant des champs de transport conditionnels au temps. Bien qu'efficaces, ces modèles perdent le biais inductif énergétique, ne fournissant aucune fonction d'énergie explicite pour évaluer, classer ou composer des configurations au moment de l'inférence.

Méthodologie : EBMol

Les auteurs introduisent EBMol, un modèle basé sur l'énergie pour la génération moléculaire 3D qui restaure le biais inductif énergétique sans nécessiter de simulation explicite pendant l'entraînement. Le cadre se compose de deux composants principaux :

1. Appariement de champ de restauration (RFM) pour l'entraînement

Pour apprendre un paysage énergétique $E_\theta(x)$ traitable sans données distribuées selon Boltzmann, les auteurs proposent l'Appariement de champ de restauration (Restoring Field Matching - RFM).

• Concept : Au lieu de traiter les points de données comme des extrémités de trajectoire (comme dans l'appariement de flux standard), le RFM étend l'intervalle d'interpolation à $t \in [-1, 1]$, passant à travers les données à $t=0$.
• Construction du champ cible : Le modèle est entraîné pour faire correspondre le gradient négatif de l'énergie, $-\nabla E_\theta(x_t)$, à un champ cible de « restauration » $u$. Ce champ pointe vers les données $x_0$ depuis les deux côtés de la variété (interpolation $t \in [0, 1]$ et extrapolation $t \in [-1, 0]$).
• Stationnarité : Une fonction de lissage $b(|t|)$ met à l'échelle le champ cible de manière à ce qu'il s'annule à $t=0$ (le point de données), garantissant que les points de données deviennent des minima locaux du paysage énergétique appris.
• Régularisation : Un terme de régularisation par atome ancre les valeurs d'énergie des points de données propres vers zéro pour éviter un biais artificiel de profondeur dans l'échantillonneur.
• Gestion de l'espace d'état : La méthode gère les états mixtes continus-discrets ($x = (c, p)$, où $c$ sont les coordonnées et $p$ les types d'atomes) en définissant des interpolants séparés : des trajectoires linéaires pour les coordonnées (alignées via un transport optimal) et des trajectoires basées sur la réflexion pour les caractéristiques catégorielles afin de rester dans le simplexe.

2. Algorithme de Langevin en miroir (MLA) pour l'échantillonnage

Pour échantillonner à partir de la distribution apprise $p_\theta(x) \propto \exp(-E_\theta(x)/\tau)$, les auteurs adaptent l'Algorithme de Langevin en miroir.

• Mises à jour unifiées : La dynamique de Langevin standard pousse les types d'atomes en dehors du simplexe de probabilité. Le MLA résout cela en effectuant des mises à jour dans un espace dual non contraint et en projetant en arrière via une application miroir.
  • Pour les coordonnées continues, le potentiel euclidien réduit le MLA à l'algorithme de Langevin non ajusté.
  • Pour les types d'atomes catégoriels, le potentiel d'entropie négative mappe le simplexe vers l'espace des logits, la fonction softmax projetant en arrière.
• Recuit parallèle : Pour pallier la difficulté de traverser des paysages énergétiques accidentés, les auteurs emploient le recuit parallèle avec plusieurs chaînes à différentes températures. Des échanges périodiques de Metropolis-Hastings permettent aux chaînes « chaudes » d'explorer largement et aux chaînes « froides » d'affiner, permettant une mise à l'échelle du calcul au moment de l'inférence.

Contributions clés

1. Appariement de champ de restauration (RFM) : Un objectif d'entraînement sans simulation qui façonne un paysage énergétique scalaire et additif par atome, avec les points de données comme minima locaux. Les auteurs fournissent des garanties théoriques concernant l'existence et l'unicité de la solution et caractérisent le comportement sous différentes forces de régularisation.
2. Cadre d'échantillonnage Langevin en miroir : Le premier cadre Langevin unifié pour géométries mixtes pour la modélisation générative profonde conjointe continu-discret, combinant les mises à jour de coordonnées et les mises à jour de types d'atomes contraintes par le simplexe en une seule étape, augmenté par le recuit parallèle.
3. Performances de pointe : EBMol obtient des résultats de pointe sur les benchmarks QM9 et GEOM-Drugs, démontrant que la modélisation basée sur l'énergie peut égaler la qualité de génération des modèles de diffusion/appariement de flux tout en offrant des capacités supplémentaires.

Résultats expérimentaux

Le modèle a été évalué sur QM9 (134k petites molécules) et GEOM-Drugs (430k molécules de type médicament).

• Qualité de génération : Sur QM9, EBMol (avec 1810 évaluations de fonctions de réseau) dépasse tous les modèles de référence en termes de stabilité des atomes et des molécules, atteignant respectivement 99,76 % et 97,09 %. Sur GEOM-Drugs, EBMol surpasse nettement le modèle de référence fort SLDM, atteignant 98,68 % de stabilité des atomes et 61,42 % de stabilité des molécules avec des budgets de calcul élevés.
• Cohérence physique : Le paysage énergétique appris sert de métrique de qualité principielle. EBMol atteint une énergie de relaxation médiane de 1,78 kcal/mol sur GEOM-Drugs, soit un ordre de grandeur inférieur aux modèles précédents à liaisons implicites et compétitif avec les modèles à liaisons explicites (Megalodon), malgré l'utilisation de 18 fois moins de paramètres et l'absence de supervision explicite des liaisons.
• Corrélation avec l'énergie physique : L'énergie apprise corrèle significativement ($r=0,65$) avec la diminution d'énergie lors du relâchement GFN2-xTB, permettant au modèle d'agir comme un filtre léger pour les géométries physiquement contraintes.
• Génération contrôlable :
  • Pilotage de la forme : En composant le potentiel appris avec des potentiels de forme différentiables (linéaire, planaire, sphérique), le modèle peut orienter la génération vers des géométries sous-représentées sans réentraînement.
  • Conception de connecteurs Zero-Shot : La nature additive de l'énergie permet une génération conditionnelle en fixant des fragments moléculaires. EBMol atteint un taux de réussite de 73,9 % dans des tâches de conception de connecteurs à travers des complexités variables sans entraînement spécifique à la tâche.

Importance et revendications

L'article revendique que EBMol démontre le potentiel des modèles basés sur l'énergie comme cadre polyvalent pour la génération moléculaire 3D. Son importance principale réside dans la restauration du biais inductif énergétique perdu par les modèles conditionnels au temps, permettant ainsi :

• Un score physiquement significatif : La capacité de classer et de filtrer les échantillons générés sur la base d'une fonction d'énergie explicite.
• Une génération composable : La capacité d'orienter la génération par composition de potentiels (par exemple, pilotage de la forme) et d'effectuer une génération conditionnelle zero-shot (par exemple, conception de connecteurs) sans réentraînement.
• Une génération de haute qualité : Atteindre une stabilité et une validité de pointe tout en générant des structures étonnamment proches des minima locaux physiques, rivalisant avec des modèles ayant une supervision explicite des liaisons.

Les auteurs notent des limitations, notamment la sensibilité de l'algorithme d'échantillonnage aux hyperparamètres et le besoin actuel de budgets de calcul importants (NFE) pour obtenir des résultats optimaux. Ils suggèrent que les travaux futurs pourraient se concentrer sur l'étalonnage de l'énergie par rapport à des grandeurs physiques (DFT/GFN2-xTB) et l'intégration d'échantillonneurs plus sophistiqués.