Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory

Depuis plus de soixante ans, les scientifiques s'appuient sur un ensemble de règles puissantes pour comprendre comment les atomes s'assemblent pour former les matériaux qui nous entourent. Ce cadre, connu sous le nom de théorie de la fonctionnelle de la densité, traite le nuage électronique entourant un atome non pas comme une collection de particules individuelles, mais comme une densité lisse et continue. Cette approche permet aux chercheurs de prédire les propriétés de tout, des nouveaux médicaments aux batteries avancées. Cependant, il y a un piège. Pour obtenir des résultats précis, la théorie nécessite un calcul complexe et répétitif qui implique de résoudre le comportement de chaque électron du système. À mesure que le système s'agrandit, le temps requis pour ces calculs explose, croissant si vite que simuler un morceau de métal de taille modeste ou une grande molécule biologique devient impossible, même sur les supercalculateurs les plus puissants. Ce goulot d'étranglement a tenu de nombreuses questions scientifiques importantes hors de portée, forçant les chercheurs à choisir entre précision et échelle.

Une équipe de chercheurs du California Institute of Technology a trouvé un moyen de briser cette barrière en apprenant à un ordinateur à sauter l'étape la plus coûteuse du calcul sans perdre en précision. Au lieu d'essayer de prédire directement la réponse finale ou d'apprendre un raccourci mathématique difficile qui échoue souvent, ils ont appris à une intelligence artificielle à effectuer l'étape spécifique et répétée qui intervient dans chaque cycle de calcul. En se concentrant sur cette étape intermédiaire, ils ont créé une méthode capable de simuler des systèmes comprenant des dizaines de milliers d'atomes sur une seule carte graphique, une tâche qui nécessitait auparavant des milliers de processeurs. Le résultat est une nouvelle façon de modéliser la matière qui est presque aussi rapide que les approximations les plus simples, tout en conservant la haute précision des méthodes les plus rigoureuses.

Le cœur du problème réside dans la manière dont la théorie traite les électrons. Pour trouver l'état stable d'un matériau, l'ordinateur doit ajuster de manière répétée une carte de la densité électronique jusqu'à ce qu'elle ne change plus. Dans la méthode traditionnelle, cet ajustement nécessite de résoudre un problème de valeurs propres massif, ce qui consiste essentiellement à trouver les motifs spécifiques de mouvement des électrons pour chaque atome du système. Cette opération est si lourde sur le plan computationnel que son coût croît de manière cubique avec le nombre d'électrons ; doubler la taille du système rend le calcul huit fois plus difficile. Pendant des décennies, les scientifiques ont tenté de contourner cela en supprimant la nécessité de suivre les motifs individuels des électrons, une stratégie appelée théorie de la densité sans orbitales. Cependant, les tentatives précédentes pour apprendre ce raccourci à l'aide de l'apprentissage automatique ont échoué. Certains ont essayé d'apprendre la relation entre la densité électronique et l'énergie, mais cette relation est mathématiquement instable et sujette à des erreurs. D'autres ont essayé de prédire directement la densité électronique finale à partir de l'arrangement des atomes, mais ces modèles échouaient lorsqu'on leur demandait de prédire des systèmes plus grands ou chimiquement différents de ceux sur lesquels ils avaient été entraînés.

Les chercheurs ont réalisé que la clé était d'arrêter d'essayer de prédire la destination finale et d'apprendre plutôt le voyage. À chaque étape du calcul traditionnel, l'ordinateur prend une carte des forces agissant sur les électrons et produit une nouvelle carte de l'endroit où les électrons sont susceptibles de se trouver. Cette étape de progression est mathématiquement stable et bien définie. L'équipe a entraîné un réseau de neurones spécialisé, un type d'intelligence artificielle conçu pour comprendre les motifs spatiaux, à apprendre précisément cette transformation. Ils ont alimenté le réseau avec des milliers d'exemples de ces cartes de forces et des cartes de densité électronique correspondantes qu'elles produisaient, tirées d'un ensemble diversifié de 8 504 molécules et matériaux solides, incluant des composés organiques, des isolants et des métaux. Le réseau a appris à prédire la nouvelle densité électronique directement à partir des forces, remplaçant efficacement le calcul lent et lourd par une prédiction rapide et apprise.

Ce qui rend cette approche unique est la manière dont elle gère la complexité de différents systèmes. Les chercheurs ont conçu le réseau pour qu'il soit invariant au domaine, ce qui signifie qu'il peut appliquer les mêmes règles apprises à une minuscule molécule et à un cristal massif sans avoir besoin d'être réentraîné pour chaque nouvelle taille. Ils ont également veillé à ce que le réseau respecte les symétries fondamentales de l'espace, de sorte qu'il comprenne que la rotation ou le déplacement d'une molécule ne change pas ses propriétés physiques. Lorsqu'il a été testé sur des systèmes qu'il n'avait jamais vus, le modèle s'est révélé remarquablement robuste. Sur un ensemble de molécules de type médicamenteux contenant des éléments et des tailles absents des données d'entraînement, la nouvelle méthode a maintené une haute précision, alors que les précédents modèles de prédiction directe voyaient leurs erreurs grimper en flèche à mesure que les molécules devenaient plus grandes. La nature auto-cohérente de la méthode lui permet de corriger ses propres erreurs à chaque étape, tel un navigateur qui vérifie fréquemment sa position plutôt que de deviner tout le trajet d'un coup.

La puissance de cette méthode est devenue évidente lorsque l'équipe l'a appliquée à des systèmes périodiques, tels que des cristaux, et à des défauts étendus dans les métaux. Dans les calculs standards pour les cristaux, l'ordinateur doit répéter le calcul lourd pour de nombreux points différents dans l'espace mathématique qui décrit le motif répétitif du cristal. La nouvelle méthode saute entièrement cette répétition, faisant évoluer la densité électronique sur la maille élémentaire sans avoir besoin d'échantillonner ces points individuellement. Cela a permis aux chercheurs de faire converger les calculs pour des systèmes métalliques et des semi-conducteurs avec un seul modèle, reproduisant les spectres électroniques et les propriétés structurelles avec la même précision que la méthode traditionnelle. De plus, le modèle a montré une capacité surprenante à transférer entre différentes approximations théoriques. Un modèle entraîné sur un type de calcul pouvait être utilisé avec une approximation différente, plus précise, pour les forces sans aucun réentraînement, simplement en échangeant les forces d'entrée.

Le test ultime est venu lorsque les chercheurs ont appliqué la méthode à une dislocation de magnésium, un type de défaut dans un cristal métallique qui est crucial pour comprendre comment les matériaux se déforment et se rompent. Ces défauts créent des champs élastiques à longue portée qui nécessitent des simulations de milliers d'atomes pour être modélisés avec précision. Les tentatives précédentes pour simuler un tel système avec une haute précision nécessitaient un supercalculateur doté de milliers de processeurs graphiques et prenaient des jours de calcul. La nouvelle méthode, tournant sur une seule carte graphique moderne, a réussi à faire converger le calcul pour un système contenant 8 250 atomes de magnésium et 82 500 électrons de valence. Le temps nécessaire pour compléter le calcul a crû de manière presque linéaire avec la taille du système, confirmant que le goulot d'étranglement computationnel avait été supprimé. Les densités électroniques résultantes étaient indiscernables de la référence de haute précision, permettant aux chercheurs de résoudre les subtiles différences d'énergie entre différents types de défauts qui déterminent la résistance du métal.

Ce travail démontre que l'apprentissage automatique peut étendre la portée des calculs de structure électronique non pas en remplaçant toute la théorie physique, mais en accélérant son opération répétée la plus coûteuse. En apprenant la carte de passage des forces vers la densité, les chercheurs ont créé un outil qui est stable, transférable à différents types de matière et capable de gérer des systèmes à une échelle qui était auparavant inaccessible. La méthode fournit un diagnostic direct de sa propre fiabilité : si le calcul ne parvient pas à converger, cela signale que le système se situe en dehors de la plage de confiance du modèle, une caractéristique dont les modèles de prédiction directe sont dépourvus. Bien que le modèle actuel prédise la densité électronique, l'énergie totale et d'autres propriétés peuvent toujours être obtenues en une seule étape finale, préservant ainsi la précision de la théorie complète. Cette approche ouvre la porte à la simulation de matériaux complexes, de systèmes biologiques et de réactions chimiques à un niveau de détail et d'échelle qui était autrefois le domaine de l'impossibilité théorique.