EScAIP, l’IA qui accélère les simulations atomistiques à grande échelle

Pour concevoir une batterie plus durable, comprendre une réaction chimique ou imaginer un nouveau semi-conducteur, les scientifiques doivent souvent répondre à une question minuscule, mais redoutablement coûteuse à calculer : comment les atomes se déplacent-ils et interagissent-ils entre eux ? Une équipe de l’université de Berkeley et du Lawrence Berkeley National Laboratory propose une réponse fondée sur l’intelligence artificielle. Son architecture, baptisée EScAIP, entend appliquer aux simulations atomistiques une idée devenue centrale dans l’IA moderne : faire croître les modèles et les données de façon méthodique pour améliorer leurs performances.

Les travaux ont été acceptés à la conférence Neural Information Processing Systems, NeurIPS 2024, l’un des grands rendez-vous internationaux du machine learning. Leur présentation est prévue le 13 décembre 2024. L’enjeu ne se limite pas à gagner quelques minutes de calcul. En réduisant fortement les besoins en mémoire et le temps d’entraînement, les chercheurs espèrent rendre des outils de simulation avancés utilisables par davantage d’équipes scientifiques.

Pourquoi simuler les interactions entre atomes ?

La matière à notre échelle est déterminée par des phénomènes qui se jouent à l’échelle atomique. La forme d’une molécule, la stabilité d’un cristal, la manière dont des ions circulent dans une batterie ou le déroulement d’une réaction chimique dépendent des forces exercées entre les atomes.

Pour les étudier, les chercheurs utilisent notamment des calculs issus de la mécanique quantique. La théorie de la fonctionnelle de la densité, généralement désignée par l’acronyme anglais DFT, est l’une des approches les plus employées. Sans entrer dans son formalisme mathématique, elle permet d’estimer l’énergie d’un système atomique et les forces qui y agissent. Ces informations servent ensuite à prédire des propriétés physiques ou chimiques.

Le problème est que la précision a un prix. À mesure que le nombre d’atomes, de configurations possibles ou d’étapes temporelles augmente, les calculs deviennent très lourds. Ils mobilisent souvent des supercalculateurs pour des recherches concernant, par exemple, les batteries et les semi-conducteurs. Or, explorer un vaste espace de matériaux ou suivre des transformations chimiques complexes exige de répéter ces calculs à très grande échelle.

C’est là qu’intervient l’apprentissage automatique. Au lieu de résoudre intégralement les équations quantiques pour chaque nouvelle configuration, un modèle peut apprendre, à partir de résultats DFT, à approximer les énergies et les forces atomiques. Une fois entraîné, il peut effectuer des prédictions bien plus rapidement. Il ne remplace pas la DFT dans tous les cas, mais peut devenir un modèle de substitution particulièrement utile pour parcourir un grand nombre de scénarios.

Les NNIP, une alternative rapide aux calculs quantiques

La famille de modèles concernée porte le nom de potentiels interatomiques basés sur des réseaux neuronaux, ou NNIP pour neural network interatomic potentials. Leur rôle est de décrire, à partir de la position et de l’environnement des atomes, les interactions qui déterminent le comportement d’un matériau ou d’une molécule.

Ces potentiels neuronaux peuvent considérablement accélérer les simulations par rapport à des calculs quantiques directs. Ils sont ainsi prometteurs pour étudier des systèmes plus grands, sur des durées plus longues, ou pour tester davantage de candidats matériaux. Mais cette vitesse ne suffit pas : il faut aussi que le modèle soit suffisamment précis, qu’il puisse apprendre sur de très grands jeux de données et que son entraînement reste réaliste du point de vue matériel.

C’est précisément le verrou visé par l’équipe de Berkeley. Selon Aditi Krishnapriyan, co-autrice de l’étude, l’élaboration d’algorithmes capables de bien passer à l’échelle pour les NNIP a été moins explorée que dans d’autres branches du machine learning. Les avancées spectaculaires des grands modèles de langage ont pourtant montré qu’augmenter de manière organisée la taille des réseaux, le volume de données et les ressources de calcul peut produire des gains importants.

Eric Qu, doctorant à l’université de Berkeley et co-auteur, explique que l’équipe a cherché à bâtir une architecture d’apprentissage automatique différente, inspirée de ces méthodes de mise à l’échelle. L’objectif est d’améliorer la modélisation des mouvements et des interactions atomiques sans faire exploser les ressources nécessaires.

EScAIP, une architecture conçue pour passer à l’échelle

EScAIP signifie Efficiently Scaled Attention Interatomic Potential. Il s’agit d’une architecture NNIP qui mobilise des mécanismes d’attention, une famille de techniques devenue célèbre avec les grands modèles de langage. Dans ce contexte, l’attention ne sert pas à traiter des phrases : elle aide le modèle à identifier les relations pertinentes au sein d’un système atomique, c’est-à-dire quels éléments de l’environnement local comptent pour prédire les interactions.

L’ambition n’est donc pas de transformer une simulation chimique en conversation avec un chatbot. L’idée est de transposer à un problème scientifique les leçons tirées de l’entraînement de très grands réseaux neuronaux : architecture adaptée, exploitation de gros volumes de données et usage intensif des processeurs graphiques, ou GPU.

Les résultats rapportés par l’équipe sont significatifs. Par rapport aux modèles existants évoqués dans l’étude, la méthode réduirait la mémoire informatique nécessaire de plus de cinq fois et permettrait d’obtenir des résultats plus de dix fois plus rapidement. Ces deux dimensions sont liées : la mémoire disponible limite fréquemment la taille des systèmes, des lots de données et des modèles qu’il est possible de traiter sur un accélérateur.

Cette efficacité est particulièrement importante à une période où les jeux de données scientifiques prennent de l’ampleur. Produire les données DFT nécessaires à l’entraînement reste coûteux et chronophage. Mais, lorsqu’un volume de données existe, il faut aussi disposer d’un modèle capable de l’exploiter. EScAIP a été conçu dans cette perspective.

Que change l’entraînement sur 100 millions de données ?

Les chercheurs indiquent que leur modèle peut être entraîné sur 100 millions de points de données en quelques jours. À titre de comparaison, l’entraînement d’un NNIP soumis à des contraintes physiques peut demander des semaines, voire des mois, à cette échelle.

Cette différence ne signifie pas que les contraintes physiques sont inutiles. Elles peuvent intégrer dans un modèle des connaissances importantes sur les symétries, les lois ou les comportements attendus d’un système. Mais elles peuvent également compliquer la montée en taille des architectures. EScAIP met davantage l’accent sur l’expressivité du modèle, c’est-à-dire sa capacité à capter des motifs complexes présents dans les données atomiques.

La possibilité de traiter un tel volume en quelques jours modifie la nature des expérimentations possibles. Les équipes peuvent tester plus rapidement des choix d’architecture, comparer des réglages d’entraînement ou étudier des jeux de données plus larges. Ce raccourcissement du cycle entre idée, calcul et analyse est crucial dans les sciences computationnelles, où les délais de calcul conditionnent souvent le nombre d’hypothèses réellement examinées.

Des applications possibles en matériaux, chimie et médicaments

L’intérêt des simulations atomistiques est transversal. En science des matériaux, elles aident à comprendre pourquoi une structure est stable, comment elle se déforme ou comment elle conduit des charges électriques. Pour les batteries, ces outils peuvent contribuer à explorer les mécanismes qui affectent les électrodes, les électrolytes et le transport des ions. Dans les semi-conducteurs, ils permettent d’étudier les propriétés de matériaux à une échelle difficilement accessible par l’expérimentation seule.

En chimie, l’objectif peut être de décrire les étapes d’une réaction et de repérer les mécanismes qui conduisent à un produit plutôt qu’à un autre. Samuel Blau, chimiste computationnel au Berkeley Lab, souligne que le modèle aide les scientifiques à déterminer les mécanismes des réactions chimiques de manière beaucoup plus efficace. Mieux comprendre la chimie de systèmes réels peut ensuite fournir de nouvelles façons de les contrôler.

Le développement de médicaments fait aussi partie des domaines susceptibles de bénéficier de simulations plus rapides. Les molécules biologiques et les interactions chimiques qui les concernent sont complexes. Une méthode informatique ne remplace ni les essais en laboratoire ni les études cliniques, mais elle peut aider à sélectionner des pistes, à formuler des hypothèses et à réduire l’espace des candidats à examiner expérimentalement.

L’apport potentiel d’EScAIP réside donc moins dans une application unique que dans son rôle d’infrastructure de calcul. S’il permet à des chercheurs de simuler davantage de configurations atomiques avec un budget informatique plus maîtrisé, il peut accélérer les phases exploratoires de nombreux projets.

Une avancée issue de la recherche publique

L’équipe souligne également la provenance du projet. EScAIP a été développé et entraîné par des chercheurs universitaires et de laboratoires nationaux, sans l’appui des grandes entreprises technologiques. Ses origines remontent à un projet de recherche mené au Berkeley Lab avec le soutien du département de l’Énergie des États-Unis.

L’accès à des ressources GPU à grande échelle a joué un rôle déterminant. Ce point rappelle une réalité de l’IA scientifique : les bonnes idées algorithmiques comptent, mais leur mise à l’épreuve exige des infrastructures de calcul capables de traiter des jeux de données volumineux. L’équipe indique avoir atteint des performances remarquables sur le jeu de données Open Catalyst, utilisé pour évaluer des modèles dédiés aux matériaux et aux interactions atomiques.

Cette dimension publique est importante pour la reproductibilité et la diffusion des méthodes. Le code d’EScAIP est mis à disposition dans un dépôt logiciel. Pour les laboratoires qui disposent des données et du matériel requis, cela facilite l’examen des choix techniques, la reproduction des expériences et l’adaptation de l’architecture à d’autres problèmes scientifiques.

Ce qu’il faut surveiller après NeurIPS 2024

La présentation à NeurIPS 2024 constitue une étape de visibilité, non un point final. La question centrale sera de savoir dans quelle mesure les gains annoncés se maintiennent sur une diversité de matériaux, de molécules et de conditions de simulation. Un modèle performant sur des jeux de référence doit aussi démontrer sa robustesse lorsqu’il est confronté à des systèmes moins familiers ou à des données expérimentales.

Il faudra également observer le coût global de la méthode. EScAIP réduit fortement les besoins en mémoire et accélère les calculs selon les résultats présentés, mais sa formation à grande échelle repose elle-même sur des données DFT et sur des GPU puissants. La démocratisation dépendra donc autant de l’accès aux infrastructures que de la disponibilité de données de qualité.

Pour Krishnapriyan et Qu, EScAIP est un premier pas vers une réflexion plus large sur les lois de mise à l’échelle appliquées aux systèmes atomiques. À mesure que les ressources computationnelles et les données scientifiques augmentent, cette approche pourrait aider les simulations à passer d’un outil spécialisé et coûteux à un levier plus largement mobilisable dans la recherche sur la matière.