Prediction Tasks in Graphs: a Framework to Control the Interpretability-Performance Trade-off

Pablo Sanchez Martin, Kinaan Aamir Khan, Isabel Valera

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Graph Neural Networks; Reinforcement Learning; Graph-level tasks; Interpretability

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: GNN framework that simultaneously optimizes performance and sparsity using RL and achieves competitive accuracy while maintaining high sparsity.

Abstract: Graph Neural Networks (GNNs) have emerged as state-of-the-art methods for solving graph-level tasks in diverse domains, such as social network analysis and bioinformatics. However, their complex nature results in a lack of human-interpretable predictions, which can hinder their practical impact. Here, we aim at improving GNN interpretability by targeting \emph{sparsity} during GNN training - i.e, by minimizing the size (and/or number) of subgraphs used to make predictions. Existing solutions in the literature suffer from two main limitations: i) they still rely on information about the entire graph; and/or ii) they do not allow practitioners to directly control the trade-off between predictive performance and sparsity. To address the above limitations, in this paper, we formulate GNN training as a bi-level optimization task, where the trade-off between interpretability and performance can be controlled by a hyperparameter. Our framework relies on reinforcement learning to iteratively maximize predictive performance and sparsity by removing edges or nodes from the input graph. Our empirical results on nine different graph classification datasets show that our method competes in performance with baselines that use information from the whole graph, while relying on significantly sparser subgraphs, leading to more interpretable GNN-based predictions.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 3265