Pith. sign in

REVIEW 3 major objections 4 minor 18 references

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

T0 review · 3 major / 4 minor · reviewed 2026-08-15 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read Message-passing neural network potentials scale linearly with layers, not cubically, enabling 100-million-atom molecular dynamics.

desk verdict A genuinely useful parallelization scheme for one-hop MPNN potentials, with an O(T) argument that holds for REANN-class models; the 'any MPNN' framing runs ahead of the evidence. read the letter →

arxiv 2505.06711 v3 pith:TP7H2H6V submitted 2025-05-10 physics.chem-ph cs.LG

Junfan Xia , Bin Jiang This is my paper · ORCID
classification physics.chem-phcs.LG
keywords messagepassingneuralnetworkpotentialmoleculardynamicsparallelizationlinearscalingrecursivelyembeddedatomspatialdecompositionghostatomsmachinelearninginteratomic
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper sets out to overturn a standard objection to message-passing neural network (MPNN) potentials: that their parallel cost grows steeply with the number of message-passing layers, effectively as $T^3$ or $T^4$. It argues that, because each layer updates an atom from its immediate one-hop neighbors within the original cutoff, a spatially decomposed simulation only needs to exchange ghost-atom feature tensors once per layer, so total cost grows linearly with $T$. This claim matters because MPNNs are among the most accurate machine-learned potentials but have been considered impractical for very large molecular dynamics. Using its recursively embedded atom neural network, the accompanying implementation demonstrates near-linear strong and weak scaling and reports simulation of more than 100 million atoms at speeds comparable to strictly local models. The paper presents this as a general framework that other MPNN architectures can adopt.

What carries the argument

The central object is the ghost-atom message exchange across subdomain boundaries in each message-passing layer. A process stores its local atoms plus ghost atoms (atoms owned by other processes that sit inside the cutoff), and each layer it receives the already-computed feature tensors of those ghosts from their owning processes rather than recomputing them. This keeps each computational subgraph intact without redundant work; gradients are handled symmetrically by exchanging only the derivative components that involve neighbor processes, using vector–Jacobian products rather than full Jacobians.

What would settle it

Take an MPNN whose layer update contains a nonlocal operation—such as global softmax attention, an all-atom normalization constant, or an electrostatic term summed over the entire box—and run it under this spatial-decomposition scheme. If the result is either wrong energies at subdomain boundaries or communication volume that grows with system size rather than with surface area times $T$, the linear-scaling claim does not generalize to that model. The test is concrete: instrument per-layer communication volume at increasing box sizes.

Watch

Extended reading notes

Core claim

On the paper's own terms, the discovery is that the previously claimed $\mathcal{O}(T^3)$—or cumulatively $\mathcal{O}(T^4)$—growth of MPNN cost with layer number is an artifact of treating message passing as if each layer rebuilt a growing neighbor list. In practice, at each layer an atom's new feature is built from features of the original neighbors inside $r_c$, so one neighbor list suffices for every layer. In parallel spatial decomposition, each process computes features for its local atoms, sends the ghost-atom subset of those features to neighboring processes, and proceeds; no process recomputes a ghost atom's features from an incomplete neighborhood. Backward force evaluation is decomposed through the same chain rule, with vector–Jacobian products limiting inter-process gradient traffic. Integrated with the REANN architecture, this yields linear scaling with $T$, strong scaling to 1,000 CPU cores and 8 GPUs, and a 100-million-atom silver simulation whose per-atom speed matches strictly local models.

Load-bearing premise

The scheme assumes that at every layer, an atom's updated feature depends only on features of its original one-hop neighbors inside the cutoff; if an MPNN uses global normalization, long-range attention, or any nonlocal term beyond that neighborhood, ghost-atom messages alone will not be enough.

Editorial extensions

If this is right

  • An MPNN with many message-passing layers is no longer inherently harder to parallelize; each added layer raises cost only linearly, so depth can be used for accuracy without sacrificing scale.
  • Invariant and equivariant MPNNs that aggregate messages from neighboring atoms—SchNet-like, NequIP-like, and MACE-like architectures—can adopt the same communication pattern without structural changes.
  • MPNN-based molecular dynamics can reach system sizes previously reserved for strictly local potentials, as shown by the over-100-million-atom silver simulation.
  • By using vector–Jacobian products rather than full Jacobians, per-process gradient tensors shrink by a factor related to the feature-vector size, cutting both compute and communication in the backward pass.
  • The parallelization is compatible with local models as well: when $T=1$, the new implementation performs comparably to existing just-in-time parallelized REANN, showing negligible added overhead.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • A direct test of the generality claim would implement this ghost-message scheme in a second, independent MPNN architecture (for example, an equivariant model) and check whether linear-in-$T$ scaling persists at the same per-atom cost; the current evidence is on one architecture.
  • The speed advantage reported against strictly local Allegro likely reflects both the parallel scheme and REANN's cheaper invariant features; separating those effects would require the same algorithm wrapped around Allegro's architecture.
  • Models with genuine long-range physics, such as electrostatics or global attention, will need an additional communication layer on top of ghost messages; the scheme as stated applies to one-hop message passing only.
  • Direct GPU-to-GPU communication, listed by the authors as future work, would likely push the strong-scaling efficiency above the ~67–88% range observed on 2–8 GPUs.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 4 minor

Summary. The paper proposes a parallelization algorithm for message passing neural network (MPNN) potentials used in molecular dynamics. The central idea is that in each message-passing layer, only ghost-atom feature tensors need to be exchanged between processes, so that the cost and communication overhead scale linearly with the number of layers T rather than with the effective cutoff volume. The authors implement this with the REANN model (REANN-MPI) and report strong and weak scaling results on CPU and GPU benchmarks, including a simulation of roughly 100 million Ag atoms. They also compare performance against Allegro, SevenNet, MACE-MP-0, and Deep Potential, and argue that the linear-in-T scaling overturns earlier estimates of O(T^3)-O(T^4) cost for MPNNs.

Significance. If the algorithmic claims hold, the contribution is valuable: it demonstrates that standard one-hop MPNN potentials can be parallelized with halo-only communication, making large-scale molecular dynamics with MPNNs as tractable as with strictly local potentials. The conceptual clarification that the growth of the effective receptive field does not require explicit enumeration of a T*r_c neighbor list is correct and important for the community. The reverse-mode differentiation on a partitioned graph with communicated vector-Jacobian products is a sound foundation, and the authors provide their implementation as open-source code with benchmark data, which is a strength. However, the paper over-claims generality beyond the tested REANN architecture, and the absolute speed comparisons with other models are not controlled, which tempers the significance of the reported speed advantages.

major comments (3)
  1. [Abstract, Section 2.2] The abstract and conclusion claim that this is a 'general parallelization framework' that can 'empower various MPNN models,' but the algorithm relies on the assumption that each MP layer aggregates only one-hop neighbor features within the original cutoff. Models with global attention (e.g., SO3KRATES), long-range electrostatics, or charge equilibration (e.g., PhysNet, CENT) violate this assumption because their updates depend on atoms outside the ghost halo. The paper's own scope clause 'messages aggregated over neighboring atoms' covers only a subset of MPNN potentials. Please either implement the scheme on a second non-REANN MPNN architecture (e.g., SchNet or NequIP) to support the generality claim, or explicitly qualify the title, abstract, and conclusion to state that the framework applies to local one-hop message-passing architectures.
  2. [Section 2.3, Figure 7] The claimed speed advantages in Figure 7 rest on uncontrolled comparisons: the Allegro and MACE-MP-0 timings are taken from Refs. 47 and 54, which used different hardware, library versions, and implementation details. Thus the 'two orders of magnitude faster' statement is not a valid measure of the parallelization efficiency and could be dominated by differences in model architecture (REANN is a relatively small invariant model). To support the claim that the parallel algorithm makes MPNNs as fast as strictly local models, the authors should measure all models on the same hardware with the same software stack, or present the comparison with explicit caveats and error margins.
  3. [Section 2.3, Figure 5] The text states that 'Figure 5 offers ambiguous numerical evidence that a properly parallelized MPNN model scales linearly with T.' This appears to be a typo for 'unambiguous,' but as written it directly contradicts the subsequent sentence and undermines the central scaling claim. Please correct this and ensure the figure caption and text consistently describe the evidence as supporting linear scaling.
minor comments (4)
  1. [Section 2.3] The exact number of atoms in the claimed 'over 100 million' simulation is not specified; 32 nodes with 'roughly 3 million atoms per node' gives about 96 million. Please provide the exact atom count (e.g., 100,663,296) and the precise system size in the text or figure.
  2. [Section 2.2, Eqs. (3)-(5)] The notation for the backward pass is terse and some terms (e.g., the dimensions of ∂ρ_i^t/∂c_j^{t-1} and the exact communication patterns for a generic MPNN) are not defined. Please clarify the tensor shapes and state explicitly which vector-Jacobian products are exchanged between which processes in each layer.
  3. [Section 2.3, Figure 6 and text] The comparison with Deep Potential (Ref. 41) uses different core counts (768 vs 3360 cores). Reporting per-core efficiency or a normalized metric would make the comparison more interpretable than raw times per step per atom.
  4. [Figure 2 caption] The distinction between 'false impression' and 'actual growth' is clear conceptually, but the caption could be more explicit: in (b), the number of contributing paths grows, but the computational graph only needs the original neighbor list. Consider adding one sentence to state this directly.

Circularity Check

0 steps flagged · score 2.0 of 10

No significant circularity: the O(T) scaling claim follows from the stated one-hop neighbor-only message-passing recurrence, and the only self-citation (REANN) is a test-bed choice, not a load-bearing premise.

full rationale

Walking the derivation chain, the paper's central claim is that an MPNN can be parallelized by exchanging only ghost-atom features in each message-passing layer, giving total cost that scales linearly with the number of layers T. This is argued directly from the architecture equations: Eq. (1) computes features from neighbors within r_c, and Eq. (2) updates node features from those one-hop features. Each layer therefore performs the same local computation over the original neighbor list, so per-layer work is proportional to the number of local edges and inter-process communication is proportional to the halo. The O(T) statement is an algorithmic consequence of the explicitly stated one-hop message-passing restriction, not a fitted parameter renamed as a prediction. The timings in Figures 5-7 are measurements of the implemented REANN-MPI code compared against external baselines; the accuracy MAEs/RMSEs are standard benchmarking results and are not used to infer the scaling law. The only notable self-citation is REANN itself (ref. 17), the authors' own architecture used as the demonstration vehicle, but the correctness of the parallelization argument does not depend on REANN's accuracy or on any uniqueness theorem from that prior work, so the self-citation is not load-bearing. The broad phrase 'general parallelization framework' is somewhat wider than the single REANN implementation actually tested, and the paper explicitly scopes its claim to 'any MPNN models with messages aggregated over neighboring atoms'; that is a scope or evidence concern, not circularity. No target result is assumed in the construction of the algorithm. Score 2 reflects only the minor, non-load-bearing self-citation.

Assumptions & free parameters 0 free parameters · 3 assumptions · 0 invented entities

The central parallelization claim introduces no new physical entities and no free parameters fitted to data. It rests on standard domain-decomposition assumptions for MD and on the one-hop message-passing structure of MPNNs. The REANN model hyperparameters affect absolute speed but are not fitted to make the scaling argument work.

assumptions (3)
  • domain assumption An MPNN layer updates a node using only the features of its immediate neighbors within the original cutoff radius at the previous layer.
    Used throughout Section 2.2 and Eq. 2; it underlies the claim that ghost-atom features can be synchronized once per layer and that no expanded neighbor list for T-hop receptive fields is needed.
  • domain assumption Spatial domain decomposition with ghost atoms within the cutoff radius captures all interactions needed for correct energy and force evaluation.
    Standard MD parallelization assumption shown in Figure 3; the proposed algorithm builds on this existing decomposition rather than replacing it.
  • standard math Backward force evaluation can be partitioned across processes by communicating only selected vector-Jacobian products, not the full computational graph.
    This relies on the chain rule of tensor differentiation (Eqs. 3-5) and on the same one-hop dependency structure used in the forward pass.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Efficient Parallelization of Message Passing Neural Network Potentials for Large-scale Molecular Dynamics." pith.science (2026). https://pith.science/paper/TP7H2H6V

@misc{pith2026250506711,
  author       = {Pith},
  title        = {Pith review of: Efficient Parallelization of Message Passing Neural Network Potentials for Large-scale Molecular Dynamics},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/TP7H2H6V}},
  note         = {Machine review of arXiv:2505.06711}
}
read the original abstract

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.

Figures

Figures reproduced from arXiv: 2505.06711 by the authors.

Figure 1
Figure 1. Schematic diagram of the REANN framework. (a) Input and MP layers, (b) each MP layer where  denotes the tensor product, and (c) the output layer after T MP iterations [PITH_FULL_IMAGE:figures/full_fig_p019_1.png] view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

18 extracted references · 2 canonical work pages

  1. [1]

    INTRODUCTION R ecent years have witnessed the revolutionary development of machine learning (ML) method s in representing potential energy surfaces (PESs) of complex systems based on ab initio calculations. 1 - 39 These machine - learned potentials (MLPs) have achieved much higher flexibility than empirical force fields with much lower computational costs...

    work page 2006

  2. [12]

    Generalized Gradient Approximation Made Simple

    DOI: 10.1145/3581784.3627041. (56) Paszke, A.; Gross, S.; Massa, F.; Lerer, A.; Bradbury, J.; Chanan, G.; Kill een, T.; Lin, Z.; Gimelshein, N.; Antiga, L.; et al. PyTorch: an imperative style, high - performance deep learning library. Advances in Neural Information Processing Systems 2019 , 32 , 8024 – 8035. (57) Lopanitsyna, N.; Fraux, G.; Springer, M. ...

    arXiv

  3. [27]

    (17) Zhang, Y.; Xia, J.; Jiang, B

    DOI: 10.1007/s41061 - 021 - 00339 - 5. (17) Zhang, Y.; Xia, J.; Jiang, B. Physically motivated recursively embedded atom neural networks: Incorporating local completeness and nonlocality. Phys. Rev. Lett. 2021 , 127 (15), 156002. DOI: 10.1103/PhysRevLett.127.156002. (18) Schütt, K.; Unke, O.; Gastegger, M. Equivariant message passing for the prediction of...

    arXiv 2021

  4. [56]

    1016/0010 - 4655(95)00042 - E

    DOI: https://doi.org/10. 1016/0010 - 4655(95)00042 - E . (52) Brandstetter, J.; Hesselink, R.; van der Pol, E.; Bekkers, E. J.; Welling, M. Geometric and physical quantities improve E(3) equivariant message passing. arXiv preprint arXiv:2110.02905 2021 . DOI: https://arxiv.org/abs/2110.02905v1 . (53) Cheng, B.; Engel, E. A.; Behler, J.; Dellago, C.; Cerio...

    arXiv

  5. [137]

    (2) Skinner, L

    DOI: https://doi.org/10.1016/S0301 - 0104(00)00179 - 8 . (2) Skinner, L. B.; Benmore, C. J.; Neuefeind, J. C.; Parise, J. B. The structure of water around the compressibility minimum. The Journal of Chemical Physics 2014 , 141 (21), 214507. DOI: 10.1063/1.4902412 (acccessed 9/4/2023). (3) Chen, W.; Ambrosio, F.; Miceli, G.; Pasquarello, A. Ab initio Elect...

    doi:10.1016/s0301 2014

  6. [890]

    (62) Grimme, S.; Antony, J.; Ehrlich, S.; Krieg, H

    DOI: 10.1103/PhysRevLett.80.890 . (62) Grimme, S.; Antony, J.; Ehrlich, S.; Krieg, H. A consistent and accurate ab initio parametrization of density functional dispersion correction (DFT - D) for the 94 elements H - Pu. J. Chem. Phys. 2010 , 132 (15), 154104. DOI: doi: http://dx.doi.org/10.1063/1.3382344 . (63) Csonka, G. I.; Perdew, J. P.; Ruzsinszky, A....

    doi:10.1103/physrevlett.80.890

  7. [1115]

    (54) Batatia, I.; Benner, P.; Chiang, Y.; Elena, A

    DOI: 10.1073/pnas.1815117116. (54) Batatia, I.; Benner, P.; Chiang, Y.; Elena, A. M.; Kovács, D. P.; Riebesell, J.; Advincula, X. R.; Asta, M.; Avaylon, M.; Baldwin, W. J. A foundation model for atomistic materials chemistry. arXiv preprin t arXiv:2401.00096 2023 . DOI: https://arxiv.org/abs/2401.00096v1 . (55) Kozinsky, B.; Musaelian, A.; Johansson, A.; ...

    arXiv

  8. [1173]

    (7) Yao, K.; Herr, J

    DOI: 10.1137/15M1054183. (7) Yao, K.; Herr, J. E.; Toth, David W.; McKintyre, R.; Parkhill, J. The Ten sorMol - 0.1 model chemistry: a neural network augmented with long - range physics. Chem. Sci. 2018 , 9 (8), 2261 -

    doi:10.1137/15m1054183 2018

Show all 18 references
  1. [1840]

    (41) Lu, D.; Wang, H.; Chen, M.; Lin, L.; Car, R.; E, W.; Jia, W.; Zhang, L

    DOI: 10.1021/acs.jctc.8b00770. (41) Lu, D.; Wang, H.; Chen, M.; Lin, L.; Car, R.; E, W.; Jia, W.; Zhang, L. 86 PFLOPS Deep Potential Molecular Dynamics simulation of 100 million atoms with ab initio accuracy. Comput. Phys. Commun. 2021 , 259 , 107624. DOI: https://doi.org/10.1...

  2. [2097]

    (26) Han, B.; Yu, K

    DOI: 10.1021/acs.jctc.3c01254. (26) Han, B.; Yu, K. Refining potential ener gy surface through dynamical properties via differentiable molecular simulation. Nat. Commun. 2025 , 16 (1), 816. DOI: 10.1038/s41467 - 025 - 56061 - z. (27) Frank, J. T.; Unke, O. T.; Müller, K. - R.;...

    2025 doi

  3. [2269]

    (8) Schütt, K

    DOI: 10.1039/c7sc04934j. (8) Schütt, K. T.; Sauceda, H. E.; Kindermans, P. - J.; Tkatchenko, A.; Müller, K. - R. SchNet – A deep learning archite cture for molecules and materials. J. Chem. Phys. 2018 , 148 (24), 241722. DOI: 10.1063/1.5019779. (9) Zhang, L.; Han, J.; Wang, H....

    2018 doi

  4. [2379]

    (36) Mausenberger, S.; Müller, C.; Tkatchenko, A.; Marquetand, P.; González, L.; Westermayr, J

    DOI: 10.1021/acs.jctc.2c0 1038. (36) Mausenberger, S.; Müller, C.; Tkatchenko, A.; Marquetand, P.; González, L.; Westermayr, J. SpaiNN: equivariant message passing for excited - state nonadiabatic molecular dynamics. Chem. Sci. 2024 , 15 (38), 15880 - 15890. (37) Alavi, S. F.;...

    arXiv

  5. [3693]

    (12) Vargas - Hernández, R

    DOI: 10.1021/acs.jctc.9b00181. (12) Vargas - Hernández, R. A.; Guan, Y. ; Zhang, D. H.; Krems, R. V. Bayesian optimization for the inverse scattering problem in quantum reaction dynamics. New J. Phys. 2019 , 21 (2), 022001. DOI: 10.1088/1367 - 2630/ab0099. (13) Drautz, R. Atom...

    2019 doi

  6. [3834]

    (33) Pinheiro, M.; Ge, F.; Ferré, N.; Dral, P

    DOI: 10.1021/acs.jpclett.0c00527. (33) Pinheiro, M.; Ge, F.; Ferré, N.; Dral, P. O.; Barbatti, M. Choosing the right molecular machine learning potential. Chem. Sci. 2021 , 12 (43), 14396 - 14413, 10.1 039/D1SC03564A. DOI: 10.1039/D1SC03564A. (34) Stark, W. G.; Westermayr, J.;...

    2021 doi

  7. [4868]

    (49) Zhang, Y.; Hu, C.; Jiang, B

    DOI: 10.1021/acs.jctc.4c00190. (49) Zhang, Y.; Hu, C.; Jiang, B. Embedded atom neural network potentials: Efficien t and accurate machine learning with a physically inspired representation. J. Phys. Chem. Lett. 2019 , 10 (17), 4962 -

    2019 doi

  8. [4967]

    (50) Thompson, A

    DOI: 10.1021/acs.jpclett.9b02037. (50) Thompson, A. P.; Aktulga, H. M.; Berger, R.; Bolintineanu, D. S.; Brown, W. M.; Crozier, P. S.; in 't Veld, P. J.; Kohlmeyer, A.; Moore, S. G.; Nguyen, T. D.; et al. LAMMPS - a flexible simulation tool for particle - based materials model...

  9. [5421]

    (15) Unke, O

    DOI: 10.1021/acs.jctc.0c00347. (15) Unke, O. T.; Chmiela, S.; Gastegger, M.; Schütt, K. T.; Sauceda, H. E.; Müller, K. - R. SpookyNet: Learning force fields with electronic degrees of freedom and nonlocal effects. Nat. Commun. 20 21 , 12 (1), 7273. DOI: 10.1038/s41467 - 021 - ...

    doi

  10. [7607]

    (46) Behler, J.; Csányi, G

    DOI: 10.1038/s41467 - 024 - 51886 - 6. (46) Behler, J.; Csányi, G. Machine learning potentials for extended systems: a perspective. Eur. Phys. J. B 2021 , 94 (7), 142. DOI: 10.1140/epjb/s10051 - 021 - 00156 - 1. (47) Musaelian, A.; Batzner, S.; Johansson, A.; Sun, L.; Owen, C....

    2021 doi

Pith tools

Reviewed August 15, 2026 · model on record in the stance chip above.