How to Explain Graph Neural Network — GNNExplainer | by Shanon Hong | Towards Data Science

Graph Neural Network(GNN) is a type of neural network that can be directly applied to graph-structured data. My previous post gave a brief introduction on GNN. Readers may be directed to this post for more details.
Many research works have shown GNN’s power for understanding graphs, but the way how and why GNN works still remains a mystery for most people. Unlike CNN, where we can extract activation of each layer to visualize the decisions of the network, in GNN it is hard to get a meaningful explanation of what features the network has learnt. Why does GNN determine a node is class A instead of class B? Why does GNN determine a graph is a chemical or molecule? It seems like GNN sees some useful structural information and determines are made upon these observations. But now the problem is, what observations does GNN see?
GNNExplainer is introduced in this paper.
Briefly speaking, it is trying to build a network to learn what a GNN has learnt.
The main principle of GNNExplainer is by reducing redundant information in a graph which does not directly impact the decisions. To explain a graph, we want to know what are the crucial features or structures in the graph that affect the decisions of a neural network. If a feature is important, then the prediction should be altered largely by removing or replacing this feature with something else. On the other hand, if removing or altering a feature does not affect the prediction outcome, the feature is considered not essential and thus should not be included in the explanation for a graph.
The primary objective for GNNExplainer is to generate a minimal graph that explains the decision for a node or a graph. To achieve this goal, the problem can be defined as finding a subgraph in the computation graph, that minimizes the difference in the prediction scores using the whole computation graph and the minimal graph. In the paper, this process is formulated as maximizing the mutual information(MI) between the minimal graph Gs and the computation graph G:
Besides, there is a secondary objective: the graph needs to be minimal. Though it was also mentioned in the first objective, we need to have a method to formulate this objective as well. The paper addresses it by adding a loss for the number of edges. Therefore, the loss for GNNExplainer is literally the combination of prediction loss and edge size loss.
Three types of explanations are discussed in the paper: the explanation for one node, the explanation for one class of nodes, and the explanation for a graph. The major difference is in the computation graphs.
For an explanation on a single node, the computation graph is its k-hops neighbours, where k is the number of convolutions in the model.
For an explanation on a class of nodes, the paper suggests selecting a reference node and using the same method to compute explanation. The reference node can be selected by taking the node whose features are closest to the mean features of all other nodes with the same class.
For an explanation of the entire graph, the computation graph becomes the union of the computation graphs for all nodes in the graph. That makes the computation graph equivalent to the entire input graph.
The learning of the minimal graph Gs is by learning a mask for edges and a mask for features. That is, for each edge in the computation graph, there exists a value in the edge_mask which determines the importance of the edge. Similarly, for each feature in the node feature, the feature_mask determines if the feature is important to the final decision.
With all those concepts, we can summarize everything for a GNNExplainer:
We need to extract the computation graph, which is the k-hops neighbour for node classification, or the entire graph for graph classification.Initialize an edge_mask for each edge in the computation graph, and a feature mask for each feature dimension.Construct a neural network that learns the edge_mask and feature_mask with loss described above.Use edge_mask and feature_mask to reduce the computation graph to a minimal graph.
We need to extract the computation graph, which is the k-hops neighbour for node classification, or the entire graph for graph classification.
Initialize an edge_mask for each edge in the computation graph, and a feature mask for each feature dimension.
Construct a neural network that learns the edge_mask and feature_mask with loss described above.
Use edge_mask and feature_mask to reduce the computation graph to a minimal graph.
That’s about everything we need to know before we can implement a GNNExplainer. To summarize up, we are trying to learn an edge_mask and a node_feature_mask which remove some edges and features from the computation graph while minimizing the difference in prediction score, the resultant graph is a minimal graph that explains the decision of the node or graph.
I’m going to implement this in Pytorch Geometric(PyG). One great advantage of PyG is that it updates very frequently and has many implementations of current models. Surprisingly I found GNNExplainer is already implemented in PyG library, which saves me a lot of time. Although it only works for node explanations, thanks to its open-source, it is not hard to change it to work for graph explanations as well.
To start, we first need to install PyG. GNNExplainer is not yet in their current release (PyG 1.4.4) but the codes are already released their Github. So to get GNNExplainer, you must clone from their Github repository and install from there.
The example code is provided on the PyG website. It is easy to follow so I’m not going to show the code in this post. But the implementation details are something we want to check out and use it for graph classification thereafter.
I’m going to trace the code based on my short summary above. The example code passes node index, together with the full feature matrix and edge list to a GNNExplainer module.
explainer = GNNExplainer(model, epochs=200)node_idx = 10node_feat_mask, edge_mask = explainer.explain_node(node_idx, x, edge_index)
What’s happening in GNNExplainer is exactly what we discussed in the previous section.
Extract the computation graph
Extract the computation graph
To Explain a node, we first need to get its k-hop computation graph. This is done by __subgraph__() method in PyG.
x, edge_index, hard_edge_mask, kwargs = self.__subgraph__(            node_idx, x, edge_index, **kwargs)
The hard_edge_mask removes all other edges outside the k-hop neighbourhood.
2. Masks are initialized by __set_mask__() method and applied to every layer of the network.
self.__set_masks__(x, edge_index)         def __set_masks__(self, x, edge_index, init="normal"):             (N, F), E = x.size(), edge_index.size(1)              std = 0.1             self.node_feat_mask = torch.nn.Parameter(torch.randn(F) * 0.1)                    std = torch.nn.init.calculate_gain('relu') * sqrt(2.0 / (2 * N))               self.edge_mask = torch.nn.Parameter(torch.randn(E) * std)                     for module in self.model.modules():                     if isinstance(module, MessagePassing):                                      module.__explain__ = True            module.__edge_mask__ = self.edge_mask
3. Initial prediction is performed using the trained model, then the prediction is used as a label to train the GNNExplainer.
# Get the initial prediction.         with torch.no_grad():                 log_logits = self.model(x=x, edge_index=edge_index, **kwargs)     pred_label = log_logits.argmax(dim=-1)          # Train GNNExplainerfor epoch in range(1, self.epochs + 1):                      optimizer.zero_grad()                 h = x * self.node_feat_mask.view(1, -1).sigmoid()                 log_logits = self.model(x=h, edge_index=edge_index, **kwargs)                  loss = self.__loss__(0, log_logits, pred_label)                 loss.backward()                 optimizer.step()
4. The loss is defined
def __loss__(self, node_idx, log_logits, pred_label):               loss = -log_logits[node_idx, pred_label[node_idx]]                m = self.edge_mask.sigmoid()               loss = loss + self.coeffs['edge_size'] * m.sum()               ent = -m * torch.log(m + EPS) - (1 - m) * torch.log(1 - m + EPS)               loss = loss + self.coeffs['edge_ent'] * ent.mean()                m = self.node_feat_mask.sigmoid()               loss = loss + self.coeffs['node_feat_size'] * m.sum()               ent = -m * torch.log(m + EPS) - (1 - m) * torch.log(1 - m + EPS)               loss = loss + self.coeffs['node_feat_ent'] * ent.mean()          return loss
The current implementation is PyG is only for node explanation. But having understanding the principles behind, it isn’t hard to rewrite the function for graph explanation.
Only A few functions we need to replace: 1) we need to replace __subgraph__ function to obtain the computation graph for the entire graph. 2) we need to set masks for the entire graph. 3) We need to change the loss function to compute loss for graphs.
Complete code implementation is available at this Github link.
GNNExplainer provides a framework to visualize what a GNN model has learnt. However, the actual explanation result may not be good enough to explain a huge graph as search space for the optimal explanation is exponentially larger than a smaller one. Instead of fitting a neural network, other search techniques may also be applied to find the optimal explanation borrowing the same concepts and the performance is yet to be proved.
GNNExplainer: Generating Explanations for Graph Neural Networks, https://arxiv.org/abs/1903.03894
Pytorch Geometric, https://pytorch-geometric.readthedocs.io/en/latest/