PROTEIN STRUCTURE REPRESENTATION LEARNING THROUGH ORIENTATION-AWARE GRAPH NEURAL NETWORKS

Abstract

By folding to 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-theart performance on various computational biology applications related to protein 3D structures.

1. INTRODUCTION

Built from a sequence of amino-acid residues, a protein performs its biological functions by folding to a particular conformation in 3D space. Therefore, untilizing such 3D structures accurately is the key for downstream analysis. While we have witnessed remarkable progress in protein structure predictions (Rohl et al., 2004; Källberg et al., 2012; Baek et al., 2021; Jumper et al., 2021) , another thread of tasks with protein 3D structures as input starts to draw a great interest, such as function prediction (Hermosilla et al., 2020; Gligorijević et al., 2021 ), decoy ranking (Lundström et al., 2001; Kwon et al., 2021; Wang et al., 2021 ), protein docking (Duhovny et al., 2002; Shulman-Peleg et al., 2004; Gainza et al., 2020; Sverrisson et al., 2021) , and driver mutation identification (Lefèvre et al., 1997; Antikainen & Martin, 2005; Li et al., 2020; Jankauskaitė et al., 2019) . Most existing works in modeling protein structures directly borrow models designed for other applications, including 3D-CNNs (Ji et al., 2012 ) in computer vision, Transformers (Vaswani et al., 2017) ) from natural language processing, and GNNs (Kipf & Welling, 2016) in data mining. Though compatible with general objects, these models have overlooked the subtleties in the fine-grained geometries, which are much more essential in protein structures. For instance, given an amino acid in the protein structure, as shown in Figure 1 , the locations of four backbone atoms (carbon, nitrogen, and oxygen) determine a local skeleton, and different residues interact with each other through performing specific orientations between their local frames, either of which have important impacts on the protein structure and its function (Nelson et al., 2008) . Recent attempts in building geometric-aware neural networks mainly focus on baking 3D rigid transformations into network operations, leading to the area of SO(3)-invariant and equivariant networks. One representative work is the Vector Neuron Network (VNN) (Deng et al., 2021) , which achieves SO(3)-equivariance on point clouds by generalizing scalar neurons to 3D vectors. Another work is the GVP-GNN (Jing et al., 2021) that similarly vectorizes hidden neurons in GNN and demonstrates better prediction accuracy on protein design and quality evaluation tasks. However, these two models can only adopt linear combinations of input vectors, which significantly limits their modeling capability. A simple example is that, given two input vector features v 1 and v 2 , the outputs w 1 v 1 + w 2 v 2 through one linear layer is constrained in the 2D plane spanned by v 1 , v 2 even after applying their scalar-product non-linearities. That is, VNN-based models are limited in perceiving orientational features, which have been proven crucial for proteins to perform their functions and interact with other partners (e.g. inner-residue torsion angles, inter-residue orientations) (Nelson et al., 2008; Voet & Voet, 2010; Xu & Berger, 2006; Alford et al., 2017) . To achieve more sensitive geometric orientation awareness, we propose a Directed Weight ( ⃗ W) perceptrons by extending not only the hidden neurons but also the weights from scalars to 3D vectors, naturally saturating the entire network with 3D structure information in the Euclidean space. Directed weights support a set of geometric-meaningful operations on both the vector neurons (vectorlist features) and the classical (scalar-list) latent features, and perform flexible non-linear integration of the hybrid scalar-vector features. As protein structures are naturally attributed proximity graphs, we introduce a new Equivariant Message Passing Paradigm on protein graphs, to connect the ⃗ Wperceptrons with the graph learning models by using rigid backbone transformations for each amino acid, which provides a versatile framework for bringing the biological suitability and flexibility of the GNN architecture. To summarize, our key contributions include: • We propose a new network unit based on the Directed Weights for capturing fine-grained geometric relations, especially for the subtle orientational details in proteins. • We construct an Equivariant Message Passing paradigm based on protein graphs. • Our overall framework, the Orientation-Aware Graph Neural Networks, is versatile in terms of compatibility with existing deep graph learning models, making them biologically suitable with minimal modifications to existing GNN models.

2. RELATED WORK

Representation learning on protein 3D structure. Early approaches rely on hand-crafted features extracted and statistical methods to predict function annotations (Schaap et al., 2001; Zhang & Zhang, 2010) . Deep learning has been found to achieve success then. 3D CNNs are first proposed to process protein 3D structures by scanning atom-level features relying on multiple 3D voxels. One of the representative works (Derevyanko et al., 2018) adopts a 3D CNN-based model for assessing the quality of the predicted structures. 3D CNNs also shed light on other tasks such as interface prediction (Townshend et al., 2019; Amidi et al., 2018) . People also extend them to spherical convolutions (Gainza et al., 2020; Sverrisson et al., 2021; Hermosilla Casajus et al., 2021) , to the Fourier space (Zhemchuzhnikov et al., 2022) and the 3D Voronoi Tessellation space (Igashov et al., 2021) . Graph Convolutional Networks (Kipf & Welling, 2016) have also been adopted to capture geometric and biochemical interactions between residues (Ying et al., 2018; Gao & Ji, 2019; Fout, 2017) , and have been shown to achieve great performance on function prediction (Li et al., 2021 ), protein design (Strokach et al., 2020) and binding prediction (Vecchio et al., 2021) . Recently, transformer-based



Figure 1: Overview (a) Each amino acid has its own rigid backbone with four heavy atoms, and can be represented by both lists of scalar and vector features. (b) Tasks associated with the protein 3D structure. Graph-level tasks consider the whole protein structures, and Node-level tasks operate on specific residues.

