Vers l'entraînement de réseaux neuronaux graphiques à milliards de paramètres pour les simulations atomiques

Les progrès récents dans les Réseaux Neuronaux sur Graphes (GNNs) pour la modélisation des simulations atomiques ont le potentiel de révolutionner la découverte de catalyseurs, une étape clé pour réaliser les percées énergétiques nécessaires à la lutte contre le changement climatique. Cependant, les GNNs qui se sont avérés les plus efficaces pour cette tâche sont intensifs en mémoire car ils modélisent des interactions d'ordre supérieur dans les graphes, telles que celles entre des triplets ou quadruplets d'atomes, ce qui rend difficile l'évolutivité de ces modèles. Dans cet article, nous présentons le Parallélisme sur Graphes, une méthode permettant de distribuer les graphes d'entrée sur plusieurs GPU, nous permettant ainsi d'entraîner des GNNs très volumineux avec des centaines de millions ou des milliards de paramètres. Nous évaluons empiriquement notre méthode en augmentant le nombre de paramètres des modèles DimeNet++ et GemNet récemment proposés d'un facteur supérieur à un ordre de grandeur. Sur le grand ensemble de données Open Catalyst 2020 (OC20), ces modèles parallélisés sur graphes entraînent des améliorations relatives de 1) 15 % sur la métrique d'erreur absolue moyenne (MAE) des forces pour la tâche S2EF et 2) 21 % sur la métrique AFbT pour la tâche IS2RS, établissant ainsi de nouveaux résultats d'état de l'art.