Paper: Multi-Scale Protein Language Model for Unified Molecular Modeling (ICLR-2024)

Reviewer comment (Experiments):
There are significant flaws in the design of code-switch protein sequences, which can potentially make the model learn meaningful insights from the unzipped sequences.

Author response:
**A1**:  Thanks for your comments. However, it seems that **the concluded significant flaws of code-switch protein sequence design could be caused by some misunderstandings.** The code-switch design and the unzip operation play crucial roles in multi-scale molecular modeling. Sorry for the confusion and we clarify it as follows: 
- **Unzipping residues also introduces non-trivial structural information for the model, aiding ms-ESM in learning diverse structural knowledge of residues.** Despite the limited types of amino acids and the fixed types of atoms within a given residue, structure-based training (pair-wise distance recovery) is conducted within an unzipped residue at the atomic scale. The structure of a single residue is thus highly diverse. Therefore, unzipping residues enables the model to acquire diverse structural information at the residue level.
- **The unzip operation allows protein language models to learn associations between proteins and molecules and reuse the knowledge acquired by the model at the residue scale.**
  - Similar to existing work [24] in cross-lingual pre-training, MLM training on code-switch sequences enables the model to capture correlations between information from the atom and residue scales.
  - By retaining information at the original residue scale, the model can fully leverage the knowledge embedded in protein language models.
  - The number of residue types does not impact the establishment of this association, because atom-scale information contains diverse structural details.
- **Evidence from experimental results and related works**:
  - Ablation experiments demonstrate the effectiveness of the unzip operation.
    - Supporting the effectiveness of this approach, ablation experiments (as shown in Table D) reveal a significant performance decrease in protein-molecule tasks when the unzip operation is omitted. This underscores the crucial significance of this approach in constructing the semantic association between the two types of information.
  - Related work also mentions the effectiveness of this approach.
    - Concurrent work titled "Generalized Biomolecular Modeling and Design with RoseTTAFold All-Atom" [25] (published on bioRxiv after our submission to ICLR) introduces a similar concept called "atomization" and notes that this approach can assist the network in learning about general atomic interactions.
In summary, both experimental results and related work affirm the effectiveness of this approach, forming the foundational basis for constructing a unified semantic representation across multiple scales.

Reviewer comment (Evaluation):
The results present in the experimental section only shows marginal improvements over the established baselines. Additionally, the unified pre-training appears to diminish performance on tasks focused solely on proteins or molecules. This compromises its viability for practical applications.

Author response:
**A3**: Please refer to the Q1 in the general response for more discussion. Briefly, i) as already mentioned in the answer to Q1 in the general response, this paper focuses on unified multi-scale modeling, aiming to obtain better results on protein-molecule tasks  instead of improving molecule/protein tasks universally. ii) We report molecule-only and protein-only results to verify that ms-ESM does not suffer from significant performance drops on separated molecular and protein modeling tasks.

Reviewer comment (Novelty):
The related work section seems to overlook several pivotal studies pertinent to molecular modeling.

Author response:
**A5**: 
Thank you for your information. Previously, as we focused more on protein and unified modeling, we did not include discussions of molecular modeling in our manuscript. Subsequently, we've included relevant discussions in the Appendix. We are open to further discussing any papers that reviewers consider pertinent to our work.
- *Molecular Modeling.* Regarding the modality of molecules, studies on molecular modeling can be categorized into three groups:
  - 1D-based methods: These represent molecules with SMILES strings and employ language modeling techniques, such as masking and contrastive self-supervision, to enhance molecular representation [1, 2, 3, 4, 5, 6].
  - 2D-based methods: These represent molecules with molecular graphs, sharing common ideas with general graph modeling. Some methods [7, 8, 9, 10] mask key substructures of molecular graphs, like motifs and functional groups, and task models with reconstructing the masked parts. Others [11, 12, 13] align views from positive pairs (corrupt versions of the same graph) and simultaneously contrast views from negative pairs (different graphs). 
  - 3D-based methods: These directly utilize the 3D structure of molecules, aligning closely with our work. Earlier studies incorporated 3D information as an auxiliary input for 2D-based methods [14, 15, 16, 17]. More recent methods focus on molecular modeling with pure 3D inputs [18, 19, 20, 21, 22, 23]. Three self-supervised techniques have been designed: geometry masking, geometry predicting, and denoising. For masking, [18] masks bond information, while [19] masks atom types, requiring models to predict masked information based on remaining context. For predicting, [18] proposes an atomic prediction task with bond information to capture global structure from local information. For denoising, models reconstruct 3D structures by adjusting corrupted structures. When corrupting structures, [19, 20, 21] add Gaussian noise to each atom of the input molecule. Several methods further introduce E(3)- and SE(3)-invariance inductive bias to the denoising technique [19, 22, 23].

Reviewer comment (Novelty):
The concept of unzipping residues into atom sequences and predicting the masked atom type appears flawed. Given the residue type or the types of adjacent atoms, deducing the type of the masked atom becomes straightforward, since the atom set is predetermined for each residue type. This indicates that unzipping residues does not introduce any unique or non-trivial information. Additionally, there is a lack of experiments to demonstrate that the masked atom type prediction loss contributes any meaningful insight.

Author response:
**A1**:  Thanks for your comments. However, it seems that **the concluded significant flaws of code-switch protein sequence design could be caused by some misunderstandings.** The code-switch design and the unzip operation play crucial roles in multi-scale molecular modeling. Sorry for the confusion and we clarify it as follows: 
- **Unzipping residues also introduces non-trivial structural information for the model, aiding ms-ESM in learning diverse structural knowledge of residues.** Despite the limited types of amino acids and the fixed types of atoms within a given residue, structure-based training (pair-wise distance recovery) is conducted within an unzipped residue at the atomic scale. The structure of a single residue is thus highly diverse. Therefore, unzipping residues enables the model to acquire diverse structural information at the residue level.
- **The unzip operation allows protein language models to learn associations between proteins and molecules and reuse the knowledge acquired by the model at the residue scale.**
  - Similar to existing work [24] in cross-lingual pre-training, MLM training on code-switch sequences enables the model to capture correlations between information from the atom and residue scales.
  - By retaining information at the original residue scale, the model can fully leverage the knowledge embedded in protein language models.
  - The number of residue types does not impact the establishment of this association, because atom-scale information contains diverse structural details.
- **Evidence from experimental results and related works**:
  - Ablation experiments demonstrate the effectiveness of the unzip operation.
    - Supporting the effectiveness of this approach, ablation experiments (as shown in Table D) reveal a significant performance decrease in protein-molecule tasks when the unzip operation is omitted. This underscores the crucial significance of this approach in constructing the semantic association between the two types of information.
  - Related work also mentions the effectiveness of this approach.
    - Concurrent work titled "Generalized Biomolecular Modeling and Design with RoseTTAFold All-Atom" [25] (published on bioRxiv after our submission to ICLR) introduces a similar concept called "atomization" and notes that this approach can assist the network in learning about general atomic interactions.
In summary, both experimental results and related work affirm the effectiveness of this approach, forming the foundational basis for constructing a unified semantic representation across multiple scales.

Reviewer comment (Evaluation):
As shown in the Tables 2 and 3, ProSmith's results with ms-ESM offer only a slight enhancement over baselines, such as He et al. (2023) in Table 2 and Kroll et al. (2023b) in Table 3.

Author response:
**A4**: Thanks for your comments. ms-ESM outperforms ESM-2+Uni-Mol in all protein-molecule tasks, and also surpasses a single strong protein model (such as ESM-1b, ProtBert) in most tasks. We will provide an analysis of the reasons from two different perspectives.
- **Comparisons with separated protein and molecule models (ESM2 + Uni-Mol)**: ms-ESM outperforms ESM-2+Uni-Mol in all protein-molecule tasks.
  - Our protein representation learning ability is similar to ESM-2, and our molecular representation learning ability is similar to Uni-Mol, but our performance on protrein-molecule joint modeling tasks is stronger than ESM-2 + Uni-Mol.
    - During pre-training, we don't utilize any protein-molecule pair data just like all baselines and adopt the same fine-tuning setting. So we believe that the comparison is very fair. 
  - It is worth noting that our model maintains an advantage over ESM-2+Uni-Mol in all Protein-Molecule task metrics, which illustrates the rationale behind our unified modeling approach.
- **Comparisons with a single strong protein model (such as ESM-1b, ProtBert)**:  Small ms-ESM can even outperform large models.
  - Under unified modeling, ms-ESM's performance can match or surpass some single representation learning models with significantly larger parameter scales (such as ESM-1b 650M) even at a much smaller scale (35M) . 
  - This illustrates that, to enhance a model's performance in protein-molecule tasks, the most crucial strategy is to construct a unified model to obtain a consistent semantic representation. This proves the significance of the unified model.

Reviewer comment (Evaluation):
Despite the possible data and information leakage in the experimental framework, the proposed method fails to outperform standard protein language models like ESM-2 in tasks such as contact and secondary structure prediction. Given that the model incorporates a broader range of pre-training loss than ESM-2, these results suggest that universal pre-training across both proteins and molecules may not offer advantages in protein-only or molecule-only tasks. This significantly weakens the paper's primary claim: the benefit of combining protein and molecule data for pre-training.

Author response:
**A9**:
- About the leakage problem:
  - Information leakage:  The input of ms-ESM in protein-only tasks only contains protein sequences at the residue level, thus there is no information leakage in these tasks. For more details, please refer to Q6.
  - Data leakage: There is almost no leakage of test data in protein-only tasks, and we also reevaluate the model's performance and find that the model performance remains almost unchanged. For more details, please refer to Q8.
- The proposed method fails to outperform ESM-2: 
  - For more details, please refer to Q1 in the general response.
    - ms-ESM retains comparable ability to ESM-2 in protein-only tasks.
    - Even though ms-ESM can't beat ESM-2 in all protein-only tasks and ms-ESM can't beat Uni-Mol in all molecule-only tasks, ms-ESM still can beat  ESM-2+Uni-Mol in ALL protein-molecule tasks when using the same fine-tuning and evaluating protocol. This demonstrates the effectiveness of unified modeling.

Reviewer comment (Experiments):
The ablation study only consider ablations on position encoding, neglecting the paper's other two significant contributions. It would be beneficial for the authors to explore comparisons with protein- or molecule-only pre-training, consider pre-training without unzipped sequences, and evaluate the impact of removing each pre-training loss.

Author response:
**A10**: Thank you very much for your helpful questions. We have added more ablation experiments (shown in Table D) and also added them to our revised manuscript. Here is the analysis of these results.
- Ablation Study on Pre-training Objectives:
  - Significantly decreased model performance is observed when either the masked atom type prediction loss or the pair-wise distance recovery loss is omitted.
  - Notably, removing the pair-wise distance recovery loss results in a greater performance loss compared to omitting the masked atom type prediction loss.
    - Explanation:
      - Without using pair-wise distance recovery loss, the training loss of ms-ESM on the atomic scale is only masked atom type prediction loss. But more than half of the atoms in a molecule are carbon atoms, which makes the model learn very little information on the atomic scale using only masked atom type prediction.
      - Without the pair-wise distance recovery loss, ms-ESM struggles to extract structural information from unzipped residues, making the training task too simplistic and hindering the acquisition of meaningful knowledge. 
  - Results indicate that, while both atom type and structural information are vital for atomic-scale details, structural information holds greater importance.
- Ablation Study on Pre-training Data:
  - A significant decrease in model performance is observed when either molecular data or protein data is excluded.
  - Interestingly, the removal of protein data leads to a more substantial performance decline than omitting molecular data.
    - Insights:
      - When the model is not trained with protein data, it quickly forgets protein-related knowledge, resulting in a notable decline in overall performance.
      - The model can still learn some atomic scale information through unzip operations without molecular data. This explains why the model performs better without molecular data compared to the scenario without protein data.
      - This experiment effectively demonstrates the beneficial impact of the unzip operation in modeling atomic-scale information.
- Ablation Study on Unzip Operation: As analyzed in Q1, the model performance drops significantly after removing the unzip operation, which illustrates its importance for modeling unified semantics.

Table D. Ablation study
|  | MSE ↓ | R^2 ↑ |
| --- | --- | --- |
| Vanilla ms-ESM | 0.627 | 0.546 |
| w/o ASPE | 0.639(+0.012) | 0.537(-0.009) |
| w/o RSPE in atoms  | 0.627(+0.0) | 0.547(+0.001) |
| w/o masked atom type loss | 0.642(+0.015) | 0.535(-0.011) |
| w/o pair-wise distance recovery loss | 0.645 (+0.018) | 0.533(-0.013) |
| w/o molecular data | 0.648 (+0.021) | 0.531(-0.015) |
| w/o protein data | 0.708(+0.081) | 0.487(-0.059) |
| w/o unzip operation | 0.638 (+0.011) | 0.538(-0.008) |

Reviewer comment (Novelty):
the technical novelty is limited; I understand code-switching is not novel but I value that it is ported to this field. However, the rest of the techniques used in this paper (like RoPE, or the transformer architecture) are not novel.

Author response:
**A1**: Please refer to the Q2 of the general response for more discussion. In brief, our work is novel for three reasons: i) we take the construction of a model (not exisiting techniques) which can process proteins and molecules uniformly as the core contribution of our work. ii) Unified molecular modeling topic is the cutting edge of AI biological application emerging recently. iii) Adopting existing techniques to our multi-scale unified molecular modeling is non-trivial

Reviewer comment (Writing):
Lack of detail on certain topics. For example, in Section 2.2, the ORDER procedure is not explained (there is a referral to Appendix A, where however I didn't find an explanation). Similarly, the "Atom Scale Position Encoding" section is not informative with lots of unintroduced symbols.

Author response:
**A2**: Thank you for bringing this to our attention, which enhances the clarity of our paper. While we briefly touched upon these details in our manuscript (Appendix A and Sections 2.2, 2.3), we will now introduce them more comprehensively in the following paragraphs and **incorporate them into the revised paper.**  
- Regarding the ORDER procedure, as indicated in Appendix A (the final line of the Hyperparameters section), we adopt the default order from PDB (protein) and SDF (molecule) files for assigning order to the atom set.  
  - To elaborate, PDB and SDF serve as standard formats for describing atomic structures of proteins and small molecules, respectively. In both formats, atoms follow specific sorting principles. In our study, we directly utilize the sorted atoms for ease of implementation.  
  - It is important to note that, given our atom-scale position encoding employs Euclidean distance to describe positional relationships, the permutation of atom order does not impact our pre-training model.  

- For atom scale position encoding, following Uni-Mol [1], we use Euclidean distance and Gaussian kernel to encode the 3D position of atoms. For $h_i, h_j \in \bar{X}$ (both $h_i$ and $h_j$ can be a residue or atom in code-switch protein sequence $\bar{X}$, Section 2.2), their atom scale position encoding $e_{ij}^{A}$ can be calculate as: 
$$
e_{ij}^{A} = \begin{cases}
0 & h_i \in \bar{R} ~or ~h_j \in \bar{R} \\\\
Gaussian(Dis(c_i, c_j)) & otherwise, c_i = Coord(h_i),c_j = Coord(h_j) \end{cases}
$$
  where $\bar{R}$ is the set of all the residues in $\bar{X}$ (Section 2.2), $Coord$ extracts coordinates from atoms (Section 2.2), $Dis$ calculates the Euclidean distance between two coordinates (Section 2.2), and $Gaussian$ is the Gaussian kernel (Section 2.3).

Reviewer comment (Theory):
The "slight modification of the Transformer" in section 2.4 appears poorly justified or at least needs more clarification. Why is $E^A$ added to the standard attention? What happens if it's not added?

Author response:
**A3**:  
- To enhance the awareness of the 3D position relationship between atoms, we incorporate $E^A$ into the attention layer of the vanilla Transformer, inspired by the approach in Uni-Mol [1], a widely recognized technique in geometric molecular modeling [2, 3, 4] known for its efficacy in various downstream tasks. Specifically, when two atoms are in close proximity in 3D space, $E^A$ boosts the attention value between them, subsequently encoding this positional relationship into the Transformer.  
- Ablation studies detailed in our paper (Section 3.4) offer empirical evidence affirming the effectiveness of this 3D position encoding. In the absence of $E^A$, ms-ESM struggles to differentiate molecules with identical atoms but distinct structures, e.g., isomer [5], resulting in a performance drop from 0.546 to 0.537 (refer to Table 7 in our paper).

Reviewer comment (Theory):
Also in Section 2.4, my guess is that the scaling by $\sqrt{d_k}$ gets disrupted by the $E^A$ term. Can you comment on this latter point?

Author response:
**A4**:  
- That's an excellent question! Currently, we haven't encountered any issues with disrupted scaling. However, we acknowledge the importance of this matter and plan to delve deeper into it in our future research.  
- It's worth noting that the inclusion of encoding (as represented by $E^A$ in our paper) into attention mechanisms is a widely adopted practice across various models in different domains. For instance, both T5 [9] and Swin Transformer [10] utilize this operation in language and vision modeling, respectively. Similarly, in the realm of molecule modeling, Transformer-M [2] and Uni-Mol [1] also incorporate this operation, demonstrating excellent performance in their respective domains.

Reviewer comment (Experiments):
The ablation study in Table 7 shows almost no improvement from vanilla ESM to the "w/o RSPE in atoms" variant.

Author response:
**A5**: Thank you for your insightful question. **The model's behavior in this ablation experiment aligns with our expectations, and residue-scale positional encoding in atoms aids the model in distinguishing atoms from different residues.** Here's an analysis of the phenomenon:  
- The model's behavior aligns with our expectations:  
  - In protein-molecule tasks, the model only requires protein sequences at the residue scale as input. Consequently, the addition of residue-scale positional encoding to atoms has no impact on the model's performance.  
- Despite the lack of performance improvement from adding residue-scale positional encoding to atoms, we choose to retain this feature in the model. This decision is driven by the understanding that residue-scale positional encoding aids the model in distinguishing atoms from different residues. In the future, we will also explore more diverse and effective approaches to integrate residue-scale information into atomic representations.

Reviewer comment (Experiments):
Did you consider the idea of representing the structure of the residue with a graph neural network? What could be the up/downsides of this approach?

Author response:
**A6**:  
- Indeed, we did consider the idea of employing a graph neural network (GNN) for representing residue structures. However, in alignment with previous research, we opted for a Transformer-based architecture in this paper. While we acknowledge the promise of using GNN for residue representation, we chose the Transformer due to its established success in molecular modeling.  
- The Transformer-based architecture is a well-established and reliable choice in molecular modeling, serving as the foundation for many successful methods such as the ESM series and Uni-Mol. Utilizing this architecture, aligned with these well-established methods, offers a more direct and secure approach. It facilitates the incorporation of their experiences into our work, mitigating the risks associated with project development. Choosing a GNN as the backbone might introduce unnecessary complexities, such as the need to train the entire model from scratch.  
- We recognize the strengths of GNN in capturing detailed structural information. For instance, GNN allows us to define edge features to capture chemical bonds in residue structures. Furthermore, it enables us to specify aggregation strategies for modeling information related to functional groups and motifs. Despite not being the focus in this paper, we believe GNN holds promise for future exploration in representing residue structures.

Reviewer comment (Evaluation):
The proposed method could not outperform baselines, as shown in Table.5 and Table.6.

Author response:
**Q1**: The proposed method could not outperform baselines, as shown in Paper Table.5 and Paper Table.6. 

**A1**: Please refer to the Q1 of the general response for more discussion. Briefly, i) as already mentioned in the answer to Q1 in the general response, this paper focuses on unified multi-scale modeling, aiming to obtain better results on protein-molecule tasks  instead of improving molecule/protein tasks universally. ii) We report molecule-only and protein-only results to verify that ms-ESM does not suffer from significant performance drops on separated molecular and protein modeling tasks.

Reviewer comment (Experiments):
Insufficient experiments regarding molecular representation learning affects the significance of the paper.

Author response:
**Q2**: Insufficient experiments regarding molecular representation learning affects the significance of the paper. Could you provide head-to-head comparisons against Unimol on molecular representation learning tasks, such as BBBP, BACE, Tox21?

**A2**: Thank you very much for your suggestions. We did a quick run to fix your concern. We don't include the results you mentioned because this paper focuses on unified multi-scale modeling for obtaining better results on protein-molecule tasks, and we only report results on  several selected representative molecule-only tasks to verify that the proposed ms-ESM does not suffer from significant performance drops on separated molecular modeling tasks. However more experimental results on molecular representation will be definitely helpful for the analysis. The added results are shown in Table A and we will also add them to the manuscript. 

Table A: Molecule-only tasks
|  | QM7 ↓ | QM8 ↓ | QM9 ↓ | BACE↑ | BBBP↑ | TOX21↑ | PCBA↑ | SIDER↑ | HIV↑ | MUV ↑ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Uni-Mol w/o H | 58.8678 | **0.0160** | **0.00540** | 83.20 | **71.52** | **78.92** | **88.12** | 57.71 | **78.3** | 72.0 |
| ms-ESM | 60.8755 | 0.0166 | 0.00590 | **83.52** | 67.41 | 75.39 | 86.15 | **63.59** | 74.9 | **72.6** |
| ms-ESM 8M | **52.8018** | 0.0173 | 0.00630 | 75.97 | 67.30 | 76.33 | 85.90 | 62.64 | 74.3 | 70.73 |

Reviewer comment (Experiments):
How about scaling the 35 million model to 650 million? Could you provide the corresponding results?

Author response:
**Q3**: How about scaling the 35 million model to 650 million? Could you provide the corresponding results? 

**A3**: We recognize the importance of understanding how changes in model size affect its performance. However, due to the substantial computational resources required to train the 650M model—an estimated 16 A100 GPUs over two weeks—we opted to train a smaller, 8M model to swiftly gather results and examine the impact of model scaling. Our intention is to extend this initial research by training larger models and including those findings in the final version of the paper.

Our observations are as follows:
- For tasks involving both proteins and molecules (as shown in Table B) and for tasks involving only proteins (as shown in Tables C and D), we found a clear trend: as the model size grows, there is a noticeable enhancement in performance.
- For tasks involving only molecules (as presented in Table A): 1) The larger model tends to perform better on the majority of these tasks, correlating increased model size with improved outcomes. 2) However, a few tasks saw superior performance from the smaller model, notably QM7 and TOX21, suggesting that the larger model may be susceptible to overfitting in some cases.

In conclusion, **we generally observed that model performance tends to improve with an increase in the number of parameters**, with a few exceptions. We plan to validate these findings by evaluating the 650M parameter model upon its completion to further test the performance of a larger-scale model.

Table B: The new results on enzyme-substrate affinity regression task
| Method | Protein Pre-training | Molecule Pre-training | MSE↓ | R^2↑ | Pearson↑ |
| --- | --- | --- | --- | --- | --- |
| XGBoost | ESM-2 35M | Uni-Mol 48M | 0.652 | 0.528 | 0.727 |
| ProSmith | ESM-2 35M | Uni-Mol 48M | 0.642 | 0.536 | 0.733 |
| XGBoost | ms-ESM 35M | ms-ESM 35M | 0.623 | 0.548 | 0.742 |
| ProSmith | ms-ESM 35M | ms-ESM 35M | **0.599** | **0.566** | **0.753** |
| XGBoost | ms-ESM 8M | ms-ESM 8M | 0.63188 | 0.54213 | 0.7269 |
| ProSmith | ms-ESM 8M | ms-ESM 8M | 0.61761  | 0.55247  | 0.7344 |

Table C: New experimental results on contact map prediction. SR: Short Range, MR: Medium Range, LR: Long Range.
|   |  SR P@L|  SR P@L/2 | SR P@L/5 | MR P@L | MR P@L/2 |  MR P@L/5 | LR P@L | LR P@L/2 | LR P@L/5 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| ESM2 35M | 0.197 | 0.293 | 0.458 | 0.223 | 0.318 | 0.452 | **0.297** | **0.385** | **0.492** |
| msESM 35M | **0.206** | **0.311** | **0.481**  | **0.225** | **0.320** | **0.453**  | 0.289 | 0.375  | 0.479 |
| msESM 8M | 0.153 | 0.213 | 0.310 | 0.143 | 0.191 | 0.259 | 0.153 | 0.197 | 0.256 |

Reviewer comment (Reproducibility):
To me, the presentation, especially the experiments part, is not clear. For example, the authors use ‘for more details of …., readers can find them in …’ many times, but this really restricts me from understanding the implementation details. A better way could be ‘following …, we fine tune … using ….’ In addition, please list the data size of downstream tasks.

Author response:
**Q1**: As many implementation details are omitted, the presentation in the experiments section is not clear enough.

**A1**: Thanks for your kind notes. Previously, we cited corresponding papers for aligned details but did not include these details in our paper. **We have added these details in the revised manuscript to make it self-contained** and list them as follows:
- Extra implementation details of protein-molecule tasks
  - *Fine-tuning dataset.* Following ProSmith [1] (section 3.1), we finetune ms-ESM and all baseline models on dataset KM [2], Davis [3], and ESP [4] for enzyme-substrate affinity regression, drug-target affinity regression, and enzyme-substrate pair classification respectively. The KM dataset contains experimental affinity constants of 11,676 enzyme-substrate pairs. The Davis dataset provides 30,056 binding affinities for pairs of 72 drugs and 442 proteins. The ESP dataset consists of 68,754 positive/negative enzyme-substrate pairs with experimental evidence. We use the standard data split provided by ProSmith in fine-tuning.
  - *Fine-tuning framework.* As mentioned in section 3.1, we use ProSmith's framework for a fair comparison. Specifically, the framework contains three main modules, i.e., molecule encoder, protein encoder, and fusion block. Two encoders extract features from proteins and molecules severally. The fusion block is a Transformer model, which is responsible for fusing protein and molecule features. The fused features are further used to regress the affinity values or predict binary affinity. We apply our model to ProSmith's framework by replacing both protein and molecule encoder with ms-ESM. We also provide the results of an XGBoost [12] variant of ProSmith, which removes the fusion block and uses simple concatenation for feature fusing. Note that we freeze both encoders in the experiments as suggested by ProSmith.
  - *Fine-tuning hyperparameters.* We directly use the hyperparameters provided by ProSmith. Specifically, the fusion block for three tasks has 6 layers of Transformer whose hidden size is 768. The epoch number is 100 and the learning rate is 1e-5. The batch sizes of the three tasks are 12, 12, and 24. We use Adam [14] as the optimizer for ProSmith and GBDT [13] with 500 iterations as the predictors for XGBoost.
- Extra implementation details for protein-only tasks
  - *Fine-tuning dataset.* Following TAPE's protocol [5], we evaluate ms-ESM on secondary structure prediction and contact prediction tasks (section 3.2). Specifically, for secondary structure prediction, we use data from [6] as training and validation sets and use CB513 [7], CASP12 [8], and TS115 [9] as test sets. The training and validation sets are filtered at the 25% sequence identity threshold with these test tests. The final training, validation and three test sets have 8678, 2170, 513, 21, 115 protein sequences, respectively. For contact prediction tasks, we use training, validation, and test sets from ProteinNet [10] with training and validation sets filtered at the 30% sequence identity threshold. The final training, validation, and test sets have 20, 24, 13945 protein sequences.
  - *Fine-tuning framework.* As suggested by TAPE, for both protein-only tasks, we use ms-ESM as the protein encoder. When doing secondary structure prediction, we use a linear output layer to predict the secondary structure each residue belongs to. When handling the contact prediction task, we use the attention from the last layer as features and then use a linear layer to predict whether these two residues have contact or not. Notably, both input of these two tasks is only protein sequences without structural information. Therefore, when using ms-ESM to handle these two tasks, we turn off the unzip.
  - *Fine-tuning hyperparameters.* We set up all the hyperparameters aligned to TAPE. For secondary structure prediction, the epoch is 5,000, batch size is 10, and learning rate is 0.001. For contact prediction, the epoch is 5, batch size 64, and learning rate is 3e-5. We use AdamW [15] as the optimizer in secondary structure prediction and Adam in contact prediction.

(To be continued.)

- Extra implementation details for molecule-only tasks
  - *Fine-tuning dataset.* We use the fine-tuning data of Uni-Mol [11] to evaluate the molecule understanding ability of ms-ESM (section 3.3). Specifically, we use QM8, QM9 datasets for molecular property regression and HIV, MUV datasets for molecular property classification, which have 21786, 133885, 41127, 93087 molecules, respectively. The data split is also provided by Uni-Mol.
  - *Fine-tuning framework.* Following Uni-Mol, a special token, i.e., [CLS], also exists in ms-ESM. Similar to NLP/CV, we simply use the representation of [CLS] to represent the whole molecule, and then use a linear head for fine-tuning on downstream tasks. For each molecule, we use the 3D conformation provided by [11] as the input of ms-ESM. In the fine-tuning stage, we do not add noises to atom coordinates.
  - *Fine-tuning hyperparameters.* For a fair comparison, we did not search the best hyperparameters. Instead, we set up all the hyperparameters aligned to Uni-Mol. Specifically, the batch sizes for these four tasks are 32, 128, 256, and 128. The learning rates are 1e-4, 1e-4, 5e-5, and 2e-5. The training epochs are 40, 40, 5, and 40. We use Adam optimizer for all the tasks.

Reviewer comment (Reproducibility):
About the ms-ESM model: during pre-training, what percentage of residues are unzipped?

Author response:
**Q3**: During pre-training, what percentage of residues are unzipped?

**A3**: Thanks for pointing it out. Due to the page limitation, we report it in Table 8 of Appendix A: 6.5% of residues are unzipped as the main experimental setting. We have added it with other details in the main text.

Reviewer comment (Experiments):
Ablation: The pair-wise distance recovery is only used at atom scale and requires atom coordinates as inputs. How about removing this loss? What is the performance? Is this term necessary?

Author response:
**Q4**: The pair-wise distance recovery is only used at atom scale and requires atom coordinates as inputs. How about removing this loss? What is the performance? Is this term necessary?

**A4**: Thank you very much for your good questions. Actually, the pair-wise distance recovery loss is the most widely used pretraining objective for captureing geometric molecular information[11][19]. As you suggested, we did a quick run by removing it for ablation and the results are shown in Table A and we also have added them to our revised manuscript.

- Results: As shown in Table A, when we do not apply pair-wise distance recovery loss, the performance of ms-ESM shows a significant decrease in the protein-molecule task.
- Analysis:
  - For molecular representation learning
    - Without using pair-wise distance recovery loss, the training loss of ms-ESM on the atomic scale is only masked atom type prediction loss. But more than half of the atoms in a molecule are carbon atoms, which makes the model learn very little information on the atomic scale using only masked atom type prediction.
  - For modeling unzipped residues
    - Without using pair-wise distance recovery loss, ms-ESM will not be able to learn structural information from the unzipped residues. ms-ESM can only learn the mapping from residues to atoms' types. Since there are only a few types of residues, and the mapping is also fixed for a given residue, this training task is too easy to make the model learn meaningful information.

Table A: Ablation study on pair-wise distance recovery loss
|  | MSE ↓ | R^2 ↑ |
| --- | --- | --- |
| Vanilla ms-ESM | 0.627 | 0.546 |
| w/o pair-wise distance recovery loss | 0.645 (+0.018) | 0.533(-0.013) |

Reviewer comment (Experiments):
The performance on protein-molecule tasks: I can’t intuitively understand why ms-ESM can outperform two separate pre-trained models (one for protein, and one for molecule)? Basically, I think the capacity of two models should be greater than a single model. Is the comparison fair? Please provide more explanation.

Author response:
**Q5**: Can’t intuitively understand why ms-ESM can outperform two separate pre-trained models (one for protein, and one for molecule)? Basically, I think the capacity of two models should be greater than a single model.

**A5**:
- Yes, the capacity of two separate models is larger than one unified model. ms-ESM performs better because **the benefits of unified modeling inside one model (with specific designs) make it learn the unified context jointly and implicitly, which outweighs its shortcomings in capacity.**
- The above observation has also been found in multi-lingual learning and multi-modal learning areas. Specifically, [18] shows that a unified multi-lingual machine translation model (one model for all language pairs) can outperform translation models trained separately. [17] demonstrates that modeling vision and language directly inside one model can beat methods that fuse the two modalities later.
- Moreover, visualization in our paper (Figure 4) also indicates that our unified ms-ESM can provide more harmonious protein/molecule features than two separate models, i.e., EMS-2 and Uni-Mol.

**Q6**: Is the comparison in protein-molecule tasks fair?

**A6**: 
- Yes, the comparison is fair. As mentioned in Section 3.1, we follow the benchmark protocol from ProSmith [1] to evaluate ms-ESM and all the baseline methods. In other words, the settings of all the models being compared are well-aligned.
- Smaller ms-ESM outperforms two larger models, which proves the benefit of unified modeling even more significantly. Specifically, ms-ESM is at a disadvantage of the model capacity compared to ESM-2 + Uni-Mol (35M < 35M + 48M). However, Table 1, 2, and 3 in our paper clearly show that ms-ESM beats ESM-2 + Uni-Mol.