Orientation-Aware Graph Neural Networks for Protein Structure Representation Learning

TL;DR

The paper addresses the need to capture fine-grained orientational information in protein structures for downstream learning. It introduces Orientation-Aware Graph Neural Networks (OAGNNs) built on Directed Weight Perceptrons that extend scalar features to 3D vectors and define geometry-aware operators, enabling $SO(3)$-equivariant processing. A global equivariant message passing framework integrates these units on protein graphs, with OA-GCN, OA-GIN, and OA-GAT variants to suit tasks. Across synthetic angular tasks and structural biology benchmarks (RES, CPD, MQA), OAGNNs demonstrate improved geometric sensing and competitive state-of-the-art performance, highlighting their potential for protein design and function prediction. The approach provides a flexible, geometry-conscious toolkit for 3D structure representation learning with practical bioinformatics impact.

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.

Figures (3)

• Figure 1: Overview (a) Each amino acid has its own rigid backbone with four heavy atoms, representing a local frame. (b) Tasks associated with the protein 3D structure. Graph-Level tasks consider the whole protein structures, and Node-Level tasks operate on specific residues.
• Figure 2: Model Details. A 1-layer DWP consists of three following modules. (a) Directed Linear Module applies multiple geometric operations to update scalar and vector features in four different ways with directed weights. (b) Non-Linearity Module employs ReLU and sigmoid functions for scalar and vector features. (c) Directed Interaction Module updates scalar and vector features by using one another as updating parameters after non-linearity.
• Figure 3: a). The synthetic study. The vector in red is the anchor vector. Other vectors falling into the cone are positive, whereas outside points are negative. b). The validation Binary Cross Entropy loss of compared methods.