Thank you very much for recognizing the novelty and contribution of our work. Your insightful comments helped us enrich the analysis a lot. With the response below, we hope to address your concerns properly.

Weakness:

W1. In fact, we do not need to train the structure encoder and the clustering model for every protein. Our structure quantization serves as a data preprocessing step. The codebook is only pre-trained on a pre-training structure encoder dataset (Appendix A.2, Dataset for training structure codebook). SaProt uses FoldSeek [1] as its structure codebook, which also needs pre-training. Since training a codebook is FoldSeek's task, it isn't explicitly detailed in SaProt's article. The primary difference between our structure quantization module and FoldSeek is that our local structure considers up to the nearest 40 residues within a 10 Å distance, whereas FoldSeek only incorporates structure information between preceding and succeeding residues.

Q1. Tokenizing a protein structure containing $L$ residues involves three steps:

Step 1. For $i = 1, 2, 3, \ldots, L$, generate a local structure for residue $r_i$, converting it into a graph $G_i$. (Left side of Figure 2c)

Step 2. For $i = 1, 2, 3, \ldots, L$, encode each graph $G_i$ into a continuous vector $e_i$ using the trained structure encoder.

Step 3. For $i = 1, 2, 3, \ldots, L$, utilize the trained clustering codebook to convert $e_i$ to a token $s_i$.

We also evaluated the time required to convert a protein into structure tokens for five proteins of different lengths using a server equipped with an RTX 1080 GPU and a Intel Xeon E5-2690 CPU. Step 2 was executed on the GPU, while steps 1 and 3 were executed on the CPU. The results are as follows.

We believe the speed is acceptable. Nevertheless, one of our important future tasks is to use parallelization techniques to optimize the speed of subgraph segmentation (Step 1) and local structure encoding (Step 2). We are working on this development and will release it in the camera-ready version.

Q2. Both our structure quantization module and FoldSeek involve two steps: training a structure encoder and a clustering model. We chose the denoising autoencoder for its focus on representation learning and simplicity, which better suits our needs. While VAE and VQ-VAE are generative models, we do not require a structure generation model. The training curve of our structure encoder is shown in Figure R9. The change of training loss and validation loss is stable.

Q3. Smaller neighborhood sizes may ignore some important neighbors. Although using neighborhood sizes is a route to accelerate average pooling, it may mislead some neighboring nodes. We selected the nearest 40 residues within a distance of 10 Å, according to [4]. As shown in Figure R1, local structures with more than 40 neighbors account for only 0.00052% (5.2e-6). Thus, 40 can cover most cases.

References

[1] Clustering predicted structures at the scale of the known protein universe.

[2] Convolutions are competitive with transformers for protein sequence pretraining.

[3] https://www.uniprot.org/uniprotkb/statistics

[4] Discovery of Novel Gain-of-Function Mutations Guided by Structure-Based Deep Learning.

We sincerely appreciate your thoughtful feedback on our work and provide a thorough response addressing your main concerns.

Weakness

W1. We agree that our work is not the first hybrid approach to protein language models, and we have cited works such as ESM-GearNet and LM-GVP in Section 2.1. However, what sets our work apart is our disentangled attention mechanisms and structure quantization module. We redesigned the self-attention mechanism in the Transformer to accommodate multiple sequence inputs, including protein sequence, structure sequence, and relative position matrix sequence. Meanwhile, the structure quantization module can encode local structure of residues with more comprehensive information. Evaluation results demonstrate the efficacy of our model across various protein understanding tasks, especially in zero-shot mutant fitness prediction. We contend that the overall model, local structure encoding and quantization, the disentangled attention mechanism, and the evaluation results jointly represent a valuable contribution to the field of computational biology.

W2. We evaluated ESM-GearNet on the fine-tuning task and compared it to ProSST. The results are as follows:

The scores of ESM-GearNet are derived from SaProt. Note that the slightly different scores compared to the manuscript are due to the updated random seed and dataset. We have re-evaluated ProSST with different seeds in order to compute the average performance. Additionally, the evaluation was conducted on the updated GO dataset by SaProt.

W3. The clustering model is trained only once. The training and inference speeds of the clustering model are as follows:

We also show the inference time of ProSST(110M), SaProt(35M), SaProt(650M) in Figure R8. ProSST(110M) is faster than SaProt (650M) and slower than SaProt (35M). The experiments are conducted on a server with one RTX 3090 GPU and two Intel 8468 CPUs. We will discuss inference speed in revised paper.

Questions

Q1. ProSST utilizes relative positional encoding, which supports inference without constraints on sequence length. For pre-training, we removed proteins longer than 2048 residues for training efficiency (extremely long sequences may cause an overflow of CUDA memory). However, we do not do any truncation during fine-tuning and zero-shot mutant effect prediction. Furthermore, we have evaluated the perplexity of ProSST on a long protein dataset containing 594 proteins longer than 2048 residues. The statistics of the sequence lengths and evaluation results are as follows:

The perplexity is 9.013, which is similar to the validation set perplexity of 9.033, suggesting that ProSST can effectively understand long protein sequences. Another notable point is that extremely long protein sequences are very scarce in the nature, as almost 99.7% of proteins are shorter than 2048 [1]. We will add the discussion about long sequences to the manuscript.

Q2.: The current implementation requires to input both sequence and structure. Extending the framework to accept sequence-only datasets is possible. Here, we provide three approaches:

Approach 1: Use AlphaFold or ESMFold to predict structures.

Approach 2: Use ProSST(MST), trained with structure masking and supporting sequence-only inputs. We will release a model ProSST(MST), an extension of the ProSST model that incorporates Masked Structure Training (MST). The MST means that during pre-training, each sample's structure sequence has a 50% probability of being replaced by a fully masked sequence [1,1,1,1,1,…,1], simulating missing protein structure. Therefore, when applying ProSST to sequence-only datasets, we need to use the masked sequence [1,1,1,1,1,…,1] as a substitute for the structure token sequence.

We have evaluated these methods on the ProteinGym benchmark, binary localization prediction (BLP) from the PEER benchmark, and perplexity on the validation set. The results are as follows:

Rows 1-2 show the performance differences between AlphaFold and ESMFold. Rows 3-4 show the performance of the new model ProSST(MST). Row 5 shows the performance of the sequence-only model. A discussion on how to deploy ProSST on sequence-only datasets will be added to the revised version.

Reference

[1] https://www.uniprot.org/uniprotkb/statistics

Thank you for your meticulous feedback. Below are our responses.

Weakness:

W1. We conducted additional experiments to analyze disentangled attention.

Experiment 1: We replaced all structure tokens in the test set with zeros or random numbers from a uniform distribution and re-evaluated ProSST. The results are:

Incorrect structure tokens decreased performance, suggesting that disentangled attention learned the sequence-structure relationship. Otherwise, performance would have been less affected.

Experiment 2: We train ProSST (K=1), where the structure tokens are replaced with a constant value of 1. This setup helps preserve the disentangled attention mechanism. If ProSST (K=1) still improves performance, it indicates that the improvement is solely due to the disentangled attention.

There is little difference between K=1 and K=0, as their perplexity curves (see Figure R3) almost overlap. This indicates that disentangled attention cannot improve performance without correct structure tokens.

Experiment 3: Visualizing the learned attentions. We show different types of attentions on Green Fluorescent Protein (GFP), including 238 residues, in Figure R4. We can see that disentangled attention learns different attention patterns, with notable differences between R2S and S2R.

W2.

(1) Dataset Split

(1.1) There is no test set for Pre-training for the structure codebook and Transformer. We use only a small validation set (without a test set) for pre-training to maximize data usage, similar to pre-trained models like ESM-2, which use 0.5% of data for validation [1].

(1.2) The fine-tuning datasets have train/valid/test splits like SaProt's and are downloaded from SaProt. Data statistics will be provided in the revised paper.

(1.3) In this zero-shot mutant effect prediction, all data are in the test set.

(2) Error and Significance Analysis

(2.1) We will add error analysis to fine-tuning datasets. We use the same hyperparameters and repeated experiments five times with different seeds. The average performance served as the metric, with the standard deviation as the error. The results are:

Note that the slightly different scores compared to the manuscript are due to the updated random seed. Additionally, the evaluation was conducted on the updated GO dataset by SaProt.

(2.2) The error in the ablation study tables will be updated in the revised version, which would not change the performance ranking.

(2.3) We used the non-parametric bootstrap method for zero-shot mutant fitness prediction to test differences between the baselines and ProSST. In all tests, the p-values were less than 0.01.

W3. The inconsistency between Table 3 and Table 4 was due to a copying error. The data for ProSST (K=2048) in Table 3 is correct, but an error occurred when copying to Table 4. The performance of ProSST in Table 4 on ProteinGym should be 0.504 for Spearman, 0.777 for NDCG, and 0.239 for Top-recall. We will correct these and fix the black marking issues in all tables.

W4. Disentangled attention cannot improve performance without correct structure tokens. We train ProSST (K=1), where the structure tokens are replaced with a constant value of 1. This setup helps preserve the disentangled attention mechanism. If ProSST (K=1) still improves performance, it indicates that the improvement is solely due to the disentangled attention. However, ProSST(K=1) is not better than ProSST(K=0). The results can be referred to response to W1.

W5. We will emphasize the reasons for selecting discrete structure tokens and using disentangled attention, and we will highlight the analysis of the relationship between them, as discussed in W1.

Questions:

Q1. We select baselines with high Spearman correlations from the ProteinGym webpage [2]. We also compared the mentioned baselines to ProSST and will update it in our paper. The results are as follows:

Since structure models often excel in the Stability subset, we additionally compared ProSST with other models in this subset. We will include other ProteinGYM subsets (Stability, Activity, Binding, Expression) in the Appendix. ProSST achieves state-of-the-art (SOTA) performance in the Stability, Binding, and Expression subsets.

Limitations:

L1. We believe our model is a valuable addition to structure-aware protein language models. Note that ESM-3 was introduced in June 2024, which was after our submission to NeurIPS in May 2024.

L2. We utilize AlphaFold to predict structures of disordered proteins. We provide the relationship between AlphaFold pLDDT and the performance of structure models like ProSST, SaProt, and ESM-IF on ProteinGYM, as shown in Figure R5-R7. There is a positive correlation between pLDDT and model performance: for ProSST, $\rho=0.30$; for SaProt, $\rho=0.31$; and for ESM-IF1 $\rho=0.42$, where $\rho$ is the Pearson correlation coefficient. This may indicate that structure models may not perform well on disordered proteins.

References

[1] Evolutionary-scale prediction of atomic-level protein structure with a language model.

[2] https://proteingym.org/benchmarks

We sincerely appreciate your thoughtful feedback on our work. We would like to address your questions and concerns as follows:

Weaknesses

W1. Feeding a "default/noise" structure input alongside the actual sequence was not included in the current version, but we have a solution for it, and this solution is not very sophisticated. We trained ProSST(MST), where MST stands for Masked Structure Training. During pre-training, each sample's structure sequence has a 50% probability of being replaced by a fully masked sequence [1,1,…,1]. This approach simulates the scenario of missing protein structures. Although ProSST(MST) outperformed ProSST(K=0) in sequence-only mode, it was inferior to both the original ProSST model with structure inputs and ProSST(MST) with structure inputs.

We offer two approaches for sequence-only proteins:

• Approach 1: Use AlphaFold or ESMFold to predict structures.
• Approach 2: Use the ProSST(MST).

We have evaluated these methods on the ProteinGym benchmark and perplexity on the validation set. The results are as follows:

Rows 1-2 show the performance of Approach 1. Rows 3-4 show the performance of the new model ProSST(MST). When using AlphaFold, ProSST (MST) is inferior to ProSST (K=2048). We will add a discussion on how to apply ProSST to the revised paper.

Questions:

Q1. The reason is that 40 covers most cases. We selected the nearest 40 residues within a distance of 10 Å, according to [1]. As shown in Figure R1, local structures with more than 40 neighbors account for only 0.00052% (5.2 x 10^-6), indicating that our choice covers most cases.

Q2. They are not the same. $L$ is the number of residues in a protein, also referred to as the length of the protein, while $l$ is the number of nodes in a graph. A protein contains $L$ residues, each corresponding to a local structure, which in turn corresponds to a graph $G$. For any arbitrary graph $G$, $l$ is the number of nodes in it.

Q3. The GVP encoder includes a six-layer message-passing graph neural network in which a geometric perceptron replaces the MLP to ensure translational and rotational invariance of the input structure. Our GVP encoder is consistent with the original GVP-GNN [2], except that we removed the residue type information. The GVP encoder is trained from scratch. We will provide detailed descriptions and parameterizations of the GVP in the Appendix.

Q4. Both choices are worse than discrete structure tokens. (1) Using the centroid embedding (K=2048) as input yields inferior results:

We believe it is because the Transformer requires a learned structure token embedding rather than a fixed embedding. (2) Directly using continuous local structure embeddings as structure inputs leads to overfitting, as shown in Figure R2. Similar results have also been observed in SaProt.

Q5. $l$ represents the total number of nodes in a graph $G$. Average pooling refers to averaging the embeddings of all nodes within the graph. $l$ is not always equal to 40. Because when there are not more than 40 residues within a 10 Å distance, $l$ will be less than 40.

Q6. The disentangled attention contains multiple regular attention mechanisms, allowing the model to learn different attention patterns for better contextual representation. The computation of the disentangled attention is also less than that of directly concatenated self-attention because the complexity of the attention mechanism increases quadratically with the sequence length. Although disentangled attention involves additional attention calculations, it does not increase the sequence length.

Q7. The $k$ in Eq.(2) should actually be $L_{max}$, which is the cutoff of relative position. We will correct it in the paper.

Q8. (1) Because our model aims to learn the contextual representation of residues rather than structure tokens, structure $s$ in Eq.(6) is left unmasked. (2) When calculating scores of variants, we utilize the structure of wild sequences because mutants only slightly alter the structure of proteins [3], and the structure of mutants is difficult to predict[4]. Other structure models, including ESM-IF, SaProt, ProteinMPNN, etc., also use the wild-type structure when scoring variant sequences.

Q9. We have compared ProSST to ESM-GearNet and ProSST(fixed parameters) in fine-tuning downstream tasks. The results are as follows:

The scores of ESM-GearNet are derived from SaProt. Note that the slightly different scores compared to the manuscript are due to the updated random seed and dataset. We have re-evaluated ProSST with different seeds to compute the average performance. Additionally, the evaluation was conducted on the updated GO dataset by SaProt.

Q10. We will correct them. Thank you.

References

[1] Discovery of Novel Gain-of-Function Mutations Guided by Structure-Based Deep Learning.

[2] Learning from Protein Structure with Geometric Vector Perceptrons.

[3] A folding space odyssey.

[4] Can AlphaFold2 predict the impact of missense mutations on structure?