Efficient Parallelization of Message Passing Neural Network Potentials for Large-scale Molecular Dynamics

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Xia, Junfan, Jiang, Bin
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866913884579823616
author Xia, Junfan
Jiang, Bin
author_facet Xia, Junfan
Jiang, Bin
contents Machine learning potentials have achieved great success in accelerating atomistic simulations. Many of them relying on atom-centered local descriptors are natural for parallelization. More recent message passing neural network (MPNN) models have demonstrated their superior accuracy and become increasingly popular. However, efficiently parallelizing MPNN models across multiple nodes remains challenging, limiting their practical applications in large-scale simulations. Here, we propose an efficient parallel algorithm for MPNN models, in which additional data communication is minimized among local atoms only in each MP layer without redundant computation, thus scaling linearly with the layer number. Integrated with our recursively embedded atom neural network model, this algorithm demonstrates excellent strong scaling and weak scaling behaviors in several benchmark systems. This approach enables massive molecular dynamics simulations on MPNN models as fast as on strictly local models for over 100 million atoms, vastly extending the applicability of the MPNN potential to an unprecedented scale. This general parallelization framework can empower various MPNN models to efficiently simulate very large and complex systems.
format Preprint
id arxiv_https___arxiv_org_abs_2505_06711
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Efficient Parallelization of Message Passing Neural Network Potentials for Large-scale Molecular Dynamics
Xia, Junfan
Jiang, Bin
Chemical Physics
Machine Learning
Machine learning potentials have achieved great success in accelerating atomistic simulations. Many of them relying on atom-centered local descriptors are natural for parallelization. More recent message passing neural network (MPNN) models have demonstrated their superior accuracy and become increasingly popular. However, efficiently parallelizing MPNN models across multiple nodes remains challenging, limiting their practical applications in large-scale simulations. Here, we propose an efficient parallel algorithm for MPNN models, in which additional data communication is minimized among local atoms only in each MP layer without redundant computation, thus scaling linearly with the layer number. Integrated with our recursively embedded atom neural network model, this algorithm demonstrates excellent strong scaling and weak scaling behaviors in several benchmark systems. This approach enables massive molecular dynamics simulations on MPNN models as fast as on strictly local models for over 100 million atoms, vastly extending the applicability of the MPNN potential to an unprecedented scale. This general parallelization framework can empower various MPNN models to efficiently simulate very large and complex systems.
title Efficient Parallelization of Message Passing Neural Network Potentials for Large-scale Molecular Dynamics
topic Chemical Physics
Machine Learning
url https://arxiv.org/abs/2505.06711