pip install torch torch_geometric) -- everything here runs on a laptop, no GPU required for the small graphs we use today;Learn AI Series):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.
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.
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.
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:
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?
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.
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.
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!
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.
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:
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:
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.
Three to chew on before next time -- and as always I'll walk through full solutions in the following episode.
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.
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.
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 ;-)
@scipio