Hierarchical Pooling and Explainability in Graph Neural Networks for Tumor and Tissue-of-Origin Classification Using RNA-seq Data

TL;DR

The study tackles cancer classification from RNA-seq data, focusing on tissue origin and tumor versus normal differentiation. It introduces a graph neural network with hierarchical pooling on a STRING-based protein interaction backbone, using Chebyshev convolutions with $K=2$ and weighted pooling to produce supernodes, followed by a small classifier. Key findings show that increasing depth via additional coarsening and convolutions does not improve performance, with the best F1-macro of $0.978$ achieved using a single pooling layer, while the model remains highly interpretable through gradient-based gene and supernode analyses. The hierarchical structure provides biologically meaningful explanations and indicates potential for cancer biomarker discovery and explainable architectures, though fixed clustering and over-smoothing present limitations to address in future work.

Abstract

This study explores the use of graph neural networks (GNNs) with hierarchical pooling and multiple convolution layers for cancer classification based on RNA-seq data. We combine gene expression data from The Cancer Genome Atlas (TCGA) with a precomputed STRING protein-protein interaction network to classify tissue origin and distinguish between normal and tumor samples. The model employs Chebyshev graph convolutions (K=2) and weighted pooling layers, aggregating gene clusters into 'supernodes' across multiple coarsening levels. This approach enables dimensionality reduction while preserving meaningful interactions. Saliency methods were applied to interpret the model by identifying key genes and biological processes relevant to cancer. Our findings reveal that increasing the number of convolution and pooling layers did not enhance classification performance. The highest F1-macro score (0.978) was achieved with a single pooling layer. However, adding more layers resulted in over-smoothing and performance degradation. However, the model proved highly interpretable through gradient methods, identifying known cancer-related genes and highlighting enriched biological processes, and its hierarchical structure can be used to develop new explainable architectures. Overall, while deeper GNN architectures did not improve performance, the hierarchical pooling structure provided valuable insights into tumor biology, making GNNs a promising tool for cancer biomarker discovery and interpretation

Figures (6)

• Figure 1: General architecture of the GNNs used in this work. In (a), the colors represent the clusters of each protein at the input graph. These nodes undergo multiple pooling and non-linearity layers (b) and their final embeddings are concatenated and used as inputs to a FC network (c).
• Figure 2: (a) Two steps of the heavy-edge matching algorithm. Starker edges represent stronger connections and the colored ellipses indicate which nodes are aggregated together. (b) A dendrogram representing the clusters of input nodes associated with the supernodes. The supernode in black corresponds to the aggregation of the reddish and blueish inputs.
• Figure 3: Performance on tumor prediction task drops slightly as the number of pooling + graph convolution layers increase, with the highest score obtained with a single layer (a). Having no graph convolutions at the initial layers leads to performance drop (b).
• Figure 4: Test results on cohort classification and the tumor prediction tasks using a multi-task model. The confusion matrix for the Pan-cancer Cohort classification is shown in (a) together with the precision and recall obtained for each category in (b), whereas the confusion matrices and precision-recall plots are shown in (c) and (d) for the Pan-cancer Tumor Prediction task.
• Figure 5: Example embeddings (a) and their saliencies (b) produced when distinguishing normal from tumorous samples. Each row corresponds to a supernode in a coarsened version of the STRING.