Graph ML: Embeddings and GNNs for Prediction

Graph Data: Why Traditional ML Fails

Traditional machine learning models, such as linear regression, Support Vector Machines (SVMs), or feed-forward neural networks, operate on fixed-size feature vectors. Each data point is treated as an independent instance, represented by a flat array of numbers. Graph data, by its nature, violates both these assumptions.

Graphs are inherently relational and irregular. A graph’s structure — the number of nodes, edges, and their connections — varies. Representing this dynamic structure as a fixed-length vector for a traditional model is problematic. For instance, an adjacency matrix for a graph with $N$ nodes is $N \times N$. Its size changes with $N$, and the ordering of nodes within the matrix is arbitrary. If two nodes are swapped, the matrix changes, yet the underlying graph structure remains identical. A traditional model would interpret these as distinct inputs.

Consider a node in a social network. Its attributes (age, location) can be vectorized. However, its connections to other users and the attributes of those neighbors are equally significant. Flattening this rich relational context into a single, fixed-size vector for the central node often results in significant information loss. Aggregating neighbor features (e.g., average neighbor age) discards specific identities and higher-order connection patterns.

Furthermore, traditional models assume data points are independent and identically distributed (i.i.d.). Graph data fundamentally violates this. A node’s features, label, and behavior are often influenced by its neighbors and the broader graph structure. For example, a user’s purchase decisions might correlate with those of their friends. Treating each node as an isolated, independent entity ignores these dependencies, which are often the most valuable signal in graph-structured problems.

These limitations mean that traditional ML models cannot directly process graph structures or capture the complex dependencies between nodes. They require feature engineering to extract graph properties into flat vectors, which is lossy and often fails to represent the full relational context.

Node Embeddings: Representing Graph Structure

Machine learning models require fixed-size numerical feature vectors as input. Graphs, however, represent data as nodes and edges, lacking a direct vector representation suitable for most ML algorithms. Node embeddings convert a node’s structural and feature information into a low-dimensional, dense vector.

These embeddings capture node similarities: nodes with similar roles or neighborhoods should have similar vectors. This transformation allows standard ML techniques, like classification or clustering, to operate on graph data by using these learned vectors as features.

Early approaches to generating node embeddings involved matrix factorization. For a graph’s adjacency matrix $A$, Singular Value Decomposition (SVD) can decompose $A$ into $U \Sigma V^T$. Truncating this decomposition to the top $k$ singular values yields a $k$-dimensional embedding for each node.

While simple, matrix factorization scales poorly with large graphs, exhibiting $O(N^3)$ complexity for $N$ nodes. This method primarily captures direct connections and struggles to represent higher-order neighborhood information efficiently.

Random walk-based methods address these scaling and expressiveness limitations by sampling node sequences. DeepWalk, for instance, generates fixed-length random walks starting from each node. These sequences are then treated as “sentences,” and nodes within a walk are “words.”

A skip-gram model, similar to Word2Vec, learns embeddings such that nodes appearing close in walks have similar vectors. This captures the local context of a node within the graph. Node2Vec extends DeepWalk by introducing parameters to control walk bias, allowing a balance between Breadth-First Search (BFS)-like local exploration and Depth-First Search (DFS)-like global exploration.

# Conceptual example of random walk generation (not DeepWalk/Node2Vec implementation)
import networkx as nx
import random

graph = nx.Graph()
graph.add_edges_from([(1, 2), (1, 3), (2, 4), (3, 5), (4, 5)])

def generate_random_walk(graph, start_node, walk_length):
    walk = [start_node]
    for _ in range(walk_length - 1):
        current_node = walk[-1]
        neighbors = list(graph.neighbors(current_node))
        if not neighbors:
            break
        walk.append(random.choice(neighbors))
    return walk

# Example walk starting from node 1
print(generate_random_walk(graph, 1, 5))

[1, 2, 4, 5, 3]

The output shows one possible random walk sequence. The actual embedding process uses many such walks and a neural network (like Word2Vec) to learn vector representations from these sequences.

Edge embeddings capture the relationship between two nodes. These are often derived from their constituent node embeddings. Common approaches include the element-wise product, sum, or concatenation of the two node vectors. The choice of aggregation depends on the specific downstream task, such as link prediction or edge classification.

Graph Neural Networks (GNNs), discussed in subsequent chapters, learn node embeddings end-to-end. They aggregate information from a node’s neighborhood iteratively, integrating both structural and feature data directly into the embedding generation process.

Graph Neural Networks: Message Passing Fundamentals

Traditional machine learning models assume data points are independent. Graph data, however, explicitly defines relationships between entities, making this independence assumption invalid. Graph Neural Networks (GNNs) learn representations that capture both node features and the graph’s topological structure by directly operating on the graph.

The core mechanism of a GNN is message passing. During each layer, a node iteratively updates its feature representation by aggregating information from its direct neighbors. This process allows nodes to incorporate local structural context into their embeddings.

This process involves two main steps for each node in a layer:

1. Aggregation: Each node collects feature vectors from its neighbors. An aggregation function (e.g., sum, mean, or max) combines these neighbor features into a single ‘message’. This step summarizes the local neighborhood information.
2. Update: The aggregated message is then combined with the node’s own current feature vector. A neural network often processes this combined input to produce the node’s new, higher-level feature representation for the next layer.

Stacking multiple GNN layers allows a node to incorporate information from increasingly distant neighbors. A GNN with k layers can learn representations that depend on the features of nodes up to k hops away, effectively expanding the node’s ‘receptive field’ across the graph.

The parameters of the aggregation and update functions, typically weights within neural networks, are learned during training. This enables the GNN to discover optimal ways to extract and combine local graph information for a given prediction task.

After several message passing layers, each node possesses an embedding that encodes both its original features and its structural context within the graph. These learned node embeddings serve as input for downstream tasks, such as node classification, link prediction, or graph-level classification.

GNN Implementation: Node Classification Example

Node classification identifies categories for individual nodes within a graph. This task uses node features and graph structure to predict a label for each node. PyTorch Geometric (PyG) provides a framework for implementing Graph Neural Networks (GNNs) efficiently.

The Cora dataset serves as a standard benchmark for this task. It contains research papers as nodes, citations as edges, and paper topics as node labels. We load the dataset and inspect its structure.

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv

# Load the Cora dataset
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

print(f'Number of nodes: {data.num_nodes}')
print(f'Number of edges: {data.num_edges}')
print(f'Number of node features: {data.num_node_features}')
print(f'Number of classes: {dataset.num_classes}')
print(f'Training nodes: {data.train_mask.sum()}')
print(f'Test nodes: {data.test_mask.sum()}')

Number of nodes: 2708
Number of edges: 10556
Number of node features: 1433
Number of classes: 7
Training nodes: 140
Test nodes: 1000

The data object holds node features (data.x), graph connectivity (data.edge_index), node labels (data.y), and masks for training, validation, and testing nodes (data.train_mask, data.val_mask, data.test_mask).

A simple two-layer Graph Convolutional Network (GCN) can perform node classification. Each GCNConv layer aggregates features from a node’s neighborhood. The model applies a ReLU activation function and dropout between layers to introduce non-linearity and prevent overfitting.

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

model = GCN(data.num_node_features, 16, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
criterion = torch.nn.CrossEntropyLoss()

The train function performs a single optimization step. It computes the loss only on the training nodes identified by data.train_mask. The test function evaluates the model’s accuracy on the test set.

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

def test():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)
    correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
    acc = int(correct) / int(data.test_mask.sum())
    return acc

for epoch in range(1, 201):
    loss = train()
    if epoch % 20 == 0:
        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')

test_acc = test()
print(f'Test Accuracy: {test_acc:.4f}')

Epoch: 020, Loss: 0.8142
Epoch: 040, Loss: 0.4429
Epoch: 060, Loss: 0.3603
Epoch: 080, Loss: 0.3113
Epoch: 100, Loss: 0.2936
Epoch: 120, Loss: 0.2801
Epoch: 140, Loss: 0.2662
Epoch: 160, Loss: 0.2520
Epoch: 180, Loss: 0.2443
Epoch: 200, Loss: 0.2372
Test Accuracy: 0.8060

This output shows the model learning to classify nodes effectively, achieving an accuracy over 80% on the test set.

GNN Training: Avoiding Common Pitfalls

Effective GNN training requires addressing specific challenges inherent to graph-structured data. Ignoring these can lead to models with inflated performance metrics or poor generalization capabilities.

Over-smoothing is a common issue in GNNs. As information propagates through multiple layers, node representations can become increasingly similar to their neighbors. After a certain depth, all nodes within a connected component may converge to nearly identical embeddings, losing their distinctiveness. This reduces the model’s ability to discriminate between nodes based on their local structure and features.

Mitigating over-smoothing often involves limiting the number of GNN layers, typically to 2-4. While this approach is simpler, it restricts the model’s receptive field, potentially losing access to important long-range dependencies. Deeper architectures can use skip connections (e.g., residual connections) or attention mechanisms to allow information to flow across layers selectively, preventing full aggregation and preserving local information.

Another pitfall involves incorrect data splitting, particularly confusing transductive and inductive learning settings. In a transductive task, the entire graph structure is known during training, but only a subset of nodes or edges have labels. The goal is to predict labels for the unlabeled nodes within the same graph. For instance, in node classification, data.train_mask, data.val_mask, and data.test_mask typically partition nodes within a single graph:

import torch
from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

print(f"Total nodes: {data.num_nodes}")
print(f"Nodes for training: {data.train_mask.sum().item()}")
print(f"Nodes for validation: {data.val_mask.sum().item()}")
print(f"Nodes for testing: {data.test_mask.sum().item()}")

Conversely, inductive learning requires the model to generalize to entirely unseen graphs or nodes that were not part of the training graph at all. Incorrectly applying a transductive split (where test nodes’ connections are visible during training) to an inductive problem inflates performance metrics. For inductive tasks, ensure the test set consists of completely separate graphs or disconnected subgraphs. This strict separation ensures the model learns truly generalizable patterns, yielding a more robust evaluation of its performance on new data.