SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We consider the following representative deep neural networks for each protein representation: GINs [56] for graphs, PointNet++ [43] 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 in 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-aware machine learning models [19, 9, 28] could 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  $\epsilon$ -neighborhood with  $\epsilon = 8.0\mathring{\mathrm{A}}$ , the voxel grids are constructed using a voxel size of  $2.0\mathring{\mathrm{A}}$  and a  $35\times 35\times 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-parameter settings.
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, we only consider a masking residue strategy, adapted from techniques used in natural language modeling [31]. The masking strategy consists of randomly sampling a set of residues to mask, by replacing the true residue type with a special mask label. Finally, we independently maximize the log likelihood of the
2With permanent DOI object identifiers issued by Zenodo. See our documentation for a list of releases.
true residue type given the structure and the remaining unmasked residues as context. We add a linear layer on top of the pre-trained base model and finetune the entire model using task-specific losses.[3]
7.2 The optimal protein representation differs depending on the task
Table 2 shows a comparison of 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 on 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/representations capture different aspects of proteins which are more or less important for solving a given task.
The tasks that are better solved with point clouds or voxel grids are Ligand Affinity, Protein Protein Interface and Structure Similarity, which heavily rely on spatial, geometric information, which is poorly modeled by a graph. We conclude that researchers should consider a priori which representation best fits their task at hand or include the representation as part of a 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 of representation depends on the task. Shown are mean and standard deviation across four runs with different seeds. The Voxel model is not applicable to residue-level tasks.
Representation Task	Graph	Point	Voxel
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 a 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 to splits based on sequence and structure similarity. Vice versa, this implies that sequence and structure similarity provide information that models can utilize to improve their predictions.
In a scenario where a model is applied to out-of-distribution (OOD) data it is advisable to construct appropriate train/test splits to correctly assess generalization capability. Most real world applications require a model to work on OOD data. Hence, whenever a model is used to predict properties of proteins that are either dissimilar in sequence or in structure to the train set, we advise to use the corresponding split of ProteinShake for evaluation.
The evaluations presented here were split on a 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.
3All code for training and pre-training is available at https://github.com/BorgwardtLab/proteinshake_models
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.
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 are 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 on supervised tasks [15, 59, 35]. 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 4). 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 [39, 2, 23], other pre-training objectives might be more appropriate for certain tasks and models. ProteinShake aims to simplify further research on topics such as this one by making tasks and model architectures directly comparable.
The availability of high quality structure predictions presents many possibilities for improving performance on supervised tasks. We strongly advise researchers to make use of this data. ProteinShake provides the SwissProt database from AlphaFoldDB readily accessible as a pre-processed dataset.
Random Split Sequence Split Structure Split
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 stages of the workflow for practitioners. The tool is highly flexible while providing rich functionality in a 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 a 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 a 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.
[1] Martín Abadi. Tensorflow: learning functions at scale. In Proceedings of the 21st ACM SIGPLAN 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.10415, 2022.
[3] Mohammed AlQuraishi. ProteinNet: a 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–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(suppl_1):D138-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-242, 2000.
[7] Brigitte Boeckmann, Amos Bairoch, Rolf Apweiler, Marie-Claude Blatter, Anne Estreicher, Elisabeth Gasteiger, Maria J Martin, Karine Michoud, Claire O'Donovan, Isabelle Phan, et al. The SWISS-PROT protein knowledgebase and its supplement TrEMBL in 2003. Nucleic acids research, 31(1):365-370, 2003.
[8] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[9] Dexiong Chen, Leslie 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):401-415, 2006.
[11] Peter JA Cock, Tiago Antao, Jeffrey T Chang, Brad A Chapman, Cymon J Cox, Andrew Dalke, Iddo Friedberg, Thomas Hamelryck, Frank Kauff, Bartek Wilczynski, et al. Biopython: freely available Python tools for computational molecular biology and bioinformatics. Bioinformatics, 25(11):1422-1423, 2009.