Paper: GTMGC: Using Graph Transformer to Predict Molecule’s Ground-State Conformation (ICLR-2024)

Reviewer comment (Novelty):
Novelty: There are plenty of previous works (not cited) implementing Transformer architectures that implement the elements of the proposed self-attention, e.g., [1] uses the adjacency matrix, [2] makes use of the distance matrix (even cited in the manuscript), while the weighted summation is ubiquitous in many fields using Transformers [3]. The initial encoding is obviously not a technical contribution. Thus, the contribution may be summarized as the integration of existing methods/approaches, augmented with the regularization loss on the distance matrix (the \beta seems to bring very minor improvement).

Author response:
As you mentioned, numerous studies have attempted to utilize $\mathbf{A}$ and $\mathbf{D}$ to enhance the Transformer network structure in molecular modeling. We believe that $\mathbf{A}$ and $\mathbf{D}$  are the most prevalent inductive biases for 3D molecular structures. Therefore, our innovation lies not in augmenting self-attention with A and D per se, but in the manner we integrate A and D into the network. 

First, we list our core formulas and the core formulas of the two works [1] [2]:

**Ours**:  $\mathbf{S}_l^\prime=\mathbf{Q}_l{\mathbf{K}_l^T}\odot{[\mathbb{E}_n+{\gamma_l^A\times{\mathbf{A}}}+\gamma_l^D\times{(\vec{\beta}^{\max}-\mathbf{D})}]}$ and $\mathbf{O}_l=\text{softmax}(\frac{\mathbf{S}_l^\prime}{\sqrt{d_k}})\mathbf{V}_l$, where $\gamma_l$ is learnable parameters.

**[1]**: $\mathbf{O}_l=(\lambda_a\text{softmax}(\frac{\mathbf{Q}_l{\mathbf{K}_l^T}}{\sqrt{d_k}})+\lambda_dg(\mathbf{D})+\lambda_g\mathbf{A})\mathbf{V}_l$, where $\lambda$ is hyperparameters and $g$ is a function.

**[2]**: $\mathbf{O}_l=\text{softmax}(\frac{\mathbf{Q}_l{\mathbf{K}_l^T}}{\sqrt{d_k}})\odot{f({\mathbf{D}}^{-1})^2}\mathbf{V}_l$,  where $f$ is a shallow, fully connected neural network.

Although our method has certain similarities with theirs, we emphasize the following differences to enhance our novelty：

1. Figure 1(b) in our manuscript demonstrates our main idea. We propose using three masks, namely $\mathbb{E}_n$, $\mathbf{A}$, and $\vec{\beta}^{\max}-\mathbf{D}$, and adjusting the different combinations between each head with learnable parameters $\gamma_l^A$ and $\gamma_l^D$. Ultimately, we aim to adjust $\mathbf{Q}_l{\mathbf{K}_l^T}$ according to the final attention-mask, ${\mathbb{E}_n+{\gamma_l^A\times{\mathbf{A}}}+\gamma_l^D\times{(\vec{\beta}^{\max}-\mathbf{D})}}$ , that has learned and integrated ‘global’, ‘nearby’, and ‘spatial’ information **in each attention head **. This approach differs significantly from [1]. In [1], the self-attention mechanism is also enhanced through the use of $\mathbf{A}$ and $\mathbf{D}$ as biases. However, [1] directly adds $\mathbf{A}$ and $g(\mathbf{D})$ to the attention-weights $softmax(\frac{\mathbf{Q}_l{\mathbf{K}_l^T}}{\sqrt{d_k}})$ with global hyperparameters $\lambda$ which means they can't learn diffrent attention patterns (as dipected in Figure 4 in our manuscript) across different  attention heads.
2. We have innovatively adopted the $\vec{\beta}^{\max}-\mathbf{D}$ approach to incorporate the hypothesis that "the farther the distance, the smaller the interaction". This approach has shown some improvement in performance in both table 3 and table 4 of our manuscript when compared to the original $\mathbf{D}$ approach, although not in all cases. It is worth noting that this approach bears some similarity to $g(\mathbf{D})$ in [1] and $f({\mathbf{D}}^{-1})^2$ in [2] with $f$ being a neural network.  But in terms of details, the importance we place on “distance” varies in different places, but [1] is the same everywhere. In addition, [2] may have parameterized different $f$ in each Trans block, but in each block, the importance they focus on  “distance” for different attention heads is still the same. And we only introduce h (num attention heads) learnable parameters in each block, while [2] introduces a fully connected neural network.
3. We strictly follow that the attention-weights $\mathbf{O}_l=\text{softmax}(\frac{\mathbf{S}_l^\prime}{\sqrt{d_k}})$ are normalized, but they might not have.

We wish that the aforementioned explanation angle could enhance our novelty. To summarize, our innovative approach to introduce $\mathbf{A}$ and $\mathbf{D}$ into self-attention is that: 

 In each attention head, we capture different attention patterns by automatically learning the weights among “global”, “nearby”, and “spatial” information. The final attention-mask $\mathbb{E}_n+{\gamma_l^A\times{\mathbf{A}}}+\gamma_l^D\times{(\vec{\beta}^{\max}-\mathbf{D})}$, which is obtained after fusion, is used to adjust the attention-score $\mathbf{Q}_l{\mathbf{K}_l^T}$. Then, through softmax, we obtain the final attention-weights $\text{softmax}(\frac{\mathbf{S}_l^\prime}{\sqrt{d_k}})$  to weight the node features $\mathbf{V}_l$. The key point is that the patterns learned in each attention head are different. We don’t need to adjust additional hyperparameters and introduce very few learnable parameters to get performance improvement (as demonstrated in tables 3 and 4). Moreover, it allows for a more interpretable attention map (as depicted in Figure 4).

Reviewer comment (Evaluation):
Strange results: According to Table 3 (ablation study) a *simple Transformer* architecture without any addition to the self-attention reaches 0.4395 (MAE) which is already better than all the other baselines. This is problematic. Also, the final improvement is only ~1%. Finally, the reported results seem to take the best-ablated model results for each metric, which is wrong.

Author response:
There is indeed such a problem, and the simple Transformer structure has indeed achieved good results. However, from an application standpoint, it is also considered an "innovation" to use such a structure for the first time in completing this task. What we would like to emphasize more is our focus on the C-RMSD indicator (as mentioned in our manuscript) since it is widely used in the field of bioinformatics to measure the difference between two spatial structures, thus making it more convincing. The D-MAE and D-RMSE indicators are derived from the Molecule3D benchmark. In order to ensure a fair comparison, we still use these two indicators as references. Overall, our proposed structure has demonstrated an improvement of more than 3% in C-RMSD compared to Simple Trans, thereby proving its effectiveness. It is possible that this task is more difficult to optimize, hence we have strongly demonstrated in Table 4 that our proposed module has been successfully applied to molecular property prediction, resulting in significant performance improvements compared to simple Trans.

In addition, we corrected our careless error reporting results in Table 1 (best model on Molecule3D Random Test, D-MAE: 0.4325; D-RMSE:0.7210; C-RMSD:0.7129).

Reviewer comment (Experiments):
Lack of comparison with other Transformer based methods.

Author response:
We compared our method with the GPS method in Table 1, which is a Graph Transformer (GT) method  (as mentioned in Sec. 4.3). In Table 6, both SE(3)-Transformer and Equiformer are GT methods. And also MAT in table4.

Reviewer comment (Experiments):
Currently, it seems the proposed contributions don't bring any advantage over a simple (large) transformer model. One needs to know the capacity of the model in order to assess the origin of the good performance. According to weakness 2, it seems the good performances are obtained *almost solely* from a large/powerful standard Transformer model. From Table 5, the model seems much bigger than other methods. Also, the discrepancies between the tables are disturbing (one single metric should be used for the model validation).

Author response:
As for the responses to Weakness 2&3, Pure Trans indeed has brought very good performance. However, as shown in the updated Table 3 and Table 4, our MSRSA module introduces extremely few extra parameters (close to 0%) and achieves 3% and 60% performance improvements on the main task, ground-state conformation prediction, and the transfer task, molecular property prediction, respectively. 

Here, we express our apologies. We used wrong calculation results in Table 5 of our previous manuscript. Therefore, we reorganized the performance of the MSRSA module on molecular property prediction in Table 5 and Table 6 in Sec. A.2. The new correct results show that we achieved considerable results on Molecule3D with fast inference time. The results on Qm9 also achieved a medium-to-high level on Gap, Homo, and Lumo. 

While our $GTMGC_{large}$ does have a relatively large parameter size, it introduces less inductive bias compared to other baselines (bond angles, dihedral angles and so on). Moreover, its inference speed is impressively efficient.

Reviewer comment (Experiments):
Ablating the initial LPE.

Author response:
we add the ablation experiments on LPE  in Table 3 and Table 4. However, the results show that the performance is very poor after removing LPE. This is because PE is an important component in the Transformer network (as is the case in NLP). Therefore, we regard PE as a member of Pure Trans.

Reviewer comment (Experiments):
It would be beneficial to have a comparison performance with (at least one) other molecule transformer-based methods such as [4,5,6] or others at a similar capacity.

Author response:
Due to time and resource constraints, we have added a comparison experiment with [2] in Table 4 (MAT). We did not fully reproduce the model in [2], but instead directly replaced the attention formula in our code with MAT to corroborate the response to weakness1.

Reviewer comment (Experiments):
It would be interesting to look at the learned weighting parameters (\gamma) of the self-attention to better understand the contribution of each (maybe even adding a weighting to the global attention).

Author response:
The learned weights $\gamma$ are demonstrated in Section A.3.3. The corresponding weights are obtained and displayed by normalizing (1, $\gamma_l^A$ and $\gamma_l^D$).

Reviewer comment (Novelty):
The assertion regarding the innovative utilization of the MoleBERT Tokenizer might be overstated, especially in light of the results presented in Table 2. Previous molecular graph papers, such as the MolGPS paper, have explored various atomic featurizations that could potentially outperform the approach presented in this work.

Author response:
Thank you for your valuable comments! We would like to clarify that our primary focus in this task is the C-RMSD evaluation metric, which is a widely accepted measure in bioinformatics that quantifies the 3D spatial difference between two structures. The D-MAE and D-RMSE metrics, which were derived from the Molecule3D benchmark, were incorporated to facilitate a fair comparison. The introduction of MoleBert, particularly in comparison with atom ids, has significantly enhanced our performance on the C-RMSD metric. We consider this to be a noteworthy innovation in terms of performance improvement. 

In response to your suggestion, we’ve conducted an additional  ablation experiment using Ogb-style embeddings. These embeddings are node features parsed by rdkit and are represented as a vector of length 9, with each element signifying a different atomic property. The outcomes can be found in the updated version of Table 2. While this method of atomic featurization shows good results in terms of D-MAE and D-RMSE, it unfortunately falls short when evaluated against the C-RMSD indicator. Despite this setback, we maintain our belief that the Molebert Tokenized IDs, which we proposed, offer a beneficial approach to predicting molecular conformation.

Reviewer comment (Novelty):
The related work should be integrated into the main text rather than relegated to the appendix. It is crucial to comprehensively cover the various attempts to construct graph transformers and elucidate why they are ill-suited for the tasks at hand.

Author response:
According to your suggestions, we have appropriately discussed some relevant literature in the Introduction section. However, considering the natural flow of the text and the impact of the notation and formulas in the preliminary sections on the subsequent Method chapter, as well as the limited amount of work directly related to our task, ground-state conformation prediction (which is not widely studied), we have still chosen to present it in the Appendix.

Nonetheless, we have also included an additional section on the related work of "graph transformers". Regarding the feasibility of other GT methods, we mainly emphasize the simplicity and elegance of our structure (as mentioned in the text), which has already demonstrated good performance. We have also compared it with GPS, a kind of GT method, in Table 1.

Reviewer comment (Experiments):
Was molecular property prediction approached as a single-task or multi-task endeavor?

Author response:
Single-task.

Reviewer comment (Presentation):
Could you clarify the rationale behind placing the molecular property prediction results in the appendix, especially considering that they do not outperform SOTA across the board?

Author response:
The purpose is to illustrate that our proposed MSRSA module can be easily migrated to molecular property prediction and achieve remarkable results. Although not outstanding, our structure is simpler and easier to implement.

We would like to apologize and clarify that the data used in Table 5 of our previous manuscript was incorrectly calculated. As a result, we have made corrections in Sec. A.2 and reorganized our molecular property prediction experiments in Tables 4, 5, and 6. The experimental results show:

1. The ablation study validates the effectiveness of the components of the MSRSA module.
2. Table 5 demonstrates that our method achieves considerable performance on large-scale macromolecular datasets while maintaining desirable inference speed.
3. On Qm9, we achieved above-average performance on three properties: Gap, Homo, and Lumo. The performance on the other three properties is not significantly different from other baselines.

This fully demonstrates that our structure, while simple and elegant, can be quickly transferred to other tasks and achieve good performance.