[Paper Review] Orientation-Aware Graph Neural Networks for Protein Structure Representation Learning

Jiahan Li, Shitong Luo+6 more|arXiv (Cornell University)|Jan 28, 2022

Machine Learning in Bioinformatics4 citations

TL;DR

This paper introduces Directed Weight (\overrightarrow{{\mathbf{W}}}) Graph Neural Networks that extend scalar weights to 3D directed vectors to better capture fine-grained geometric relations in protein 3D structures, such as backbone orientations and side-chain torsion angles. By integrating \overrightarrow{{\mathbf{W}}}-perceptrons into a novel SO(3)-equivariant message passing framework on protein graphs, the model achieves state-of-the-art performance on multiple protein structure prediction tasks while maintaining geometric expressivity and biological plausibility.

ABSTRACT

By folding into particular 3D structures, proteins play a key role in living beings. To learn meaningful representation from a protein structure for downstream tasks, not only the global backbone topology but the local fine-grained orientational relations between amino acids should also be considered. In this work, we propose the Orientation-Aware Graph Neural Networks (OAGNNs) to better sense the geometric characteristics in protein structure (e.g. inner-residue torsion angles, inter-residue orientations). Extending a single weight from a scalar to a 3D vector, we construct a rich set of geometric-meaningful operations to process both the classical and SO(3) representations of a given structure. To plug our designed perceptron unit into existing Graph Neural Networks, we further introduce an equivariant message passing paradigm, showing superior versatility in maintaining SO(3)-equivariance at the global scale. Experiments have shown that our OAGNNs have a remarkable ability to sense geometric orientational features compared to classical networks. OAGNNs have also achieved state-of-the-art performance on various computational biology applications related to protein 3D structures. The code is available at https://github.com/Ced3-han/OAGNN/tree/main.

Motivation & Objective

• To address the limitation of classical and existing equivariant GNNs in capturing subtle geometric features like amino acid orientations and torsion angles in protein 3D structures.
• To develop a geometrically expressive neural network framework that supports rich geometric operations on both scalar and vector features.
• To design an SO(3)-equivariant message passing paradigm on protein graphs that preserves global geometric invariance without constraining model architecture.
• To enable seamless integration of the proposed framework into existing GNN architectures with minimal modifications.
• To demonstrate superior expressivity and performance on protein structure representation tasks compared to classical and equivariant baselines.

Proposed method

• Proposes \overrightarrow{{\mathbf{W}}}-perceptrons that generalize scalar weights to 3D directed vectors, enabling geometric operations on both scalar and vector features.
• Introduces a new message passing paradigm that uses local backbone orientations ({\bm{O}}_{u}) to transform features into local coordinate systems for equivariant processing.
• Employs multi-layer \overrightarrow{{\mathbf{W}}}-perceptrons (\mathcal{F} and \mathcal{H}) to process aggregated node and edge features, with transformations back to global coordinates.
• Applies the framework to multiple GNN variants: DW-GCN, DW-GIN, and DW-GAT, with attention mechanisms adapted for scalar and vector representations.
• Uses Frobenius and standard inner products for vector and scalar attention scores, respectively, enabling selective neighbor aggregation.
• Enforces SO(3)-equivariance globally by leveraging the rigid-body transformation of each amino acid residue, proven formally in the appendix.

Experimental results

Research questions

• RQ1 Can extending scalar weights to 3D directed vectors significantly improve the modeling of fine-grained geometric features in protein 3D structures?
• RQ2 How can SO(3)-equivariance be effectively maintained in a general GNN framework without imposing architectural constraints?
• RQ3 To what extent do \overrightarrow{{\mathbf{W}}}-perceptrons outperform classical and existing equivariant GNNs in capturing protein-specific geometric relations like side-chain torsion angles?
• RQ4 Can the proposed framework be flexibly integrated into diverse GNN architectures while preserving geometric expressivity and performance?
• RQ5 Does the equivariant message passing paradigm enhance performance across multiple protein structure learning tasks?

Key findings

• The proposed \overrightarrow{{\mathbf{W}}}-GNN framework demonstrates superior expressivity in modeling geometric relations compared to classical scalar-weighted networks and existing SO(3)-equivariant models.
• On a synthetic dataset, the model successfully captures angular relationships—such as side-chain torsion angles—that are poorly modeled by VNN-based approaches.
• The framework achieves state-of-the-art performance on multiple protein structure learning tasks, including function prediction, decoy ranking, protein docking, and driver mutation identification.
• The DW-GAT variant, which incorporates learnable attention for scalar and vector features, enables selective neighbor aggregation, improving representation learning on heterogeneous protein interactions.
• The SO(3)-equivariant message passing paradigm ensures global geometric consistency without restricting the design of \mathcal{F} and \mathcal{H}, enabling flexible and powerful architectures.
• All models and code will be published upon acceptance, supporting reproducibility and future extension to other 3D learning tasks.