Par Ali Kore, Amrit Krishnan, David Emerson

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

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

Au cœur du projet AtomGen se trouve Atomformer, un modèle transformateur uniquement encodeur, adapté pour travailler avec des structures atomistiques tridimensionnelles.  Inspiré par l’architecture Uni-mol+ , Atomformer intègre directement l’information spatiale 3D dans l’architecture du modèle à l’aide d’embeddings positionnels par paires gaussiens.  Plus précisément, cette méthode exploite les distances euclidiennes entre les atomes et un noyau gaussien conscient de type paire qui est projeté et ajouté au masque d’attention.  Il intègre aussi des métadonnées intégrées (masse atomique, rayon, valence, etc.) basées 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 de données diversifié d’atomes avec énergie+Forces

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

Figure 1. « Visualisations des structures atomiques à partir des composants du jeu 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 (2x MPtrj et 3x upsampling SPICE), offrant une gamme diversifiée d’environnements atomistiques pour l’entraînement. Il inclut les nombres atomiques, les coordonnées 3D, les forces, l’énergie de formation, l’énergie totale et un drapeau booléen indiquant la présence de données valides sur l’énergie de formation.  L’ampleur et la diversité de l’ensemble de données, combinées à l’apprentissage multitâche des énergies et des forces, permettent un apprentissage complet à travers différents types d’environnements chimiques.

Traitement et intégration des données

Une contribution technique significative du projet AtomGen est le développement de pipelines de traitement de données efficaces pour chaque jeu de données source. Ces pipelines gèrent la conversion de divers formats de données en une structure unifiée compatible avec la bibliothèque HuggingFace Datasets. Les pipelines intègrent un traitement parallèle pour de grands ensembles de données, gèrent des formats de données compressés, convertissent des unités, calculent des propriétés dérivées et gèrent le bourrage dynamique ainsi que le batching de 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 de jeux de données pour le traitement par lots. »

Intégration de HuggingFace

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

AtomGen exploite l’écosystème HuggingFace pour le développement de modèles et la gestion de jeux de données. Le projet inclut des implémentations personnalisées d’AtomformerConfig, AtomformerModel, ainsi que divers modèles spécifiques à chaque tâche tels que AtomformerForMaskedAM et Structure2EnergyAndForces. Ces implémentations permettent une intégration fluide avec l’API Trainer et le hub de modèles de HuggingFace, facilitant la reproductibilité et la facilité d’utilisation. 

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

Le cours DataCollatorForAtomModeling est un élément technique clé d’AtomGen, offrant une flexibilité dans la préparation des données et soutenant diverses techniques de pré-entraînement. Ce collateur de données personnalisé gère les complexités du batching des données, incluant le bourrage dynamique des séquences d’entrée, le calcul de masques d’attention et le calcul optionnel des codages positionnels laplaciens.  Ces encodages sont utilisés avec l’implémentation TokenGT comme identifiants de nœud.

Le collateur de données supporte plusieurs objectifs pré-entraînement auto-supervisés.  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 masquée des atomes (MAM) : similaire à BERTCette technique masque aléatoirement les identités atomiques pour que le modèle puisse les prédire, l’aidant à apprendre le contexte atomique et les interactions.  Mettre l’argument mam à un nombre en virgule flottante masquera aléatoirement cette fraction d’atomes.
2. Réduction du bruit des coordonnées : Cela implique de perturber les coordonnées atomiques et de confier au modèle la récupération des positions originales, améliorant ainsi sa compréhension de la géométrie moléculaire.  Le degré de perturbation peut être ajusté comme un hyperparamètre coords_perturb comme un nombre à virgule flottante qui représente l’écart-type d’une distribution normale de bruit centrée autour de zéro ajoutée aux coordonnées.
3. Objectif MaskGIT : Inspirée par l’approche MaskGIT pour la génération d’images, cette technique applique une stratégie de masquage plus complexe aux atomes, améliorant potentiellement la compréhension de la structure globale du système par le modèle.  Cet objectif est activé en mettant l’argument « mam » sur Vrai, ce qui permet cette stratégie au lieu de l’objectif standard MAM.

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

Approche de pré-entraînement et d’ajustement fin

Bien que le collecteur de données supporte diverses techniques de pré-entraînement auto-supervisées, comme discuté précédemment, celles-ci n’ont pas été utilisées dans les expériences actuelles de pré-entraînement. L’exploration de ces diverses approches auto-supervisées demeure une voie passionnante pour la recherche future.  AtomGen utilise une approche en deux étapes : préentraînement sur l’ensemble de données S2EF-15M, suivi d’un ajustement fin 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 4xA40 pendant deux semaines, totalisant 1344 heures GPU. 

Pour l’ajustement fin, AtomGen utilise la suite de benchmarks ATOM3D, qui comprend une gamme de tâches de modélisation atomistique telles que les propriétés des petites molécules (SMP), la prédiction de stabilité des mutations (MSP) et l’affinité de liaison des ligands (LBA). Chacune des 8 tâches d’ATOM3D est prétraitée et formatée pour être compatible avec le collecteur de données et la bibliothèque de jeux de données HuggingFace.

Performance et scalabilité

Le projet AtomGen inclut des optimisations pour l’entraînement sur des ensembles de données à grande échelle, incluant le checkpointing de gradient pour un entraînement efficace en mémoire, l’entraînement en précision mixte pour un calcul plus rapide, et le support de l’entraînement distribué pour les configurations multi-GPU.  

Sur la tâche SMP, le modèle Atomformer affiné a atteint un MAE de test de 1,077, comparativement à 1,13 pour un modèle entraîné à partir de zéro, démontrant l’efficacité de l’approche de préentraînement.  Nous fournissons également la répartition des EMA sur les 20 cibles :

Cibles	Pré-entraîné (évaluation)	Pré-entraîné (test)	Évaluation (Scratch)	Scratch (test)
Constante de rotation A [GHz]	0.1216	18.1177/0.1566*	0.1621	18.1576/0.1972*
Constante de rotation B [GHz]	0.0541	0.0677	0.0657	0.0794
Constante rotationnelle C [GHz]	0.0283	0.0404	0.0296	0.0421
Moment dipôle [D]	0.3044	0.3014	0.4128	0.4107
Polarisabilité isotrope [a0^3]	0.3388	0.3392	0.534	0.5358
Énergie de l’HOMO [Ha]	0.0104	0.0104	0.0107	0.0106
Énergie de LUMO [Ha]	0.014	0.0143	0.0174	0.0174
Gap (LUMO-HOMO) [Ha]	0.0165	0.0169	0.0194	0.0197
Étendue spatiale électronique [a0^2]	2.0202	1.9656	2.2719	2.2126
Énergie vibratoire à point zéro [Ha]	0.0006	0.0006	0.0007	0.0007
Énergie interne à 0 K [Ha]	0.0697	0.0783	0.1307	0.1291
Énergie interne à 298,15 K [Ha]	0.0697	0.0783	0.1307	0.1291
Enthalpie à 298,15 K [ha]	0.0697	0.0783	0.1307	0.1291
Énergie libre à 298,15 K [Ha]	0.0697	0.0783	0.1307	0.1291
Capacité thermique à 298,15 K [cal/(mol· K)]	0.1531	0.1518	0.2325	0.2337
Énergie interne à 0 K (thermochimie) [Ha]	0.0117	0.0116	0.0185	0.0183
Énergie interne à 298,15 K (thermochimie) [Ha]	0.0117	0.0116	0.0185	0.0183
Enthalpie à 298,15 K (thermochimie) [Ha]	0.0117	0.0116	0.0186	0.0183
Énergie libre à 298,15 K (thermochimie) [Ha]	0.0116	0.0116	0.0185	0.0183
Capacité thermique à 298,15 K (thermochimie) [cal/(mol· K)]	0.1519	0.151	0.2323	0.2324

Tableau 1. La répartition des 20 cibles de la tâche SMP pour le modèle pré-entraîné et le modèle a été affinée à partir de zéro.  L’astérisque (∗) indique l’EMA après filtrage d’un échantillon d’indice de 1469 provenant de l’ensemble de test avec une constante rotationnelle A atypique de 232e3.

Conclusion

Le projet AtomGen représente un effort important pour appliquer des modèles basés sur des transformateurs aux tâches de modélisation atomistique. En se concentrant sur un traitement efficace des données, la conception d’architectures de modèles et l’intégration avec des cadres populaires d’apprentissage profond, AtomGen offre une base solide pour la recherche en chimie computationnelle et en science des matériaux. Les contributions du projet à l’agrégation de données, au traitement et au développement de modèles, ainsi que son soutien à diverses techniques de pré-entraînement, ont le potentiel 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 pour un pré-entraînement efficace sur des ensembles de données à grande échelle, incluant le checkpoint de gradient, l’entraînement en précision mixte et de nombreux objectifs de pré-entraînement. Ces optimisations démontrent les meilleures pratiques pour pré-entraîner et affiner les modèles de transformateurs moléculaires sur un ensemble diversifié d’ensembles de données. En appliquant avec succès ces techniques à grande échelle, AtomGen offre des perspectives précieuses pour surmonter les défis pratiques liés à l’application de l’apprentissage profond à des systèmes moléculaires complexes.