General-purpose Pre-trained Model Towards Cross-domain Molecule Learning

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.

1. Authors did not provide ablation studies to directly show the effectiveness of the main components, e.g., MoDE, of BIT.
2. 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.
3. The results of Transformer-M are not included in virtual screening and molecular property prediction benchmarks.

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.

Thank you for pointing out our unclear description. We will revise the writing to provide a more comprehensive description.

In section 3.1, authors mentioned domain type embedding but without further description.

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.

Also in section 3.1, authors introduce two special nodes, i.e., [M_VNode] and [P_VNode]. It is unclear how to build up the input representation with these two nodes.

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.

What is the specific value of noise scale controlling hyperparameter $\sigma$

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

In section 3.3, why only add noise to molecules when doing pre-training?

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.

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?

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.

In table 3, why the PCBA dataset is not included in the table?

For molecular property prediction, we compare against the most representative 2D and 3D pretraining baselines, and report their official results from [4]. We exclude PCBA since most baselines have not reported results on this dataset. However, BIT still achieves the state-of-the-art performance among available records, as shown below:

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

We hope our response can alleviate your concerns. Please let us know if you have any additional questions.

I appreciate your detailed responses. The additional experiments have alleviated my concerns that the model's performance is solely due to increased parameters and data. Furthermore, I concur that developing a general-purpose model for both protein pockets and small molecules is an innovative concept. However, I still maintain that BIT lacks technical novelty due to the following reasons:

• Unified Tokenization Approach: The tokenization method used in BIT is simply a standard approach that involves treating molecules at the atom level. This approach has been employed in previous studies, such as Uni-Mol and Transformer-M. Therefore, I do not consider it to be a novel technique. Additionally, I was unable to locate a detailed explanation of the tokenization approach in the paper.

• Unified Self-Supervised Pre-training Objectives: Similar to the tokenization approach, the pre-training objectives of BIT are aligned to Transformer-M. It is hard to say the objectives are novel.

• General Methodology to Unify Domain-specific Information: Because MoDE is very similar to MOE, the general methodology should be credited to MOE, which was proposed earlier.

• Reformulating Virtual Screening as a Retrieval Task: DrugCLIP [1] has previously reformulated virtual screening as a retrieval task, and has provided a more comprehensive evaluation of this approach.

In conclusion, I believe that BIT is a well-implemented model, but it lacks a deep understanding. I therefore maintain my original score.

[1] Gao B, Qiang B, Tan H, et al. DrugCLIP: Contrastive Protein-Molecule Representation Learning for Virtual Screening. NeurIPS 2023

We deeply appreciate the reviewer for the insightful and constructive comments!

In the comparison tables, the models (usually proposed BIT) with the best mean have values given in boldface. Are all the models trained on the same data-sets? For some predicted features, the differences between BIT and some other models are large, even though some of them are Transformers too. Are the optimal values for BIT the result of the proposed BIT architecture, the fine-tuning on different modes (the multi-modality), different data sets, better pre-training, or something other? The discussion on the comparison to other approaches is needed in the conclusions/discussion section.

Thanks for pointing this out. You are right that our method and other pre-training baselines are inconsistent in terms of the pre-training data and pre-training objectives. Being able to leverage more cross-domain data is an advantage of our proposed BIT. For downstream tasks, we report the official results of baselines and follow the same experimental settings as baselines to ensure a fair comparison. We will discuss this issue in the revised paper.

You provided us with an excellent suggestion to better isolate the benefit of our backbone modeling approach through controlled experiments. Thus, we introduce ablation experiments to investigate the factors influencing performance. Using the same pre-training data, we remove select pre-training objectives or train different backbones with identical procedures to quantify their impact. 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. As summarized in the following tables, we find 1) our two pre-training objectives are necessary and effective, and 2) introducing MoDE significantly improves performance. In the revision, we will incorporate these ablation studies to provide a controlled comparison.

Please correct the spelling of some words. Just as well, please correct the editing of mathematical expressions. E.g. in equations (1) and (2) the equal signs = should be aligned ;-)

Thanks for your patience and helpful reminder. We will carefully revise our writing.

• Some generalization of the model to other areas would be welcome.

• Your paper and model is strictly molecule learning oriented. Do you think that the general approach can be used in other sciences? E.g. in biological experiments on cancerous cells and the impact of some sort of certain treatments, which would imply modifications in the course of the operation?

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.

We deeply appreciate the reviewer for the insightful and constructive comments!

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?

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.

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.

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.

We deeply appreciate the reviewer for the insightful and constructive comments!

The pretraining process appears to be confined to proteins and small molecules. This limited scope raises questions about the model's applicability to a broader range of biomolecules and interactions.

Thanks for pointing this out. A key advantage of our proposed method over existing works is indeed its broader applicability. Existing methods typically use different tokens for different biomolecules, such as amino acids for proteins and atoms for small molecules. This requires designing specialized encoders for each type of biomolecule, limiting the applicability. Our method instead employs the same token (i.e., atom) for all biomolecule types and uses Mixture-of-Domain-Experts (MoDE) to capture domain specificity. By simply adding corresponding domain expert FFNs, our approach can be readily applied to any biomolecule and interaction.

The reason we focus on proteins and small molecules in this paper is that protein-molecule binding-related tasks are typical applications with sufficient structural data. This aligns with the main focus of most related works in structure-based drug design and enables a direct comparison to the previous methods in the field. However, as the reviewer points out, investigating more diverse, high-quality biomolecules for pre-training would be valuable future work to further boost the performance and applicability of our model.

The proposed Transformer backbone is invariant to geometric transformations, which may limit its expressive power compared to SE(3) or E(3) equivariant architectures.

Thanks for your thoughtful reminder. We have exactly the same question ourselves when designing the backbone structures. Ultimately, we choose our Transformer backbone over the equivariant architectures for the following reasons: While equivariant architectures efficiently encode rotational and translational symmetries for molecules, their computational overhead becomes prohibitive for large-scale pre-training. For our goal of pre-training a sizable model using atom tokens at scale, training scalability and stability are more crucial. Furthermore, sufficient training data equipped with augmentation techniques can compensate for the lack of built-in equivariance for the Transformer. Empirical results across tasks confirm our backbone's superior performance as expected. Though promising, equivariant networks' current challenges make our backbone a better fit for now.

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.

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.