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: The authors propose TAGE, a task-agnostic explanation method for explaining GNNs. TAGE explains GNN embedding models without downstream tasks and allows the explanation of multi-task models. This paper maximizes the mutual information of masked graph embedding and masked subgraph embedding as the objective function. The regularization term is added to restrict the number of edges. Empirical results focus on three aspects: (1) Improvement in fidelity scores brought by TAGE; (2) Visualizations on explanations to the GNN model; and (3) Comparison of computational time cost among GNN explainers.