Graph Neural Networks

Apply GNNs to fraud detection, social networks, and molecular design with message passing, graph convolutions, and pooling.

Graph Neural Networks (GNNs) are a class of deep learning models designed to operate on graph-structured data, where nodes represent entities and edges represent relationships between them. GNNs process information through graph topology, making them suitable for tasks like node classification, link prediction, and graph classification. They matter when data naturally forms networks or when relationships between entities carry predictive value.

Training a Graph Neural Network

To train a GNN model, you typically use a framework that supports graph operations. A common approach uses PyTorch Geometric, which provides torch_geometric.nn.GCNConv for graph convolution layers. The training loop requires handling adjacency matrices and node features as tensors.

Example training setup:

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

class GCN(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(num_features, 16)
        self.conv2 = GCNConv(16, num_classes)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

model = GCN(num_features=10, num_classes=3)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

The key tensor inputs are:

Training step example:

model.train()
optimizer.zero_grad()
out = model(data.x, data.edge_index)
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()

Failure mode: If edge_index contains duplicate edges or self-loops without proper handling, gradients can become unstable or the model fails to learn meaningful representations.

Node Classification with GNNs

For node classification tasks, you must provide both the graph structure and labels for a subset of nodes. The graph is represented as an adjacency matrix or edge list, and features are stored per node.

In PyTorch Geometric, node classification uses data objects like Data(x, edge_index, y) where:

Model training requires a mask to identify which nodes are used for training:

train_mask = torch.tensor([True, False, True, False, ...])

The model output must match the number of classes specified in the final layer. Mismatched dimensions cause runtime errors during loss computation.

Failure mode: Incorrectly setting the y tensor shape or missing labels for nodes in the training mask leads to IndexError or ValueError during backward pass.

Link prediction involves predicting whether an edge exists between two nodes. Common approaches include computing node embeddings and using dot products or MLPs to score potential connections.

In PyTorch Geometric, link prediction often uses a scoring function:

def link_pred_score(z, edge_index):
    # z: node embeddings from GNN
    # edge_index: edges for which we compute scores
    return torch.sigmoid(torch.sum(z[edge_index[0]] * z[edge_index[1]], dim=1))

The input graph must be undirected for symmetric scoring. Directed graphs require separate scoring for forward and reverse edges.

Failure mode: Using the same node embeddings for both ends of an edge without considering directionality leads to poor performance on directed graph tasks.

Graph Classification Tasks

Graph classification involves predicting a label for an entire graph rather than individual nodes. This requires aggregating information from all nodes in the graph.

A common approach is to use a global pooling layer like GlobalAttentionPool or Set2Set to reduce the node-level representations to a single graph-level vector.

Example architecture:

from torch_geometric.nn import global_mean_pool

class GraphClassifier(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super(GraphClassifier, self).__init__()
        self.conv1 = GCNConv(num_features, 16)
        self.conv2 = GCNConv(16, 32)
        self.classifier = torch.nn.Linear(32, num_classes)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        x = global_mean_pool(x, batch)  # Aggregate over nodes
        x = self.classifier(x)
        return F.log_softmax(x, dim=1)

The batch tensor maps each node to its corresponding graph in the batch.

Failure mode: Incorrectly handling the batch tensor causes incorrect aggregation, leading to misaligned predictions across different graphs.

Handling Large Graphs

When working with large graphs, memory usage becomes critical. Techniques include mini-batch training using DataLoader or sampling subgraphs.

In PyTorch Geometric:

from torch_geometric.data import DataLoader

loader = DataLoader(dataset, batch_size=32, shuffle=True)
for data in loader:
    # Process batch of graphs
    out = model(data.x, data.edge_index, data.batch)

Sampling methods like NeighborSampler help reduce memory requirements by processing only neighborhoods.

Failure mode: Not properly handling the batch tensor in mini-batch scenarios leads to incorrect graph-level aggregations or loss computation errors.

Model Evaluation Metrics

Evaluation metrics for GNNs depend on the task type. For node classification, common metrics include accuracy, precision, recall, and F1-score.

For link prediction, metrics like ROC-AUC and average precision are used.

Example evaluation code:

def evaluate(model, data):
    model.eval()
    with torch.no_grad():
        out = model(data.x, data.edge_index)
        pred = out.argmax(dim=1)
        correct = (pred == data.y).sum()
        accuracy = correct / data.num_nodes
        return accuracy.item()

Failure mode: Using incorrect metrics or misinterpreting results due to class imbalance in node classification can lead to misleading performance assessments.

Common Configuration Issues

GNNs require careful attention to hyperparameters such as learning rate, number of layers, hidden dimensions, and dropout rates. The optimal configuration depends heavily on the graph structure and task.

Example hyperparameter tuning setup:

# Learning rate schedule
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)

Failure mode: Using a fixed learning rate without warmup or decay leads to poor convergence or overfitting.

Memory and Performance Considerations

Memory usage grows with graph size and model depth. For large graphs, consider using sparse matrix operations or memory-efficient training techniques.

PyTorch's torch.sparse module supports sparse adjacency matrices:

# Convert dense adjacency to sparse format
adj_sparse = torch.sparse_coo_tensor(indices, values, size)

Failure mode: Dense adjacency matrices in memory-intensive tasks cause out-of-memory errors, especially when the graph has many nodes.

Debugging GNN Training

Debugging GNN training involves monitoring loss curves, gradient norms, and model outputs. Common issues include vanishing gradients or exploding gradients.

Gradient clipping:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Failure mode: Not clipping gradients leads to unstable training with NaN losses or divergent weights.

Graph Data Preparation

Graph data must be converted into tensors suitable for GNN operations. This includes:

Example preprocessing:

# Convert edge list to tensor format
edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)

Failure mode: Incorrect tensor dimensions or data types cause runtime errors during forward passes.

Edge Cases and Gotchas

Some common edge cases include graphs with isolated nodes, disconnected components, or varying degrees of sparsity.

Isolated nodes can be handled by ensuring they have appropriate feature vectors:

# Add dummy features for isolated nodes
x = torch.cat([x, torch.zeros(num_isolated, num_features)])

Failure mode: Ignoring isolated nodes results in incorrect embedding dimensions and downstream errors.

Model Serialization

Trained GNN models should be saved using standard PyTorch methods:

torch.save(model.state_dict(), 'gcn_model.pth')
# Later load with:
model.load_state_dict(torch.load('gcn_model.pth'))

Failure mode: Saving only the model architecture without state dict causes loss of trained weights.

Batch Processing Considerations

When processing multiple graphs in batches, ensure that all graphs are properly batched and that pooling operations align correctly.

Example batch handling:

# Ensure batch tensor is correct
batch = torch.tensor([0, 0, 1, 1, 1])  # First two nodes belong to graph 0, last three to graph 1

Failure mode: Incorrect batch indexing causes misaligned node embeddings and incorrect graph-level predictions.

This page was rewritten on 15 October 2026. It replaced a templated version that was not specific to its own title. How these pages are written and checked. If anything here is wrong, tell us at [email protected] and we will correct it.