Paper: Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning (ICLR-2024)

Reviewer comment (Experiments):
The paper includes comparisons with several models, such as SchNet, DimeNet++, EGNN, and ET. However, there are other models that exhibit superior performance, including GemNet, NeurIPS, Allegro, MACE, etc. The paper would benefit from a more comprehensive comparison with existing methods. While the experimental results demonstrate the effectiveness of the proposed method, a more detailed comparison with other state-of-the-art models would strengthen the paper's claims. Additionally, it is important to cite these related works, as they are highly relevant to the topic of equivariant GNNs.

Author response:
Thanks for the valuable suggestion! We appreciate your recommendation to conduct a more comprehensive comparison with existing state-of-the-art models. In response to your feedback, we have augmented our experimental comparisons by including four strong baselines, namely GemNet, Equiformer, MACE, and LEFTNet. MACE was selected due to its superior performance among the suggested baselines according to its paper [A]. Our GET consistently outperforms these additional strong baselines, reaffirming its effectiveness and superiority. We have also expanded the discussion in the related works section to provide readers with a comprehensive understanding of the landscape of existing methods and emphasize the unique contributions and strengths of our proposed approach.


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

Reviewer comment (Writing):
Typo: The model compared with GET is TorchMDNet (ET), not TorchMD.

Author response:
We are sorry for our choice of ambiguous abbreviation. We have corrected it in the revision.

Reviewer comment (Novelty):
Equation (3) is not well defined (the MLP/RBF used); it seems closely related to Shnet or Physnet [3]. Also, learning/refining the self-attention has been investigated (e.g. [1]).

Author response:
> Q1: Equation (3) is not well defined (the MLP/RBF used) it seems closely related to Shnet or Physnet [3].

We apologize for the unclear presentation. The MLP is a 2-layer multi-layer perceptron with SiLU [A] activation. The RBF function embeds distances into $d_{\text{rbf}}$-dimensional vectors with the following definition:

$$\text{RBF}(d)[k] = u(\frac{d}{c})* \exp(-\frac{|d-\mu_k|}{2 \gamma}), 1\leq k \leq d_{\text{rbf}}$$

where $\mu_k$ is uniformly distributed in $[0, c]$($c=7.0$ in our paper), $\gamma = \mu_k - \mu_{k-1} = \frac{c}{d_{\text{rbf}}}$, and $u(x)$ is the polynomial envelope function:

$$u(x) = 1 - \frac{(p+1)(p+2)}{2} x^p + p(p+2)x^{p+1} - \frac{p(p+1)}{2}x^{p+2}$$

We set $p=5$ in our paper. The reason why we apply the RBF function is inspired by [A][B] for its promising performance. We have added the definition to the revision.

[A] Schütt, K., Kindermans, P. J., Sauceda Felix, H. E., Chmiela, S., Tkatchenko, A., & Müller, K. R. (2017). Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30.

[B] Gasteiger, J., Groß, J., & Günnemann, S. (2020). Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123.

> Q2: Also, learning/refining the self-attention has been investigated (e.g. [1]).
> 
> [1] Geometric transformer for end-to-end molecule properties prediction.

Thanks for the comment. We would like to point out that our GET is directly implemented on 3D coordinates with equivariance, while [1] actually uses pairwise distance matrix to represent the geometry to achieve invariance. Moreover, we inject the inductive bias of custom hierarchy into message passing, while [1] is implemented on atomic level only. We have updated the related work with more discussion on the difference between our work and previous literature on geometric transformers/GNNs.

Reviewer comment (Novelty):
It is unclear how equation (5) is different from Jin et al.'s pooling.

Author response:
> Q3: Unclear how (5) is different from (Jin and al.)'s pooling.

We are sorry we did not express the difference clear enough. Equation (5) is the pooling on atomic pairwise relations to obtain the relation between two blocks, while (Jin and al.)'s pooling is implemented on atoms within each block to obtain the block-level representations. In particular, (Jin and al.)'s method first pools atom-level features into residue-level features, and then conducts message passing on the residue-level graphs. In constrast, our bilevel method retains the information of both atom-level and residue-level features simultaneously in each layer, namely, allowing both attention-based message passing between the atoms winthin the same residue and the neighbor residues.

Reviewer comment (Theory):
There is some lack of motivation as to why one needs to keep track of the coordinates (which requires all the equivariant design of the non-self-attention layers) rather than working with equivariant measures such as the interatomic pairwise distances as most works do, and using other metrics too (e.g. C-RMSD [3,4]).

Author response:
> Q4: There is some lack of motivation as to why one needs to keep track of the coordinates (which requires all the equivariant design of the non-self-attention layers) rather than working with equivariant measures such as the interatomic pairwise distances as most works do (and using other metrics too, e.g. C-RMSD [3,4]).
> 
> [3] Molecular geometry prediction using a deep generative graph neural network
> 
> [4] Diffusion-based molecule generation with informative prior bridges

Thank you for raising the insightful question regarding equivariant designs. We would like to answer this question from two aspects: theoretical expressivity and practical efficiency.

**Theoretical Expressivity:**

Existing theoretical work [A] (Theorem 8) has established that equivariant methods, which involve keeping track of coordinates, offer greater expressivity compared to invariant methods relying on pairwise distances (please see Theorem 8 in [A]). Specifically, the study reveals that equivariant Graph Neural Networks (GNNs) share expressivity with the Geometric Weisfeiler-Leman Test (GWL), while invariant counterparts align with the Invariant GWL (IGWL). The crucial distinction emerges in scenarios commonly encountered in molecular representation learning, where detecting isomorphic geometric graphs pose a challenge for invariant methods due to the lack of directional geometry in pairwise distances. In contrast, equivariant counterparts, such as our proposed GET, effectively handle these cases, enhancing the model's capacity to discern intricate geometric relationships within molecular structures.

**Practical Efficiency:**

Furthermore, the practical efficiency of our approach is necessary for addressing the challenges associated with scalability in molecular interaction graphs. Fully connected graphs, often implemented in invariant methods, exhibit $O(n^2)$ complexity ($n$ being the number of nodes), making them less viable for large graphs. Most existing methods relying on pairwise distances use $n \times n$ distance matrices, limiting their application to small molecular graphs with fewer than 50 nodes [B][C]. In our experiments, molecular interaction graphs are considerably larger, ranging from 230 to 872 atom nodes on average, as they incorporate extensive biomolecular structures. The equivariant nature of our GET, with linear complexity relative to the number of nodes, ensures its scalability to these larger graphs in practical settings.

We hope this clarification provides a thorough understanding of our motivation for choosing equivariant designs in GET.

[A] Joshi, C.K., Bodnar, C., Mathis, S.V., Cohen, T. &amp; Lio, P.. (2023). On the Expressive Power of Geometric Graph Neural Networks. Proceedings of the 40th International Conference on Machine Learning, in Proceedings of Machine Learning Research 202:15330-15355 Available from https://proceedings.mlr.press/v202/joshi23a.html.

[B] Molecular geometry prediction using a deep generative graph neural network

[C] Diffusion-based molecule generation with informative prior bridges

Reviewer comment (Experiments):
In order to have a good comparison with the baselines, one should add the models' complexity and capacity comparison.

Author response:
> Q7: In order to have a good comparison with the baselines, one should add the models' complexity and capacity comparison. Thus, while comparing to old baselines such as Shnet or Dimenet++ (Tables 2 and 3), one should compare the difference in capacity and complexity of the models which still perform well and should be much more efficient.

Thank you for the valuable suggestion!

In our revised manuscript, we have thoughtfully addressed the comparison of model complexity and capacity. Detailed information on the number of parameters and training times for all baselines, both old and new, is available in Table 9 in Appendix G.

When comparing GET to simpler yet weaker baselines (e.g., SchNet and EGNN), it is evident that GET may have more parameters and a slower training speed. However, it's crucial to note that GET exhibits competitive parameter and computation efficiency when compared with more complex yet stronger baselines, such as Equiformer, MACE, and LEFTNet. It's worth mentioning that a significant portion of parameters in GET is attributed to the Feedforward Neural Network (FFN) that projects latent features to higher dimensions in intermediate layers, aligning with the structure of vanilla Transformers. Without this part, GET has the least parameters among all the models.

Moreover, the throughput of GET is comparable to both atom-level and hierarchical counterparts. This aligns with the model's design, which considers both block-level and atom-level geometry. To further elucidate the efficiency of our approach, we provide a complexity analysis of GET in Appendix D, showing its linear complexity concerning the number of nodes in the graph. This characteristic indicates favorable scalability to large graphs in practical settings.

In conclusion, our model maintains a moderate number of parameters and efficient training speed while consistently outperforming baselines across multiple evaluation experiments. We believe these detailed analyses contribute to a comprehensive understanding of the efficiency and effectiveness of our proposed model.

Reviewer comment (Experiments):
While comparing to old baselines such as Shnet or DimeNet++ (Tables 2 and 3), one should compare the difference in capacity and complexity of the models, which still perform well and should be much more efficient.

Author response:
> Q7: In order to have a good comparison with the baselines, one should add the models' complexity and capacity comparison. Thus, while comparing to old baselines such as Shnet or Dimenet++ (Tables 2 and 3), one should compare the difference in capacity and complexity of the models which still perform well and should be much more efficient.

Thank you for the valuable suggestion!

In our revised manuscript, we have thoughtfully addressed the comparison of model complexity and capacity. Detailed information on the number of parameters and training times for all baselines, both old and new, is available in Table 9 in Appendix G.

When comparing GET to simpler yet weaker baselines (e.g., SchNet and EGNN), it is evident that GET may have more parameters and a slower training speed. However, it's crucial to note that GET exhibits competitive parameter and computation efficiency when compared with more complex yet stronger baselines, such as Equiformer, MACE, and LEFTNet. It's worth mentioning that a significant portion of parameters in GET is attributed to the Feedforward Neural Network (FFN) that projects latent features to higher dimensions in intermediate layers, aligning with the structure of vanilla Transformers. Without this part, GET has the least parameters among all the models.

Moreover, the throughput of GET is comparable to both atom-level and hierarchical counterparts. This aligns with the model's design, which considers both block-level and atom-level geometry. To further elucidate the efficiency of our approach, we provide a complexity analysis of GET in Appendix D, showing its linear complexity concerning the number of nodes in the graph. This characteristic indicates favorable scalability to large graphs in practical settings.

In conclusion, our model maintains a moderate number of parameters and efficient training speed while consistently outperforming baselines across multiple evaluation experiments. We believe these detailed analyses contribute to a comprehensive understanding of the efficiency and effectiveness of our proposed model.

Reviewer comment (Reproducibility):
How to construct the bilevel representation should be elaborated (e.g., what is the K of the KNN graph?).

Author response:
We apologize for the missing details. Each node in the graph consists of a set of atoms in the block. Suppose there are $n_i$ atoms in block $i$, then it is assigned with a feature matrix $\mathbf{H}_i \in \mathbb{R}^{n_i \times d}$ and a coordinate matrix $\vec{\mathbf{X}}_i \in \mathbb{R}^{n_i \times 3}$, where each row corresponds to one atom. And edges are deduced from the k-nearest neighbors of each node, where distance is defined as the minimum pairwise distances between atoms in two blocks. In this paper we use $K=9$ which is enough for supporting the competitive performance of our model. Overall, the topology of the graph is defined on the block-level geometry while atom-level instances are maintained via extensions of node features from single vectors to dynamic matrices. We have updated the manuscript for a clearer presentation in Section 3.1.

Reviewer comment (Writing):
Typo: (In Sec. 1) "In this paper, we approache ..."

Author response:
Thanks for pointing this out. We have fixed the typo in the revision.

Reviewer comment (Experiments):
Compared to KNN, are there better ways to construct a block?

Author response:
Thanks for the comment. The most common ways to construct the graph are KNN graph, radius graph (adding edges by distance cutoff), and complete graph (each node connects to all other nodes). Practically we choose KNN graph because its complexity (i.e. Number of edges) is linear to the number of nodes, which shows desired scalability to larger graphs used in our paper. It is worth mentioning that achieving good balance between complexity and performance through better construction of the graph is also a valuable research question. Existing literature has been exploring combination of different constructions [A]. This might also be related to sparse attention [B] considering that the topology of graphs can be represented as adjacent matrices.

[A] Zhang, Z., Xu, M., Jamasb, A., Chenthamarakshan, V., Lozano, A., Das, P., & Tang, J. (2022). Protein representation learning by geometric structure pretraining. arXiv preprint arXiv:2203.06125.

[B] Tay, Y., Dehghani, M., Bahri, D., & Metzler, D. Efficient transformers: A survey. arXiv 2020. arXiv preprint arXiv:2009.06732.

Reviewer comment (Experiments):
The reliance on domain-specific knowledge for the construction of building blocks may limit the universality of the proposed GET model (e.g., each residue in the proteins is one node). This may render the comparison with hierarchical models (e.g., GVP) a bit unfair.

Author response:
Thank you for the comment. There could be some misunderstandings here. We provide the explanations below.

First, we would like to emphasize that our GET model can accommodate arbitrary construction of building blocks, using wheter domain-specific kowledge or purely data-driven approaches. To show this universality, we did implement GET-PS in section 4.2, Table 3, which uses a data-driven method (i.e. principal subgraphs in the PS-VAE paper [A]) to build blocks in small molecules, and it indeed achieves brilliant performance. 

Second, we want to clarify that one significant strength of our GET is the capability of capturing both coarse-grained (blocks) and fine-grained (atoms) geometry. In section 4.2, for fair comparison, all the block-level baselines and the hierarchical baselines use exactly the same building blocks as our GET. However, their inability to retain bilevel geometry simultaneously makes them suboptimal compared to our proposed GET.

[A] Kong, X., Huang, W., Tan, Z., & Liu, Y. (2022). Molecule generation by principal subgraph mining and assembling. Advances in Neural Information Processing Systems, 35, 2550-2563.

Reviewer comment (Experiments):
The paper would be strengthened by a more thorough comparison with state-of-the-art methods, e.g., GemNet, Equiformer, EquiformerV2, and LEFTNet.

Author response:
Thank you for your insightful recommendation to conduct a more thorough comparison with state-of-the-art methods. We have expanded our experimental comparisons by incorporating four suggested baselines: GemNet, Equiformer, MACE (suggested by Reviewer KrxA), and LEFTNet in Table 2. Besides, owing to their commendable performance and high efficiency, MACE and LEFTNet are further selected for testing on the mix-domain data augmentation in Figure 3 and zero-shot generalization ability in Table 3. Delightfully, our GET consistently outperforms all these strong baselines in all cases, demonstrating its general effectiveness and superiority. We have updated the manuscript with the results of these additional experiments.

Reviewer comment (Experiments):
The experiments conducted on PPA and PBA datasets may not fully demonstrate the effectiveness of the proposed model, given their relatively small scale.

Author response:
Thanks for the comment.

We admit that PPA and LBA are of small scale, but these two datasets are our suitable choice to verify the effectiveness of GET in learning unified representations between different types of molecular interactions, namely, protein and protein in PPA, and protein and small molecule in LBA.  Actually, in the domain of molecular interaction, how to address the scarcity of data is a long-term problem.  While it is arduous to collect more data due to the high cost of wetlab experiments, we have tried to enhance the robustness and reliability of the empirical evaluations by running parallel experiments with different random seeds and data splits.

Besides the experiments on  PPA and LBA, we additionally evaluate the zero-shot ability of our GET on DNA/RNA-ligand binding affinity prediction in section 4.3, as this is strong evidence for the exceptional generalizability and high data efficiency of the proposed method. It indicates our model's proficiency in adapting to completely unseen data, which is particularly valuable in addressing the challenge of data scarcity.

We also want to emphasize the promising strategy outlined in our paper, which involves better utilizing all available data from different domains through the proposed unified molecular representation. This strategy aligns with the broader goal of overcoming data scarcity challenges by maximizing the information gleaned from diverse datasets.

Reviewer comment (Experiments):
I do not see a particular design of modeling molecular interactions in the GET architecture. It appears that it is also possible to evaluate the GET model on a broader range of molecular settings, especially bare molecules.

Author response:
Thanks for the suggestion! In this paper, we mainly focus on proposing a unified molecular representation for different types of molecular complexes, as well as a corresponding model for effective representation learning. The motivation behind is that a unified molecular representation is critical for modeling interaction geometry between multiple molecules, and endows the model with inspiring generalizability. While it is possible to evaluate GET on bare molecules, it is out of the main scope of this paper and be better left for further exploration. We have discussed this point in the limitation section (now in Appendix L due to the space limit) in the revision.

Reviewer comment (Evaluation):
The equivariant transformers and attention mechanisms raise potential concerns for practical application due to their high computational complexity, as noted in Table 9.

Author response:
Thanks for the comment. We appreciate your attention to this matter, and we have carefully considered your comments.

We acknowledge that equivariant transformers and attention mechanisms can introduce computational complexity. However, we would like to emphasize that our approach remains highly competitive in terms of efficiency when compared to new strong baselines (GemNet, Equiformer, MACE, and LEFTNet). In practical terms, our model offers a compelling balance between performance and computational speed, as demonstrated by the results in Table 9. Also, we have provided complexity analysis in Appendix D, which shows good scalability (linear instead of quadratic) with respect to the number of atoms in the complex graph.

Moreover, despite any inherent complexity, our approach remains significantly faster than traditional computational tools (such as rosetta) in real-world applications, supporting its practicality. 

We have updated the manuscript to reflect these clarifications and additions to Table 9. We believe that these revisions enhance the overall quality and understanding of our work.

Reviewer comment (Presentation):
Including the graphical illustration (Figure 4) within the main text could provide readers with a more immediate and clear understanding of the architecture and operational flow of GET.

Author response:
Thanks for the advice! While it might be hard to directly include figure 4 in the main text due to the space limit, we have added a reference to it at the beginning of section 3 (Method).

Reviewer comment (Experiments):
Adding statistics for the PDBBind benchmark within the paper would offer a more concrete assessment for readers unfamiliar with the benchmark.

Author response:
Thanks for the advice! We have added necessary statistics for the PDBBind benchmark in the revision (Section 4.1 and Appendix F).