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 \emph{training} approach that jointly i) finds the most predictive subgraph by removing edges and/or nodes -- -\emph{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.
翻译:图神经网络(GNN)在解决图分类任务中已取得最先进性能。然而,大多数GNN架构从图中所有节点和边聚合信息,无论这些元素与当前任务的相关性如何,从而阻碍了其预测的可解释性。与以往工作不同,本文提出一种GNN训练方法,该方法联合实现以下目标:i)通过移除边和/或节点来找到最具预测性的子图——无需对子图结构做出假设——同时ii)优化图分类任务的性能。为此,我们利用强化学习解决由此产生的双层优化问题,并采用基于共形预测的奖励函数来反映分类器当前训练中的不确定性。在九个不同图分类数据集上的实验结果表明,我们的方法在性能上与基线方法相当,同时依赖显著更稀疏的子图,从而产生更具可解释性的GNN预测。