ProteinShake: Building datasets and benchmarks for deep learning on protein structures
Authors: Tim Kucera
Carlos Oliver
Dexiong Chen
Karsten Borgwardt
Venue: NeurIPS 2023 Datasets and Benchmarks Poster (2023)
Keywords: protein structure
benchmarking
datasets
TL;DR: Datasets and evaluation tasks for protein 3D structure data.

Abstract: We present ProteinShake, a Python software package that simplifies dataset
creation and model evaluation for deep learning on protein structures. Users
can create custom datasets or load an extensive set of pre-processed datasets from
biological data repositories such as the Protein Data Bank (PDB) and AlphaFoldDB.
Each dataset is associated with prediction tasks and evaluation functions covering
a broad array of biological challenges. A benchmark on these tasks shows that pre-
training almost always improves performance, the optimal data modality (graphs,
voxel grids, or point clouds) is task-dependent, and models struggle to generalize
to new structures. ProteinShake makes protein structure data easily accessible
and comparison among models straightforward, providing challenging benchmark
settings with real-world implications.

ProteinShake is available at: https://proteinshake.ai

ProteinShake
Building datasets and benchmarks for deep learning on protein structures
Tim Kucera 1,2,3,, Carlos Oliver 1,2,3, Dexiong Chen 1,2,3 and Karsten Borgwardt 1,2,3
DDpaarmment of Biosystems Science and Engineering, ETH Zürich, Switzerland SIB, Swiss Institute of Bioinformatics, Basel. Switzerland
Max Planck Institute of Biochemistry, Martinsried, Germany "These authors contributed equally. (kucera, oliver, dchen, borgvardt) @bieeeee...hpppeeeeeeeeeeeeeeee. waateemmmmp.... de
Abstract
We present ProteinShake, a Python software package that simplifies dataset creation and model evaluation for deep learning on protein structures. Users can create custom datasets or load an extensive set of pre-pros pre-processed datasets from biological data repositories such as the Protein Data Bank (PDB) and AlphaFoldDB. Each dataset is associated with prediction tasks and evaluation functions covering broad array of biological challenges. benchmark on these tasks shows that pre- training almost always improves performance, the optimal data modality (graphs, voxel grids, or point clouds) is task dask-dependent, and models struggle to generalize to new structures. ProteinShake makes protein structure data easily accessible and comparison among models straightforward, providing challenging benchmark settings with real-world implications.
ProteinShake is available uthttps://proteinshatee
Introduction
Over the decades, data describing protein sequence, structure, and function has been amassed in databases such as UniProt [3] and RCSB PDB K as well as in a plethora of domain-specific
databases 5 38 12 This abundance of data presents prime application for deep learning models and an opportunity to generate new insights with machine learning in challenging real-world scenarios. The recent addition of millions of high quality structure predictions from AlphaFold BC and other structure prediction algorithms are now cager utilize.
32 adds valuable resource to this pool of data, which researchers
However, the heterogeneity of protein data constitutes a problem for reproducible machine learning. Structures stored in heavy formats not developed for deep learning models, metadata is often scattered across different databases, and files can be corrupted or of poor quality. Additionally, the diverse approaches and network architectures in geometric deep learning require different represen- tations the protein structure, such as point clouds, graphs, or voxel grids. This results in highly variable data processing steps and evaluation schemes across publications, and eventually prevents comparison.
Ideally, datasets and evaluations should be standardized. This includes well-defined and transpar- ent pre-processing steps of the data and reproducible data splits in the evaluation, with selected, domain-appropriate metrics. As modern deep learning models tend to be generalists (i.e. operate the data curation process should be harmonized across differ- in multiple contexts and tasks)
37th Conference on Neural Information Processing Systems (NeurIPS 2023) Track on Datasets and Benchmarks.

---

ent biological tasks, such that model architectures can be assessed by how well they capture the underlying fundamental biological concepts. Such efforts have already propelled research across
different disciplines (e.g. ImageNet [14], OGB [22 ) including biological domains (e.g. TAPE [4]]
MoleculeNet [55), ProteinNet 3 rnaglib 34
similar effect can be expected in the recently
expanding domain of structural biology data.
We therefore propose ProteinShake harmonize the pre-processing and evaluation steps of protein structure data from various databases, converting them to deep-learning-read formats while supporting all major deep learning frameworks and architectures. Protein structures are natively available in different representations such as point clouds, voxel grids, and graphs. For evaluation, we provide standardized data splits based on sequence and structure similarity, with appropriate metrics. The tasks span a wide variety of biologically important prediction targets and machine learning problems, ranging from functional label classification, over pairwise interaction prediction and structural similarity search, to regression on ligand affinity.
The library is easily extended to include (and explicitly designed for) user-contributed datasets and tasks. The source code and our release scripts are fully open-source and regularly updated, such that datasets remain up-to-date and can be reproduced anywhere. All datasets and tasks are hosted in pre-processed form on a public dataset repository to speed up model development. Additionally, we host public leaderboard as an opportunity for researchers to compare their models.
Our design goal is to provide a simple user interface that abstracts away boilerplate code and removes the need make decisions that require expert domain knowledge, while maintaining full reproducibility and customizability. We anticipate that ProteinShake will be a useful resource for researchers working at the intersection of machine learning and structural biology, and that it will serve them as basis for reproducibility and comparability.
2
Related work
Previous works have tackled various aspects of the data preparation workflow for protein structures. ProteinShake mainly distinguishes from these works through its accessibility, culminating in single-line usage of datasets and evaluators.
We required number of features for our library under which will also briefly assess the related work: Amenability. Users with all backgrounds should be able to use it without worrying about domain-speciftt parameters that require expert biological knowledge. Availability. All data and splits should be fully reproducible on any machine, without relying on a central database. Transparency. The processing steps from raw data to the final data objects should be open-source and self-contained. Extendability. Based on the previous points, it should be easy to extend the database with new datasets and tasks. Compatibility. The data should be compatible with all popular geometric deep learning approaches. Recency. Finally, the library should provide the means to update datasets with newly available data, which requires data versioning.
One of the first efforts was ATOM3D 500 which proposed six biologically relevant protein tasks and provided Nomain-specifie splits for each. The work is accompanied by Python package which facilitates access to these tasks, and allows users to work with their own data. Datasets imported to ATOM3D are endowed with several utilities such saving, fetching, splitting, transforms and filters. However, because the library only accepts pre-built datasets, this leaves the challenge of collecting, parsing and curating the user, thereby limiting amenability, availability, and transparency. Notably, datasets large enough to use for pre-training are not included, although pre-training is becoming the basis for state-of-the-as models recently 451
TorchDrug (1]. a machine learning platform originally developed for drug discovery, is recently expanding towards protein data. The API exposes several pre-built datasets and provides highly tun- able and composable processing blocks for custom datasets, with a focus on graph representations of proteins. However, obtaining a ready-to-use dataset object requires composing series of processing steps (amenability). Further, the processing logic used to derive the datasets from primary data is not public, limiting availability, transparency and extendability.
https://prooeeinsaake.cuu at/leedeerrdeedeeebbbrrrrd
2

---

Graphein 25 is primarily a library for converting protein structure files into graphs. It provides many customizations and utility functions, and currently expanding provide pre-processed datasets (without evaluations). The focus of this library however lies on the graph conversion, leaving many of our requirements unapproached. Similar to Graphein, PyUUL [41] converts protein structure files to voxel grids and point clouds. biopandas [6 and biopython fundamental processing on structure files such as parsing and cleaning.
perform some
As none of these works have fully addressed the requirements we had as machine learners, we decided create ProteinShake with a focus on simplicity of use. The following sections detail the architecture of the framework and the datasets and tasks we provide in the current release.
3 Contribution
ProteinShake manages the full data and evaluation workflow, allowing the user to simply "plug in" model (Figure D The library provides datasets and associated evaluation tasks for benchmarking models. Each dataset object contains the logic for every processing step, starting at the download of raw data from primary resources (availability and transparency). Parameters customizable, but set with appropriate defaults, such that data can be used out of the box (amenability). ProteinShake also converts data for all common deep learning frameworks and protein structure representations (compatibility). Task and dataset interfaces are easily customizable such that users can extend elements unique to their application (extendability). The current release hosts database of eight annotated datasets (Table [ a large-scale pre-training dataset based on AlphaFoldDB, and ten prediction tasks. The releases are versioned and allow for continuous updates of the datasets (recency). Additionally, we demonstrate the use of ProteinShake by conducting experiments to assess the importance of three key features: protein representation, pre-training and data splitting.
ProteinShake
ProteinShake
Your Model
Data
Tasks
Data Integration
Representation
Similarity Splits
Evaluation
p
Graph
p
VS
Prediction
Target
pa
Point Clond
pa
Performaance
pa
Voxel Grid
H Similarity
Residine Property
Train/Test
Modal
Figure 1: ProteinShake simplifies deep learning structural biology workflows from data preparation to model evaluation. We offer annotated datasets which can be converted to several representations (point clouds, graphs, voxel grids) and then cast several popular computation frameworks (Py Torch, TensorFlow, NumPy/IAX, Pytorch Geometric, DGL, NetworkX). Datasets also serve as the basis of prediction tasks for which provide domain-specific splits and evaluation metrics.
Building datasets
core class in ProteinShake is the Dataset. holds a collection of protein objects containing the structure (I, y, coordinates and atom/residue identities) and associated annotations at the atom, residue, and protein level (e.g., binding status of a residue functional class of a protein). The protein objects also accommodate meta-information, including quality-related features such as B-factor or PL.DTT values. Other relevant features can added easily, as the protein object implemented as flexible, dictionary- -trpa storage.

---

Each dataset is constructed by implementing three steps: downloading raw data from a source database, filtering according to quality measures, and gathering annotations (possibly from other sources). All surrounding processing work such as API querying, parsing, cleaning, encoding, and storing are taken care of by the library.
Currently, ProteinShake provides annotated datasets from eight different areas of protein biology (Table for which we compile ten supervised model evaluation tasks (Section 5. We provide two large-scale (unannotated) datasets collected from the RCSB databank and AlphafoldDB which are primarily intended for pre-training and self-supervised learning. The former contains a large set of experimentally solved protein structures from the RCSB Protein Data Bank [6] (ca. 36,000 structures), the latter contains computational predictions for several organisms as well as the SwissProt database (ca. 500,000 structures) B0 77.
Table 1: Datasets currently hosted by ProteinShake. The library is designed to be straightforwardly extended, such that new datasets can be easily integrated in the future.
Dataset
Protein count
Area of protein biology
Pfam
31'109
Evolutionary relationships
Gene Ontology
32'633
Functions, Components, Pathways
Enzyme Commission
15'603
Reaction catalysis
Ligand binding
4'642
Small molecule binding
Protein-protein interfaces
2'839
Protein binding
Structure similarity
1'000
Structure alignment
Structural class
10'066
Geometric relationships
Virtual screening
38
Drug discovery
RCSB monomers
36'576
Experimental structures
AlphaFold monomers
541'143
Predicted structures
4.1 Representations and frameworks
A crucial decision when developing deep learning model for protein structures is the choice of an appropriate representation of the structure which maps raw coordinates and annotations to objects that can be used in the model. This mapping can be done in several ways, each of which have biological implications and also dictate the type model to be used.
At its core, a protein structure is simply a set of coordinates 3D space (a point cloud). This representation abstracts away the fact that each of the points (atoms or residues) physically interact with each other in a distinct way. One can represent these interactions by constructing a geometric graph, where each edge represents likely interaction as given by spatial proximity. Lastly, one might prefer a more coarse-grained representation where the protein structure is rasterized on regular voxel grid. Each of these representations rooted in a different field of machine learning and is typically used in dedicated model architectures.
Apart from the general representation, ie., those atoms and residues that lie on the interface to the solvent and directly interact with the environment. The surface carries special significance for the protein function, it constitutes the direct interface in binding events. This information can be utilized in modeling.
important nuance of protein structure is its surface,
With this in mind, one of ProteinShake 's core features is to allow easy conversions between com- monly used protein representations. ProteinShake currently supports converting any protein dataset to e-neighborhood or k-ncarest-ncighbor graphs based on inter-atomic or inter-residue distances, voxel grids which represent the occupancy of a regular 3D grid that is laid over the protein structure, and point clouds of atoms residues which are sets of labeled points in 3D space. Irrespective of the representation, all proteins are annotated with their solvent accessibility scores (computed with freesasa 33 (,)), which can either be used feature, or filter reduce the protein to its exposed surface residues.
The representations are available for the major deep learning frameworks PyTorch 42
TensorFlow
D. NumPy/JAX g and the graph learning frameworks SyTorrh--Gemeetri 16
DGL 53
and
NetworkX 2

---

Protein Classification
Residue Classification
Sclf-Supervision eig. Pretraining AlphaFold
Pairwise Residue
e.g. Gene Ouiology Engyme
e.g. Biuding Pocket Prodiction
egg. Binding Inlerface Prodiction
Class?
Mask
Pairwise Protein
Retrieval
Ranking
e.g. Structure Aligument or Protein-Proteinn Interaction
egg. Similaz Structure
egg. Drug Screening
Search
?
Figure 2: Overview of the supported prediction tasks in the current release. ProteinShake is able to model wide range of biological research questions and machine learning problems, covering regression, ranking, retrieval, classification, pre-training, metric learning, and more. ProteinShake designed drastically simplify the addition of new tasks, such that research questions be addressed quickly.
Building benchmarks
The second core object in the library is the Task which is the central utility for benchmarking. A task extends a dataset with the logic for splitting data into train/validattinntess splits and metrics to evaluate of predictions. We created and host several tasks based annotated datasets which we describe detail below and illustrate in Figure 2
5.1 Sequence and structure dataset splits
commonly employed splitting method for protein data is sequence-similaritt split. Proteins first clustered by sequence similarity (e.g., with CD-Hit (17) and the splits are then based on the clusters, ensuring similar instances end up in the same split. This ensures that model generalization can properly be assessed.
It is however known that protein structures more conserved than protein sequences e which implies that dissimilar sequences might still adopt similar structures. Hence, a sequence-based split does not guarantee structural dissimilarity between train and test data.
We create structure-similarity split by utilizing the fast structure alignment software Foldseek [51] A complete clustering is computationally prohibitive so we resort to an efficient heuristic; test and validation sets are built iteratively by randomly sampling an initial protein and then retrieving all similar instances with Foldseek, constructing structure-based clusters analogously to the requence-based split. Structural similarity defined by thresholds on the Local Distance Difference Test (LDDT) metric. Clusters are then sampled until the desired set size is reached. This split procedure serves as a novel and challenging setting that tests the ability of models to generalize in terms of structure and not just sequence homology.
All tasks are available with random, sequence, and structure splits, providing different levels of difficulty for generalization. The splits are pre-computed with similarity thresholds between 30%- 90% for the sequence split and 50%-90% for the structure split, with default of 70%. Users who wish more fine-grained control over the similarity splits can change the parameters when processing locally. ProteinShake also offers the ability to add custom split procedures, for example integrate existing benchmarks that follow other splitting strategies.
5.2 Tasks
To showcase the advantage standardization and the ability ProteinShake to integrate various types of prediction problems, we implement range of tasks covering different application areas
5

---

of biology which we categorize in three major areas: function prediction, geometric reasoning, and physical interaction modeling (Figure
The tasks cover various machine learning settings: regression, multi-class and multi-label classifica- tion, retrieval, ranking, metric learning and self-supervised learning. The inputs the model can be single proteins, pairs of proteins, or proteins paired with other instances such as chemical compounds. Predictions can be made on the atom, residue, protein level. For each task, several metries are available, and we designate one metric as the default according to literature where possible.
5.2.1 Structure-functiion relationships
There exist large, expertly curated ontologies that organize proteins into hierarchies of functional roles [ 12 [S] In an effort to better understand the relationship between protein structure and the protein's 's role(s) in the cell we build three classification tasks. In this setting, the input is a protein structure and the prediction target is one of the following classifications:
Enzyme Class (multi-class): given protein, predict the chemical reaction it catalyzes. Built using the Enzyme Commission hierarchy Default metric accuracy
4 The target is given by the top level of the hierarchy,
20
Gene Ontology (multi-labell) given a protein, predict its Gene Ontology 12 labels. The Gene Ontology describes various aspects of molecular function, cellular components, and biological processes. Default metric is Fmax 20
Protein Family (multi-class): given a protein, predict its protein family (Pfam). The Pfam annotations [S] classify proteins according to the functional attributions of their domains.
Default metric is accuracy.
Successful models these tasks have the potential deepen our understanding of structure-function relationships in proteins
20
5.2.2 Geometric reasoning
Next, we design tasks to test the ability of models to learn geometric relationships between structures. Here we ask whether the learned representations a model reflect structural properties. We define three prediction targets:
Structure Class (multi-class): given a protein, predict the correct structural class of a protein. This task built on the Structural Classification of Proteins (SCOP) 33 Default metric is accuracy.
Structural Similarity (regression): given an unaligned pair of proteins, predict the (aligned) Local Distance Difference Test (LDDT) of the structures. Target values computed after aligmment with TM-align [58] for all pairs of 1000 randomly sampled single-chain proteins. Default metric is Spearman rank correlation.
Structural Search (retrieval): retrieve a set of proteins structurally similar to a query 51 Simi- larity defined by LDDT>=0.8 after alignment with TM-align. Default metric is precision @k
Models that can learn good representation of protein structure geometries have applications in fast structure alignment [60], efficient structure search [51 33 and motif discovery 44
5.2.3 Modeling physical interactions
Physical interactions between proteins and other ligands such as small molecules and RNA form a large part of protein function. The PDBbind-CN [54] database collects and curates set of experimen- tally derived protein structures with known binding partners along with additional annotations such as binding affinity and dissociation constants. It has previously been used in binding site prediction
and ligand affinity prediction problems (29] (27 (26]. We build three tasks on top of the PDBbind -NN data and the DUDE-Z virtual screening benchmark
Ligand Affinity (regression): given protein and a chemical compound (represented as a SMILES string), predict the affinity (dissociation constant K₂) between the two molecules. Default metric is Pearson correlation 27

---

Protein Protein Interface (binary): predict the binary contact matrix between a pair of bound proteins chains. Contact threshold is set to 6.0 A. Chains are independently centered and randomly rotated simulate unbound form. Default metric median AUROC E 50 although we also provide metrics for unbalanced data.
Binding Site Detection (binary): given a protein residue, predict whether it belongs to a small molecule binding cavity. Binding site residues are those within the binding pocket provided by PDBBind [54. Default metric is Matthew's Correlation [8.
Virtual Screening (retrieval): given protein and a set of small molecules (provided as SMILES strings) containing both active ligands and non-ligands (decoys), rank the small molecules such that actives are ranked higher than the decoys. Default metric is the enrichment factor [TO]
Successful models in this family of tasks have potential to accelerate drug discovery improve our understanding of the interactome [57]
[2 and to
6 Contributing and maintenance
ProteinShake classes handle the boilerplate code of dataset and task creation such that users can implement their custom datasets and tasks by only providing the elements specific to their application. We welcome contributions. See the documentation for more information how to create datasets and tasks, and how to contribute them.
The preparation of a dataset, even on compute cluster, can easily take multiple hours up to several days. The pre-processed datasets are therefore hosted providing ready-to-us datasets in the matter of minutes. The ProteinShake database has fully automated, versioned releases to incorporate newly added protein structures and annotations. Additionally, the classes contain all processing code, such that they can be built by anyone from scratch, guaranteeing reproducibility and availability.
7 Experiments
For future reference we provide a benchmark of baselines for each task, representation, and data split, We observe that: i) the optimal protein representation differs depending on the task, ii) structure-based splits are harder to generalize to, and iii) pre-training with AlphaFoldDB enhances performance for most models and tasks.
7.1 Experimental setup
We consider the following representative deep neural networks for each protein representation: GINs [56] for graphs, PointNet++ [33 for point clouds, and a 3D CNN for voxels. These models serve as our base models for obtaining residue embeddings, which are then either used directly, pooled, or combined with other protein/molecule embeddings to perform predictions at residue level, protein level, or protein pair level, respectively. Note that the voxel model cannot easily be applied residue-level tasks, as each voxel may contain several residues. Our models are implemented in PyTorch 42] and PyG [16]. For each model, we use minimal hyper-parameter tuning to show proof-of-concept results. More advanced structure-an aware machine learning models potentially boost performances. We cordially invite researchers to contribute their own models to our leaderboard. The data parameters are left to the default settings of ProteinShake. Graphs are constructed using an c-neighborhood with e = 8.0 A, the voxel grids are constructed using a voxel size of 2.0 A and 35 35 35 voxel grid size. We report model size, runtime and memory usage in Appendix Table 6 We refer the interested reader to our code base for the details of model architectures and hyper- paremmett settings.
could
Additionally to supervised training on the respective task training data, we pre-train models on structures from AlphaFoldDB which are provided by ProteinShake. For simplicity, only consider a masking residue strategy, adapted from techniques used in natural language modeling (31.. The masking strategy consists of randomly sampling set of residues mask, by replacing the true residue type with a special mask label. Finally, we independently maximize the log likelihood the
With permanent DOl object identifiers issued by Zenodo. See documentation for list of releases.
7

---

true residue type given the structure and the remaining unmasked residues as context. We add linear layer on top of the pre-trained base model and finetune the entire model using task-specific losses
7.2 The optimal protein representation differs depending on the task
Table shows a comparison model performances on different tasks and representations (based on the random split, results for the other splits can be found in the Appendix). While the graph model generally performed the best with large margin, some tasks are better modeled with point clouds or voxel grids.
The tasks describe different physical processes and are heterogeneous in their prediction targets. The data representations the other hand capture different aspects of protein structure and also the corresponding model architectures process information in very different ways. We hence hypothesize that the different models/representttion capture different aspects of proteins which are more less important for solving a given task.
The tasks that are better solved with point clouds or voxel grids are Ligand Affinity, Protein Protein In- terface and Structure Similarity, which heavily rely on spatial, geometric information, which poorly modeled by graph. We conclude that researchers should consider priori which representation best fits their task at hand include the representation as part of hyperparameter search.
Table 2: Comparison of models trained with different representations of protein structure across various tasks, on a random data split. The optimal choice representation depends on the task. Shown are mean and standard deviation across four runs with different seeds. The Voxel model is not applicable residue-leved tasks.
Representation
Graph
Point
Voxel
Task
Binding Site
0.721 + 0.010
0.609 = 0.006
Enzyme Class
0.790 0.007
0.712 = 0.016
0.643 + 0.026
Gene Ontology
0.704 + 0.001
0.580 = 0.002
0.602 + 0.018
Ligand Affinity
0.670 0.019
0.683 = 0.003
0.689 = 0.013
Protein Family
0.728 + 0.004
0.609 = 0.004
0.668 + 0.005
Protein-Protein Interface
0.883 0.050
0.974 0.003
Structural Class
0.495 + 0.012
0.293 = 0.013
0.337 + 0.011
Structural Similarity
0.598 0.018
0.627 = 0.006
0.645 + 0.020
7.3 Generalization is harder in structure-based splits
Figure 3 shows comparison of data splits, across tasks and representations. Generally, models decrease in performance going from a random split, to a sequence split, to a structure split. This indicates that models struggle to generalize splits based sequence and structure similarity, Vice versa, this implics that sequence and structure similarity provide information that models can utilize to improve their predictions.
In scenario where a model applied to out-of-distribution (OOD) data is advisable construct appropriate train/test splits to correctly assess generalization capability. Most real world applications require model to work on OOD data. Hence, whenever model is used to predict properties of proteins that are either dissimilar in sequence in structure to the train set, we advise to use the corresponding split ProteinShake for evaluation.
The evaluations presented here were split on moderate 70% similarity threshold, where the effect is already drastic. across tasks and models. We emphasize that real world applications regularly present such distribution shifts. The data split type should be carefully considered, and in case of doubt the most stringent structure split should be used.
All code for training and pre-training is available at httpps///gsthab..com/Boorrrrrrr proteinshake_model

---

Random Split
Sequence Spiit
Structure Split
Enzyme Class
Protein Family
Ligand Affinity
Binding Site
0.8
0.6
Accuracy
0.6
Accuracy
Spearman R
0.4
MCC
0.2
0.2
0.0
0.0
Graph Point Voxel
Graph Point Voxel
Grash Point Voxel
Graph
Point
Structure Similarity
Structural Class
Gene Ontology
Protein-Protein Interface
1.00
0.6
(median)
Speamann
Accuracy
0.75
0.4
0.25
0.0
0.0
0.0
0.00
Graph
Point
Voxel
Graph
Point
Voxel
Grash
Point
Voxel
Graph
Point
Figure 3: Comparison of random, sequence, and structure splits across tasks and representations. Models generalize less well to sequence and structure splits, respectively. The split should be chosen according to the prediction problem.
Graph
Point
LA
SC
EC
Taak
Tbok
GO
Took
G0
LAA
60
PP1
85
Improvement [56]
Improvement (56
Improvement [56]
Figure 4: Relative improvement due to pre-training across tasks and representations. In most settings, performance is substantially improved by pre-training with AlphaFoldDB. Tasks abbreviated with their initials. Values are relative to the metric values obtained from the supervised model without pre-training.
7.4 Pre-training with AlphaFoldDB enhances performance across models and tasks
Unsupervised pre-training with unlabeled data has become a standard strategy to improve model performance supervised tasks [1 59 K5]
We can also see this effect for our prediction tasks on experimentally determined protein structures, even though we used predicted structures from AlphaFoldDB for pre-training (Figure D. Pre-training substantially improves model performance in the majority of test cases, however there are some model-task combinations where initialization with pre-trained weights is even detrimental. This means that, in some cases, the parameters learned under the pre-training objective (here residue masking) are less informative for solving the task than random initializations. The mechanisms of pre-training are still discussed pre-tre pre-training objectives might be more appropriate for certain tasks and models. ProteinShake aims simplify further research topics such this one by making tasks and model architectures directly comparable.
39 p 2 other
The availability of high quality structure predictions presents many possibilities for improving perfor- mance on supervised tasks. We strongly advise researchers make use of this data. ProteinShake provides the SwissProt database from AlphaFoldDB readily accessible pre-processed dataset.

---

8 Conclusion
We presented ProteinShake, an open-source Python package which unifies data preparation for machine learning on protein 3D structures, eliminating one of the most tedious and error-prone prone stages of the workflow for practitioners. The tool highly flexible while providing rich functionality in simple, accessible and extensible way. Our experiments demonstrate this flexibility and provided some imperatives for model development. We also showcased the capabilities of the library as a framework for building datasets and benchmarks by implementing diverse set of tasks, covering many biological research questions and machine learning problems. A limitation of these benchmarks is that they are simplified approximations to complex biological problems, made with inter-task comparability in mind (e.g. a uniform splitting strategy across tasks). Future releases of our benchmarks will include further task-specific considerations and will be complemented with re-implementations of existing benchmarks.
We hope to have laid foundation for researchers to now integrate both established and novel datasets and benchmarks. Importantly ProteinShake was built explicitly to facilitate community collaboration, we therefore invite researchers to contribute custom datasets, tasks, and leaderboard submissions to the library. We believe ProteinShake will improve reproducibility and greatly facilitate machine learning model development in structural protein biology.
References
[1] Martín Abadi. Tensorflow: learning functions at scale. In Proceedings of the 21 ACM
SIGPIAN International Conference on Functional Programming, 2016.
[2] Ahmed Alajrami and Nikolaos Aletras. How does the pre-training objective affect what large language models learn about linguistic properties? arXiv preprint arXiv:2203. A.NV::203./0415. 2022.
[3] Mohammed AlQuraishi. ProteinNet: standardized data set for machine learning of protein structure. BMC bioinformatics, 20(1):1-10, 2019.
[4] Amos Bairoch. The ENZYME database in 2000. Nucleic Acids Research, 28(1):304-300 305, 2000.
[5] Alex Bateman, Lachlan Coin, Richard Durbin, Robert D Finn, Volker Hollich, Sam Griffiths- Jones, Ajay Khanna, Mhairi Marshall, Simon Moxon, Erik LL Sonnhammer, et al. The Pfam protein families database. Nucleic Acids Research, 32(supp__I:DD13-_ D141, 2004.
[6] Helen M Berman, John Westbrook, Zukang Feng, Gary Gilliland, Talapady N Bhat, Helge Weissig, Ilya N Shindyalov, and Philip E Bourne. The protein data bank. Nucleic Acids Research, 28(1):235-244 2000.
[7] Brigitte Boeckmann, Amos Bairoch, Rolf Apweiler, Marie-Claude Blatter, Anne Estreicher, Elisabeth Gasteiger, Maria Martin, Karine Michoud, Claire O'Donovan, Isabelle Phan, et al. The SWISS-PROT protein knowledgebase and its supplement TrEMBL in 2003. Nucleic acids research, (11:33665-37) -60-37 2003.
[8] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Minn Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[9] Dexiong Chen, Leslie O' O'Bray, and Karsten Borgwardt. Structure-aware transformer for graph representation learning. In International Conference on Machine Learning, pages 3469-3489. PMLR, 2022.
[10] Hongming Chen, Paul D Lyne, Fabrizio Giordanetto, Timothy Lovell, and Jin Li. On evaluating molecular-docking methods for pose prediction and enrichment factors. Journal of Chemical Information and Modeling, 46(1) 46(1):401-415 2006.
[11] Peter JA Cock, Tiago Antao, Jeffrey T Chang, Brad A Chapman, Cymon Cox, Andrew Dalke, Iddo Friedberg, Thomas Hamelryck, Frank Kauff, Bartek Wilezynski, et al. Biopython: freely available Python tools for computational molecular biology and bioinformatics. Bioinformatics, 1212:-::422112333 2009.
10

---

[12] Gene Ontology Consortium. The gene ontology (go) project in 2006. Nucleic acids research, [=ssupp_)):D3222211 2222, 2006.
[13] UniProt Consortium. Uniprot: hub for protein information. Nucleic Acids Research, 43(DI):D204-DD) D212, 2015.
[14] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large- scale hierarchical image database. 2009 IEEE conference on computer vision and pattern recognition, pages 248-255 Ieee, 2009.
[15] Dumitru Erhan, Aaron Courville, Yoshua Bengio, and Pascal Vincent. Why does unsupervised pre-training help deep learning? In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 201-208. JMLR Workshop and Conference Proceedings, 2010.
[16] Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. arXiv preprint arXiv: /903.02428, 2019.
[17] Limin Fu, Beifang Niu, Zhengwei Zhu, Sitao Wu, and Weizhong Li. CD-HIT: accelerated for clustering the next-generation sequencing data. Bioinformatics, 28(23):310-33152 2012.
[18] Tiziano Gallo Cassarino, Lorenza Bordoli, and Torsten Schwede. Assessment of ligand binding site predictions in CASP10. Proteins: Structure, Function, and Bioinformatics, 82:154-166 2014.
[19] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning. pages 1272-1272. PMLR, 2017.
[20] Vladimir Gligorijević, Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Berenberg, Tommi Vatanen, Chris Chandler, Bryn C Taylor, Ian M Fisk, Hera Vlamakis, al. Structure- protein function prediction using graph convolutional networks. Nature Communications, 12(1):3166 2021.
[21] Aric A. Hagberg, Daniel A. Schult, and Pieter Swart. Exploring network structure, dynamics, and function using networkx. In Gaël Varoquaux, Travis Vaught, and Jarrod Millman, editors, Proceedings of the 7th Python in Science Conference, pages 11-15, Pasadena, CA USA, 2008.
[22] Weihua Hu, Matthias Fey, Marinka Zitiik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benehmark: Datasets for machine learning on graphs. Advances in Neural Information Processing Systems, 33:22118-22111 2020.
[23] Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zimik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre- vre-tranning graph neural networks. arXiv preprint arXiv: arXiv: /2265, 2019.
[24] Kristoffer Illergärd, David H Ardell, and Ame Elofsson. Structure is three to ten times more conserved than sequence- study of structural response in protein cores. Proteins: Structure, Function, and Bioinformatics, 77(3):499-500 2009.
[25] Arian Jamasb, Ramon Viñas Torné, Eric Ma, Yuanqi Du, Charles Harris, Kexin Huang, Dominic Hall, Pietro Lió, and Tom Blundell. Graphein-a python library for geometric deep learning and network analysis on biomolecular structures and interaction networks. Advances Neural Information Processing Systems, 35:27153-27167, 2022.
[26] José Jiménez, Stefan Doerr, Gerard Pertinez-Rosell Alexander S Rose, and Gianni De Fabritiis. DeepSite: protein-binding site predictor using 3D-convolutional neural networks. Bioinformat- ics, 33(19)):006--3442 2017.
[27] José Jiménez, Miha Skalic, Gerard Martinez-Rosell, and Gianni De Fabritiis. KDEEP: protein- ligand absolute binding affinity prediction via 3D-connolutional neural networks. Journal of Chemical Information and Modeling, 58(2):287-296, 2018.

---

[28] Bowen Jing. Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learn- ing from protein structure with geometric vector perceptrons. arXiv preprint wxXv:2009.01411, 2020.
[29] Derek Jones, Hyojin Kim, Xiaohua Zhang, Adam Zemla, Garrett Stevenson, WF Drew Bennett, Daniel Kirshner, Sergio E Wong, Felice C Lightstone, and Jonathan E Allen. Improved protein-ligand binding affinity prediction with structure-based deep fusion inference. Journal of Chemical Information and Modeling, 614):115331152, 2021.
[30] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, al. Highly accurate protein structure prediction with AlphaFold. Nature, (96(7873)5583-888 2021.
[31] Jacob Devlin Ming- Wei Chang Kenton and Lee Kristina Toutanova. BERT: Pre-training deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, 2019
[32] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scau prediction of atomic-level protein structure with language model. Science, 379(6637):1123 1130, 2023.
[33] Yang Liu, Qing Ye, Liwei Wang, and Jian Peng. Learning structural motif representations for efficient protein structure search. Bioinformatics, 34(17):i773-i7)) 2018.
[34] Vincent Mallet, Carlos Oliver, Jonathan Broadbent, William L Hamilton, and Jérôme Waldispühl. RNAglib: a python package for RNA 2.5 graphs. Bioinformatics, 38(5): 1459, 2022.
[35] Huanru Henry Mao. A survey on self-supervised pre-training for sequential transfer learning in neural networks. arXiv preprint arXiv:2007.00800 2020.
[36] Simon Mitternacht. FreeSASA: An open source C library for solvent accessible surface area calculations. F1000Research, 5, 2016.
[37] Simon Mitternacht.
FreeSASA-Python.
httpss//githubb.om//rreesss
freesasa-pythhn 2022.
[38] Alexey G Murzin, Steven E Brenner, Tim Hubbard, and Cyrus Chothia. SCOP: structural classification of proteins database for the investigation of sequences and structures. Journal of molecular biology, 247(4):536-54)) 540, 1995.
[39] Alejandro Newell and Jia Deng. How useful self superviedd pretraining for visual tasks? In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
[40] Carlos Oliver, Dexiong Chen, Vincent Mallet, Pericles Philippopoulos, and Karsten Borgwardt. Approximate network motif mining via graph learning. arXiv preprint 0tti::20060010000 2022.
[41] Gabriele Orlando, Daniele Raimondi, Ramon Duran-Roon Romaña, Yves Moreau, Joost Schymkowitz, and Frederic Rousseau. PyUUL provides an interface between biological structures and deep learning algorithms. Nature Communications, 13(1):1-9, 2022.
[42] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32, 2019.
[43] Charles Ruizhongtai Qi, Yi, Hao Su, and Leonidas Guibas. PointNet+4:: Deep hierarchical feature learning on point sets in metric space. In Advances in Neural Information Processing Systems, 2017.
[44] Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Peter Chen, John Canny, Pieter Abbeel, and Yun Song. Evaluating protein transfer learning with TAPE, Advances Neural Information Processing Systems, 32, 2019.
12

---

[45] Roshan Rao, Joshua Meier, Tom Sercu, Sergey Ovchinnikov, and Alexander Rives. Transformer protein language models are unsupervised structure learners. Biorxiv, pages 2020-12, 2020.
[46]
Sebastian Raschka. BioPandas: Working with molecular structures in pandas DataFrames. Open Source Softw., 2(14):279, 2017.
[47] Reed M Stein, Ying Yang, Trent E Balius, Matt J O O'Meara, Jiankun Lyu, Jennifer Young, Khanh Tang, Brian K Shoichet, and John Irwin. Journal of Chemical Information and Modeling, 61(2):699 691,4449,,,, 2021.
unmatchhd decoys in docking benchmarks.
[48] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pages 10347-10357, 2021.
[49] Raphael Townshend, Rishi Bedi, Patricia Suriana, and Ron Dror. End-to-end learning on 3D protein structure for interface prediction. Advances in Neural Information Processing Systems, 32, 2019.
[50] Raphael JL Townshend, Martin Vögele, Patricia Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Bowen Jing, Brandon Anderson, Stephan Eismann. et al. A1OM3D: Tasks on molecules three dimensions. arXiv preprint arXiv:2012.04035, 2020.
[51] Michel van Kempen, Stephanie S Kim, Charlotte Tumescheit, Milot Mirdita, Jeongjae Lee, Cameron LM Gilchrist, Johannes Söding, and Martin Steinegger. Fast and accurate protein structure search with Foldseek. Nature Biotechnology. pages 1-4, 2023.
[52] Kaili Wang, Renyi Zhou, Yaohang Li, and Min Li. DeepDTAF: deep learning method to predict pignei binding affinity. Briefings in Bioinformatics, 22b5bbbaa 2021.
[53] Minjie Yu Wang. Deep Graph Library: Towards efficient and scalable deep learning graphs. ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
[54] Renxiao Wang, Xueliang Fang, Yipin Lu, Chao- Yie Yang, and Shaomeng Wang. The PDBbind database: methodologies and updates. Journal of Medicinal Chemistry, 48 48(12):411-4411, 2005.
[55] Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. MoleculeNet: benchmark for molecular machine learning. Chemical Science, 9():513-530, 2018.
(56) Keyulu Xu, Weihua Hu, Jure Leskovee, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint
2018
[57] Mengying Zhang, Qiang Su, Yi Lu. Manman Zhao, and Bing Niu. Application of machine learning approaches for protein proteen interactions prediction. Medicinal Chemistry, 13(::5506- 514, 2017
[58] Yang Zhang and Jeffrey Skolnick. TM-align: a protein structure alignment algorithm based on the TM-score. Nucleic Acids Research, 33(7):2302-2309 2005.
[59] Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining, arXiv preprint arXiv:2203.06115. 2022.
[60] Jianwei Zhu, Sheng Wang, Dongbo Bu, and Jinbo Xu. Protein threading using residue co- variation and deep learning. Bioinformatics, 34(13):1233-1273, 2018.
[61] Zhaocheng Zhu, Chence Shi. Zuobai Zhang, Shengchao Liu, Minghao Xu, Xinyu Yuan, Yang- tian Zhang, Junkun Chen, Huiyu Cai, Jiarui Lu. et al. TorchDrug: A powerful and flexible machine learning platform for drug discovery. arXiv preprint arxiv:2202.08833 08320, 2022.
13


Reviews:
[{"review_id": "lqqQbWAt1o", "reviewer": "Reviewer_Da4b", "summary": "", "questions": "", "limitations": "Yes", "rating": 9, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "Major strengths of this work are: 1) the ability to split data using either sequence or structure-based similarity; the latter is particularly important for structure-based tasks and is often not used by those in the field. 2) The generation and benchmarking of three different ways of representing protein structures (graphs, point clouds, or voxels). 3) Well-described benchmarks on popular tasks.", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "mDLe4lS2FF", "reviewer": "Reviewer_j1GE", "summary": "", "questions": "", "limitations": "This work is a great step towards enabling a larger community to work with protein data. Naturally, also due to the relatively short conference format, it cannot replace a rigorous study of the subject, which might be in many cases necessary to make substantial progress.", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "- The paper proposes tasks for models with different objectives, modelling structure-function relationships, geometric protein structure and physical interactions. Similarly, the paper offers different splits to reflect real-world settings and to adjust the difficulty of the task in terms of generalization.\n- The pipeline seems easy to work with and is accessible with many different frameworks (e.g., PyTorch/TensorFlow, PyG/DGL, etc.).\n", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "9N7CvZLA3n", "reviewer": "Reviewer_8GvG", "summary": "", "questions": "", "limitations": "-\tUsers cannot add standard test datasets in the field into the system to control the generation of the training datasets such that the similarity between the training datasets and the standard test datasets is below a threshold.\n-\tDifferent prediction tasks may need different evaluation metrics. However, there is little discussion about the evaluation metrics. Some evaluation metrics such as AUROC for protein interface prediction task may not be appropriate because of the highly imbalance distribution of the positive and negative classes. The protein interface prediction task describe in this work is a very challenging task. But the  AUROC-mean score (0.974) reported in this work is very high, probably because it does not consider how well the method can predict the positive cases.\n-\tSome prediction tasks may be ill defined. For instance, predicting if a pair of residues from two proteins are in their interaction interface is very similar to the inter-protein contact prediction widely used in the field. However, the relationship between the two is not discussed. The latter appears to be better defined and easier to understand. The practical value of another task (StructureSearchTask (retrieval): retrieve a set of proteins structurally similar to a query) is not articulated provided that FoldSeek can quickly retrieve a set of proteins structurally similar to a query. \n-\tThis website https://borgwardtlab.github.io/proteinshake/#leaderboard does not provide secured access.  \n-\tThere is lack of rationale why freesasa was chosen over other tools such as DSSP to calculate solvent accessibility scores.", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "+ Allow users to use three different representations of protein structures and provide some guidance for users to choose different representations for different tasks\n+ The datasets are in the format ready for being used by some standard deep learning packages.\n+ Three ways of splitting datasets are useful.\n+ Clustering proteins to create splits is useful.\n+ Including the model evaluation into the pipeline is valuable.\n+ Pretraining models on AlphaFold predicted structures is very useful.", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "HAXUDVYPs9", "reviewer": "Reviewer_XEHS", "summary": "", "questions": "", "limitations": "The authors have not explicitly addressed limitations of their proposed benchmark set.", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "## Significance and relevance to the broader research community\nMachine Learning for biology is a subfield of increasing importance. One of the limiting factors for ML researchers to enter the field is  that considerable domain-expertise is necessary to construct meaningful datasets and tasks. Providing benchmarks like the one proposed here removes some of these barriers and is thus potentially very valuable to the field.\n\n## quality of the research\nThe approach taken by the authors is systematic and seems well thought trough. The authors set up a series of criteria for benchmarks/datasets to be useful to the community: amenability, availability, transparency, extendibility and compatibility, and address these systematically in the design of their software library.  Generally, the tasks seem meaningful, and care has been taken to use relevant splits for training/test.\n\n## clarity of paper\nThe paper is very clearly written and supported by high quality figures.\n\n## Accessibility, Accountability, and transparency\nMy impression from the paper and website is that care has been taken to ensure accessibility and transparency. The authors claim that the database will be \"regularly updated\", and provides code for releasing new versions of the datasets, suggesting that this is more than an empty promise.", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}]

Author rebuttal:

Meta-review: The reviewers are generally enthusiastic about this work and consider this dataset and benchmark make it easy for the community to develop machine learning methods to study proteins. The concerns raised by the reviewers were well addressed by the authors during the rebuttal period.
Decision: Accept (Poster)