Par Ali Kore, Amrit Krishnan, David Emerson

Le projet AtomGen se concentre sur l'amélioration des capacités de modélisation atomistique grâce à des techniques avancées d'apprentissage automatique. En intégrant des modèles d'apprentissage profond à 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 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 à codeur seul adapté au travail avec des structures atomistiques tridimensionnelles. Inspiré par l'architecture Uni-mol+, Atomformer incorpore des informations spatiales 3D directement dans l'architecture du modèle à l'aide d'encastrements positionnels gaussiens par paire. Plus précisément, cette méthode exploite les distances euclidiennes entre les atomes et un noyau gaussien tenant compte des paires, qui est projeté et ajouté au masque d'attention. Elle incorpore également des encastrements de métadonnées (masse atomique, rayon, valence, etc.) basés sur les espèces atomiques. 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ées d'atomes avec énergie+forces

Pour entraîner le modèle de base d'Atomformer, l'équipe d'AtomGen a compilé l'ensemble de données S2EF-15M, une agrégation à grande échelle de structures atomiques, de leurs forces et de leurs énergies. Ce jeu 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 présente 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 (2x MPtrj et 3x SPICE upsampling), fournissant une gamme variée d'environnements atomistiques pour la formation. Il comprend les numéros 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'échelle et la diversité de l'ensemble de données, associé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

Une contribution technique importante du projet AtomGen est le développement de pipelines de traitement de données efficaces pour chaque ensemble 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 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 la mise en 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 de HuggingFace

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

AtomGen s'appuie sur l'écosystème HuggingFace pour le développement de modèles et la gestion de jeux de données. Le projet comprend des implémentations personnalisées de AtomformerConfig, AtomformerModel et divers modèles spécifiques à des tâches telles que 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 facilité d'utilisation. 

Collecte des données et techniques de pré-entraînement

La classe DataCollatorForAtomModeling est un composant technique clé d'AtomGen, offrant une flexibilité dans la préparation des données et prenant en charge diverses techniques de pré-entraînement. Ce collecteur de données personnalisé gère les complexités de la mise en lot des données, y compris le remplissage dynamique des séquences d'entrée, le calcul des masques d'attention et le calcul optionnel des encodages positionnels du Laplacien. Ces encodages sont utilisés avec l'implémentation de TokenGT comme identifiants de nœuds.

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 d'atomes masqués (MAM): Semblable à la modélisation du langage masqué de l'ORET, cette technique masque aléatoirement les identités atomiques pour que le modèle puisse les prédire, ce qui l'aide à apprendre le contexte et les interactions atomiques. La définition de l'argument mam à un nombre à virgule flottante permet de masquer aléatoirement cette fraction d'atomes.
2. Le débruitage des coordonnées: Il s'agit de perturber les coordonnées atomiques et de demander au modèle de retrouver les positions d'origine, afin d'améliorer sa compréhension de la géométrie moléculaire. Le degré de perturbation peut être ajusté en tant qu'hyperparamètre coords_perturb, sous la forme d'un nombre à virgule flottante représentant l'écart-type d'une distribution normale de bruit centrée autour de zéro qui est ajoutée aux coordonnées.
3. Objectif de MaskGIT: Inspirée de l'approche MaskGIT dans la génération d'images, cette technique applique une stratégie de masquage plus difficile aux atomes, ce qui peut améliorer la compréhension par le modèle de la structure globale du système. Cet objectif est activé en fixant l'argument "mam" à True, ce qui active cette stratégie au lieu de l'objectif MAM standard.

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

Approche de préformation et de mise au point

Bien que le collecteur de données prenne en charge diverses techniques de pré-entraînement auto-supervisé, comme nous l'avons vu précédemment, celles-ci n'ont pas été utilisées dans les expériences de pré-entraînement actuelles. L'exploration de ces diverses approches auto-supervisées reste une piste intéressante 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 réglage 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 de GPU. 

Pour le réglage fin, AtomGen utilise la suite de référence 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 la stabilité des mutations (MSP) et l' affinité de liaison du ligand (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 HuggingFace Datasets.

Performance et évolutivité

Le projet AtomGen comprend des optimisations pour l'entraînement sur des ensembles de données à grande échelle, notamment le point de contrôle du gradient pour un entraînement efficace en termes de 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 obtenu une MAE de test 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 la répartition de la MAE sur l'ensemble des 20 cibles :

Objectifs	Pré-entraîné (eval)	Pré-entraîné (test)	Scratch (eval)	Gratter(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 de rotation C [GHz]	0.0283	0.0404	0.0296	0.0421
Moment dipolaire [D]	0.3044	0.3014	0.4128	0.4107
Polarisation isotrope [a0^3]	0.3388	0.3392	0.534	0.5358
Énergie du HOMO [Ha]	0.0104	0.0104	0.0107	0.0106
Énergie du LUMO [Ha]	0.014	0.0143	0.0174	0.0174
Écart (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 au 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. Ventilation des 20 cibles de la tâche SMP pour le modèle pré-entraîné et le modèle affiné à partir de zéro. L'astérisque (∗) indique le MAE après filtrage d'un échantillon à l'indice 1469 de l'ensemble de test avec une constante de rotation A aberrante de 232e3.

Conclusion

Le projet AtomGen représente un effort significatif dans l'application de modèles basés sur des transformateurs à des tâches de modélisation atomistique. En se concentrant sur le traitement efficace des données, la conception de l'architecture du modèle et l'intégration avec des cadres d'apprentissage profond populaires, AtomGen fournit une base solide pour la recherche en chimie computationnelle et en science des matériaux. Les contributions du projet en matière d'agrégation de données, de traitement et de 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, y compris le point de contrôle du gradient, l'entraînement en précision mixte et de nombreux objectifs de pré-entraînement. Ces optimisations démontrent les meilleures pratiques pour le pré-entraînement et le réglage fin des modèles de transformateurs moléculaires sur un ensemble varié de données. En appliquant avec succès ces techniques à grande échelle, AtomGen fournit des indications précieuses pour surmonter les défis pratiques de l'application de l'apprentissage profond à des systèmes moléculaires complexes.