🤖 AI

NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation

L'article présente NEAT, un nouveau transformateur d'ensembles autorégressif guidé par les voisinages qui atteint une qualité de génération moléculaire 3D de pointe sur les ensembles de données QM9 et GEOM-Drugs, avec des avantages significatifs en termes de vitesse, en garantissant l'invariance par permutation au niveau des atomes grâce à une stratégie d'entraînement agnostique de l'ordre.

Auteurs originaux : Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

Publié 2026-05-07

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayiez d'enseigner à un robot comment construire une structure 3D complexe, comme un modèle moléculaire composé de boules de différentes couleurs (atomes) reliées par des bâtonnets (liaisons). L'objectif est que le robot crée de nouvelles structures valides à partir de rien, et non pas simplement qu'il copie des structures existantes.

Ce papier présente un nouvel architecte-robot appelé NEAT (Neighborhood-Guided, Efficient, Autoregressive Set Transformer). Voici comment il fonctionne, expliqué simplement :

Le Problème : La Confusion de la « File d'Attente »

La plupart des robots précédents tentaient de construire des molécules en suivant une file d'attente stricte, comme des élèves attendant d'entrer dans une classe un par un. Ils devaient décider : « Qui passe en premier ? Qui passe en second ? »

Le Problème : Dans une molécule, les atomes n'ont pas réellement de place « première » ou « seconde ». Ils existent tous simultanément dans un amas 3D. Les forcer dans une file d'attente, c'est comme essayer de décrire un flocon de neige en listant ses cristaux dans un ordre spécifique. C'est arbitraire et confus. Si le robot se trompe d'ordre, la molécule se brise.
L'Ancienne Méthode : Les méthodes précédentes tentaient de résoudre cela en inventant un ordre « canonique » (standard), mais cela biaisait le robot. Il apprenait mieux l'ordre de la liste que la forme de la molécule.

La Solution : La Stratégie de la « Voisinage »

NEAT change la donne. Au lieu d'une file, il traite la molécule comme une foule.

Comment ça marche : Imaginez que vous construisez une maison brique par brique. Au lieu de demander « Quelle brique vient ensuite dans la liste ? », NEAT regarde le mur actuel que vous avez construit et demande : « Qui sont les voisins se tenant juste à côté de ce mur ? »
La Magie : NEAT examine la « frontière » de la structure actuelle. Il sait que tout atome touchant la structure actuelle est un candidat valide pour être ajouté ensuite. Il ne se soucie pas de l'ordre ; il ne se soucie que du voisinage.
Le Résultat : Cela rend le robot « invariant par permutation ». Peu importe si vous regardez la molécule par la gauche, par la droite ou à l'envers ; le robot voit les mêmes voisins et construit la même structure valide.

Le Processus de Construction : Deux Étapes à la Fois

Lorsque NEAT décide d'ajouter un nouvel atome, il fait deux choses simultanément :

Qu'est-ce que c'est ? Il devine le type d'atome (par exemple : « Cela devrait être une boule de Carbone »).
Où est-il ? Il utilise une technique mathématique spéciale appelée Flow Matching pour déterminer exactement où dans l'espace 3D cet atome devrait flotter. Pensez-y comme un courant de rivière qui pousse doucement une feuille (l'atome) depuis un point de départ aléatoire vers son endroit parfait et stable dans la molécule.

Pourquoi NEAT est Spécial ?

L'article affirme que NEAT est une mise à niveau majeure pour trois raisons :

Il est Rapide : Parce qu'il n'a pas à simuler des milliers d'étapes bruyantes comme les autres robots (qui sont comme essayer de sculpter une statue en taillant un bloc de pierre pendant des heures), NEAT construit des molécules beaucoup plus vite. L'article indique qu'il est 5 à 16 fois plus rapide que ses concurrents.
Il est Flexible : Parce qu'il ne dépend pas d'une file d'attente fixe, vous pouvez lui donner un « kit de démarrage » (une molécule partielle) et lui demander de terminer le travail. Il peut compléter la structure peu importe la façon dont vous lui avez remis le kit de démarrage. D'autres robots échouent souvent si le kit de démarrage ne correspond pas à leurs règles spécifiques de « file d'attente ».
Il est Précis : Sur des tests standards (ensembles de données appelés QM9 et GEOM-Drugs), NEAT construit des molécules qui sont chimiquement valides et uniques, égalant ou surpassant les meilleures méthodes existantes.

Le Bémol (Limites)

Les auteurs sont honnêtes sur ce que NEAT ne peut pas encore faire :

Limites de Taille : Il ne peut construire que des molécules aussi grandes que celles qu'il a vues pendant l'entraînement. Il ne peut pas extrapoler pour construire une molécule géante s'il n'a appris que sur des petites.
Ambiguïté des Liaisons : Il dépend de la connaissance des atomes connectés (liaisons) pour déterminer le voisinage. Si vous lui donnez une structure cristalline étrange où les liaisons sont floues, il pourrait se confondre.
Pas de « Magie » des Règles 3D : Il n'a pas de règles physiques intégrées pour la rotation 3D (équivariance SE(3)), mais les auteurs ont constaté que simplement faire tourner aléatoirement les données d'entraînement suffisait à ce qu'il fonctionne bien de toute façon.

Résumé

En bref, NEAT est une nouvelle façon d'enseigner aux ordinateurs à concevoir des molécules 3D. Au lieu de les forcer à suivre une liste rigide et artificielle, il leur apprend à regarder leurs voisins immédiats et à construire naturellement. Il est plus rapide, plus flexible et tout aussi précis que les méthodes actuelles de l'état de l'art.

Résumé Technique : NEAT – Transformeur d'Ensemble Autoregressif, Efficace et Guidé par le Voisinage pour la Génération Moléculaire 3D

Énoncé du Problème

Les flux de travail computationnels actuels pour la découverte de molécules bioactives reposent fortement sur le criblage virtuel ou la génération de novo. Bien que le criblage virtuel soit limité par la taille finie des bibliothèques chimiques et par des coûts computationnels croissants, la modélisation générative offre une alternative prometteuse. Cependant, la génération de structures moléculaires 3D présente des défis uniques. Les molécules sont intrinsèquement tridimensionnelles, et leurs propriétés sont souvent déterminées par leur conformation.

Les modèles de génération 3D les plus avancés existants utilisent principalement des techniques de diffusion et d'appariement de flux (flow-matching). Ces méthodes souffrent souvent d'une mise à l'échelle quadratique dans les mises à jour de passage de messages, manquent d'adaptabilité inhérente aux tailles de molécules variables (nécessitant souvent des nombres d'atomes fixes déterminés à l'initialisation) et peinent à gérer les insertions dynamiques d'atomes pendant le débruitage.

Les modèles autoregressifs offrent une alternative flexible qui prend naturellement en charge des nombres d'atomes variables et ne nécessite qu'un seul appel d'encodeur par token. Cependant, les transformeurs autoregressifs standards requièrent un ordre séquentiel des tokens. Dans les graphes moléculaires, les atomes n'ont pas d'ordre séquentiel inhérent. Des travaux antérieurs ont résolu ce problème en utilisant des ordonnancements canoniques d'atomes, mais ces approches ne sont pas invariantes par permutation des atomes. Elles biaisent la prédiction du token suivant vers des conventions d'ordonnancement spécifiques, limitant la capacité du modèle à généraliser ou à effectuer des tâches telles que la complétion de préfixe où l'ordre des atomes de départ peut ne pas correspondre à l'ordre canonique de la molécule finale. À l'inverse, les approches basées sur les Réseaux de Neurones à Graphes (GNN) qui assurent l'invariance par permutation produisent souvent des molécules à faible validité chimique.

Méthodologie

Les auteurs proposent NEAT (Neighborhood-Guided, Efficient, Autoregressive Set Transformer), le premier modèle transformeur autoregressif pour la génération moléculaire 3D qui respecte l'invariance par permutation des atomes sans s'appuyer sur des ordonnancements canoniques d'atomes.

Architecture Principale

NEAT traite les graphes moléculaires comme des ensembles d'atomes plutôt que comme des séquences. L'architecture se compose de trois composants principaux :

Module d'Augmentation de Données (Guidage par le Voisinage) : Au lieu des encodages de position séquentiels, le modèle utilise une stratégie d'échantillonnage pour imiter la croissance moléculaire. Étant donné un graphe moléculaire, un sous-graphe connecté (ensemble source) est échantillonné. L'ensemble cible est défini comme les « nœuds frontières » (voisins à 1 saut) de l'ensemble source. Cela fournit un ensemble de continuations de tokens valides pour l'étape suivante, indépendamment de l'ordre dans lequel les atomes ont été ajoutés.
Tronc Transformeur d'Ensemble : L'ensemble source est encodé à l'aide d'un transformeur d'ensemble (basé sur NanoGPT/GPT-2 avec attention bidirectionnelle et sans encodages de position). Les types d'atomes sont encodés linéairement, et les positions sont encodées à l'aide de caractéristiques de Fourier. Un regroupement additif agrège ces représentations de nœuds en une représentation au niveau de l'ensemble ( $z_n$ ). Cette conception assure l'invariance par permutation de l'ordre des atomes dans l'ensemble source.
Tête d'Appariement de Flux : La représentation au niveau de l'ensemble $z_n$ $z_{n}$ est utilisée pour prédire le token suivant, factorisée en type d'atome et position : $p(a_{n+1}, x_{n+1} | z_n) = p(a_{n+1} | z_n) \cdot p(x_{n+1} | a_{n+1}, z_n)$ $p (a_{n + 1}, x_{n + 1} ∣ z_{n}) = p (a_{n + 1} ∣ z_{n}) \cdot p (x_{n + 1} ∣ a_{n + 1}, z_{n})$ .
- Type d'Atome : Modélisé comme une distribution catégorielle via une couche linéaire et un softmax.
- Position d'Atome : Modélisée via l'Appariement de Flux Conditionnel (CFM). Le modèle apprend un champ vectoriel pour transporter une distribution de position initiale aléatoire vers les positions atomiques réelles.

Stratégie d'Entraînement

Guidage par le Voisinage : Pendant l'entraînement, le modèle n'est pas supervisé par un seul token suivant, mais par l'ensemble de tous les nœuds frontières valides. Cela élimine le besoin de canonisation.
Forçage de l'Enseignant (Teacher Forcing) : La prédiction du type d'atome et la récupération de la position sont découplées. Le modèle prédit le type, et la tête d'appariement de flux apprend la position conditionnée à ce type.
Augmentation : Pour améliorer la robustesse face à l'accumulation d'erreurs, du bruit gaussien est injecté dans une fraction des atomes de l'ensemble source pendant l'entraînement, et les positions sont recentrées par rapport à l'ensemble source bruité.
Fonction de Perte : La perte totale est la somme de la perte d'entropie croisée pour les types d'atomes et de la perte d'appariement de flux pour les positions. Un problème d'affectation linéaire est résolu pour coupler les distributions source et cible, minimisant les coûts de transport et évitant les croisements de trajectoires.

Inférence

La génération commence par un atome initial aléatoire. À chaque étape, l'ensemble actuel d'atomes est encodé, le type du prochain atome est prédit, et une position initiale aléatoire est échantillonnée. Le champ vectoriel d'appariement de flux est ensuite intégré (en utilisant les méthodes d'Euler ou d'Euler-Maruyama) pour déterminer les coordonnées 3D finales. Ce processus se répète jusqu'à ce qu'un token d'arrêt soit prédit.

Contributions Clés

Génération Autoregressive Invariante par Permutation : NEAT introduit une stratégie d'entraînement novatrice utilisant le guidage par le voisinage et les transformeurs d'ensemble pour réaliser une génération moléculaire 3D invariante à l'ordonnancement des atomes, surmontant les limitations de l'ordonnancement canonique dans les modèles autoregressifs.
Efficacité : En traitant l'entrée comme un ensemble et en utilisant un seul appel d'encodeur par token, NEAT évite la mise à l'échelle quadratique du passage de messages basé sur la diffusion et le besoin de multiples étapes d'intégration par étape de génération d'atome trouvé dans certaines bases de référence d'appariement de flux.
Capacité de Complétion de Préfixe : La nature agnostique de l'ordre de NEAT permet une complétion de préfixe de haute qualité (génération de molécules autour d'un échafaudage donné) sans exiger que le préfixe s'aligne sur un ordonnancement canonique spécifique, une tâche où les modèles basés sur des séquences comme QUETZAL peinent.
Performances de l'État de l'Art : Le modèle obtient des résultats compétitifs ou supérieurs sur des benchmarks standards (QM9 et GEOM-Drugs) par rapport à la fois aux modèles autoregressifs et aux bases de référence de diffusion/appariement de flux.

Résultats Expérimentaux

Le modèle a été évalué sur les ensembles de données QM9 (petites molécules organiques) et GEOM-Drugs (molécules de type médicament).

Qualité de Génération : Sur QM9, NEAT a surpassé les modèles autoregressifs précédents (G-SchNet, G-SphereNet, Symphony, Geo2Seq) en termes de validité et d'unicité xyz2mol. Il a atteint une validité xyz2mol de 95,2 % et une unicité de 91,6 %, se classant deuxième en validité et unicité de recherche. Sur GEOM, NEAT a atteint la validité de recherche la plus élevée (99,7 %) et la deuxième meilleure validité xyz2mol (93,9 %) et unicité (91,6 %) parmi les modèles autoregressifs, égalant les performances des modèles d'appariement de flux de l'état de l'art comme FlowMol3.
Complétion de Préfixe : Dans une tâche générant 1 000 molécules à partir de 100 préfixes divers, NEAT a considérablement surpassé QUETZAL. NEAT a atteint 97,8 % de validité de recherche et 92,7 % de validité xyz2mol, tandis que QUETZAL n'a atteint que 27,8 % et 28,1 % respectivement. Cela démontre la capacité de NEAT à généraliser les caractéristiques structurelles indépendamment des conventions d'ordonnancement des atomes.
Efficacité Temporelle : NEAT démontre des avantages de vitesse significatifs. Il s'exécute environ 5 fois plus vite que SemlaFlow et 16 fois plus vite que FlowMol3. Dans les scénarios de pire cas, la complexité de NEAT est $O(BN^3)$ (où $N$ est le nombre d'atomes), tandis que les modèles de diffusion/appariement de flux évoluent comme $O(BTN^2)$ (où $T$ est le nombre d'étapes d'intégration). Puisque $N$ est généralement plus petit que $T$ pour les molécules de type médicament, NEAT offre un avantage de vitesse pratique.

Importance et Revendications

L'article revendique que NEAT fournit une base solide pour exploiter les architectures de transformeurs dans des tâches complexes de conception de médicaments, spécifiquement pour la génération conditionnelle de molécules 3D. En éliminant la dépendance aux ordonnancements canoniques d'atomes, le modèle permet des tâches flexibles comme la génération basée sur un échafaudage (complétion de préfixe) qui étaient auparavant difficiles pour les approches autoregressives.

Les auteurs notent que si l'architecture actuelle n'est pas explicitement équivariante SE(3), l'utilisation d'une augmentation par rotation aléatoire pendant l'entraînement compense efficacement cela, résultant en une variance mineure dans les prédictions. Le travail suggère que l'intégration de NEAT avec un transformeur équivariant SE(3) est une direction prometteuse pour la recherche future.

Les auteurs soulignent que si le modèle génère des structures 3D de haute qualité, la génération de molécules n'est qu'une étape dans un pipeline complexe nécessitant une validation extensive, une synthèse et des tests. Ils reconnaissent également des limitations, telles que la dépendance aux informations de bord pour le guidage par le voisinage (ce qui peut limiter l'applicabilité aux cristaux inorganiques) et l'absence d'extrapolation au-delà du nombre maximal d'atomes observé pendant l'entraînement.