Graph embedding refers to learning low-dimensional representations of nodes in a network; the representation encodes structural information and node features. Such a representation can then be used for various downstream machine learning tasks like node classification, link prediction, visualization, etc. The downstream tasks often work with dynamic or evolving networks such as social networks, recommendation networks etc. The majority of the graph embedding methods assume fixed graph structure and thus are unsuitable for dynamic networks. The GraphSAGE embedding method overcomes this limitation by incorporating two changes, sampling and aggregating features, in the graph convolutional networks (GCNs) that are used for fixed structure networks. These changes, explained below, make GraphSAGE not only computationally efficient to scale to very large graphs but also permit embeddings to be generated for those nodes of a graph not seen before. This inductive capability to work with unseen data makes GraphSAGE a highly valuable node embedding tool suitable for a large number of applications. In this blog post, I will highlight the aggregation and sampling used in GraphSAGE and provide an example of its usage. Before reading on, you may want to go over my post on graph convolutional networks.
How Does GraphSAGE Perform Aggregation?
We will assume each node has a feature vector associated with it. The associated feature vector with a node, for example, may be the profile of the person represented by the node. In absence of explicit features, structural features of nodes such as the node degree are used in GraphSAGE.
The key component of GraphSAGE algorithm is that embedding of a node is obtained by aggregating embeddings in a layer-wise manner by slowly expanding the local neighborhood of the node. At layer-0, the embedding of a node is simply its feature vector. Thus, the embedding of the target node A at layer-0 in the graph shown below is $\bf X_A$.
Figure 1. Input Graph for Embedding |
The layer-1 embedding of a node is the aggregate embedding of all of its neighboring nodes away at one link. Similarly, the layer-2 embedding of a node is obtained by aggregating embeddings of its 2-neighborhood nodes. Figure 2 below shows this layer-wise embedding process for the target node A of the input graph of Figure 1.
For k = 2, we get the embedding vector $\bold{Z}_A for the target node in the above graph. The matrices W and B are the trainable matrices. These trainable matrices are learned by defining a suitable loss function. Both supervised and unsupervised learnings are possible.
While the above description for aggregation appears similar to that used in graph convolutional networks, the difference is that the aggregation function is learned during training in GraphSAGE and it is predefined and fixed in GCNs. It is this difference that makes GraphSAGE an inductive learner as opposed to GCNs being transductive learners.
Neighborhood Sampling in GraphSAGE
In graph convolution networks, every neighboring node of the target node at the specified neighborhood size participates in message passing and contributes towards the computation of the embedded vector of the target node. When the neighborhood size is enlarged, the number of nodes contributing to the embedded vector computation can grow exponentially for certain graphs. This problem gets exacerbated when the neighborhood of a target node includes a hub or celebrity node having millions of connections. To avoid exponential computation growth, GraphSAGE algorithm randomly selects only a sample of neighboring nodes. This allows GraphSAGE to be used for extremely large graphs with billions of nodes. You can see an illustration of neighborhood sampling in Figure 3 below where the No icon shows the nodes not being sampled.
Figure 3. Illustration of Neighborhood Sampling |
Now, let's use GraphSage to generate node embeddings for link prediction.
Applying GraphSAGE to Perform Link Prediction
In link prediction, the objective is to predict whether a link exists between a pair of nodes. Link prediction has many applications; for example, it is used for recommending friends on social media networks.
We will use the Cora dataset, widely used in graph machine learning similar to the usage of MNIST dataset in deep learning. The Cora dataset consists of 2708 scientific publications classified into one of seven classes. Each publication is a node in the network, and is described by a 0/1-valued word vector of 1433 dimensions. The binary vector values represent the absence/presence of the corresponding word from the dictionary of 1433 unique words. The citation network has 10566 edges.
The implementation will be done using the Deep Graph Learning (DGL) library, built for easy implementation of graph neural network models on top of existing deep learning frameworks such as PyTorch, MXNet and TensorFlow. Each node of a network is represented by a unique integer, called its node ID, in DGL. An edge is represented by a pair of integers corresponding to the IDs of its end nodes. You will need to install DGL before proceeding any further.
We begin by importing all necessary libraries. We also specify PyTorch as the backend for DGL.
import itertoolsimport osos.environ["DGLBACKEND"] = "pytorch"import dglimport dgl.dataimport numpy as npimport scipy.sparse as spimport torchimport torch.nn as nnimport torch.nn.functional as F%matplotlib inline
We are going to treat the link prediction problem as a binary classification problem. When a pair of nodes is connected, there is an edge linking them. We treat this as an example of positive class; the absence of a link between a node pair is treated as negative class example. A sample of positive and negative class examples will be extracted from Cora dataset with the rest of the data being treated as test data. We begin by loading the data. We pick 10% of the edges as positive class examples and the same number of edges as negative examples.
dataset = dgl.data.CoraGraphDataset()g = dataset[0]# Split edge set for training and testingu, v = g.edges()
eids = np.arange(g.num_edges())eids = np.random.permutation(eids)test_size = int(len(eids) * 0.1)train_size = g.num_edges() - test_sizetest_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]]train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]
# Find all negative edges and split them for training and testingadj = sp.coo_matrix((np.ones(len(u)), (u.numpy(), v.numpy())))adj_neg = 1 - adj.todense() - np.eye(g.num_nodes())neg_u, neg_v = np.where(adj_neg != 0)
neg_eids = np.random.choice(len(neg_u), g.num_edges())test_neg_u, test_neg_v = ( neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]],)train_neg_u, train_neg_v = ( neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]],)
Before training, we need to convert our original graph to a subgraph without test nodes.
Now, the question is how do we perform link prediction using the embeddings that will be generated by our GraphSAGE model. Remember that link prediction operates on node pairs. Since we can describe a pair of nodes with an edge connecting them, we can set up our training and test data based on edges. To do so, we create four graphs: positive train graph, positive test graph, negative train graph, and negative test graph. All these graphs have same number of nodes as in the original data; however, the edges in each graph correspond to node pairs of positive and negative training and test examples. In DGL, we can do this by the following code.
One way to predict the likelihood of a link between a pair of nodes is to compute the similarity of the respective node vectors. We will use the dot product to compute node similarities. The corresponding DGL code for this is as follows.
In epoch 0, loss: 0.7079574465751648 In epoch 5, loss: 0.6902216076850891 In epoch 10, loss: 0.6735929250717163 In epoch 15, loss: 0.6262624263763428 In epoch 20, loss: 0.5796554684638977 In epoch 25, loss: 0.555548369884491 In epoch 30, loss: 0.5190547108650208 In epoch 35, loss: 0.5012964606285095 In epoch 40, loss: 0.48360729217529297 In epoch 45, loss: 0.46244215965270996 In epoch 50, loss: 0.4450105130672455 In epoch 55, loss: 0.4255025088787079 In epoch 60, loss: 0.41155683994293213 In epoch 65, loss: 0.396867573261261 In epoch 70, loss: 0.3816315531730652 In epoch 75, loss: 0.36709561944007874 In epoch 80, loss: 0.3521302044391632 In epoch 85, loss: 0.33679577708244324 In epoch 90, loss: 0.32077470421791077 In epoch 95, loss: 0.3040713667869568 AUC 0.8397133936793872