ProtEx: A Retrieval-Augmented Approach for Protein Function Prediction

Thank you for reviewing our work!

"It is to be discussed whether the paper is novel enough for a machine learning conference. pro: the task of protein function prediction with retrieval-augmentation is not well-explored today, and the paper provides meaningful discussion. con: the methodology itself is not novel."

We believe the simplicity of our approach is an advantage, as we show it outperforms more complex approaches from prior work (e.g. Table 4). Our method is carefully designed and we provide ablations showing the value of each component of our method (e.g. the novel pretraining strategy).

Novelty of the methods aside, we believe our findings are likely to be of significant interest to at least two groups within the ICLR community:

1. Researchers interested in retrieval-augmented modeling methods and applications. Protein function prediction has many properties (rare classes, sequences far from the training data, large label space) that we show are very amenable for retrieval-augmented methods and also offer new challenges for researchers in this area. Identifying an area where such methods can have such a significant impact will likely be interesting to this community. Thus, we believe our results can help introduce methodological researchers to an emerging application area.
2. Researchers and practitioners interested in improving the ability to characterize proteins. We believe this is a very impactful application of ML, and we note that some of the prior work we compare with has been published in ML venues, e.g. ProtST (ICML 2023), PromptProtein (ICLR 2023), GearNet (ICLR 2023). We believe that our key finding that retrieval-augmented methods can enable significant improvements over the current state-of-the-art is novel and well supported by our experiments and analysis, and can inspire future work.

"Could you provide an analysis of how ProtEx performs when fewer exemplar sequences are available, or even when limited to a small number of retrieved sequences? For example the orphan protein. This could help in understanding ProtEx’s robustness and effectiveness when retrieval options are sparse."

We agree this is a crucial question. While this can be a challenging case for any ML-based method, we find that the relative performance of ProtEx is actually strongest in cases where fewer reference examples are available for a given class.

In the original submission, we provided analysis of how well ProtEx performs when few exemplars are available for a given class. Per Figure 5, the relative performance of ProtEx is actually strongest for the bucket representing test sequences belonging to classes with only 1-17 reference sequences.

To further analyze this, we have also performed a similar analysis for EC-based prediction below, focusing on the clustered EC split. We analyze performance when we restrict the computation of the max micro-averaged F1 score to classes in the dev set with less than a given number of training sequences. Notably, conditioning on exemplars is particularly effective for these rare classes relative to our baseline without exemplars and ProteInfer (from Sanderson et al., 2023), highlighting a strength of a retrieval-augmented approach such as ProtEx.

Thank you for reviewing our work!

"For broader readers, the authors may consider use a few illustrative figures or tables to justify the significance of retrieval-augmentation for protein function prediction"

Thank you, this is a great idea, and we indeed hope to make the paper interesting and accessible to readers who may have some familiarity with retrieval-augmented models (e.g. from NLP) but not protein function prediction, and vice versa. We will look to extend Figure 1 to more clearly contrast our approach with standard methods for protein function prediction that do not condition predictions on retrieved exemplars in our camera ready version.

"The performance of ProtEx heavily depends on the quality and comprehensiveness of the reference database. In cases where exemplars are sparse or unavailable for certain protein functions, the retrieval mechanism may underperform, potentially limiting the model’s utility in practical applications with limited exemplar coverage. Only BLAST is tested for similarity search, which may weaken the soundness of this study."

Please see our new results below!

"Although ProtEx aims to improve generalization, its effectiveness for predicting completely novel functions is not well-supported in the experiments. The evaluation focuses primarily on tasks with annotated classes, leaving uncertainty about the model’s robustness in generalizing to truly unseen or hypothetical functions, a critical aspect of protein function prediction. As a matter of fact, a lot of EC classes, for example, are labeled by expert largely based on the sequence similarity and that explains why the BLAST baseline alone yields already good results."

We agree with your point in general, as predicting completely novel functions is an interesting problem but beyond the scope of this paper. However, we argue that more accurately predicting functional labels within existing ontologies such EC numbers, GO terms, and Pfam domains can still be useful, and indeed is a widely studied problem. Furthermore, we have shown that ProtEx can generalize to new functions not seen at training time (Table 6). This is interesting because new labels are frequently added to these ontologies, and ProtEx can adapt to these new labels without any re-training (although it does require at least 1 positive exemplar sequence for the new label).

We also evaluated ProtEx on settings where the sequence similarity between test and train sequences is limited. Table 2 shows performance on clustered splits that restrict similarity between train and test sequences, Table 3 shows performance on specifically curated challenge sets, Table 4 partitions performance by sequence similarity, and Table 5 also shows performance on a clustered split that restricts similarity between train and test sequences. In all cases ProtEx demonstrates significant improvements over methods from prior work.

(Response continued in next comment.)

Thanks for the detailed answer for addressing my question. I hope the discussion did help improve the quality of manuscript :) I have adjusted my score to 6.

Thank you for reviewing our work!

"In Sec. 3, the proposed method is conditioned on label and labels are enumerated. What's the advantage for the conditioning?"

We are not sure whether you are asking about the advantage of conditioning the selection of positive and negative exemplars based on a label, or the advantage of including the label as an explicit input to the model. So we will address both here.

Overall, the main advantage of making a prediction specific to a given label is so that the positive and negative exemplars can be selected based on a particular label. Some of the rationale for this is discussed in the beginning of section 3.1. It is important to note that because our model only has a limited context window, we can only condition each prediction on a limited number of exemplars. Selecting these exemplars conditioned on a label lets us focus each prediction on the exemplars that are most relevant towards understanding the class boundary of a specific label, as illustrated in Figure 1. For the same query sequence, but another label, we can then select a different set of exemplars for understanding the class boundary. Otherwise, if we selected a single set of exemplars per sequence, most classes would likely not be well represented in this set.

Whether we include the label itself as an input to the model (as shown in Figure 1) or not does not have a significant impact on performance empirically (see the ablation shown in Table 13 in the appendix). This is an interesting finding, because without the label explicitly represented, the model must rely entirely on the information provided by the positive and negative exemplars to determine the relevance of the given class. This ability is further validated by the results in Table 6.

However, including the label as input allows for a clearer ablation of the impact of conditioning on the exemplars. Understanding the impact of conditioning on exemplars is central to the paper and this ablation is included in all of our main result tables. Without the label or exemplars, the prediction task would be underspecified.

"In Tab. 6, if the candidate labels are not incorporated, how the prediction is done, since the model cannot take it as input?"

This is probably clearest to explain by considering the example given in Figure 1. Even without explicitly considering the candidate label (e.g. EC:1.2.3.4), it is still possible to make a binary judgment based on whether the input sequence is more similar to the positive or the negative exemplars.

"Since ProtEx uses BLAST to produce the pre-training and fine-tuning datasets for exampler, it would be reasonable for ProtEx to outperform BLAST. Could the integration of BLAST results with ProtEx be considered to enhance performance? If not, maybe ProtEx totally absorbs the knowledge from BLAST"

We note that it is non-trivial to outperform BLAST on the tasks that we study. On many of the tasks, BLAST outperforms the previous best results from neural models. Given this observation, the core motivation for ProtEx is to combine the relative strengths of pre-trained neural networks with similarity-search based methods such as BLAST.

Notably, it would be relatively trivial for a retrieval-augmented model to match the BLAST baseline. The model could simply learn the heuristic of copying the prediction associated with most similar sequence retrieved by BLAST. What is interesting about ProtEx is that it consistently outperforms BLAST. Further, Figure 4 demonstrates that ProtEx can outperform BLAST at all points of their respective PR curves. Therefore, ProtEx is not only recapitulating BLAST. ProtEx can aggregate information across multiple exemplars using a learned non-linear model, transfer knowledge learned during pre-training, and learn task-specific notions of similarity between query and exemplar sequences during fine-tuning. As we show below, our proposed pre-training and fine-tuning procedures are critical for ProtEx to outperform BLAST.

Moreover, we note that on the Pfam task (Table 5), BLAST performs poorly, perhaps due to the construction of the split which ensures that the sequences in the test set have at most 25% sequence identity to the train set). ProtEx, on the other hand, achieves SOTA performance, indicating the robustness of our approach.

Finally, while we primarily focused on BLAST as the source of exemplars in this paper, as it is a widely used and understood tool, our method is not specific to BLAST. We offer some new results using Foldseek and MMSeqs2 as alternatives to BLAST and still demonstrate strong performance (please see our response to Reviewer nQPw).

(Response continued in next comment.)