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

Authors: eric  qu
aditi s krishnapriyan
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:50:57

Abstract: Scaling has been a critical factor in improving model performance and generalization across various fields of 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 scaling other types of machine learning models, the study of scaling in Neural Network Interatomic Potentials (NNIPs) remains limited. NNIPs act as surrogate models for ab initio quantum mechanical calculations, predicting the energy and forces between atoms in molecules and materials based on atomic configurations. The dominant paradigm in this field is to incorporate numerous physical domain constraints into the model, such as symmetry constraints like rotational equivariance. We contend that these increasingly complex domain constraints inhibit the scaling ability of NNIPs, and such strategies are likely to cause model performance to plateau in the long run. In this work, we take an alternative approach and start by systematically studying NNIP scaling properties and strategies. Our findings indicate that scaling the model through attention mechanisms is both 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 novel 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 speed up in inference time, 5x less in memory usage---compared to existing NNIP models. 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 towards developing general-purpose NNIPs that achieve better expressivity through scaling, and continue to scale efficiently with increased computational resources and training data.

Main paper: 
 introduction: 
 In recent years, the principle of scaling model size, data, and compute has become a key factor for improving performance and generalization in machine learning (ML), across fields from natural language processing (NLP) [Kaplan et al., 2020] to computer vision (CV) [Dosovitskiy et al., 2021 , Zhai et al., 2022] . Scaling in ML is, in a large part, defined by the ability to best exploit GPU computing capabilities. This typically involves efficiently increasing model sizes to large parameter counts, as well as optimizing model training and inference to be optimally compute-efficient.Parallel to these developments, ML models have also been rapidly developing for atomistic simulation, addressing problems in drug design, catalysis, materials, and more [Deringer et al., 2019 , Unke et al., 2021] . Among these, machine learning interatomic potentials, and particularly neural network interatomic potentials (NNIPs), have gained popularity as surrogate models for computationally intensive ab initio quantum mechanical calculations like density functional theory. NNIPs are designed to predict the energies and forces of molecular systems with high efficiency and accuracy, allowing downstream tasks such as geometry relaxations or molecular dynamics to be carried out on systems that would be intractable to simulate directly with density functional theory.Current NNIPs are predominantly based on graph neural networks (GNNs). The atomistic system is represented as a graph, where nodes correspond to atoms and edges representing interactions between atoms. Many effective models in this field have increasingly tried to embed physically-inspired constraints into the model, often justified by the belief that these constraints improve accuracy and data efficiency. Common constraints include incorporating predefined symmetries into the NN architecture, such as rotational equivariance, as well as using complex input feature sets.NNIP models that integrate symmetry constraints [Batzner et al., 2022 , Batatia et al., 2022 , Liao et al., 2024] often rely on computationally intensive tensor products of rotation order L [Geiger and Smidt, 2022] to maintain rotational equivariance. Although recent advancements have reduced the computational complexity of these operations [Passaro and Zitnick, 2023a, Luo et al., 2024] , the remaining computational overhead still significantly limits their scalability. Other approaches [Gasteiger et al., 2020a [Gasteiger et al., , 2021 [Gasteiger et al., , 2022] ] use basis expansions of the edge directions, angles, and dihedrals as features. Generally, incorporating these constraints tends to be compute-inefficient. As a result, many of the models in the field remain highly-constrained and small, despite the availability of larger datasets [Chanussot et al., 2021 , Jain et al., 2013] and more computational resources. We contend that these increasingly complex domain constraints inhibit the scaling ability of NNIPs, and such strategies are likely to plateau over time in terms of model performance. As the scale of the models increase, we hypothesize that imposing these constraints hinders the learning of effective representations, restricts the model's ability to generalize, and impedes efficient optimization. Many of these feature-engineered approaches are not optimized for efficient parallelization on GPUs, further limiting their scalability and efficiency, especially when applied to larger systems.In many other fields of ML, general-purpose architectures that best exploit computing capabilities outperform models with handcrafted, domain-specific constraints [Dosovitskiy et al., 2021 , Zhai et al., 2022] . These observations motivate us to ask: How can we develop principled methods and design choices that enable the creation of general-purpose neural network interatomic potentials that scale effectively with increased computational resources and training data?To answer this question, we conduct an initial ablation study to identify which components in NNIPs are most conducive to scaling. In NNIPs with built-in rotational equivariance, it is commonly believed that increasing the rotation order (L) improves model performance, even though it incurs additional computational cost. However, our investigations show that increasing the rotation order also adds more parameters to the model, and NNIPs are not always adjusted to account for this difference in parameter count. Our investigations also show that how parameters are added to the model is critical, as different types of parameter increases can differently impact the model's expressivity. We find that increasing the parameters of other components of the model besides the rotation order--particularly those involved in attention mechanisms--greatly improves model performance.Based on these insights, we develop the Efficiently Scaled Attention Interatomic Potential (EScAIP), an NNIP architecture explicitly designed for scaling by incorporating highly optimized attention mechanisms. To the best of our knowledge, our model is the first to leverage attention mechanisms on the neighbor representations of atoms rather than only the nodes, resulting in more expressivity. We also leverage advancements in attention mechanisms [Lefaudeux et al., 2022] , which have computational and memory efficiencies for scaling on large datasets.Our model achieves the best performance on a wide range of chemical applications, including the top performance on the Open Catalyst 2020 (OC20), Open Catalyst 2022 (OC22), SPICE molecules, and Materials Project (MPTrj) datasets. It also demonstrates a 10x speed up in inference time and 5x less in memory usage compared to existing NNIP models. To analyze rotational equivariance, on the validation set, we 1) predict forces on a set of atomistic systems (A), 2) rotate the atomistic systems and predict forces (B), and then 3) compute the cosine similarity between Inference Memory (GB/Sample)  
related works: 
 Neural Network Interatomic Potentials. There have been significant advancements in the development of neural network interatomic potentials (NNIPs), and we give a very general overview of the field. These models are usually trained to predict the system energy and per-atom force based on system properties, including atomic numbers and positions. We classify these models into two categories: (1) models that are based on Group Representation node features, and (2) models that are based on node features represented by Cartesian Coordinates. In the former, the node features are equivariant to different groups acting on the atomic positions, such as rotations and translations. In the latter, most architectures obey basic group symmetries, such as rotation and translation invariance.• Group Representation Architectures. The first model that used group representation node features was the Tensor Field Network [Thomas et al., 2018] , followed by an improved version, NequIP [Batzner et al., 2022] . Then, MACE [Batatia et al., 2022] incorporated the Atomic Cluster Expansion [Drautz, 2019] into the architecture. SCN [Zitnick et al., 2022] used spherical functions to represent equivariant node features, followed by an efficiency improvement in the tensor products, eSCN [Passaro and Zitnick, 2023b] . Equiformer [Liao and Smidt, 2022, Liao et al., 2024] incorporated graph attention into the architecture.• Cartesian Coordinates Architectures. SchNet [Schütt et al., 2017] was the initial work that only used edge distances as input to maintain invariant node features. DimeNet [Gasteiger et al., 2020a,b] and GemNet [Gasteiger et al., 2021 [Gasteiger et al., , 2022] ] added invariant bond direction feature sets as input. They designed an output head that maintains rotational equivariance with invariant node features. Another line of work tries to maintain equivariant features in Cartesian space by explicitly modeling spherical functions [Frank et al., 2022 , Bekkers et al., 2024 , Chen and Ong, 2022 , Cheng, 2024 , Haghighatlari et al., 2022] .There has also been a growing focus in the NNIP domain on generating larger datasets with quantum mechanical simulations, and using this to train models. These datasets span domains such as molecules [Eastman et al., 2023 , Smith et al., 2020] , catalysts [Chanussot et al., 2021 , Tran et al., 2023] , and materials [Barroso-Luque et al., 2024 , Yang et al., 2024 , Merchant et al., 2023 , Jain et al., 2013 , Choudhary et al., 2020] . [Liao et al., 2024] on the OC20 2M dataset. Energy (eV) and force (eV/Å) mean absolute error (MAE) are reported, along with the model's parameter counts. The leftmost column shows the original results from [Liao et al., 2024] , where different L had a different number of trainable parameters. We look at scaling parameters through the attention mechanisms (AT) and spherical channels (SC) for the original L = 2 and L = 4 models, such that the number of parameters is approximately equal to the original L = 6 model. Scaling parameters in different ways affects the overall energy and forces error, and increasing attention parameters is particularly effective in improving model performance (More AT). We also modify the architecture to be invariant (L = 0), allowing us to examine the effects of excluding rotational equivariance while controlling for the number of parameters (BOO). After controlling for parameter counts, many of the models have comparable error to the original L = 6 model.There has been a trend of incorporating physicallyinspired constraints into NNIP model architectures, such as all Group Representation Architectures that incorporate symmetry constraints into the model. However, there have been other lines of work that do not try to build in symmetry directly into the NN, and instead either try to "approximate" the symmetry [Pozdnyakov and Ceriotti, 2023 , Wang et al., 2022 , Finzi et al., 2021] or learn the symmetry via data augmentation techniques [Puny et al., 2022 , Duval et al., 2023] . 
investigation on how to scale neural network interatomic potentials: 
 We systematically investigate strategies for scaling neural network interatomic potential (NNIP) models through an ablation study. We examine how higher-order symmetries (rotation order L) impact scaling efficiency and identify the most effective methods for increasing model parameters ( §3.1). We also assess the importance of incorporating directional bond features ( §3.2). We conduct experiments using a leading NNIP architecture, the EquiformerV2 model [Liao et al., 2024] , on the Open Catalyst 2020 (OC20) Dataset [Chanussot et al., 2021] 2M split to evaluate the performance of different scaling strategies.A prevalent approach to improve the capability of NNIP models with group representation features is to increase the order of representations (L). Liao et al. [2024] did a study on the EquiformerV2 model, varying L to examine its impact on model performance. However, they did not control for the total number of trainable parameters in the model. This variation introduces discrepancies that can confound the true effect of L on the model's performance.Ablation Study Settings. To clarify the impact of increasing L on model performance and determine the most effective strategy for increasing parameters in NNIP models, we conduct a parametercontrolled experiment using the EquiformerV2 model on the OC20 S2EF 2M dataset. We standardize the number of trainable parameters across different values of L to isolate the effects of increasing L, and systematically add parameters to different components of the original L = 2 and L = 4 EquiformerV2 models from Liao et al. [2024] . Our approach targets four distinct configurations: increasing parameters solely in the attention mechanisms (AT), solely in the spherical channels that act on all group representations in the NN (SC), evenly across both attention mechanisms and spherical channels (AT + SC), and a configuration where spherical channels are reduced while significantly boosting attention parameters (More AT).Results of Ablation Study. The comparative analysis reveals a clear hierarchy in performance gains with different parameter scaling strategies. The More AT configuration yields the highest performance improvement, followed by AT, AT + SC, and SC. The results, summarized in Fig. 2 , show that once the number of parameters across models are controlled, many of the models have comparable error to the original L = 6 model. Increasing the parameters of the attention mechanisms is most beneficial and provides more substantial improvements than simply adding more parameters across all components.We explore what the most minimal representations are of the atomistic system to enable the model to learn a scalable, data-driven feature set, and find that incorporating bond directional information is useful for NNIP models. As opposed to other domains, such as social networks, the edges (or bonds) in molecular graphs possess distinct geometric attributes, i.e., pairwise directions. However, the raw value of the bond direction changes with the rotation and translation of the molecule, making it challenging to directly utilize these features in NNIP models.We propose a straightforward and data-driven approach to embed the bond directional information.To avoid the computational inefficiency of taking a tensor product, we aim to use the simplest possible representation of bond direction that is rotationally invariant. Inspired by Steinhardt et al. [1983] ,we use an embedding of the Bond-Orientational Order (BOO) to represent the directional features.Formally, for a node v, the BOO of order l is,EQUATIONwhere duv is the normalized bond direction vector between node v and u, n v is the number ofEQUATIONm is the spherical harmonics of order l and degree m. This can be interpreted as the minimum-order rotation-invariant representation of the l-th moment in a multipole expansion for the distribution of bond vectors ρ bond (n) across a unit sphere. In other words, BOO is the simplest way to encode the neighborhood directional information in a rotationally invariant manner. The BOO features BOO(v) ∈ R L+1 for a node v is the concatenation of BOO(v) (l) . In theory, the BOO feature contains all the directional information of the neighborhood, and the embedding network can learn to extract such information.Testing the bond-orientational order (BOO) features. We conduct a study to test the BOO features. We modify the EquiformerV2 model to be L = 0 and replace the spherical harmonics directional features with embeddings of the BOO features. The results are in Fig. 2 . The L = 0 model achieves comparable results with the L = 6 model. This finding suggests that the BOO features are a straightforward and effective way to incorporate bond directional information in NNIP models, and that it is also possible to learn additional information solely through scaling. 
efficiently scaled attention interatomic potential (escaip): 
 We introduce a new NNIP architecture, Efficiently Scaled Attention Interatomic Potential (EScAIP), which leverages highly optimized self-attention mechanisms for expressivity, with design choices centered around scalability and efficiency. To avoid costly tensor products, we operate on scalar features that are invariant to rotations and translations. This enables us to take advantage the optimized self-attention mechanisms from natural language processing, making the model substantially more time and memory efficient than equivariant group representation models such as EquiformerV2 [Liao et al., 2024 ]. An illustration of our model is shown in Fig. 3 . We describe the key components of the model and the motivation behind their design: Input Block. The input to the model is a radius-r graph representation of the molecular system. We use three attributes from the molecular graph as input: atomic numbers [Zitnick et al., 2022] , Radial Basis Expansion (RBF) of pairwise distances [Schütt et al., 2017] , and Bond Order Orientation (BOO) features from §3.2. The atomic numbers embeddings are used to encode the atom type information, while the RBF and BOO embeddings are used to encode the spatial information of the molecular system. These input attributes are the minimal representations of the system, enabling the model to learn a scalable, data-driven feature set. We also note that the attributes can be pre-computed, requiring minimal computational cost. The input features are then passed through a feed forward neural network (FFN) to produce the initial edge and node features.(b) (c) (d) (e) (f) (f)Graph Attention Block. The core component of the model is the graph attention block, illustrated in Fig. 4 . It takes node features and molecular graph attributes as input. All the features are projected and concatenated into a large message tensor of shape (N, M, H), where N is the number of nodes, M is the max number of neighbor, and H is the message size. The message tensor is then processed by a multi-head self-attention mechanism. The attention is parallelized over each neighborhood, where M is the sequence length. By using customized Trition kernels [Tillet et al., 2019 , Lefaudeux et al., 2022] , the attention mechanism is highly optimized for GPU acceleration. The output of the attention mechanism is aggregated back to the atom level. The aggregated messages are then passed through the node-wise Feed Forward Network (FFN) to produce the output node features. To the best of our knowledge, this attention mechanism is unique because it acts on a neighborhood level, which is more expressive than the graph attention architectures that only act on the node level.Readout Block. We use two readout layers for each graph attention block, which follows GemNet-OC [Gasteiger et al., 2022] . The first readout layer takes in the unaggregated messages from the graph attention block and produces edge readout features. The second readout layer takes in the output node features from the node-wise FFN and produces node readout features. The node and edge readout features from all graph attention blocks are concatenated and passed into the output block for output prediction.Output Block. The output block takes the concatenated readout features and predicts the per-atom forces and system energy. The energy prediction is done by an FFN on the node readout features. The force prediction is divided into two parts: the force magnitude is predicted by an FFN on the node readout features, and the force direction is predicted by a transformation of the unit edge directions with an FFN on the edge readout features. As opposed to GemNet [Gasteiger et al., 2022] , the transformation is not scalar but vector-valued. Thus, the predicted force direction is not equivariant to rotations of the input data. In our experiments, we found this symmetry-breaking output block made the model perform better. The reason could be that this formulation has more degrees of freedom and so is easier to optimize. We note that though the force direction is initially not equivariant, the trained model is able to learn this symmetry from the data (See §5.4).We also note that predicting the force magnitude from node readout features is very helpful for energy prediction. The reason could be that the energy prediction is a global property of the molecular system, while the force magnitude is a local property of the atom. By guiding the model towards a fine-grained force magnitude prediction, the model can learn a better representation of the system, which can in turn help it predict the system energy more accurately. 
experiments: 
 We conduct experiments on a wide range of chemical systems, including catalysts (OC20 and OC22) §5.1, materials (MPTrj) §5.2, and molecules (SPICE and MD22) §5.3, §B.2.Dataset. We evaluate the performance of our EScAIP model on the Open Catalyst dataset [Chanussot et al., 2021 , Tran et al., 2023] , which consists of 172 million systems with 73 atoms on average. All reported results are tested on NVIDIA V100 32G.We evaluate on the S2EF task, which is the prediction of system energy and per-atomic force from atomistic structure.Settings. We use three variants of the EScAIP model: Small (83M), Medium (146M), and Large (317M). The models are trained to predict the energy and forces of each sample (S2EF task). We train the model on the OC20 All+MD, OC20 2M, and OC22 splits. We evaluate the performance on the four validation sets and test sets (both have 4M samples in total) and compare the results with EquiformerV2 [Liao et al., 2024] , eSCN [Passaro and Zitnick, 2023b] , SCN [Zitnick et al., 2022] , and GemNet-OC [Gasteiger et al., 2022] , the best performing models on this dataset.Results. The results of EScAIP on the Open Catalyst dataset are summarized in Tab. 1, where EScAIP achieves state-of-the-art performance across all splits: OC20 2M, OC20 All+MD, and OC22.We note that we exclude models that train with a denoising objective, as we focus on the performance of the model architecture itself. There is a clear trend that increasing the model size improves the performance of EScAIP. Notably, even the Small model achieves competitive performance against other models while remaining significantly more efficient, making it suitable for downstream, practical applications. More results on the scalability of the EScAIP model can be found in the Appendix B.1.Efficiency Comparisons. We provide the runtime and memory usage of EScAIP and other baseline models on the OC20 dataset in Tab. 2. EScAIP is approximately 10x faster and has 5x less memory usage than an EquiformerV2 [Liao et al., 2024] model of comparable size. Additionally, as shown in Fig. 5 , EScAIP scales effectively with batch size while maintaining minimal performance variation across batches. This consistency is because EScAIP's input is padded to the maximum system size, enabling efficient use of PyTorch's compile feature. These qualities could make EScAIP well-suited for practical applications.Dataset. We evaluate EScAIP's performance on the Matbench-Discovery benchmark [Riebesell et al., 2023] , a widely recognized benchmark for assessing models in new materials discovery. The model is trained on the MPTrj dataset [Deng et al., 2023] , which consists of 1.6 million samples. This approach adheres to the "compliant" setting of the Matbench-Discovery benchmark.Settings. Given the relatively small dataset size, we use a small version of EScAIP with 45M parameters. The model is trained to predict the energy, force, and stress of each sample. After training for 100 training epochs, we increase the energy coefficient in the loss function and fine-tune the model for another 50 epochs. We evaluate the performance on the Matbench-Discovery benchmark and compare the results with EquiformerV2 [Liao et al., 2024 , Barroso-Luque et al., 2024] , ORB MPTrj [Orbital-Materials, 2024], SevenNet [Park et al., 2024] , and MACE [Batatia et al., 2023 [Batatia et al., , 2022 ]--the top compliant models on this benchmark. We note that we exclude models that train with a denoising objective, as we focus on the performance of the model architecture itself1 .Results. The results of EScAIP on the Matbench-Discovery benchmark are summarized in Tab. 3. EScAIP achieves state-of-the-art performance on this benchmark, outperforming other models. In the code release, we also provide the EScAIP model before the energy fine-tuning step, as this may be more relevant for applications such as molecular dynamics simulations.Dataset. We evaluate the EScAIP model's performance on the SPICE dataset [Eastman et al., 2023] Results. A summary of EScAIP's results on the SPICE dataset is provided in Tab. 4, where it outperforms MACE-OFF23 in predicting the energy and forces on the different test sets.Settings. To assess whether EScAIP learns rotational equivariance after training on various datasets, we design the following procedure: first, a batch is randomly sampled from the validation dataset and passed through the trained model to obtain a force prediction (A). Next, we rotate the batch by a random angle and obtain a second force prediction (B) from the model. To quantify rotational ) . This process is repeated for 128 batches, and we report the average cosine similarity.Results. The results of the rotational equivariance analysis are presented in Tab. 5, and the cosine similarity is consistently ≥ 0.99. These findings indicate that though EScAIP is not initially rotationally equivariant, after training it is able to correctly map input system rotations to output system predictions. This could suggest that these symmetries can be effectively learned from the data.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: All datasets are available on Github. We also provide all code and model checkpoints.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusions: 
 We have investigated scaling strategies for developing neural network interatomic potentials (NNIPs) on large-scale datasets. Based on our investigations, we introduced a new NNIP architecture, Efficiently Scaled Attention Interatomic Potential (EScAIP), that leverages highly optimized selfattention mechanisms for scalability and expressivity. We demonstrated the effectiveness of EScAIP on a wide range of chemical datasets (OC20, OC22, MPTrj, SPICE) and showed that EScAIP achieves top performance on these different prediction tasks, while being much more efficient in training and inference runtime, as well as memory usage. We highlight some important takeaways from our work and the future of machine learning interatomic potentials more broadly:The "sweet lesson." We note that our line of investigation in this work follows some of the general principles of the bitter lesson [Sutton, 2019] . That is, strategies that focus on scaling and compute tend to outperform those that try to embed domain knowledge into models. However, in this field, we prefer to think of this as a "sweet lesson." Training large, constrained models requires significantly more computational resources, making this feasible for only a limited number of researchers. Efficient scaling strategies thus democratize large-scale training and make it accessible to a broader community.We're still not giving enough credit to the data. Thus far, much of the effort in the NNIP field has concentrated on model development. However, atomistic systems are far more complex than the domain-specific information being embedded into models. Predefined symmetry constraints and handcrafted features offer only a simplistic representation of this complexity. A path forward to capture these complexities is to focus on generating comprehensive datasets, ideally accompanied by relevant evaluation metrics, allowing NNs to learn the rest of the information through gaining expressivity via scaling.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Yes, we provide this in the experiments section. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. Justification: Yes, this is also discussed in the introduction. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This is not relevant to this work. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Yes, we cite the relevant datasets we used.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes] Justification: Yes, we provide a new model and include implementation details.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: We don't have this.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: As datasets continue to grow, training models from scratch on small datasets will likely become unnecessary. While constraints may be beneficial in the very small data regime (though data augmentation techniques can also help here), leveraging the representation of a pre-trained large model can serve as a starting point for fine-tuning on smaller datasets. This could make the very small dataset regime essentially a non-factor in the future, and it is likely that the need for NNIPs with built-in hard-constraints becomes even less necessary. Beyond focusing on data generation, other techniques are likely to gain importance in the NNIP domain. These include model distillation, general training and inference strategies that are model agnostic and can be applied to any NNIP, and approaches to better connect with experimental results. Finally, more comprehensive strategies will be important for evaluating NNIP accuracy and utility. We provide additional details on our investigations from §3.To further investigate how scaling efficiency varies as a function of training dataset size, we train the parametercontrolled Equiformer V2 models with different amounts of training data. The results in Fig. 6 show that the scaled L = 2 and L = 4 models exhibit a steeper performance improvement (log-log slope) compared to the original L = 6 model. In particular, the More AT configuration (more attention) has a steeper log-log slope compared to the AT + SC configuration and the original L = 6 model. This suggests that increasing the complexity of the attention mechanisms is a more effective strategy for scaling with increasing training dataset size, rather than increasing L.B.1 Catalysts (OC20 and OC22)To illustrate EScAIP's scalability, we train the model on varying sizes of training data and model configurations. The results, shown in Fig. 8 , indicate a clear trend: as model and data sizes grow, EScAIP's performance continues to improve. We also include results to complement Fig. 1 : the same efficiency, performance, and scaling comparisons between EScAIP and baseline models on the Open Catalyst dataset for Energy MAE (meV ↓) vs. Inference Speed (Sample/Sec ↑) and Energy MAE vs. Memory (GB/Sample ↓). The trend is similar to the forces MAE results, and EScAIP achieves better performance with smaller time and memory cost.Dataset. We evaluate the performance of our EScAIP model on the MD22 dataset [Chmiela et al., 2023] , which consists of seven molecular systems with varying sizes. It consists of energy and force labels calculated from DFT simulations.Settings. We use an EScAIP model 15M parameters on each system and evaluate the performance on the test set (train-test split 95:5). The model is compared with MACE [Batatia et al., 2022] , VisNet-LSRM [Li et al., 2024] , and sGDML [Chmiela et al., 2023] . We note that there were discrepancy of results of VisNet-LSRM in the MACE [Kovács et al., 2023b] paper and the VisNet-LSRM [Li et al., 2024] , thus we reported both in the table. We use the same train/validation splits as the baselines [Chmiela et al., 2023] , where the training set is hundreds to thousands of samples, and also apply data augmentation (randomly rotating each training sample 16 times). Inference Memory (GB/Sample) After we train the model, we also run molecular dynamics (MD) simulations to check the stability of the potential and evaluate how well the simulation recovers the distribution of interatomic distances, h(r), in the simulation [Raja et al., 2024 , Fu et al., 2023] . We use the simulation setup from Fu et al.[2023] and run the simulation for 200000 steps (100 ps) using the Langevin integrator with a friction coefficient of 0.5. The temperature is set to 500 K.Results. The results of EScAIP on the MD22 dataset are summarized in Tab. 6. EScAIP outperforms other models in both energy and force prediction, especially for large molecules. The low h(r) error in the MD simulation also indicates that the model is able to capture this observable accurately. Interestingly, MD22 is not a particularly large dataset: the training dataset sizes are in the thousands. Despite this, a scalable architecture with high parameter counts is still able to achieve good performance.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Based on the ORB technical report orb[2024], it is possible that the ORB MPTrj model result reported here was pre-trained with a denoising objective. 


Decision: This paper proposes EGAP, an efficient graph attention-based model for learning interatomic potentials. The reviewers generally agree that the paper tackles an important problem (PKom: "The problem of efficiently predicting complex molecule properties is important in practice.") and presents a well-written and easy-to-follow method (PKom: "The paper is easy to follow and written and presented well.").

In the initial review and during the subsequent rebuttal, a few reviews raised concern regarding the experiment validation:
- Inconsistent dataset splits and usage: wbCz pointed out significant inconsistencies in the MD22 dataset splits, making comparisons unfair ("By utilizing a 95:5 train-test split, the authors have made their results incomparable to baseline models."). Similarly, PKom noted unfair comparisons due to different dataset sizes used for EGAP and baselines in OC20. 
- Discrepancy with previous published results: Furthermore, wbCz found discrepancies between reported results and those in cited papers, raising concerns about accuracy and fairness.
- Limited experimental scope: wbCz and VH7W suggested evaluating on larger datasets like OC20 All+MD to better demonstrate scalability and efficiency.
- Lack of scalability experiments: L3rz and wbCz highlighted the absence of experiments demonstrating performance scaling with model size, questioning the scalability claims.

The authors addressed most of these concerns in their rebuttal by:
- Revising the MD22 experiments to use consistent dataset splits.
- Providing results on OC20 All+MD and OC22, demonstrating state-of-the-art performance and improved efficiency.
- Conducting scaling experiments with varying model sizes and dataset sizes.
- Clarifying and correcting the reported results.

Despite the initial weaknesses, the authors' comprehensive rebuttal and the updated results significantly strengthen the paper. The reviewers generally acknowledged the improvements and raised their scores accordingly (PKom, VH7W, L3rz, wbCz), leading to an concensus decision of acceptance.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: see above
limitations: limitations: not discussed explicitly
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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)
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: limitations: See above.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: 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.

Comments: title: Response to reviewer VH7W
Reply to rebuttal
Response to review PKom
Rebuttal by Authors
comment: Thank you for your response!
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.
Thank you for your response.
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.
Can you double check? The results are different from their latest version.
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.
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.
The response makes some important clarifications and adds useful experiments, therefore I raise the score.
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.