Paper: Evaluating Representation Learning on the Protein Structure Universe (ICLR-2024)

Reviewer comment (Experiments):
It would have been nice to see baselines for all of the downstream tasks in the benchmark using the tools in this software suite (or, lacking that, even scores copied from their respective papers if needs be). Certain tasks like "reaction class prediction" are currently missing. Part of the value-add here is the ease of running experiments on all of these tasks, and that isn't currently demonstrated in the current version of the manuscript. Also, consolidated baseline scores are useful sanity checks for reproduction experiments down the line.

Author response:
Please stay tuned for the updated manuscript containing more downstream results. We will upload this ASAP.

Reviewer comment (Evaluation):
Please provide details about the confidence intervals in Table 2.

Author response:
We report the standard deviation over three runs across different seeds. We will mention this in the revised paper.

Reviewer comment (Experiments):
Efforts to mitigate potential data leakage are crucial to ensuring the integrity of benchmarking results, as such leakage could introduce misleading elements into the research findings. Have you considered the removal of overlapping sequences between the pretraining datasets and the downstream testing datasets to further safeguard against such issues?

Author response:
Thank you for raising this point. Please consider our counter argument against why using a part of the downstream test dataset for pre-training neural networks is okay.

We believe that utilizing a portion of the downstream data, **without labels**, to find a suitable model initialization that improves performance on the end task is acceptable in our setting. Our primary goal with using a large corpus of AlphaFold DB structures is to boost downstream task accuracy, and not directly use the label information, which contains key details. This is in line with models in natural language processing like BERT (and even SciBERT), the full Wikipedia datasets were utilized for BERT pre-training without carefully removing any overlap with fine-tuning sets. 

Now, protein data may be a special case since exactly identical sequences and structures exist across datasets, unlike in language data. We somewhat mitigate this by using the non-redundant set of AFDB which are predicted structures which consider idealised geometry. There is still a semantic gap between crystal structures from the PDB (our downstream tasks) and AFDB. They are never identical.

However, we do agree that when it comes to the domain generalization properties of pre-training models, it might be better to explicitly exclude certain samples in the pre-training phases.

Another counterargument is that it’s fairly straightforward to obtain and utilize predicted protein structure information from tools like AlphaFold during pre-training in real-world settings. Researchers working on real-world protein modelling and design problems are actively starting to use AlphaFold-predicted structures to augment their datasets and are obtaining powerful gains from it. Two notable examples include:
- [Learning inverse folding from millions of predicted structures](https://www.biorxiv.org/content/10.1101/2022.04.10.487779v1), Hsu et al. (ICML oral, 2022)
- [RosettaFold2](https://www.biorxiv.org/content/10.1101/2023.05.24.542179v1), which uses AlphaFold-predicted structures to train the latest RosettaFold version.

In sum, as long as the pre-training data does not include labels and is consistent across models, we believe it is a reasonably fair methodology for improving downstream performance, and the particular scenario of AlphaFold-predicted structures is already being used in real-world scenarios to the best of our knowledge.

Reviewer comment (Experiments):
In terms of featurization, the paper seems to primarily focus on simple feature extraction methods. It might be valuable if the authors explored the integration of additional featurization techniques, such as incorporating 3D presentations like Uni-Mol (available at https://github.com/dptech-corp/Uni-Mol).

Author response:
Uni-Mol is focused on protein binding pocket and ligand (small molecule) interactions, which is **not directly applicable** to our setting. Quoting the repository, “The framework comprises two models: a molecular pretraining model that has been trained using 209M molecular 3D conformations, and a pocket pretraining model that has been trained using 3M candidate protein pocket data.”

ProteinWorkshop focuses on protein-only representation learning and self-supervised pre-training at the node and graph level. Critically, we are interested in tasks that consider the **entire** protein structure, as opposed to the setting of considering a cropped patch (the protein pocket) and its interaction with a ligand.

In the future, it would be interesting to extend our framework to biomolecular complexes, as the PDB remains the main data source. (See [related response](https://openreview.net/forum?id=sTYuRVrdK3&noteId=Ha3Lo64kpL) to Reviewer W2RU.)

Reviewer comment (Experiments):
In terms of featurization, it might also be valuable to utilize pretrained models trained on 2D data, like ESM-1b and ESM-2. Considering that Table 1 showcases results using ESM as a backbone model, which produced superior outcomes on the fold dataset, it could be beneficial to include the results for other ESM models in the table.

Author response:
At present, we chose the 650M variant of ESM2 because this is the parameter count/scale at which the ESM2 authors observed significant improvements in the model’s ability at structure-related tasks. (See Figure 1 in [the Science paper](https://www.science.org/doi/10.1126/science.ade2574) introducing ESM2.) Our results further reinforce their observation – we find that combining ESM2 650M with our structural featurisation yields extremely strong results on Fold classification tasks. 

Our goal with the benchmark was to demonstrate the utility of our ProteinWorkshop framework and to provide new insights from fairly comparing standard architecture choices. It should be fairly straightforward to swap the exact version of the ESM language model used for featurisation within ProteinWorkshop, but we believe we have already used one of the best ESM variant for structural tasks. While we cannot be certain without having done the experiment, it is a reasonable guess that a variant of ESM2 with fewer parameters, or the older generation of ESM1b models, will not perform any better than ESM2 650M nor would this change the outcome of our benchmarks.

We included ESM2 as a reference, but our focus is on structure-based geometric GNN models and scaling them up.

Reviewer comment (Experiments):
Table 1 predominantly presents results for two specific tasks. However, several other downstream tasks have not been included in this experiment. Many of these unreported downstream tasks are of substantial importance and are conspicuously absent from the paper. Is there a particular reason for not considering these tasks, and could they potentially be included to provide a more comprehensive view of the research findings?

Author response:
Thank you for raising this point. Please stay tuned for the updated manuscript containing more downstream results. We will upload this ASAP.

Reviewer comment (Experiments):
In the current version, the authors didn’t provide results on all downstream tasks.

Author response:
Please stay tuned for the updated manuscript containing more downstream results. We will upload this ASAP.

Reviewer comment (Experiments):
I am not sure whether the hyperparameters are well-searched, since the best results reported here are still worse than some existing methods. For example, the best results on Fold (considering both with and without auxiliary tasks) are still worse than ProNet [1] and CDConv [2] which don’t use any auxiliary tasks.

Author response:
Given the large number of models and featurisation schemes, we did try our best to do a consistent and fair hyperparameter search. See Appendix B5 for details -- essentially, we used the fold classification task to select the best learning rate and dropout per model and featurisation scheme. We fixed a consistently large hidden dimensionality across models, as we also wanted to focus on scaling model size and dataset size via pre-training.

While our best performing models do not outperform the best reported results on Fold, we have obtained these results in a consistent and fair experimental setup. Our goal was to demonstrate the utility of our benchmarking framework and uncover the impact of architectural considerations such as featurisation schemes, geometric GNN models, and pre-training/auxiliary tasks. We are confident that the insights from our benchmarking study will subsequently help us build models that outperform the state of the art across both node and graph level representation learning tasks. 

Finally, it is worth stating that both CDConv and ProNet are multi-scale architectures with a **highly specialised design tailored to graph/global-level prediction tasks**. It is not trivial to generalise such architectures for node-level tasks, unlike the suite of generic as well as protein-specific geometric GNNs considered in this work.

In summary, while the subject of this particular paper is to introduce ProteinWorkshop as a platform for evaluating protein representation learning, we are now actively working on building state of the art models for node and graph level protein tasks based on the insights gained!

Reviewer comment (Reproducibility):
Do ESM results in the table use pre-trained ESM weights?

Author response:
Yes, we use a frozen pre-trained ESM2 650M parameter model. We apologize for the oversight and have added this point to the paper. We chose the 650M variant of ESM2 because this is the parameter count/scale at which the ESM2 authors observed significant improvements in the model’s ability at structure-related tasks. Our results further reinforce their observation – we find that combining ESM2 650M with our structural featurisation yields extremely strong results on Fold classification tasks.

Reviewer comment (Experiments):
Some other pretraining strategies are used in Geom3D [3] and ESM-GearNet [4].

Author response:
Re. Geom3D and ESM-GearNet: We will cite these papers in the revised version. Do note that Geom3D strategies seem to be specialised for small molecules (as stated in section 4.7 of their paper) but some of them are reasonably close to those studied by us, distance/angle prediction. 

Let us provide further justification of how we selected our pre-training tasks.

We focussed on pre-training tasks that roughly fall under the category of denoising (eg. sequence, coordinates). We were particularly interested in self-supervised objectives that were:
1. Extremely scalable, so as to pre-train on the large scale AlphaFold Protein Structure Database (AFDB) of 2.4M structures; and 
2. Train protein representations at the fine-grained node level, so as to be general-purpose across the downstream tasks considered. 

Naturally, we would have liked to continue exploring more pre-training strategies, but could not do so due to the significant cost of running each pre-training experiment on 2.4M AFDB structures (total cost of GPU hours roughly being $50,000). Notably, the structural clustering of AFDB we use results in a dataset approx. 3x larger than what is used for pre-training GearNet and its variants. Thus, the scalability of the pre-training objectives was a critical consideration in selecting them for our benchmarking study.

We did not yet consider other possible tasks that have also been used in the literature, such as contrastive learning and generative modelling-inspired objectives. Most of these objectives are (1) computationally heavier and more cumbersome to set up than corruption-type objectives, making them harder to scale up, and (2) only train protein representations for the global/graph level and do not operate at the node level.

We would also like to highlight that, ultimately, ProteinWorkshop is a platform through which we certainly hope to integrate future pre-training objectives such as those from complementary studies like Geom3D. The goal of the benchmarking study in this paper was to **demonstrate the utility of our framework** and making pre-training on the largest non-redundant corpus of protein structures accessible to the community. We believe this makes ProteinWorkshop an ideal framework to advance self-supervised protein representation learning in the future.

Reviewer comment (Evaluation):
The work is missing an explanation of the limitations of the featurization schemes and pre-training tasks.

Author response:
Thank you for the suggestion – we will now include a paragraph each on the limitations of current featurization schemes as well as pre-training tasks in the revised version (which we are preparing with results on new datasets). 

**Limitations of featurization schemes:**
- All of our featurization schemes are at the residue level, with the position of the alpha Carbon atom being assigned the position of a particular node. This is conventional practice in the field due to efficiency, as all atom representations require very high GPU memory/make models extremely slow.
- The extent to which each scheme is a ‘complete’ representation of the geometry and structure of the protein residue it represents is variable. For instance, backbone-only featurizations simply ignore the orientations of the side chain atoms in the residue, so the geometric GNN must account for this information implicitly.
- However, the extent to which providing complete information about all atoms and side chain orientations is debatable, as the exact coordinates from PDB files are known to contain artifacts from structure determination via crystallography (see [ProteinMPNN](https://www.science.org/doi/10.1126/science.add2187)). This is also one of the interesting outcomes of our benchmark – letting the model implicitly learn about side chain orientation performs better or equally well as explicitly providing complete side chain atomic information.

**Limitations of pre-training tasks:**
- We only focussed on pre-training tasks that roughly fall under the category of corrupting information in the input (eg. sequence, coordinates) and tasking the model with producing the uncorrupted input. We were particularly interested in self-supervised objectives that were (1) extremely scalable, so as to pre-train on the large scale AlphaFold Protein Structure Database (AFDB) of 2.4M structures; and (2) train protein representations at the fine-grained node level, so as to be general-purpose across the downstream tasks considered.
- We did not consider other possible tasks that have also been used in the literature, such as contrastive learning and generative modelling-inspired objectives. Most of these objectives are (1) computationally heavier and more cumbersome to set up than corruption-type objectives, making them harder to scale up, and (2) only train protein representations for the global/graph level and do not operate at the node level.

Reviewer comment (Evaluation):
Would be beneficial to include a discussion about the generalizability of the benchmark results to the overall protein structure space, and how this translates to proteins not included in the current dataset.

Author response:
Indeed, we will elaborate on this further. By using the AFDB as a pre-training corpus, models trained using ProteinWorkshop should, in principle, exhibit strong generalisation to the currently known (and predicted) natural protein structure universe. The reason for this is that afdb_rep_v4, the FoldComp-curated clustering of the AFDB we used for pre-training, contains a non-redundant set of 2.4M structures, ie. each of the structures in the dataset has below 50% TM score with one another, indicating that these are unique folds. Any model trained on afdb_rep_v4 would have ‘seen’ any new natural protein fold in its pre-training corpus.

Reviewer comment (Evaluation):
Missing a discussion about how geometric models may be improved to surpass sequence-based models.

Author response:
We ultimately believe that future protein representation learning models may blur the boundaries of sequence-based vs geometric/structure-based modelling. Concurrent works submitted to ICLR have already been to explore combining geometric graph neural networks with pre-trained protein language models for improved performance on specific downstream tasks, and we believe this trend will continue for the foreseeable future. Moreover, another avenue worth considering for future research on geometric models, one particularly suited to such models, is generative modeling-based pre-training objectives. Some early works have begun to demonstrate the usefulness for downstream tasks of incorporating diffusion-based objectives as a pre-training regime for geometric graph neural networks. As more flexible generative modeling frameworks emerge, we anticipate this trend will also continue, as we will likely see methods such as geometric flow matching begin to supersede diffusion pre-training by instead pre-training models to learnably interpolate between one relevant biomolecular data domain to another. Since our benchmark is built with modularity in mind, our work is poised to quickly take advantage of new generative modeling advances for enhanced geometric pre-training.

Reviewer comment (Evaluation):
Missing (i) aggregation of methods for improving model performance; and (ii) computation of uncertainties in evaluations.

Author response:
(i) In order to extract the best performance out of models, it is natural that end-users of ProteinWorkshop may want to perform ensembling of the predictions of various methods. However, our goal with this paper was to rigorously compare architectural choices around featurization, GNN layers, and pre-training/auxiliary tasks in a fair manner, as a way of demonstrating the utility of ProteinWorkshop. We did not set out with the goal of obtaining SOTA performance on all datasets. 

(ii) Could the reviewer please clarify what they meant by uncertainties in evaluation? We did tune the model hyperparameters fairly across all models, and performed each experiment with multiple random seeds in order to obtain standard deviations across our results. This allows us to make statistically significant takeaways from our benchmarking experiments.