Thank you for your constructive feedback and valuable suggestions! We address your concerns as follows:

Question: dataset split and inconsistent dataset usage of MD22

Thank you for the suggestions. We have revised our MD22 experiments to have the same train/val split as the baselines. Specifically, we follow the procedure used in MACE [1]: we use the number of samples used in training sgdml [2] as training+validation set (95:5), and the rest as test set. In the updated experiment, our model is better than state-of-the-art models like MACE [3] on almost all of the molecules (see general comment for details).

Question: limited experimental scope (need to test on OC20-All+MD) and inconclusive performance difference on OC20-2M

We have updated the results of EGAP on OC20 2M and All+MD split, with comparisons of relevant baselines. Our model is now state-of-the-art on both OC20 All+MD and OC20 2M, and 10x faster at training, 30x faster at inference than the previous best-performing models (EquiformerV2 [4]) (see general comment for details).

Question: lack of scalability experiments

We have added an experiment demonstrating the scalability of EGAP in the attached pdf (Figure 1). The details can be found in the overall response. From the figure, we see that EGAP scales well when data size and number of parameters increases. It also utilizes computation resources more efficiently, as it takes 10x less training time than EquiformerV2 and achieves better results.

Question: result discrepancies

Thank you for pointing this out. We took the numbers directly from the MACE paper [1], as we are using the MACE data split in the training. We have added a new row in Table 2 of the pdf for the recalculated results in the VisNet-LSRM paper [5].

---

[1] Kovács, D. P., Batatia, I., Arany, E. S., & Csányi, G. (2023). Evaluation of the MACE force field architecture: From medicinal chemistry to materials science. The Journal of Chemical Physics, 159(4).

[2] Batatia, I., Kovacs, D. P., Simm, G., Ortner, C., & Csányi, G. (2022). MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35, 11423-11436.

[3] Chmiela, S., Vassilev-Galindo, V., Unke, O. T., Kabylda, A., Sauceda, H. E., Tkatchenko, A., & Müller, K. R. (2023). Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9(2), eadf0873.

[4] Liao, Y. L., Wood, B. M., Das, A., & Smidt, T. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. In The Twelfth International Conference on Learning Representations.

[5] Li, Y., Wang, Y., Huang, L., Yang, H., Wei, X., Zhang, J., ... & Liu, T. Y. Long-Short-Range Message-Passing: A Physics-Informed Framework to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation. In The Twelfth International Conference on Learning Representations.

Thank you for your constructive feedback and valuable suggestions! We address your concerns as follows:

Question: scalability of the model, and performance

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

Question: reason for the higher efficiency than EquiformerV2

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

Question: performance on OC20

We have updated the results of EGAP on OC20 2M and All+MD split, with comparisons of relevant baselines. Our model is now state-of-the-art on both OC20 All+MD and OC20 2M, and 10x faster at training, 30x faster at inference than the previous best-performing models (EquiformerV2) (see general comment for details).

Question: investigation of higher-order tensor beyond EquiformerV2

First, 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.

Second, 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.

Finally, 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.

Question: MD22 experiment

We 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).

Question: details on the experiment hyper parameters

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

Question: details of the output block

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

Question: JMP has better performance on MD22

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

Question: training procedures and statistics

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

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

Thank you for your positive feedback and valuable suggestions! We address your concerns as follows:

Question: more experiments and baselines

We 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).

Question: improve Table 1

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

Question: effect of L in BOO

We 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:

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

Question: inference time comparison

We have included the training time, inference speed, and FLOPs count in Table 1 of the attached pdf.

Question: typo

Thank you for pointing this out. We will fix it in the final version of the paper.

Thank you for your constructive feedback and valuable suggestions! We address your concerns as follows:

Question: inefficient attention over edges

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

Question: different training and validation splits

Thank 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).

Question: Equiformer not compared in MD22

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

Question: Runtime and memory for different molecule size

We have tested our 86M model with different molecule size (atoms are randomly positioned), the results are in the table below:

Question: Ablation over model components

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

Question: Ablation over model size and data size

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

Question: Minor issues

Thank you for pointing these out. We will fix them in the final version of the paper.

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