MbExplainer: Multilevel Bandit-Based Explanations for Downstream Models with Augmented Graph Embeddings

Abstract

Graph Neural Networks (GNNs) are a highly useful tool for performing various machine learning prediction tasks on graph-structured data. In many industrial applications, it is common that the graph embeddings generated from training GNNs are used in an ensemble model where the embeddings are combined with other tabular features, e.g., original node or edge features, in a downstream machine learning task. The tabular features may even arise naturally if, e.g., one tries to build a graph such that some of the node or edge features are stored in a tabular format. In this paper we address the problem of explaining the output of such ensemble models for which the input features consist of learned neural graph embeddings combined with additional tabular features. Therefore, we propose MbExplainer, a model-agnostic explanation approach for downstream models with augmented graph embeddings. MbExplainer returns a human-comprehensible triple as an explanation for an instance prediction of the whole pipeline consisting of three components: a subgraph with the highest importance, the topmost important node features, and the topmost important augmented downstream features. A game-theoretic formulation is used to take the contributions of each component and their interactions into account by assigning three Shapley values corresponding to their own specific games. Finding the explanation requires an efficient search through the local search spaces corresponding to each component. MbExplainer applies a novel multilevel search algorithm that enables simultaneous pruning of local search spaces in a computationally tractable way. In particular, three interweaved Monte Carlo Tree Searches are utilized to iteratively prune the local search spaces. MbExplainer also includes a global search algorithm that uses contextual bandits to efficiently allocate pruning budget among the local search spaces. We demonstrate the effectiveness of MbExplainer by presenting a set of comprehensive numerical examples on multiple public graph datasets for node classification, graph classification, and graph regression tasks.