Multimodal Distillation of Protein Sequence, Structure, and Function

Dear Reviewer G223,

We are grateful for your thorough review. Your comments are highly valued, and we would like to express our heartfelt gratitude. We do our utmost to address the questions you have raised:

Q1 It is unclear whether ProteinSSA makes use of pretrained embedding model, especially for the teacher model. The paper mentions training ProteinBERT with additional modalities, but generally claims that ProteinSSA does not require large-scale pretraining. Can you describe how you obtain the embeddings for each modality? Do you use pretrained models for some or all modalities?

A1 Thank you for your valuable feedback! (1) ProteinBERT [1] is not a part of ProteinSSA, and ProteinSSA does not make use of pre-trained embedding model for both the teacher model and the student model. (2) ProteinBERT is the first model to utilize the protein annotation information, using 8943 frequent Gene Ontology (GO) annotations. Compared with ProteinBERT, we use 2752 GO annotations from the GO dataset. We have modified this in the revised manuscript to decrease the misleading. The current state-of-the-art (SOTA) sequence-function pre-trained model is the KeAP [2], but the SOTA model KeAP fails to generate protein embeddings that surpass those obtained from the sequence pre-trained model ESM-1b [10]. This observation demonstrates the limitations of the sequence-function pre-trained model and motivates this work. (3) The main objective of the teacher model is to learn embeddings that contain functional information and provide intermediate supervision during knowledge distillation for the student model. Therefore, the complete training of the teacher model is not our primary concern. (4) Pre-training involves learning general representations from a large dataset. In this work, the teacher model is trained on a relatively small dataset consisting of about 30 thousand proteins with 2752 GO annotations. This is in comparison to ProteinBERT, which uses 106 million proteins, and KeAP, which uses 5 million triplets. (5) There are three protein modalities that we used, including sequence, structure, and function. Both the student and teacher networks use the proposed protein graph message-passing and pooling layers to process the sequence and structure. Thus, the sequence-structure embeddings are obtained from the designed graph neural network. In the teacher model, the function information is encoded by the annotation encoder to get the functional embeddings.

Q2 The paper does not compare results against larger scale protein models for relevant tasks, including the ones mentioned in related work (e.g., ESM, KeAP, ProtST). It would be good to get a sense of much model size affects performance on the studied tasks. Can you describe how large your model ends up being in terms number of trainable parameters?

A2 Thank you for your informative reviews! (1) Different from those pretraining or self-supervised learning works, our method is to obtain comprehensive embeddings for the student model, which have the functional information learned by the distillation loss. To show the effectiveness, we compare the proposed ProteinSSA to pretraining or self-supervised learning methods: DeepFRI [4], ESM-1b [10], ProtBERT-BFD [5], LM-GVP [6], amino-acid-level IEConv [7], GearNet-based methods [8], ESM-2 [3], KeAP [2], and ProtST [9] on these four tasks, including protein fold classification, enzyme reaction classification, Gene Ontology (GO) term prediction, and enzyme commission (EC) number prediction. Param. means the number of trainable parameters (B: billion; M: million; K: thousand).

As shown in these tables, without any pre-training or self-supervised learning, our proposed framework, ProteinSSA, achieves comparable accuracy with those methods with less trainable parameters and even outperforms them on fold and enzyme reaction classification.

Q3 The GNN architecture is not fully described.

A3 As mentioned in Section 3.3, a protein sequence is composed of $n$ residues, which are represented as graph nodes. We combine the one-hot encoding of residue types with their physicochemical properties, and they serve as the features for each graph node, denoted as $x_i$. In Section 3.1, we introduced the local coordinate system (LCS) [11] as the geometric features for protein 3D structures. The edge feature $e_{ij}$, defined in Eq. (2), represents the concatenation of the geometric features for protein 3D structures used LCS. We define the sequential distance as $l_{ij} = |i-j|$ and the spatial distance as $d_{ij}=|| P_i-P_j||$, and we establish the edge conditions using Eq. (3). Therefore, the protein's sequential and structural information is captured through the graph node and edge features. In Appendix B.2, we provide detailed information about the graph neural network model. We employ two message-passing layers followed by one average sequence pooling layer. After the pooling layer, the number of residues is halved, and we update the edge conditions before performing another round of message passing and pooling, as illustrated in Figure 1. Thus, the final GNN architecture includes eight message-massing and four pooling layers, which are sufficient for achieving satisfactory results. By leveraging message-passing and pooling layers, we integrate the graphs from different modalities (sequence and structure) into joint embeddings.

Q4 Can you add the performance of the teacher model into your results tables?

A4 In terms of the teacher model's performance, as we have stated, we use about 30K proteins with 2752 GO annotations from the GO dataset to train the teacher model, without further division into the biological process (BP), molecular function (MF), and cellular component (CC). The $F_{\mathrm{max}}$ is used as the evaluation metric. We got the $\mathrm{F}_{\mathrm{max}}$ for the teacher model is 0.489 overall.

Dear Reviewer 1uow,

We are grateful for your thorough review. Your comments are highly valued, and we would like to express our heartfelt gratitude. We do our utmost to address the questions you have raised:

Q1 The paper lacks some important baselines. For example, the paper didn't report the teacher model's performance and the performance of removing the KL divergence term.

A1 Thank you for your valuable feedback! (1) The proposed method is compared with existing popular protein representation learning methods. We have reported the performance of removing the KL divergence term in the ablation study, shown in Table 5, w/o Teacher means removing the teacher model, which leads to substantial performance drops across all tasks compared to the full ProteinSSA. (2) In terms of the teacher model's performance, as we have stated, we use approximately 30K proteins from the GO dataset to train the teacher model, without further division into the biological process (BP), molecular function (MF), and cellular component (CC). The $F_{\mathrm{max}}$ is used as the evaluation metric. Overall, we achieved a $F_{\mathrm{max}}$ score of 0.489 for the teacher model. (3) If we do not consider the student model, the inputs of the teacher model consist of sequence, structure, and function. Therefore, we evaluated the teacher model using the GO term dataset and calculated $F_{\mathrm{max}}$ scores for BP, MF, and CC. The results are as follows:

The goal of the Gene Ontology (GO) task is to identify protein functions. Molecular Function (MF) describes activities that occur at the molecular level, with 1,943 classes. Biological Process (BP) represents larger processes, categorized into 489 classes. Cellular Component (CC) describes the parts of a cell or its extracellular environment, with 320 classes. From the provided table, it is evident that incorporating function information as input significantly enhances performance, particularly for MF and CC. These two classes have fewer categories and are more accessible, resulting in higher scores.

Q2 Actually, I don't really get the reason why the author needed to train a student model, which seems redundant. In this paper, the student model is not smaller than the teacher model. Instead, the student model shares parameters with the teacher model. It seems the author just needs to finetune the ProteinBERT involving the additional GO information constraint.

A2 Thank you for your informative reviews! (1) At present, not even 1% of sequenced proteins have functional annotation [1,2]. While the teacher network requires extra functions as input, such information is not always available. As discussed in Section 3.2, we highlight the limitations of pre-training and the absence of a well-learned protein functional encoder to encode functional information. To address these challenges and make better use of functional information without extensive pre-training, we propose the multimodal knowledge distillation framework. (2) Compared with the teacher model, the student model is smaller than the teacher model, due to the deprecation of the annotation encoder. (3) The student model is trained by a hybrid loss, present in Eq.(10) in the revised manuscript, it does not share parameters with the teacher model. The Kullback-Leibler (KL) divergence loss is to regularize the parameters of the student model indirectly through the distribution of its embeddings. (4) ProteinBERT [3] is the first model to utilize the protein annotation information, using 8943 frequent GO annotations. The current state-of-the-art (SOTA) sequence-function pre-trained model is the KeAP [4], but the SOTA model KeAP fails to generate protein embeddings that surpass those obtained from the sequence pre-trained model ESM-1b [5]. This observation demonstrates the limitations of the sequence-function pre-trained model and motivates the development of ProteinSSA. (5) In Section 3.2, we have included a preliminary experiment to highlight the limitations of existing sequence-function pre-training models. This experiment serves as the motivation for proposing the multimodal knowledge distillation framework, ProteinSSA, which aims to embed protein sequence, structure, and function.

Q3 Many parts of the paper make me feel confused, especially the KL divergence part. For example, what do $P_S(G_S,A)$ and $P(Z_S|G_S,A)$ mean? Are these VAE model? If it's true, then expanding the $P_S(G)$ to $P(G|z)P(z)$, I don't think the assumption that "$E_{p_S(G)}\left[K L\left[p_S(y \mid G) P_T() y \mid G\right]\right]$ does not depend on z" holds. By the way, I don't really get what the source domain and target domain mean. It seems they are the same domain in the exception that source has an additional constraint on GO.

A3 (1) In the context of domain adaptation, a domain refers to a specific distribution of data. The main goal is to overcome the distribution shift between the source and target domains, which can lead to a decrease in performance when directly applying a model trained on the source domain to the target domain. (2) In section 3.3, we have clarified that the reason for employing domain adaptation strategies is to develop a sample-independent method, particularly for protein samples without available functional labels. The goal is to align the student's latent distributions obtained from sequence-structure embeddings to the distributions of the teacher model's latent space, thereby bridging the gap between different domains and enhancing the performance of the student model. The source domain is the distribution of teacher model's embeddings, and the target domain is the distribution of student model's embeddings. (3) $P(Z_S|G_S,A)$ mean the distribution of $z_S$ in the source domain. In Section 3.1, we have given the definitions. There is a source domain $S$ with the data distribution $p_S(z_S| G_S, A)$, and there is also a target domain $T$ with the data distribution $p_T(z_T|G_T)$, the latent embeddings $z_S, z_T$ are from the teacher and student network for a protein. $A$ with $k$ terms are the GO annotations. $G_S, G_T$ are the input graphs in the source and target domains. Thus, $P_S(G_S,A)$ means the distribution of the input graph $G_S$ and annotation $A$. (4) This is not a VAE model. (5) In Eq.(8) in the revised manuscript, it defines the misalignment in the representation space between the source and target domains. The theoretical derivations are shown in Appendix D. $E_{p_S(G)}\left[KL\left[p_S(y|G)\parallel p_T(y|G)\right]\right]$ is fixed and and often small [7] (not dependent on the representation $z$, and $y$ is the protein class), there is no $z$ in this term, which is not dependent on the representation $z$. From Proposition 2 shown in Appendix D, we can see the conditional misalignment in the representation space is bounded by the conditional misalignment in the input space.

Q4 In Equation 5, why the author directly add the $h_S$ to $h_A$ without any transformation?

A4 The annotations associated with $G_S$ serve as the input for the annotation encoder, resulting in the extraction of feature vector $h_A$. Notably, both $h_S$ (the feature vector obtained from processing $G_S$) and $h_A$ possess identical feature dimensions, enabling a direct addition of these vectors.

Q5 Removing the AE-T doesn't influence the performance much.

A5 (1) In our study, we conducted ablation experiments to analyze the impact of excluding the annotation encoder in the teacher model. This is denoted as "w/o AE-T" in Table 5. The results showed a slight decrease in performance because function information was incorporated into the loss function without using the annotation encoder. However, completely removing the teacher model (w/o Teacher) resulted in significant performance drops across all tasks compared to the full ProteinSSA approach. These ablations highlight the importance of utilizing the annotation encoder and incorporating function information in the loss function and the teacher model for optimal results. (2) By removing the annotation encoder in the teacher model but still incorporating function information into the loss function, it can be viewed as a multi-task learning method. This approach leverages knowledge transfer across tasks, leading to improved generalization performance. (3) The auxiliary annotation encoder in the teacher model is implemented as a multi-layer perceptron. Further details about the model can be found in Appendix B.2. Additionally, the GO encoder can be enhanced by utilizing large language models or graph neural networks, as well as incorporating more GO terms to capture a broader range of functional information.

Thank you again!

[1] Torres, M., Yang, H., Romero, A. E. & Paccanaro, A. Protein function prediction for newly sequenced organisms. Nat. Mach. Intell. 3, 1050–1060 (2021).

[2] Ibtehaz, et al. "Domain-PFP allows protein function prediction using function-aware domain embedding representations." Communications Biology 6.1 (2023): 1103.

[3] Brandes, Nadav, et al. "ProteinBERT: a universal deep-learning model of protein sequence and function." Bioinformatics 38.8 (2022): 2102-2110.

[4] Hong-Yu Zhou, et al. Protein representation learning via knowledge enhanced primary structure modeling, 2023.

[5] Zeming Lin, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. BioRxiv, 2022.

[6] Hehe Fan, et al. Continuous-discrete convolution for geometry-sequence modeling in proteins. In The Eleventh International Conference on Learning Representations, 2023.

[7] Ben-David, Shai, et al. "A theory of learning from different domains." Machine learning 79 (2010): 151-175.

Dear Reviewer 1uow,

We genuinely appreciate the time and effort you've dedicated to reviewing our work and your acknowledgment of the answers we've conducted.

(1) I have carefully read the author's response and other reviewers' comments. I do agree with reviewer KkPB's opinion that there are too many elements incorporated into this method without a cohesive story, making the paper confusing. For example, I still feel unclear about the motivation of training a student model instead of directly finetuning ProteinBERT with additional functional constraints. Therefore, I'd like to keep my current score.

A1 Thanks for your criticisms. We have made significant efforts to improve the manuscript and have just completed the latest round of revisions. The submitted version has undergone multiple iterations, with each revision incorporating valuable feedback and making necessary improvements.

At present, not even 1% of sequenced proteins have functional annotation [1,2]. While the teacher network requires extra functions as input, such information is not always available. As discussed in Section 3.2, we highlight the limitations of pre-training and the absence of a well-learned protein functional encoder to encode functional information. To address these challenges and make better use of functional information without extensive pre-training, we propose the multimodal knowledge distillation framework.

In case our answers have justifiably addressed your concerns, we respectfully thank you that support the acceptance of our work. As you know, your support holds great significance for us. Also, please let us know if you have any further questions. Look forward to further discussions!

[1] Torres, M., Yang, H., Romero, A. E. & Paccanaro, A. Protein function prediction for newly sequenced organisms. Nat. Mach. Intell. 3, 1050–1060 (2021).

[2] Ibtehaz, Nabil, Yuki Kagaya, and Daisuke Kihara. "Domain-PFP allows protein function prediction using function-aware domain embedding representations." Communications Biology 6.1 (2023): 1103.

Q3 The different elements are well described, but what isn’t clear exactly is how they will be synthesized into something new and exciting, as well as why the choices (e.g edge representation of 3.1, sequence function representation of 3.2) were made, and why (beyond them being SOTA at one point in time).

A3 Thank you for your suggestions! We have carefully reviewed and incorporated your feedback into the manuscript, especially for section 3, we add a summary of the subsections at the beginning, and describe what they functions. For example, (1) we define the problem and notations, and introduce the local coordinate system (LCS) [7] as the geometric features for protein 3D structures. LCS features are widely adopted in protein structure modeling methods [4,8], as they provide informative and comprehensive descriptions of protein structures by considering distance, direction, and orientation. (2) In Section 3.2, we have included a preliminary experiment to highlight the limitations of existing sequence-function pre-training models. This experiment serves as the motivation for proposing the multimodal knowledge distillation framework, ProteinSSA, which aims to embed protein sequence, structure, and function. (3) In Section 3.3, we have presented the overall framework of ProteinSSA, which involves training a teacher model and a student model through iterative knowledge distillation. We have provided a detailed explanation of the network architecture, highlighting the use of the same Graph Neural Network (GNN) architecture in both the teacher and student models. Additionally, we have introduced the concept of protein domain adaptation and derived the loss for the student model. The message-passing mechanism used in both the teacher and student models is inspired by CDConv, which enables better modeling of protein sequences and structures. We have clarified that the reason for developing a sample-independent method, for proteins without available functional labels. The goal is to align the student's latent distributions obtained from sequence-structure embeddings to the distributions of the teacher model's latent space, thereby bridging the gap between different domains and enhancing the performance of the student model.

Q4 It isn't yet explained why the authors think knowledge distillation is the best way to incorporate knowledge from annotations. Why not just use the teacher network directly?

A4 Thank you for your reviews! At present, not even 1% of sequenced proteins have functional annotation [9,10]. While the teacher network requires extra functions as input, such information is not always available. Another way is the pre-training and fine-tuning paradigm, as discussed in Section 3.2, we highlight the limitations of pre-training and the absence of a well-learned protein functional encoder to encode functional information. To address these challenges and make better use of functional information without extensive pre-training, we propose the multimodal knowledge distillation framework.

Q5 Both the sentence preceding equation 8 and the sentences that follow are overly wordy, but without the benefit of clarity. Words about “reduce the bound in the representation spaces” or “KL divergence matches distributions…” is a bit misleading;

A5 Thank you for your informative reviews! We appreciate your feedback and have taken steps to clarify wordy sentences. Regarding Eq.(6)-Eq.(10), we would like to clarify that these equations were used to illustrate the process of deriving an ideal loss for the student model in a theoretical manner, where the goal is to align the student embedding distribution with the teacher embedding distribution. We apologize for any confusion caused and have made the necessary rectifications.

Q6 Table 3 reports only point estimates of max accuracy.

A6 Thank you for highlighting this aspect! In Section 4.1, we present the performance is measured with averaged values over three initializations, this means that we report the mean terms of three independent runs instead of point estimation on these four tasks, including protein fold classification, enzyme reaction classification, Gene Ontology (GO) term prediction, and enzyme commission (EC) number prediction. For the variance or confidence intervals of reported results, we follow the results of baselines reported in [4]. For the task of fold and reaction classification, the performance is measured as mean accuracy. For GO Term and EC number prediction, the $\mathrm{F}_{\mathrm{max}}$ is used as the evaluation metric. Here, we report the mean (variance) for the proposed ProteinSSA:

Q7 The phrase ‘grammar of life’ isn’t a helpful metaphor.

A7 Thank you for your comment! We highly appreciate your feedback and acknowledge that the knowledge acquired by these models cannot always be distilled into explicit rules for the compositional orientation of elements in the protein language. In the introduction, we discuss protein language models (PLMs) that have demonstrated the ability to learn the 'grammar of life' from large numbers of protein sequences. In works such as [12,13], language models are applied to model protein sequences, highlighting the common characteristics shared between human languages and proteins. For instance, the hierarchical organization [14] suggests that the four levels of protein structures can be analogized to letters, words, sentences, and texts in human languages to a certain degree. In light of these considerations, it may be appropriate to modify the phrase "learn the grammar of life" to "learn the certain grammar of life".

Q8 Equation 5 has a term $\alpha$ that controls "the isotropic of protein representations". What does this mean?

A8 Thank you for your question! In this context, the term "isotropic" refers to the uniformity or symmetry of protein representations. It may be better to change the word isotropic to balance. The hyper-parameter $\alpha$ controls the balance between the contribution of the annotation encoder ($h_A$) and the graph-level protein embeddings ($h_S$) in the combined representation ($z_S$). By adjusting the value of $\alpha$, we can control the extent to which the two components contribute to the final representation. A higher value of $\alpha$ would give more weight to the graph-level protein embeddings, while a lower value would amplify the influence of the annotation encoder. This allows us to control the isotropy or symmetry of the resulting protein representations, determining whether they are more influenced by the annotation or the structural information captured in the graph-level embeddings.

Q9 A few minor typos in the first sentence on page 6.

A9 Thanks for your careful reviews! We sincerely appreciate your detailed and valuable suggestions, and we thoroughly revise our paper based on your constructive comments.

Q10 I don’t think you need to invoke the CLT here to model the distribution of the embeddings as Gaussian, you can just assume it to be true.

A10 Thank you for your suggestion! It is true that we can assume protein embeddings follow a Gaussian distribution without mentioning the central limit theorem [15]. We have modified this in the revised version.

Q11 why the teacher embeddings are concatenated with a separate functional embedding rather than using function as an extra term in the loss function for classification. How come?

A11 Thank you for insightful reviews! (1) We conducted ablation experiments, shown in Table 5, where we excluded the annotation encoder in the teacher model, denoted as "w/o AE-T" in Table 5. This means that we incorporated function information into the loss function for the teacher models, while deprecating the use of the annotation encoder. As a result, there was a slight decrease in performance observed. These ablations demonstrate the importance of using the annotation encoder and incorporating function information in the loss function for optimal results. (2) The input functions used in the annotation encoder can be considered as part of the Label-Augmented method. Previous works have successfully utilized label-augmented techniques to enhance model training [16,17]. This technique involves encoding labels and combining them with node attributes through concatenation or summation. By doing so, it improves feature representation and enables the model to effectively utilize valuable information from labels. The Label-Augmented method is applicable to various types of graphs and can be employed in different domains. Thus, we choose to incorporate functional embeddings, concatenated with graph-level protein sequence-structure embeddings. (3) Another approach to improve the model's performance is by incorporating function as an additional term in the loss function, which is commonly seen in multi-task learning (MTL) methods. The objective of MTL is to leverage knowledge transfer across tasks and achieve better generalization. However, when tasks are diverse, a simple shared MTL model can suffer from task interference [18]. Previous work [19] has explored the relationship and mutual adaptability of protein tasks. Considering the concept of prompt learning [18] in the MTL domain, it holds promise as a potential direction for improvement.

Q12 Why generalization bound arguments are being used here;

A12 Thanks for your valuable comments. There exists a misalignment for KL-guided domain adaptation method in the representation space between source domain and target domain because of data differences. Theoretically, the student model should have a lower generalization bound when transferring knowledge from the teacher model, as shown in Eq.(8) in the revised manuscript. In order to approximate the ideal target domain loss, we reduce the terms inside the radical sign, because the second term in Eq.(8) is often small and fixed. Thus, we only need to leverage Kullback-Leibler (KL) divergence to align the distributions between source domain and target domain, i.e., $KL[p_S(z)|| p_T(z)]$.

Q13 Section 4.1 begins by describing ProteinBERT and how it is pre-trained. Is this part of ProteinSSA?

Q13 Thanks for your question! ProteinBERT [1] is not a part of ProteinSSA. It is the first model to utilize the protein annotation information, using 8943 frequent GO annotations. Compared with ProteinBERT, we use 2752 GO annotations from the GO dataset. We have modified this in the revised manuscript to decrease the misleading.

Thank you again for all the efforts that helped us improve our manuscript! In case our answers have justifiably addressed your concerns, we respectfully thank you that support the acceptance of our work. As you know, your support holds great significance for us. Also, please let us know if you have any further questions. Look forward to further discussions!

[1] Brandes, Nadav, et al. "ProteinBERT: a universal deep-learning model of protein sequence and function." Bioinformatics 38.8 (2022): 2102-2110.

[2] Zhou, Hong-Yu, et al. "Protein Representation Learning via Knowledge Enhanced Primary Structure Modeling." bioRxiv (2023): 2023-01.

[3] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Sal Candido, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. BioRxiv, 2022.

[4] Hehe Fan, Zhangyang Wang, Yi Yang, and Mohan Kankanhalli. Continuous-discrete convolution for geometry-sequence modeling in proteins. In The Eleventh International Conference on Learning Representations, 2023.

[5] Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. In International Conference on Learning Representations, 2023.

[6] Varadi, Mihaly, et al. "AlphaFold Protein Structure Database: massively expanding the structural coverage of protein-sequence space with high-accuracy models." Nucleic acids research 50.D1 (2022): D439-D444.

[7] John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph-based protein design. Advances in neural information processing systems, 32, 2019.

[8] Anishchenko, Ivan, et al. "De novo protein design by deep network hallucination." Nature 600.7889 (2021): 547-552.

[9] Torres, M., Yang, H., Romero, A. E. & Paccanaro, A. Protein function prediction for newly sequenced organisms. Nat. Mach. Intell. 3, 1050–1060 (2021).

[10] Ibtehaz, Nabil, Yuki Kagaya, and Daisuke Kihara. "Domain-PFP allows protein function prediction using function-aware domain embedding representations." Communications Biology 6.1 (2023): 1103.

[11] Zeming Lin, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. BioRxiv, 2022.

[12] Asgari, E. and Mofrad, M. R. K. (2015). Continuous distributed representation of biological sequences for deep proteomics and genomics. PLOS ONE.

[13] Yang, K. K., Wu, Z., Bedbrook, C. N., and Arnold, F. H. (2018). Learned protein embeddings for machine learning. Bioinformatics.

[14] Ferruz, N. and Höcker, B. (2022). Controllable protein design with language models. Nature Machine Intelligence, pages 1–12.

[15] Oliver Johnson. Information theory and the central limit theorem. World Scientific, 2004.

[16] Samy Bengio, et al. 2010. Label embedding trees for large multi-class tasks. Advances in Neural Information Processing Systems 23 (2010).

[17] Xu Sun, Bingzhen Wei, Xuancheng Ren, and Shuming Ma. 2017. Label embedding network: Learning label representation for soft training of deep networks. arXiv preprint arXiv:1710.10393 (2017)

[18] Wang, Zeyuan, et al. "Multi-level Protein Structure Pre-training via Prompt Learning." The Eleventh International Conference on Learning Representations. 2022.

[19] Hu, Fan, et al. "A Multimodal Protein Representation Framework for Quantifying Transferability Across Biochemical Downstream Tasks." Advanced Science (2023): 2301223.

Dear Reviewer KkPB,

Thanks for your appreciation and detailed review. We respond to your comments with a heart full of gratitude. We try our best to response the questions below:

Q1 ProteinSSA does not require pretraining, and does not make use of annotations. This is only true of the student model. Table 1 is thus misleading.

A1 Thank you for your valuable feedback! (1) It is indeed true that ProteinSSA uses annotations for the teacher model, its objective is to learn embeddings in the latent space that contain functional information and provide intermediate supervision during knowledge distillation for the student model. Therefore, the complete training of the teacher model is not our primary concern. Our main focus is to obtain comprehensive embeddings for the student model, which is trained using distillation loss and task loss without the annotations input. (2) Pre-training involves learning general representations from a large dataset, while fine-tuning adapts the pre-trained model to a specific task or dataset. In the teacher-student framework, the teacher model is usually a well-learned model that serves as a source of knowledge for the student model. The student model aims to mimic the behavior or predictions of the teacher model. In this work, the teacher model is trained on a relatively small dataset consisting of about 30 thousand proteins with 2752 GO annotations. This is in comparison to ProteinBERT [1], which uses 106 million proteins, and KeAP [2], which uses 5 million triplets. It can still be considered as training rather than pre-training for the teacher model, and we mainly want to obtain a better student model. We have clarified this information and modified Table 1 in the revised version for better understanding.

Q2 What problem does section 3.2 address? Or what part of the final architecture is being discussed here? It’s not clear, and would benefit from being rewritten. How does CDConv relate to KeAP and ESM-1b? These are very briefly described, but not in sufficient detail to tell the reader why each was chosen, and how they relate to each other & to ProteinSSA as a whole. It’s only when you read to the bottom of 3.2 ... , lead with this, and then describe the limitations of other sequence-based models that require extensive pre-training afterwards.

A2 Thank you for your informative reviews! We have carefully considered your suggestions and made the following revisions to the manuscript. (1) In Section 3.2 of the revised manuscript, we now begin with a summary and central idea of the subsection, followed by the preliminary experiment. We emphasize that the performance of pre-trained models is influenced by various factors, such as model size, dataset scale, and choice of pre-training tasks. To illustrate this, we have conducted an experiment. (2) In this experiment, we want to highlight that the current state-of-the-art (SOTA) sequence-function pre-trained model, KeAP [2], fails to generate protein embeddings that surpass those obtained from the sequence pre-trained model ESM-1b [3]. This observation demonstrates the limitations of the sequence-function pre-trained model and motivates the development of ProteinSSA. (3) As mentioned in Section 3.2, CDConv [4] introduces an effective fundamental operation to encapsulate protein structure and sequence without pre-training or self-supervised learning. CDConv achieves superior performance compared to pre-training methods like GearNet-Edge-IEConv [5]. CDConv is currently the most effective publicly available method for modeling protein sequence and structure. In the field of protein pre-training, KeAP is the current SOTA model for downstream tasks, pre-trained on protein sequences and functions, while ESM-1b is the most prevalent sequence pre-training model. By incorporating the embeddings from KeAP and ESM-1b to enhance the embeddings obtained from CDConv, we can compare the quality and performance of the embeddings from the two pre-trained models, evaluating their ability using the solid sequence-structure base model, CDConv. (4) The student model in the ProteinSSA framework is designed to model protein sequence and structure together, learning functional information from the teacher model without pre-training. Additionally, the message-passing mechanism formulated in Eq.(4) is inspired by CDConv, which convolves node and edge features from sequence and structure simultaneously.

(4) Regarding the authors answer to Q8 ('isotropic'), this description that the authors provide I think is the clearest description of what the parameter controls, and maybe should be what appears in the text. Regarding the choice to incorporate functional annotations through label encoding rather than as a separate loss function, I appreciate the justification via Bengio et al and Sun et al. I think this justification should also appear in the text (and please add both citations to the manuscript).

A4 Thanks for your suggestions! We have added these contents to the revised manuscript.

(5) I find the writing here is still quite confusing. Generalization bound arguments are typically made with respect to a given model on different datasets. This application here is to try and bound the loss of the student by the loss of the teacher network plus some distributional alignment terms of the student and teacher model latent variables. But the notation is overloaded here in confusing ways: (a) What does $Y$ specifically refer to here? Is it shared by both models? I did not think the student model had any notion of labeled data necessarily. (b) Mutual information is defined on random variables, not on their realizations; so equation (15) and it's description is very confusing for me

A5 Thanks for your valuable reviews! The teacher and student models can be used in different datasets. To reduce the generalization bound, we focus on optimizing marginal misalignment and obtain the derived loss for the student model, which is Eq. 9. (a) $y$ represents the underlying functional label for the student model. Although the student model may not have these functional labels, but we can assume that they exist for theoretical reasons. In Appendix D, we have mentioned we assume the source and target domains have the same support set (training set) to derive this theory, i.e, the difference is that the teacher has functional labels, while the student does not. The derived misalignment Eq. 8 and the derived loss Eq. 9 are based on the assumption that the source and target domains have the same support set. Thus, the loss of Eq. 9 can be used in an unsupervised way for the student to predict functions. However, the student model is applied to different downstream tasks, like classification, which have classification classes. Thus, we add the supervised student loss $L_\mathrm{student}$ and the knowledge distillation loss the $L_{kd}$ as the final hybrid loss for the student to improve the performance on the classification task. (b) Mutual information is typically used to compute the relationship between random variables. However, if we have a way to discretize or represent the non-variable or actual values as random variables, we may be able to apply mutual information to measure their relationship. This would involve defining appropriate probability distributions for the variables and estimating the probabilities based on the observed data. In Eq. 15, the variable z represents the latent embedding, which is a learned representation of the input graph G. The variable y represents the data label or target variable, which provides information about the class or category to which the input data belongs. The variable G represents the input data itself. G is typically considered as the independent variable in the context of the equation. The mutual information term, denoted as I_S(·, ·), is calculated on the source domain. It quantifies the amount of information shared between the variables z and y (or G and y) in the source domain. It measures the dependence or correlation between these variables in the context of the source domain data.

Thank you again for all the efforts that helped us improve our manuscript! In case our answers have justifiably addressed your concerns, we respectfully thank you that support the acceptance of our work. As you know, your support holds great significance for us. Also, please let us know if you have any further questions. Look forward to further discussions!

[1] Ibtehaz, Nabil, Yuki Kagaya, and Daisuke Kihara. "Domain-PFP allows protein function prediction using function-aware domain embedding representations." Communications Biology 6.1 (2023): 1103.