PaperlayerView exact revision on arXiv

Source-linked AI summary

Protein Representation Learning by Geometric Structure Pretraining

Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, Jian Tang

arXiv:2203.06125v5cs.LG

TL;DR

Protein representation learning needs methods that exploit 3D structure, because sequence-based pretraining does not explicitly capture structural information linked to protein function. The paper introduces GearNet with geometric pretraining based on contrastive and self-prediction tasks, achieving competitive or better results than sequence-based encoders with less pretraining data.

  • Problem

    Sequence-based protein pretraining uses abundant unlabeled sequences but does not explicitly leverage structural information known to determine protein functions.

  • Method

    GearNet learns geometric protein representations with relational and edge message passing, pretrained through multiview contrastive learning and masked prediction of geometric or physicochemical attributes.

  • Results

    Across function-prediction and fold-classification benchmarks, the proposed methods outperform or match state-of-the-art sequence-based encoders while using much less pretraining data.

  • Takeaways & Limitations

    Structure-based self-supervised pretraining provides a basis for protein representation learning across diverse property-prediction tasks.

  • Takeaways & Limitations

    Evaluation uses specified dataset splits, including sequence-identity constraints, and contrastive learning depends on choosing substructures large enough to retain information but small enough to avoid trivial views.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Protein function annotation remains difficult as sequence data expand, while existing sequence-based pretraining does not explicitly leverage structural information. The paper introduces GearNet and geometric pretraining to learn structure-based representations for diverse protein-property tasks.

  • Low-cost sequencing has produced massive protein-sequence volumes, but annotating new sequences remains costly and time-consuming.
  • Sequence-based protein encoders trained on millions of unlabeled sequences do not explicitly capture structural information known to determine protein functions.
  • Existing structure-based encoders do not explicitly model interactions between edges, while structure-based pretraining remains limited by scarce experimentally determined structures.
  • GearNet encodes sequential and structural edges in protein residue graphs and uses relational message passing with sparse edge-level message passing.
  • Geometric pretraining combines multiview contrastive learning with masked self-prediction of residue types, distances, angles, and dihedral angles.
  • Across EC, GO, fold, and reaction benchmarks, GearNet-Edge outperforms existing protein encoders on most supervised tasks, while pretraining on fewer than a million samples matches or exceeds sequence-based encoders trained on million- or billion-scale datasets.

2 RELATED WORK

Protein representation learning has used sequences, MSAs, and structures, with structure-based methods encoding spatial information through CNNs and GNNs. Structure-based pretraining remains relatively unexplored, and this paper combines relational and edge message passing with contrastive and self-prediction objectives.

  • Prior protein representation methods use amino acid sequences, multiple sequence alignments, or structures for downstream tasks including function prediction.
  • Structure-based methods encode spatial information using 3D CNNs or graph neural networks, motivated by the relationship between protein structure and function.
  • AlphaFoldDB versions 1 and 2 provided the largest protein structure database available before March 2022 for the paper’s pretraining.
  • Some approaches incorporate structural information into pretrained sequence models during pretraining or finetuning.
  • Few works focus on structure-based protein pretraining; this paper adds relational and edge message passing, novel substructure augmentations, and four self-prediction tasks.

3 STRUCTURE-BASED PROTEIN ENCODER

GearNet is a structure-based protein encoder that represents residues and their spatial relationships with relational graphs, while GearNet-Edge additionally models interactions between edges.

  • 3.1 GEOMETRY-AWARE RELATIONAL GRAPH NEURAL NETWORK: GearNet represents proteins as residue-level relational graphs whose nodes are alpha carbons and whose edge features include residue types, sequential distances, and spatial distances.The graph uses 3D coordinates together with sequential, radius, and K-nearest-neighbor edges.
  • 3.1 GEOMETRY-AWARE RELATIONAL GRAPH NEURAL NETWORK: Relational graph convolution shares one learnable kernel within each edge type, balancing model capacity against memory cost.Neighboring node features are aggregated separately by relation type, followed by batch normalization, activation, and a residual update.
  • 3.2 EDGE MESSAGE PASSING LAYER: GearNet-Edge adds sparse edge message passing to model dependencies between spatial or sequential interactions involving neighboring residues.The layer is motivated by prior geometric encoders that explicitly model interactions between edges.
  • 3.2 EDGE MESSAGE PASSING LAYER: The edge graph links two original edges when they share an endpoint without being identical, assigning relation types from the angle between them.Angles are discretized into 8 bins to reduce the memory cost of computing many kernel matrices.
  • 3.2 EDGE MESSAGE PASSING LAYER: Relational message passing models different spatial interactions among residues and is presented as the first edge-message-passing approach for macromolecular representation learning.Compared with AlphaFold2 triangle attention, the method uses angular information for sparse edge interactions.

4 GEOMETRIC PRETRAINING METHODS

This section presents geometric pretraining that learns protein representations from structurally derived views and self-prediction objectives. Multiview contrastive learning aligns biologically correlated substructures, while four prediction tasks target physicochemical and geometric properties.

  • Multiview contrastive learning: The contrastive framework preserves similarity between biologically related substructures before and after mapping them into a low-dimensional latent space.The motivation is that structural motifs and substructures reflect protein evolution history and functions.
  • Multiview contrastive learning: Subsequence cropping samples contiguous residues to capture recurring protein domains, whereas subspace cropping uses three-dimensional structural information.The two schemes are designed to extract biologically meaningful protein substructures.
  • Multiview contrastive learning: Multiview contrastive learning constructs two augmented views of each protein residue graph and brings their latent representations closer than representations of negative samples.Views use subsequence or subspace cropping together with identity transformation or random edge masking.
  • Multiview contrastive learning: Sampling size affects EC performance non-monotonically: results first rise and then drop as cropped substructures become larger.The authors examine subsequence length and spatial cutoff while using identity transformations.
  • Results: Pretrained representations show clear separation by protein familial classification in a visualization based on domain annotations.The authors report this separation as evidence of the effectiveness of the pretraining method.
  • Self-prediction methods: Four self-prediction tasks predict residue types, distances, angles, and dihedrals from masked protein structures.They apply masked prediction to single residues, residue pairs, triplets, and quadruples; angular values are discretized because they are more sensitive to structural errors.

5 EXPERIMENTS

Experiments evaluate GearNet variants and pretraining methods across four downstream protein tasks. Structure-based encoders and unlabeled-structure pretraining generally outperform relevant baselines, with edge message passing and multiview contrast providing especially strong results.

  • Experimental setup: Experiments cover EC number prediction, GO term prediction, fold classification, and reaction classification using AlphaFold2-derived structures.Pretraining uses 365K proteome-wide predictions and 440K Swiss-Prot predictions; evaluation reports Fmax for EC and GO and mean accuracy for fold and reaction classification.
  • Results: Structure-based encoders outperform all baselines without pretraining on 7 of 8 datasets.GearNet-Edge leads on several function-prediction tasks, while GearNet-Edge-IEConv achieves the best fold-classification results.
  • Results: Pretraining with unlabeled structures substantially improves structure-based encoders, with Multiview Contrast best on 7 of 8 datasets.Multiview Contrast reaches state-of-the-art results on EC, GO-BP, GO-MF, Fold, and Reaction tasks.
  • Results: Pretrained structure-based encoders match or exceed sequence-based encoders despite using fewer than one million structures.The structure-based models also perform well across all four evaluated task types, whereas sequence-based models perform poorly on fold classification.
  • Ablation studies: Relational graph convolution significantly improves results over graph convolution with a shared kernel across edge types.The ablation supports treating different edge types distinctly in the encoder.
  • Ablation studies: Edge message passing consistently improves performance over GearNet without edge message passing.The mechanism explicitly models interactions between edges, complementing distance and angle features.

6 CONCLUSIONS AND FUTURE WORK

The paper presents GearNet, a structure-based protein encoder with relational message passing and explicit edge interactions, together with five self-supervised pretraining methods. Across multiple benchmarks, it outperforms prior encoders from scratch and matches or exceeds sequence-based methods while using much less pretraining data.

  • Contributions: GearNet performs relational message passing on protein residue graphs to learn structure-based protein representations.Its sparse edge message passing mechanism explicitly models interactions between edges.
  • Contributions: Five self-supervised pretraining methods use contrastive learning and self-prediction frameworks for protein structures.The methods target representation learning from unlabeled 3D structures.
  • Results: Across multiple benchmark tasks, GearNet outperforms previous encoders trained from scratch and matches or exceeds state-of-the-art sequence-based baselines with much less data.The conclusion frames structure-based self-supervised learning as a step toward protein structure understanding.

A.1 SEQUENCE-BASED METHODS FOR PROTEIN REPRESENTATION LEARNING

Protein representation learning has expanded from sequence models to structural and graph-based approaches, but the supplied discussion emphasizes both their capabilities and limitations. Structure-aware methods represent proteins through residue-level, atom-level, or surface graphs, while self-supervised graph pretraining offers several possible objectives.

  • Sequence-based protein models learn biochemical and co-evolutionary information from large corpora using self-supervised objectives before downstream fine-tuning.
  • Structure-based approaches encode proteins through residue-level or atom-level graphs and protein surfaces for tasks including design, classification, quality assessment, and function prediction.
  • Graph neural network pretraining uses self-supervised tasks such as edge prediction, context prediction, attribute reconstruction, and contrastive learning.
  • Structural pretraining can leverage unlabeled protein structures, but the paper notes that only limited prior work had explored it and that available experimentally determined structures were scarce.

C.1 PROTEIN GRAPH CONSTRUCTION

The protein graph is constructed as a relational residue graph combining sequential and spatial connectivity, with optional edge-aware message passing to increase model capacity. The design uses lightweight residue features and filters spatial edges to emphasize long-range interactions.

  • Graph construction: The graph construction adds sequential, radius, and KNN edges, while separate edge message passing can improve model capacity.
  • Graph construction: Radius and KNN edges are filtered when residues are too close in sequence, excluding local spatial connections below the long-range cutoff.
  • Graph construction: The implementation uses dseq = 3, dradius = 10.0, k = 10, and dlong = 5, yielding seven edge types.
  • Why spatial edges: KNN-only graphs have nearly constant node degrees, whereas radius-only graphs can be highly sparse, motivating the combination of both spatial edge types.
  • Features: Residue nodes use one-hot residue types, while edge features concatenate endpoint features, edge type, sequential distance, and spatial distance.
  • IEConv enhancement: The IEConv enhancement maps edge features to kernel matrices for neighbor aggregation and improves fold-label prediction despite weaker function-prediction performance.

D SELF-PREDICTION METHODS

The paper introduces self-prediction objectives that reconstruct masked protein properties from the remaining structure. These objectives cover residue identity, distances, angles, and dihedrals, while contrastive learning links biologically correlated substructures.

  • Residue Type Prediction masks residue node features and predicts the missing residue types using a structure-based encoder.
  • Distance Prediction masks graph edges and predicts Euclidean distances between the representations of their endpoint residues.
  • Angle and Dihedral Prediction mask adjacent edge pairs or triplets and classify their geometric values into eight bins over [0, π].

E.1 DATASET STATISTICS

The evaluation covers function, fold, and reaction prediction using datasets with varied sizes and label structures, alongside sequence- and structure-based baselines. Function tasks use protein-centric and pair-centric metrics, while model outputs are passed through a three-layer MLP.

  • Datasets: The downstream evaluation includes EC prediction, GO prediction, fold classification, and reaction classification.
  • Datasets: Fold Classification contains 16,712 proteins assigned to 1,195 SCOPe 1.75 folds.
  • Datasets: Reaction Classification contains 37,428 proteins categorized into 384 reaction classes, with sequence-similarity clusters kept within individual splits.
  • Function metrics: EC and GO prediction are formulated as multiple binary classification tasks asking whether proteins have particular functions.
  • Function metrics: Fmax averages per-protein precision and recall across thresholds and selects the maximum F-measure.
  • Function metrics: AUPRpair is the average precision over all protein-function pairs, equivalent to micro-average precision for multiple binary classification.
  • Implementation: For all models, predictions are produced by a three-layer MLP whose hidden-layer dimension matches the model-output dimension.
  • Baselines: The benchmark compares sequence encoders, structure encoders, pretrained protein language models, and models combining sequence and structural information.

F ADDITIONAL EXPERIMENTAL RESULTS ON EC AND GO PREDICTION

Additional experiments test robustness across sequence-identity cutoffs, pretraining datasets, backbones, sampling schemes, and evaluation metrics. The results generally support robust geometric pretraining, while revealing metric and sampling sensitivities.

  • Robustness to sequence-identity cutoffs: At lower sequence-identity cutoffs, the model remains strongest among non-pretrained models and achieves comparable or better results than pretrained ESM-1b.Lower cutoffs create test sets with less similarity to training proteins.
  • Alternative evaluation metric: AUPR results show the model achieves the best EC-prediction performance in both from-scratch and pretrained settings, despite gaps from state-of-the-art results.The authors attribute the gaps partly to inconsistency between Fmax and AUPR.
  • Self-prediction sampling: Randomly sampling four nodes for dihedral prediction lowers Fmax from 0.859 to 0.821 compared with sampling three consecutive edges.The result supports learning residue representations that capture local spatial structure.
  • Pretraining datasets: Pretraining on AlphaFold Database v1, AlphaFold Database v2, or experimentally determined PDB structures yields comparable EC-prediction performance.The PDB dataset contains 123,505 structures and 305,265 chains selected at resolutions between 0.0 and 2.5 angstroms.
  • Backbone generality: All five pretraining methods substantially improve an EGNN backbone on the EC dataset, with Distance Prediction and Multiview Contrast among the strongest methods.This experiment tests whether the pretraining benefits extend beyond the GearNet backbone.
  • Sequence–structure comparison: Combining sequence representations from ESM-1b with a structure-based encoder provides an additional comparison between sequence and structural information.The combined encoder is evaluated on EC prediction under the same training configurations.

I COMBINE NEURAL AND RETRIEVAL-BASED METHODS

The paper compares neural structure representations with retrieval-based structure search and evaluates their combination. Retrieval is strong for proteins with similar structures in training, while ensembling improves several prediction metrics.

  • Scope of ensembling: Foldseek ensembles are omitted for fold classification because Foldseek performs poorly in that setting.The comparison covers EC, GO, and fold-classification tasks.
  • Neural and retrieval comparison: Foldseek is a strong baseline on EC and GO prediction, outperforming GearNet-Edge without pretraining and competing with pretrained GearNet-Edge on several tasks.Its advantage is especially associated with proteins having similar structures in the training set.
  • Ensembling: Ensembling neural and Foldseek predictions significantly improves both metrics on EC, GO-BP, and GO-MF compared with either method separately.The paper identifies combining retrieval and neural methods as a direction for future work.
  • Structure-based search: On SCOPe40 structure search, GearNet-Edge-IEConv representations achieve the best average sensitivity among the evaluated baselines.Similarity is measured by cosine distance between learned protein representations.

J LATENT SPACE VISUALIZATION

UMAP visualization of pretrained GearNet-Edge embeddings shows that proteins from the same superfamily tend to cluster together, with several superfamilies clearly separated.

  • Visualization setup: UMAP maps embeddings for all AlphaFold Database v1 proteins into two dimensions and colors the 20 most common superfamilies differently.The embeddings come from GearNet-Edge pretrained with Multiview Contrast.
  • Latent-space organization: Proteins from the same superfamily tend to group together while proteins from different superfamilies are separated.The visualization qualitatively evaluates the organization of pretrained protein embeddings.
  • Latent-space organization: The visualization clearly separates the Protein kinase, Cytochrome P450, and TRAFAC class myosin-kinesin ATPase superfamilies.These examples provide the paper’s strongest qualitative evidence of superfamily discrimination.
  • Residue-level interpretation: Integrated Gradients is applied to obtain residue-level interpretations for ATP-binding and Heme-binding predictions.Residues are colored by attribution scores, with higher scores shown in red and lower scores in blue around the ligand-binding complexes.
Loading 2203.06125v5…