Graph Neural Networks: GCN, MPNN, and GAT, Explained Simply

Graph Neural Networks: GCN, MPNN, and GAT, Explained Simply A visual guide to how graph neural networks work under the hood Introduction Neural networks are an incredible innovation. Since a long period of time and up until now, they have been used as a key component in solving complex AI problems. Under the hood, neural networks learn a sophisticated mathematical function that transforms input data into a desired target. However, by default, normal neural networks do not use any knowledge about the relationship between the parts of the input data. For instance, to process images, convolutions are commonly used as a way to combine each pixel with its neighbouring pixels, because they are related to each other. Otherwise, a neural network would not know if a pixel at position N is related to a pixel at position N + 1. This extra context can improve the performance of a model. The same is true for graphs, which represent a set of objects along with the relationships between them. There are many objects that can be represented by graphs, such as molecules, social networks, players during a soccer match, traffic, or metro maps. Graphs can contain valuable context and it is important to understand how one can exploit their full potential. For that reason, there exist graph neural networks (GNN) that, as the name suggests, apply neural networks to graph structures. Applications A great thing about GNNs is that once trained, they can be applied to new graphs with other structures. For example, if a GNN is trained on molecules of certain types, we can still use that GNN to perform a classification task by giving it a molecule whose graph contains a completely new, unseen structure. That is how, for instance, there has been using a popular use-case of GNN consisted of training a model for antibiotic discovery. Apart from it, a GNN can also be used to classify individual nodes or edges. GNN's output can also be used to classify a graph as a whole. Graph Convolutional Networks (GCN) Concept Let's go back to convolutions. As we know, they take a pixel and its neighbourhood as input, and combine them to produce a new value for the pixel. This approach assumes there is a relationship between adjacent pixels and allows the model to take into account the local context around the pixel. We can naturally apply this idea to graphs: by picking up a node with its adjacent nodes, our method will combine them, and produce a new node with new features. The described approach is presented in the section "Update rule". In addition, what makes this idea interesting is that graphs can be seen as a generalization of images. In fact, each pixel in an image is connected to up to 4 adjacent pixels. There, there are common semantic similarities in convolution processes in both cases. Layers In general, a GNN contains a small number of layers (usually between 2 and 4). A higher number of layers is usually avoided, as it might cause an oversmoothing problem, which is described later in this article. Each layer transforms a feature vector from the previous layer using aggregation functions applied to it and its neighbours. This process is applied in parallel to each node independently, and the resulting feature vectors might have a different shape than the one from the previous layer. As a result, the shape of the feature vectors from the last GNN layer can differ from the input shape on the first layer. Update rule To describe the update rule, we would need three matrices: A - adjacency matrix (A[i][j] = A[j][i] = 1 if vertices i and j are connected, and A[i][j] = A[j][i] = 0 otherwise). H - feature matrix. The i-th row of the matrix represents a feature vector of the i-th node. W - learnable linear transformation used by the GNN. This matrix is shared across all nodes of the graph. By multiplying A by H, we get a neighbour-feature sum matrix. In other words, for each node in A, AH sums the feature values defined in H only for the nodes that are adjacent to it. For non-adjacent nodes, the feature value is ignored (multiplied by 0). Let's have a look at the example below. By taking the result of AH, we can then multiply it by the matrix W which is learned by a neural network. As a last step, we apply a non-linear transformation σ. As a result, the update rule can be written as: For the non-linear function σ, ReLU or LeakyReLY is usually chosen in GNN. Given that matrix multiplication is associative, for optimization purposes, specifically to reduce computational cost, when calculating AHW, HW is computed first and then multiplied by A on the left side. However, there are several issues with the current approach that we need to address in the next sections. Central node First of all, during the computation done for each node, it does not take into account any information about the node itself. For instance, we can clearly see that when we obtained the element (AH)[1][1] for the first node, the feature value corresponding to that node (3) was multiplied by zero, because in the adjacency matrix we had A[1][1] = 0. This problem can be easily solved by adding ones to the diagonal elements of A: Given that, the update formula becomes: Feature normalization Secondly, by performing matrix multiplication, the scale of features changes. To fix this, a normalization is performed using the degree matrix D obtained from A, where D[i][i] equals the number of neighbours of node i (including itself), while D[i][j] = 0 for i ≠ j. For example, for the graph in the example above, the matrix D would have had the following form: The update rule becomes: This formula can also be rewritten in node-wise level (which is also called mean-pooling update formula): Symmetric normalization Another popular way to fix the scale in GCN is to use symmetric normalization (Kipf & Welling, ICLR 2017), where the inverse square root of D is applied on both sides of Ā: Or, on the node level, the formula can be rewritten as follows: Training & Inference A great thing about GNNs is that they can generalize to new graph structures. The training logic is not applied only to the graph that was used for training. GNNs learn transformations that are applied individually to nodes, regardless of how many nodes or edges the graph has. All they need is a learned, shared matrix W that transforms the feature vector of any node across layers. For example, this idea is very different from fully connected neural networks, where the number of weights is tied to the input size. Nevertheless, it is important to understand that GNN inference on a new graph usually works well when its structure is still similar to the original graph the GNN was trained on. If a new graph during inference is completely different from the original graph, the performance might become worse. Speaking of training, backpropagation in GNNs works in a similar way to normal neural networks. A GNN can be trained either on a single large graph or on multiple graphs at the same time. Typically, when a GNN is trained on multiple graphs, it generalizes better to new graphs. It is also important to know that a GNN produces node embeddings, which are then usually passed to a separate, smaller model to perform a downstream task (for example, node, edge, or graph classification). In this setup, the GNN acts as an intermediate feature extractor, and the labels used to compute the loss value, and thus to train the GNN, come from the downstream task. However, there are rare cases where this is not true, and the GNN can directly produce the final predictions in the system. As mentioned before, the dimension of feature vectors at each layer of a GNN can differ across layers, and it is one of the main hyperparameters of a GNN. Advantages Like CNNs, GCNs successfully use the local context around a given node, which boosts the overall model's performance. Apart from that, a nice property of GCNs is that their computations are linear with respect to the graph size (O(|V| + |E|)). Because the weight matrix W is shared across graph nodes, the number of parameters of convolutions does not depend on the input graph size. For a particular graph structure, GCNs treat nodes with different importance based on their adjacency to other nodes. With all the advantages that GCN can offer, let's now have a look at two more advanced graph networks that go even further to reach the maximum potential of GNNs. Message Passing Neural Networks (MPNN) We have just seen how GCN uses information about the graph structure. However, it mostly operates only on node features. We can go one step further and also make it possible to operate on graph edges. For that, we can introduce the concept of message passing, which we will use during the aggregation process. A message is an abstract concept describing a value that flows along an edge during computation. More concretely, let's imagine a pair of connected nodes i and j, connected via an edge e[i][j]. A message sent from node i to j can be described mathematically as the following function (fₑ is called a message function): The next step consists of aggregating all messages entering a given node (fᵥ is called a readout function): Below we can see a visualisation of the process showing how the message function fₑ and the readout function fᵥ combine nodes and edges to get the next graph state: On one side, MPNNs are powerful but require a lot of computation and memory. In practice, they are usually used with small graphs. In practice, fₑ and fᵥ are usually small MLP (multi-layer perceptrons). Graph Attention Networks (GAT) GAT is a generalization of GCN. They work in the same way as GCN, except that instead of using raw values of node degrees in the computations, the network learns importance values by itself. That is why the concept is called attention, similar to what is done in Transformers, which can decide the importance of pairwise elements in a given input sequence by themselves. By modifying the original update formul