PersGNN: Applying Topological Data Analysis and Geometric Deep Learning to Structure-Based Protein Function Prediction

Graph Neural Networks (GNN), which can be motivated through spectral graph convolutions, are a popular form of graph representation learning that provide a framework for extending traditional deep learning techniques to non-Euclidean, graphical data [7, 8]. GNNs learn context-aware node embeddings through successive rounds of local neighborhood aggregation (or "message passing") in a graph. These node embeddings can be combined to create global representations for entire graphs. Previous work has used 3D Convolutional Neural Networks (3D CNNs) to extract useful information from protein structure [9, 10]. However, this technique is both memory and computationally inefficient because proteins are sparse in 3D space and 3D CNNs will perform many convolutions over empty space. More recent work proposed to overcome some of these inefficiencies with a GNN [5]. In this method, 3D protein structures are represented by contact maps, which threshold the pairwise distances between the alpha-carbons of each residue in the protein. These contact maps are then fed into the GNN, which embeds structural information that can used to make function predictions.

Our GNN model architecture is inspired by the one of Gligorijevic et al., which uses the graph convolutional layers from Kipf and Welling, along with language modeling [5, 7]. In contrast, our work focuses on learning from structure. Thus, we omit the language modeling component and instead opt for a simple one-hot representation of amino acids. Other modeling differences, such as model depth and learning rate, were optimized during the validation stage.

We use persistent homology, an area of topological data analysis, to construct a topological representation of the protein structure. In this section, we briefly describe the approach and refer the interested reader to a more thorough survey [11].

We start with the 3D atomic coordinates of a protein structure and take the union of balls around the atoms. We vary the radii of these balls, and track how the topology of this union changes: these are represented as alpha shapes [12]. By doing this and sweeping across all radii, we get an increasing sequence of nested alpha shapes called the filtration. By following the changes in the sequence, one can keep track of the appearances and disappearances of topological features in the filtration.

We record the pairs of radii when such topological features appear and disappear, where the appearance of a new feature is the birth and the disappearance of this same feature (such as when it gets merged with another feature) is the death. The difference between the birth-death pairs, i.e. death - birth, is called the persistence of the feature. The larger the persistence, the more prominent the topological feature. The full set of all birth-death pairs is called a persistence diagram. In this work, we specifically look at the 1-dimensional and 2-dimensional persistence diagrams of each protein structure, where 1-dimensional features correspond to channels (“loops”) and 2-dimensional features correspond to voids (“cavities”) in the protein structure.

Feeding persistence diagrams, in this case the topological summaries of all the channels and voids in a protein structure, into a neural network architecture requires a transformation of the diagram points. An approach to do this was first proposed by Hofer et al. [13]. Expanding on this, Carrière et al. [14] proposed a layer for neural network architectures based on the DeepSet architecture [15] to encode a vector representation of the persistence diagram through point-wise transformations. We implement a modification of this approach to process both 1D and 2D persistence diagrams of the protein structure, which we call “PersNet” in this work.

The problem of describing protein shape was one of the early motivations behind persistent homology [16]. Persistent homology has also been used previously to construct topological representations to describe protein structures for protein classification problems [17, 18]. However, these approaches turn topological information into feature vectors and rely on the construction of ad hoc handcrafted summaries. More recently, vectorizing persistence diagrams in ways that require minimal processing and retain more of the information, such as to persistence images [19], have been used in machine learning applications for scientific domains [20, 21, 22, 23]. However, these vectorized representations are static — in contrast, the approach outlined here processes the topological information through an automatically differentiable layer, thereby continuing to learn the representation over the training process.

Using a GNN in place of a 3D CNN significantly reduces computation time and memory requirements of learning on protein structures, but also requires representing the protein structure with a contact map. This simplification is necessary for the GNN to work but can also lead to a loss of key structural information. To overcome this, we propose incorporating persistent homology into the learning process by converting the 3D atomic coordinates of the protein structure to persistence diagrams and processing this information through a layer in the neural network architecture, as described in 2.2. We note that previous work has used persistent homology together with 3D CNNs to classify protein folds [23]; however, as mentioned earlier, they used a static vectorization of the persistence diagrams of the protein structure rather than a neural network processing layer that learns the representation through training and did not use GNNs — both things that we found very important in our approach to achieving high accuracy. There has also been previous work combining persistent homology and GNNs [24], but here they incorporated persistent homology into the graph neural network architecture itself by using persistent homology to reweight messages passed between the graph nodes during convolutions. Our method allows us to build on the successes of applying GNNs to protein data (and capitalize on their computational efficiency) while also adding additional structural features that are potentially not captured by a protein’s contact map, such as more global topological information.

Our hybrid model, shown in Figure 1, concatenates the output from the GNN with the output from the PersNet models that process both the 1D and 2D topological features (channels and voids respectively). This hybrid representation is then passed into a Multi-Layer Perceptron (MLP) to perform the classification task. The architecture of the GNN and PersNet models is the same as the ones described in Sections 2.1 and  2.2. PersGNN is trained end-to-end on a protein’s contact maps, amino acid labels, and persistence diagrams and all sub-components are trained simultaneously. This enables PersGNN to learn a complex representation of protein structure efficiently and with minimal manual processing.

In addition to benchmarking our hybrid model against its component models, we trained a simple three-layer MLP to create a more complete picture of how these models are learning from structure. This is distinct from many assessments of protein functional annotations, such as CAFA, which use the alignment-based BLAST as a baseline. Since this investigation is focused on extracting information from protein structure, rather than its amino acid sequence, we decided not to use the sequence-only BLAST (or any other sequence-based methods) as a baseline. The baseline model is given all the same information as the GNN: for each residue, the model is given the amino acid label and the number of residue “contacts,” which are computed from the same contact maps that are fed into the GNN.

Our dataset consists of 33,007 proteins taken from PDB that belong to the Arabidopsis, Celegans, E.coli, Fly, Human, Mouse, Yeast and Zebrafish species. This is randomly split into train, validation and test sets at a $75\%$, $12.5\%$, $12.5\%$ ratio. Random splitting can sometimes be overly optimistic and not representative of realistic challenges [25]. Previous work, for example, has split their datasets to minimize $\%$ sequence identity between training and tests set, which we leave for future work.

For labels, we represent protein function with molecular function (MF) gene ontology (GO) terms. For each term in PDB, we obtains its corresponding MF GO terms using the Structure Integration with Function, Taxonomy and Sequence (SIFTS) database [26, 27]. Consistent with previous studies, we filter out GO terms with less than 25 representative proteins. After the filtration, our label space is made up of 730 MF GO terms. Each protein can also have multiple labels, making this a multi-label multi-classification problem.

Each protein entry in PDB is described by a list of amino acid residues and their 3D atomic coordinates. For the GNN and MLP models, the 3D atomic coordinates are converted to $C_{\alpha}-C_{\alpha}$ contact maps. To construct these contact maps, we compute pairwise distances between the $\alpha$-carbon of each residue and denote a "contact" between residues if their $\alpha$-carbons are within eight angstroms. The GNN and MLP models also take a one-hot encoded representation of amino acids, which the GNN uses as node labels. We also construct 1D and 2D persistence diagrams from the 3D atomic coordinates of each protein structure, which are then processed through the persistence network layers.

All models are trained using a weighted binary cross-entropy loss function [5], the ADAM optimizer [28], for 300 epochs and with a learning rate of $10^{-5}$. To normalize our models, we employ weight normalization and dropout regularization (p = 0.1) throughout [29, 30]. For each method, we create an ensemble by independently training 10 models and taking the average bit-score (per GO term) across the 10 models at evaluation time.