Graph Neural Networks: Graph Classification (Part III)

When It Comes to Labeling Whole Graphs, Not Just Nodes

Many real-life situations can be modeled as graphs, but turning the relational structure of these graphs into valuable information that can help solve complex tasks is a real challenge.

Take the example of drug discovery. In the early stages of drug development, scientists need to screen large libraries often composed of hundreds of thousands of compounds (drug candidates) against targets (biological events). This requires the use of an arsenal of tools such as robotics, data processing and control software, and sensitive detectors [1]. Such approaches can be very time consuming and expensive, with a very low hit rate (a typical hit rate is less than 1% in most assays!).

And, yet, the molecules' atomic composition and arrangement can already tell us a lot about their biological behavior.

Objective

This article focuses on using graph neural networks for graph classification. It also explores explainability techniques for these models. 

Figure 5 —Mini-batching of graphs, illustration by Lina Faik, inspired by [2]

This approach is the one implemented in PyTorch. It has two main advantages:

1. It does not require changing GNN operators using a message passing scheme as, by construction, messages are not exchanged between two nodes of different graphs.
2. There is no risk of having a computational or memory overload since the adjacency matrices are saved sparsely (only the non-zero entries which correspond to the edges are kept).

Model Architecture

The model learns to classify graphs using three main steps:

1. Embed nodes using several rounds of message passing.
2. Aggregate these node embeddings into a single graph embedding (called readout layer). In the code below, the average of node embeddings is used (global mean pool).
3. Train a classifier based on graph embeddings.

Figure 6 — GNN model architecture, illustration by Lina Faik

You can find below the code used. 

What About Edge Features?

The type of bond that links two atoms in a given molecule holds valuable information about the molecule such as its stability, its reactivity, the presence of some chemically organic functional groups, etc. Therefore, including this feature has the potential to improve model performance.

But how can edge features be used when training the model?

If we take the example of GCN, it can easily be done by replacing the zeros and ones of the adjacency matrix with the edge weights, as illustrated in Figure 7. In this context, each message-passing iteration through the GCN updates the hidden embedding of nodes based on the aggregated and now weighted information of their neighborhood.

Figure 7 — Including edge weights in a GNN model via the adjacency matrix, illustration by Lina Faik

Other more sophisticated approaches to include the weights or categorical features of features exist. For more information, you can watch this short video here from DeepFindr which gives a large overview of the possibilities.

2.2. Application to the Detection of Toxic Molecules

Methodology

Target. As seen in the first section, the molecules are labeled depending on the experiment. For this reason, each experiment outcome is taken as an individual classification task.

Train / Test sets. The data is divided into two datasets: a training set, which contains around 70% of the total number of graphs, and a test set that contains the rest of the graph. This split is randomly done three times for each model.

Model evaluation. The models are trained using cross-entropy loss with class weights. They are evaluated according to the mean accuracy measured on the test sets, as well as other common classification metrics.

Hyperparameters. For each target, multiple parameters are tested such as the type and number of GNN convolutions used for node embedding (e.g., GCN, GAT, etc.), the latent dimension of node embedding, and the learning rate.

Features. Node features are all used when training the model. Concerning the edge features, the possibility of using the type of bonds (single, double, or triple) was also tested.

Results

The results show satisfactory accuracy on average. However, the accuracy depends to a large extent on the number of labeled molecules (the higher the better) and the percentage of positive outcomes (the lower the worse).

As for the hyperparameters, the best combination also differs depending on the target. Hence, it is difficult to draw a general rule.

 
Figure 8 — Accuracy of models on test set depending on the target (experiment)

Figure 9 — Accuracy of models on test set vs. % positive class, a large dot size means a high number of labeled graphs for the experiment

3. What About Explainability for GNN?

3.1. A Quick Overview of Approaches

Getting good performance is one thing, but having confidence in the prediction to take action is another. To trust the prediction of a model, one can examine the reasons why the model generated it. Sometimes, these explanations can be more important than the results themselves as they reveal the hidden patterns that the model has detected and better guide the decision-making.

For graphs, explicability is about three questions:

• Which nodes and features were relevant to making the prediction?
• How relevant were they?
• How relevant were the node and edge features of the graph?

What approaches are used to answer these questions?

The paper [5] contains a survey of current methods of GNN explicability. It classifies them in this tree which gives a good overview of the different types of approaches.

Figure 10 — Classification of the explainability approaches for graphs, illustration by Lina Faik, inspired by [5]

First, it is important to distinguish between:

• Instance-level methods, that provide explanations at the level of individual predictions
• Model-level approaches, that give explanations at the level of the whole model

Let’s explore explanations at the instance level:

• Gradient or features-based methods: They rely on the gradients or hidden feature maps to approximate input importance. Gradients-based approaches compute the gradients of target prediction with respect to input features by back-propagation whereas feature-based methods map the hidden features to the input space via interpolation to measure importance scores. In this context, larger gradients or feature values mean higher importance.
• Perturbation-based methods: They examine the variation in the model predictions with respect to different input perturbations. This is done by masking nodes or edges and observing the results for instance. Intuitively, predictions remain the same when important input information is kept.
• Decomposition methods: They decompose prediction into the input space. Layer by layer the output is transferred back until the input layer is reached. The values then indicate which of the inputs had the highest importance on the outputs.
• Surrogate: Train a simple and interpretable surrogate model to approximate the predictions of the model in the neighboring area of the input.

3.2. Application of GNNExplainer to the Use Case

GNNExplainer falls into the category of perturbation-based methods.

Without going into technical details, it basically consists of learning soft masks for edges and node features. To do so, it starts by randomly initializing soft masks and combining them with the original graph via element-wise multiplications. Then, the masks are optimized by maximizing the mutual information between the predictions of the original graph and the predictions of the newly obtained graph.

More information is available in the original paper [6].

Figure 11 shows the explanations obtained for a molecule of the test set. The more the links between the atoms tend towards the red, the more they played an important role in the prediction.

Figure 11 — Explanations obtained for a given prediction using GNNExplainer

Strong knowledge of organic chemistry is needed to investigate the consistency of the results. However, one might note that the presence of SO2 was detected as important. This is the formula of sulfur dioxide which is a chemical compound known to be toxic.

4. Key Takeaways

✔ Graph Neural Networks, GNNs, can be used to classify entire graphs. The idea is similar to node classification or link prediction: learning an embedding of graphs (instead of nodes) using the structural properties of these graphs.

✔ When it comes to understanding the outcome of a model for a given instance, many approaches exist. They can rely on the gradient or features, use perturbation techniques, decompose the outcome, or use surrogate models.

✔ ️The real-life applications are multiple. This article was about the detection of the toxicity of molecules. Predicting other molecular properties lends itself well to this graph-based approach.

✔ ️However, biology is not the only industry. It can be used for instance in the retail industry: this article explains how GNNs can be applied to generate recommendations based on the graph of users’ sessions.

References