Towards Training GNNs using Explanation Directed Message Passing
Abstract
With the increasing use of Graph Neural Networks (GNNs) in critical real-world applications, several post hoc explanation methods have been proposed to understand their predictions.
However, there has been no work in generating explanations on the fly during model training and utilizing them to improve the expressive power of the underlying GNN models.
In this work, we introduce a novel explanationdirected neural message passing framework for GNNs, E XPASS (EXplainable message PASSing), which aggregates only embeddings from nodes and edges identified as important by a GNN explanation method.
E XPASS can be used with any existing GNN architecture and subgraph-optimizing explainer to learn accurate graph embeddings.
We theoretically show that E XPASS alleviates the oversmoothing problem in GNNs by slowing the layer-wise loss of Dirichlet energy and that the embedding difference between the vanilla message passing and E XPASS framework can be upper bounded by the difference of their respective model weights.
Our empirical results show that graph embeddings learned using E XPASS improve the predictive performance and alleviate the oversmoothing problems of GNNs, opening up new frontiers in graph machine learning to develop explanation-based training frameworks.
Introduction
Graph Neural Networks (GNNs) are increasingly used as powerful tools for representing graphstructured data, such as social, information, chemical, and biological networks [1, 2].
With the deployment of GNN models in critical applications (e.g., financial systems and crime forecasting [3, 4]), it becomes essential to ensure that the relevant stakeholders understand and trust their decisions.
To this end, several approaches [5–13] have been proposed in recent literature to generate post hoc explanations for predictions of GNN models.
In contrast to other modalities like images and texts, generating instance-level explanations for graphs is non-trivial.
In particular, it is more challenging since individual node embeddings in GNNs aggregate information using the entire graph structure, and, therefore, explanations can be on different levels (i.e., node attributes, nodes, and edges).
While several categories of GNN explanation methods have been proposed: gradient-based [5, 10, 14], perturbation-based [8, 9, 11, 13, 15], and surrogatebased [7, 12], their utility is limited to generating post hoc node- and edge-level explanations for a given pre-trained GNN model.
Thus, the capability of GNN explainers to improve the predictive performance of a GNN model lacks understanding as there is very little work on systematically analyzing the reliability of state-of-the-art GNN explanation methods on model performance [16].
[Table]
To address this, recent works have explored the joint optimization of machine learning models and explanation methods to improve the reliability of explanations [17, 18].
Zhou et al.
[18] proposed
DropEdge as a technique to drop random edges (similar to generating random edge explanations) during training to reduce overfitting in GNNs.
More recently, Spinelli et al.
[17] used meta-learning frameworks to generate GNN explanations and show an improvement in the performance of specific
GNN explanation methods.
While these works make an initial attempt at jointly optimizing explainers and predictive models, they are neither generalizable nor exhaustive.
They fail to show improvement in the downstream GNN performance [17] and degree of explainability [18] across diverse GNN architectures and explainers.
Further, there is little to no work done on either theoretically analyzing the effect of GNN explanations on the neural message framework in GNNs or on important GNN properties like oversmoothing [19].
Present work.
In this work, we introduce a novel explanation-directed neural message passing framework, E XPASS , which can be used with any GNN model and subgraph-optimizing explainer to learn accurate graph representations.
In particular, E XPASS utilizes GNN explanations to steer the underlying GNN model to learn graph embeddings using only important nodes and edges.
E XPASS aims to define local neighborhoods for neural message passing, i.e., identify the most important edges and nodes, using explanation weights, in the k -hop local neighborhood of every node in the graph.
Formally, we augment existing message passing architectures to allow information flow along important edges while blocking information along irrelevant edges.
We present an extensive theoretical and empirical analysis to show the effectiveness of E XPASS on the predictive, explainability, and oversmoothing performance of GNNs.
Our theoretical results show that the embedding difference between vanilla message passing and E XPASS frameworks is upperbounded by the difference between their model weights.
Further, we show that embeddings learned using E XPASS relieve the oversmoothing problem in GNNs as they reduce information propagation by slowing the layer-wise loss of Dirichlet energy (Section 4.2).
For our empirical analysis, we integrate E XPASS into state-of-the-art GNN models and evaluate their predictive, oversmoothing, and explainability performance on real-world graph datasets (Section 5).
Our results show that, on average, across five GNN models, E XPASS improves the degree of explainability of the underlying
GNNs by 39.68%.
Our ablation studies show that for an increasing number of GNN layers, E XPASS achieves 34.4% better oversmoothing performance than its vanilla counterpart.
Finally, our results demonstrate the effectiveness of using explanations during training, paving the way for new frontiers in GraphXAI research to develop explanation-based training algorithms.
Related works
Graph Neural Networks.
Graph Neural Networks (GNNs) are complex non-linear functions that transform input graph structures into a lower dimensional embedding space.
The main goal of
GNNs is to learn embeddings that reflect the underlying input graph structure, i.e., neighboring nodes in the graph are mapped to neighboring points in the embedding space.
Prior works have proposed several GNN models using spectral and non-spectral approaches.
Spectral models [20–24] leverage Fourier transform and graph Laplacian to define convolution approaches for GNN models.
However, non-spectral approaches [25–29] define the convolution operation by leveraging the local neighborhood of individual nodes in the graph.
Most modern non-spectral models are messagepassing frameworks [30, 31], where nodes update their embedding by aggregating information from k -hop neighboring nodes.
Post hoc Explanations.
With the increasing development of complex high-performing GNN models [25–29], it becomes critical to understand their decisions.
Prior works have focused on developing several post hoc explanation methods to explain the decisions of GNN models [5, 7, 9, 11–13, 32].
More specifically, these explanation methods can be broadly categorized into i) gradient-based methods [5] that leverage the gradients of the GNN model to generate explanations; ii) perturbation-based methods [9, 11, 13] that aim to generate explanations by calculating the change in GNN predictions upon perturbations of the input graph structure (nodes, edges, or subgraphs); and iii) surrogate-based methods [7, 12] that fit a simple interpretable model to approximate the predictive behavior of the given GNN model.
Finally, recent works have introduced frameworks to theoretically and empirically analyze the behavior of state-of-the-art GNN explanation methods with respect to several desirable properties [16, 33].
Preliminaries
Notations.
Let G =
( V , E , X ) denote an undirected graph comprising of a set of nodes V and a set of edges
E .
Let X
= { x 1 , x 2 , . . .
, x N } denote the set of node feature vectors for all nodes in V , where x v ∈ R
d captures the attribute values of a node v and N = |V| denotes the number of nodes in the graph.
Let A ∈ R N ×
N
be the graph adjacency matrix,
where element A uv = 1 if there exists an edge e ∈ E between nodes u and v and A uv = 0 otherwise.
We use N u to denote the set of immediate neighbors of node u , i.e., ,
N u
= { v ∈ V|
A uv = 1 } .
Finally, the function deg : V (cid:55)→ Z > 0 is defined as deg ( v ) =
|N v | and outputs the degree of a node v ∈ V
Graph Neural Networks (GNNs).
Formally, GNNs can be formulated as message passing networks [30] specified by three key operators M SG , A GG , and U PD .
These operators are recursively applied on a given graph G for a L -layer GNN model defining how neural messages are shared, aggregated, and updated between nodes to learn the final node representations in the L th layer of the GNN.
Commonly, a message between a pair of nodes ( u, v ) in layer l is characterized as a function of their hidden representations h ( l − 1) u and h ( l − 1) v from the previous layer: m ( l ) uv =
M SG ( h ( l − 1) u , h
( l − 1) v ) .
The A GG operator retrieves the messages from the neighborhood of node u and aggregates them as: m ( l ) u =
A GG ( m ( l ) uv | v ∈ N u ) .
Next, the U PD operator takes the aggregated message m ( l ) u at layer l and combines it with h ( l − 1) u to produce node u ’s representation for layer l as h ( l ) u
= U PD ( m ( l ) u , h
( l − 1) u ) .
Lastly, the final node representation for node u is given as z u = h ( L ) u .
Graph Explanations.
In contrast to other modalities like images and texts, an explanation method for graphs can formally generate multi-level explanations.
For instance, in a graph classification task, the explanations for a given graph prediction can be with respect to the node attributes M x ∈ R d , nodes M n ∈ R N , or edges M e ∈ R N × N .
Note that these explanation masks are continuous but can be discretized using specific thresholding strategies [33].
Oversmoothing.
Cai et al.
[34] and Zhou et al.
[35] defined bounds for analyzing oversmoothing for a GNN using Dirichlet Energy.
For a graph G with adjacency matrix A and degree matrix D , we define ˜
A
=
A + I N and ˜ D
= D
+ I N as the adjacency and degree matrices respectively of the graph G with self-loops.
We also define the augmented normalized Laplacian of G as ˜∆= I N − ˜ D − 12 ˜
A ˜ D − 12 , and P =
I N − ˜∆ .
Our Framework: E XPASS
Here, we describe E XPASS , our proposed explainable message-passing framework that aims to learn accurate and interpretable graph embeddings.
In particular, E XPASS incorporates explanations into the message-passing framework of GNN models by only aggregating embeddings from key nodes and edges as identified using an explanation method.
Problem formulation (Explanation Directed Message Passing).
Given a graph G =
( V , E , X ) ,
E XPASS aims to generate a d -dimensional embedding z u ∈ R d for each node u ∈ V using an explanation-directed message passing framework that filters out the noise from unimportant edges and improves the expressive power of GNNs.
4.1 Explanation Directed Message Passing
The central idea of E XPASS is to propose a novel method for improving the neural message passing scheme of GNN models by utilizing explanations during model training and aggregating important neural messages along edges in graph neighborhoods.
Next, we describe the existing message-passing scheme in GNNs and our explainable counterpart.
Message Passing.
As described in Section 3, each GNN layer can be described using the M SG , A GG , and U PD operators.
For each node u ∈ V , the ( l + 1) th layer embeddings h ( l +1) u is computed using a
GNN operating on the node’s neighboring attributes.
Formally, the GNN layer can be formulated as:
UPDATE
AGGREGATEMESSAGE
[Equation]
s u
v
≈
k=1 u k=
u b c vvvvvvvvvv 5 s u v
≈ k=1: 1-hop neighborhood
Neural Message Passing
Explanation-directed Neural Message Passing (EXPASS) s uv i : importance of node v i on node u
G =
( V , E , X )
Graph Machine Learning Model
Y
GraphExplainer
[Figure]
weights 0.
where h ( l +1) u represents the updated embedding of node u , ψ is the M SG operator, (cid:76) is the A GG operator ( e.g., summation), ϕ is an U PD function ( e.g., any non-linear activation function), and h ( l ) u represents the embedding of node u from the previous layer.
We obtain an embedding z u for node u by stacking L GNN layers.
Finally, the node embeddings Z ∈ R are then passed to a R EADOUT function to obtain an embedding for the graph.
E XPASS .
Here, we describe our proposed explainable message-passing scheme that incorporates explanations into the message-passing step in individual GNN layers on the fly during the training process.
Given an explanation method, which generates an importance score s uv ∈ M eu for every edge e uv ∈ E , we can weight the edge contribution in the neighborhood N u of node u as:
[Equation]
Note that E XPASS is agnostic to explanation types and can also incorporate explanations on node attributes and node level.
For instance, the importance scores for individual nodes can be computed by averaging the outgoing scores s uv for all v ∈ N u .
Subsequently, we can replace the s uv score by using the average score s u to weight edges in the E XPASS layers, and for node attributes, we can multiply the node attribute explanation M au to the original node attribute vector.
To enable explainable message passing and only retain the important embeddings for node u , E XPASS removes the knowledge of irrelevant nodes and edges from the local neighborhood N u of node u using its explanations.
For instance, if node v is considered important to node u , E XPASS transforms the aggregated messages of node u using the node importance scores s uv .
Note that since the explanations of node u include important nodes and edges in the L -hop neighborhood of node u , even though node u is only locally modified, the change will spread through all the nodes in every GNN layer.
Furthermore, to avoid spurious correlations, we ensure that explanations are only generated for correctly classified nodes and graphs.
Explanation weights infuse information from higher-order neighborhoods into each layer of the GNN model, specifically, from as many L -hop neighbors because explanation weights within each layer are computed using the L -layer GNN model.
To illustrate this, we next show the weight computations for a GNN explanation method.
Without loss of generality, let us consider GNNExplainer as our explanation method whose mask for the selected graph is formulated as: G mask =
( X ′ , A ′ )
=
( X ⊙ σ ( M x ) , A ⊙ σ ( M e )) , where
W =
[ M x , M e ] are the explainers parameters, σ is the sigmoid function, and ⊙ denotes elementwise multiplication.
Here, s uv represents the element in row v and column u of M e .
Gradient descent-based optimization is used to find the optimal values for the masks minimizing the following objective:
[Equation]
, where f θ is the L -layer GNN model and
C is the total number of classes.
This shows that a L -hop neighborhood is used to compute s uv .
Formally, it minimizes the uncertainty of the predictive model when the GNN computation is limited to the explanation subgraph.
This uncertainty is minimized as a proxy of the maximization of the mutual information between the prediction with the unmasked graph and masked graph.
4.2 Theoretical Analysis
Here, we provide a detailed theoretical analysis of our proposed E XPASS framework.
In particular, we (i) provide a theoretical upper bound on the embedding difference obtained from a vanilla message passing and E XPASS framework and (ii) show that graph embeddings learned using E XPASS relieves the oversmoothing problem in GNNs by reducing information propagation.
Theorem 1 (Differences between E XPASS and Vanilla Message Passing) .
Given a non-linear activation function σ that is Lipschitz continuous, the difference between the node embeddings between a vanilla message passing and E XPASS framework can be bounded by the difference in their individual weights, i.e.,
[Equation]
where W ( l ) a and W ′ ( l ) a are the weights for node u in layer l of the vanilla message passing and
E XPASS framework and W ( l ) n and W ′ ( l ) n are their weight matrix with neighbors of node u at layer l .
Proof Sketch.
In Theorem 1, we prove that the ℓ 2 -norm of the differences between the embeddings of vanilla message passing and E XPASS framework at layer l is upper bounded by the difference between their weights and the embeddings of node u and its subgraph.
See Appendix A for more details.
Definition 1 (Dirichlet Energy for a Node Embedding Matrix [35]) .
Given a node embedding matrix
H ( l ) =
[ h ( l ) 1 , . . .
, h ( l ) n ] T learned from the GNN model at the l th layer, the Dirichlet Energy E ( H ( l ) ) is defined as:
[Equation]
where a ij are elements in the adjacency matrix ˜ A and deg i , deg j is the degree of node i and j , respectively.
Cai et al.
[34] extensively show that higher Dirichlet energies correspond to lower oversmoothing.
Furthermore, they show that the removal of edges or, similarly, the reduction of edge weights on graphs helps alleviate oversmoothing.
Proposition 1 (E XPASS relieves Oversmoothing) .
E XPASS alleviates oversmoothing by slowing the layer-wise loss of Dirichlet energy.
The complete proof is provided in Appendix A.
Experiments
Next, we present experimental results for our E XPASS framework.
More specifically, we address the following questions: Q1) Does E XPASS enable GNNs to learn more accurate embeddings and improve their degree of explainability?
Q2)
How does E XPASS affect the oversmoothing and predictive performance of GNNs with an increasing number of layers?
Q3)
Does E XPASS depend on the quality of explanations for improving the predictive and oversmoothing performance of GNNs and are they better than attention weights?
Q4)
How does E XPASS helps in the evolution of explanation during the training of the GNN model?
5.1 Datasets and Experimental setup
We first describe the datasets used to study the utility of our proposed E XPASS framework and then outline the experimental setup.
Datasets.
We use real-world molecular chemistry datasets to evaluate the effectiveness of E XPASS w.r.t.
the performance of the underlying GNN model and understand the trade-off between explainability and accuracy for a graph classification task.
We consider four benchmark datasets, which includes Mutag [36], Alkane-Carbonyl [37], DD [38], and Proteins [39].
See Appendix B.1 for a detailed overview of the datasets.
GNN Architectures and Explainers.
To investigate the flexibility of E XPASS , we incorporate it into five different GNN models: GCN [40], GraphConv [41], LEConv [42], GraphSAGE [28], GAT [43], and GIN [27].
We use GNNExplainer [13] as our baseline GNN explanation method to generate edge-level explanations for most of our experiments.
In addition, we use Integrated Gradients [44] and PGMExplainer [12], a node-level explanation method, to demonstrate E XPASS ’s sensitivity to the choice of explainers.
Implementation details.
We consider DropEdge [45] as our baseline method for comparing the oversmoothing performance of E XPASS as DropEdge randomly removes edges from the input graph at each training epoch, acting like a message passing reducer.
Across all experiments, we use topK (k=40%) node features/edges, and use them to generate explanations for all explanation methods.
All other hyperparameters of the explanation and baseline methods were set following the author’s guidelines.
For all our experiments (unless mentioned otherwise), we use the baseline architectures with three GNN layers followed by ReLU layers and set the hidden dimensionality to 32.
Finally, we use a single linear layer to transform the graph embeddings to their respective classes.
See
Appendix B.2 for more details.
Performance metrics for GNN Explainers.
To measure the reliability of GNN explanation methods, we use the graph explanation faithfulness metric [16]: GEF (ˆ y u , ˆ y u ′ )
= 1 − exp − KL
(ˆ y u
|| ˆ y u ′ ) , where ˆ y u is predicted probability vector using the whole subgraph and ˆ y u ′ is the predicted probability vector using the masked subgraph, where we generate the masked subgraph by only using the topK features identified by an explanation and the Kullback-Leibler (KL) divergence score (denoted by “ || ” operator) quantifies the distance between two probability distributions.
Note that GEF is a measure of the unfaithfulness of the explanation.
So, higher values indicate a higher degree of unfaithfulness.
Performance metrics for Oversmoothing.
Zhou et al.
[18] introduced the Group Distance Ratio (GDR) metric to quantify oversmoothing in GNNs.
It measures the ratio between the average of pairwise representation distances between graphs belonging to different (inter) and same (intra) groups.
Formally, one would prefer to reduce the intra-group class representations and increase the inter-group distance to relieve the over-smoothing issue.
Hence, lower GDR values denote higher oversmoothing in GNNs.
Burn-in period.
We defined the burn-in period as a number n of epochs during training in which no explanations are used.
The burn-in period is necessary to avoid feeding spurious explanations to the model.
The length of the burn-in period (i.e., the number of epochs) was treated as a hyperparameter and fine-tuned using the validation set.
At the end of the burn-in period, a predefined percentage of correctly predicted graphs per batch is randomly sampled and their explanations are used in the model training.
The percentage of correctly predicted graphs sampled in each batch was set to 0.4 for all our experiments.
See Appendix C.2 for ablation on burn-in periods.
5.2 Results
Q1)
E XPASS improves the predictive performance and explainability of GNNs.
To measure the predictive performance and degree of explainability of GNNs trained using E XPASS , we compute
[Figure]
G r oup
D i s t an c e R a t i o ( )
[Figure]
r
[Figure]
D i s t
[Figure]
c e R a t i o ( )
[Figure]
their average predictive performance (using AUROC and F1-score) and fidelity (using Graph Explanation Faithfulness) using different GNN models and datasets.
Across four datasets and five GNN architectures, we find that E XPASS -augmented GNNs learn graph embeddings that are more accurate (higher AUROC and F1-score) and result in more faithful explanations (lower Graph Explanation
Faithfulness score) than their vanilla counterparts.
On average, E XPASS improves the AUROC and F1-score by 1.51% and 1.05%, respectively.
In particular, we observe that E XPASS improves the predictive behavior of high-performing models like GIN (+2.06% in AUROC and +2.50% in
F1-score)
but shows little to no improvement in the case of LeConv, which utilizes a node-scoring mechanism through the similarity between a node and its neighbors’ embeddings.
Finally, we find that E XPASS -augmented GNNs significantly improve the explainability of a GNN and achieve a 39.68% better faithfulness score as compared to vanilla GNNs (Table 1).
See Appendix C.1 for results on node classification graph downstream tasks.
Q2)
E XPASS relieves Oversmoothing in GNNs
We examine the oversmoothing (using the Group
Distance Ratio metric [18]) and predictive performance of GNNs trained using E XPASS with their vanilla counterparts.
The oversmoothing problem in GNNs shows that the representations of nodes converge to similar vectors as the number of layers increases.
Therefore, we analyze the oversmoothing of the GNNs for an increasing number of layers and find that, on average, across two architectures,
E XPASS improves the group distance ratio by 34.4% (Figure 2).
Further, we also analyzed the oversmoothing behavior of E XPASS for node classification tasks (in Appendix C.1) and our results indicate an inherent trade-off between oversmoothing and predictive performance of GNNs (Figures 5-7).
Q3)
Ablation studies.
We conduct ablations on several components of E XPASS with respect to its oversmoothing and predictive performance.
E XPASS for different TopK Explanations.
We investigate the oversmoothing and predictive performance of GNNs for different topK explanations (i.e., topK edges identified by a GNN explanation) chosen in the message passing.
Results show that E XPASS alleviates oversmoothing by using only the topK edges to learn graph embeddings and explicitly filter out the noise from unimportant edges.
In particular, we observe that the GDR values decrease (denoting higher oversmoothing) with the increase in the use of topK edges (Figure 3).
More specifically, we find that the GDR value at topK=0.1 is 11.92% higher than vanilla message passing (i.e., using all edges in the graph).
E XPASS vs. DropEdge.
We compare the predictive and oversmoothing and predictive performance of E XPASS and DropEdge.
Here, we show that message passing using optimized explanation-directed information outperforms random edge removal.
We find that E XPASS outperforms DropEdge across both oversmoothing and accuracy metrics.
In particular, on average, across different topK values,
E XPASS improves the oversmoothing, AUROC, and F1-score performance of vanilla message passing by 71.16%, 9.53%, and 12.63%, respectively (Figure 3).
E XPASS using Node Explanations.
We investigate the effect of the choice of the baseline explanation method on the performance of E XPASS with respect to the vanilla message passing framework.
More specifically, we evaluate the predictive and explainability performance of E XPASS -augmented GNNs when trained using node explanations generated using Integrated Gradients (IG) [44].
Similar to
[Table]
ROTEINS
[Table]
the results of E XPASS with GNNExplainer as the baseline explanation method (Table 1)
[Table]
we find that E XPASS trained using IG explanations also improves the AUROC (+2.80%), F1-score (+1.11%), and GEF (+23.67%) of the vanilla GNN model.
Our results show that the choice of explainer can make a difference in the E XPASS performance, depending on the dataset.
For instance, IG is a node-masking explainer that is not considered a strong explanation method and its effects are variable across datasets [33].
We recommend using graph-specific explainers that optimize for fidelity and sparsity on the edges of the input graph, which would be a best fit to increase the performance of the network.
See Appendix C.3 for results using PGMExplainer.
Further, our results show that
E XPASS is a model- and explainer-agnostic framework that can improve the downstream task and explainability performance across different GNN architectures using diverse GNN explainers.
[Figure]
D i s t
[Figure]
a t i o (
[Figure]
Dataset
Method
AUROC
[Table]
E XPASS vs. Attention.
We demonstrate the utility of using explanations vs. attention weights in the message passing step using GAT [43] model architecture.
On average, across four datasets, we find that E XPASS achieves higher AUROC (+3.85%) and F1-score (+2.24%) than the attention-based
GAT model (Table 3).
In addition, GNNExplainer [13] demonstrated that post hoc GNN explainers generate better explanations than attention weights, which further highlights the benefits of E XPASS .
In comparison to E XPASS , GAT can be considered as a special case of our framework, where attention weights replace explanations.
On the other hand, E XPASS has larger benefits since it can be applied to any existing GNN architectures that lack explainability.
Q4)
Visualizing explanations.
Here, we visualize how the explanation develops over the training process of the GNN model.
In particular, we visualize the generated explanations from E XPASS -GCN with GNNExplainer trained on the MUTAG dataset at different epochs during the training process and find that the explanations converge to the ground-truth explanation of a non-mutagenic molecule (i.e., the absence of a carbon ring alongside the highlighted NO 2 molecules) as the training progresses (Figure 8).
Further, we compare the generated explanations for a vanilla GCN and its E XPASS counterpart and find that the explanation for vanilla GCN falsely identifies the carbon-carbon bonds as important (Figure 4).
This qualitative analysis provides further evidence for the observed higher faithfulness results (Table 1) of explanations generated using our proposed E XPASS framework.
Conclusion and Discussion
In this work, we propose the problem of learning graph embeddings using explanation-directed message passing in GNNs.
To this end, we introduce E XPASS , a novel message-passing framework that can be used with any existing GNN model and subgraph-optimizing explainer to learn accurate embeddings by aggregating only embeddings from nodes and edges identified as important by a
GNN explainer.
We perform an extensive theoretical analysis to show that E XPASS relieves the oversmoothing problem in GNNs, and the embedding difference between the vanilla message passing
[Table]
[Figure]
[Table]
framework and E XPASS can be upper bounded by the difference of their respective layer weights.
Our empirical results on benchmark datasets show that E XPASS improves the explainability of the underlying GNN model without sacrificing its predictive performance.
However, the training of
E XPASS depends on the choice of explanation method, the number of data points to explain, and the dataset of choice, which is computationally more expensive than its vanilla counterparts.
We find that the training time of E XPASS can be improved by using techniques like batch processing and efficient sampling of correctly-classified nodes and graphs.
Further, adapting post-hoc explainers to generate subgraphs utilizing the embedding space would also improve the computation time of E XPASS .
Our proposed method and findings open exciting new avenues to learn graph representations by jointly training models and explanation methods.
We anticipate that E XPASS could open new frontiers in graph machine learning for developing explanation-based training frameworks.
Acknowledgements
The authors would like to thank the anonymous reviewers for their helpful feedback that helped improve the work.
CA would like to thank Lasse Mohr and Samuele Firmani for the helpful discussions at the beginning of the project and LOGML Summer School for connecting with the students.
The views expressed here are those of the authors and do not reflect the official policy or position of the affiliated company.
A Proofs for Theorems in Section
Theorem 1.
Given a non-linear activation function σ that is Lipschitz continuous, the difference between the node embeddings between a vanilla message passing and E XPASS framework can be bounded by the difference in their individual weights, i.e.,
[Equation]
where W ( l ) a and W ′ ( l ) a are the weights for node u in layer l of the vanilla message passing and
E XPASS framework and W ( l ) n and W ′ ( l ) n are their respective weight matrix with the neighbors of node u at layer l .
Proof.
For a given node u , the node representation output by layer l of the GNN is given by:
[Equation]
where we consider the A GG operator as a fully-connected layer, U PD to be a sigmoid activation function σ ( · ) , W ( l ) a is the weights for node u in layer l and W ( l ) n is the weight matrix with the neighbors of node u at layer l .
Let us consider an edge in-hoc explanation that generates a binary mask highlighting the important edges for the prediction of node u .
Note that using the edge mask, we can also get a node-level mask signifying the importance of neighboring nodes.
Let us denote that node explanation mask as s v where s v
= 1 if the node is important, otherwise s v = 0 .
Formally, the corresponding message passing equations for E XPASS can be written as:
[Equation]
where h ′ ( l ) u and h ′ ( l ) v represents the embeddings of node u and v using the feedback explanation, and W ′ ( l ) a and W ′ ( l ) n represents the corresponding weights at layer l for GNN model trained using
E XPASS .
The difference between the node embeddings obtained after the message-passing in layer l from
Equations 4-5 is given as:
[Equation]
Taking the ℓ 2 -norm on both sides and assuming a normalized Lipschitz non-linear sigmoid activation, i.e., ∥ σ ( b ) − σ ( a
[Equation]
∥
[Equation]
≤ ∥ b − a ∥ 2 , we get:
[Equation]
(Using Triangle Inequality and Faithfulness property of explanations)
Given a faithful explanation, the node embeddings for node u using the vanilla message passing network are equivalent to that E XPASS since most explainers optimize the mask to approximate the input embedding.
More specifically, for a given node embedding h ′ ( l − 1) u
[Equation]
a faithful explanation bounds the ϵ u to zero.
In addition to faithfulness, a GNN using vanilla message passing and E XPASS can predict a node u to the same class only if both frameworks generate similar node embeddings (Proposition 1 in Agarwal et al.
Using Matrix-norm and Triangle Inequality for the sum in the neighborhood, we get:
[Equation]
Again, using the faithfulness property of explanations, the contribution of node embeddings from node v ∈ N ( u ) ∩ s v = 0 is irrelevant to the final embedding and can be removed.
Finally, using Matrix-norm inequality on the first term, we get:
[Equation]
Thus, we observe that the embedding difference at layer l between a vanilla message passing network and the E XPASS is purely based on the difference between their weights and the embeddings of node u and its subgraph.
Definition 2 (Dirichlet Energy for a Node Embedding Matrix [35]) .
Given a node embedding matrix h ( l ) =
[ h ( l ) 1 , . . .
, h ( l ) n ] T learned from the GNN model at the l th layer, the Dirichlet Energy E ( h ( l ) ) is defined as:
[Equation]
where a ij are elements in the adjacency matrix ˜ A and d i ,
d j is the degree of node i and j , respectively.
Cai et al.
[34] extensively show that higher Dirichlet energies correspond to lower oversmoothing.
Furthermore, they show that the removal of edges or ,similarly, reduction of edge weights on graphs help alleviate oversmoothing.
Proposition 1 (E XPASS relieves Oversmoothing).
E XPASS alleviates oversmoothing by slowing the layer-wise loss of Dirichlet energy.
Proof Sketch.
Here, we show the capabilities of E XPASS as a framework that alleviates the oversmoothing problem in GNNs.
To this end, we utilize the bounds on the Dirichlet energy of the
E XPASS embeddings at the l th layer of the GNN model by Zhou et al.
[Equation]
[Equation] where λ 1 , λ 0 are the non-zero eigenvalues of the symmetric normalized Laplacian ˜∆ that is closest to 1 and 0, respectively, and s ( l ) min , s ( l ) max are the squares of the minimum and maximum singular values of weight W ( l ) , respectively.
Since E XPASS reduces the input graph to its specific explanation, we argue that it can alleviate oversmoothing by reducing the information propagation along irrelevant nodes and edges.
From the perspective of Dirichlet energy, we know from [19] that, for Erd˝os-Rényi graphs, λ 0 converges to 1 as the graph becomes denser.
Oono et al.
[19] state that GNNs oversmooth on sufficiently large graphs (similar to Erd˝os-Rényi graphs).
Under this assumption, E XPASS , by definition introduces sparsity inside the ˜∆ of the input graph by using a smaller set of topK important edges for learning embeddings and, thus, reduces λ 0 to tighten the upper-bound in Equation 8.
In practice, the choice of explainer used in E XPASS can reduce λ 0 to varying degrees.
More specifically, explainers that promote sparsity would push λ 0 closer to zero and slow down the decrease of Dirichlet energy in subsequent GNN layers.
Finally, we know from Cai et al.
[34] that higher values of
Dirichlet energy per layer correspond to lower oversmoothing, we assert that E XPASS alleviates oversmoothing.
B Experiment B.1 Datasets
Mutag.
The MUTAG [36] dataset contains 188 graph molecules labeled into two different classes according to their mutagenic properties, i.e., effect on the Gram-negative bacterium S. Typhimuriuma .
Kazius et al.
[36] identifies several toxicophores - motifs in the molecular graph - that correlate with mutagenicity.
Alkane-Carbonyl.
The Alkane-Carbonyl [37] dataset contains 1,125 molecular graphs categorized into two classes where an instance in the positive group indicates a molecule that contains an unbranched alkane and a carbonyl (C=O) functional group.
DD.
The DD [38] dataset was derived from [46] and contains 1,178 protein graphs where nodes represent individual amino-acids and edges represent their spatial proximity.
The task is to predict whether a given protein is an enzyme or not.
Proteins.
The Proteins [39] dataset was derived from [46] and contains 1,113 protein graphs where nodes represent secondary structure elements and edges indicate neighborhood in the amino-acid sequence or the 3D space.
The task is to predict whether a given protein is an enzyme or not.
PubMed.
The PubMed dataset [47] is a citation network from the PubMed database, with overmillion nodes and edges respectively.
It contains a bag-of-words representation of documents and citation links between documents.
The task is to predict a node’s class among 3 classes.
B.2 Implementation details
GNN libraries and models.
All our models were implemented using PyTorch Geometric (2.1.0) and PyTorch (1.11.0).
For our experiments, we used baseline GNN architectures with three layers followed by ReLU layers and set the hidden dimensionality to 32.
Finally, we used a single linear layer to transform the graph embeddings to their respective classes.
We selected Adam as our optimizer and a weighted Cross Entropy Loss to train both vanilla and E XPASS frameworks.
All models were trained over three independent runs with a learning rate of 0.01 for 200 epochs for DD and Proteins datasets and 150 epochs for Alkane and MUTAG datasets.
E XPASS .
We define the burn-in period as a number of epochs during training in which no explanations are used.
The burn-in period is necessary to avoid feeding spurious explanations to the model since an untrained model can lead to unfaithful explanations.
The length of the burn-in period was treated as a hyperparameter and fine-tuned during the model fine-tuning phase.
After fine-tuning, we found that a burn-in period in the range [5, 15] worked best, whereas most E XPASS models outperformed their vanilla counterparts using a burn-in period of 5 and 10 epochs in our experiments.
We generated explanations for a specific percentage of correctly predicted graphs sampled in each batch and were set to 0.4 for all our experiments.
The generated explanations are normalized to [0, 1] and hard-masked over the topK most relevant edges, where topK is a percentage of the total number of edges in the input graph and was set to topK ∈ [0 . 3 , 0 . 4] for our experiments.
GNN explanation methods.
At each epoch, the model weights were frozen to generate explanations, which were calculated as the median over n independent runs of GNNExplainer, in order to obtain consistent explanations.
Then, the model weights were trained using generated explanations.
We chose the median instead of the mean to prevent the individual outliers from significantly changing the final explanations.
The number of individual runs of an explainer was treated as a hyperparameter and was set to five for GNNExplainer and one for Integrated Gradients (as it generates consistent explanations over multiple runs).
In each run, the GNNExplainer was trained for 200 epochs (150 in the case of Alkane) with a learning rate of 0.01.
All other hyperparameters of the explanation methods were set using the author’s guidelines.
Note that these multiple iterations of the explainers are not required for E XPASS to perform well when using other stable GNN explainers.
To summarize, the following parameters were treated as hyperparameters: the learning rate of the model, the learning rate of the explanation method, the number of epochs the explanation method was trained for, the number of times the GNNExplainer was computed at each epoch for each sampled graph, the percentage of correctly classified graphs that were randomly sampled to compute the explanations and the percentage of top edges/nodes that were selected as the most relevant.
On the other hand, in the case of vanilla models, the learning rate was fine-tuned during the tuning phase.
Dataset.
The train, validation, and test split was at 80%, 10%, and 10% for Alkane, Proteins, and DD following prior works.
In the case of MUTAG, no validation set was used due to the smaller dataset size, and the train and test split was at 80% and 20%.
GNN performance metric.
The GEF scores were evaluated as the mean over the individual scores of all generated explanations on the test dataset, where the explanations were hard-masked with topK = 0 .
1 for GNNExplainer, and topK = 0 .
25 for Integrated Gradients/PGMExplainer.
Note that, since Integrated Gradients/PGMExplainer generates a node mask instead of an edge mask, we required a higher topK value to generate a non-empty hard mask over the input graphs, since we retain the topK most relevant nodes in the explanation mask.
[Figure]
[Figure]
C Additional results C.1 Node classification results
We extend our proposed framework to GNN models trained on different graph downstream tasks.
In particular, we conduct additional experiments to obtain the over-smoothing and predictive behavior of E XPASS for node-level tasks.
We train five state-of-the-art GNN models and their E XPASS counterparts on the PubMed node classification dataset.
Our results show that E XPASS alleviates the over-smoothing effect in GNNs for models with higher depths (Figure 7) and achieves on-par or higher predictive performance (Table 4).
We find that, on average, E XPASS augmented GCN achieves 19.53% better over-smoothing performance for node-classification GNN models with higher depths.
C.2 Burn-in period
The burn-in period was treated as a hyperparameter and fine-tuned for each dataset and architecture.
An example of the effect of the burn-in period on the AUROC and F1-score is reported in Table 5, where the change in performance is evaluated for the Proteins dataset when using a lag of 5, 10, and 15.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
C.3 PGMExplainer results
Here, we showcase the flexibility of the E XPASS with respect to different GNN explainers.
In
Table 6, we show the effectiveness of E XPASS with respect to different explainers by further utilizing
PGMExplainer [12] as the explanation generator.
We utilize the graph explanations of PGMExplainer as node masks on the input graph (as they cannot generate edge-level masks) and incorporate them using our explanation-aware message-passing scheme with GCN as our underlying architecture.
On average, across three datasets, we find that E XPASS trained using PGM-Explainer achieves higher
GEF (+36.56%) than their vanilla counterparts (Table 6).
Also, it achieves a boost of 11.27% in
AUROC for MUTAG and a 1% improvement in the F1-score for the Alkane dataset.
We observe that PGMExplainer, similar to Integrated Gradients, produces node masks, which lack detail and do not provide finer changes to the underlying GNN model, like edge masks.
We hypothesize that this contributes to the large variation in the predictive performance across datasets.
[Table] [Table]
[Figure] [Figure]