ProteinRPN: Towards Accurate Protein Function Prediction with Graph-Based Region Proposals

Protein sequences are represented as graphs $G(V,E)$, where the vertices $V$ correspond to the protein residues, and the edges $E$ represent the proximity of residues in three-dimensional space. The adjacency matrix $A\in\mathbb{R}^{N\times N}$ for an $N$-residue protein graph is defined by calculating the contact map, where an edge is added between two nodes if the distance between their $C_{\alpha}$ atoms is less than 10 Å. In this work, we use $G(V,E)$ and $G(Z,A)$ interchangeably, where $V$ and $E$ denote the set of vertices and edge list, while $Z\in\mathbb{R}^{|V|\times D}$ and $A\in\mathbb{R}^{|V|\times|V|}$ represent the node features and adjacency matrices, respectively, and $D$ is the chosen dimension for residue features. The goal of ProteinRPN is to predict a probability vector $\hat{\mathbf{y}}_{i}^{(j)}\in\mathbb{R}^{l_{j}}$, where $l_{j}$ denotes the number of GO terms associated with subontology $j\in\{\text{BP},\text{CC},\text{MF}\}$. The vector $\hat{\mathbf{y}}_{i}^{(j)}$ represents the predicted probabilities for the $l_{j}$ GO terms, reflecting the likelihood of each protein being associated with multiple GO terms across all subontologies.

Inspired by object detection models in computer vision, we propose a strategy analogous to region proposal networks in Faster R-CNN (Ren et al. 2015), adapted for protein function prediction in graphs. By targeting regions containing functional residues, which are often a small subset of the protein, this approach improves functional understanding. To the best of our knowledge, this is the first work to introduce graph region proposals, applied specifically to protein function prediction.

The proposed Region Proposal Network employs $k$ layers of Graph Convolutional Networks (GCNs) (Kipf and Welling 2017) to process protein graphs $G(Z,A)$. In particular, let $H^{(0)}=Z$ represent the initial hidden node embedding matrix. The hidden embeddings $H$ are updated iteratively as:

The second step in the region proposal module involves localizing regions that are likely to contain functional residues. This is formulated as a node classification task, where the goal is to predict whether the anchor centred around each node contains a functionally relevant region. More precisely, given the node embeddings $Z_{1}$ after $k$ GCN layers, the classification of each node $v_{i}$ in the transformed graph $G^{\prime}(Z_{1},A)$ is performed using a Graph Attention Network (GAT) convolution (Veličković et al. 2018). The output for each node $v_{i}$ can be formulated as:

Nodes predicted as functional are selected, and $k$-hop subgraphs (anchors) centered around these nodes are extracted. This results in a collection of anchors enriched for functional regions, ensuring high recall but with room for precision improvement. To address this, we introduce a pruning step that selectively retains the most functionally relevant subgraphs within the larger anchors. This pruning leverages a novel node-drop pooling layer that incorporates domain knowledge alongside a hierarchy-aware attention mechanism. Rather than relying on conventional dot product attention, we utilize penumbral cone attention (Tseng et al. 2023) for modeling the inherent hierarchical relationships in proteins. These hierarchies span multiple levels, from the arrangement of secondary and tertiary structures to the organization of functional domains and motifs, all the way up to the interactions of subunits within protein complexes. By capturing these complex dependencies, penumbral cone attention enables a more precise focus on critical functional regions within the protein graph, refining our predictions and improving precision without compromising recall.

Node drop pooling layers are commonly used to reduce graph size by selectively removing lower-scoring nodes while retaining higher-scoring ones based on their importance or features. Instead of removing low-scoring nodes, we evaluate the functionality of each node in the Node Drop Pooling layer. Subsequently, we leverage the Functional Attention Layer to enrich representations of high-scoring nodes, ensuring that lower-scoring nodes contributing valuable contextual information are retained to preserve the overall graph structure.

Once candidate functional residues are identified, their representations are refined through a functional attention layer. This layer assigns weights to edges based on their connectivity to predicted functional nodes, allowing the model to emphasize relationships critical to protein function. By incorporating multistage refinement, we iteratively enhance the accuracy of functional node identification. The edge-centric approach helps preserve the structural integrity of selected subgraphs, avoiding the fragmentation that can occur when individual nodes are selected in isolation, in line with insights from PDBSite.

We feed the original residue features $Z\in\mathbb{R}^{N\times D}$ as the node feature matrix for enrichment. For each edge $(i,j)$ in the graph, the model computes an attention score ${e}_{ij}$ using the concatenation of the feature vectors ${Z}_{i}$ and $Z_{j}$, followed by a learnable weight vector ${a}\in\mathbb{R}^{2D_{1}\times 1}$ and a ReLU activation function, that would reduce all negative scores to zero, i.e., ${e}_{ij}=\text{ReLU}\left({a}^{\top}[{Z}_{i}\,\|\,{Z}_{j}]\right)$. This attention score is then adjusted based on the node type ${z}_{j}\in\{0,1\}$ of the target node $j$, modifying the score as follows:

where $\alpha_{FA}\geq 1$ and $\beta_{FA}<1$. This adjustment increases the attention for functional nodes while reducing it for contextual ones. For the purpose if this study, we use $\alpha_{FA}=1,\beta_{FA}=0.5$ in order to explicitly ensure focus on functional nodes. The attention coefficients $\alpha_{ij}$ are obtained by normalizing $e_{ij}$ across all neighbors. They determine how much influence a neighboring node $i$ has on the target node $j$.

Finally, the updated feature vector for node $j$, ${Z}_{2j}$, is computed by aggregating the messages from its neighbors applying, weighted by the corresponding attention coefficients, i.e., ${Z_{2j}}=\sum_{i\in\mathcal{N}(j)}\alpha_{ij}\cdot W\cdot{Z}_{i}$, where the transformation matrix $W\in\mathbb{R}^{D_{2}\times D}$ is a learnable parameter as in the GAT layer and helps transform the initial extracted features of the nodes (residues) into a new space where relationships between residues can be more effectively captured. As a result of this operation, subgraphs surrounding functional residues—those likely to be critical for protein function—get more attention and influence the final node representations more significantly. This approach enhances the model’s ability to capture the rich, context-aware interactions between residues, leading to a more comprehensive understanding of the protein’s functional regions.

In the final step, the enriched representations $Z_{2}$ are fed into a Graph Multiset Transformer (GMT) layer, which transforms node-level embeddings into a comprehensive graph-level representation by capturing both local interactions and global structure. The GMT layer introduces learnable super-nodes to capture long-distance structural information and aggregates this information into a unified graph representation.

Our model is optimized using a multi-component loss function that integrates cross-entropy loss $\mathcal{L}_{\text{CE}}$ for multilabel classification, contrastive loss $\mathcal{L}_{\text{con}}$, and a penalty term $\mathcal{L}_{\text{penalty}}$ to minimize the number of disconnected components in the functional attention layer.

The contrastive loss, $\mathcal{L}_{\text{con}}$, is a combination of supervised contrastive (SupCon) loss (Khosla et al. 2021) and self-supervised noise contrastive estimation (InfoNCE) loss (van den Oord, Li, and Vinyals 2019). SupCon encourages the model to cluster representations of proteins with similar GO terms, while InfoNCE ensures that the representations are robust to noise by maximizing the similarity between original and perturbed embeddings. The combined contrastive loss for a batch of $B$ proteins is defined as:

where $z_{i}$ and $z_{j}$ are the embeddings of proteins $i$ and $j$, $\text{sim}(z_{i},z_{j})$ represents their cosine similarity, and $\tau$ is a temperature parameter. The indicator function $\mathbf{1}\{y_{i}\cap y_{j}\neq\emptyset\}$ ensures that only pairs with shared GO terms contribute to the SupCon loss, im order to adapt it to the multilabel case. The InfoNCE loss optimizes the similarity between the original and perturbed embeddings $z_{i}$ and $z_{i}^{\prime}$. The hyperparameters $\alpha_{\text{SupCon}}$ and $\alpha_{\text{NCE}}$ control the contributions of the SupCon and InfoNCE losses.

To ensure that identified functional regions are structurally cohesive, reflecting the biological reality of interconnected functional residues, we introduce a connected components penalty. This discourages the formation of disconnected components in the functional attention layer and is defined as:

The overall loss function is formulated as:

This comprehensive loss function guides the model toward producing biologically meaningful predictions by leveraging robust, context-aware graph representations while maintaining the structural coherence of the predicted functional subgraphs.

PDBSite (Ivanisenko, Grigorovich, and Kolchanov 2000) is a comprehensive dataset comprising biologically active sites derived from the Protein Data Bank (PDB) (Berman et al. 2000). The dataset encompasses 4,723 active sites belonging to 197 different functions located within 603 proteins. PDBSite stands out among annotation databases due to its diverse representation of functional categories, enabling broad analysis across various protein functions. We leverage PDBSite to guide our model architecture and pretrain the model on predicting functional sites.

For protein function prediction, we utilize a dataset curated by (Gu et al. 2023), originally developed to train their model, HEAL, which serves as our baseline. This dataset is an adapted version of the DeepFRI dataset (Gligorijević et al. 2021), comprising 36,629 sequences sourced from the PDB database (Berman et al. 2000) and 42,994 from the SWISS-MODEL repository (Bienert et al. 2016). Further details can be found in the Appendix.

We begin by training ProteinRPN on the PDBSite which is split into training and validation sets with an 80:20 ratio, with the goal of predicting all functional sites within a protein. The details of the pretraining can be found in the Appendix. Then we train the entire framework on the comprehensive protein function prediction task using the HEAL dataset.

We have conducted comprehensive experimennts to comapre ProteinRPN’s [erformance to SOTA models. Those methods encompass sequence-based models such as BLAST (Altschul et al. 1990) and FunFams (Das et al. 2015), sequence and PPI-based models like DeepGO (Kulmanov, Khan, and Hoehndorf 2017), and sequence and structure-based models such as DeepFRI (Gligorijević et al. 2021) and HEAL (Gu et al. 2023).

We conduct ablation studies to assess the significance of each model component, including the impact of secondary structure, coordinate information, and contrastive learning losses. Additional studies to test the efficacy of other modules can be found in the Appendix.

Model predictions are evaluated using the standard Critical Assessment of Functional Annotation (CAFA) evaluator (Jiang et al. 2016). Protein-centric Fmax, the maximum F1 score over all prediction thresholds ranging from 0 to 1 with a step size of 0.1, is utilized. Smin, representing the semantic distance between predicted and actual annotations, considers the information content of each function. The function-centric AUPR is employed as a robust measure for situations with high class imbalance. Further details on the formulas and implementation are available in (Jiang et al. 2016), and comprehensive information on model training and hyperparameters can be found in the Appendix.

Table 1 presents the performance metrics of ProteinRPN in comparison to all baseline models on the HEAL dataset. ProteinRPN consistently outperforms the baselines across all metrics, showing notable improvements over the HEAL model. Specifically, ProteinRPN achieves higher Fmax scores, with gains of 6.4% in Biological Process (BP), 2.7% in Cellular Component (CC), and 7.1% in Molecular Function (MF) ontologies. Beyond Fmax, ProteinRPN also demonstrates superior performance in Smin and Area Under the Precision-Recall Curve (AUPR), highlighting its effectiveness in predicting protein function GO terms.

Moreover, as shown in Table 2, the ablation study reveals that both contrastive learning and the incorporation of domain knowledge positively contribute to the model’s overall performance.

During pretraining, the region proposal module exhibits strong performance, achieving an ROC of 0.95 in the anchor functionality prediction task and 0.85 in the pruning task. Although direct comparison is limited due to the absence of established baselines, the module’s effectiveness is evident in downstream functional prediction tasks.

Overall, the results demonstrate that enabling the model to detect and focus on residues critical for function significantly enhances its performance. This improvement is primarily driven by the multistage refinement approach, which efficiently localizes functional residues within protein structures.

We evaluate the functional residue predictor in ProteinRPN by analyzing specific proteins. For example, on protein 2BCC (B chain, 422 residues, 10 functional), ProteinRPN accurately identifies 8 functional residues, with region proposals covering subgraphs of 28 residues, as shown in Fig. 2(a). Functional residues predicted correctly are highlighted in green, while missed ones are marked in red.

Similarly, for protein 2CHG (A chain, 226 residues, 11 functional), the model successfully identifies 9 functional residues, with region proposals covering 43 residues. As shown in Fig. 2(b), the correctly identified residues are closely clustered within the structure, while the missed residues are located farther from the cluster. These results demonstrate its ability to accurately identify and localize constellations of functional residues.