[Paper Review] Contrastive Representation Learning for 3D Protein Structures

Pedro Hermosilla, Timo Ropinski|arXiv (Cornell University)|May 31, 2022

Protein Structure and Dynamics25 citations

TL;DR

The paper introduces a domain-specific contrastive learning framework to pre-train representations of 3D protein structures, enabling improved performance on protein similarity, fold classification, function prediction, and protein-ligand binding affinity tasks. Pre-trained models often outperform from-scratch baselines and achieve state-of-the-art or competitive results across multiple tasks.

ABSTRACT

Learning from 3D protein structures has gained wide interest in protein modeling and structural bioinformatics. Unfortunately, the number of available structures is orders of magnitude lower than the training data sizes commonly used in computer vision and machine learning. Moreover, this number is reduced even further, when only annotated protein structures can be considered, making the training of existing models difficult and prone to over-fitting. To address this challenge, we introduce a new representation learning framework for 3D protein structures. Our framework uses unsupervised contrastive learning to learn meaningful representations of protein structures, making use of proteins from the Protein Data Bank. We show, how these representations can be used to solve a large variety of tasks, such as protein function prediction, protein fold classification, structural similarity prediction, and protein-ligand binding affinity prediction. Moreover, we show how fine-tuned networks, pre-trained with our algorithm, lead to significantly improved task performance, achieving new state-of-the-art results in many tasks.

Motivation & Objective

• Motivate learning from limited labeled 3D protein structure data by leveraging unlabeled structures from the Protein Data Bank (PDB).
• Develop a contrastive learning framework that produces invariant representations of random sub-structures from the same protein.
• Design a protein-specific graph encoder that captures spatial and orientation information for 3D structures.
• Demonstrate that pre-trained representations improve downstream tasks such as structural similarity, fold classification, function prediction, and binding affinity.

Proposed method

• Represent proteins as graphs where nodes are amino acid Cα positions with local orientation frames and multiple features.
• Use domain-specific sub-structure sampling to create positive pairs from the same protein and negative samples from other proteins.
• Apply a two-branch encoder (graph encoder E and MLP projector P) to obtain latent representations h and z, and optimize cosine-based contrastive loss with temperature τ.
• Compute edge features (relative position in local frame, frame orientation, shortest-path distance) to build a spatially aware graph convolution operation.
• Implement a hierarchical protein encoder with ResNet-like blocks and graph pooling that propagates information while preserving translation invariance and rotation equivariance.
• Adopt a smooth receptive field with distance-based gating to manage neighborhood influence and avoid discontinuities at the cutoff distance d.

Experimental results

Research questions

• RQ1 Can unsupervised contrastive learning on unlabeled 3D protein structures produce useful representations for diverse downstream tasks?
• RQ2 Do domain-aware sub-structure sampling and a protein-specific graph encoder improve representation quality compared to generic graph methods?
• RQ3 To what extent do pre-trained representations transfer to tasks like protein structural similarity, fold classification, function prediction, and binding affinity prediction?

Key findings

• Pre-trained representations improve performance on protein structural similarity tasks across multiple datasets and metrics.
• Fine-tuning the pre-trained protein encoder generally yields the best results on several tasks, sometimes outperforming from-scratch models, while in some similarity tasks fixed representations with an MLP also show strong gains.
• The learned representations support state-of-the-art or competitive results in fold classification, GO-term prediction, enzyme reaction classification, and binding affinity prediction across varying sequence similarity conditions.
• The proposed framework enables faster similarity computations (dot product in latent space) compared to traditional methods, with substantially reduced runtime.