Link Prediction and Graph Classification
Edge prediction task, negative sampling, graph-level pooling, GINConv for graph classification.
Link Prediction and Graph Classification is a free Learn AI with Python lesson on CoddyKit — lesson 4 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the Learn AI with Python learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.
Two New Graph Tasks
Beyond classifying nodes, GNNs handle:
- Link prediction: will an edge exist between two nodes? (friend suggestions, drug interactions)
- Graph classification: assign a label to an entire graph (is this molecule toxic?)
Link Prediction Setup
In link prediction we first compute node embeddings with a GNN, then score candidate node pairs. A high score means the model believes an edge should connect them.
Scoring an Edge
A common edge score is the dot product of the two node embeddings: score = dot(h_u, h_v). Similar embeddings yield a high dot product, predicting a likely link.
h = gnn(data.x, data.edge_index) # node embeddings
score = (h[u] * h[v]).sum(dim=-1) # dot product per pairNegative Sampling
The graph only lists edges that exist (positives). To train a classifier we also need non-edges. Negative sampling randomly picks node pairs that are not connected as negative examples, balancing the training set.
from torch_geometric.utils import negative_sampling
neg_edge_index = negative_sampling(
edge_index=data.edge_index,
num_nodes=data.num_nodes,
num_neg_samples=data.edge_index.size(1),
)BCEWithLogitsLoss
Link prediction is binary (edge or no edge). We score positive and negative pairs, label them 1 and 0, and train with BCEWithLogitsLoss, which combines a sigmoid with binary cross-entropy in a numerically stable way.
import torch
pos = (h[pos_u] * h[pos_v]).sum(-1)
neg = (h[neg_u] * h[neg_v]).sum(-1)
scores = torch.cat([pos, neg])
labels = torch.cat([torch.ones_like(pos), torch.zeros_like(neg)])
loss = torch.nn.functional.binary_cross_entropy_with_logits(scores, labels)Switching to Graph Classification
For graph classification we need a single vector per graph, not per node. After the GNN layers produce node embeddings, we pool them into one graph-level representation.
global_mean_pool
global_mean_pool averages all node embeddings in a graph to produce one fixed-size vector, regardless of graph size. A batch index tells it which nodes belong to which graph when graphs are batched together.
from torch_geometric.nn import global_mean_pool
h = gnn(x, edge_index) # [num_nodes, dim]
hg = global_mean_pool(h, batch) # [num_graphs, dim]
logits = classifier(hg)Why Pooling Matters
Pooling makes the model invariant to node ordering and graph size: two isomorphic graphs yield the same pooled vector. Mean pooling is simple; sum and max pooling are alternatives with different sensitivities.
GINConv
GINConv (Graph Isomorphism Network) is a more expressive convolution. It uses an MLP and sum aggregation specifically designed to maximize the discriminative power of message passing for graph-level tasks.
from torch_geometric.nn import GINConv
import torch
mlp = torch.nn.Sequential(
torch.nn.Linear(in_dim, hid),
torch.nn.ReLU(),
torch.nn.Linear(hid, hid),
)
conv = GINConv(mlp)The Weisfeiler-Leman Connection
GIN is designed to be as powerful as the Weisfeiler-Leman (WL) test, a classic algorithm for distinguishing non-isomorphic graphs. Many simpler GNNs cannot tell certain graphs apart; GIN can, up to the limits of the WL test, making it strong for graph classification.
Choosing the Right Tool
Match the architecture to the task:
- Link prediction: GNN embeddings + dot-product scoring + negative sampling + BCE loss
- Graph classification: expressive convs like GINConv + global pooling + classifier
Quick Check
Test your knowledge.
Recap
You learned link prediction and graph classification:
- Edge score =
dot(h_u, h_v), trained with negative sampling and BCEWithLogitsLoss - global_mean_pool turns node embeddings into a graph-level vector
- GINConv is highly expressive, matching the Weisfeiler-Leman test
Frequently asked questions
Is the “Link Prediction and Graph Classification” lesson free?
Yes — the full text of “Link Prediction and Graph Classification” is free to read here on the web, and the Learn AI with Python course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the Learn AI with Python course, upgrade to CoddyKit PRO.
What will I learn in “Link Prediction and Graph Classification”?
Edge prediction task, negative sampling, graph-level pooling, GINConv for graph classification. You practise Learn AI with Python with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.
Do I need any experience to start Learn AI with Python?
No prior experience is required. Learn AI with Python on CoddyKit is structured for beginners through advanced learners; this is — lesson 4 of 4, so you can start here or from the beginning and move at your own pace.
How long does the “Link Prediction and Graph Classification” lesson take?
Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.
Can I write and run code in this Learn AI with Python lesson?
Yes. Every Learn AI with Python lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.
All lessons in this course
- Graph Theory for Machine Learning
- Graph Convolutional Networks (GCN)
- Node Classification with GNN
- Link Prediction and Graph Classification