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 2024 poster (2024)
Keywords: Neural Network Interatomic Potentials
Machine Learning Force Fields
Scaling
Graph Neural Networks
Attention
Primary area: machine_learning_for_physical_sciences
TL;DR: We develop a neural network interatomic potential architecture that is optimized for scalability and efficiency, achieving state-of-the-art results on a wide range of chemical systems including OC20, OC22, MPTrj, and SPICE.

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).
 After training EScAIP, we test its ability to learn rotational equivariance by predicting forces on new, unseen atomistic systems before and after rotation. The model's force predictions exactly match the rotated forces,  suggesting that it has precisely learned rotational equivariance.
Finally, 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.

The Importance of Being Scalable:
Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains
Eric Qu UC Berkeley ericqu@berkeley.co edu
Aditi S. Krishnapriyan UC Berkeley, LBNL aditik1@berkeley edu
Abstract
Scaling has been a critical factor improving model performance and generalization across various fields of machine learning. It involves how model's performance changes with increases in model size or input data, as well as how efficiently computa- tional 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. NNTPs act surrogate models for ab initio quantum mechanical calculations, predicting the energy and forces between atoms molecules and materials based 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 in- creasingly 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 scal- ing properties and strategies. Our findings indicate that scaling model through at- tention 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 novel multi- head self-attention formulation within graph neural networks, applying attention at the neighbor-level representations. Implemented with highly-optimized attention GPU kermels, ESc achieves substantial gains in efficiency-------eeeeee least 10x speed up in inference time, 5x less in memory usage- usage-compare compared to existing NNIP models. ES- cAIP also achieves state-of- offte--------trt--- performance on a wide range of datasets including catalysts (OC20 and OC22), molecules (SPICE), and materials (MPTrj). After train- ing EScAIP, we test its ability to learn rotational equivariance by predicting forces on new, unseen atomistic systems before and after rotation. The model's force predic- tions exactly match the rotated forces, suggesting that it has precisely learned rota- tional equivariance. Finally, emphasize that our approach should be thought of a philosophy rather than a specific model, representing a veloping general-purpos. NNIPs that achieve better expressivity through scaling, and continue to scale efficiently with increased computational resources and training data.
towards de-
1 Introduction
In recent years, the principle of scaling model size, data, and compute has become key factor for improving performance and generalization
machine learning (MI), across fields from natural language processing (NLP) [Kaplan et al., 2020) to computer vision (CV) [Dosovitskiy et 2021 Zhai et al. [022]. Scaling in ML. is, in 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 optimally compute-efficient
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

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 er 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 NNIPss are predominantly based on graph neural networks (GNNs). The atomistic system is represented as a graph, where nodes correspond atoms and edges representing interactions between atoms. Many effective models in this field have increasingly tried to embed physically ihysically-inspired 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 rotation order [Geiger and Smidi [0022] 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 Gasttiggrr al., 2020a 2021 2022] use basis expansions of the edge directions. angles, and dihedrals as features. Generally, incorporating these constraints tends to be compute-ingg -nntute-inffficcent. As result, many of the models in the field remain highly-constrained and small, despite the availability of larger Jain et al. 2013) and more computational resources.
datasets [Chanussot et al.) 2021
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 s ability to generalize, and impedes efficient optimization. Many of these nntture-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 MI., seeneal--urposs architectures that best exploit computing capabilities outperform models with handcrafted, domain-specifi constraints (Dosovitskly et al.. 2021 et al.) develop principled methods and design
2022]. These observations motivate to ask: How choices that enable the creation of general-pun -arrrp-e 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 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- particularyyy those involved in attention mechanisms- -greally 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 al. 2022 which have computational and memory efficiencies for scaling on large datasets.
Our model achieves the best performance 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 evaluate how well EScAIP has learned rotational equivariance, held-out validation set not seen during training, we 1) predict forces set of atomistic systems (A), 2) rotate the atomistic systems and predict forces (B), and then 3) compute
wide range of chemical applications. including the

---

20
216M
EquilormerV2
20
216M
eSCN
EquiformerV2
Force MAE (meV/A)
GemNNet-OC EScAJP
Force MAE (meViA)
GemNet-OC
18
18
EScAIP
200M
31M
31M
16
83M
16
153M
153M
14
14
317M
146M
146M
31
12
60
80
100
120
2
7
Inference Speed (Sample/Sec)
Inference Memory (CB/Sample)
Figure 1: Efficiency, performance, and scaling comparisons between EScAIP and baseline models on the Open Catalyst dataset (OC20). Force MAE (me 4) vs. Inference Speed (Sample/Sec f) and Force MAE vs. Memory (GB/Sample L) is reported. Results with Energy MAE can be found in the Appendix Fig. 7 EScAIP achieves better performance with smaller time and memory cost.
the cosine similarity between the force predictions (B) and the rotated version of force predictions (A). After training EScAIP on different datasets, we find that EScAIP accurately predicts the forces on the rotated systems, as indicated by high cosine similarity scores (2 0.999). This suggests that has successfully learned and captured rotational equivariance. We also provide evidence that EScAIP scales well with compute, and is designed in such a way that it will further improve in efficiency as advances in GPU computing continue to increase. Our code and model checkpoints are publicly available at wtt.ttts://// http:::///tthutt com A
Related Works
Neural Network Interatomic Potentials. There have been significant advancements in the devel- opment of neural network interatomic potentials (NNIPs), and give a very general overview of the current state 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 the current class of 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 fea- tures 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, NequlP [Baznner et al.) 2022). Then, MACE [Batatia et al. [2022 incorporated the Atomic Cluster Expansion Drautz]] 2019 into the architecture. SCN ZZinnikk et al. 2002 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 al. 2017 is an example of 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 2022 added invariant bond direction feature sets as input. They designed an output head that maintains rotational equivariance with invariant node features. TorchMD-No Net added attention mechanisms over the atomic graph [Thölke and Fabritiis
2022 Pelaez et al.) [024]. Another line of work tries to maintain equivariant features in Cartesian
space by explicitly modeling spherical functions [Frank al.
2022 Bekkers et al. 2024 Chen and
Ong 2022
Cheng 2024 Haghighatlari et al. 2022
Liu et al.
2022
Datasets for NNIP training. 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 [Bastman et al. 2023 Smith et al. 2020 Anstine et al. 2024
catalysts [Chanussot et al. [002] Tran et al. 2023 and materials (Barroso-Luquu al. 2024
Tran
Yang
al. 2024 Merchant et al.
2023 Jain et al.
2013
Choudhary et al. 2020].

---

28
305
Force MAE (meV/A)
26.24
19M
AT: Attention
SC: Spherical Channels
L-0
--4
Energy MAB (meV)
300
19M
298
AT: Attention
SC: Spherical Channels
L-0
L-4
L=2
L=6
295
24
59
83M
22
21.37
283
20
20. 46
05:M
20.85
21 63
20.13
8IM
05M
20.54
280
284
83M
280
061
278
279
275
275
18
03M
85M
275
270
B8M
277
84M
275
84M
85M
276
85M
95M
275
et al., AT
Invariant B00
al., AT
More AT
lnvariant
2024
SC
AT
2024
SC
B00
Figure 2: Results of ablation study of EquiformerV2 [Lao et al. 2024) the OC20 2M dataset. Energy (eV) and force (eV/A) 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 I. 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= and L= models, such that the number of parameters is approximately equal to the original L= 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 (Invariant BOO). After controlling for parameter counts, many the models have comparable error to the original L 6 model.
Constrained vs. Unconstrained Architectures. There has been a trend of incorporating physically-inspired constraints into NNIP model architectures, such all Group Representation Architectures that incorporate symmetry constraints into the model. However, there have been other lines of work that do not try build in symmetry directly into the NN, and instead either try to "approximate" the symmetry (Pozdnyakov and Ceriotti, [023 Wang et al. [2022] Finzi et al. 2021
or learn the symmetry via data augmentation techniques Puny al. 2022
Duval et al. 2023
3 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 EguiformerV2 model [Liao et al. 2024], on the Open Catalyst 2020 (OC20) Dataset [Chanussot et al. [021] 2M split evaluate the performance of different scaling strategies.
3.1 Optimal Components for Scaling Neural Network Interatomic Potentials
A prevalent approach to improve the capability of NNIP models with group representation features is to increase the order of representations (T.). Liao et al. 2024) did a study on the EquiformerV? model, varying L 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 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 parameter-controlle 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 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 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 sealing strategies. The More A7' 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.
3.2 Bond Directional Features
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 useful for NNIP models. As opposed other domains, such 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 straightforward and data-driven approach to embed the bond directional information. To avoid the computational inefficiency of taking tensor product, we aim to use the simplest possible representation of bond direction that is rotationally invariant. Inspired by Sttinnarrtt et al. [1883]. we use an embedding of the -oin-rrrenttiional Order (BOO) to represent the directional features. Formally, for node v, the BOO of order / is,
(O)
Σ m==
4m
Y (dum)
21+1
weNei(e)
(000v))Connce (IBO0
where duy is the normalized bond direction vector between node e and u. ny is the number of neighbors vf Nei(w) is the neighbors of v. y(!) is the spherical harmonics of order I and degree This can be interpreted as the minimum-order rotation-invarian representation of the I-th moment in a multipole expansion for the distribution of bond vectors Pboond (n) across 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) E RR 4+11 for node the concatenation of BOO(e) (1), 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 oond-onientational order (BOO) features. We conduct a study to test the BOO features. We modify the EquiformerV2 model to be 0 and replace the spherical harmonics directional features with embeddings of the BOO features. The results are in Fig. 2l The == 0 model achieves comparable results with the L= 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 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 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 EquiformerVV V2 (Liao et al., 2024]. An illustration of our model is shown in Fig. We describe the key components of the model and the motivation behind their design:
Input Block. The input to the model is radius-r graph representation of the molecular system. We use three attributes from the molecular graph as input: atomic numbers [Zinick et al. 2022). Radial
5

---

ES:IP Architecture
(b) Graph Attention
(c) Output Block
Input Block
Input
Node Featurns
Gruph Altriluato
Node
Edge Rondeent
Grupă: Attributio
Resslant
XB
Liner
Edge Distance
Timeer
Linor
ldge Distance
RMS Norm
Alussic Number
Fend Forward
Linar
Alamic Number
Gmaph
IOO Peetures
Linear
8000
Attention
(b)
(
Prone
Asarapate
Festures
Readout Edge
Concat
Concat
Foael Furooudd
Forward
Foed
Linerr
Agrnegate
Linese
RMS Norm
Cunani
Fona Magrittrt
Porcc Dicectin
(c) Feed
(f) Readout
Feed Forward
Multi-head Self-Atteatio
Feed Forourd
Forward
&
Linner
Linerr
Readout Node
Coneat
Aggrapate
Pooling
RMMS
Node
Edge
Restem
Per-stom
Norm
Festur
Featurea
Enengy
Piete
Output Block
(c)
Figure 3: Illustration of the Efficiently Scaled Attention Interatomic Potential (EScAIP) model architecture. The model consists of graph attention blocks (dashed box), each of which contains graph attention layer, feed forward layer, and two readout layers for node and edge features. The concatenated readouts from each block are used to predict per-atom forces and system energy.
Basis Expansion (RBF) of pairwise distances [Schütt et al, 2017 and Bond Order Orientation (BOO) features from 3.2 The atomic numbers embeddings 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 feed forward neural network (FFN) to produce the initial edge and node features.
Graph Attention Block. The component of the model is the graph attention block, illustrated in Fig.[4] takes node features and molecular graph attributes as input. All the features are projected and concatenated into large message tensor of shape (N,M, (N,,HH) 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 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-OG [Gatteige et al. 2022). The first readout layer takes the unaggregated messages from the graph atten- tion block and produces edge readout features. The second readout layer takes in the output node fea- tures 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.
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 foree prediction is divided into two parts: the force magnitude predicted by an FFN on the node readout features, and the force direction is predicted by transformation of the unit edge directions with an FFN on the edge readout features. As opposed to GemNet (Gatteiger et al.) 2022), the transformation is not scalar but vector- valued. Thus, the predicted force direction is not equivariant rotations of the input data. our experiments, we found this symmetry -braakknn output block made the model perform better. The reason could be that this formulation has more degrees of freedom and
Output Block.

---

Mulfi-Head Self-Attention
Messago
Mcssage
Aggregation
Message Size
Hidden Size
Message Size
Masked Mean
Message Size
* of Nodes
Batch
* of Nodes
of
Size
Nodes
Message
Max # of Neighbors
Sequence
Max + of Neighbors
Length
Figure 4: Detailed illustration of the graph attention block. The input attributes are projected and concatenated into a large message tensor. The tensor is fed into optimized multi-head self-attention computation, where the max number of neighbors dimension is the sequence length dimension.
Table EScAIP performance on the OC20 AII+MD, OC20 2M, and OC22 datasets. The results are reported Energy (eV) and Force (eV/À) mean absolute error (MAE). EScAIP generally achieves the best Energy and Force MAE among all current models. Due to its efficiency, EScAIP requires less training time compared the other models.
Validation
Test
# Parameters
Energy MAE (meV)
Force MAE (meV/A)
Energy MAE (me
Force MAE
Dataset Model
(meV/A)
GemNet-OCC-LF
216M
252
19.99
241
19.01
OC20 AII+MD
cSCN K=20
200M
17.09
228
15.60
EquifomerVV
31M
232
16,26
228
15.50
EquifomerV2 (Ag 2)
153M
230
14.60
227
13.80
EquiformerV2 (Ag=4)
153M
227
15.04
219
14.20
EScAIP-Small
15.07
233
15.73
EScAIP-Medium
146M
12.82
221
13.19
EScAIP-I Large
317M
211
12.17
215
12.65
GemNet-dT
31M
358
29.50
GemNet-OC
38M
286
25.70
OC22 OC20 2M
SCN
126M
279
21.90
cSCN
20.50
EquifornerV2
85M
20.46
EScAIP-Small
83M
263
20.15
EScAIP-Medium
146M
254
19.08
GemNet-OC
39M
707
35,0
EquifornerV2
122M
531
26.79
462
27.1
EScAIP- Medium
24.32
473
25.73
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 global property of the molecular system, while the force magnitude is local property of the atom. By guiding the model towards a fine-grained foree magnitude prediction, the model can learn better representation of the system, which can in turn help predict the system energy more accurately.
5
Experiments
We conduct experiments on a wide range chemical systems, including catalysts (OC20 and OC22) materials (MPTr)) $5.2, and molecules (SPICE and MD22) $5.
5.1
5.1
Catalysts (OC20 and OC22)
Dataset. We evaluate the performance of our EScAIP model on the Open Catalyst dataset [Chanussot al.) 2021]] Tran al. 2023), which consists of 172 million systems with 73 atoms on average. We evaluate on the S2EF task, which is the prediction of system energy and per-atee per-atomic force from atomistic structure.

---

Table 2: EScAIP efficiency comparisons with bascline models on the OC20 dataset. All reported results are measured on NVIDIA V100 32G.
# of
Training Speed (Sample/See)
Training Memory
Inference Speed (Sample/Sec)
Inference Memory
Dataset
Model
Parameters
(GB/Sample) 4
(GB/Sample) !
GemNet-OC
216M
9.44
3,40
23.67
2.31
LS6 K=20
200M
2.16
6,90
6.3
6.15
EquiformerV2
31M
14.22
1.63
34.55
1.61
EquiformerV2
153M
2.85
6.33
7.92
5.24
EScAIP-Small
83M
35.84
1.23
107.04 04
1.09
EScAlP-Medium
146M
25.36
1.54
74.77
1.28
EScAIP-Large
312M
12.88
2.78
40.56
2.28
100
Atoms
1000 Atoms
100
Atoms
1000 Atoms
ESoAIP-S 93M
30
1.0
ESoAIP-L 312M
25
0.8
EquiformeV22 31M
Dquiformer22 153M
Memory (GB)
Time (5)
cSCN 200M
20
15
0.4
10
0.2
5
0.0
1
13
17
21
13
17
21
Batch Size
Batch Size
Figure 5: Inference runtime and memory usage comparison of EScAlP and baseline models on the OC20 dataset. Mean and standard deviation (shown as error bars) are reported across 16 randomly sampled batches per batch size. Grey lines indicate the cumulative number of atoms in the batch. EScAIP not only scales efficiently with batch size, but also exhibits minimal variation in performance across different batches. All reported results are tested on NVIDIA V100 32G.
Settings. We use three variants of the EScAIP model: Small (83M), Medium (146M), and Large (317M). The models are trained 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 EquiformerV V2 [Lao et al. 2024 eSCN [Passaro and Zitmick 2023b), SCN [Zitnick et al. 2022 and GemNet- OC Gasteiger et 2022 the best performing models on this dataset.
Results. The results of EScAIP on the Open Catalyst dataset are summarized in Tab. 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 ES.AIP. Notably, even the Small model achieves competitive performance against other models while remaining significantly more efficient, making suitable for downstream, practical applications. More results on the scalability of the EScAIP model can be found the Appendix B..
Efficiency Comparisons. We provide the runtime and memory usage of EScAIP and other baseline models on the OC20 dataset in Tab EScAlP is approximately 10x faster and has 5x less memory usage than an EquiformerV [Liao et al. 2024] model comparable size. Additionally, as shown in Fig ES- cAIP scales effectivey with batch size while maintaining minimal performance variation across batches. This consistency because FSScIIP's input is padded the maximum system size, enabling efficient use of Torch's 's compile feature. These qualities could make EScAIP well-suited for practical applications.
5.2
Materials (MPTrj)
Dataset. We evaluate EScAIP's performance on the Matbench-Discoverry benchmark [Riebesell 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-Discowee Settings. Given relatively small dataset size, we use small version of EScAIP with 45M
al. 2023
benchmark.
parameters. The model is trained to predict the energy, force, and stress of each sample. After training

---

Table 3: EScAIP performance on the Matbench- Dissooerr benchmark. Mcan absolute error (MAE) and Root Mean Squared Error (RMSE) are reported in e/atom. We only include models trained without denoising objective, focusing on the performance of the model architecture itself.
Model
F
DAF
Precision
Rocall
Accuracy
TPR
FPR
TNR
FNR
MAE
RMSE
R21
MACE
0.669
3.777
0.577
0.795
0.878
0.796
0.107
0.893
0.204
101
0.697
SevenNet
0.724
4,252
0.65
0.818
0.904
0.818
0.081
0.919
0.182
48
92
0.75
ORB MPuj
0.765
4.702
0.922
0:941
0.756
EquiformerV2 V2
0.77
4,64
0.709
0.841
0.926
0.841
0.063
0.937
0.159
42
87
0.778
EScATP
0.781
4.734
0.724
0.848
0.927
0.847
0.058
0.941
0.152
84
0.792
Table EScAIP performance on the SPICE dataset. The results are reported in Energy (me V/ato)) and Force (me) mean absolute error (MAE).
DES370K Monomers
DES370K Dimers
Solvated Amino Acids
Model
Metric
PubChem
Dipeptides
Water
QMugs
Force MAE
14.75
6.58
6.62
10.19
19.43
13.57
16.93
MACE
Energy MAF
0.88
0.59
0.54
0.42
0.98
0.83
0.45
Force MAE
5.86
3,48
5.21
11.52
10.31
8.74
EScAIP
Energy MAE
0.53
0.41
0.31
0.61
0.41
for 100 training epochs, we increase the energy coefficient the loss function and fine-tune the model for another 50 epochs. We evaluate the performance the Mathench-Discovee
benchmark and
compare the results with EquiformerV2 [Liao et al. [2224] Barroso- Luque et al. [2024). ORB MPTri
Orrital-Materiall 2024], SevenNet [Park et al.) 20224, and MACE [Batatia 2023 2022] the lop compliant models on this benchmark. We note that we exclude models that train with denoising objective, as we focus the performance of the model architecture itsel
Results. The results of EScAIP the Mathench-Discover EScAIP achieves state-of-the-an art performance on this benchmark, outperforming other models. We note that for this geometry relaxation task, models trained with direct forces predictions can be effective. We also provide models trained with gradient-bart based forces, which are better suited for applications such constant-energy molecular dynamics simulations.
benchmark are summarized in Tab.
5.3 Molecules (SPICE)
We evaluate the EScAIP model's performance on the SPICE dataset [aatmaa et al. 2023 which consists of approximately one million molecules across seven different categories. To ensure comparability, we adopt the same training and evaluation settings as used for the MACE-OFF23 model [Kováes et al. [202a]]
Dataset.
Settings. We use a smaller EScAIP model with 45M parameters, trained to predict the energy and forces of each sample. The model's performance then evaluated the different SPICE test datasets and compared directly with MACE-OFF23 Kovácc et al. 2023a
Results. A summary of EScAIP's results the SPICE dataset is provided in Tab. where outperforms MFCE-OFF23 in predicting the energy and forces on the different test sets.
5.4 Rotational Equivariance Test and Analysis
Settings. To assess whether EScAIP learns rotational equivariance after training on various datasets, we design the following procedure: first, we use a held-out validation set that was not seen during training. We randomly sample batch from this validation dataset and pass it 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 equivariance, we calculate the cosine similarity between prediction (B) and the rotated version of prediction (A). This process repeated for 128 batches, and report the average cosine similarity.
Based on the ORB technical orb pre-trained with denoising objective.
2024 it possible that the ORB MPTr model result reported here

---

Table 5: To analyze rotational equivariance, on atomistic systems unseen during training, 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 the force predictions (B) and the rotated version of force predictions (A), After training EScAIP on different datasets, we find that the cosine similarity is consistently 2 0.99, meaning EScAIP is essentially always predicting the rotations correctly.
Dataset
OC20 AIl+MD
MPTrj 45M
SPICE 45M
# of Params.
83M
146M
312M
Before Training
0.2109
0.2940
0.2132
0.2287 0.9999
0.2364
After Training
0.9981
0.9987
0.9994
0.9999
Results. The results of the rotational equivariance analysis are presented in Tab. 3 and the cosine similarity is consistently 2 0.99. These findings indicate that though ESc is not initially rotationally equivariant, after training it is able to correctly map input system rotations to output system predictions. This suggests that these symmetries can be effectively learned from the data. We note that for OC20, where systems consist of surface and adsorbate, the model only trained on systems aligned in the z- z-direction. However, it accurately predicts forces even when the new systems are rotated in various other directions as part of our rotational equivariance test.
6 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 self-attention mechanisms for scalability and expressivity. We demonstrated the effectiveness of EScAIP on 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, 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 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 hand- crafted features offer only simplistic representation of this complexity. A forwar to capture these complexities is to focus on generating comprehensive datasets, ideally accompanied by relevant evalua- tion metrics, allowing NNs to learn the rest of the information through gaining expressivity via scaling.
Future of NNIPs. As datasets continue to grow, training models from scratch on small datasets will likely become unnecessary. While geometric constraints may be beneficial in the very small data regime (though data augmentation techniques can also help here), leveraging the representation of 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 quilt-in geometric 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.
10

---

Acknowledgments and Disclosure of F Funding
This work was initially supported by Laboratory Directed Research and Development (LDRD) funding under Contract Number DEEAAC0--D5CHI1231. It was then supported in part by the Office of Naval Research (ONR) under grant N00014-23- 1-2587. This research used resources the National Energy Research Scientific Computing Center (NERSC), a U.S. Department Energy Office of Science User Facility located at Lawrence Berkeley National Laboratory. operated under Contract No. DE-ACO2- A5CHI1231, We thank Ishan Amin, Sam Blau, Xiang Fu, Rasmus Hoegh, Mit Kotak, Toby Kreiman, Jennifer Listgarten, Ryan Liu, Sanjeev Raja, and Brandon Wood for helpful discussions and comments.
References
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint 2001v:2001.08361, 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL wttp://rweeee.etttttoorri///aoee
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12104-12113, 2022.
Volker L Deringer, Miguel A Caro, and Gábor Csányi. Machine learning interatomic potentials as emerging tools for materials science. Advanced Materials, 3 31(46):1902765, 2019.
Oliver T Unke, Stefan Chmicla, Huziel E Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schütt, Alexandre Tkatchenko, and Klaus-Robas Müller. Machine learning force fields. Chemical Reviews, (21(16):00122-111 2021.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E (3)-cquivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information P Processing Systems, 35:1 1423 1436, 2022.
Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. Equiformery2: Improved equivariant transformer for scaling to higher-degree representations. In The Twelfth International Conference on Learning Representations, 2024. URL/ http://wwww....eeeeooooodd/dddaa
Mario Geiger and Tess Smidt. c3nn: Euclidean neural networks. arXiv preprint 22X.:m22.r4..222000999333 2022.
Saro Passaro and C Lawrrence Zitnick. Reducing so (3) convolutions to so (2) for efficient equivariant gnns. In International Conference Machine Learning, pages 27420-27438. PMLR, 2023a.
Shengjie Luo, Tianlang Chen, and Aditi S. Krishnapriyan. Enabling efficient equivariant operations in the fourier basis via gaunt tensor products. In The Twelfth International Conference on Learning Representations, 2024. URL httph::/ http:ppaaarreeeeeeeeeeeerrrrrrrr met:/foruur1dd-mhyyx3333333
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for
molecular graphs. In International Conference on Learning Representations, 2020a. URL
http:///wrrrveeeeeeeeeeerrrrraaa httpp:://pperevvier hett:/forumm.dd-telbb..tt..
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790-6802, 2021.

---

Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ward Ulissi, C LLwwennce Zitnick, and Abhishek Das. Gemnet-oc: Developing graph neural networks for large and diverse molecular simulation datasets. Transactions on Machine Learning Research, 2022.
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. Open catalyst 2020 (oc20) dataset and community challenges. ACS Catalysis, (1100):6059-667) 2021. doi:
Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, et al. Commentary: The materials project: A materials genome approach to accelcrating materials innovation. materials, I(I),
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xformers: A modular and hackable transformer modelling library.
ooomeerss 2022.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation -qquivironn-quvvriiann neural networks for 3d point clouds. arXiv preprint arX1v:1802.08211, 2018.
Ralf Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Physical Review B, 99(1):014104,, 2019.
Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Advances in Neural Information Processing Systems, 14.8054-8067, 2022.
Saro Passaro and C. Lawrence Zitnick. Reducing so(3) convolutions so(2) for efficient equivariant gnns. In International Conference on Machine Learning 2023b.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2022.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robers Müller. Schnet: continuous- convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
Johannes Gasteiger, Shankari Giri, Johannes T Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprins mrXi::2011.14115. 2020b.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations, 2022. URL
https://openrevieuu http//ormmiiddzllzz222000
Raul P Pelaez, Guillem Simeon, Raimondas Galvelis, Antonio Mirarchi, Peter Eastman, Stefan Doerr, Philipp Tholke, Thomas Markland, and Gianni De Fabritiis. Torchmd-net 2.0: Fast neural network potentials for molecular simulations. Journal of Chemical Theory and Computation, 2024.
Thorben Frank, Oliver Unke, and Roburt Müller. So3krates: Equivariant attention for interactions on arbitrary length- scales in molecular systems. Advances Neural Information Processing Systems, 35:29400-29413, 2022.
Erik J Bekkers, Sharvaree Vadgama, Rob D Hesselink, Putri A van der Linden, and David W Romero. Fast, expressive (n) equivariant networks through weight-sharing in position-orient cotitioonoiientatioon space, 2024.
Chi Chen and Shyue Ping Ong. A universal graph deep learning interatomic potential for the periodic table. Nature Computational Science, 2(11):718-7288 2022.
12

---

Bingging Cheng. Cartesian atomic cluster expansion for machine learning interatomic potentials, 2024.
Mojtaba Haghighatlari, Jie Li, Xingyi Guan, Oufan Zhang, Akshaya Das, Christopher Stein, Farna Heddar-Zadeh, Meili Liu, Martin Head- Gordon,,,,, Luke Bertels, et al. Newtonnet: a newtonian message passing network for deep learning of interatomic potentials and forces. Digital Discovery, 1(3):333-343, 2022.
Yi Liu, Limei Wang. Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference on Learning Representations (ICI,R),002 2022.
Peter Fastman, Pavan Kumar Behara, David I. Dotson, Raimondas Galvelis, John E. Herr, Josh T' Horton, Yuezhi Mao, John D Chodera, Benjamin P Pritchard, Yuanqing Wang, et al. Spice, a dataset of drug- like molecules and peptides for training machine learning potentials. Scientific Data, 10(1):11 2023,
Justin S Smith, Roman Zubatyuk, Benjamin Nebgen, Nicholas Lubbers, Kipton Barros, Adrian Roitberg, Olexandr Isayev, and Sergei Tretiak. The ani- Lecx and ani-Ix data sets, coupled-cluster and density functional theory properties for molecules. Scientific data, 7(1):134, 2020.
Dylan Anstine, Roman Zubatyuk, and Olexandr Isayev, Aimnet2: A neural network potential to meet your neutral, charged, organic, and eeemental-oraanii needs, 2024.
Richard Tran, Janice Lan, Muhammed Shuaibi, Brandon M Wood, Siddharth Goyal, Abhishek Das, Javier Gerras-Domingo Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, et al. The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysts, ACS Catalysis, 13(5):3066-3000 3084, 2023.
Luis Barroso-Luque, Muhammed Shuaibi, Xiang Fu, Brandon M. Wood, Misko Dzamba, Meng Gao, Ammar Rizvi, C. Lawrence Zimick, and Zachary W. Ulissi. Open materials 2024 (omat24) inorganic materials dataset and models, 2024.
Han Yang, Chenxi Hu, Yichi Zhou, Xixian Liu, Yu Shi, Jielan Li, Guanzhi Li, Zekun Chen, Shuizhou Chen, Claudio Zeni, Matthew Horton, Robert Pinsler, Andrew Fowler, Daniel Zügner, Tian Xie, Jake Smith, Lixin Sun, Qian Wang, Lingyu Kong, Chang Liu, Hongxia Hao, and Ziheng Lu. Mattersim: A deep learning atomistic model across elements, temperatures and pressures, 2024.
Amil Merchant, Simon Batzner, Samuel S Schoenholz, Muratahan Aykol, Gowoon Cheon, and Ekin Dogus Cubuk. Scaling deep learning for materials discovery. Nature, 624(7990):80-85, 2023.
Kamal Choudhary, Kevin F Garrity, Andrew CE Reid, Brian DeCost, Adam Biacchi, Angela R Hight Walker, Zachary Trautt, Jason Hattrick-Simpers, A Gilad Kusne, Andrea Centrone, et al. The joint automated repository for various integrated simulations (jarvis) for data-driven materials design. npj computational materials, 6(1): 2020.
Sergey Pozdnyakov and Michele Ceriotti. Smooth, exact rotational symmetrization for deep learning on point clouds. Advances in Neural Information Processing Systems, 36, 2023.
Rui Wang, Robin Walters, and Rose Yu. Approximately equivariant networks for imperfectly symmetric dynamics. In International Conference on Machine Learning, pages 23078-23091, PMLR, 2022.
Mare Finzi, Gregory Benton, and Andrew Wilson. Residual pathway priors for soft equivariance constraints. Advances in Neural Information Processing Systems, 34:30037-30049, 2021.
Omri Puny, Matan Atzmon, Edward Smith, Ishan Misra, Aditya Grover, Heli Lipman. Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations, 2022. UuRliittpp::/oppereeiiee wwt::frummdddzzlUuuuaaaoo 55nXR
and Yaron
Alexandre Duval, Victor Schmidt, Alex Hernandez Garcia, Santiago Miret, Fragkiskos D. Malliaros, Yoshua Bengio, and David Rolnick. Faenet: Frame averaging equivariant gnn for materials modeling, 2023.
Paul Steinhardt, David R Nelson, and Marco Ronchetti. Bond-orientational order in liquids and glasses. Physical Review B, 28(2) 1983.
13

---

Philippe Tillet, Hsiang-Tssi Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10-19, 2019.
Janosh Riebesell, Rhys EA Goodall, Anubhav Jain, Philipp Benner, Kristin A Persson, and Alpha A Lee. Matbench discovery-an evaluation framework for machine learning crystal stability prediction. arXiv preprint arXiv:2308. 74920, 2023.
Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher J Bartel, and Gerbrand Ceder. Chgnet as a pretrained universal neural network potential for charge-informed atomistic modelling. Nature Machine Intelligence, 5()):1031-1041, 2023.
Orbital-Materials. Orb forcefield models from orbital material forcefield models from orbital mate- 2024. Accessed: 2024-10-29)
rials. http:///tthuu..ocoooootttttaaaaa
Yutack Park, Jaesun Kim, Seungwoo Hwang, and Seungwu Han. Scalable parallel algorithm for graph neural network interatomic potentials in molecular dynamics simulations. Journal of Chemical Theory and Computation, 2024.
Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M. Elena, Dávid P. Kováes, Janosh Riebesell, Xavier R. Advincula, Mark Asta, William J. Baldwin, Noam Bernstein, Arghya Bhowmik, Samuel M. Blau, Vlad Cärare, James Darby, Sandip De, Flaviano Della Pia, Volker Deringer, Rokas Elijošius, Zakariya El-Machachi. Edvin Fako, Andrea C. Ferrari, Annalena Genreith-Schriever, Janine George, Rhys E. A. Goodall, Clare Grey, Shuang Han, Will Handley, Hendrik H. Heenen, Kersti Hermansson, Christian Holm, Jad Jaafar, Stephan Hofmann, Konstantin S. Jakob, Hyunwook Jung, Venkat Kapil, Aaron Kaplan, Nima Karimitari, Namu Kroupa, Jolla Kullgren, Matthew Kuner, Domantas Kuryla, Guoda Liepuoniute, Johannes T. Margraf, Ioan- Bogdan Magdāu, Angelos Michaelides, J. Harry Moore, Aakash A. Naik, Samuel Niblett, Sam Walton Norwood, Niamh O'Neill, Christoph Ortner, Kristin A. Persson, Karsten Reuter, Andrew S. Rosen, Lars L. Schaaf, Christoph Schran, Eric Sivonxay, Tamás K. Stenczel, Viktor Svahn, Christopher Sutton, Cas van der Oord, Eszter Yarga-Umbrich, Tejs Vegge, Martin Vondrák, Yangshuai Wang, William C. Witt, Fabian Zills, and Gábor Csányi. A foundation model for atomistic materials chemistry. 2023.
Technical blog:
Introducing
the orb ai-based interatomic com.iaacs....ooaaaaoooommm..ooot /pttt/
potential.
technical-blog- how.ogginnrrouuiig-tteeeee orbe--orb-ae ai -baeett-immeerrttooo----teaaaall Accessed: 2024-10-29.
2024.
Dávid Péter Kovács, J Harry Moore, Nicholas J Browning, Ilyes Batatia, Joshua T Horton, Venkat Kapil, William C Witt, Ioan-Bol Bogdan Magdāu, Daniel J Cole, and Gábor Csányi. Mace-off23; Transferable machine learning force fields for organic molecules. arXiv preprint arXiv:2312.15211, 2023a.
Richard Sutton. The bitter lesson. Incomplete Ideas (blog), 13(1):38, 2019.
Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T Unke, Adil Kabylda, Huziel E Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, (2)=eadf0873, 2023.
Yunyang Li. Yusong Wang, Lin Huang, Han Yang. Xinran Wei, Jia Zhang, Tong Wang, Zun Wang, Bin Shao, and Tie-Yan Liu. Long-short-range message-passing: A physics- informed framework simulation. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreviie: nwt. www.mriddddddrddtt...
capture non-local interaction for scalable molecular dynamics
Dávid Péter Kovács, Ilyes Batatia, Eszter Sára Arany, and Gábor Csányi. Evaluation of the mace force field architecture: From medicinal chemistry materials science. The Journal of Chemical Physics, 159(4). 2023b.
Sanjeev Raja, Ishan Amin, Fabian Pedregosa, and Aditi S Krishnapriyan. Stability- awarr training of neural network interatomic potentials with differentiable boltzmann estimators. e arXiv preprint arXiv:2402. 3984, 2024.
Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafael Gomez-Bombarell and Tommi S Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. Transactions on Machine Learning Research, 2023. 2023.
14

---

Lez Singem-0.3
AT 1-2 Slopam-611
Foree MAE (meVVA)
x 100
More AT L-2 Stope 0.40
Origtnal L-1 Slogo- 0.28
Slape- 0.32
Siope-0.33
Energy MAE (mee))
3.2x 10
3.1x10
Morw AT L-4 Slopem0.00 0.50
Criginal Slopem
Morw AT 1-2 Slope-0.11
Slapam-6.06
2.9
2.8 10
2
25%
50%
75%
100%
25%
50%
75%
100%
Portion of Training Set
Portion of Training Set
Figure 6: Force MAE vs. Training Dataset Size for EquiformerV2 ablation study on the OC20 2M dataset. Slope is fitted by linear regression. We scale the parameters of the original L= and L = 4 models from Liao et al. 2024] through the attention mechanisms and/or spherical channels, such that the number of parameters is approximately equal to the original L= model. As the training dataset size increases, the scaled L=2 and L= 4 models have a steeper slope, indicating faster performance improvement with increasing data.
A Additional Details on Investigations
We provide additional details on our investigations from
Results Ablation Study comparing Force MAE vs. Training Dataset Size. To further investigate how scaling efficiency varies function of training dataset size, we train the parameter- controlled Equiformer V2 models with different amounts of training data. The results in Fig. Fig. show that the scaled = 2 and L = 4 models exhibit steeper performance improvement (log-log slope) compared to the original L= 6 model. In particular, the More AT configuration (more attention) has 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
Additional Details and Results on Experiments
B.1
Catalysts (OC20 and OC22)
To illustrate EScAIP's scalability, train the model varying sizes of training data and model configurations. The results, shown in Fig. g indicate a clear trend: as model and data sizes grow, EScAIP's performance continues improve. We also include results complement Fig. [ 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 1) 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.
B.2 Large Molecules (MD22)
We evaluate the performance of our EScAIP model on the MD22 dataset [Chmiela et al.) 2023), which consists seven molccular systems with varying sizes. It consists of energy and force labels calculated from DFT simulations.
Dataset.
Settings. We use an EScAIP model 15M parameters on each system and evaluate the performance
on the test set (train-test test split 95:5). The model is compared with MACE [Batatia et al.) [002].
VisNet-LSRM et al. 2024], and sGDML [Chmiela et al. 2023 note that there were discrepancy of results of VisNer-LSRM
the MACE Kovács al. 2023b paper and the VisNet-LSRM [Li 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 hundreds to thousands samples, and also apply data
augmentation (randomly rotating each training sample 16 times).
15

---

260
260
216M
EquiformerV2
216M
EquiormeerV2
Epergy MAE (zue
250
200M
GomNet OC EScAIP
eSCN
Emergy MAE (meV)
250
GemNet.CC OC:
eSCN
EScIIP
240
240
31M
31M
153M
83M
83M
230
230
146M
220
146M
317M
210
210
Inference Speed (Sample/Sec)
Inference Memory (GB/Sample)
Figure 7: Efficiency, performance, and scaling comparisons between EScAIP and baseline models on the Open Catalyst dataset. Energy MAE (meV 4) vs. Inference Speed (Sample/Sec t) and Energy MAE vs. Memory (GB/Sample J) reported. EScAIP achieves better performance with smaller time and memory cost.
EScATP 45M
ESCAIP 45M
EScAIP 83M
325
EScAIP 146M
ESCAIP 146M
Force MAE (me)/AA
EquifonmerV2
EquiformerV2 31M
Energy MAE (me V
Equiformer 83M
Equi Formmeer 31M
25
Equitomerv2 134M
Equi For 134M
EquiformerV
EquiformerV2 153M
20
250
15
225
172M All+DD
500k
2M
172M
All+
Training Data Size
Training Data Size
Figure 8: Scaling experiment of EScAIP on OC20. Forces MAE (/À) and Energy (me across 4 validation splits are reported. For 500k, 1M, and 2M split, the ESCAIP model is trained for 30 epochs: for All+MD, the EScAIP model is trained for 8 epochs. Force and Energy MAE consistently decreases as model size and training data size increases.
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 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. Interest- ingly, MD22 is not particularly large dataset: the training dataset sizes are in the thousands. Despite this, a scalable architecture with high parameter counts is still able achieve good performance.
16

---

Table 6: EScAIP performance on the MD22 dataset. The results are reported in Energy (meV V/atom), Force V/À) and h(r) (unitless) mean absolute error (MAE).
Tetra- peptide
Fatty acid
Tetra- saccharide
Nucleic acid (AT-AT)
Nuclcic acid (AT-AT-CG-CG)
Buckyball cutcher
Double-wald walled
Metric
nanotube
of Atoms
42
56
87
60
118
148
370
Energy
0.40
1.0
2.0
0.52
0.52
0.34
0.47
sGDML
Force
33
29
31
Energy
0.064
0.102
0.062
0.079
0.058
0.141
0.194
MACE
Force
2.8
3.8
VisNet-LSRM
Energy
0.080
0.058
0.044
0.055
0.049
0.124
0.117
(MACE)
Force
5.7
3.6
5.0
5.2
8.3
11.6
28.7
VisNet-LSRM
Energy
0.068
0.068
0.053
0.056
0.039
(Paper)
Force
3,9
2.5
3.3
3.4
4.6
Energy
0.053
0.052
0,041
0.062
0.042
0.112
0.095
EScAIP
Force
3,3
3.2
3.1
3.8
4.3
2.5
8.1
M(r)
0.09
0.11
0.15
0.05
0.21
17

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide a short (1 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist one of the factors in their evaluation. While "[Yes] generally preferable to "[No] ", is perfectly acceptable to answer "[No] provided proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). ). general, answering "[No] " or "[NA] is not grounds for rejection. While the questions are phrased in binary way, acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper the supplemental material, provided appendix. you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
The claims made in the abstract and introduction are detailed in the results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss some limitations in
18

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should diseuss possible limitations of their approach to address problems privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
We do not have a theoretical derivation.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question:
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that affects the main claims conclusons the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, we provide all details of our method and a schematic.
Guidelines:
The answer NA means that the paper does not include experiments.
19

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible important, regardless of whether the code and data are provided or not.
If the contribution is 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 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 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 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. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors 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 (c.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access data and code
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.
Please see the NeurIPS code and data submission guidelines (https hwtpabbliiiggddeeeeeeeee
//tips.cc/public//
for more details.
While we encourage the release of code and data, we understand that this might not pos- sible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central the contribution (e.g.. for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https://nips.. http:/iisscccppu///p/ httpccc/ppbllccc http://ss...pp.iiiciiieeee/////s CodeeSuum iemiisssioni wwwiisssiiopooooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
20

---

6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperpa- rameters, how they were chosen, type of optimizer, etc.) necessary 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 the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance the experiments?
Answer: [No]
Justification: experiment.
Due to computational cost, we are not able provide error bars for every
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative еттог rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification:
Yes, we provide computational details.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well as estimate the total compute.
21

---

The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary failed experiments that didn make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code
Answer: [Yes]
Justification:
We have reviewed this and the research conforms to this code.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consideration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, this is also discussed the introduction.
Guidelines:
The answer NA means that there 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 negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeem 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 partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This is not relevant to this work.
Guidelines:
The answer NA means that the paper poses no such risks.
22

---

Released models that have 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 dataset.
The authors should state which version of the asset is used and, if possible, include 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.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, ppaperswithodee com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For cxisting 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 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Yes, provide new model and include implementation details. Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details the dataset/code/model as part 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL include an anonymized zip file.
14. 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 't have this.
Guidelines:
23

---

The answer NA means that the paper does not involve crowdsoureing 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 much detail possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve this.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equívalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "zR1fRDHwtJ", "reviewer": "Reviewer_PKom", "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.", "questions": "see above", "limitations": "not discussed explicitly", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The problem of efficiently predicting complex molecule properties is important in practice.\n2. The paper has interesting observations about scaling behaviors.\n3. The paper is easy to follow and written and presented well.\n4. 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?\n\n2. 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?\n\n3. Why EquiformerV2 is not compared to in Table 2:?\n\n4. Can the authors provide runtime and memory comparison for different molecule sizes during testing? This can give practitioners important insights.\n\n5. 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.\n\n6. 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.\n\nOther minor issues:\n\n- L62 - broken reference\n- fig. 1 must be referenced in the text\n- fig. 1 \"Right: Force MAE vs. Datasize. \" is not correct or the plot xlabel is not correct\n- In Table 2, for Fatty acid h(r) is missing for EGAP", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback and valuable suggestions! We address your concerns as follows:\n\n**Question: inefficient attention over edges**\n\nWe clarify the details about the architecture of our model: in the realm of NNIPs, most of the molecular graph is constructed by a radius graph with a “max number of neighbors\" constraint, which essentially resembles a k-nearest neighbor graph with a cut-off distance. This means the number of nodes in each neighborhood is near constant. Thus, we organize the edge features as $[N, K, H]$ (with padding), where $N$ is the number of nodes, $K$ is the max number of neighbors (usually chosen as 20-30 in OC20), and $H$ is the hidden dimension. The attention is parallelized over each neighborhood, i.e. the “sequence length” in the attention is $K$. Thus, the attention component of our model scales as $O(NK^2)$ in time efficiency, rather than $O(E^2)$. Compared with Graphormer, we don’t need to construct a $N\\times N$ attention bias matrix to inject edge features (which is a fully connected graph), making our model more memory efficient as well.\n\n**Question: different training and validation splits**\n\nThank you for the suggestions. We have revised our MD22 experiments to have the same train/val split as the baselines. For the OC20 experiments, we used the same split as released in the official dataset, which is the same split that all the models we compare to use as well. We also added results for OC20 All+MD. All the added results can be found in the pdf attached in the overall response. As seen here, **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 [1]).** In the updated MD22 experiments, our model is better than state-of-the-art models like MACE [2] on almost all of the molecules (see general comment for details).\n\n**Question: Equiformer not compared in MD22**\n\nThere is no official implementation of Equiformer on the MD22 experiment. We have tried to train EquiformerV2 [1] on the MD22 dataset, but it didn't produce reasonable results. We think it's not fair to compare our model with our unofficial implementation of EquiformerV2 [1]. Thus, we didn't include it here. \n\n**Question: Runtime and memory for different molecule size**\n\nWe have tested our 86M model with different molecule size (atoms are randomly positioned), the results are in the table below:\n\n| Molecule Size | Time (ms) | Memory (GB) | Time wo/ Flash Atten  (ms) | Memory wo/ Flash Atten (GB) |\n|-|-|-|-|-|\n| 50            |   134   |    4.3    |         462             |             5.7           |\n| 100           |   146   |    4.9    |         518             |               6.2         |\n| 150           |  153    |     5.2   |          573            |                6.8        |\n| 200           |   176   |    6.3    |           635           |                 7.4       |\n| 250           |   195   |    7.1    |              697        |                  8.2      |\n\n\n**Question: Ablation over model components**\n\nWe have included the speed comparison of Flash attention and vanilla attention in the table above. It's obvious that the Flash attention kernel is helpful for both memory and speed of our model. We want to emphasize that the model relying on higher order group representation, such as EquiformerV2, is not possible to use such kernels. This is because they have to use a special attention mechanism to maintain the node features in the group representation space. For our model, since the node features are scalars, we are able to apply optimizations from other fields such as NLP and CV.\n\n**Question: Ablation over model size and data size**\n\nWe have added a scaling experiment of EGAP (see the general response for details, and pdf Fig. 1), which includes an ablation over model size and data size.\n\n**Question: Minor issues**\n\nThank you for pointing these out. We will fix them in the final version of the paper.\n\n[1] 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.\n\n[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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j4vgsyaVni", "reviewer": "Reviewer_VH7W", "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.", "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.\n- How does the model inference time compare to other models?\n\n(Typo in line 124)", "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.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The proposed model achieves SOTA performance on MD22, and comparable performance to spherical equivariant model such as EquiformerV2, by means of invariant features.\n- Experiments and ablations are extensive\n- The model stability is probed\n- 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\n- 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'.\n- The authors could have used the extra space they have to perform some other experiment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and valuable suggestions! We address your concerns as follows:\n\n**Question: more experiments and baselines**\n\nWe have updated the results of EGAP on OC20 2M and All+MD split, with comparisons of relevant baselines. We also added another experiment on the OC22 dataset. **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)** (see general comment for details).\n\n**Question: improve Table 1**\n\nIn table 1, the NequIP model is the representation of the line of work that uses equivariant group representations as node features. Other models, such as Equiformer and eSCN, all have the same formulation as NequIP. The GemNet model represents the line of work with invariant node features, such as DimeNet and SchNet. These two categories cover most of the NNIP architectures. Following your suggestion, we will add the formulation of MACE in Table 1, which is a special case of group representation models. MACE used the local directional info in the construction of Atomic Cluster Expansion, but its core message update step is still the same as NequIP.\n\n**Question: effect of L in BOO**\n\nWe have performed an ablation with different values of L in BOO features on OC20 2M. The model all have about 86M parameters (there are slight differences in parameter count due to different embedding layers for BOO features, but it's on the order of thousands and negligible). All of them are trained with identical hyper parameters for 30 epochs. The results are as follows:\n\n|L in BOO feature|L=12|L=6|L=0|\n|-|-|-|-|\n|Force MAE (meV/A)|0.02034|0.02148|0.02317|\n\nThe conclusion is that higher L in BOO features indeed helps with the model performance. This could because higher order directional features help the model learn a better representation of the neighborhood.\n\n**Question: inference time comparison**\n\nWe have included the training time, inference speed, and FLOPs count in Table 1 of the attached pdf.\n\n**Question: typo**\n\nThank you for pointing this out. We will fix it in the final version of the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sSEDvRD5OD", "reviewer": "Reviewer_L3rz", "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.", "questions": "- An apples-to-apples comparison to the other models in Table 2, standardizing the size of the training dataset, would be welcome.\n- More details on the way that the attention mechanism is scaled would be welcome. Where, exactly, are the extra parameters spent?\n- 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.\n- 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.\n- 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": "See above.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n- 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.\n- 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?\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback and valuable suggestions! We address your concerns as follows:\n\n**Question: scalability of the model, and performance**\n\nThank you for the suggestions. We agree with your definition of scalability, and we will clarify it further in the introduction of the final version of the paper. Our definition of a “scalable\" model follows that the model performance improves with scale: the performance increases as the size of the dataset, model parameters, and computational resources increases. A “scalable” model architecture means that the model can efficiently utilize resources (data and compute). 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. \n\n**Question: reason for the higher efficiency than EquiformerV2**\n\nThe efficiency gain compared with Equiformer mostly comes from the change to invariant node features. Equiformer’s node features are equivariant to rotations, and they have to use costly tensor products to maintain it. In EGAP, the node features are invariant scalar, which not only eliminates the time-consuming tensor products, but also enables us to apply optimizations from other fields such as NLP and CV.\n\n**Question: performance on OC20**\n\nWe 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) (see general comment for details).\n\n**Question: investigation of higher-order tensor beyond EquiformerV2**\n\nFirst, we have performed the same scaling experiments on eSCN [1], which also utilizes higher order tensors as node features. We trained different versions of eSCN on OC20 2M, where we add parameters on “Hidden\" dimensions, “Spherical” channels, or “Both\". The results also indicate that lower order representations could achieve comparable performance to the higher order ones. We also controlled the number of parameters to be the same as the EquiformerV2 model in the investigation. The main difference between eSCN and EquiformerV2 is the message passing step, where eSCN uses linear layers and EquiformerV2 uses attention. The better performance of EquiformerV2 also indicates that the attention mechanism is vital for scaling.\n\n| L | Params | Energy MAE | Force MAE|Component|\n|-|-|-|-|-|\n| 6 | 83M | 273 | 20.4 | Both |\n| 2 | 83M | 282 | 23.3 | Both |\n| 4 | 83M | 281 | 21.1 | Both |\n| 2 | 83M | 279 | 22.6 | Hidden |\n| 4 | 83M | 277 | 20.7 | Hidden |\n| 2 | 83M | 287 | 24.3 | Spherical |\n| 4 | 83M | 284 | 21.7 | Spherical |\n\nSecond, the tensor product for higher-order tensors in EquiformerV2 is an efficient implementation. Vanilla tensor products (used in MACE and NequIP) scale $O(L^6)$ in time, where L is the rotation order. In EquiformerV2, the authors adopted the SO2 Convolution operation to reduce this time complexity to $O(L^3)$. However, it is still 30x slower than EGAP.\n\nFinally, our model is now state-of-the-art on both OC20 All+MD and OC20 2M. This further verifies that scaling the attention is more effective than scaling the order of tensors.\n\n**Question: MD22 experiment**\n\nWe have revised our MD22 experiments to have the same train/val split as the baselines. Our model is better than state-of-the-art models like MACE on almost all of the molecules (see general comment for details).\n\n**Question: details on the experiment hyper parameters**\n\nThe way to scale the attention mechanism is by increasing the hidden dimension of the attention and the number of heads. We will add a detailed table for the model parameters used in the experiments in the appendix. \n\n**Question: details of the output block**\n\nThe design of the output block follows GemNet, but with a few modifications. In GemNet’s output head, it predicts a magnitude for each edge direction in the neighborhood, and aggregates them together. This makes it recover equivariance for force predictions. In EGAP, we predict a transformation for each edge direction that alters the direction, which breaks equivariance. In our experiments, we found that this symmetry breaking output block made the model perform better. The reason could be that this formulation is more expressive and easier to optimize. We also checked the model's equivariance to rotation after training. The procedure is as follows: We take N random batches from the dataset and pass them through the model to get output A. Then, we randomly rotate the input coordinates and pass them through the model again to get output B. Finally, we rotate the output B inversely and check the force cosine similarity with output A. In OC20 All+MD, the randomly initialized EGAP model with 134M parameters has a cosine similarity of 0.1384. After 4 epochs of training, it becomes 0.9627. This shows that the model learns approximate equivariance through scaling, without architectural constraints.\n\n**Question: JMP has better performance on MD22**\n\nWe note that the JMP model is pre-trained on much more data (over 100M data points), and then fine-tuned on MD22. We will clarify that we are state-of-the-art among all models solely trained on MD22.\n\n**Question: training procedures and statistics**\n\nOur training procedure follows the EquiformerV2 model. We will add a detailed table for the training hyper-parameters in the appendix. We included some training statistics (training time, speed, FLOPs) in Table 1 of the attached pdf.\n\n[1] Passaro, S., & Zitnick, C. L. (2023, July). Reducing SO (3) convolutions to SO (2) for efficient equivariant GNNs. In International Conference on Machine Learning (pp. 27420-27438). PMLR.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Zh7uivbHf7", "reviewer": "Reviewer_wbCz", "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.", "questions": "1. Why didn't the authors follow the same dataset split procedure used by baseline models, especially for the MD22 dataset?\n2. Can the authors provide scalability experiments showing how EGAP performs as model size increases?\n3. 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?\n4. Can the authors explain the discrepancies between their reported results for baseline models and those in the original papers, particularly for the MD22 dataset?\n5. 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": "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.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper addresses an important challenge in the field of NNIPs by focusing on scalability and efficiency.\n2. The presentation of the model illustration is clear and appealing, aiding in understanding the proposed architecture.\n3. 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.\n\n2. 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:\n   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.\n   b) The authors have effectively used a different train/validation/test set compared to baseline models, rather than just using additional training data.\n   c) The claim about performance differences with dataset size is not supported by evidence in the paper.\n\n3. 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.\n\n4. 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.\n\n5. 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.\n\n6. 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.\n\n\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Overall Response to All Reviewers

We thank all the reviewers for their constructive feedback. In our general response, we address some common issues discussed by the reviewers and provide updated experimental results for these points:

**State-of-the-art results on OC20-2M and OC20-All+MD.** To showcase the scalability of our model, we have trained our model on the full OC20 All+MD split (172M data points). We also tuned our model further on the OC20 2M split. The results are presented at Table 1 in the attached pdf. **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 model (EquiformerV2 [1]).**

**Scaling experiments on OC20.** We have conducted scaling experiments by training our model with different numbers of parameters on different numbers of training samples. The results can be found in Figure 1 at the attached pdf. We control the training epochs to be 30 for 500k, 1M, and 2M training samples, and 4 epochs for All+MD split. There is an obvious trend of improved performance as the number of parameters and data size increases. EGAP also scales more efficiently in terms of the number of parameters compared with the EquiformerV2 in the 2M split, while taking less training time.

**New experiment: state-of-the-art results on OC22.** To further show the effectiveness of our model, we have added another experiment on the OC22 dataset (around 10M data points of training data). Our model is state-of-the-art on OC22 (see results in pdf Table 1).

**Revised MD22 experiment with same training and validation split.** We retrained our model with the same train/val split as the baselines (5x-10x smaller than the original training set size). To accommodate the smaller training set, we decreased the model parameters (from 45M to 15M) and applied data augmentation (randomly rotating each training sample 16 times). The updated results can be found in Table 2 in the pdf attached. **In this scenario, EGAP is better than state-of-the-art models like MACE [2] on almost all of the molecules.**

[1] 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.

[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.

Meta-review: 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.
Decision: Accept (poster)