Library

The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains

Scaling has been critical in improving model performance and generalization in machine learning. It involves how a model's performance changes with increases in model size or input data, as well as how efficiently computational resources are utilized to support this growth. Despite successes in other areas, the study of scaling in Neural Network Interatomic Potentials (NNIPs) remains limited. NNIPs act as surrogate models for ab initio quantum mechanical calculations. The dominant paradigm here is to incorporate many physical domain constraints into the model, such as rotational equivariance. We contend that these complex constraints inhibit the scaling ability of NNIPs, and are likely to lead to performance plateaus in the long run. In this work, we take an alternative approach and start by systematically studying NNIP scaling strategies. Our findings indicate that scaling the model through attention mechanisms is efficient and improves model expressivity. These insights motivate us to develop an NNIP architecture designed for scalability: the Efficiently Scaled Attention Interatomic Potential (EScAIP). EScAIP leverages a multi-head self-attention formulation within graph neural networks, applying attention at the neighbor-level representations. Implemented with highly-optimized attention GPU kernels, EScAIP achieves substantial gains in efficiency--at least 10x faster inference, 5x less memory usage--compared to existing NNIPs. EScAIP also achieves state-of-the-art performance on a wide range of datasets including catalysts (OC20 and OC22), molecules (SPICE), and materials (MPTrj). We emphasize that our approach should be thought of as a philosophy rather than a specific model, representing a proof-of-concept for developing general-purpose NNIPs that achieve better expressivity through scaling, and continue to scale efficiently with increased computational resources and training data.

cs.LGChemistryComputer Science
Source

Paper

Similar papers

Data efficiency and extrapolation trends in neural network interatomic potentialsJoshua A Vita, Daniel Schwalbe-Koda2023 · Machine Learning: Science and Technology · 24 citations In Library
Beyond Scaling: Chemical Intuition as Emergent Ability of Universal Machine Learning Interatomic PotentialsShinnosuke Hattori, Kohei Shimamura, A. Nakano et al.2025 · 1 citation In Library
Neural scaling of deep chemical modelsNathan C Frey, Ryan Soklaski, Simon Axelrod et al.2023 · Nature Machine Intelligence · 178 citations In Library
Scalable Parallel Algorithm for Graph Neural Network Interatomic Potentials in Molecular Dynamics Simulations.Yutack Park, Jaesun Kim, Seungwoo Hwang et al.2024 · Journal of Chemical Theory and Computation · 259 citations In Library
Implementing a neural network interatomic model with performance portability for emerging exascale architecturesSaaketh Desai, S. Reeve, J. Belak2020 · Computer Physics Communications · 14 citations In Library
Injecting Domain Knowledge from Empirical Interatomic Potentials to Neural Networks for Predicting Material PropertiesZeren Shui, Daniel S. Karls, Mingjian Wen et al.2022 · Neural Information Processing Systems · 10 citations In Library
The Design Space of E(3)-Equivariant Atom-Centered Interatomic PotentialsIlyes Batatia, Simon L. Batzner, D. Kov'acs et al.2022 · Nature Machine Intelligence · 361 citations In Library
Egret-1: Pretrained Neural Network Potentials for Efficient and Accurate Bioorganic SimulationE. Mann, Corin C. Wagen, Jonathon E. Vandezande et al.2025 · 15 citations In Library
E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic PotentialsSimon L. Batzner, Albert Musaelian, Lixin Sun et al.2021 · Nature Communications · 2.2k citations In Library
Benchmarking chemically scalable machine-learning interatomic potentials for large-scale simulations of multicomponent alloysShuang Fei, Penghua Ying, Kai Liu et al.2026 · PHYSICAL REVIEW MATERIALS · 0 citations In Library

Peer review

Accept (poster)NeurIPS 2024 · via OpenReview
Reviewer wbCz5/10 · confidence 5/52024-07-11

Summary

This paper introduces the Efficient Graph Attention Potential (EGAP), a new architecture for Neural Network Interatomic Potentials (NNIP) designed for scalability and efficiency. The authors investigate scaling strategies for NNIP models and propose a model that leverages optimized self-attention mechanisms. They evaluate EGAP on the MD22 and OC20 datasets, claiming state-of-the-art performance and improved efficiency compared to existing models.

Strengths

1. The paper addresses an important challenge in the field of NNIPs by focusing on scalability and efficiency. 2. The presentation of the model illustration is clear and appealing, aiding in understanding the proposed architecture. 3. The authors provide an investigation into scaling strategies for NNIP models, which could be valuable for future research in this area.

Weaknesses

1. Dataset split inconsistency: There is a significant issue with the dataset split on the MD22 dataset. The authors mention using a 95:5 train-test split, stating: "We use an EGAP model with six blocks and 48M parameters on each system and evaluate the performance on the test set (train-test split 95:5)." However, this deviates from the splits used in previous works [1]. In prior studies, the training data was predefined (e.g., the AT-AT target utilized 3k training data, which is 15% of the total data). The 95:5 split in previous works was applied to the training data for train-validation splitting, not for train-test splitting. For instance, the AT-AT target data was split into 2,850 train, 150 validation, and 17,001 test samples. By utilizing a 95:5 train-test split, the authors have made their results incomparable to baseline models. 2. Inconsistent dataset usage: The authors claim, "The total dataset was used in the training and evaluation of EGAP, but the other models only used a subset of the dataset. We do not anticipate a significant difference in the performance of the models because the dataset is composed of very similar equilibrium structures, and improvement from increasing dataset size saturates quickly." This statement is problematic for two reasons: a) The baseline models used the remaining data as a test set, not just a subset of the training data. By including more of the dataset in their training data, the authors have reduced the size of the test set. b) The authors have effectively used a different train/validation/test set compared to baseline models, rather than just using additional training data. c) The claim about performance differences with dataset size is not supported by evidence in the paper. 3. Limited experimental scope: Experiments on the OC20 dataset are conducted only on the 2M split of S2EF, rather than the larger S2EF-All or S2EF-All+MD splits. Given that the main motivation is efficiency, it would be more fitting to test on the larger set. This limitation significantly weakens the paper's claims about scalability and efficiency. 4. Inconclusive performance difference: On the OC20 2M split, EGAP underperforms on the force metric while outperforming on the energy metric. Since there is a regulation between these metrics during training, prioritizing one over the other can result in these outcomes even for the same model. This renders the performance difference unconvincing. 5. Lack of scalability experiments: The authors didn't provide any scalability experiments in terms of model size, which leaves their scalability claims unsupported. Without demonstrating how the model performs as its size increases, it's difficult to evaluate the true scalability of EGAP. 6. Result discrepancies: The results presented don't align with those provided in [1]. For example, in [1], the Tetrasaccharide target for VisNet-LSRM is reported as 0.1055 energy (kcal/mol) and 0.0767 forces (kcal/mol/Å), which equals 4.574914 energy (meV) = 0.05258521839 (meV/atom) and 3.326028 force (meV). However, the authors report these as 0.044 (meV/atom) and 5.0 force (meV). These discrepancies raise concerns about the accuracy of the reported results and the fairness of comparisons. [1] Li, Y., Wang, Y., Huang, L., Yang, H., Wei, X., Zhang, J., Wang, T., Wang, Z., Shao, B., and Liu, T.Y. (2024). Long-short-range message-passing: A physics-informed framework to capture non-local interaction for scalable molecular dynamics simulation. In The Twelfth International Conference on Learning Representations.

Questions

1. Why didn't the authors follow the same dataset split procedure used by baseline models, especially for the MD22 dataset? 2. Can the authors provide scalability experiments showing how EGAP performs as model size increases? 3. Why didn't the authors test EGAP on larger splits of the OC20 dataset (S2EF-All or S2EF-All+MD) to better demonstrate its efficiency advantages? 4. Can the authors explain the discrepancies between their reported results for baseline models and those in the original papers, particularly for the MD22 dataset? 5. How do the authors justify their claim that increasing dataset size wouldn't significantly affect model performance, given that they used a different train/test split than baseline models?

Rating

5

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

The paper does not adequately address the limitations of their approach, particularly regarding the dataset split inconsistencies and the lack of scalability experiments. These limitations significantly impact the comparability and validity of their results.

Authorsrebuttal2024-08-07

Rebuttal by Authors

Thank you for your constructive feedback and valuable suggestions! We address your concerns as follows: **Question: dataset split and inconsistent dataset usage of MD22** Thank you for the suggestions. We have revised our MD22 experiments to have the same train/val split as the baselines. Specifically, we follow the procedure used in MACE [1]: we use the number of samples used in training sgdml [2] as training+validation set (95:5), and the rest as test set. In the updated experiment, our model is better than state-of-the-art models like MACE [3] on almost all of the molecules (see general comment for details). **Question: limited experimental scope (need to test on OC20-All+MD) and inconclusive performance difference on OC20-2M** We have updated the results of EGAP on OC20 2M and All+MD split, with comparisons of relevant baselines. Our model is now state-of-the-art on both OC20 All+MD and OC20 2M, and 10x faster at training, 30x faster at inference than the previous best-performing models (EquiformerV2 [4]) (see general comment for details). **Question: lack of scalability experiments** We have added an experiment demonstrating the scalability of EGAP in the attached pdf (Figure 1). The details can be found in the overall response. From the figure, we see that EGAP scales well when data size and number of parameters increases. It also utilizes computation resources more efficiently, as it takes 10x less training time than EquiformerV2 and achieves better results. **Question: result discrepancies** Thank you for pointing this out. We took the numbers directly from the MACE paper [1], as we are using the MACE data split in the training. We have added a new row in Table 2 of the pdf for the recalculated results in the VisNet-LSRM paper [5]. --- [1] Kovács, D. P., Batatia, I., Arany, E. S., & Csányi, G. (2023). Evaluation of the MACE force field architecture: From medicinal chemistry to materials science. The Journal of Chemical Physics, 159(4). [2] Batatia, I., Kovacs, D. P., Simm, G., Ortner, C., & Csányi, G. (2022). MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35, 11423-11436. [3] Chmiela, S., Vassilev-Galindo, V., Unke, O. T., Kabylda, A., Sauceda, H. E., Tkatchenko, A., & Müller, K. R. (2023). Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9(2), eadf0873. [4] Liao, Y. L., Wood, B. M., Das, A., & Smidt, T. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. In The Twelfth International Conference on Learning Representations. [5] Li, Y., Wang, Y., Huang, L., Yang, H., Wei, X., Zhang, J., ... & Liu, T. Y. Long-Short-Range Message-Passing: A Physics-Informed Framework to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation. In The Twelfth International Conference on Learning Representations.

Reviewer L3rz5/10 · confidence 3/52024-07-11

Summary

Following an investigation into the scaling of an attention-based NNIP model (EquiformerV2), this paper finds that it is better to increase the number of parameters in the attention mechanism than to increase the number of parameters using higher-order tensors. Based on these results, the paper proposes a model architecture that significantly reduces memory usage and accelerates training while scoring well on OC-20 and MD-22.

Strengths

- The model achieves good results despite restricting itself to fairly limited physical data. This is evidence in support of the paper's claim that scaling alone is sufficient to learn physical symmetries. - BOO appears a promising method of encoding rotation-invariant bond information.

Weaknesses

- The paper appears to treat "efficient" and "scalable" as synonyms throughout, but there would appear to be a meaningful distinction, one that gets to a significant limitation of this paper. A non-scalable model performs badly on large molecules. An inefficient model takes excessive resources to achieve its performance. There is no demonstration that this model improves the scalability of current architectures. Thus, the contribution of this paper appears to be limited to a potential reduction in training times. There is no enhancement to the accuracy or scope of pre-existent NNIP models. If the paper means to claim that its proposed architecture scales better, it should show that as its number of parameters increase, its performance continues to improve where the performance of similar models plateaus. Various implementations of this model with different sizes are not described, so it would appear inaccurate to assert the superior scalability of this model. - Limited ablation studies in general. It is not clear which change gives the model its improvement in efficiency over EquiformerV2. Is it the use of BOO? Advancements in attention mechanism efficiency borrowed from NLP and CV? Simply that the model does not waste flops on higher-order tensors that do not improve its performance? Something else? - The model marginally outperforms EquiformerV2 on OC20, but it has marginally more parameters. It is not clear whether there is any advancement in architecture, or whether the model is simply better because it has slightly more space to learn. - It is not clear whether higher-order tensors are inefficient only as employed in EquiformerV2, or whether they are inefficient in general. Thus, the conclusion that scaling attention is superior to scaling the order of the tensors appears premature. In the current case, the investigation of scaling strategies is no more powerful than an extended ablation study of EquiformerV2.

Questions

- An apples-to-apples comparison to the other models in Table 2, standardizing the size of the training dataset, would be welcome. - More details on the way that the attention mechanism is scaled would be welcome. Where, exactly, are the extra parameters spent? - The details of the output block are not completely clear, but there seems to be no need to lose the rotation equivariance of the force predictions at any point here. Instead of feeding the edge unit vectors through a standard FFNN, it seems that they could be fed through an equivariant neural network as in https://arxiv.org/pdf/2102.03150. This would retain the equivariance of the model's force predictions as a whole. - JMP-L (https://arxiv.org/pdf/2310.16802) significantly outperforms this model on MD-22, yet the abstract claims that this model is SOTA. Clarifying that this model is SOTA among a certain class of models would be worthwhile. - More details on the training procedure and statistics are welcome, since a significant contribution of the paper is its reduction in the computer memory and training time requirements of NNIP models.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

See above.

Reviewer VH7W7/10 · confidence 4/52024-07-15

Summary

Authors explore scaling strategies for ML interatomic potentials, which are alternatives to the increase of spherical order L, which is costly. By leveraging the insights they extract from ablations, and avoiding the inefficiency of equivariance, they build a new invariant model based on features containing angular power spectra. Experiments show that this new model achieves SOTA on MD22.

Strengths

- The proposed model achieves SOTA performance on MD22, and comparable performance to spherical equivariant model such as EquiformerV2, by means of invariant features. - Experiments and ablations are extensive - The model stability is probed - The investigation of alternatives to equivariance is a necessary topic

Weaknesses

- Some other experiment showcasing the model accuracy of some other benchmark dataset with more models would have been beneficial - Table 1 is a bit poor, perhaps authors could use a bit of extra space to include some other ways of dealing with 'bond directional features'. - The authors could have used the extra space they have to perform some other experiment.

Questions

- It is not clear to me up to which L the BOO is summed. In the same way, though not affecting my evaluation, I am curious how the inclusion of different Ls impacts the accuracy. - How does the model inference time compare to other models? (Typo in line 124)

Rating

7

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors say they discuss some limitations in section 3. However, and given that they have spare space, I would expect a small limitations paragraph.

Reviewer PKom6/10 · confidence 2/52024-07-18

Summary

The paper presents a novel approach to predict properties of some 3d molecules based on efficient transformer-based graph neural networks. The paper investigates the computational bottlenecks of the previous approaches and addresses them by using simpler features and designing a novel graph attention architecture. The method shows good performance and is efficient, which is shown on two datasets.

Strengths

1. The problem of efficiently predicting complex molecule properties is important in practice. 2. The paper has interesting observations about scaling behaviors. 3. The paper is easy to follow and written and presented well. 4. The results are mostly convincing.

Weaknesses

1. The self-attention described in L174 is computed over edges which seems unusual in the graph transformer methods that compute self-attention over nodes. The number of edges is much bigger so it's surprising that the model is still efficient. Does the model have to store a square E x E matrix at some point? What's the motivation of computing self-attention over edges? Why not to add edge features as some kind of biases in the self-attention matrix over nodes like in Graphormer? 2. In L206, "The total dataset was used in the training and evaluation of EGAP, but the other models only used a subset of the dataset." makes the comparison not fair even if the authors "do not anticipate a significant difference". Using a smaller dataset for EGAP or a full dataset for the baselines is expected. On the OC20 dataset, do all the methods use the same train/val/test splits? 3. Why EquiformerV2 is not compared to in Table 2:? 4. Can the authors provide runtime and memory comparison for different molecule sizes during testing? This can give practitioners important insights. 5. It would be important to ablate which components of the proposed method contribute to better runtime and memory. For example, how much FlashAttention helps, how much using a small order L helps, etc. I believe some of these components like FlashAttention can be easily incorporated to the baselines. 6. Other ablations such as EGAP model size and dataset size vs performance would be interesting to see to better understand method's behavior under different regimes. Other minor issues: - L62 - broken reference - fig. 1 must be referenced in the text - fig. 1 "Right: Force MAE vs. Datasize. " is not correct or the plot xlabel is not correct - In Table 2, for Fatty acid h(r) is missing for EGAP

Questions

see above

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

not discussed explicitly

Reviewer PKom2024-08-12

The response makes some important clarifications and adds useful experiments, therefore I raise the score.

Reviewer L3rz2024-08-13

I acknowledge that I have read the rebuttal. Can you also check the results of ViSNet-LSRM on MD22 is correct? I will increase my score to borderline accept.

Authorsrebuttal2024-08-13

Thank you for your response. We took the reported numbers of VisNet-LSRM from the original paper and the MACE paper (in separate rows). During training, we used the same train/val/test dataset splits as MACE (the VisNet-LSRM data splits for MD22 are not available). There is also a unit conversion between the two papers, which we accounted for in our row in the table in the PDF. Please let us know if that clarifies your question.

Reviewer L3rz2024-08-13

Can you double check? The results are different from their latest version.

Authorsrebuttal2024-08-13

For the MD22 table in the attached pdf, in the VisNet-LSRM (MACE) row, we took the results of VisNet-LSRM from the MACE paper [1]. In the VisNet-LSRM (Paper) row, we took the results from the [ICLR submission of VisNet-LSRM](https://openreview.net/forum?id=rvDQtdMnOl) [2] (the MD22 results are the same their [latest arXiv submission](https://arxiv.org/abs/2304.13542)). Our conversion equations are as follows (e.g. in Tetrapeptide / Ac-Ala3-NHMe): Energy MAE (mEV/atom) = Energy MAE (kcal/mol) / (mEV to kcal/mol) / (Number of Atoms) = 0.0654 / 0.02306052 / 42 = 0.06752418667 Force MAE (mEV/Å) = Force MAE (kcal/mol/Å) / (mEV to kcal/mol) = 0.0902 / 0.02306052 = 3.9114469231 We note that our calculation of Tetrasaccharide matched with the results from Reviewer wbCz. Please let us know if that clarifies your question. --- [1] Kovács, D. P., Batatia, I., Arany, E. S., & Csányi, G. (2023). Evaluation of the MACE force field architecture: From medicinal chemistry to materials science. The Journal of Chemical Physics, 159(4). [2] Li, Y., Wang, Y., Huang, L., Yang, H., Wei, X., Zhang, J., ... & Liu, T. Y. Long-Short-Range Message-Passing: A Physics-Informed Framework to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation. In The Twelfth International Conference on Learning Representations.

Reviewer VH7W2024-08-13

Reply to rebuttal

I thank the authors for the rebuttal. Given their satisfactory reply to my concern and, in my opinion, the other reviewers’ ones, as well as the new results provided, I raise my score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC