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

IntroductionNeural 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.ApplicationsA 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.Diagram showing the input graph state G(X, A) and output graph state G(H, A) produced by a GNN, where A is the adjacency matrix, and X and H are embedding matrices whose i-th rows, x[i] and h[i], correspond to the i-th node's feature vector. The GNN preserves the original graph structure but transforms the node states from x[i] to h[i]. The resulting graph G(H, A) can then be used for downstream tasks such as node classification, edge classification, or graph classification, and can even generalize to graphs with structures different from those seen during training.Graph Convolutional Networks (GCN)ConceptLet'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.LayersIn 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.An example of a node h[i] aggregating information via convolution from its adjacent nodes h[1], h[2], and h[3] to produce the next layer's representation h'. Each layer preserves the same node and edge structure as the previous layer, with updated vectors h[i]. GNNs typically contain between 2 and 4 layers. Backpropagation flows in the opposite direction, from the last layer to the first.Update ruleTo 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.The product of matrices A and H, computed for the graph shown on the left. We can clearly see that the zero values in the adjacency matrix cause the matrix multiplication to ignore the features of non-adjacent nodes (shown in red).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:Update rule: A is the adjacency matrix, H is the feature value matrix, W is the learnable shareable matrix, and H' is the updated feature value matrix at the next layer. σ is a non-linear function.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 nodeFirst 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:Adding the identity matrix to A accounts for the central node itself during the update.Given that, the update formula becomes:Updated ruleFeature normalizationSecondly, 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:Degree matrix D depicted for the graph on the left. D[i][i] contains the number of adjacent nodes for node i including itself.The update rule becomes:Update formula including feature normalizationThis formula can also be rewritten in node-wise level (which is also called mean-pooling update formula):Mean-pooling update formula (node-wise)Symmetric normalizationAnother 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 Ā:Update formula including symmetric normalizationOr, on the node level, the formula can be rewritten as follows:Node-wise update formula including symmetric normalizationTraining & InferenceA 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.AdvantagesLike 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):Message passing formula, including the aggregation of feature vectors of nodes i and j along with the edge feature vector e[i][j] between them. m[i][j] is a vector, not a scalar.The next step consists of aggregating all messages entering a given node (fᵥ is called a readout function):The updated feature value at the next layer is obtained by aggregating the current node's feature vector h[i] with another aggregation function applied to all the adjacent messages m[j][i] flowing into node i.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:Computations visualized for graph node h[1]. First, the information about h[1] itself, along with its adjacent edges and nodes, is passed into a message function fₑ. This produces message vectors m[2][1] and m[3][1] (each corresponding to an edge e[2][1] and e[3][1] respectively). These messages are then combined with the initial vector h[1] in a readout function, which finally produces the feature vector h'[1] at the next layer.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 formula from GCN with attention weights, the update formula now becomes:Node-wise update formula including attention weights α[i][j] instead of fixed values defined by the adjacency matrix A. a[i][j] is a scalar value, not vector.The learned weight α[i][j] can be literally interpreted as how important node i is to node j. In comparison to GCN, where the coefficients aij were explicitly defined as 1 / √(|Ni| ⋅ |Nj|).Each attention weight α[i][j] can be calculated using an attention function a, which takes as input information about h[i], h[j], and e[i][j]. The attention weights are then normalized using a softmax function.An advantage of GATs is that they require less memory, because the learned coefficients α[i][j] are simply scalar values for each edge, while in MPNNs, the computed messages were learned vectors for each edge.Similar to Transformers, GATs often use multiple heads to capture more signals and further boost model performance.Diagram showing the use of multi-head attention. The three arrows flowing between nodes represent separate heads capturing various signals in the network. Image adapted by the author. Source: Graph Attention Networks | Petar Veličković, Guillem Cucurul | ICLR 2018OversmoothingOversmoothing is a problem where, when there are too many stacked layers in GNNs, node feature representations become nearly identical. This tends to happen in deep GNNs, because with repeated aggregation functions (for example, taking an average), each node gradually absorbs more and more information from its neighbours and converges towards them.Common techniques to reduce oversmoothing include adding skip connections, where a node feature vector is passed directly to the next layer, or edge dropping, where, in a similar way to the dropout technique, randomly chosen edges are removed during training to reduce information overhead.Oversmoothing is one of the reasons why, in practice, GNNs usually have a small number of layers (for example, 2 to 4).ConclusionIn this article, we've seen how GNNs work under the hood and explored the main architectures. As it turns out, there's nothing especially complicated about them: they operate just like standard neural networks (including backpropagation), except for how the convolution operation is redefined.GNNs are particularly well suited to problems involving graph-structured data. By learning the linear transformation W (and, in the case of GAT, attention weights as well), they can automatically identify the most important relationships within the graph. Based on the problem requirements, graph size, and desired complexity, any of these options, GCN, MPNN, GAT, or another variation, can be chosen.ResourcesAll images unless otherwise noted are by the author