Task-Agnostic Graph Neural Explanations

1 INTRODUCTION . Graph neural networks ( GNNs ) ( Kipf & Welling , 2017 ; Veličković et al. , 2018 ; Xu et al. , 2019 ) have achieved remarkable success in learning from real-world graph-structured data due to their unique ability to capture both feature-wise and topological information . Extending their success , GNNs are widely applied in various research fields and industrial applications including quantum chemistry ( Gilmer et al. , 2017 ) , drug discovery ( Wu et al. , 2018 ; Wang et al. , 2020 ) , social networks ( Fan et al. , 2019 ) , and recommender systems ( Ying et al. , 2018 ) . While multiple approaches have been proposed and studied to improve GNN performance , GNN explainability is an emerging area and has a smaller body of research behind it . Recently , explainability has gained more attention due to an increasing desire for GNNs with more security , fairness , and reliability . Being able to provide a good explanation to a GNN prediction increases model reliability and reduces the risk of incorrect predictions , which is crucial in fields such as molecular biology , chemistry , fraud detection , etc . Existing methods adapting the explanation methods for convolutional neural networks ( CNNs ) or specifically designed for GNNs have shown promising explanations on multiple types of graph data . A recent survey ( Yuan et al. , 2020 ) categorizes existing explanation methods into gradient-based , perturbation , decomposition , and surrogate methods . In particular , perturbation methods involve learning or optimization ( Ying et al. , 2019 ; Luo et al. , 2020 ; Yuan et al. , 2021 ; Lin et al. , 2021 ) and , while bearing higher computational costs , generally achieve state-of-the-art performance in terms of explanation quality . These methods train post-hoc explanation models on top of the prediction model to be explained . Earlier approaches like GNNExplainer ( Ying et al. , 2019 ) require training or optimizing an individual explainer for each data instance , i.e. , a graph or a node to be explained . In contrast , PGExplainer ( Luo et al. , 2020 ) performs inductive learning , i.e. , it only requires a one-time training , and the explainer can be generalized to explain all data instances without individual optimization . Compared to other optimization-based explanation methods , PGExplainer significantly improves the efficiency in terms of time cost without performance loss by learning . However , even state-of-the-art explanation methods like PGExplainer are still task-specific at training and hence suffer from two crucial drawbacks . First , current methods are inefficient in explaining multitask prediction for graph-structured data . For example , one may need to predict multiple chemical properties in drug discovery for a molecular graph . In particular , ToxCast from MoleculeNet has 167 prediction tasks . In these cases , it is common to apply a single GNN model with multiple output dimensions to make predictions for all tasks . However , one is unable to employ a single explainer to explain the above model , since current explainers are trained specifically to explain one prediction task . As a result , in the case of ToxCast , one must train 167 explainers to explain the GNN model . Second , in industry settings , it is common to train GNN models in a two-stage fashion due to scaling , latency , and label sparsity issues . The first stage trains a GNN-based embedding model with a massive amount of unlabeled data in an unsupervised manner to learn embeddings for nodes or graphs . The second stage trains lightweight models such as multilayer perceptrons ( MLPs ) using the frozen embeddings as input to predict the downstream tasks . In the first stage , the downstream tasks are usually unknown or undefined , and existing task-specific explainers can not be applied . Also , there can be tens to hundreds of downstream tasks trained on these GNN embeddings , and training a separate explainer for each task is undesirable and downright impossible . To address the above limitations , we present a new task-agnostic explanation pipeline , shown in Figure 1 , where we decompose a prediction model into a GNN embedding model and a downstream model , designing separate explainers for each component . We design the downstream explainers to cooperate with the embedding explainer . The embedding explainer is trained using a self-supervised training framework , which we dub Task-Agnostic GNN Explainer ( TAGE ) , with no knowledge of downstream tasks , models , or labels . In contrast to existing explainers , the learning objective for TAGE is computed at the graph or node embeddings without involving task-related predictions . In addition to eliminating the need for downstream tasks in TAGE , we argue that the self-supervision performed on the embeddings can bring additional performance boost in terms of the explanation quality compared to existing task-specific baselines such as GNNExplainer and PGExplainer . We summarize our contributions as follows : 1 ) We introduce the task-agnostic explanation problem and propose a two-stage explanation pipeline involving an embedding explainer and a downstream explainer . This enables the explanation of multiple downstream tasks with a single embedding explainer . 2 ) We propose a self-supervised training framework TAGE , which is based on conditioned contrastive learning to train the embedding explainer . TAGE requires no knowledge of downstream tasks . 3 ) We perform experiments on real-world datasets and observe that TAGE outperforms existing learning-based explanation baselines in terms of explanation quality , universal explanation ability , and the time required for training and inference . 2 TASK-AGNOSTIC EXPLANATIONS . 2.1 NOTATIONS AND LEARNING-BASED GNN EXPLANATION . Our study considers the attributed graph G with node set V and edge set E. We formulate the attributed graph as a tuple of matrices ( A , X ) , where A ∈ { 0 , 1 } |V |×|V | denotes the adjacency matrix and X ∈ R|V |×df denotes the feature matrix with feature dimension of df . We assume that the prediction model F that is to be explained operates on graph-structured data through two components : a GNN-based embedding model and lighter downstream models . Denoting the input space by G , a node-level embedding model En : G → R|V |×d takes a graph as input and computes embeddings of dimension d for all nodes in the graph , whereas a graph-level embedding model Eg : G → R1×d computes an embedding for the input graph . Subsequently , the downstream model D : Rd → R computes predictions for the downstream task based on the embeddings . Typical GNN explainers consider a task-specific GNN-based model as a complete unit , i.e. , F : = D ◦ E . Given a graph G and the GNN-based model F to be explained , our goal is to identify the subgraph Gsub that contributes the most to the final prediction made by F . In other words , we claim that a given prediction is made because F captures crucial information provided by some subgraphGsub . The learning-based ( or optimization-based ) GNN explanation employs a parametric explainer Tθ associated with the GNN model F to compute the subgraph Gsub of the given graph data . Concretely , the explainer Tθ computes the importance score for each node or edge , denoted as wi or wij , or masks for node attributes denoted as m. It then selects the subgraph Gsub induced by important nodes and edges , i.e. , whose scores exceed a threshold t , and by masking the unimportant attributes . In our study , we follow Luo et al . ( 2020 ) , focusing on the importance of edges to provide explanations to GNNs . Formally , we have Gsub : = ( V , Esub ) = Tθ ( G ) , where Esub = { ( vi , vj ) : ( vi , vj ) ∈ E , wij ≥ t } . 2.2 TASK-AGNOSTIC EXPLANATIONS . As introduced in Section 1 , all existing learning-based or optimization-based explanation approaches are task-specific and hence suffer from infeasiblity or inefficiency in many real-application scenarios . In particular , they are of limited use when downstream tasks are unknown or undefined and fail to employ a single explainer to explain a multitask prediction model . To enable the explanation of GNNs in two-stage training and multitask scenarios , we introduce a new explanation paradigm called the task-agnostic explanation . The task-agnostic explanation considers a whole prediction model as an embedding model followed by any number of downstream models . It focuses on explaining the embedding model regardless of the number or the existence of downstream models . In particular , the task-agnostic explanation trains only one explainer T ( tag ) θ to explain the embedding model E , which should satisfy the following features . First , given an input graph G , the explainer T ( tag ) θ should be able to provide different explanations according to specific downstream tasks being studied . Table 1 compares the properties of common GNN explanation methods and the desired task-agnostic explainers in multitask scenarios . Second , the explainer T ( tag ) θ can be trained when only the embedding model is available , e.g. , at the first stage of a two-stage training paradigm , regardless of the presence of downstream tasks . When downstream tasks and models are unknown , T ( tag ) θ can still identify which components of the input graph are important for certain embedding dimensions of interest . 3 THE TAGE FRAMEWORK . Our explanation framework TAGE follows the typical scheme of GNN explanation introduced in the previous section . It provides explanations by identifying important edges in a given graph and removing the edges that lead to significant changes in the final prediction . Specifically , the goal of the TAGE is to predict the importance score for each edge in a given graph . Different from existing methods , the proposed TAGE breaks down typical end-to-end GNN explainers into two components . We now provide general descriptions and detailed formulations to the proposed framework . 3.1 TASK-AGNOSTIC EXPLANATION PIPELINE . Following the principle of the desired task-agnostic explanations , we introduce the task-agnostic explanation pipeline , where a typical explanation procedure is performed in two steps . In particular , we decompose the typical end-to-end learning-based GNN explainer into two parts : the embedding explainer TE and the downstream explainer Tdown , corresponding to the two components in the two-stage training and prediction procedure . We compare the typical explanation pipeline and the two-stage explanation pipeline in Figure 1 . The embedding explainer and downstream explainers can be trained or constructed independently from each other . In addition , the embedding explainer can cooperate with any downstream explainers to perform end-to-end explanations on input graphs . The downstream explainer aims to explain task-specific downstream models . As downstream models are usually lightweight MLPs , we simply adopt gradient-based explainers for downstream explainers without training . The downstream explainer takes a downstream model and the graph or node embedding vector as inputs and computes the importance score of each dimension on the embedding vector . The importance scores then serve as a condition vector input to the embedding explainer . Given the condition vector , the embedding explainer explains the GNN-based embedding model by identifying an important subgraph from the input graph data . In other words , given different condition vectors associated with different downstream tasks or models , the embedding explainer can provide corresponding explanations for the same embedding model . Formally , we denote the downstream explainer for models from D by Tdown : D × Rd → Rd , which maps input models and embeddings into importance scores m for all embedding dimensions . We denote the embedding explainer associated with the embedding model E by TE : Rd × G → G , which maps a given graph into a subgraph of higher importance , conditioned on the embedding dimension importance m ∈ Rd . The training procedures of the embedding explainer are independent of downstream tasks or downstream explainers . In particular , the downstream explainer is obtained from the downstream model only , and the training of the embedding explainer only requires the embedding model and the input graphs . As downstream models are usually constructed as stacked fully connected ( FC ) layers and the explanation of FC layers has been well studied , our study mainly focuses on the non-trivial training procedure and design of the embedding explainer .

Summary: In this manuscript, a new explainer for GNNs is proposed. The newly proposed method aims to provide a task-agnostic explanation for the embedding GNNs rather than a specific downstream task. The motivation of the proposed method is that the modern GNNs are typically trained in a two-stage manner, where the embedding GNN is trained in the first stage and then the task-specific lightweight MLP. The proposed method is trained based on the MI between the whole graph and sub-graph generated by the explainer.