Paper: GeoMFormer: A General Architecture for Geometric Molecular Representation Learning (ICLR-2024)

Reviewer comment (Evaluation):
The model achieves strong performance in many benchmarks but also contains more # parameters than most baseline models.

Author response:
Thank you for the suggestion. We have carefully followed your advice and conducted preliminary experiments with our GeoMFormer of different sizes on the OC20 IS2RE task. To investigate the scalibility of our GeoMFormer, we vary the number of layers in [6, 12, 18] and keep all other hyperparameters the same for a fair comparison. The results are shown in Table 1. 

**Table 1. Performance of GeoMFormer of different sizes on OC20 IS2RE.**


| #Layer | ID MAE$\downarrow$ | OOD Ads. MAE$\downarrow$ | **OOD Cat. MAE$\downarrow$** | **OOD Both MAE$\downarrow$** | **Average MAE$\downarrow$** | **ID EwT$\uparrow$** | OOD Ads. EwT$\uparrow$ | OOD Cat. EwT$\uparrow$ | OOD Both EwT$\uparrow$ | **Average EwT$\uparrow$** |
| :----: | :----------------: | :----------------------: | :--------------------------: | :--------------------------: | :-------------------------: | :------------------: | :--------------------: | :--------------------: | :--------------------: | :-----------------------: |
|   6    |       0.3986       |          0.4771          |            0.4164            |            0.4266            |           0.4297            |        10.41         |          6.63          |          9.32          |          6.30          |           8.17            |
|   12   |       0.3883       |          0.4562          |            0.4037            |            0.4083            |           0.4141            |        11.26         |          6.70          |          9.97          |          6.42          |           8.59            |



From Table 1, we can see that scaling GeoMFormer up brings consistent performance gains on the challenging OC20 IS2RE task: (1) Average MAE: 3.6% relative improvements; (2) Average EwT: 5.1% relative improvements. Moreover, due to the time limits and computational resource constraints, the experiments on our 18-layer GeoMFormer are still running, while we have observed gains comparing its validation curve with the 12-layer GeoMFormer. We will definitely follow your suggestion to conduct further experiments on the effects of the scaling up on other tasks and include these results in our paper as supporting evidence on the scalibility of our GeoMFormer.

Reviewer comment (Experiments):
Following the weakness, training GeMFormer of different sizes can be useful to elaborate how much performance gain stems from scaling up # parameters. Since the proposed method designs an equivariant Transformer that can be better scaled, such experiments can further demonstrate the advantage of GeMFormer.

Author response:
Thank you for the suggestion. We have carefully followed your advice and conducted preliminary experiments with our GeoMFormer of different sizes on the OC20 IS2RE task. To investigate the scalibility of our GeoMFormer, we vary the number of layers in [6, 12, 18] and keep all other hyperparameters the same for a fair comparison. The results are shown in Table 1. 

**Table 1. Performance of GeoMFormer of different sizes on OC20 IS2RE.**


| #Layer | ID MAE$\downarrow$ | OOD Ads. MAE$\downarrow$ | **OOD Cat. MAE$\downarrow$** | **OOD Both MAE$\downarrow$** | **Average MAE$\downarrow$** | **ID EwT$\uparrow$** | OOD Ads. EwT$\uparrow$ | OOD Cat. EwT$\uparrow$ | OOD Both EwT$\uparrow$ | **Average EwT$\uparrow$** |
| :----: | :----------------: | :----------------------: | :--------------------------: | :--------------------------: | :-------------------------: | :------------------: | :--------------------: | :--------------------: | :--------------------: | :-----------------------: |
|   6    |       0.3986       |          0.4771          |            0.4164            |            0.4266            |           0.4297            |        10.41         |          6.63          |          9.32          |          6.30          |           8.17            |
|   12   |       0.3883       |          0.4562          |            0.4037            |            0.4083            |           0.4141            |        11.26         |          6.70          |          9.97          |          6.42          |           8.59            |



From Table 1, we can see that scaling GeoMFormer up brings consistent performance gains on the challenging OC20 IS2RE task: (1) Average MAE: 3.6% relative improvements; (2) Average EwT: 5.1% relative improvements. Moreover, due to the time limits and computational resource constraints, the experiments on our 18-layer GeoMFormer are still running, while we have observed gains comparing its validation curve with the 12-layer GeoMFormer. We will definitely follow your suggestion to conduct further experiments on the effects of the scaling up on other tasks and include these results in our paper as supporting evidence on the scalibility of our GeoMFormer.

Reviewer comment (Reproducibility):
Is cutoff of interatomic distances applied when implementing attention between atoms in GeMFormer? If yes, how is the cutoff determined?

Author response:
In our experiments, we do not apply the cutoff of interatomic distances, which means that each atom can attend to all other atoms. It is a good point when we deal with large molecular systems, and our model can naturally adopt the cutoff of interatomic distances when necessary by (1) set a cutoff of interatomic distances; (2) use this cutoff to determine the neighborhood of each atom; (3) generate attention masks according to the neighborhood of each atom.

Reviewer comment (Experiments):
The work has included benchmarks on multiple standard molecular property prediction tasks. How does the model perform on QM9. Though it's a bit older compared with some benchmarks, it's still widely used. It would be helpful to see where GeMFormer sits on QM9 as well.

Author response:
Thank you for your suggestion. We also agree with you on the popularity of the QM9 dataset. This dataset offers quantum chemical properties (e.g., HOMO-LUMO gap, total energy, etc) at Density Functional Theory level for a relevant, consistent, and comprehensive chemical space of 134k small organic molecules with equilibrium structures. Therefore, the tasks of the QM9 dataset are invariant prediction tasks. In our experiments, we have actually adopted several popular datasets that evaluate molecular model in a similar way. For example, the Molecule3D dataset consists of around 3.8 million molecules in total and the task is to predict the HOMO-LUMO energy gap given the equilibrium structures, which is similar to the QM9 task and is much larger in terms of scales. Additionally, the experiments on MD17, which is another fairly popular dataset, also examine the model's ability to perform the energy prediction task. Moreover, we also follow your suggestion to further conduct experiments on the QM9 dataset. Due to the time and computational resources limits, we will update the results on the QM9 dataset to further improve the quality of our work.

Reviewer comment (Experiments):
In appendix D.7, the authors present ablation studies on N-body simulation task. However, as the work focuses mostly on molecular property predictions, it would be better to include some ablations on molecular benchmarks to further validate the design choices.

Author response:
Thank you for the suggestion. We have carefully followed your advice and conducted additional ablation study on the MD17 dataset, which is widely used to verify the effectiveness of molecular models. We present our results in Table 2. The results further demonstrate the effectiveness of our design choice. Specifically, removing any of the attention module consistently induces performance drop. In addition, the gains brought by cross-attn modules are extremely significant for both energy prediction (invariant) and force prediction (equivariant): 20.8% relative improvement for energy and 60.8% relative improvement for force gaining form adding the Inv-Cross-Attn and Equ-Cross-Attn, which highlights the effectiveness of the key idea of our design philosophy. We will update this new additional ablation study in the revised version of our paper.

Reviewer comment (Experiments):
I'm generally satisfied with the extensive experiments, but it would be better to benchmark GeoMFormer on QM9 due to its popularity.

Author response:
Thank you for the suggestion. We also agree with you that the QM9 dataset has been adopted as a popular benchmark to evaluate the effectiveness of geometric molecular models. This dataset provides quantum chemical properties (e.g., HOMO-LUMO gap, total energy, etc) at Density Functional Theory level for a relevant, consistent, and comprehensive chemical space of 134k small organic molecules with equilibrium structures. Thus, the tasks of the QM9 dataset are invariant prediction tasks. In our experiments, we actually adopted several popular datasets that evaluate molecular model in a similar way. For example, the Molecule3D dataset consists of around 3.8 million molecules in total and requires the model to predict the HOMO-LUMO energy gap given the equilibrium structures, which is similar to the QM9 task and is much larger. Besides, the experiments on MD17 also examine the model's ability to perform the energy prediction task, which is also very popular. Moreover, we also follow your advice to further conduct experiments on the QM9 dataset. Due to the time and computational resources limits, we will update the results on the QM9 dataset to further improve the quality of our work.

Reviewer comment (Experiments):
For experiments on PCQM4MV2, it seems only GeoMFormer and Uni-Mol+ are directly comparable, as both of them essentially utilize more supervision information.

Author response:
We would like to clarify the setting of the PCQM4Mv2 dataset. In the training set of the PCQM4Mv2 dataset, each molecule is associated with both its 2D graph structure and 3D geometric structure obtained by using DFT-level relaxation for equilibrium states. In the validation set, each molecule is associated with only its 2D graph structure. It is a reasonable setting because the DFT-level relaxation is extremely expensive. In real-world applications, if a deep learning model need to take expensive equilibrium structures as the input, it has no more advantages compared to classical methods when we need to perform prediction on new molecules. Different from previous baselines focusing on graph learning, recent works including Transformer-M, Uni-Mol+ have noticed the value of these 3D equilibrium structures provided in the training time and pushed the frontier of this challenging task. In our experiment, we follow the training strategy of Uni-Mol+. Compared to these strong baselines, our framework is demonstrated to be effective by the empirical results and remains efficient compared to the sophisticated Uni-Mol+ architecture. From these points, we believe the PCQM4Mv2 task is a good evaluation protocal.

Reviewer comment (Experiments):
In the experiment, I wonder why different invariant tasks use different neural architectures as baseline methods, e.g., Section 5.2 and 5.3. The same thing happens for equivariant tasks.

Author response:
To ensure a fair comparison, we follow previous works to choose strong baselines for each dataset and reported the official results from the original paper if available for most of the experiments. In fact, few previous works have benchmarked their proposed model across all the extensive and diverse datasets presented in our paper. Due to this, we have different baselines in different tasks like Section 5.2 and 5.3. This phenomenon further demonstrates the generality and effectiveness of our framework, as our GeoMFormer is able to simultaneously perform well on a wide range of tasks involving both invariant and equivariant prediction and consistently achieves strong performance compared to these specialized baselines.

Reviewer comment (Evaluation):
I wonder if the author could report the results of multiple runs, e.g., standard deviation to enhance the empirical results.

Author response:
Thank you for the suggestion. Tasks in our experiments can be divided into two categories according to their dataset sizes: (1) large-scale datasets, including OC20, PCQM4Mv2 and Molecule3D; (2) small-scale datasets, including N-Body simulation and MD17. For large-scale datasets, we observe that the results are rather stable across different runs, i.e., the standard deviation is extremely small. For small-scale datasets, we actually follow the convention of previous works to perform multiple runs for fair comparisons. For N-body simulation, we perform five runs with different random seeds and report the averaged results across these five runs. For MD17, we averaged the results over three random splits. Following the convention of previous works, we report the averaged results for both the datasets.

Reviewer comment (Experiments):
In N-BODY SIMULATION in Section 5.4, what are the input feature and groundtruth? input is 3D position, output is velocity? Please add more details. Thanks!

Author response:
Thank you for the suggestion. We provide more details here. In the N-Body simulation experiment, each data point is a particle system, which consists of 5 particles that carry a positive or negative charge and have an initial position and an initial velocity. The system is controlled by physical rules involving attractive and repulsive forces. The task requires the model to predict the positions of the five particles after 1,000 timesteps. Specifically, the input feature includes (1) the initial positions of the five particles $p^0=\{p^0_1,p^0_2,\cdots,p^0_5\}\in\mathbb{R}^{5\times 3}$ , (2) the initial velocities of five particles $v^0=\{v^0_1,v^0_2,\cdots,v^0_5\}\in\mathbb{R}^{5\times 3}$ and (3) the charges (positive or negative) of five particles $c=\{c_1,c_2,\cdots,c_5\}\in\{-1,1\}^5$. The output is the predicted positions of the five particles after 1,000 timesteps $\hat{p}^t=\{\hat{p}^t_1,\hat{p}^t_2,\cdots,\hat{p}^t_5\}\in\mathbb{R}^{5\times 3}$ and the groundtruth is the true positions of the five particles after 1,000 timesteps $p^t=\{p^t_1,p^t_2,\cdots,p^t_5\}\in\mathbb{R}^{5\times 3}$ . we will incorporate additional details in the revised version of the paper.

Reviewer comment (Theory):
The intuition and rationale of modeling is not clear. There is a lack of analysis of the difference between an invariant representation and an equivariant representation in terms of the information they contain.

Author response:
Thank you for spending time reviewing our paper. Here are our responses to your questions:

> **Provide analysis of the difference between invariant and equivariant representations in terms of the information they contain (W1 & Q1)**

In the Introduction section of our paper, we actually provide discussion of the strengths/weaknesses between modeling invariant and equivariant representations, which demonstrates the necessity of developing a powerful model that can perform both invariant and equivariant prediction with strong performance at the same time. Here we would like to further elucidate the disparity between invariant and equivariant representations in terms of the information they encapsulate, substantiated by additional supporting evidence.

From the view of neural network design, equivariant representations preserve more complete structural information, but the equivariant constraints restrict the choice of operations that can be used. On the other hand, invariant representations offer more flexibility for non-linear operations, but retain less structural information.

First, we provide more explantions and supporting evidence on why equivariant representations preserve more complete structural information. In the literature, existing invariant models commonly use invariant features of molecules (e.g., interatomic distances, bond angles, dihedral/torsion angles, improper angles and so on) which are extracted from the raw 3D geometric structures of molecules (the position of each atom in the molecule). Although these invariant features are important to describe the properties of molecules, the mapping from raw geometric structures to invariant features still introduces information loss:
- In a recent study on the expressive power of geometric graph neural networks [1], the authors provided thorough theoretical analysis on the expressive power of invariant layers and equivariant layers commonly used in geometric graph neural networks by leverging the Geometric Weisfeiler-Lehman test. Each invariant layer in geometric graph neural networks corresponds to the proposed Invariant Geometric Weisfeiler-Lehman test (IGWL). In their framework, previous mentioned invariant features can be categorized into different classes in terms of the least number of nodes (atoms in the context of molecular modeling) that are involved to compute it. For example, interatomic distances involve two atoms, while bond angles involve three atoms. This features are then called as $k$-body scalars. Based on this framework, the authors carefully characterized the expressive power of invariant layers and prove that "any number of iterations of IGWL cannot distinguish any 1-hop identical geometric graphs $\mathcal{G}_1$ and $\mathcal{G}_2$ where the underlying attributed graphs are isomorphic" and "GWL (corresponds to equivariant layer) can distinguish by propagating information from the geometric structure", which indeed indicates the information loss from raw geometric structures to invariant features. 
- In PaiNN [2], the authors also provided intuition on the point that there exist limits of invariant representations. Firstly, the authors demonstrated that to express enough information (like bond angles of all neighbors of an atom), using invariant representations would requires computations with higher complexity than using equivariant representations (see the analysis of Table 1 in [2]). Moreover, the authors also showed that equivariant representations allow to propagate crucial geometrical information beyond the neighborhood, which is not possible in the invariant case (see the analysis of Figure 1 in [2]). These evidence motivated the authors to develop models based on equivariant representations.

Reviewer comment (Evaluation):
The experimental results cannot support the advantage of fusing two kinds of representation. The ablation studies in Table 8 are interesting. The model without any cross-attn can still achieve competitive performance. The cross-attn module, the main contribution of this paper, may contribute little to performance in experiments.

Author response:
> **Regarding the ablation study (W3 & Q2)**

We would like to clarify the significance of our ablation study (Table 8) with updated experiment results. As can be seen from the baseline performance of the N-Body simulation task, the absolute number of the MSE criterion is very low, which means that we should focus on the relative error reduction when comparing the performance of two models. In Table 8, we actually have the following conclusions on the ablation study: (1) removing any cross-attn module (Inv-Cross-Attn, Equ-Cross-Attn) consistently induces performance drop; (2) the gains brought by cross-attn modules are significant: 7.8% relative improvement for Inv-Cross-Attn, 13.0% relative improvement for Equ-Cross-Attn, 17.5% relative improvement for both Inv-Cross-Attn and Equ-Cross-Attn.

Moreover, we further conducted ablation study on the MD17 benchmark as a double-check. The results are presented in the following table. From these results, we can see the same trend as in the N-Body simulation task: (1) removing any cross-attn module (Inv-Cross-Attn, Equ-Cross-Attn) consistently induces performance drop; (2) the gains brought by cross-attn modules are extremely significant for energy prediction (invariant): 18.7% relative improvement for Inv-Cross-Attn, 9.8% relative improvement for Equ-Cross-Attn, 20.8% relative improvement for both Inv-Cross-Attn and Equ-Cross-Attn; (3) the gains brought by cross-attn modules are also extremely significant for force prediction (equivariant): 28.0% relative improvement for Inv-Cross-Attn, 43.9% relative improvement for Equ-Cross-Attn, 60.8% relative improvement for both Inv-Cross-Attn and Equ-Cross-Attn. 

Based on these thorough results, we indeed demonstrate the significance and effectiveness of our framework, which is strongly supported and should be definitely taken into account.

Reviewer comment (Experiments):
Can you give empirical evidence that invariant and equivariant representations have different useful information for molecular learning?

Author response:
Thank you for spending time reviewing our paper. Here are our responses to your questions:

> **Provide analysis of the difference between invariant and equivariant representations in terms of the information they contain (W1 & Q1)**

In the Introduction section of our paper, we actually provide discussion of the strengths/weaknesses between modeling invariant and equivariant representations, which demonstrates the necessity of developing a powerful model that can perform both invariant and equivariant prediction with strong performance at the same time. Here we would like to further elucidate the disparity between invariant and equivariant representations in terms of the information they encapsulate, substantiated by additional supporting evidence.

From the view of neural network design, equivariant representations preserve more complete structural information, but the equivariant constraints restrict the choice of operations that can be used. On the other hand, invariant representations offer more flexibility for non-linear operations, but retain less structural information.

First, we provide more explantions and supporting evidence on why equivariant representations preserve more complete structural information. In the literature, existing invariant models commonly use invariant features of molecules (e.g., interatomic distances, bond angles, dihedral/torsion angles, improper angles and so on) which are extracted from the raw 3D geometric structures of molecules (the position of each atom in the molecule). Although these invariant features are important to describe the properties of molecules, the mapping from raw geometric structures to invariant features still introduces information loss:
- In a recent study on the expressive power of geometric graph neural networks [1], the authors provided thorough theoretical analysis on the expressive power of invariant layers and equivariant layers commonly used in geometric graph neural networks by leverging the Geometric Weisfeiler-Lehman test. Each invariant layer in geometric graph neural networks corresponds to the proposed Invariant Geometric Weisfeiler-Lehman test (IGWL). In their framework, previous mentioned invariant features can be categorized into different classes in terms of the least number of nodes (atoms in the context of molecular modeling) that are involved to compute it. For example, interatomic distances involve two atoms, while bond angles involve three atoms. This features are then called as $k$-body scalars. Based on this framework, the authors carefully characterized the expressive power of invariant layers and prove that "any number of iterations of IGWL cannot distinguish any 1-hop identical geometric graphs $\mathcal{G}_1$ and $\mathcal{G}_2$ where the underlying attributed graphs are isomorphic" and "GWL (corresponds to equivariant layer) can distinguish by propagating information from the geometric structure", which indeed indicates the information loss from raw geometric structures to invariant features. 
- In PaiNN [2], the authors also provided intuition on the point that there exist limits of invariant representations. Firstly, the authors demonstrated that to express enough information (like bond angles of all neighbors of an atom), using invariant representations would requires computations with higher complexity than using equivariant representations (see the analysis of Table 1 in [2]). Moreover, the authors also showed that equivariant representations allow to propagate crucial geometrical information beyond the neighborhood, which is not possible in the invariant case (see the analysis of Figure 1 in [2]). These evidence motivated the authors to develop models based on equivariant representations.

Reviewer comment (Evaluation):
In Table 8, the model without cross-attn can still achieve competitive performance. From this point, is it very important to fuse two kinds of representation on such a powerful Transfermor model?

Author response:
> **Regarding the ablation study (W3 & Q2)**

We would like to clarify the significance of our ablation study (Table 8) with updated experiment results. As can be seen from the baseline performance of the N-Body simulation task, the absolute number of the MSE criterion is very low, which means that we should focus on the relative error reduction when comparing the performance of two models. In Table 8, we actually have the following conclusions on the ablation study: (1) removing any cross-attn module (Inv-Cross-Attn, Equ-Cross-Attn) consistently induces performance drop; (2) the gains brought by cross-attn modules are significant: 7.8% relative improvement for Inv-Cross-Attn, 13.0% relative improvement for Equ-Cross-Attn, 17.5% relative improvement for both Inv-Cross-Attn and Equ-Cross-Attn.

Moreover, we further conducted ablation study on the MD17 benchmark as a double-check. The results are presented in the following table. From these results, we can see the same trend as in the N-Body simulation task: (1) removing any cross-attn module (Inv-Cross-Attn, Equ-Cross-Attn) consistently induces performance drop; (2) the gains brought by cross-attn modules are extremely significant for energy prediction (invariant): 18.7% relative improvement for Inv-Cross-Attn, 9.8% relative improvement for Equ-Cross-Attn, 20.8% relative improvement for both Inv-Cross-Attn and Equ-Cross-Attn; (3) the gains brought by cross-attn modules are also extremely significant for force prediction (equivariant): 28.0% relative improvement for Inv-Cross-Attn, 43.9% relative improvement for Equ-Cross-Attn, 60.8% relative improvement for both Inv-Cross-Attn and Equ-Cross-Attn. 

Based on these thorough results, we indeed demonstrate the significance and effectiveness of our framework, which is strongly supported and should be definitely taken into account.