Graph Neural Networks¶
What This Is¶
A graph neural network (GNN) is a model that learns over data whose natural structure is a graph: a set of nodes and edges between them. Social networks, molecules, citation networks, transportation systems, knowledge bases, and protein interaction maps all have graph structure that breaks plain MLPs (no fixed input size) and CNNs (no regular grid).
The single idea at the core of every modern GNN is message passing: each node updates its representation by aggregating information from its neighbours, then mixing that with its own state. Stack a few message-passing layers and a node "sees" a larger neighbourhood; unroll across the graph and the whole graph can be pooled into a single embedding.
When You Use It¶
- nodes carry features and the relationships between them carry signal (citation networks, molecules, recommendations)
- the task is node classification, edge prediction, or whole-graph classification
- the structure is irregular (variable node count, variable neighbourhood size) in a way that CNNs/MLPs cannot handle naturally
- you need to exploit locality in a non-grid space
Do Not Use It When¶
- the graph has almost no structure — if the edges are near-random, a plain MLP on node features wins
- the graph is tiny and classical methods (node2vec, spectral clustering) already solve it
- you are treating a generic tabular problem as a "graph" because it sounds fancy (common temptation; rarely justified)
The Three Task Types¶
| task type | what you predict | example | loss |
|---|---|---|---|
| node classification | a label for each node | classify papers by topic in a citation graph | cross-entropy per node |
| edge prediction (link prediction) | whether an edge should exist | "are these two users friends?" | binary cross-entropy on positive vs. sampled negative edges |
| graph classification | one label for the whole graph | is this molecule toxic? | cross-entropy after pooling node embeddings |
The shape of the loss and the evaluation split change with the task. For edge prediction you must carefully split edges, not nodes, or you leak.
Message Passing In One Equation¶
For a node v with feature h_v and neighbours N(v), a single GNN layer computes:
h_v^{(l+1)} = UPDATE( h_v^{(l)}, AGGREGATE( { h_u^{(l)} : u ∈ N(v) } ) )
Every GNN family is a choice of AGGREGATE and UPDATE:
| family | AGGREGATE |
UPDATE |
|---|---|---|
| GCN (Kipf & Welling) | normalized sum (using degree) | linear + activation |
| GAT (Veličković) | attention-weighted sum over neighbours | linear + activation |
| GraphSAGE | mean / max / LSTM over a sampled neighbourhood | linear + activation on concatenation |
| MPNN (general) | learned message function | learned update function (e.g., GRU) |
| GIN (Graph Isomorphism Network) | sum (maximally expressive) | MLP |
For most first-look problems, GCN or GraphSAGE is a strong baseline. GAT earns its complexity when edges have variable importance. GIN earns its complexity when theoretical expressiveness actually matters.
A Minimal GCN Layer In PyTorch¶
Written out to make the aggregation shape unambiguous:
import torch
from torch import nn
class GCNLayer(nn.Module):
"""Single GCN layer with symmetric normalization."""
def __init__(self, in_dim, out_dim):
super().__init__()
self.linear = nn.Linear(in_dim, out_dim)
def forward(self, x, adj_norm):
# x: (N, in_dim) — node features
# adj_norm: (N, N) sparse — D^{-1/2} (A + I) D^{-1/2}
return adj_norm @ self.linear(x)
class GCN(nn.Module):
def __init__(self, in_dim, hidden_dim, out_dim):
super().__init__()
self.gc1 = GCNLayer(in_dim, hidden_dim)
self.gc2 = GCNLayer(hidden_dim, out_dim)
def forward(self, x, adj_norm):
x = torch.relu(self.gc1(x, adj_norm))
x = self.gc2(x, adj_norm)
return x
For production work, PyTorch Geometric (torch_geometric.nn.GCNConv, GATConv, SAGEConv) gives you efficient sparse implementations and data batching utilities. Use it once you understand the equations above.
Depth Is Different Here¶
In a CNN, making the network deeper gives it a larger receptive field. Same intuition is true in GNNs — k layers let each node see its k-hop neighbourhood — but there is a sharp limit that is specific to GNNs:
- over-smoothing: after a few message-passing rounds, node representations converge to the same vector because every node has been averaged with a large share of the graph. By layer 5–10 in many graphs, node features become indistinguishable.
- over-squashing: a long-range dependency must be compressed through the low-capacity message vectors between distant nodes; information bottlenecks at articulation points.
Counter-intuitively, most strong node-classification results use 2–3 GNN layers. Deeper usually hurts. When depth is needed (e.g., molecular property prediction with long-range interactions), specialized tricks (residual connections, jumping knowledge, graph transformers) are used.
Edge Prediction — The Splitting Trap¶
Splitting a graph for link prediction is not the same as splitting rows of a table:
- the model sees a partial graph during training
- "positive" edges that should exist go in the evaluation set and are removed from the training graph
- "negative" edges are sampled from the node pairs that have no edge
Three failures that are textbook here:
- leaking test edges into the training graph (you trained on the answer)
- using unbalanced negative sampling (model predicts "no edge" for everything and scores 99%)
- evaluating on a temporally later graph while having trained on a random split (the comparable topic: see Splitter Choice Under Ambiguity)
The right split for a deployed link-prediction model is usually time-aware: train on edges present by time T, evaluate on edges that appeared after T.
Pooling For Graph Classification¶
Whole-graph tasks need a permutation-invariant pooling step at the end:
| pooling | behavior | good for |
|---|---|---|
| sum | preserves size information | molecule size correlates with label |
| mean | size-invariant average | most default cases |
| max | picks out highest-activation features | chemistry-style "does any fragment look like X?" |
| attention pooling | learned weights per node | heterogeneous graphs |
| hierarchical pooling (DiffPool, SortPool) | coarsens the graph layer by layer | more flexible; more expensive |
Start with mean pooling and go fancier only when evaluation says so.
What To Inspect¶
- over-smoothing diagnostic: compute the mean pairwise distance between node embeddings at each layer. If the distance collapses after 2–3 layers, you are over-smoothing.
- degree distribution in train vs. test — models often learn degree-based shortcuts
- sampled neighbourhood size in GraphSAGE — stability vs. compute trade
- class balance per node — some graphs are extremely imbalanced (one dominant topic); use class-balanced sampling
- isolated nodes — nodes with no edges collapse to their own features; handle or drop explicitly
- message gradient norms — should be healthy; a vanishing message gradient means deeper layers add nothing
Failure Pattern¶
The defining failure of a deep GNN is silent over-smoothing. The loss keeps decreasing, but the model has converged to "every node is the same vector." Node-level accuracy collapses. The fix: stop adding layers; use residuals or jumping-knowledge connections; consider a graph transformer if you really need long-range signal.
A second common failure is evaluation on a random node split when the deployment is temporal or inductive. Transductive benchmarks (Cora, CiteSeer, PubMed with full-graph visibility) inflate numbers that do not transfer to production inductive settings (brand-new nodes, brand-new subgraphs).
Common Mistakes¶
- treating every table as a graph — graphs earn their complexity only when the structure carries signal
- stacking 6–10 layers and wondering why performance plateaus (over-smoothing)
- training on the full graph for link prediction (edge leakage)
- using a transductive metric on an inductive deployment
- ignoring edge weights / types when they matter (use relational GCNs, heterogeneous networks)
- skipping normalization — GCN math only works if the adjacency is properly normalized
- evaluating a graph classifier on an unbalanced test set without reporting macro metrics
Decision: GCN vs. GAT vs. SAGE vs. Graph Transformer¶
| option | when it wins | trade |
|---|---|---|
| GCN | small to medium graphs, strong baseline | struggles with heterogeneous or long-range |
| GraphSAGE | large graphs needing neighbour sampling | sampling hyperparameters matter |
| GAT | edges vary in importance; attention helps | more compute |
| GIN | graph classification with theoretical needs | not always best empirically |
| Graph transformer | long-range interactions; dense attention across graph | expensive; needs careful positional encoding |
Start with GCN or SAGE. Escalate if evaluation says so.
Practice¶
- Load the Cora citation dataset (PyTorch Geometric bundles it). Train a 2-layer GCN. Report node classification accuracy.
- Stack to 6 layers without residuals. Plot the mean pairwise node-embedding distance per layer. See over-smoothing in action.
- Swap the GCN for a GraphSAGE with mean aggregation. Compare accuracy and training speed.
- On a small link-prediction dataset, split edges into train/val/test. Verify your validation positive edges are not present in the training adjacency. Train a link predictor using dot-product of node embeddings.
- Sample negative edges 1:1 with positive edges. Report ROC AUC on validation. Then redo with 10:1 negative sampling and verify what the metric does. Explain why.
- On a molecule graph dataset (e.g., MUTAG), compare mean, max, and sum pooling for graph classification. Defend a choice.
Runnable Example¶
pip install torch-geometric (plus the matching torch build). The GCN snippet above runs on tiny synthetic graphs; for real benchmarks, PyG's Planetoid datasets are the standard entry point. Running this in the browser runner is not supported.
Longer Connection¶
GNNs sit adjacent to several academy topics:
- Attention and Transformers — GAT is attention on a graph; graph transformers are transformers where the adjacency matrix informs the attention mask
- Self-Supervised and Representation Learning — graph contrastive methods (GraphCL, DGI) are a clean application of SSL ideas
- Splitter Choice Under Ambiguity — the inductive vs. transductive split decision is a concrete instance of that clinic
- Honest Splits and Baselines — the link-prediction splitting trap is the textbook example of "the split must mirror the deployment"
For the big-picture decision about whether to use a GNN at all, the pattern is the same as elsewhere in the academy: try the simplest model on node features alone. If it wins, ship it. If it loses by a small margin, suspect the structure carries signal. Only then does a GNN earn its complexity.