Machine Learning Interatomic Potentials: Advancing Open-Source Software for Efficient and Scalable Molecular Simulation

Cet article présente mlip v2, une nouvelle génération de logiciels open source qui améliore l'efficacité, l'évolutivité et la flexibilité des potentiels interatomiques d'apprentissage automatique grâce à une API modulaire repensée, un backend équivariant haute performance et des fonctionnalités avancées telles que l'architecture eSEN et une gestion améliorée des interactions électrostatiques.

L'essentiel

Imaginez que vous essayez de simuler le comportement d'une machine complexe composée de milliards de petits engrenages en mouvement (les atomes). Pour obtenir l'image la plus précise, vous devez utiliser les lois de la physique quantique, mais le faire revient à essayer de calculer la trajectoire de chaque engrenage individuel en utilisant un superordinateur qui met des années à terminer une seule seconde de simulation. C'est trop lent pour être utile.

Voici les Potentiels Interatomiques par Apprentissage Automatique (MLIPs). Considérez-les comme un « raccourci intelligent ». Ce sont des modèles d'IA entraînés sur les résultats de ces calculs physiques parfaits mais lents. Une fois entraînés, ils peuvent prédire le mouvement des atomes presque instantanément, avec une précision quasi identique à celle du superordinateur, mais en une fraction du temps.

Cependant, jusqu'à présent, utiliser ces raccourcis intelligents revenait à essayer de conduire une voiture de course haute performance avec un volant cassé et une carte qui ne fonctionne que pour une ville spécifique. Les outils étaient dispersés, difficiles à mettre à l'échelle et rigides.

Ce papier présente mlip v2, une mise à niveau majeure de la boîte à outils logicielle qui alimente ces simulations. Voici ce qu'ils ont construit, expliqué simplement :

1. La Nouvelle Salle des Machines (Le Framework Logiciel)

Les auteurs ont complètement redessiné la « salle des machines » du logiciel.

• L'Ancienne Façon : Imaginez une boîte à outils où chaque outil était collé à une poignée spécifique. Si vous vouliez changer la poignée, vous deviez casser l'outil.
• La Nouvelle Façon (mlip v2) : Ils ont construit un système modulaire où chaque outil (traitement des données, entraînement, simulation) s'emboîte comme des briques LEGO de haute qualité. Vous pouvez échanger des pièces facilement sans briser toute la structure. Cela rend beaucoup plus facile pour les scientifiques de personnaliser le logiciel selon leurs besoins spécifiques.

2. Le Turbo (Backend e3j)

L'un des plus grands goulots d'étranglement dans ces simulations réside dans l'exécution de mathématiques complexes liées aux formes 3D (appelées « opérations équivariantes »).

• L'Analogie : Imaginez essayer de faire tourner un objet 3D dans votre tête. Le faire pour des millions d'atomes est épuisant.
• La Solution : Ils ont intégré un nouveau moteur haute vitesse appelé e3j. C'est comme donner au logiciel un turbo spécifiquement conçu pour les mathématiques 3D. Le papier montre que cela permet au logiciel de fonctionner jusqu'à 3 fois plus vite sur les puces informatiques modernes (GPU et TPU).

3. Nouvelles Superpuissances

La mise à jour n'a pas seulement accéléré les choses ; elle a donné au logiciel de nouvelles capacités qu'il ne possédait pas auparavant :

• Le Système « Expert » (Mélange d'Experts) :

  • Le Problème : Entraîner un seul cerveau géant sur chaque type de molécule (de l'eau aux médicaments complexes) est difficile. Il se perd souvent.
  • La Solution : Ils ont introduit une architecture appelée eSEN qui agit comme une équipe de spécialistes. Au lieu d'un seul cerveau essayant de tout savoir, le système achemine différents problèmes vers différents « experts » au sein du modèle. Cela lui permet d'apprendre à partir d'ensembles de données massifs et désordonnés sans être submergé.

• Comprendre l'Électricité (Électrostatique) :

  • Le Problème : Les atomes portent souvent des charges électriques. Les modèles précédents avaient du mal à gérer les systèmes où la charge totale changeait, conduisant à des prédictions inexactes.
  • La Solution : La nouvelle version « écoute » explicitement la charge totale du système. C'est comme donner à l'IA une boussole qui sait toujours où se trouve le « Nord » (la charge totale), lui permettant de modéliser beaucoup plus précisément les systèmes chargés (comme les ions dans une batterie ou l'eau salée).

• Sentir la Courbe (Étiquettes Hessiennes) :

  • Le Problème : Savoir comment les atomes se déplacent (forces) revient à connaître la pente d'une colline. Mais pour prédire comment une balle roule et vibre, vous devez aussi connaître la courbure de la colline.
  • La Solution : Le logiciel peut désormais être entraîné à prédire cette « courbure » (appelée Hessienne). Cela aide l'IA à mieux comprendre la forme du paysage énergétique, conduisant à des prédictions plus précises de la façon dont les molécules vibrent et réagissent.

• Trouver le Chemin (Recherche d'État de Transition) :

  • Le Problème : Lorsque des produits chimiques réagissent, ils doivent traverser un « col de montagne » à haute énergie (état de transition) pour atteindre l'autre côté. Trouver ce col revient à chercher une aiguille dans une botte de foin.
  • La Solution : Ils ont ajouté un outil intégré appelé NEB (Nudged Elastic Band) qui étire automatiquement un élastique d'atomes entre un point de départ et un point d'arrivée pour trouver ce col de montagne efficacement.

• Espace pour Respirer (Ensembles NPT) :

  • Le Problème : Dans le monde réel, les liquides et les solides se dilatent et se contractent lorsque la pression ou la température change. Les anciennes simulations maintenaient souvent la taille du conteneur fixe, ce qui n'est pas réaliste.
  • La Solution : Le nouveau logiciel peut désormais simuler des systèmes où la taille du conteneur change pour maintenir une pression constante (NPT), tout comme un vrai ballon qui se dilate dans l'air chaud.

4. Le Résultat

Les auteurs ont publié des modèles pré-entraînés (les « cerveaux » déjà instruits sur un ensemble de données massif de molécules) prêts à l'emploi. Ils ont testé ces modèles et ont constaté qu'ils étaient très précis pour prédire l'énergie, les forces et même les charges électriques des atomes.

En résumé : Les auteurs ont pris un outil puissant mais lourd pour simuler les atomes et l'ont transformé en une plateforme élégante, modulaire et ultra-rapide. Ils ont ajouté de nouveaux « muscles » (vitesse), de nouveaux « sens » (conscience de la charge et de la courbure) et de nouveaux « outils » (recherche de voies de réaction), rendant possible la simulation de systèmes chimiques complexes et réels qui étaient auparavant trop difficiles ou trop lents à modéliser. Le logiciel est open-source, ce qui signifie que n'importe qui peut le télécharger et commencer à l'utiliser immédiatement.

Résumé technique

Résumé technique : mlip v2 – Faire progresser les logiciels open source pour des simulations moléculaires efficaces et évolutives

Énoncé du problème

Les potentiels interatomiques basés sur l'apprentissage automatique (MLIP) offrent une voie pour atteindre une précision proche de l'ab initio dans les simulations atomistiques, et ce à une fraction du coût computationnel des méthodes de structure électronique comme la théorie de la fonctionnelle de la densité (DFT). Cependant, leur adoption plus large est entravée par un outil fragmenté, une scalabilité limitée et des conceptions logicielles rigides qui peinent à soutenir une simulation efficace, une facilité d'utilisation pour les applications et une innovation méthodologique rapide. Bien que la première version de la bibliothèque mlip (v1) ait établi un cadre unifié basé sur JAX pour l'entraînement et le déploiement de MLIP, elle a été conçue comme un système de première génération. Elle présentait des limites en matière de composition, de contrôle sur les pipelines de bout en bout et de capacité à gérer des fonctionnalités scientifiques avancées telles que l'électrostatique complexe, la recherche d'états de transition et l'entraînement à grande échelle sur plusieurs jeux de données.

Méthodologie et architecture

L'article présente mlip v2, une refonte complète de la bibliothèque mlip qui répond à ces limitations grâce à une refonte ciblée de l'API et à l'intégration de nouveaux composants haute performance.

1. Cadre logiciel unifié et refonte de l'API

• Classe Graph unifiée : La bibliothèque remplace l'ancien jraph.GraphsTuple par une nouvelle classe Graph unifiée. Celle-ci sert de structure de données centrale pour les entrées, les sorties et les caractéristiques latentes intermédiaires des modèles, standardisant l'interface entre tous les composants du modèle (Graph $\to$ Graph). Cela élimine les dépendances envers des projets archivés et facilite une composition et une extensibilité plus propres.
• Architecture modulaire : La conception découple les blocs de construction fondamentaux (traitement des données, entraînement, inférence) via des interfaces minimales et clairement définies. Cela permet une personnalisation flexible des flux de travail, y compris l'entraînement sur plusieurs jeux de données et le fine-tuning multi-têtes.
• Stratégie de migration : Malgré les refactorisations internes, la bibliothèque maintient une interface familière pour les flux de travail principaux afin de minimiser les changements disruptifs pour les utilisateurs existants, avec l'appui d'un guide de migration complet.

2. Backend haute performance (e3j)

Pour optimiser l'exécution sur divers matériels, mlip v2 intègre e3j, un nouveau backend open source haute performance pour les opérations équivariantes.

• Implémentation : e3j fournit des noyaux dédiés pour les opérations équivariantes utilisant à la fois Pallas (pour les TPU) et CUDA (pour les GPU).
• Modèles ciblés : Il accélère spécifiquement les modèles reposant sur les produits tensoriels de Clebsch-Gordan, tels que MACE et NequIP, qui constituent souvent des goulots d'étranglement computationnels.
• Performance : Les benchmarks indiquent des accélérations d'exécution allant jusqu'à 3 fois par rapport à l'implémentation v1.

3. Capacités scientifiques élargies

Le cadre introduit plusieurs nouvelles méthodologies pour élargir la portée des applications MLIP :

• Architecture eSEN avec mélange d'experts (MoE) : La bibliothèque intègre l'architecture eSEN, qui utilise une formulation MoE. Cela permet un entraînement évolutif sur de grands jeux de données diversifiés tout en préservant une inférence efficace. Le mécanisme de routage permet de contracter des experts spécialisés en un seul noyau dense lors de l'inférence.
• Électrostatique avancée et modélisation des charges :
  • Prédiction des charges partielles : Tous les modèles prennent désormais en charge la prédiction des charges partielles atomiques.
  • Interactions à longue portée : Un terme d'interaction de Coulomb modifié (suivant la formulation PhysNet) est implémenté pour gérer l'électrostatique à longue portée, incluant une régularisation à noyau doux pour éviter les divergences.
  • Conditionnement par charge globale : Pour améliorer la précision sur les systèmes à charges globales variables, les modèles incorporent une embedding de la charge totale du système, concaténée avec les embeddings des numéros atomiques.
• Entraînement par étiquettes Hessian : La bibliothèque prend en charge l'entraînement avec des dérivées d'ordre deux (Hessiennes) de l'énergie. Pour gérer les coûts computationnels, elle emploie une stratégie d'échantillonnage (produits vecteur-Jacobien) où seules certaines composantes de force sont différenciées par rapport à toutes les coordonnées atomiques. Cela facilite l'entraînement de modèles de base sur des informations de courbure sans le coût prohibitif de la rétropropagation de la Hessienne complète.
• Recherche d'états de transition : Un moteur personnalisé implémentant la méthode Nudged Elastic Band (NEB) (incluant la variante climbing image) est intégré, en interface avec ASE pour localiser les états de transition.
• Simulations d'ensemble NPT : La bibliothèque introduit le support pour les simulations isotherme-isobare (NPT) via un barostat Monte Carlo (MC) basé sur JAX couplé à un intégrateur de Langevin. Cette approche évite les évaluations de contraintes coûteuses requises par d'autres barostats (par exemple, Berendsen ou Parrinello-Rahman) en utilisant un critère de Metropolis basé sur les changements d'énergie potentielle.

4. Fine-tuning multi-têtes

Un cadre unifié pour le fine-tuning multi-têtes est introduit, permettant aux modèles pré-entraînés sur de grands jeux de données d'être spécialisés pour des tâches en aval (par exemple, chimies spécifiques ou niveaux de théorie) sans oubli catastrophique. Cela est réalisé grâce à un squelette équivariant partagé associé à des têtes de lecture spécifiques aux jeux de données et à des tables d'énergie atomique.

Résultats et validation

Les auteurs fournissent une validation extensive de la nouvelle bibliothèque et des modèles pré-entraînés (MACE, NequIP, ViSNet et eSEN) entraînés sur un sous-ensemble curé du jeu de données OMOL25 (spécifiquement le sous-ensemble SPICE2, contenant environ 1,76 million de structures).

• Précision : Les modèles pré-entraînés ont été évalués sur sept sous-ensembles moléculaires de SPICE2. L'architecture eSEN a obtenu les erreurs absolues moyennes (MAE) les plus faibles pour l'énergie et les forces sur la plupart des sous-ensembles.
• Fidélité physique : L'évaluation utilisant MLIPAudit a montré que toutes les architectures ont obtenu des scores quasi parfaits sur les distributions de longueurs de liaison, la planarité des cycles et la stabilité des géométries de référence. eSEN a obtenu le score global le plus élevé (0,716), suivi de ViSNet (0,699).
• Électrostatique et charges : Les modèles avec embedding de charge globale ont démontré une précision de prédiction d'énergie significativement améliorée pour les systèmes chargés globalement par rapport à ceux sans. Les prédictions de charges partielles étaient précises sur tous les sous-ensembles.
• Entraînement par Hessienne : Une étude contrôlée a démontré que l'entraînement avec des étiquettes Hessienne réduisait considérablement l'erreur dans les fréquences vibrationnelles prédites par rapport à un modèle de base entraîné uniquement sur les énergies et les forces.
• Validation NPT : L'intégrateur NPT basé sur JAX a montré un excellent accord avec les implémentations de référence ASE (Berendsen et Parrinello-Rahman) en termes de température, de compressibilité isotherme et de fonctions de distribution radiale, tout en offrant des accélérations de 2,2 à 4,0 fois.
• Temps d'exécution : Les benchmarks ont confirmé que l'intégration de e3j et du backend optimisé a entraîné des accélérations constantes sur les modèles MACE et NequIP, la bibliothèque prenant en charge les simulations par lots sur des appareils uniques.

Importance et affirmations

L'article positionne mlip v2 comme une fondation évolutive et adaptable pour la simulation moléculaire basée sur l'apprentissage automatique. Son importance principale réside dans le comblement du fossé entre la recherche en ML et l'application pratique en :

1. Unifiant la pile : Fournissant un cadre unique et extensible reliant le traitement des données, l'entraînement des modèles et la simulation moléculaire.
2. Améliorant la scalabilité : Permettant un entraînement efficace sur de grands jeux de données diversifiés grâce aux formulations MoE et aux backends haute performance (e3j).
3. Élargissant l'applicabilité : Introduisant des fonctionnalités permettant la modélisation de systèmes complexes, réactifs et hors équilibre, y compris les espèces chargées, les états de transition et les environnements de phase condensée dans des conditions thermodynamiques réalistes (NPT).
4. Accessibilité open source : Publiant la bibliothèque sous la licence Apache 2.0 avec des modèles pré-entraînés et une documentation complète pour réduire la barrière à l'entrée tant pour les chercheurs appliqués que pour les développeurs de méthodes.

Les auteurs soulignent que, bien que la bibliothèque fasse progresser considérablement l'état de l'art en matière d'infrastructure logicielle, les résultats présentés sont indicatifs des performances réalisables avec la bibliothèque plutôt qu'un benchmark définitif entre architectures, notant qu'il est difficile de définir des paramètres d'hyperparamètres comparables entre différentes familles de modèles.