La méthode « IA à grande échelle » accélère les simulations atomistiques pour les scientifiques

Les calculs quantiques des systèmes moléculaires nécessitent souvent des quantités extraordinaires de puissance de calcul ; ces calculs sont généralement effectués sur les plus grands superordinateurs du monde afin de mieux comprendre les produits du monde réel tels que les batteries et les semi-conducteurs.

Aujourd'hui, les chercheurs de l'UC Berkeley et du Lawrence Berkeley National Laboratory (Berkeley Lab) ont développé une nouvelle méthode d'apprentissage automatique qui accélère considérablement les simulations atomistiques en améliorant l'évolutivité du modèle. Cette approche réduit de plus de cinq fois la mémoire informatique requise pour les simulations par rapport aux modèles existants et fournit des résultats dix fois plus rapidement.

Leurs recherches ont été acceptées au Neural Information Processing Systems (NeurIPS) 2024, un lieu de conférence et de publication sur l'intelligence artificielle et l'apprentissage automatique. Ils présenteront leurs travaux lors de la conférence du 13 décembre, et une version de leur article est disponible sur le site arXiv serveur de préimpression.

« Nous voulions créer un type différent d'architecture d'apprentissage automatique en utilisant des méthodes généralement appliquées aux grands modèles de langage », a déclaré Eric Qu, étudiant diplômé de l'UC Berkeley et co-auteur du document de recherche. « Grâce à notre approche, les chercheurs peuvent cartographier plus efficacement la façon dont les atomes se déplacent et interagissent les uns avec les autres. »

Comprendre ce qui arrive aux plus petits éléments constitutifs de la nature peut ouvrir la voie à une compréhension plus approfondie de la science des matériaux, de la chimie et du développement de médicaments, entre autres sujets scientifiques fondamentaux.

« Ce modèle peut aider les scientifiques à déterminer les mécanismes de réaction chimique de manière beaucoup plus efficace », a déclaré Samuel Blau, chimiste informaticien au Berkeley Lab. « Si vous parvenez à comprendre la chimie complexe des systèmes du monde réel, vous pourrez découvrir comment les contrôler de nouvelles manières. »

Balances affamées

Au cours de la dernière décennie, les scientifiques et les ingénieurs ont construit de grands modèles de langage comme ChatGPT en utilisant des ensembles de données massifs et une stratégie appelée mise à l'échelle. La mise à l'échelle consiste à rendre ces modèles plus grands et plus intelligents en augmentant systématiquement le nombre de paramètres dans les réseaux de neurones. La manière dont vous augmentez ces paramètres est importante : différents paramètres contribuent aux performances du modèle de manière distincte, et l'optimisation de ce processus peut conduire à des améliorations significatives.

Les chercheurs peuvent également concevoir de nouvelles opérations ou de nouveaux composants au sein de l’architecture du réseau neuronal, tels que de nouveaux mécanismes d’attention, plus expressifs, permettant ainsi d’augmenter davantage les paramètres tout en maintenant ou en améliorant l’efficacité.

Mais ce n’est pas seulement une question de taille ; la mise à l'échelle signifie également trouver des moyens de rendre ces modèles plus efficaces, en utilisant des algorithmes plus intelligents pour économiser du temps et de la puissance de calcul pendant la formation et l'utilisation. Au lieu de se concentrer uniquement sur la puissance de traitement brute, les chercheurs mesurent souvent l’efficacité en fonction du temps réellement nécessaire pour entraîner ou exécuter ces modèles, en donnant la priorité aux performances réelles.

Cependant, les principes de mise à l’échelle n’ont pas été largement appliqués à un autre type de modèle d’apprentissage automatique particulièrement utile pour les scientifiques : les potentiels interatomiques des réseaux neuronaux (NNIP). Les NNIP constituent une alternative efficace aux simulations de mécanique quantique coûteuses en termes de calcul, qui permettent aux chercheurs de prédire beaucoup plus rapidement les propriétés moléculaires et matérielles.

« Les NNIP deviennent rapidement l'approche la plus puissante pour la simulation moléculaire ou matérielle », a déclaré Aditi Krishnapriyan, co-auteur de l'article, professeur adjoint à l'UC Berkeley et chercheur à la division de mathématiques appliquées et de recherche informatique du laboratoire de Berkeley. « Auparavant, la conception d'algorithmes intelligents à grande échelle était principalement développée dans d'autres domaines de l'apprentissage automatique, tels que les grands modèles de langage, et dans une moindre mesure pour l'étude des matériaux, de la chimie ou de la physique. »

L’équipe de Berkeley a ainsi développé une architecture NNIP qui pourrait être mise à l’échelle efficacement. L'architecture, connue sous le nom d'EScAIP (Efficiently Scaled Attention Interatomic Potential), représente une avancée significative dans la mise à l'échelle des modèles d'apprentissage automatique pour les applications scientifiques, a déclaré Krishnapriyan.

Il pleut des données

Alors que les grands modèles de langage tels que ChatGPT sont formés sur du texte, pour lequel des milliards d'exemples existent sur Internet, les NNIP s'appuient sur des données générées par une technique courante dans la recherche informatique appelée théorie fonctionnelle de la densité (DFT).

DFT est une approche numérique basée sur la physique qui utilise la mécanique quantique pour prédire comment les atomes interagissent dans les molécules et les matériaux. Bien que les simulations DFT soient très puissantes, elles sont également coûteuses en termes de calcul, et générer une grande quantité de données d'entraînement DFT peut prendre beaucoup de temps. L'apprentissage automatique a le potentiel d'accélérer ces simulations en agissant comme un modèle de substitution pour le DFT. Ce n'est que récemment que des ensembles de données DFT contenant 100 millions de points de données ont été publiés, les ensembles de données précédents atteignant un maximum d'environ 1 à 2 millions, ouvrant la voie à une mise à l'échelle cruciale pour les NNIP.

Cependant, les modèles NNIP actuels qui intègrent des contraintes physiques nécessitent souvent du matériel informatique, de la mémoire et du temps de traitement importants, et peuvent également ajouter des complications dans l'optimisation facile des paramètres du réseau neuronal. En revanche, ESCAIP n'inclut pas beaucoup de contraintes physiques intégrées ; au lieu de cela, il s’efforce de rendre le modèle d’apprentissage automatique aussi expressif que possible, notamment en concevant un nouveau mécanisme d’attention personnalisé pour le contexte atomistique.

Cette approche permet à EScAIP de capturer des modèles complexes dans les données et d'apprendre des informations physiques clés directement à partir des données elles-mêmes, en contournant le besoin de contraintes explicites. Par exemple, après une formation, EScAIP peut, sur de nouveaux systèmes atomiques inédits, cartographier avec précision n'importe quelle orientation atomique par rapport aux forces prédites, capturant ainsi une symétrie connue sous le nom d'équivariance rotationnelle.

« Le nouveau modèle EScAIP peut s'entraîner sur 100 millions de points de données en quelques jours, alors qu'un NNIP physiquement contraint nécessiterait des semaines ou des mois », a déclaré Blau. En conséquence, le nombre de groupes de recherche capables de former ces modèles augmente considérablement.

« Nous croyons vraiment qu'il faut aider les gens à poursuivre leurs objectifs scientifiques en utilisant des outils auparavant moins accessibles », a ajouté Qu. « L'EScAIP donne cette chance aux scientifiques disposant de différentes quantités de ressources. »

EScAIP constitue une amélioration significative par rapport aux NNIP de pointe précédents, avec des modèles entraînés atteignant les meilleures performances sur des ensembles de données de référence NNIP communs couvrant divers systèmes chimiques, y compris les catalyseurs (tels que le projet Open Catalyst), les matériaux (tels que le projet Materials). Project) et des molécules (telles que SPICE).

Sur des ensembles de données comme Open Catalyst, il s'agit également du premier modèle en tête du classement qui a été développé et formé uniquement par des chercheurs universitaires et des laboratoires nationaux, plutôt que par des équipes de grandes entreprises technologiques. Cependant, Qu et Krishnapriyan estiment que ce modèle doit être considéré comme un premier pas dans une nouvelle direction.

« Nous disons à la communauté scientifique : 'Hé, regardez par ici, explorons cette idée davantage' », a déclaré Krishnapriyan. « EScAIP est une première preuve de concept sur la manière de réfléchir à la mise à l'échelle des modèles d'apprentissage automatique dans le contexte de systèmes atomistiques, et représente désormais une « limite inférieure » de ce qui est possible. Nous pensons que c'est la direction dans laquelle nous devrions réfléchir. sur le terrain alors que nous entrons dans un avenir avec plus de données et de ressources informatiques.

Selon Krishnapriyan, l'EScAIP trouve ses origines dans un projet de recherche et développement dirigé par le laboratoire de Berkeley (LDRD), Développement de nouvelles méthodes d'apprentissage automatique basées sur la physique, qui a contribué à façonner ses idées fondamentales. Elle souligne que l'exploitation des vastes ressources GPU du Centre national de calcul scientifique de recherche énergétique (NERSC) du ministère de l'Énergie (DOE) était cruciale pour développer et former des modèles sur des ensembles de données à grande échelle. En utilisant plusieurs GPU simultanément, l'équipe a atteint des performances optimales sur l'ensemble de données Open Catalyst, une réussite remarquable, d'autant plus qu'elle est la seule équipe d'une entreprise non technologique à le faire avec beaucoup moins de ressources.