AtomGen : rationalisation de la modélisation atomistique grâce à l’intégration d’ensembles de données et de benchmarks

Par Ali Kore, Amrit Krishnan et David Emerson

Le projet AtomGen vise à améliorer les capacités de modélisation atomistique grâce à des techniques d'apprentissage automatique avancées. En intégrant des modèles d'apprentissage profond à de vastes ensembles de données atomistiques, AtomGen ambitionne d'améliorer la prédiction et la compréhension des propriétés des systèmes à différentes échelles, des petites molécules organiques aux systèmes adsorbat-catalyseur.

Atomformer : Un transformateur pour les systèmes atomistiques 3D

Au cœur du projet AtomGen se trouve Atomformer, un modèle de transformateur composé uniquement d'un encodeur, adapté aux structures atomistiques tridimensionnelles. Inspiré de l'architecture Uni-mol+ , Atomformer intègre directement les informations spatiales 3D dans l'architecture du modèle grâce à des plongements positionnels gaussiens par paires. Plus précisément, cette méthode exploite les distances euclidiennes entre les atomes et un noyau gaussien prenant en compte le type de paire, projeté et ajouté au masque d'attention. Elle intègre également des plongements de métadonnées (masse atomique, rayon, valence, etc.) basés sur l'espèce atomique. Cette approche permet au modèle de capturer les relations spatiales entre les atomes, ce qui est essentiel pour une prédiction précise des propriétés.

S2EF-15M : Un ensemble diversifié de données d'atomes avec énergie et forces

Pour entraîner le modèle de base d'Atomformer , l'équipe d'AtomGen a compilé l'ensemble de données S2EF-15M , un ensemble à grande échelle de structures atomiques, de leurs forces et de leurs énergies. Cet ensemble de données regroupe des informations provenant de l' Open Catalyst Project (OC20, OC22, ODAC23), du Materials Project Trajectory Dataset (MPtrj) et de SPICE 1.1.4 . La figure 1 ci-dessous illustre un ensemble d'échantillons provenant de chacun des ensembles de données constitutifs.

Figure 1. « Visualisations des structures atomiques des composants de l'ensemble de données S2EF-15M : ODAC23 (en haut à gauche), MPtrj (en haut à droite), SPICE (au centre), OC20 (en bas à gauche) et OC22 (en bas à droite). »

L'ensemble de données S2EF-15M contient plus de 15 millions de systèmes (suréchantillonnage 2x MPtrj et 3x SPICE), offrant une grande diversité d'environnements atomiques pour l'entraînement. Il comprend les numéros atomiques, les coordonnées 3D, les forces, l'énergie de formation, l'énergie totale et un indicateur booléen signalant la présence de données d'énergie de formation valides. L'ampleur et la diversité de cet ensemble de données, combinées à l'apprentissage multitâche des énergies et des forces, permettent un apprentissage complet dans différents types d'environnements chimiques.

Traitement et intégration des données

L'une des contributions techniques majeures du projet AtomGen réside dans le développement de chaînes de traitement de données efficaces pour chaque ensemble de données source. Ces chaînes assurent la conversion de divers formats de données en une structure unifiée compatible avec la bibliothèque HuggingFace Datasets. Elles intègrent le traitement parallèle pour les grands ensembles de données, gèrent les formats de données compressés, convertissent les unités, calculent les propriétés dérivées et gèrent le remplissage dynamique et le traitement par lots des systèmes de taille variable.

Figure 2. « Initialisation et utilisation d'un pipeline de données de modélisation atomique utilisant AtomTokenizer et DataCollatorForAtomModeling pour l'objectif MaskGIT, avec génération aléatoire d'ensembles de données pour le traitement par lots. »

Intégration HuggingFace

Tous les ensembles de données et les modèles préentraînés sont téléchargés sur la plateforme huggingface.

AtomGen exploite l'écosystème HuggingFace pour le développement de modèles et la gestion des ensembles de données. Le projet comprend des implémentations personnalisées d' AtomformerConfig , d'AtomformerModel et de divers modèles spécifiques à certaines tâches, comme AtomformerForMaskedAM et Structure2EnergyAndForces . Ces implémentations permettent une intégration transparente avec l'API Trainer et le hub de modèles de HuggingFace, facilitant ainsi la reproductibilité et la simplicité d'utilisation. 

Techniques de collecte de données et de préentraînement

La classe DataCollatorForAtomModeling est une composante technique essentielle d'AtomGen. Elle offre une grande flexibilité dans la préparation des données et prend en charge diverses techniques de pré-entraînement. Ce collecteur de données personnalisé gère la complexité du traitement par lots, notamment le remplissage dynamique des séquences d'entrée, le calcul des masques d'attention et, en option, le calcul des encodages positionnels laplaciens. Ces encodages sont utilisés comme identifiants de nœuds avec l'implémentation TokenGT .

Le collecteur de données prend en charge plusieurs objectifs de pré-entraînement auto-supervisé. Les objectifs de préentraînement basés sur les coordonnées et les atomes peuvent être utilisés simultanément ou séparément.

1. Modélisation atomique masquée (MAM) : Semblable à la modélisation du langage masqué de BERT , cette technique masque aléatoirement les identités atomiques que le modèle doit prédire, l’aidant ainsi à apprendre le contexte et les interactions atomiques. Définir l'argument `mam` sur un nombre à virgule flottante masquera aléatoirement la fraction d'atomes correspondante.
2. Débruitage des coordonnées : Cette technique consiste à perturber les coordonnées atomiques et à demander au modèle de retrouver leurs positions initiales, améliorant ainsi sa compréhension de la géométrie moléculaire. Le degré de perturbation est ajustable via l’hyperparamètre ` coords_perturb` , un nombre à virgule flottante représentant l’écart type d’une distribution normale de bruit centrée sur zéro, ajouté aux coordonnées.
3. Objectif MaskGIT : Inspirée de l’approche MaskGIT en génération d’images, cette technique applique une stratégie de masquage plus complexe aux atomes, ce qui peut améliorer la compréhension de la structure globale du système par le modèle. Cet objectif est activé en définissant l’argument « mam » sur True, ce qui permet d’utiliser cette stratégie à la place de l’objectif MAM standard.

Ces diverses techniques de pré-entraînement permettent à Atomformer d'apprendre des représentations riches des structures atomistiques à partir de multiples perspectives, ce qui peut potentiellement améliorer ses performances sur les tâches en aval.

Approche de pré-entraînement et de mise au point

Bien que le collecteur de données prenne en charge diverses techniques de préentraînement autosupervisé, comme mentionné précédemment, celles-ci n'ont pas été utilisées dans les expériences de préentraînement actuelles. L'exploration de ces différentes approches autosupervisées constitue une piste de recherche prometteuse. AtomGen utilise une approche en deux étapes : un pré-entraînement sur l'ensemble de données S2EF-15M, suivi d'un ajustement précis sur des tâches spécifiques. La phase de pré-entraînement utilise le modèle Structure2EnergyAndForces, qui prédit à la fois les forces par atome et l'énergie par système. Le modèle a été entraîné sur un système 4xA40 pendant deux semaines, pour un total de 1 344 heures de calcul GPU. 

Pour un réglage précis, AtomGen utilise la suite de tests ATOM3D , qui comprend diverses tâches de modélisation atomistique telles que les propriétés des petites molécules (SMP), la prédiction de la stabilité des mutations (MSP) et l'affinité de liaison des ligands (LBA). Chacune des huit tâches d'ATOM3D est prétraitée et formatée pour être compatible avec l'outil de collecte de données et la bibliothèque de jeux de données HuggingFace.

Performances et évolutivité

Le projet AtomGen comprend des optimisations pour l'entraînement sur des ensembles de données à grande échelle, notamment la mise en place de points de contrôle de gradient pour un entraînement économe en mémoire, l'entraînement en précision mixte pour un calcul plus rapide et la prise en charge de l'entraînement distribué pour les configurations multi-GPU.  

Sur la tâche SMP, le modèle Atomformer affiné a atteint une erreur absolue moyenne (MAE) de 1,077, contre 1,13 pour un modèle entraîné à partir de zéro, ce qui démontre l'efficacité de l'approche de pré-entraînement. Nous fournissons également les détails de la MAE pour l'ensemble des 20 cibles :

Tableau 1. Répartition des 20 cibles de la tâche SMP pour le modèle préentraîné et le modèle peaufiné à partir de zéro. L'astérisque (*) indique l'erreur absolue moyenne (MAE) après filtrage d'un échantillon à l'indice 1469 de l'ensemble de test présentant une constante de rotation A aberrante de 232e3.

Conclusion

Le projet AtomGen représente un effort considérable pour appliquer les modèles basés sur les transformateurs à la modélisation atomistique. En privilégiant le traitement efficace des données, la conception de l'architecture des modèles et leur intégration avec les cadres d'apprentissage profond les plus utilisés, AtomGen fournit une base solide pour la recherche en chimie computationnelle et en science des matériaux. Ses contributions en matière d'agrégation, de traitement et de développement des données, ainsi que sa prise en charge de diverses techniques de pré-entraînement, sont susceptibles d'accélérer les progrès dans des domaines tels que la découverte de médicaments et la conception de matériaux.

En intégrant de nombreuses techniques de préentraînement efficaces sur des ensembles de données à grande échelle, notamment le contrôle de gradient, l'entraînement en précision mixte et divers objectifs de pré-entraînement, AtomGen démontre les meilleures pratiques pour le pré-entraînement et l'ajustement fin de modèles de transformateurs moléculaires sur un ensemble diversifié d'ensembles de données. L'application réussie de ces techniques à grande échelle apporte des informations précieuses pour surmonter les défis pratiques liés à l'application de l'apprentissage profond aux systèmes moléculaires complexes.