Skip to content

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

  1. Load the Cora citation dataset (PyTorch Geometric bundles it). Train a 2-layer GCN. Report node classification accuracy.
  2. Stack to 6 layers without residuals. Plot the mean pairwise node-embedding distance per layer. See over-smoothing in action.
  3. Swap the GCN for a GraphSAGE with mean aggregation. Compare accuracy and training speed.
  4. 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.
  5. 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.
  6. On a molecule graph dataset (e.g., MUTAG), compare mean, max, and sum pooling for graph classification. Defend a choice.

Runnable Example

GNN practice is best done with PyTorch Geometric. Install with 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:

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.