Scalable foundation interatomic potentials via message-passing pruning and graph partitioning
Source abstract
Abstract
Atomistic foundation models (AFMs) have great promise as accurate interatomic potentials, and have enabled data-efficient molecular dynamics simulations with near quantum mechanical accuracy. However, AFMs remain markedly slower at inference and are far more memory-intensive than conventional interatomic potentials, due to the need to capture a wide range of chemical and structural motifs in pre-training datasets requiring deep, parameter-rich model architectures. These deficiencies currently limit the practical use of AFMs in molecular dynamics (MD) simulations at extended temporal and spatial scales. To address this problem, we propose a general workflow for accelerating and scaling AFMs containing message-passing architectures. We find that removing low-contribution message-passing layers from AFM backbones serves as an effective pruning method, significantly reducing the parameter count while preserving the accuracy and data-efficiency of AFMs. Once pruned, these models become more accessible for large scale simulations via a graph-partitioned, GPU-distributed strategy, which we implement and demonstrate within the AFM fine-tuning platform MatterTune. We show that this approach supports million-atom simulations on both single and multiple GPUs, and enables task-specific large-scale simulations at nanosecond timescales with AFM-level accuracy.
Source: openalex_inverted_index · Quality: baseline
Research context
OpenAlex topics: Machine Learning in Materials Science · Block Copolymer Self-Assembly · Quantum many-body systems
SciVX discovery tags are terminology-based. OpenAlex topics are source-provided classifications.
References (80)
Show source-provided references
1. Cheng, T., Xiao, H. & Goddard III, W. A. Full atomistic reaction mechanism with kinetics for co reduction on cu (100) from ab initio molecular dynamics free-energy calculations at 298 K. Proc. Natl. Acad. Sci. USA 114, 1795–1800 (2017).10.1073/pnas.1612106114 ↗
2. Le, J.-B. et al. Modeling electrified pt(111)-had/water interfaces from ab initio molecular dynamics. JACS Au 1, 569–577 (2021).10.1021/jacsau.1c00108 ↗
3. Cheng, B., Engel, E. A., Behler, J., Dellago, C. & Ceriotti, M. Ab initio thermodynamics of liquid and solid water. Proc. Natl. Acad. Sci. USA 116, 1110–1115 (2019).10.1073/pnas.1815117116 ↗
4. Mo, Y., Ong, S. P. & Ceder, G. First principles study of the Li10GeP2S12 lithium super ionic conductor material. Chem. Mater. 24, 15–17 (2012).10.1021/cm203303y ↗
5. Islamov, M. et al. High-throughput screening of hypothetical metal-organic frameworks for thermal conductivity. npj Comput. Mater. 9, 11 (2023).10.1038/s41524-022-00961-x ↗
6. Richards, W. D. et al. Design and synthesis of the superionic conductor Na10Snp2S12. Nat. Commun. 7, 11009 (2016).10.1038/ncomms11009 ↗
7. Mo, Y., Ong, S. P. & Ceder, G. Insights into diffusion mechanisms in p2 layered oxide materials by first-principles calculations. Chem. Mater. 26, 5208–5214 (2014).10.1021/cm501563f ↗
8. Hohenberg, P. & Kohn, W. Inhomogeneous electron gas. Phys. Rev. 136, B864 (1964).10.1103/PhysRev.136.B864 ↗
9. Kohn, W. & Sham, L. J. Self-consistent equations including exchange and correlation effects. Phys. Rev. 140, A1133 (1965).10.1103/PhysRev.140.A1133 ↗
10. Raghavachari, K., Trucks, G. W., Pople, J. A. & Head-Gordon, M. A fifth-order perturbation comparison of electron correlation theories. Chem. Phys. Lett. 157, 479–483 (1989).10.1016/S0009-2614(89)87395-6 ↗
11. Beck, T. L. Real-space mesh techniques in density-functional theory. Rev. Mod. Phys. 72, 1041 (2000).10.1103/RevModPhys.72.1041 ↗
12. Payne, M. C., Teter, M. P., Allan, D. C., Arias, T. & Joannopoulos, aJ. Iterative minimization techniques for ab initio total-energy calculations: molecular dynamics and conjugate gradients. Rev. Mod. Phys. 64, 1045 (1992).10.1103/RevModPhys.64.1045 ↗
13. Kresse, G. & Furthmüller, J. Efficiency of ab-initio total energy calculations for metals and semiconductors using a plane-wave basis set. Comput. Mater. Sci. 6, 15–50 (1996).10.1016/0927-0256(96)00008-0 ↗
14. Marx, D. & Hutter, J. Ab initio molecular dynamics: theory and implementation. Mod. Methods Algorithms Quantum Chem. 1, 141 (2000).
15. Tuckerman, M. E. Ab initio molecular dynamics: basic concepts, current trends and novelapplications. J. Phys. Condens. Matter 14, R1297 (2002).10.1088/0953-8984/14/50/202 ↗
16. Bartók, A. P., Payne, M. C., Kondor, R. & Csányi, G. Gaussian approximation potentials: the accuracy of quantum mechanics, without the electrons. Phys. Rev. Lett. 104, 136403 (2010).10.1103/PhysRevLett.104.136403 ↗
17. Schutt, K. et al. Schnet: a continuous-filter convolutional neural network for modeling quantum interactions. In Proc. 31st International Conference on Neural Information Processing Systems 992–1002 (ACM, 2017).
18. Schutt, K., Unke, O. & Gastegger, M. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In Proc. International Conference on Machine Learning 9377–9388 (PMLR, 2021).
19. Wang, H., Zhang, L., Han, J. et al. Deepmd-kit: a deep learning package for many-body potential energy representation and molecular dynamics. Comput. Phys. Commun. 228, 178–184 (2018).10.1016/j.cpc.2018.03.016 ↗
20. Drautz, R. Atomic cluster expansion for accurate and transferable interatomic potentials. Phys. Rev. B 99, 014104 (2019).10.1103/PhysRevB.99.014104 ↗
21. Chen, C., Ye, W., Zuo, Y., Zheng, C. & Ong, S. P. Graph networks as a universal machine learning framework for molecules and crystals. Chem. Mater. 31, 3564–3572 (2019).10.1021/acs.chemmater.9b01294 ↗
22. Chen, C. & Ong, S. P. A universal graph deep learning interatomic potential for the periodic table. Nat. Comput. Sci. 2, 718–728 (2022).10.1038/s43588-022-00349-3 ↗
23. Batzner, S. et al. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nat. Commun. 13, 2453 (2022).10.1038/s41467-022-29939-5 ↗
24. Batatia, I., Kovacs, D. P., Simm, G., Ortner, C. & Csányi, G. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Adv. Neural Inf. Process. Syst. 35, 11423–11436 (2022).10.52202/068431-0830 ↗
25. Musaelian, A. et al. Learning local equivariant representations for large-scale atomistic dynamics. Nat. Commun. 14, 579 (2023).10.1038/s41467-023-36329-y ↗
26. Liao, Y.-L. & Smidt, T. Equiformer: equivariant graph attention transformer for 3D atomistic graphs. In International Conference on Learning Representations (ICLR) (2023).
27. Thompson, A. P. et al. Lammps-a flexible simulation tool for particle-based materials modeling at the atomic, meso, and continuum scales. Comput. Phys. Commun. 271, 108171 (2022).10.1016/j.cpc.2021.108171 ↗
28. Schutt, K. et al. Schnetpack: a deep learning toolbox for atomistic systems. J. Chem. Theory Comput. 15, 448–455 (2018).10.1021/acs.jctc.8b00908 ↗
29. Novikov, I. S., Gubaev, K., Podryabinkin, E. V. & Shapeev, A. V. The mlip package: moment tensor potentials with mpi and active learning. Mach. Learn. Sci. Technol. 2, 025002 (2020).10.1088/2632-2153/abc9fe ↗
30. Kapil, V. et al. i-pi 2.0: A universal force engine for advanced molecular simulations. Comput. Phys. Commun. 236, 214–223 (2019).10.1016/j.cpc.2018.09.020 ↗
31. Batatia, I. et al. A foundation model for atomistic materials chemistry. Preprint at https://doi.org/10.48550/arXiv.2401.00096 (2023).10.48550/arXiv.2401.00096 ↗
32. Merchant, A. et al. Scaling deep learning for materials discovery. Nature 624, 80–85 (2023).10.1038/s41586-023-06735-9 ↗
33. Zhang, D. et al. Dpa-2: a large atomic model as a multi-task learner. npj Comput. Mater. 10, 293 (2024).10.1038/s41524-024-01493-2 ↗
34. Shoghi, N. et al. From molecules to materials: pre-training large generalizable models for atomic property prediction. In International Conference on Learning Representations (ICLR) (2024).
35. Yang, H. et al. Mattersim: a deep learning atomistic model across elements, temperatures and pressures. Preprint at https://doi.org/10.48550/arXiv.2405.04967 (2024).10.48550/arXiv.2405.04967 ↗
36. Neumann, M. et al. Orb: a fast, scalable neural network potential https://arxiv.org/abs/2410.22570 (2024).
37. Rhodes, B. et al. Orb-v3: atomistic simulation at scale https://arxiv.org/abs/2504.06231 (2025).
38. Liao, Y.-L., Wood, B. M., Das, A. & Smidt, T. EquiformerV2: Improved equivariant transformer for scaling to higher-degree representations. In International Conference on Learning Representations (ICLR) (2024).
39. Fu, X. et al. Learning smooth and expressive interatomic potentials for physical property prediction. Preprint at https://doi.org/10.48550/arXiv.2502.12147 (2025).10.48550/arXiv.2502.12147 ↗
40. Kong, L., Shoghi, N., Hu, G., Li, P. & Fung, V. Mattertune: an integrated, user-friendly platform for fine-tuning atomistic foundation models to accelerate materials simulation and discovery. Digital Discovery 4, 2253–2262 (2025).10.1039/D5DD00154D ↗
41. Radova, M., Stark, W. G., Allen, C. S., Maurer, R. J. & Bartók, A. P. Fine-tuning foundation models of materials interatomic potentials with frozen transfer learning. npj Comput. Mater. 11, 237 (2025).10.1038/s41524-025-01727-x ↗
42. Kaur, H. et al. Data-efficient fine-tuning of foundational models for first-principles quality sublimation enthalpies. Faraday Discuss. 256, 120–138 (2025).10.1039/D4FD00107A ↗
43. Du, X. et al. Machine-learning-accelerated simulations to enable automatic surface reconstruction. Nat. Comput. Sci. 3, 1034–1044 (2023).10.1038/s43588-023-00571-7 ↗
44. Jumper, J. et al. Highly accurate protein structure prediction with alphafold. Nature 596, 583–589 (2021).10.1038/s41586-021-03819-2 ↗
45. Shi, T. et al. High active material loading in all-solid-state battery electrode via particle size optimization. Adv. Energy Mater. 10, 1902881 (2020).10.1002/aenm.201902881 ↗
46. Holstun, T. et al. Accelerating the electrochemical formation of the δ phase in manganese-rich rocksalt cathodes. Adv. Mater. 37, 2412871 (2025).10.1002/adma.202412871 ↗
47. Li, Q., Han, Z. & Wu, X.-M. Deeper insights into graph convolutional networks for semi-supervised learning. In Proc. AAAI Conference on Artificial Intelligence Vol. 32 (AAAI, 2018).10.1609/aaai.v32i1.11604 ↗
48. Oono, K. & Suzuki, T. Graph neural networks exponentially lose expressive power for node classification. Preprint at https://doi.org/10.48550/arXiv.1905.10947 (2019).10.48550/arXiv.1905.10947 ↗
49. Cai, C. & Wang, Y. A note on over-smoothing for graph neural networks. Preprint at https://doi.org/10.48550/arXiv.2006.13318 (2020).10.48550/arXiv.2006.13318 ↗
50. Chen, D. et al. Measuring and relieving the over-smoothing problem for graph neural networks from the topological view. In Proc. AAAI Conference on Artificial Intelligence, Vol. 34, 3438–3445 (AAAI, 2020).10.1609/aaai.v34i04.5747 ↗
51. Fuchs, P., Chen, W., Thaler, S. & Zavadlav, J. chemtrain-deploy: a parallel and scalable framework for machine learning potentials in million-atom md simulations. J. Chem. Theory Comput. 21, 7550–7560 (2025).10.1021/acs.jctc.5c00996 ↗
52. Bradbury, J. et al. Jax: composable transformations of python+ numpy programs. http://github.com/jax-ml/jax Accessed: 2025-09-25 (2018).
53. Han, K., Deng, B., Farimani, A. B. & Ceder, G. Distmlip: a distributed inference platform for machine learning interatomic potentials. Preprint at https://doi.org/10.48550/arXiv.2506.02023 (2025).10.48550/arXiv.2506.02023 ↗
54. Yu, X., Bates, J., Jellison, G. & Hart, F. A stable thin-film lithium electrolyte: lithium phosphorus oxynitride. J. Electrochem. Soc. 144, 524 (1997).10.1149/1.1837443 ↗
55. Bates, J., Dudney, N., Neudecker, B., Ueda, A. & Evans, C. Thin-film lithium and lithium-ion batteries. Solid State Ion. 135, 33–45 (2000).10.1016/S0167-2738(00)00327-1 ↗
56. Westover, A. S. et al. Plasma synthesis of spherical crystalline and amorphous electrolyte nanopowders for solid-state batteries. ACS Appl. Mater. Interfaces 12, 11570–11578 (2020).10.1021/acsami.9b20812 ↗
57. Li, W., Ando, Y., Minamitani, E. & Watanabe, S. Study of Li atom diffusion in amorphous Li3PO4 with neural network potential. J. Chem. Phys. 147, 214106 (2017).10.1063/1.4997242 ↗
58. Larsen, A. H. et al. The atomic simulation environment-a Python library for working with atoms. J. Phys. Condens. Matter 29, 273002 (2017).10.1088/1361-648X/aa680e ↗
59. AlKaabi, K., Prasad, D. L., Kroll, P., Ashcroft, N. & Hoffmann, R. Silicon monoxide at 1 atm and elevated pressures: crystalline or amorphous? J. Am. Chem. Soc. 136, 3410–3423 (2014).10.1021/ja409692c ↗
60. Erhard, L. C., Rohrer, J., Albe, K. & Deringer, V. L. Modelling atomic and nanoscale structure in the silicon–oxygen system through active machine learning. Nat. Commun. 15, 1927 (2024).10.1038/s41467-024-45840-9 ↗
61. Tauster, S., Fung, S., Baker, R. & Horsley, J. Strong interactions in supported-metal catalysts. Science 211, 1121–1125 (1981).10.1126/science.211.4487.1121 ↗
62. Liu, S. et al. Ultrastable au nanoparticles on titania through an encapsulation strategy under oxidative atmosphere. Nat. Commun. 10, 5790 (2019).10.1038/s41467-019-13755-5 ↗
63. Gao, C., Lyu, F. & Yin, Y. Encapsulated metal nanoparticles for catalysis. Chem. Rev. 121, 834–881 (2020).10.1021/acs.chemrev.0c00237 ↗
64. Polo-Garzon, F. et al. In situ strong metal–support interaction (smsi) affects catalytic alcohol conversion. ACS Catal. 11, 1938–1945 (2021).10.1021/acscatal.0c05324 ↗
65. Karypis, G. & Kumar, V. A fast and high quality multilevel scheme for partitioning irregular graphs. SIAM J. Sci. Comput. 20, 359–392 (1998).10.1137/S1064827595287997 ↗
66. Behnel, S. et al. Cython: the best of both worlds. Comput. Sci. Eng. 13, 31–39 (2010).10.1109/MCSE.2010.118 ↗
67. Falcon, W. A. Pytorch lightning. https://github.com/Lightning-AI/pytorch-lightning Accessed: 2025-09-25 (2019).
68. Li, S. et al. Pytorch distributed: experiences on accelerating data parallel training. Preprint at https://doi.org/10.48550/arXiv.2006.15704 (2020).10.48550/arXiv.2006.15704 ↗
69. Loshchilov, I. & Hutter, F. Decoupled weight decay regularization. Preprint at https://doi.org/10.48550/arXiv.1711.0510 (2017).10.48550/arXiv.1711.0510 ↗
70. Kresse, G. & Furthmüller, J. Efficient iterative schemes for ab initio total-energy calculations using a plane-wave basis set. Phys. Rev. B 54, 11169 (1996).10.1103/PhysRevB.54.11169 ↗
71. Kresse, G. Hafner j. ab inito molecular dynamics for metals. Phys. Rev. B 47, 558–561 (1993).10.1103/PhysRevB.47.558 ↗
72. Perdew, J. P., Burke, K. & Ernzerhof, M. Generalized gradient approximation made simple. Phys. Rev. Lett. 77, 3865 (1996).10.1103/PhysRevLett.77.3865 ↗
73. Kresse, G. & Joubert, D. From ultrasoft pseudopotentials to the projector augmented-wave method. Phys. Rev. B 59, 1758 (1999).10.1103/PhysRevB.59.1758 ↗
74. Blöchl, P. E. Projector augmented-wave method. Phys. Rev. B 50, 17953 (1994).10.1103/PhysRevB.50.17953 ↗
75. Sun, J., Ruzsinszky, A. & Perdew, J. P. Strongly constrained and appropriately normed semilocal density functional. Phys. Rev. Lett. 115, 036402 (2015).10.1103/PhysRevLett.115.036402 ↗
76. Monkhorst, H. J. & Pack, J. D. Special points for brillouin-zone integrations. Phys. Rev. B 13, 5188 (1976).10.1103/PhysRevB.13.5188 ↗
77. Dwivedi, V. P. et al. Benchmarking graph neural networks. J. Mach. Learn. Res. 24, 1–48 (2023).
78. Rusch, T. K., Bronstein, M. M. & Mishra, S. A survey on oversmoothing in graph neural networks. Preprint at https://doi.org/10.48550/arXiv.2303.10993 (2023).10.48550/arXiv.2303.10993 ↗
79. Kornblith, S., Norouzi, M., Lee, H. & Hinton, G. Similarity of neural network representations revisited. In Proc. International Conference on Machine Learning 3519–3529 (PMlR, 2019).
80. Hirata, A. et al. Atomic-scale disproportionation in amorphous silicon monoxide. Nat. Commun. 7, 11591 (2016).10.1038/ncomms11591 ↗