Abstract
Graph Neural Networks (GNNs) have achieved state-of-the-art performance insolving graph classification tasks. However, most GNN architectures aggregateinformation from all nodes and edges in a graph, regardless of their relevanceto the task at hand, thus hindering the interpretability of their predictions.In contrast to prior work, in this paper we propose a GNN \emph{training}approach that jointly i) finds the most predictive subgraph by removing edgesand/or nodes -- -\emph{without making assumptions about the subgraph structure}-- while ii) optimizing the performance of the graph classification task. Tothat end, we rely on reinforcement learning to solve the resulting bi-leveloptimization with a reward function based on conformal predictions to accountfor the current in-training uncertainty of the classifier. Our empiricalresults on nine different graph classification datasets show that our methodcompetes in performance with baselines while relying on significantly sparsersubgraphs, leading to more interpretable GNN-based predictions.