[Paper Review] Evaluating representation learning on the protein structure universe

Arian R. Jamasb, Alex Morehead+9 more|PubMed|Jun 19, 2024

Protein Structure and Dynamics1 references6 citations

TL;DR

The paper introduces ProteinWorkshop, a comprehensive benchmark for pretraining and evaluating Geometric Graph Neural Networks on large-scale protein structures, showing that AlphaFold-based pretraining and denoising tasks improve representations, with equivariant models gaining more from pretraining, and ESM-2-650M with structural features matching state-of-the-art on some tasks.

ABSTRACT

We introduce <i>ProteinWorkshop</i>, a comprehensive benchmark suite for representation learning on protein structures with Geometric Graph Neural Networks. We consider large-scale pre-training and downstream tasks on both experimental and predicted structures to enable the systematic evaluation of the quality of the learned structural representation and their usefulness in capturing functional relationships for downstream tasks. We find that: (1) large-scale pretraining on AlphaFold structures and auxiliary tasks consistently improve the performance of both rotation-invariant and equivariant GNNs, and (2) more expressive equivariant GNNs benefit from pretraining to a greater extent compared to invariant models. We aim to establish a common ground for the machine learning and computational biology communities to rigorously compare and advance protein structure representation learning. Our open-source codebase reduces the barrier to entry for working with large protein structure datasets by providing: (1) storage-efficient dataloaders for large-scale structural databases including AlphaFoldDB and ESM Atlas, as well as (2) utilities for constructing new tasks from the entire PDB. <i>ProteinWorkshop</i> is available at: github.com/a-r-j/ProteinWorkshop.

Motivation & Objective

• Provide a unified, modular benchmark to evaluate protein structure encoders across tasks and model choices.
• Enable pretraining on large, diverse structure corpora (experimental and predicted) to learn informative local and global representations.
• Assess how different featurisations and symmetry constraints (invariant vs equivariant) affect downstream performance.
• Investigate the benefit of auxiliary denoising tasks and pretraining on both node-level and graph-level protein tasks.

Proposed method

• Benchmark a suite of rotation-invariant and equivariant Geometric GNNs (SchNet, EGNN, TFN, MACE, GCPNet, GearNet) across varying input featurisations.
• Use AlphaFoldDB as the primary pretraining corpus and evaluate storage-efficient dataloaders for large structure datasets.
• Apply five pretraining tasks including structure/sequence denoising, inverse folding, and pLDDT prediction, with noising schemes to corrupt coordinates or angles.
• Evaluate a range of downstream tasks spanning node-level (e.g., inverse folding, PPI site, metal binding, PTM site) and graph-level (fold, GO, reaction class, antibody developability) annotations.
• Compare against a sequence-based baseline (ESM-2-650M with structural features) to assess cross-modality effectiveness.

Figure 1: Overview of ProteinWorkshop , a comprehensive benchmark suite for evaluating pre-training and representation learning of Geometric GNNs on large-scale protein structure data.

Experimental results

Research questions

• RQ1 How do invariant vs equivariant GNNs perform across diverse protein-structure tasks?
• RQ2 What featurisation schemes and levels of structural detail (Cα, backbone, sidechains) best support representation learning for proteins?
• RQ3 To what extent do pretraining tasks and auxiliary denoising tasks improve downstream performance, and which model classes benefit most?
• RQ4 Can large-scale pretraining on AlphaFold structures close the gap between structure-based encoders and state-of-the-art sequence-based models?
• RQ5 How well do the learned representations transfer from pretraining to both local residue-level and global protein annotations?

Key findings

• Equivariant GNNs generally outperform invariant GNNs across tasks, with protein-specific architectures also showing strong performance.
• Featurising with Cα atoms, virtual angles, and backbone torsions yields strong results on many task-model combinations, suggesting backbone information aids generalization.
• Auxiliary denoising tasks (sequence and structure) consistently improve performance across most combinations, stabilising training in some cases.
• Pretraining on AlphaFoldDB with denoising-based tasks enhances downstream performance, especially for equivariant models.
• Augmenting sequence-based ESM-2-650M with structural features matches or surpasses some state-of-the-art GNNs on (super)family fold and GO prediction tasks.

Figure 2: Ranking analysis of Gene Ontology-Biological Process (GO-BP) test performance across different encoders, feature sets, and auxiliary tasks.