Learn AI Series (#131) - Graph Neural Networks

avatar

Learn AI Series (#131) - Graph Neural Networks

variant-b-09-blue.png

What will I learn

  • You will learn why some data is fundamentally graph-shaped and loses its most valuable signal the moment you flatten it into a table or a sequence;
  • message passing -- the single, disarmingly simple operation that every modern graph neural network is built on top of;
  • Graph Convolutional Networks (GCN), which generalise the convolution idea from episode #45 to neighbourhoods that are irregular in stead of a tidy pixel grid;
  • Graph Attention Networks (GAT), where the attention machinery from episode #51 decides which neighbour actually matters;
  • the three levels of graph prediction (node, edge, graph) and where each one shows up in the real world -- molecules, fraud rings, recommender systems;
  • PyTorch Geometric, the library that makes all of this practical instead of a weekend of index-juggling ;-)

Requirements

  • A working modern computer running macOS, Windows or Ubuntu;
  • An installed Python 3(.10+) distribution with PyTorch (pip install torch torch_geometric) -- everything here runs on a laptop, no GPU required for the small graphs we use today;
  • You've been through episode #45-#46 (convolutional networks, because a GCN is basically convolution set free from the grid) and especially episode #51 (attention mechanisms, which GAT leans on directly). A quick refresher on episode #24 (the curse of dimensionality) will not hurt either.

Difficulty

  • Beginner

Curriculum (of the Learn AI Series):

Learn AI Series (#131) - Graph Neural Networks

I ended episode #130 with a question left deliberately hanging in the air. We spent that whole episode drawing little causal graphs by hand -- severity pointing at treatment, treatment pointing at recovery -- and treating the graph as a map we reason over. And then I asked: if a graph is such a natural way to write down how things relate, could a neural network learn to operate on graph-structured data directly, passing signal along the edges instead of forcing everything into a flat row of numbers? Today we answer that with a resounding yes ;-)

Think about a social network for a second. You have users, and those users follow each other, reply to each other, upvote each other. You could flatten all of that into a table -- one row per user, columns for follower count, post count, average votes received -- and a gradient boosting model (episode #19) would happily chew on it. But you would be throwing away the most valuable thing you own: the structure of who connects to whom. A user with 500 followers who are all sleepy bots is a completely different animal from a user with 500 followers who are all active curators. The table cannot tell them apart. The graph can. That is the whole motivation in a single sentence.

Solutions to episode #130's exercises

Before we open the graph toolbox, let's clear last time's causal homework -- as promised, full code.

Exercise 1 -- crank the confounding up. The task was to regenerate the Simpson's Paradox dataset but make the sickest patients far more likely to be treated (treatment_prob = 0.1 + 0.8 * severity), then refit both the naive and the severity-adjusted logistic regressions. The thing you are meant to see is that the naive coefficient gets more misleading as the confounding strengthens.

import numpy as np
from sklearn.linear_model import LogisticRegression

np.random.seed(42)
n = 2000
severity = np.random.beta(2, 5, n)

# Stronger triage: the sickest are now overwhelmingly the treated ones.
treatment_prob = 0.1 + 0.8 * severity
treatment = np.random.binomial(1, treatment_prob)

recovery_prob = 0.6 - 0.4 * severity + 0.2 * treatment
recovery = np.random.binomial(1, np.clip(recovery_prob, 0.01, 0.99))

naive = LogisticRegression().fit(treatment.reshape(-1, 1), recovery)
correct = LogisticRegression().fit(np.column_stack([treatment, severity]), recovery)
print(f"naive coefficient      : {naive.coef_[0][0]:+.3f}")   # more negative than before
print(f"severity-adjusted coeff: {correct.coef_[0][0]:+.3f}") # still positive, near +0.2's effect

The one-sentence answer I was fishing for: the harder the confounder pushes treatment onto the sickest patients, the more the naive comparison mixes up "sick" with "treated", so the naive coefficient slides further into nonsense while the adjusted one stays honest. Confounding is not noise you can average away with more data -- more data just gives you a more confident wrong answer.

Exercise 2 -- check your overlap. After fitting propensity scores, print the score distributions for treated and control separately, then estimate the ATT with n_neighbors=1 and n_neighbors=5.

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import NearestNeighbors

X = severity.reshape(-1, 1)
ps = LogisticRegression().fit(X, treatment).predict_proba(X)[:, 1]

for name, mask in [("treated", treatment == 1), ("control", treatment == 0)]:
    s = ps[mask]
    print(f"{name}: min={s.min():.2f} mean={s.mean():.2f} max={s.max():.2f}")

def att(scores, treatment, outcome, k):
    t_idx, c_idx = np.where(treatment == 1)[0], np.where(treatment == 0)[0]
    nn = NearestNeighbors(n_neighbors=k).fit(scores[c_idx].reshape(-1, 1))
    diffs = []
    for i in t_idx:
        _, m = nn.kneighbors(scores[i].reshape(1, -1))
        diffs.append(outcome[i] - outcome[c_idx[m[0]]].mean())
    return np.mean(diffs)

print(f"ATT (k=1): {att(ps, treatment, recovery, 1):+.3f}")
print(f"ATT (k=5): {att(ps, treatment, recovery, 5):+.3f}")

The two estimates should land close to each other and both stay positive. Poor overlap would mean the treated group lives almost entirely at high propensity scores while the controls cluster at low ones -- so when you go looking for an untreated twin of a high-score treated patient, the nearest control is not remotely similar, and the "match" is a fiction. When the two distributions barely share a region, matching is estimating an effect for people who have no real counterpart in the data, and you should not trust it. k=5 smooths the estimate (less variance) but risks reaching for worse matches (more bias) -- the classic trade-off from episode #13.

Exercise 3 -- cause or effect? For a churn model: (a) account age -- a plausible cause, older accounts churn less, keep it; (b) support tickets last month -- borderline, usually a leading symptom of dissatisfaction rather than a cause, defensible either way as long as it is genuinely known before the prediction window; (c) subscription cancelled -- that is churn, it is the label wearing a moustache, pure leakage, drop it; (d) monthly spend -- a plausible cause (engaged payers stay), keep it; (e) retention calls the company made -- an effect of an internal churn-risk flag, so it leaks the very thing you are trying to predict, drop it unless you can prove it was logged independently. The reasoning was the whole point -- if you justified each in one line, you nailed it.

What is a graph, formally?

A graph G = (V, E) is a set of nodes V (also called vertices) joined by edges E. Each node can carry a feature vector; each edge can carry attributes too (a weight, a type, a direction). Graphs come in flavours -- directed or undirected, weighted or unweighted -- but the ingredients are always those two: things, and the connections between things.

In matrix form we write connectivity as an adjacency matrix A, where A[i][j] = 1 if there is an edge from node i to node j. Node features stack into a matrix X, one row per node. Having said that, nobody serious stores the full A for a large graph -- it is almost all zeros, a criminal waste of memory. Instead we use COO format: two parallel lists of indices, one for source nodes and one for targets.

import torch

# A tiny social network: 5 users.
# Edges: 0-1, 0-2, 1-2, 2-3, 3-4 (undirected, so both directions are listed).
edge_index = torch.tensor([
    [0, 1, 0, 2, 1, 2, 2, 3, 3, 4],  # source nodes
    [1, 0, 2, 0, 2, 1, 3, 2, 4, 3],  # target nodes
], dtype=torch.long)

# Node features: [post_count, avg_vote, account_age_years]
x = torch.tensor([
    [120, 0.85, 3.0],
    [45,  0.62, 1.5],
    [300, 0.91, 5.0],
    [15,  0.40, 0.5],
    [80,  0.70, 2.0],
], dtype=torch.float)

print(f"Nodes: {x.shape[0]}, Features per node: {x.shape[1]}")
print(f"Edges: {edge_index.shape[1] // 2} (undirected)")

This COO layout is exactly how PyTorch Geometric stores graphs, and it scales to millions of edges without breaking a sweat -- because you only ever store the edges that exist, not the yawning ocean of ones that don't.

The key idea: message passing

Every modern GNN architecture -- and I mean every single one -- boils down to one principle: message passing. Each node updates its own representation by aggregating information from its neighbours. After one round, a node knows about its direct neighbours. After two rounds, it knows about neighbours-of-neighbours. After k rounds, information from k hops away has diffused into it. That is the whole trick, and it is beautiful in its simplicity.

The general framework has three steps:

  1. Message -- each node builds a message to send along its edges (usually derived from its current features);
  2. Aggregate -- each node collects the messages arriving from its neighbours (sum, mean, or max);
  3. Update -- each node folds the aggregated messages back into its own representation.

This is conceptually the cousin of convolution on images (episode #45), except the "neighbourhood" is no longer a fixed grid. A pixel always has 8 neighbours. A graph node might have 2 neighbours, or 2,000. That irregularity is exactly what makes graphs expressive -- and exactly what makes the code a little more interesting.

import torch.nn.functional as F

def simple_message_pass(x, edge_index):
    """One round of message passing: mean of each node's neighbour features."""
    num_nodes = x.shape[0]
    src, dst = edge_index  # a source sends its features to a destination

    agg = torch.zeros_like(x)
    agg.index_add_(0, dst, x[src])  # accumulate source features at the destination

    deg = torch.zeros(num_nodes, dtype=torch.float)
    deg.index_add_(0, dst, torch.ones(src.shape[0]))
    deg = deg.clamp(min=1).unsqueeze(1)  # avoid dividing by zero

    return agg / deg  # mean of neighbour features

h = simple_message_pass(x, edge_index)
print(f"Node 0 original: {x[0]}")
print(f"Node 0 after MP: {h[0]}")  # now the average of nodes 1 and 2

That is the entire mechanism in its most naked form. Node 0 is wired to nodes 1 and 2, so after one round its representation becomes the mean of their features. What we just built is essentially a graph convolution -- and making it learnable is one linear layer away. Wowzers, right?

Graph Convolutional Networks (GCN)

Kipf and Welling's Graph Convolutional Network (2017) took that hand-rolled averaging and dressed it up into a clean, differentiable layer. The GCN update rule for node i, in words, is: transform the mean of your neighbours' features (yourself included) with a learnable weight matrix W, then pass it through a non-linearity.

Two details earn their keep here. First, the node includes itself in the aggregation (we add self-loops), otherwise a node would forget who it is after every round. Second, the normalisation is a touch cleverer than a plain mean -- it uses the geometric mean of node degrees (the famous D^{-1/2} A D^{-1/2}), which empirically behaves better than dividing by degree alone.

import torch.nn as nn

class GCNLayer(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim)

    def forward(self, x, edge_index):
        num_nodes = x.shape[0]
        src, dst = edge_index

        # Add self-loops: every node also sends a message to itself.
        self_loops = torch.arange(num_nodes, device=x.device)
        src = torch.cat([src, self_loops])
        dst = torch.cat([dst, self_loops])

        # Degree, counted WITH the self-loops, for symmetric normalisation.
        deg = torch.zeros(num_nodes, device=x.device)
        deg.index_add_(0, dst, torch.ones(src.shape[0], device=x.device))
        deg_inv_sqrt = deg.pow(-0.5)
        norm = deg_inv_sqrt[src] * deg_inv_sqrt[dst]  # D^{-1/2} A D^{-1/2}

        # Transform first, then aggregate the normalised messages.
        h = self.linear(x)
        out = torch.zeros_like(h)
        out.index_add_(0, dst, h[src] * norm.unsqueeze(1))
        return out

Stacking GCN layers buys you reach: a 2-layer GCN sees 2-hop neighbourhoods, a 3-layer sees 3-hop, and so on. But here is the catch, and it is a genuinely counter-intuitive one -- go too deep and you hit over-smoothing. After many rounds of averaging, every node's representation drifts toward the same vector, the structure gets washed out, and the model turns to grey mush. Unlike CNNs and transformers, which love depth, GCNs have a sweet spot that is frustratingly shallow: 2 to 4 layers, usually. Nota bene: this is one of the few places in deep learning where "add more layers" is actively bad advice.

Graph Attention Networks (GAT)

Back in episode #51 we saw how attention lets a model learn which inputs deserve focus. GCN, for all its elegance, treats every neighbour as equally important (weighted only by degree). But your closest collaborator and some random account that followed you once both sit in your neighbourhood -- and they carry wildly different signal. Averaging them uniformly throws that away.

Graph Attention Networks (Velickovic et al., 2018) fix this by letting each neighbour earn a learned attention weight. In stead of a flat average, the node computes how much to listen to each neighbour, based on their actual features.

class GATLayer(nn.Module):
    def __init__(self, in_dim, out_dim, heads=4):
        super().__init__()
        self.heads = heads
        self.head_dim = out_dim // heads
        self.W = nn.Linear(in_dim, out_dim, bias=False)
        self.attn = nn.Parameter(torch.randn(heads, 2 * self.head_dim))
        nn.init.xavier_uniform_(self.attn.unsqueeze(0))

    def forward(self, x, edge_index):
        num_nodes = x.shape[0]
        src, dst = edge_index
        h = self.W(x).view(num_nodes, self.heads, self.head_dim)

        # Attention score for every edge.
        edge_feat = torch.cat([h[src], h[dst]], dim=-1)  # (E, heads, 2*head_dim)
        scores = (edge_feat * self.attn).sum(dim=-1)       # (E, heads)
        scores = F.leaky_relu(scores, 0.2)

        # Softmax PER destination node (normalise over each node's neighbours).
        scores_max = torch.zeros(num_nodes, self.heads, device=x.device)
        scores_max.index_reduce_(0, dst, scores, 'amax', include_self=False)
        scores = scores - scores_max[dst]  # numerical stability
        exp_scores = scores.exp()
        sum_exp = torch.zeros(num_nodes, self.heads, device=x.device)
        sum_exp.index_add_(0, dst, exp_scores)
        alpha = exp_scores / (sum_exp[dst] + 1e-8)  # the attention weights

        # Weighted aggregation of neighbour messages.
        out = torch.zeros(num_nodes, self.heads, self.head_dim, device=x.device)
        out.index_add_(0, dst, h[src] * alpha.unsqueeze(-1))
        return out.reshape(num_nodes, -1)

The multi-head mechanism (same idea as the transformers of episode #52) lets different heads attend to different kinds of relationship -- one head might learn to weight high-activity neighbours, another might chase neighbours with similar posting patterns -- and their outputs get concatenated. GAT's edge over GCN is that its weights are computed per edge, per layer, and they shift with the actual feature values rather than being frozen into the graph's shape. That dynamism costs a little more compute, but on many tasks it earns a few points of accuracy in return.

Building a node classifier with PyTorch Geometric

Right, enough index gymnastics -- let's do it the grown-up way with PyTorch Geometric (PyG), the standard library for graph deep learning. PyG ships optimised message-passing layers, graph data structures, and the classic benchmark datasets, so you never have to hand-roll an index_add_ again unless you want to.

from torch_geometric.nn import GCNConv, GATConv
from torch_geometric.datasets import Planetoid

# Cora: a citation network. 2708 papers, 7 topic classes, 5429 edges.
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
print(f"Nodes: {data.num_nodes}, Edges: {data.num_edges}")
print(f"Features: {data.num_node_features}, Classes: {dataset.num_classes}")
print(f"Train: {data.train_mask.sum()}, Test: {data.test_mask.sum()}")

class GNN(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim, use_gat=False):
        super().__init__()
        if use_gat:
            self.conv1 = GATConv(in_dim, hidden_dim, heads=4, concat=True)
            self.conv2 = GATConv(hidden_dim * 4, out_dim, heads=1, concat=False)
        else:
            self.conv1 = GCNConv(in_dim, hidden_dim)
            self.conv2 = GCNConv(hidden_dim, out_dim)
        self.dropout = nn.Dropout(0.5)

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

model = GNN(dataset.num_node_features, 64, dataset.num_classes, use_gat=False)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

for epoch in range(200):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

model.eval()
pred = model(data).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()
print(f"Test accuracy: {acc:.4f}")

Cora is the "MNIST of graph learning" -- a citation network where each paper is a node, citations are edges, and the job is to classify papers by topic. With just two GCN layers you will typically land around 80-82% accuracy. Flip use_gat=True and you nudge to 82-84%. And here is the number that should make you sit up: a plain MLP that ignores the graph and looks only at each paper's own features gets roughly 55-60%. That gap -- twenty-odd points -- is pure structure. The citations, the who-references-whom, carry more signal than the words in the papers themselves. Holy Macaroni!

Beyond node classification

Node classification is the "hello world" of GNNs, but it is not the only game in town. There are three levels of prediction on a graph:

Node-level -- predict something for each node (our Cora example). On a social graph: is this user a bot?

Edge-level -- predict whether an edge exists, or what its properties are. On a social graph: will these two users become friends? In drug discovery: will this molecule bind to that protein?

Graph-level -- predict a property of the entire graph. In chemistry: is this molecule toxic? In network analysis: is this subgraph a fraud ring?

For graph-level tasks you need a readout (or pooling) step that squashes all the node representations into a single graph vector.

from torch_geometric.nn import global_mean_pool

class GraphClassifier(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.conv1 = GCNConv(in_dim, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, hidden_dim)
        self.fc = nn.Linear(hidden_dim, out_dim)

    def forward(self, x, edge_index, batch):
        x = F.relu(self.conv1(x, edge_index))
        x = F.relu(self.conv2(x, edge_index))
        x = global_mean_pool(x, batch)  # average all node features per graph
        return self.fc(x)

The batch tensor is the clever bit: it tells PyG which nodes belong to which graph in a mini-batch. Under the hood PyG concatenates all the graphs in a batch into one big disconnected graph and uses batch to keep track of membership -- so message passing never leaks across graph boundaries, and you still get the speed of batched tensor ops. Simple and efficient, exactly the kind of trick I love.

Where GNNs actually shine

None of this is academic decoration. GNNs are quietly earning their keep in production, in a handful of places where the graph structure is the value:

  • Drug discovery and molecular property prediction. Molecules are natively graphs -- atoms are nodes, bonds are edges. GNNs predict toxicity, solubility, and binding affinity straight from structure, and real pharma outfits use GNN-based virtual screening to narrow millions of candidates down to a handful worth testing in a lab.
  • Social network analysis. Community detection, link prediction, spotting bot swarms. The structure encodes relationships that no amount of per-user tabular feature engineering can recover.
  • Recommendation systems. User-item interaction graphs let GNNs capture collaborative-filtering signal across multi-hop neighbourhoods (this is the causal-vs-correlational recommender tension we flagged back in episode #130 -- structure helps, but mind the popularity bias). Pinterest famously runs graph-based recommenders at web scale.
  • Knowledge graphs. Entities and relations form giant graphs; GNNs reason over them for question answering and knowledge-base completion. This ties straight back to the RAG systems of episodes #64-65 -- a knowledge graph can serve as structured, hop-able context instead of a flat pile of documents.
  • Fraud detection. Fraudsters cluster in graph space. A fraud ring's transaction graph looks structurally wrong even when every individual account has been carefully groomed to look normal -- and that structural wrongness is precisely what a GNN can smell.

The limits of GNNs

Honesty time, because I promised you honesty across this whole series. GNNs have real, load-bearing limitations, and you should know them before you deploy one:

  • Over-smoothing. As we saw, more than 4-5 layers and all node representations converge to mush. The sweet spot is annoyingly shallow.
  • Expressivity. Standard message-passing GNNs are provably bounded by the Weisfeiler-Leman graph isomorphism test -- there exist genuinely different graphs that a vanilla GCN maps to identical representations. It literally cannot tell them apart, no matter how much you train it.
  • Scalability. For graphs with millions of nodes, full-batch training does not fit in memory, and mini-batching is hard because neighbourhoods overlap and explode. Sampling tricks like GraphSAGE and Cluster-GCN exist, but they trade exactness for feasibility.
  • Dynamic graphs. Most of the literature assumes a static graph, yet real networks change every second. Temporal graph learning is an active, less-mature research frontier.

None of these are reasons to avoid GNNs -- they are reasons to reach for them with your eyes open, knowing what the tool can and cannot do.

Round-up

  • Graphs represent relational data -- social networks, molecules, knowledge bases, even the causal DAGs from last episode -- that loses its most important signal the moment you flatten it into a table;
  • message passing is the one mechanism underneath everything: each node aggregates information from its neighbours, over multiple rounds, reaching k hops after k layers;
  • GCN applies symmetric normalisation and a learned transform to neighbour aggregation -- simple, differentiable, effective, but prone to over-smoothing if you stack it too deep;
  • GAT adds learned attention weights so each neighbour contributes differently, computed per edge and per layer, often buying a few points of accuracy;
  • PyTorch Geometric gives you production-grade GCN, GAT and dozens of other layers with memory-efficient sparse representations -- on Cora, two layers beat a graph-blind MLP by twenty points;
  • prediction lives at three levels -- node, edge, and graph (the last one needs a readout/pooling step) -- covering fraud rings, molecules, friend suggestions and more;
  • and mind the limits: over-smoothing, the Weisfeiler-Leman expressivity ceiling, and the pain of scaling to giant, ever-changing graphs.

Exercises

Three to chew on before next time -- and as always I'll walk through full solutions in the following episode.

  1. Watch over-smoothing happen. Take the Cora GNN above and rebuild it with 2, then 4, then 8 GCN layers (same hidden size throughout). Train each for 200 epochs and record the test accuracy. Then, for the 8-layer model, grab the final-layer node embeddings and compute the average pairwise cosine similarity between them. Write two sentences linking what you see in the accuracy to what you see in that similarity number.

  2. GCN versus GAT, head to head. Train both variants (use_gat=False and use_gat=True) on Cora five times each with different random seeds, and report the mean and standard deviation of test accuracy for both. Does GAT reliably beat GCN, or do their error bars overlap? Add one sentence on whether the extra compute of attention looks worth it on this particular dataset.

  3. Build a graph by hand and reason about hops. Construct a small edge_index for a 6-node "friend graph" of your own invention (draw it on paper first). Run the simple_message_pass function from this episode twice in a row on random node features, and confirm that after two rounds a node's representation has been influenced by nodes exactly two hops away. Then answer in one line: for a node whose farthest neighbour is 4 hops away, how many message-passing rounds does it take before that distant node's signal can possibly reach it?

We have spent this whole episode celebrating data that is proudly, irreducibly graph-shaped -- and rightly so. But let me leave you with an uncomfortable admission: the overwhelming majority of the data sitting in the world's databases is not a molecule or a social graph at all. It is boring, flat, tabular stuff -- rows and columns, the spreadsheets and SQL tables that quietly run every bank, shop and hospital on the planet. And here is the twist that ought to nag at you: on that kind of data, the fancy deep networks we have been building for a hundred episodes very often lose to a humble gradient-boosted tree from episode #19. Why does deep learning struggle exactly where there is the most data of all -- and what do you actually reach for when the problem is a plain old table? That is where we go next ;-)

De groeten!

@scipio

stem stemsocial steemstem python programming
0
0
0.000
0 comments