E2EGraph: An End-to-end Graph Learning Model for Interpretable Prediction of Pathlogical Stages in Prostate Cancer

Prostate cancer is one of the deadliest cancers worldwide. An accurate prediction of pathological stages using the expressions and interactions of genes is effective for clinical assessment and treatment. However, identification of interactions using biological procedure is time consuming and prohibitively expensive. A graph is a powerful representation for the complex interactome of genes, their transcripts, and proteins. Recently, Graph Neural Networks (GNNs) have gained great attention in machine learning due to their capability to capture the graphical interactions among data entities. To leverage GNNs for predicting pathological stage stages, we developed an end-to-end graph representation and learning model, namely E2EGraph, which can automatically generate a graph representation using gene expression data and a multi-head graph attention network to learn the strength of interactions among genes and make the prediction. To ensure the reliability of model prediction, we identify critical components of graph representation and GNN model to interpret prediction results from multiple perspectives at gene and patient levels. We evaluated E2EGraph to predict pathological stages of prostate cancer using The Cancer Genome Atlas (TCGA) data. Our experimental results demonstrate that E2EGraph reaches the state-of-art prediction performance while being effective in identifying marker genes indicated by interpretability. Our results point to a direction where adaptive graph construction and attention based GNNs can be leveraged for various prediction tasks and interpretation of model prediction in a variety of data domains including disease prediction.