Graph Convolutional Networks (GCN): All You Need to Know & Code Implementation

Intro

Machine learning models for images, audio, and tabulated data are well known and have been around for quite some time, but what about graph-structured data like social interactions or molecules? This task is more recent. To deal with it Graph Neural Networks (GNN) were developed. In this article, I am going to explain how one of the simplest GNN models — Graph Convolutional Network (GCN) — works. I will talk about both the intuition behind it with simple examples and rigorous formulas and show you how to use it on a real example of organic-chemical data. Let’s dive into it!

Theory

Shortly about graph representation

A graph is often stored in an adjacency matrix (denoted as A), where each row corresponds to a node in a graph. If a particular node is connected to another node then it has one on the corresponding entry in the adjacency matrix, otherwise the corresponding entry is zero. Consider the following graph with its adjacency matrix.

A graph and its adjacency matrix. In salmon, I marked the node number corresponding to rows/columns

It is evident that one cannot feed this adjacency matrix into a usual MLP hoping to get some meaningful result, at least because the same graph has many different representation in terms of adjacency matrices depending on the ordering of the nodes:

Example from Sanchez-Lengeling, Benjamin, et al. “A Gentle Introduction to Graph Neural Networks.” Distill, 2 Sept. 2021, https://distill.pub/2021/gnn-intro/

So the AI model which processes graph-structured data must give the same output regardless of the ordering of the nodes. To tackle this challenge message passing was created.

Message Passing

The idea behind message passing is simple. Each node has some characteristics (features) which are stored in the form of vectors. For each node, we simultaneously generate a so-called message based on the features of its direct neighbors via some permutation invariant function (e.g. summation). Then we update the feature vectors of each node based on its message via another function (it may be MLP or simply summation). This way each node accumulates some information about its neighbors. After two iterations of such message passing the nodes will contain the information about the neighbors of their neighbors and so on. To better understand it consider the following example with the graph we just defined (in the example we focus specifically on node 3):

Example of message passing for node 3

I showed in the picture how features of node 3 can be updated if we use simple summation for calculating a message and summation for updating the node’s feature vector (also called the state of the node).

Mathematically we can rewrite it in the following way:

m is the message vector, M_t is the function producing the message, e is the edge features (in the example we did not have it), h_v is the state of the node, U is the function for state update, N(v) is the set of neighbors of v

We considered the most simple option, where U is a summation and M is the identity for h_w. So we just added up the feature vectors of the neighbors to the feature vector of the node. But then we can rewrite it in a matrix form:

Here H_l is the matrix of node features on the l-th iteration; A tilde is the adjacency matrix where each node is considered connected to itself (i.e. self-loops are induced). To better grasp it consider our previous example in the matrix form:

Message passing step in a matrix form. The left-hand matrix is the Adjacency matrix with self-loops, it is multiplied by the feature matrix, and the result is the new feature matrix

Surprisingly, this is exactly the idea behind Graph Convolutional Networks. Additionally, we would want to multiply the result by a matrix of learnable weights. Also, you might have noticed that the entries of the feature matrix increase after such message passing. To prevent them from blowing up, we scale down the result by multiplying the adjacency matrix by the squared inverse of degree matrix D. The definition of D as well as A tilde and square inverse is below:

Putting everything together we get the original GCN initially described in the original research paper (Kipf, 2017):

GCN layer is a matrix form

That’s it! While the formula might seem intimidating, the idea is a simple exchange of states between neighboring nodes. Let’s now practice with some code implementation of GCN.

Practice

I prepared a Google Collab with the code, feel free to follow via this link:

[embed]Google Colab Edit descriptioncolab.research.google.com

For GCN implementation I will use the Torch Geometric framework which has everything you need to work with graph-structured data and which works similarly to the usual Pytorch. I selected the ZINC dataset consisting of small organic molecules, which are naturally stored in terms of a graph. The target variable is penalized logP (also called constrained solubility). That means that the task is graph regression. First, we import the dataset and create data loaders:

from torch_geometric.datasets import ZINC
import torch
import torch_geometric
from torch import nn
from torch_geometric.loader import DataLoader

train_dataset = ZINC('datasets/ZINC', subset=True, split='train')
val_dataset = ZINC('datasets/ZINC', subset=True, split='val')
test_dataset = ZINC('datasets/ZINC', subset=True, split='test')

train_loader = DataLoader(train_dataset, 16, shuffle=True)
val_loader = DataLoader(val_dataset, 16, shuffle=True)
test_loader = DataLoader(test_dataset, 16, shuffle=False)

Second, we define the GCN model. Take a look at this code:

from torch_geometric.nn import GCNConv
from torch_geometric.nn import global_add_pool
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class GCN(nn.Module):
    def __init__(self, num_node_features, hidden_dim, output_dimension):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(num_node_features, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, hidden_dim)
        self.conv3 = GCNConv(hidden_dim, hidden_dim)
        self.readout = nn.Linear(hidden_dim, output_dimension)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index)
        x = self.conv2(x, edge_index)
        x = self.conv3(x, edge_index)
        x = global_add_pool(x, batch)
        x = self.readout(x)
        return x

Here are three main things to notice about this code snippet:

• troch geometrics impelments the GCN layer (torch_geometric.nn.GCNConv), so we can simply use it. I inserte dthree consecutive GCN layers.
• Since the task is of graph layer, at the final stage we accumulate the updated states of each node using global-add-pooling (torch_geometric.nn.global_add_pool)
• As the final layer, I inserted a simple linear layer, which serves as a readout layer, i.e. based on the aggregated node features it makes the predictions.

And now training function. It is a basic training function one uses with PyTorch. I selected MSE for both the metrics and loss function. As an optimizer I use Adam. Take a look at it before I make any comments:

def train(model, max_epochs, train_loader, val_loader):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    metrics = nn.MSELoss()
    loss = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    for epoch in range(max_epochs):
        model.train()
        train_loss = 0

        for data in train_loader:
            data = data.to(device)
            data.x = data.x.float()
            optimizer.zero_grad()

            pred = model(data.x, data.edge_index, data.batch)
            pred = pred.squeeze(1)
            l = loss(pred, data.y)

            l.backward()
            optimizer.step()
            train_loss += l.item()

        train_loss = train_loss / len(train_loader)

        model.eval()
        val_loss = 0

        for data in val_loader:
            with torch.no_grad():
                data = data.to(device)
                data.x = data.x.float()

                pred = model(data.x, data.edge_index, data.batch)
                pred = pred.squeeze(1)
                metric = metrics(pred, data.y)

                val_loss += metric.item()

        val_loss = val_loss / len(val_loader)
        print(f"Epoch: {epoch}, training loss: {train_loss}, validation MSE: {val_loss}")

Please note, that unlike with usual MLP models Graph Neural Network srequire not only a feature matrix (data.x in our case) but also a matrix showing graph connectivity (data.edge_index). Moreover, very often graphs have edge features along with node features.

NB: In fact, the ZINC dataset also has node features indicating the number of bonds between the atoms, but GCN does not work with edge features, so I ignored them. In the upcoming articles I will review the models, which do process these edge features, so we can hope for a better result!

Finally, we can train our model.

model = GCN(num_node_features=1, hidden_dim=16, output_dimension=1).to(device=device)
train(model=model, max_epochs=20, train_loader=train_loader, val_loader=val_loader)

Notice, that the parameters except for hidden_dim were NOT selected arbitrarily. num_node_features is the number of features corresponding to each node in a graph. output_dimension is the number of features we predict for each graph. Since the graphs in this dataset are not large, we can even check each of their components. Consider the following code:

node_features = train_dataset[0].x
edge_index = train_dataset[0].edge_index
edge_attr = train_dataset[0].edge_attr
y = train_dataset[0].y
print("Node features: ", node_features)
print("Edge indices", edge_index)
print("Edge features", edge_attr)
print("Target: ", y)In the output we can see 4 components of each graph:

When you run this code and get the output note the following:

• Node features. Recall that in this dataset each node has only one feature. You may see it exactly from the node features output.
• Edge indices. It is very important to notice that to spare the memory graphs are not stored in adjacency matrices but in COO (coordinate) format. That is a format of a matrix of the shape (2, number_of_edges). The first row contains source nodes, the second contains corresponding target nodes.
• Edge features. The features of each edge. Exactly the thing that we ignored in this article.
• Target. The thing we wanted to predict. In our case it is a scalar, for this reason, we put output_dimension=1.

If you trained the model, you would see an MSE of approximately 3.0 — 3.5. Spoiler: it is not a good result, but to make sure, you can make 10 inferences on the training set and compare them to actual values:

for i in range(10):
  pred = model(test_dataset[i].x.float().to(device), test_dataset[i].edge_index.to(device), torch.zeros_like(test_dataset[i].x.squeeze(1), dtype=torch.int64).to(device))
  print(f"Predicted: {pred.item()}; Actual: {test_dataset[i].y.item()}")

You would get output similar to this:

Summary

So, is out model really garbage? Probably with this simple setup and the fact that we ignored the edge features, yes. I did not want to oversimplify things and give unrealistically small datasets, thus I gave a real-world dataset on which researchers are working even nowadays. But GCN remains a classical model, which is excellent to start studying GNN with. In the next articles, I will show you more sophisticated models such as GIN, which hopefully can get a better result! So make sure to subscribe!

P.S. Love data, science, and data science!

Sources

[embed]Semi-Supervised Classification with Graph Convolutional Networks We present a scalable approach for semi-supervised learning on graph-structured data that is based on an efficient…arxiv.org

[embed]A Gentle Introduction to Graph Neural Networks What components are needed for building learning algorithms that leverage the structure and properties of graphs?distill.pub

[embed]torch_geometric.nn.conv.GCNConv - pytorch_geometric documentation cached (, optional) - If set to , the layer will cache the computation of (\mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}}…pytorch-geometric.readthedocs.io