Paper: General-purpose Pre-trained Model Towards Cross-domain Molecule Learning (ICLR-2024)

Reviewer comment (Novelty):
Contribution is minor. BIT combines the architecture of Transformer-M and Mixture-of-Domain-Experts technique, which both are from existing methods [1, 2, 3]. In addition, the way of combining the two is also a simple adaption. In summary, I appreciate that authors provide a strong method but I also believe that the contribution of this paper is not enough for acceptance.

Author response:
We deeply appreciate the reviewer for the insightful and constructive comments!

> Contribution is minor. BIT combines the architecture of Transformer-M and Mixture-of-Domain-Experts technique, which both are from existing methods [1, 2, 3]. In addition, the way of combining the two is also a simple adaption. 

We respectfully disagree with the reviewer that our contribution is minor. We acknowledge that the Transformer-M architecture and Mixture-of-experts technique have been verified in previous works, which we have extensively cited in section 3.2. We did not claim these or their combination as our contribution. 

We want to clarify that our main contribution is designing a general-purpose self-supervised framework for biomolecule modeling, as summarized in the Introduction. To the best of our knowledge, we are the first to propose a viable solution for this objective, and to learn fine-grained inter-molecular interactions with a simple backbone. Along the path to unified self-supervised pre-training, we make a series of key contributions:

- A unified tokenization approach across diverse molecular domains, enabling the straightforward combination architecture of Transformer-M and MoDE. Existing works either train in one particular molecule domain, limiting the model robustness, or train separate specialized backbones for different data and token types, restricting model scalability. 
- Unified self-supervised pre-training objectives applicable to diverse pre-training data sources and downstream tasks. This enhances model versatility and performance across tasks. 
- A general methodology to unify the diverse structural and domain-specific information of biomolecules, enriching the training dataset to enhance model performance.
- Reformulating virtual screening as a retrieval task, achieving state-of-the-art performance at ultra-fast inference speeds.

In summary, we are the first to demonstrate the feasibility and potential of cross-domain molecular self-supervised learning with a single backbone. The efficiency and broad applicability of our framework allow it to outperform existing approaches. We believe this self-supervised pre-training framework is novel and valuable to the community.

Reviewer comment (Experiments):
Authors did not provide ablation studies to directly show the effectiveness of the main components, e.g., MoDE, of BIT.

Author response:
Thanks for your suggestion. We have conducted further experiments to investigate the impact of MoDE following the suggestion. We compare three backbones, all trained on the same samples including both small molecules and protein-ligand complexes. One backbone is the vanilla Transformer-M. The second is Transformer-M with a larger FFN (2 $\times$ hidden size), having the same parameters as BIT. The third is our proposed BIT. Note that we selected two representative datasets (i.e., a small Tox21 dataset and a large HIV dataset) from MoleculeNet for the ablation study on molecular property prediction, due to time constraints. The results are presented in the table below. We found that introducing MoDE significantly improves the performance, especially on PDBbind where both small molecules and proteins could be encoded simultaneously. This verifies the benefit of introducing MoDE. 

| Backbone | HIV (AUC) $\uparrow$ | Tox21 (AUC) $\uparrow$ | PDBbind (MAE) $\downarrow$ | DUD-E (AUC) $\uparrow$ |
|---------|---------|---------|--------|--------|
| Transformer-M | 78.3 | 76.9 | 0.977 | 98.0 |
| Transformer-M (2 $\times$ hidden size) | 77.8  | 77.5 | 0.952 | 97.8 |
| BIT | **80.2** | **78.1** | **0.927** | **98.4** |

Reviewer comment (Experiments):
The comparison between BIT and the main baseline, i.e., Transformer-M, may not be fair enough. As BIT adopts MoDE technique, it has more trainable parameters than the vanilla Transformer-M. Moreover, BIT uses not only small molecule data but also protein-ligand complex data in the pre-training stage, while Transformer-M only uses small molecule data.

Author response:
Thanks for your suggestion. We have conducted further experiments to investigate the impact of MoDE following the suggestion. We compare three backbones, all trained on the same samples including both small molecules and protein-ligand complexes. One backbone is the vanilla Transformer-M. The second is Transformer-M with a larger FFN (2 $\times$ hidden size), having the same parameters as BIT. The third is our proposed BIT. Note that we selected two representative datasets (i.e., a small Tox21 dataset and a large HIV dataset) from MoleculeNet for the ablation study on molecular property prediction, due to time constraints. The results are presented in the table below. We found that introducing MoDE significantly improves the performance, especially on PDBbind where both small molecules and proteins could be encoded simultaneously. This verifies the benefit of introducing MoDE. 

As for the results of Transformer-M not being included in virtual screening and molecular property prediction benchmarks, this is because Transformer-M did not provide the specific experimental results mentioned above, so we were unable to make a direct fair comparison with it.

| Backbone | HIV (AUC) $\uparrow$ | Tox21 (AUC) $\uparrow$ | PDBbind (MAE) $\downarrow$ | DUD-E (AUC) $\uparrow$ |
|---------|---------|---------|--------|--------|
| Transformer-M | 78.3 | 76.9 | 0.977 | 98.0 |
| Transformer-M (2 $\times$ hidden size) | 77.8  | 77.5 | 0.952 | 97.8 |
| BIT | **80.2** | **78.1** | **0.927** | **98.4** |

Reviewer comment (Experiments):
The results of Transformer-M are not included in virtual screening and molecular property prediction benchmarks.

Author response:
As for the results of Transformer-M not being included in virtual screening and molecular property prediction benchmarks, this is because Transformer-M did not provide the specific experimental results mentioned above, so we were unable to make a direct fair comparison with it.

Reviewer comment (Writing):
In section 3.1, authors mentioned domain type embedding but without further description.

Author response:
We apologize for our unclear description. The domain type embedding, i.e. the molecule domain embedding and protein domain embedding $M_{type}, P_{type} \in R^d$, are two learnable vectors used to indicate the domain type of each atom token. They have the same length as atom feature embedding and structural positional embedding. These three embeddings are summed up as input to the self-attention module.

Reviewer comment (Reproducibility):
In section 3.1, authors introduce two special nodes, i.e., [M_VNode] and [P_VNodes]. It is unclear how to build up the input representation with these two nodes.

Author response:
We apologize for our unclear description. These two special node representations function similarly to the [CLS] token in BERT. However, unlike BERT, each domain expert in our model has its own distinct indicator token. [M_VNode] is inserted at the beginning of the molecule representation sequence, while [P_VNode] is inserted at the beginning of the protein representation sequence. We make the connection between [M_VNode] and each molecule atom individually, and also make the connection between [P_VNode] and each protein atom individually. Both special node embeddings are randomly initialized and updated during training like regular node representation.

Reviewer comment (Reproducibility):
What is the specific value of noise scale controlling hyperparameter $\sigma$?

Author response:
$\sigma$ is set to 0.2, the same value used in Transformer-M.

Reviewer comment (Experiments):
In section 3.3, why only add noise to molecules when doing pre-training?

Author response:
Since proteins tend to be relatively rigid in most cases, generative models typically assume knowledge of the bound protein structure [1]. This is a widely accepted assumption in current molecular docking research [2]. Therefore, we adopt this approach in our work and maintain a fixed protein pocket during ligand recovery, as done in related studies [2,3].

[1] Pagadala, Nataraj S., Khajamohiddin Syed, and Jack Tuszynski. "Software for molecular docking: a review." Biophysical reviews 9 (2017): 91-102.

[2] Corso, Gabriele, et al. "DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking." The Eleventh International Conference on Learning Representations. 2023.

[3] Guan, Jiaqi, et al. "3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction." The Eleventh International Conference on Learning Representations. 2023.

Reviewer comment (Experiments):
In section 3.4.1, why not use [P_VNode] or the combination of [M_VNode] and [P_VNode] as the representation of protein-ligand complexes?

Author response:
In the representation of the protein-ligand complex, [M_VNode] is positioned as the initial token, whereas the placement of [P_VNode] is more flexible. Therefore, using [M_VNode] simplifies engineering implementation. Additionally, early empirical studies demonstrated that using [M_VNode] either individually or in combination with [P_VNode] resulted in similar learning performance.

Reviewer comment (Experiments):
Some generalization of the model to other areas would be welcome.

Author response:
Thank you for acknowledging and supporting our work. Our current model focuses on general-purpose biomolecule learning. It operates at the molecular resolution level and could potentially be adopted for problems determined by physical structure. However, progress remains to be made in order to handle more general biological problems that cannot be fully explained through molecule interactions.

The idea of multi-domain modeling is powerful and could be applied to other scientific domains to capture correlations between different data domains, such as precision therapy (cell lines and molecules) and text-guided protein design (text and proteins). 
 
We plan to pursue more general AI capabilities step by step. The first step is building a foundation model for biomolecules and biochemical reactions. After this, we intend to integrate additional multi-domain data including images and text, which could enable solving more general biomedical problems. This feedback has been invaluable for identifying opportunities and challenges moving forward.

Reviewer comment (Experiments):
The main novelty of the paper is the introduction of MoDE in order to process data from diverse molecular domains. In my opinion, the paper lacks a proper evaluation of what these experts learn. For example, do molecule and pocket experts learn similar weights, or are there significant differences? What is the performance of this model when only one type of feed-forward layer is used (like in Transformer-M), but the same pretraining procedure is applied?

Author response:
Thanks for pointing this out. To further investigate the impact of MoDE, we introduce additional experiments. In these experiments, we use the same pre-training procedure to train two additional backbones. One backbone omits MoDE and uses a shared FFN, while the other also omits MoDE but utilizes a larger FFN (2 $\times$ hidden size) to match BIT in parameter scale. Note that we selected two representative datasets (i.e., a small Tox21 dataset and a large HIV dataset) from MoleculeNet for the ablation study on molecular property prediction, due to time constraints. The results are presented in the table below. We find that introducing MoDE significantly improves performance, especially on PDBbind where both small molecules and proteins could be encoded simultaneously. This aligns with our motivation for introducing MoDE. Furthermore, the performance gains cannot be attributed solely to increased parameters, as using a larger FFN alone did not consistently improve performance across tasks.

| Backbone | HIV (AUC) $\uparrow$ | Tox21 (AUC) $\uparrow$ | PDBbind (MAE) $\downarrow$ | DUD-E (AUC) $\uparrow$ | 
|---------|---------|---------|--------|--------|
| BIT | **80.2** | **78.1** | **0.927** | **98.4** |
| w/o MoDE | 78.3 | 76.9 | 0.977 | 98.0 |
| w/o MoDE (2 $\times$ hidden size) | 77.8  | 77.5 | 0.952 | 97.8 |

Reviewer comment (Experiments):
The Authors propose to use two pre-training objectives. It would be interesting to see what is the impact of each of them. Why did the Authors decide to use a different pre-training procedure than used in Transformer-M? The experimental tables are missing the results achieved for the non-pretrained model.

Author response:
Thanks for your suggestion. To examine the impact of the two pre-training objectives, we introduce additional experiments by eliminating select objectives to study their effectiveness. As shown in the table below, the pretrain-then-finetune paradigm is significantly superior to direct supervised training for a Transformer model. We also observe a notable performance decrease when removing either pre-training objective, particularly coordinate denoising for PDBbind (2D+3D) and masked token denoising for molecular property prediction (2D). These results indicate that both pre-training objectives are crucial and yield positive outcomes.

We did not use the pre-training procedure from Transformer-M for two reasons: 1) supervised pre-training based on the homo-lumo gap prediction may lead to negative transfer, and 2) there is extremely limited data with both 3D co-crystal structures and experimentally determined binding affinity values (approximately 23,000 in PDBbind). Our aim is to build a more versatile model using more unsupervised data, so we adopt a purely self-supervised pre-training approach.

| Pre-training objectives | HIV (AUC) $\uparrow$ | Tox21 (AUC) $\uparrow$ | PDBbind (MAE) $\downarrow$ | DUD-E (AUC) $\uparrow$ | 
|---------|---------|--------|--------| --------|
| BIT | **80.2** | **78.1** | **0.927** | **98.4** |
| w/o Masked token denoising | 78.2 | 77.6 | 0.939 | 96.5 |
| w/o Coordinate denoising | 78.5 | 76.1 | 1.016 | 97.1 |
| w/o pre-training | 70.9 | 75.1 | 1.114 | 95.7 |

Reviewer comment (Experiments):
The choice of the models used in the experiments seems arbitrary. For example, why are different models used for binding affinity prediction and molecular property prediction? GROVER could be used in both scenarios. Why are some of these models not pre-trained, while the original works provide pre-training procedures (and sometimes also the pre-trained weights), e.g. GROVER and MAT?

Author response:
We apologize for the unclear description of the baseline. The choice of baselines follows these rules: for molecular property prediction, we choose the most representative 2D and 3D pre-train baselines and report their official results from [1]; for binding affinity prediction, we report the official results of state-of-the-art GNN-based baselines from [2], and also include the recently published pre-trained methods Transformer-M and MBP [3]. We will provide a clear description in the corresponding tables. Neither GROVER nor MAT provided the specific experimental results mentioned above, so we were unable to make a direct fair comparison with them.

[1] Yu, Qiying, et al. "Unified Molecular Modeling via Modality Blending." arXiv preprint arXiv:2307.06235v1 (2023). 

[2] Li, Shuangli, et al. "Structure-aware interactive graph neural networks for the prediction of protein-ligand binding affinity." SIGKDD. 2021.

[3] Yan, Jiaxian, et al. "Multi-task Bioassay Pre-training for Protein-ligand Binding Affinity Prediction." arXiv preprint arXiv:2306.04886 (2023).

Reviewer comment (Experiments):
Another benchmark for structure-based virtual screening would be helpful in assessing the performance of the proposed method. It has been shown, that decoys contain hidden biases that can impact the performance of deep learning methods [1].

Author response:
Thank you for the recommendation. To further validate BIT’s performance, we follow the suggestion in the comment and test the BIT model, which is fine-tuned on DUD-E, on the unbiased external test set LIT-PCBA [4]. Note that we exclude all the targets present in LIT-PCBA from the training set to prevent data leakage. BIT still outperforms the conventional docking method Glide-SP, the supervised deep learning method Planet, and the pre-trained method DrugCLIP, as illustrated below:

|  | AUROC $\uparrow$ |  $EF_{0.005}$ $\uparrow$ | $EF_{0.01}$ $\uparrow$ | $EF_{0.05}$ $\uparrow$ |
|---------|---------|--------|--------| --------|
| Glide-SP | 53.15 | 3.17 | 3.41 | 2.01 |
| Planet   | 57.31 | 4.64 | 3.87 | 2.43 |
| DrugCLIP | 57.17 | 8.56 | 5.51 | 2.27 |
| BIT      | **61.44** | **10.07** | **5.80** | **2.87** |

[4] Tran-Nguyen, Viet-Khoa, Célien Jacquemard, and Didier Rognan. "LIT-PCBA: an unbiased data set for machine learning and virtual screening." Journal of chemical information and modeling 60.9 (2020): 4263-4273.

Reviewer comment (Evaluation):
For binding affinity prediction, did you consider measuring non-linear correlation, e.g. using the Spearman correlation coefficient? This evaluation metric could be better at showing which methods can correctly prioritize compounds with strong affinity.

Author response:
We have considered using non-linear correlation to evaluate the ‘scoring power' on PDBbind. However, as established in the standard comparative assessment of scoring functions (CASF) benchmark [5], such metrics are better suited to quantifying 'ranking power’, where multiple ligands are ranked against the same target protein. Since PDBbind contains only one ligand per unique target protein, the non-linear correlation is not applicable in this case. Thus, we quantify 'scoring power' by reporting key regression metrics - Root Mean Square Error (RMSE), Mean Absolute Error (MAE), Pearson's correlation coefficient (R), and the standard deviation (SD), following precedents from previous work [6], to ensure a fair comparison.

[5] Su, Minyi, et al. "Comparative assessment of scoring functions: the CASF-2016 update." Journal of chemical information and modeling 59.2 (2018): 895-913.

[6] Stepniewska-Dziubinska, Marta M., Piotr Zielenkiewicz, and Pawel Siedlecki. "Development and evaluation of a deep learning model for protein–ligand binding affinity prediction." Bioinformatics 34.21 (2018): 3666-3674.

Reviewer comment (Reproducibility):
Do you plan to publish the code for better reproducibility of your method?

Author response:
Yes. The code will be made publicly available upon acceptance of the paper.

Reviewer comment (Experiments):
The model only considers the binding pocket segment of proteins—while computationally efficient, may not be practically feasible without costly simulations to identify these segments. This reliance on prior knowledge or expensive computations could limit the accessibility and scalability in practical applications.

Author response:
This is true. Pocket searching falls outside the scope of this paper, which is focused on training a cross-domain pre-trained model for biomolecules. For the binding-related task, we follow the standard protocol from mainstream approaches [1,2] and the competing methods [3], which assume the binding pocket as given. Pocket searching remains an active research area that may require distinct methodological considerations. Machine learning also has demonstrated potential in this domain, as evidenced by recent approaches [4,5]. 

[1] Trott, Oleg, and Arthur J. Olson. "AutoDock Vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading." Journal of computational chemistry 31.2 (2010): 455-461.

[2] Koes, David Ryan, Matthew P. Baumgartner, and Carlos J. Camacho. "Lessons learned in empirical scoring with smina from the CSAR 2011 benchmarking exercise." Journal of chemical information and modeling 53.8 (2013): 1893-1904.

[3] Li, Shuangli, et al. "Structure-aware interactive graph neural networks for the prediction of protein-ligand binding affinity." Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining. 2021.

[4] Krivák, Radoslav, and David Hoksza. "P2Rank: machine learning based tool for rapid and accurate prediction of ligand binding sites from protein structure." Journal of cheminformatics 10 (2018): 1-12.

[5] Gainza, Pablo, et al. "Deciphering interaction fingerprints from protein molecular surfaces using geometric deep learning." Nature Methods 17.2 (2020): 184-192.

Reviewer comment (Theory):
The strategy of pre-training the model on both equilibrium structures of molecules and higher-energy protein-ligand complexes may be conceptually problematic. These two types of data represent vastly different energy states, and it is unclear what meaningful semantic learning can be achieved by pretraining them together. This may lead to a model that does not adequately distinguish between the distinct energetic landscapes of the two systems. Also, the scales of the two training sources are not balanced, where PCQM4Mv2 is significantly larger than Q-BioLiP, but the model performance on molecular property prediction is not very significant compared to the baselines, which warrants a further examination.

Author response:
We understand the reviewer’s concern is that the equilibrium structures of molecules and protein-ligand complexes may inhabit different energy states, as the same molecule may have different conformations in different biomolecule contexts, which may confuse the model. However, this is not necessarily problematic. Even though these are two distinct energy landscapes, our pre-training objective remains consistent across different types of pre-training data in restoring the low-energy structures from the corrupted ones depending on the biomolecule context. The issue is analogous to polysemy in NLP where a word can have multiple semantic meanings depending on the context. With sufficient training data, the model can learn context-aware representations automatically. Currently, without NLP-scale data, our solution is using separate domain experts for different biomolecules. We have also experimented with a unified Transformer architecture without MoDE, but the results were unsatisfactory as demonstrated below - the reviewer's observation may explain why. Note that we selected two representative datasets (i.e., a small Tox21 dataset and a large HIV dataset) from MoleculeNet for the ablation study on molecular property prediction, due to time constraints.

| Backbone | HIV (AUC) $\uparrow$ | Tox21 (AUC) $\uparrow$ | PDBbind (MAE) $\downarrow$ | DUD-E (AUC) $\uparrow$ | 
|---------|---------|---------|--------|--------|
| BIT | **80.2** | **78.1** | **0.927** | **98.4** |
| w/o MoDE | 78.3 | 76.9 | 0.977 | 98.0 |

Regarding unbalanced data sources, prior work in vision-language pre-training has shown that adding even small amounts of paired image-text data provides value over text-only data, especially for downstream vision-language tasks [6]. Similarly, the inclusion of Q-BioLiP primarily benefits binding-related tasks, with marginal gains for property prediction. This outcome is unsurprising, given the relatively smaller number of molecules available for training in Q-BioLiP, coupled with the fact that PCQM4Mv2 is also utilized by recent baselines (i.e., MoleculeSDE and MoleBLEND).

We appreciate these insightful comments. To enhance the model's capabilities, we aim to accumulate more diverse biomolecular training data and scale up model size. This feedback has been invaluable for identifying opportunities and challenges.

[6] Su, Weijie, et al. "VL-BERT: Pre-training of Generic Visual-Linguistic Representations." International Conference on Learning Representations. 2020.

Reviewer comment (Reproducibility):
There is a lack of a detailed explanation for how positional encodings of 2D and 3D structures are integrated into the model, which leaves a gap in understanding the full architecture and mechanics of the model.

Author response:
We apologize for not providing a detailed description of the positional encodings. We will revise the writing to provide a more comprehensive description. To clarify, the positional encodings for each data format (2D or 3D) comprise two components:

- Pairwise encodings: shortest path distance and edge information for 2D graphs [7]; Euclidean distances encoded with the Gaussian Basis Kernel function for 3D structures [8]. They are added to the self-attention module as bias terms.

- Atom-wise encodings: degree information for 2D graphs; the sum of distances to all other atoms for 3D structures. They are added to the input atom features.

The 2D and 3D positional encodings are integrated into the model through two separate channels. The channel that matches the input format (2D, 3D, or 2D+3D) is activated, while the non-corresponding channel remains disabled. 

[7] Ying, Chengxuan, et al. "Do transformers really perform badly for graph representation?." Advances in Neural Information Processing Systems 34 (2021): 28877-28888.

[8] Shi, Yu, et al. "Benchmarking graphormer on large-scale molecular modeling datasets." arXiv preprint arXiv:2203.04810 (2022).