EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations

Summary

The authors propose EquiformerV2, which is an improvement over the original Equiformer architecture. The main improvement is using a more efficient parameterization of the tensor products used in Equiformer which is computationally expensive for higher order representations. The more efficient parameterization involves using SO(2) linear layers instead of SO(3) tensor products. Moreover, they have three architectural improvements – adding an extra layer norm during attention, S2 activations instead of gates and "separable" layer normalization which differs from the previous layer normalization in terms of the denominator used. They perform experiments on OC20 to show improvements compared to other models in the literature.

Strengths

Overall, the idea of using computing tensor products more efficiently using SO(2) linear layers makes good sense in its application to Equiformer in order to scale it up via higher order SO(3) representations. Empirical work around improving the architecture by appropriate use of different type of layer norms and activation functions is also valuable. The components feeding into this model aren't always original, but that shouldn't detract from the significance of putting together the full model and showing that it improves performance on a benchmark. The clarity of the writing is usually good but there are places where it can be improved (see below).

Weaknesses

The main weakness is insufficient comparison against the original Equiformer architecture. There are four differences versus the V1 architecture as far as I can tell: (1) the more efficient parameterization of tensor products via the ideas of eSCN, (2) an extra layer norm applied to scalar features in the attention module, (3) separable S2 activation instead of gates and (4) layer normalization with a different way of computing the denominator (called "separable" layer normalization in the paper). - There is an ablation in Table 1(a) on one of the datasets looking at the effects of changes (2), (3), and (4). It is unclear how much of the difference in performance (specially, on energies) is statistically significant. I understand training each model is expensive, but since the experiments in this table are on the smaller S2EF-2M dataset, isn't it possible to have error bars here? - Why isn't there a comparison in Table 1(a) of the effect of incorporating change (1) in EquiformerV2? The SO(2) linear layers replacing the tensor product are in principle as expressive as the tensor products but the optimization dynamics can be different since the parameterization of the architecture is different. It would be good to see a side-by-side comparison of Equiformer's V1 and V2 when keeping other architectural hyperparameters fixed (e.g. number of channels, maximum representation order etc.) - It would be good to include EquiformerV1 performance in the other sub-tables of Table 1 too, in order to see the cummulative effect of the changes to the original architecture. - Similarly, it would be useful to see EquiformerV1 performance in Table 2 and 3 as well (possibly with a lower $L_\text{max}$ but with other hyperparameters tuned). - How does the setup for IS2RE in Table 2 differ from EquiformerV1's Table 3? Is one with relaxations after training on S2EF and the other is through direct energy prediction? A side-by-side comparison under the same setup would help a lot here. - How necessary are the higher orders compared to increasing the number of layers or channels? SO(2) linear layers instead of tensor products will improve the efficiency for lower orders too, so can we make up for the performance of higher orders by instead increasing the number of layers or channels? - Have you benchmarked EquiformerV2 on QM9 where EquiformerV1 has already been benchmarked? It's a dataset that is different from OC20 in various ways, so it would be useful to see a comparison. For other suggestions for improvements, see the questions below. Equiformer(V1): https://arxiv.org/abs/2206.11990

Questions

- The explanation of separable S2 activation was unclear: what does line 186 mean? "The activation first converts vectors of all degrees to point samples on a sphere for each channel, applies unconstrained functions F to those samples, and finally convert them back to vectors." - Have you done direct IS2RE prediction as done in the EquiformerV1 paper with/without an auxiliary denoising loss? - Can you justify "higher degrees can better capture angular resolution and directional information"? Perhaps via an example? - Was stochastic depth used in EquiformerV2? Related work section mentions it was used in EquiformerV1? - Line 69: "apply typical functions to rotated features during message passing" was unclear to me. - Line 142: "eSCN convolutions go one step further and replace the remaining non-trivial paths of the SO(3) tensor product with an SO(2) linear operation to allow for additional parameters of interaction between $\pm m_o$ without breaking equivariance." Can you explain this in the context of Appendix A.3? - Line 175: why would $f^{(0)}_{ij}$ be less well-normalized? - Line 178: there is a slight overload in notation with $a$ and $a_{ij}$. - Line 210: why is there no centering in the layer norm formulation? Wouldn't that continue to preserve equivariance as the mean vector belongs to the same representation? There is centering in the 0th degree activations and traditional layer norm. - Line 210: do you use a small $\epsilon$ in the new separable layer normalization denominator for numerical stability? - Line 280: "Finally, these modifications enable training for longer without overfitting (row 7)" don't both models improve performance when trained for longer, not just the model in row 7? - On comparison of speed: are all the models using the same training/inference pipeline such that the comparison is fair? It would be useful to report inference cost in FLOPs or a similar metric as well.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Limitations and broader impact are adequately addressed.

Summary

This paper provides a new equivariant graph neural networks named EquiformerV2 to enhance the original Equiformer performance. It uses four new modules. The first module is to use the convolution in the eSCN (https://arxiv.org/abs/2302.03655) to replace the depth-wise tensor production accelerating the speed. The second module is the separable $S^{2}$ activation using the spherical grids in SCN (https://arxiv.org/abs/2206.14331) to encourage the non-linearity. The third module is the separable layer normalization which uses the variance of all equivariant feature $\ell \geq 1$ to do normalization. The fourth module is the attention re-normalization using a layer normalization before the non-linear function of the attention score branch.

Strengths

1. This paper is well organized and written. The figure clearly shows the modification on the model architecture. 2. The experiments in Table 2 supports the proposed model architecture can achieve SOTA performance on the OC20 All training sets as well as OC20 All+MD. For example, on OC20 All, test energy MAE is improved 13meV in S2ET test and 25meV in IS2RE test. Such improvement is great. These two datasets usually take extensive training time to perform experiments on them. From the Throughput metric, the EquiformerV2 has better training efficiency compared to current baselines.

Weaknesses

1. As a suggestion, it will be better if efficiency study includes both training time and inference time. Computational complexity metric such as FLOPs or MACs can help measure the inference complexity. 2. The description of spherical grid is not very clear in the paper. Although it is introduced in the SCN paper, I think a brief introduction can help people understand why such operation can enhance the non-linearity.

Questions

1. For the speed up of DFT calculations, could you briefly introduce the AdsorbML and explain why EquiformerV2 can be used to accelerate the DFT? 2. Is it possible to apply such efficient equivariant architecture in directly accelerating the DFT calculations such as the SchNorb (https://www.nature.com/articles/s41467-019-12875-2), PhiSNet (https://arxiv.org/abs/2106.02347), QHNet (https://arxiv.org/abs/2306.04922) and recent QH9 dataset (https://arxiv.org/abs/2306.09549)? Since the computational complexity is also a problem for these equivariant networks. 3. For the equivariance, since the equiformerV2 uses the spherical grids, I am curous about whether equiformerV2 is rigorous equivariant or approximate equivariant. If this architecture is approximate equivariant, I think it will be great if there are some experiments to verify the equivariance.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

1. As an suggestion, to comprehensively study the improvement on original Equiformer, it will be better if there is a comprehensive experiments on the original datasets such as QM9.

Summary

In this paper, the authors proposed EquiformerV2, which is an equivariant network for 3D molecular modeling. The EquiformerV2 is built on the Equiformer with several architectural modifications: 1) replace SO(3) convolutions (tensor product operations) with efficient SO(2) counterparts from eSCN; 2) Attention Re-normalization; 3) Separable S2 activation; 4) separate Layer Normalization. These changes enable EquiformerV2 to achieve good performance on the large-scale OC20 benchmark and also the new AdsorbML dataset.

Strengths

1. The targeted problem is of great interest to the community. The EquiformerV2 provides another attempt to enlarge the maximum degree of irreducible representations and obtain performance gains on large-scale DFT benchmarks. 2. Good empirical performance. On the OC20 benchmark, the EquiformerV2 achieves state-of-the-art performance on the Structure-to-Energy-Force task. The model trained on this task further serves as a force-field evaluator to achieve strong performance on IS2RS and IS2RE tasks. The EquiformerV2 outperforms the compared baselines on all these tasks, especially on force prediction. Additionally, it also improves the success rate a lot on the AdsorbML dataset.

Weaknesses

1. The novelty of integrating the eSCN convolution and S^2 activation into the Equiformer is limited. Among the proposed architectural changes, the eSCN convolution is the key component to enable Equiformer to use irreducible representations of higher degrees, and the S^2 activation also replaces all non-linear activations. However, these design strategies should be mainly credited to the eSCN work. 2. The motivation for the other architectural modifications should be thoroughly clarified. First, lines 174-175 in Section 4.2 suggest the "less well-normalized" issue, which motivated the authors to propose re-normalization and Separable Layer Normalization. It is better to provide further quantitive evidence to reveal how such an issue affects the model's performance, and why these modifications could remove or mitigate such effect. Second, the authors proposed Separable S^2 Activation because the original S^2 activation would make the training process diverge. However, it is hard to understand why such separable modifications could make the training process stable. Is there any further essential reason behind such a phenomenon? It is suggested to provide further analysis on such modifications. 3. More analyses are necessary if the computational resources are acceptable. First, the authors did not provide a performance comparison between whether using the eSCN convolution. Although the eSCN convolution is equivalent to the SO(3) counterpart, the computation processes of these two parameterizations are different. Second, it is suggested to further report the number of model parameters and memory costs for the EquiformerV2 and the compared baselines in Table 2. Third, how does EquiformerV2 perform on the IS2RS and IS2RE using the direct setting (use or not the denoising setting) that is the same as EquiformerV1? Overall, the major weakness of this work lies in the novelty, unclear motivations, and incomplete analyses. If the authors could well address the above concerns, I would like to increase my scores.

Questions

1. Towards the claim "Higher degrees can better capture angular resolution and directional information, which is critical to accurate prediction of atomic energies and forces". [Lines 32-33, also see Lines 97-98], could you further provide explanations and evidence from deep learning molecular models? 2. How is the S^2 activation implemented? Is there any sampling process inducing randomness or incompleteness to make the module not strictly equivariant? If so, is there any measurement of such approximated equivariance? 3. Could you briefly introduce how the EquiformerV2 model accelerates the DFT calculation in AdsorbML? Are there further quantitative results on the error measure (e.g., MAE) between the EquiformerV2 predictions and DFT labels?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors carefully discuss the broader impact and limitations of this work.

Summary

This paper propose EquiformerV2, a advance verison based on Equiformer and eSCN structure extend to higher degree representations, which achieve better performance in force and energy tasks.

Strengths

This paper is well-written and organized, presenting a clear and coherent structure throughout. The introduced EquiformerV2, an upgraded version of the original Equiformer, including three architectural improvements: attention re-normalization, separable S^2 activation, and separable layer normalization. These enhancements contribute to the SOTA performance in OC20 dataset. The proposed model achieves a high degree representation with efficiency, as outlined in the paper. The authors also provide a comprehensive ablation study to support the necessity of these modifications, effectively highlighting their respective contributions to the overall performance improvement.

Weaknesses

* A few spelling errors. For instance, in Section 6, the word "acknolwdge" etc. Along with any other mistakes found throughout the manuscript. * The experiments conducted in this study primarily utilize the OC20 dataset. While this dataset is relevant, it is essential to note that there are various other DFT-based datasets available that could provide a more comprehensive evaluation of the proposed architecture. Such as OC22, OQMD[1,2], SPICE[3], and PCQM4Mv2 etc. [1] Saal, J. E., Kirklin, S., Aykol, M., Meredig, B., and Wolverton, C. "Materials Design and Discovery with High-Throughput Density Functional Theory: The Open Quantum Materials Database (OQMD)", JOM 65, 1501-1509 (2013). [2]. Kirklin, S., Saal, J.E., Meredig, B., Thompson, A., Doak, J.W., Aykol, M., Rühl, S. and Wolverton, C. "The Open Quantum Materials Database (OQMD): assessing the accuracy of DFT formation energies", npj Computational Materials 1, 15010 (2015). [3] SPICE, A Dataset of Drug-like Molecules and Peptides for Training Machine Learning Potentials

Questions

* Would it be possible for you to include the results from the small dataset, as you mentioned? This would provide valuable insights into the performance and scalability of your proposed approach. * More details about structural relaxations, specifically, the relax trajectories and time efficiency.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

same above.

Thanks for the response

Thanks for the response -- I've increased my score to a 6.

Updated Results of EquiformerV2 on QM9 Dataset and Overall Comparison to Equiformer(V1)

> 1. Updated results of EquiformerV2 on QM9 dataset. We provide the updated results of EquiformerV2 on all 12 tasks of the QM9 dataset below. EquiformerV2 achieves **better results** than Equiformer(V1) **on 9 out of the 12 tasks**, same result on 1 task and slightly worse on the other 2 tasks. | | $\alpha$ | $\Delta \varepsilon$ | $\varepsilon_{HOMO}$ | $\varepsilon_{LUMO}$ | $\mu$ | $C_{\nu}$ | $U_0$ | $U$ | $H$ | $G$ | $R^2$ | $ZPVE$ | |----------------|----------|----------------------|----------------------|----------------------|----------|-----------|----------|----------|----------|----------|----------|----------| | Equiformer(V1) | **.046** | 30 | 15 | 14 | .011 | **.023** | 6.59 | 6.74 | 6.63 | 7.63 | .251 | **1.26** | | EquiformerV2 | .050 | **29** | **14** | **13** | **.010** | **.023** | **6.17** | **6.49** | **6.22** | **7.57** | **.186** | 1.47 | Additionally, we note that in the PDF in our **General Response**, we have complete comparisons between Equiformer(V1) and EquiformerV2 on (1) QM9 dataset, (2) OC20 S2EF-2M dataset, (3) OC20 IS2RE dataset without IS2RS auxiliary task, and (4) OC20 IS2RE dataset with IS2RS auxiliary task. Moreover, EquiformerV2 achieves **better results** than Equiformer(V1) on **(1) QM9 dataset**, **(2) OC20 S2EF-2M dataset**, and **(4) OC20 IS2RE dataset with IS2RS auxiliary task**. We would like to repeat again that the performance gain brought by higher degrees can depend on tasks and datasets. For larger datasets like OC20 S2EF-2M dataset, higher degrees clearly improves (Table VI and Table VIII). For smaller datasets like QM9 or OC20 IS2RE, higher degrees do not always bring significant improvement, but for these cases, we can pre-train on larger datasets or incorporate auxiliary tasks [1] or denoising objectives [2] to leverage the better expressivity (see the section of limitations (Line 347-Line 352)). Indeed, we show that for OC20 IS2RE dataset, using higher degrees results in overfitting, but when we use IS2RS auxiliary task, EquiformerV2 clearly improves upon Equiformer(V1) as in Table V. [1] Godwin et al. Simple GNN Regularisation for 3D Molecular Property Prediction and Beyond. ICLR 2022. [2] Zaidi et al. Pre-training via Denoising for Molecular Property Prediction. ICLR 2023.