Paper: Multimodal Distillation of Protein Sequence, Structure, and Function (ICLR-2024)

Reviewer comment (Reproducibility):
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. This appears inconsistent and further clarification would be helpful (significance, clarity).

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

Reviewer comment (Experiments):
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 how much model size affects performance on the studied tasks (quality, significance).

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

| Method | Param.| Pretraining Dataset |Fold | SuperFamily | Family | Enzyme Reaction |
| :--- | :---: | :---: | :---: | :---: | :---: |:---: |
|DeepFRI | 6.2M| Pfam (10M)| 15.3 | 20.6 | 73.2 | 63.3 |
|ESM-1b | 650M| UniRef50 (24M)| 26.8 | 60.1 | 97.8 | 83.1 |
|ProtBERT-BFD| 420M| BFD (2.1B)| 26.6 | 55.8 | 97.6 | 72.2 |
| IEConv (amino level)| 36.6M | PDB (476K)| 50.3 | **80.6** | 99.7 |88.1|
|GearNet-Edge-IEConv (Multiview Contras)| 42M |AlphaFoldDB (805K)|54.1 | 80.5 | **99.9** | 87.5 |
|GearNet-Edge-IEConv (Residue Type Prediction)| 42M|AlphaFoldDB (805K)|48.8 | 71.0 | 99.4 | 86.6|
|GearNet-Edge-IEConv (Distance Prediction)|42M |AlphaFoldDB (805K)|50.9 | 73.5 | 99.4 | 87.5|
|GearNet-Edge-IEConv (Angle Prediction)|42M |AlphaFoldDB (805K)|56.5 | 76.3 | 99.6 | 86.8|
|GearNet-Edge-IEConv (Dihedral Prediction)|42M |AlphaFoldDB (805K)|51.8 | 77.8 | 99.6 |87.0 |
| ProteinSSA | 100M |-|**60.5** | 79.4 | 99.8 | **89.4** |

| Method |Param. |Pretraining Dataset|GO-BP | GO-MF | GO-CC | EC |
| :--- | :---: | :---: | :---: | :---: |:---: |:---: |
|DeepFRI |6.2M| Pfam (10M)| 0.399 | 0.465 | 0.460 | 0.631 |
|ESM-1b | 650M|UniRef50 (24M)| 0.470 | 0.657 | 0.488 | 0.864 |
|ProtBERT-BFD| 420M| BFD (2.1B)| 0.279 | 0.456 | 0.408 | 0.838 |
|LM-GVP| 420M|UniRef100 (216M)| 0.417 | 0.545 | **0.527** | 0.664|
| IEConv (amino level)| 36.6M |PDB (476K)| 0.468 | 0.661 | 0.516 |-|
|GearNet-Edge-IEConv (Multiview Contras)| 42M |AlphaFoldDB (805K)| **0.490** | 0.654 | 0.488 | 0.874 |
|GearNet-Edge-IEConv (Residue Type Prediction)| 42M |AlphaFoldDB (805K)|0.430 | 0.604 | 0.465 | 0.843|
|GearNet-Edge-IEConv (Distance Prediction)|42M | AlphaFoldDB (805K)|0.448 | 0.616 | 0.464 | 0.839 |
|GearNet-Edge-IEConv (Angle Prediction)| 42M |AlphaFoldDB (805K)|0.458 | 0.625 | 0.473 | 0.853 |
|GearNet-Edge-IEConv (Dihedral Prediction)|42M | AlphaFoldDB (805K)|0.458 | 0.626 | 0.465 | 0.859 |
|KeAP| 420M |ProteinKG25 (5M)| 0.466 | 0.659 | 0.470 | 0.845 |
|ESM-2| 650M |UniRef50 (24M)| 0.472 | 0.662 | 0.472 | 0.874 |
|ProtST-ESM-1b| 759M | ProtDescribe (553K)| 0.480 | 0.661 | 0.488| **0.878** |
|ProtST-ESM-2| 759M | ProtDescribe (553K)| 0.482 | **0.668** | 0.487| **0.878** |
| ProteinSSA |100M |-|0.464 | 0.667 | 0.492 | 0.857 |

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.

Reviewer comment (Writing):
The GNN architecture is not fully described (clarity).

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

Reviewer comment (Experiments):
**Lack of baselines:** 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.

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

| Method| GO-BP | GO-MF | GO-CC | 
| :--- | :---: | :---: | :---: | 
| ProteinSSA (student) | 0.464 | 0.667| 0.492 | 
| ProteinSSA (teacher) | 0.769 | 0.832| 0.849 | 

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.

Reviewer comment (Novelty):
**The motivation is unclear:** 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.

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

Reviewer comment (Writing):
**The writing is confusing:** 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)}[KL[p_S(y|G)P_T()y|G]]$ 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.

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

Reviewer comment (Theory):
In Equation 5, why the author directly add the $h_S$ to h_A without any transformation? It seems they are from different semantic spaces.

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

Reviewer comment (Experiments):
Removing the AE-T doesn't influence the performance much. Does that mean this additional GO encoder didn't add to much benefit to the whole model?

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

Reviewer comment (Presentation):
The authors stress that ProteinSSA does not require pretraining, and does not make use of annotations. This is only true of the *student* model, since the teacher does clearly make use of annotations where such exist. Table 1 is thus misleading.

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

Reviewer comment (Presentation):
The title of section 3.2 isn’t very informative. What problem does it address? Or what part of the final architecture is being discussed here? It’s not clear, and would benefit from being rewritten. For example, 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 that you discover that this section is all about establishing that pre-trained models are limited in different ways, and *that’s* why ProteinSSA was made. Please, lead with this, and then describe the limitations of other sequence-based models that require extensive pre-training afterwards.

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

Reviewer comment (Presentation):
Reading through the subsections of section 3, it’s hard to put my finger on what the focus of this paper is. 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). I think section 3 would be clearer, and benefit the reader if it had a summary of the subsections at the beginning, and for each subsection to describe one part of the whole model.

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

Reviewer comment (Writing):
Reading through to Section 3.3, 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? The answer is (I think from subsequent sections) that functional annotation are only sometimes available, so by instead aligning the latent space of the student with that of the teacher, the student derives the benefit of additional knowledge. This is not clear from reading section 3, but should be clearly spelled out somewhere within (or within the introduction).

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

Reviewer comment (Writing):
Both the sentence preceding equation 8 and the sentence that follow are overly wordy, but without the benefit of clarity. It’s clear that the addition of a KL regularization term KL($P_{S}(z_{S})$, $P_{T}(z_{T})$) will force the student embedding distribution to become like the teacher distribution, thus affecting the student embedding state. Words about “reduce the bound in the representation spaces” or “KL divergence matches distributions…” is a bit misleading; all that’s intended here is an intention to regularize the parameters of the student model indirectly through the distribution of its embeddings.

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

Reviewer comment (Evaluation):
Table 3 reports only point estimates of max accuracy. Max accuracies are very difficult to parse in a meaningful way. The improvements of ProteinSSA would be better qualified if the distribution of accuracies from multiple runs, especially that of fold classification, was reported. Even if the alternatives cannot be re-run, ProteinSSA results could be reported more faithfully.

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

| Method | Fold | SuperFamily | Family | Enzyme Reaction | GO-BP | GO-MF | GO-CC | EC |
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
| ProteinSSA | 60.5(0.60) | 79.4(0.96) | 99.8(0.04) | 89.4(0.43) | 0.464(0.007) | 0.667(0.003) | 0.492(0.004) | 0.857(0.011) |

Reviewer comment (Writing):
Equation 5 has a term $\alpha$ that controls "the isotropic of protein representations". The meaning of this is unclear.

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

Reviewer comment (Theory):
Page 6 in section 3.3 invokes the CLT. Invoking the CLT is not necessary here to model the distribution of the embeddings as Gaussian, as it can simply be assumed to be true. Additionally, it’s not clear that the different batch-derived embeddings are independent.

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

Reviewer comment (Experiments):
The ablation study of section 4.5 is welcome, but does not address one of the key choices of the paper (raised in the *Protein Domain Adaptation* paragraph of section 3.3), which is 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.

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

| Method | Fold | SuperFamily | Family | Enzyme Reaction | GO-BP | GO-MF | GO-CC | EC |
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
| ProteinSSA | 60.5 | 79.4 | 99.8 | 89.4 | 0.464 | 0.667 | 0.492 | 0.857 |
| w/o AE-T | 60.4 | 79.1 | 99.7 | 88.9 | 0.454 | 0.664 | 0.490 | 0.854 |

Reviewer comment (Presentation):
Section 4.1 begins by describing ProteinBERT and how it is pre-trained. It is unclear whether this is part of ProteinSSA. If so, ProteinSSA cannot really claim (as in table 1) that it is not pre-trained. If not, it is unclear whether mentioning ProteinBERT here is relevant.

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