[Paper Review] Protein Representation Learning by Geometric Structure Pretraining

Zuobai Zhang, Minghao Xu+5 more|arXiv (Cornell University)|Mar 11, 2022

Machine Learning in Bioinformatics40 citations

TL;DR

The paper introduces GearNet, a geometry-aware relational graph neural network pretrained on protein structures to learn structure-based representations, outperforming many baselines on function and fold-related tasks with substantially less pretraining data. It combines multiview contrastive learning and self-prediction tasks to leverage AlphaFold-predicted structures.

ABSTRACT

Learning effective protein representations is critical in a variety of tasks in biology such as predicting protein function or structure. Existing approaches usually pretrain protein language models on a large number of unlabeled amino acid sequences and then finetune the models with some labeled data in downstream tasks. Despite the effectiveness of sequence-based approaches, the power of pretraining on known protein structures, which are available in smaller numbers only, has not been explored for protein property prediction, though protein structures are known to be determinants of protein function. In this paper, we propose to pretrain protein representations according to their 3D structures. We first present a simple yet effective encoder to learn the geometric features of a protein. We pretrain the protein graph encoder by leveraging multiview contrastive learning and different self-prediction tasks. Experimental results on both function prediction and fold classification tasks show that our proposed pretraining methods outperform or are on par with the state-of-the-art sequence-based methods, while using much less pretraining data. Our implementation is available at https://github.com/DeepGraphLearning/GearNet.

Motivation & Objective

• Motivate learning protein representations from 3D structures rather than solely from sequences.
• Develop a structure-based encoder that captures spatial and chemical interactions in proteins.
• Pretrain the encoder using multiview contrastive learning and self-prediction tasks to leverage unlabeled structures.
• Demonstrate that structure-based pretraining can match or exceed sequence-based methods with fewer pretraining samples.

Proposed method

• Propose GearNet, a geometry-aware relational graph neural network that builds residue graphs with sequential, radius, and KNN edges and applies relational graph convolutions with edge-type conditioning.
• Enhance GearNet with an Edge Message Passing Layer (GearNet-Edge) that performs sparse message passing between edges via a line-graph construction and angular relations.
• Use multiview contrastive learning over biologically meaningful substructures (subsequence and subspace crops) to align representations, with a cosine-similarity-based InfoNCE loss.
• Introduce four self-prediction pretraining tasks (residue type, distances, angles, dihedrals) predicting masked geometric/physicochemical properties at varying levels of residue groups.
• Pretrain on AlphaFold DB structures (365K proteome-wide + 440K Swiss-Prot) and fine-tune on downstream tasks including EC number prediction, GO term prediction, fold classification, and reaction classification.

Experimental results

Research questions

• RQ1 Can a geometry-aware GNN pretrained on protein structures yield representations that improve downstream function and fold predictions compared to sequence-based pretraining?
• RQ2 Do edge-focused message passing and contrastive pretraining on structural motifs provide significant gains over baseline structure-based and sequence-based encoders?
• RQ3 How do subsequence and spatial substructure crops influence the effectiveness of multiview contrastive pretraining?
• RQ4 What is the impact of combining GearNet with IEConv layers on fold classification versus function prediction tasks?

Key findings

• GearNet-based encoders outperform several baselines on function prediction and fold classification without pretraining.
• Edge message passing (GearNet-Edge) yields strong improvements, particularly on EC, GO-BP, and GO-MF tasks, and remains competitive on GO-CC.
• Pretraining with the proposed methods substantially boosts performance, often surpassing or matching state-of-the-art sequence-based encoders that are trained on orders of magnitude more data.
• Multiview contrastive pretraining (subsequence and subspace crops) achieves best overall results across EC, GO, Reaction, and Fold tasks.
• Pretrained structure-based encoders can perform on par with or better than sequence-based models despite using far fewer pretraining structures (less than one million).