Endowing Protein Language Models with Structural Knowledge

The main reported results in the paper (Table 1), where the model's performance is compared to ESM-Gearnet MVC, show only a slight advantage, and this advantage is likely due to the use of a stronger backbone model (ESM-2 vs. ESM-1b). Therefore, the results may not be very persuasive in demonstrating a good performance improvement for the PST model.

In our new results, detailed under point 3 in the general comments, we have observed that PST models with fixed representations outperform all fine-tuned models, including ESM-Gearnet MVC, across all function prediction tasks. This finding combined with the enhanced practicality in terms of reduced computational resource requirements, clearly underscores the advantage of PST over fine-tuned models like ESM-Gearnet-MVC.

Moreover, based on the results in a very recent preprint (https://arxiv.org/pdf/2303.06275.pdf), that we cite in our revised manuscript, we observe that the performances of ESM-Gearnet-MVC based on ESM-1b are very similar or even better than the one obtained using ESM-2, on most of the considered tasks.

The approach of incorporating protein structural information into the ESM model using a GNN has been previously proposed in other papers[1]. While this paper just applies this approach to the field of protein representation learning, it lacks novelty, and it does not discuss the differences between these methods.

First, it is worth highlighting that our proposed architecture is significantly different from LM-Design, as elaborated in point 2.i of the general comment. Second, the central objective of LM-Design centers around protein design such as inverse folding, whereas our work primarily aims to enhance protein representations specifically for various tasks related to protein function prediction.

Table 1 does not report the performance of ESM-2 under end-to-end training conditions, and it does not specify important hyperparameters of the compared models, such as model size, which is a crucial factor affecting relative model performance.

We have added the performance of ESM-2 under end-to-end training, and we notice that its performance is inferior compared to the PST (finetuned). Moreover, we have added the number of parameters of the ESM-2 model and of the corresponding PST model, as well as the ones of the baselines we run ourselves, in the appendix B.4. The number of parameters of the other models are not reported in the Gearnet paper.

Has the paper explored the impact of different backbone models on performance? For example, ESM-1b vs. ESM-2.

We appreciate the reviewer's suggestion to explore the impact of various backbone models. Nonetheless, our current focus is on the state-of-the-art PLM, namely ESM-2, as it presents more relevance and potential for advancement compared to its predecessor, ESM-1b. Moreover, our investigations into the impact of model size have indicated that the quality of the base sequence model is directly proportional to the performance of the corresponding PST in general. Hence, it is a reasonable inference that a PST model built upon ESM-2 would likely surpass one based on ESM-1b.

Since there are pre-trained GNN models designed for protein structural graph, such as GearNet[1], did the paper investigate the performance impact of using such pre-trained GNN models?

We would like to point out that in contrast to previous approaches that use a deep GNN to encode the structural information and add the encoded features atop a PLM, PST adopts a different approach. Specifically, we use a different two-layer GIN model (as described in section 4.1.2) at each self-attention layer within the base ESM-2 model. Therefore, the weights learned in each GIN model should differ from one layer to another, making the application of a deep, pre-trained GNN model inapplicable in this context. We have clarified this by adding a brief introduction of our PST method at the beginning of Section 3.2 in the revised manuscript. See also our general comments above.

A controlled comparison between ESM-2 and PST might have considered continuing MLM pre-training for ESM-2 for an equivalent number of steps as PST pre-training. However, the experiments showing that tuning only the structure-related parameters (section 4.4) is sufficient to improve performance partially address this.

The ESM-2 models have been pre-trained on the Uniref-50, a dataset much larger than the structure dataset (i.e. Swissprot) used for pre-training the PST models. Consequently, further MLM pre-training of ESM-2 on Swissprot could potentially lead to overfitting, negatively affecting model's performance in downstream tasks. Moreover, we believe that "updating only the structure-related parameters is sufficient to improve performance" should be interpreted as a strength of our method, rather than a weakness, as elaborated in the point 2.iv in the general comments.

While the improvement from PST over ESM-2 is largest for smaller model sizes, the efficiency improvement is partially offset by the need to predict a structure.

While we acknowledge this is a good point, it is worth noting that structures can generally be sourced from precomputed databases (such as ESMFold and AlphaFoldDB), and are reusable for various tasks. Therefore, we argue that the additional computational costs due to structure predictions are quickly compensated by the reduced costs due to the data- and parameter efficiency of our models. Moreover, the effectiveness of our model based on fixed embeddings, without any additional finetuning, significantly reduces computational demands compared to fine-tuned models. This further strengthens our argument regarding the cost-effectiveness of our approach.

nit: The BERT citation is not formatted correctly.

Thank you for noticing this point, we have addressed the formatting issue in the revised manuscript.

While the integration of a structure extractor with a PLM is presented as a unique proposition for function prediction tasks, it exits parallel studies in the protein representation field. For instance, LM-Design [1] proposes a similar structure adapter into PLMs that endows PLMs with structural awareness, and the structure adapter could access an arbitrary additional structure encoder (GNNs, ProteinMPNN etc.). As a result, this might raise concerns about novelty in comparison to existing methods.

First, it is worth highlighting that our proposed architecture is significantly different from LM-Design, as elaborated in point 2.i of the general comments. Second, the central objective of LM-Design centers around protein design such as inverse folding, whereas our work primarily aims to enhance protein representations specifically for various tasks related to protein function prediction.

A limitation in the study is the lack of an empirical exploration regarding the selection of different structure encoder. While the paper presents results based on a specific structure extractor module, it would have been insightful to see comparisons or benchmarks against various other structure encoding methodologies.

Thank you for the suggestion. Our research draws significant inspiration from Chen et al. 2022 [1], leading us to choose GIN as our structural extractor due to its balance of performance and efficiency. Additionally, we believe that the quantity of structural information used in the structural extractor is more influential than the specific type of structural extractor chosen. We thus explored the impact of varying amounts of structural information in our study, as demonstrated in Section 4.4 and Figure 2.

The paper primarily focuses on the integration of a structure extractor with the ESM-2 model. However, with the availability of larger models like ESM2-3b, have there been empirical studies to assess the impact and advantages of the structure extractor? Specifically, as the scale of the PLM increases, does the benefit of adding a structure extractor diminish or remain consistent? It would be crucial to understand the interplay between model size and the structure extractor's efficacy.

Our work includes scalability studies for various model sizes that can be accommodated by most GPUs, as depicted in Figures 2, 3, and 4. These studies indicate that the benefits of utilizing structural extractors are maintained across different scales. Although the advantages might diminish slightly for very large models, the utilization of structural information is still expected to offer benefits. Regarding extremely large models like ESM2-3b, we were only able to perform inference with it, as fitting a PST model of this size exceeded our capabilities. Interestingly, in the Variant Effect Prediction (VEP) task, our results show that ESM2-3b is outperformed by smaller versions of ESM2 models. For instance, ESM2-650M and the corresponding PST achieved a higher mean Spearman's correlation (0.489 and 0.501 respectively) compared to ESM2-3b's 0.456, highlighting the effectiveness of smaller models in certain applications.

Reference

[1] Dexiong Chen, Leslie O’Bray, and Karsten Borgwardt. Structure-aware transformer for graph representation learning. ICML 2022.

Overall speaking, the novelty of this paper is enough. There are plenty of works that integrated the graph representation in protein representations, even alphafold. The authors mentioned about the minimal training cost and parameter efficient, while I acknowledge the good point. This is hard to be a strong strength.

Please see the point "Novelty and impact of our work" in the general comments.

The framework actually is a GIN + ESM model, with the ESM freezed and GIN and head tuned. In Figure 1, however, the GIN is not clearly presented. If GNN in the figure is the extractor, it means the GNN would be processed multiple L layers, this is somehow not necessary or why for this design?

Note that in the transformer architecture, the node embeddings used for each GNN corresponds to the residue embeddings of the respective layer, and these embeddings can vary from one layer to another. Consequently, the query, key, and value matrices have to be re-computed at each layer. Please see also our first point of the general comments for a clarification of the PST architecture.

Small question, the GIN is randomly initialized, what's the difference between zero-initialized as you mentioned in the paper? Besides, is the node embedding in GIN is initialized by ESM token embedding?

See Clarification of the PST architecture in our general response, where we clarified our general approach, and should address general questions about the methodology. Now to answer the reviewer's questions: i) the weights in GIN, i.e., $\theta$ is randomly initialized while $W_s$ is initialized to zero such that the PST model without training makes the same prediction as the base ESM-2 model. ii) The node embedding in the GIN before each self-attention layer is given by the residue embedding at the respective layer. We have added a clarification of PST in Section 3.2 of the revised version.

The paper does not provide a detailed analysis of the limitations of the proposed framework or potential areas for improvement.

We kindly direct the reviewer's attention to the limitations discussed in section 4.4 and Appendix A.

1. The paper writing style needs to be further improved.

We thank the reviewer for the suggestion and have updated the manuscript accordingly (please see Clarification of the PST architecture in our general comments). As a result, a thorough description of the ESM-2 model is crucial as it lays the groundwork for the mechanism discussed above.

2. The introduction of structural knowledge of 3D protein structures is limited.

We acknowledge the reviewer's point about the importance of using additional geometric information in several structure prediction tasks. However, its impact on protein function prediction, the primary focus of our work, remains less evident. To elucidate how the incorporation of different amounts of geometric information impacts PST models, we direct the reviewer's attention to our ablation study presented in Section 4.2 under "Amount of structural information required for refining ESM-2". Our results suggest that while incorporating more structural features, such as the distance information used in PiFold [1], improves training accuracy, it negatively affects the model's performance in downstream tasks, as shown in Figure 2.

3. The paper only adopts the subset of AlphaFoldDB.

We did not train our PST model on ESMFold, which is out of the computational capacity in our group. However, our results indicate that fine-tuning of ESM-2 with structural data, even on a relatively modest-sized structure dataset such as Swissprot (~550k structures) can already enhance the quality of protein representations. We are eager to explore the full potential of PST on larger datasets like ESMFold in future work.

4. The downstream tasks adopted in this work are limited.

• First, our method was evaluated on a wide range of tasks, such as enzyme commission number prediction and GO term prediction given a protein as used to evaluate Gearnet and DeepFRI. To extend our analysis, we use the novel ProteinShake tasks to evaluate the quality of the embeddings on other biologically meaningful tasks such as binding site prediction, structural class prediction and protein family prediction. We believe that these tasks collectively provide a thorough overview of a wide range of protein-related applications.
• Second, our focus was primarily on protein function prediction task, which is a common focus of many representation learning approaches such as Gearnet and DeepFRI, allowing us to benchmark our representations against existing approaches.
• Third, we found that the suggested benchmarks did not contain UniProt IDs which prevented us from retrieving the exact corresponding structures from AlphaFoldDB needed for direct comparison on the Ankh datasets. Instead we were able to apply the Ankh model to our datasets for the EC and GO tasks (see updated Table 1). Regarding xTrimPGLM, we could not easily find the model weights online.

5. Besides the ablation studies claimed in the pre-training strategy, the author should add Seq-only to get rid of the reason solely brought by the sequence-based training.

First, we want to clarify that in the "pretraining strategies" section, the "Struct Only" strategy specifically consists of updating only the structure extractors in the PST model. In this strategy, the weights of the base ESM-2 model remain unchanged during training, allowing the trained PST model to also generate the sequence representations (which is the same as the base ESM-2's representations) at inference, by bypassing the structure extractors in the trained PST. The "Struct Only + Seq" strategy simply takes the average of both structure and sequence representations. This was fully detailed in our submitted version. We have further included Figure 6 to illustrate the difference between these strategies in the revised manuscript.

As for the "sequence-only training" referenced by the reviewer, we believe this corresponds to the ESM-2 model. We encourage the reviewer to refer to Section 4.3 of our paper, where we present an extensive comparison between PST and ESM-2.

3. Performance compared to finetuned models: In our submitted version, we opted for simplicity and efficiency by reducing the representation dimensions to 2048 using PCA, and did not perform any hyper-parameter tuning on the downstream tasks. This did not lead to a fair comparison with ESM-Gearnet which exhibited extensive hyper-parameter tuning. Therefore, our revised version eliminates the dimension reduction, using the full-scale representations instead. Additionally, we have incorporated dropout layers in the MLP classification head to prevent overfitting. New results are updated in Table 1 of the revised manuscript as well as in the following table (with Fmax/AUPR as the metrics). Our new findings underscore that PST with fixed representations surpasses all fully fine-tuned models in terms of AUPR, notably in the GO-BP and GO-CC task. This outcome holds considerable importance for two reasons: i) it contrasts with trends in NLP and computer vision, where fine-tuning typically prevails over fixed representation models; ii) The enhanced practicality of PST with fixed representations, especially in terms of reduced computational resource requirements, signifies a clear contribution in the field.