Paper: Endowing Protein Language Models with Structural Knowledge (ICLR-2024)

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
Has the paper explored the impact of different backbone models on performance? For example, ESM-1b vs. ESM-2.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Presentation):
nit: The BERT citation is not formatted correctly.

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

Reviewer comment (Presentation):
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?

Author response:
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.

Reviewer comment (Reproducibility):
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 initialized by ESM token embedding?

Author response:
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.

Reviewer comment (Evaluation):
The paper does not provide a detailed analysis of the limitations of the proposed framework or potential areas for improvement.

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

Reviewer comment (Writing):
The paper writing style needs to be further improved. The Method spends a lot of pages introducing the principle of ESM-2, which is not the contribution of the paper. On the contrary, the introduction of the PST framework is vague. How the node embedding output by GNN is further incorporated into the residue embedding of the original ESM? Did they just add, concentrate? or viewed as the individual token embedding which is conducted with self-attention to the residue embedding of the original ESM? I highly suggest the authors to add more details.

Author response:
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.

Reviewer comment (Novelty):
The introduction of structural knowledge of 3D protein structures is limited. This paper just compressed the 3D structures into 2D graphs without considering the 3D geometric features, such as the SE(3)-Equivariant features, which have already been confirmed to be critical in modeling the 3d protein structures or the protein-protein docking patterns by many works, such as Alphafold2, EquiDock, etc. The only considered feature is the distances between nodes severed as the edge attributes while performing less-satisfied on the downstream tasks.

Author response:
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.

Reviewer comment (Experiments):
The paper only adopts the subset of AlphaFoldDB. I wonder if the author ever tried other structural databases such as metagenomic protein databases predicted by ESMFold? How do the data quality and quantity affect the performance of PST?

Author response:
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.

Reviewer comment (Experiments):
The downstream tasks adopted in this work are limited. As the protein structure information decides the physical and chemical properties of proteins, the downstream tasks should not only be performed on protein structure prediction tasks, but also on other tasks, like protein solubility prediction, secondary structural prediction, etc. Like xTrimoPGLM (chen et al.) and Ankh (Elnaggar, et al) do.

Author response:
- 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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Evaluation):
"While PST typically surpasses its base ESM-2 counterpart at similar model sizes, this performance gain tapers off with increasing model size." Does this observation indicate that if we adopt a huge amount of protein sequence data and also the large scale of PLMs, the structural knowledge can be well-captured on the sequence-based pertaining, like xTrimoPGLM with 100B dominates on most of the downstream tasks?

Author response:
- First, this observation mentioned does not imply that "the largest models always perform the best in all tasks". For instance, in tasks such as structural class and protein family prediction (shown in Figure 3), models incorporating structural knowledge can outperform even very large sequence-based models. Therefore, the benefit of incorporating structural knowledge remains significant.  
- Second, We acknowledge the relevance of huge protein language models like xTrimoPGLM, and have added a discussion in Section 2. While scaling up sequence models might seem advantageous, it also brings practical challenges, notably in terms of data and computational requirements. In contrast, structure-focused models like PST are more efficient in both data usage (e.g. PST was only trained on Swissprot) and computational resources. Within modest GPU memory constraints, PST using fixed representations generally surpasses the performance of state-of-the-art sequence models like ESM-2.