2024/04/18 by Pablo Sánchez-Martín, Sanchez-Martin, Pablo, Kinaan Aamir Khan +3 · 2 citations
Computer Science · #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Social and Information Networks (cs.SI) #Topic Modeling
paper · pdf · doi:10.48550/arxiv.2404.12356
openalex publication_date 2024/04/18 · openalex created_date 2024/04/20 · openalex updated_date 2026/08/01
Graph Neural Networks (GNNs) have achieved state-of-the-art performance in solving graph classification tasks. However, most GNN architectures aggregate information from all nodes and edges in a graph, regardless of their relevance to the task at hand, thus hindering the interpretability of their predictions. In contrast to prior work, in this paper we propose a GNN training approach that jointly i) finds the most predictive subgraph by removing edges and/or nodes -- -without making assumptions about the subgraph structure -- while ii) optimizing the performance of the graph classification task. To that end, we rely on reinforcement learning to solve the resulting bi-level optimization with a reward function based on conformal predictions to account for the current in-training uncertainty of the classifier. Our empirical results on nine different graph classification datasets show that our method competes in performance with baselines while relying on significantly sparser subgraphs, leading to more interpretable GNN-based predictions.