Distributional Priors Guided Diffusion for Generating 3D Molecules in Low Data Regimes

Ce travail présente GODD, un modèle de diffusion guidé par des priors structuraux qui permet de générer des molécules 3D dans des régimes de données rares en apprenant à partir de régions denses et en généralisant efficacement aux décalages structurels hors distribution.

L'essentiel

Imaginez que vous êtes un chef cuisinier très talentueux, mais qui n'a appris à cuisiner qu'avec des ingrédients très courants : des pommes de terre, du riz et des carottes. Vous êtes devenu un expert pour créer des plats délicieux avec ces ingrédients.

Maintenant, imaginez qu'un client vous demande de créer un plat à base d'un ingrédient très rare, que vous n'avez jamais vu, comme une baie exotique qui ne pousse que dans une jungle lointaine. Si vous essayez de cuisiner comme d'habitude, vous allez probablement échouer, car votre cerveau ne connaît pas les règles de cette baie.

C'est exactement le problème que les scientifiques rencontrent avec la création de nouvelles molécules (les briques de base de la vie et des médicaments) par ordinateur.

Le Problème : La Cuisine des Molécules

Les chercheurs utilisent des intelligences artificielles (des "chefs robots") pour inventer de nouvelles molécules. Ces robots apprennent en regardant des millions de molécules existantes. Le problème, c'est que la plupart des molécules connues sont très similaires (comme nos pommes de terre).

Si le robot essaie de créer une molécule avec une structure très rare (comme notre baie exotique), il échoue. Il essaie de copier ce qu'il connaît, et le résultat est soit une molécule qui n'existe pas physiquement, soit un échec total. C'est ce qu'on appelle un problème de "région de données sparse" (une zone où il y a très peu d'exemples).

La Solution : Le "GODD" (Le Chef avec une Carte Magique)

Les auteurs de cette paper ont créé un nouveau système appelé GODD. Voici comment cela fonctionne avec une analogie simple :

Au lieu d'essayer d'apprendre à cuisiner la baie exotique directement (ce qui est impossible sans l'avoir jamais vue), le GODD utilise une carte mentale (un "prior structurel").

1. L'Apprentissage Asymétrique (Le Dessinateur et le Peintre) :
   Imaginez un duo d'artistes :

   • Le Dessinateur (l'encodeur) : Il regarde seulement la forme de la baie exotique (la structure rare) et la transforme en un croquis abstrait, une "idée" de forme. Il ne regarde pas la couleur ou le goût, juste la géométrie.
   • Le Peintre (le décodeur) : Lui, il est un expert pour peindre des molécules complètes à partir de n'importe quel croquis.

   Le secret ? Le Dessinateur ne voit que la partie rare, mais le Peintre doit reconstruire toute la molécule. Cela force le système à comprendre la logique de la forme, pas juste à mémoriser des recettes.

2. La Boussole (Le Guide) :
   Une fois que le système a appris à comprendre la "forme" de ces structures rares, il utilise cette compréhension comme une boussole.

   Quand on lui demande de créer une molécule avec la baie exotique, il ne devine pas au hasard. Il prend la "boussole" (la carte mentale de la forme rare) et guide son processus de création pour s'assurer que le résultat final respecte cette forme, même s'il n'a jamais vu cette baie auparavant.

Pourquoi c'est génial ?

• Pas besoin de voir l'impossible : Le système n'a pas besoin de voir des milliers d'exemples de la baie exotique pour la comprendre. Il suffit de lui montrer la forme, et il déduit comment l'intégrer dans un plat complet.
• Résultats concrets : Dans leurs tests, ce système a réussi à créer des molécules valides et nouvelles avec des structures rares beaucoup mieux que les autres méthodes (environ 12,6% de succès en plus !).
• Utilité pour la médecine : Cela ouvre la porte à la découverte de nouveaux médicaments. Souvent, les médicaments les plus efficaces ont des structures très complexes et rares que les méthodes actuelles ne peuvent pas imaginer. GODD permet d'explorer ces territoires inconnus.

En résumé

C'est comme si vous donniez à un architecte une photo d'un château de sable très étrange que personne n'a jamais construit. Au lieu de lui dire "construis-le", vous lui donnez une règle géométrique qui explique pourquoi ce château tient debout. Grâce à cette règle, l'architecte peut construire ce château, même s'il n'a jamais vu de sable comme celui-là.

Le GODD est cette règle géométrique intelligente qui permet aux ordinateurs d'imaginer de nouvelles molécules là où les humains et les autres intelligences artificielles échouaient jusqu'ici.

Résumé technique

1. Problématique et Contexte

La génération de molécules 3D de novo repose souvent sur des modèles génératifs (comme les modèles de diffusion) qui nécessitent d'énormes quantités de données pour apprendre la distribution des structures moléculaires. Cependant, ces modèles peinent à générer des échantillons valides dans des régions de l'espace chimique peu denses (sparse regions), c'est-à-dire pour des structures rares ou absentes dans l'ensemble d'entraînement.

Ce défi est formulé comme un problème de génération hors distribution (Out-of-Distribution ou OOD). Contrairement aux travaux antérieurs qui se concentrent principalement sur les décalages de propriétés (ex: changer la solubilité), cet article aborde les décalages structurels (structural shifts), tels que l'apparition de nouveaux squelettes moléculaires (scaffolds), de nouveaux cycles (rings) ou de fragments rares. Les méthodes existantes, y compris les modèles conditionnels basés sur des prédicteurs de propriétés scalaires, échouent souvent à généraliser à ces nouvelles structures complexes lorsque les données sont rares.

2. Méthodologie : GODD (Geometric OOD Diffusion Model)

Les auteurs proposent GODD, un cadre novateur basé sur la diffusion qui utilise des priors structurels distributionnels pour guider la génération vers des régions de données rares sans nécessiter d'entraînement sur ces données spécifiques.

A. Auto-encodeur Asymétrique Équivariant (EAAE)

Le cœur de l'approche est un auto-encodeur asymétrique conçu pour capturer les priors structurels :

• Asymétrie : L'encodeur ($E$) opère uniquement sur une sous-structure (ex: un squelette ou un fragment OOD), tandis que le décodeur ($D$) reconstruit la molécule complète. Cela force le modèle à apprendre une représentation latente de la sous-structure qui est généralisable à des molécules entières.
• Équivariance SE(3) : Pour respecter les symétries physiques des molécules (invariance par translation et rotation), l'architecture repose sur des Réseaux de Neurones Graphiques Équivariants (EGNN).
  • L'encodeur produit des caractéristiques latentes équivariantes pour les coordonnées ($f_x$) et invariantes pour les features atomiques ($f_h$).
  • Le théorème 3.1 prouve que la fonction de perte de cet auto-encodeur est invariante sous le groupe SE(3), garantissant que les conditions extraites respectent les contraintes de symétrie nécessaires au modèle de diffusion.

B. Modèle de Diffusion Guidé par les Priors (DSDM)

Le modèle de diffusion conditionnel est entraîné pour générer des molécules en utilisant les latents $f_x$ et $f_h$ issus de l'EAAE comme conditions :

• Conditionnement : Les features latentes sont concaténées aux features des nœuds du graphe moléculaire au cours du processus de débruitage.
• Contrainte d'invariance : Le modèle de diffusion est conçu pour satisfaire une condition d'invariance stricte : la distribution de probabilité doit rester inchangée sous rotation et translation des conditions d'entrée.
• Fonction de perte : La perte totale combine la perte de reconstruction de l'auto-encodeur ($L_{EAAE}$) et la perte de diffusion ($L_{DSDM}$), formant une borne inférieure variationnelle invariante SE(3) de la vraisemblance (Théorème 3.2).

3. Contributions Clés

1. Première approche OOD structurelle : C'est la première étude à formuler la génération de molécules 3D dans des régions de données rares comme un problème de génération OOD sous décalage structurel, prouvant théoriquement l'équivariance des priors extraits.
2. Architecture EAAE : Conception d'un auto-encodeur asymétrique équivariant capable d'encoder des sous-structures (squelettes, cycles) en priors distributionnels généralisables, sans avoir besoin de données OOD pour l'entraînement.
3. Performance supérieure : Le cadre GODD améliore significativement le taux de réussite (success rate) par rapport aux modèles de diffusion unconditionnels, conditionnels et aux méthodes spécifiques OOD existantes.
4. Application au design de médicaments : Démonstration de l'efficacité du modèle sur des tâches de conception de fragments (linker design) dans des contextes OOD, surpassant les méthodes basées sur des fragments classiques.

4. Résultats Expérimentaux

Les expériences ont été menées sur les jeux de données QM9 et GEOM-DRUG, avec des tâches de génération de molécules basées sur des squelettes (scaffolds) et des cycles (rings) rares.

• Génération de Cycles (Ring-Structure) :
  • Sur QM9, pour des molécules avec 4 à 8 cycles (régions très rares), GODD atteint un taux de réussite de 40,5 %, contre moins de 26,2 % pour les meilleures méthodes de référence (comme CGD ou MOOD).
  • Les modèles de base (EDM, GeoLDM) échouent presque totalement (< 7 %) sur ces structures complexes.
• Génération de Squelettes (Scaffold) :
  • Pour les squelettes très rares (OOD II, < 10 occurrences), GODD génère des molécules avec une proportion de squelettes cibles > 95 % et une couverture de squelettes de 85,7 %, là où les méthodes fragmentaires (DiffLinker, LinkerNet) chutent drastiquement.
  • GODD améliore le taux de réussite de 12,6 % par rapport aux méthodes basées sur des fragments.
• Design de Linkers (Fragment-Based Drug Design) :
  • Sur la tâche de liaison de fragments (linker design) avec des structures OOD, GODD atteint un taux de validité de 65,2 % (contre < 50 % pour les concurrents) et montre de meilleures propriétés de "drug-likeness" (QED) et d'accessibilité synthétique (SA).

5. Signification et Impact

Ce travail démontre que l'intégration de priors structurels distributionnels via un auto-encodeur asymétrique équivariant permet de surmonter la barrière de la rareté des données en génération moléculaire.

• Généralisation : Le modèle apprend à transférer la connaissance des structures denses vers des structures rares, éliminant le besoin d'ensembles de données OOD massifs pour l'entraînement.
• Découverte de médicaments : La capacité à générer des molécules valides et novatrices avec des squelettes ou des fragments rares ouvre de nouvelles voies pour la découverte de médicaments, en particulier pour explorer des espaces chimiques sous-représentés.
• Robustesse théorique : La preuve formelle de l'invariance SE(3) assure que le modèle respecte les lois physiques fondamentales, ce qui est crucial pour la fiabilité des modèles génératifs en chimie computationnelle.

En résumé, GODD représente une avancée majeure pour la génération de molécules 3D en régimes de faible données, offrant une solution robuste aux défis posés par les décalages structurels complexes.