Graph Neural Networks
1. Message Passing Framework
1.1 Graph Representation
A graph consists of:
- Nodes with features
- Edges with features
- Adjacency matrix
1.2 Message Passing Neural Network (MPNN)
The general message passing framework (Gilmer et al., 2017) consists of three steps:
Message computation:
Aggregation:
Update:
where is the neighbourhood of node and denotes a permutation-invariant operation.
1.3 Aggregation Functions
| Aggregation | Definition | Properties |
|---|---|---|
| Mean | Smooth, invariant to degree | |
| Sum | Expressive, sensitive to degree | |
| Max | Robust to outliers | |
| Attention | Learned importance |
2. Graph Convolutional Network (GCN)
2.1 Spectral Perspective
The graph Fourier transform uses the eigenvectors of the graph Laplacian:
where is the degree matrix.
The normalised Laplacian:
has eigendecomposition .
2.2 Spectral Convolution
A spectral filter acts on graph signal as:
where is a diagonal matrix of filter coefficients.
2.3 GCN Approximation
Kipf & Welling (2017) approximated the spectral filter with a first-order Chebyshev polynomial:
Imposing :
With renormalisation trick , :
2.4 GCN Properties
Permutation equivariance: If is a permutation matrix:
Over-smoothing: After layers, node features converge:
where is the dominant eigenvector of .
3. GraphSAGE
3.1 Inductive Learning
GraphSAGE (Hamilton et al., 2017) enables inductive learning by sampling and aggregating:
where denotes a fixed-size sample.
3.2 Aggregation Variants
Mean aggregator:
LSTM aggregator:
where is a random permutation.
Pooling aggregator:
where .
4. Graph Attention Network (GAT)
4.1 Attention Mechanism
GAT (Veličković et al., 2018) computes attention weights for edges:
where denotes concatenation and is a learnable attention vector.
4.2 Multi-Head Attention
where is the number of attention heads.
4.3 GAT Complexity
This is linear in the number of edges, making GAT scalable.
5. Over-Squashing
5.1 The Problem
Over-squashing (Alon et al., 2021) occurs when information from distant nodes is compressed through a bottleneck:
In graphs with exponentially growing neighbourhoods, the information from nodes must be compressed into a fixed-dimensional vector.
5.2 Information-Theoretic Analysis
The mutual information between a target node and distant nodes decreases exponentially with distance:
With each layer, the information is diluted across more neighbours.
5.3 Solutions
Graph rewiring: Add long-range edges to reduce diameter.
Jumping knowledge: Concatenate features from all layers:
Graph transformers: Use global attention to bypass the neighbourhood bottleneck.
6. Temporal Graph Networks
6.1 Dynamic Graphs
Temporal graphs have edges that appear/disappear over time:
6.2 Temporal Message Passing
where is the timestamp of the last edge between and .
6.3 Graph Transformers
Graphormer (Ying et al., 2021):
where encodes structural information (centrality, spatial encoding).
GPS (Rampášek et al., 2022): Combines local message passing with global attention.
7. Code Examples
7.1 GCN Implementation
import torch
import torch.nn as nn
import torch.nn.functional as F
class GCNLayer(nn.Module):
"""Graph Convolutional Network layer."""
def __init__(self, in_features, out_features, bias=True):
super().__init__()
self.linear = nn.Linear(in_features, out_features, bias=bias)
self.reset_parameters()
def reset_parameters(self):
nn.init.xavier_uniform_(self.linear.weight)
if self.linear.bias is not None:
nn.init.zeros_(self.linear.bias)
def forward(self, x, adj):
"""
Forward pass.
Parameters
----------
x : Tensor (N, in_features) - node features
adj : Tensor (N, N) - adjacency matrix (with self-loops)
Returns
-------
h : Tensor (N, out_features) - updated node features
"""
# Degree normalisation
deg = adj.sum(dim=1, keepdim=True)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
# Normalised adjacency: D^{-1/2} A D^{-1/2}
adj_norm = deg_inv_sqrt * adj * deg_inv_sqrt.T
# Aggregate and transform
support = self.linear(x)
output = torch.spmm(adj_norm, support) if adj.is_sparse else adj_norm @ support
return output
class GCN(nn.Module):
"""Multi-layer GCN for node classification."""
def __init__(self, in_features, hidden_features, out_features, num_layers=2, dropout=0.5):
super().__init__()
self.layers = nn.ModuleList()
# Input layer
self.layers.append(GCNLayer(in_features, hidden_features))
# Hidden layers
for _ in range(num_layers - 2):
self.layers.append(GCNLayer(hidden_features, hidden_features))
# Output layer
self.layers.append(GCNLayer(hidden_features, out_features))
self.dropout = dropout
def forward(self, x, adj):
"""
Forward pass through all layers.
Parameters
----------
x : Tensor (N, in_features)
adj : Tensor (N, N)
Returns
-------
logits : Tensor (N, out_features)
"""
for i, layer in enumerate(self.layers[:-1]):
x = layer(x, adj)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.layers[-1](x, adj)
return x
def compute_over_smoothing(self, x, adj, num_layers=10):
"""Analyse over-smoothing by computing feature similarity."""
similarities = []
for layer in self.layers[:num_layers]:
x = layer(x, adj)
x = F.relu(x)
# Compute average cosine similarity between all pairs
x_norm = F.normalize(x, dim=1)
sim = torch.mm(x_norm, x_norm.T)
avg_sim = (sim.sum() - sim.diag().sum()) / (x.size(0) * (x.size(0) - 1))
similarities.append(avg_sim.item())
return similarities
# Example: GCN
gcn = GCN(in_features=128, hidden_features=64, out_features=10)
# Create dummy graph
N = 100
x = torch.randn(N, 128)
adj = torch.randn(N, N)
adj = (adj > 0.5).float() # Binary adjacency
adj = (adj + adj.T) / 2 # Symmetric
adj += torch.eye(N) # Self-loops
logits = gcn(x, adj)
print(f"Input shape: {x.shape}")
print(f"Output shape: {logits.shape}")
print(f"Parameters: {sum(p.numel() for p in gcn.parameters()):,}")
# Analyse over-smoothing
similarities = gcn.compute_over_smoothing(x, adj, num_layers=8)
print(f"\nOver-smoothing analysis:")
for i, sim in enumerate(similarities):
print(f" Layer {i+1}: avg similarity = {sim:.4f}")
7.2 GAT Implementation
class GATLayer(nn.Module):
"""Graph Attention Network layer."""
def __init__(self, in_features, out_features, num_heads=8, dropout=0.6, concat=True):
super().__init__()
assert out_features % num_heads == 0
self.num_heads = num_heads
self.head_dim = out_features // num_heads
self.concat = concat
self.linear = nn.Linear(in_features, out_features, bias=False)
self.attention = nn.Parameter(torch.Tensor(num_heads, 2 * self.head_dim))
if bias:
self.bias = nn.Parameter(torch.Tensor(out_features))
else:
self.bias = None
self.dropout = nn.Dropout(dropout)
self.leaky_relu = nn.LeakyReLU(0.2)
self.reset_parameters()
def reset_parameters(self):
nn.init.xavier_uniform_(self.linear.weight)
nn.init.xavier_uniform_(self.attention)
if self.bias is not None:
nn.init.zeros_(self.bias)
def forward(self, x, adj, return_attention=False):
"""
Forward pass with multi-head attention.
Parameters
----------
x : Tensor (N, in_features)
adj : Tensor (N, N) - adjacency matrix
Returns
-------
h : Tensor (N, out_features)
attention : Tensor (num_heads, N, N), optional
"""
N = x.size(0)
# Linear transformation
h = self.linear(x) # (N, num_heads * head_dim)
h = h.view(N, self.num_heads, self.head_dim) # (N, num_heads, head_dim)
# Compute attention scores
# For each edge (i, j), compute score
attn_input = torch.cat([h.unsqueeze(2).expand(-1, -1, N, -1),
h.unsqueeze(1).expand(-1, N, -1, -1)], dim=-1)
# Attention: (num_heads, N, N)
attn = (attn_input * self.attention.unsqueeze(0).unsqueeze(2).unsqueeze(3)).sum(-1)
attn = self.leaky_relu(attn)
# Mask with adjacency
mask = (adj == 0)
attn = attn.masked_fill(mask.unsqueeze(0), float('-inf'))
# Softmax
attn = F.softmax(attn, dim=-1)
attn = self.dropout(attn)
# Aggregate
h = torch.einsum('hnk,nkd->hnd', attn, h) # (num_heads, N, head_dim)
if self.concat:
h = h.permute(1, 0, 2).contiguous().view(N, -1) # (N, num_heads * head_dim)
else:
h = h.mean(dim=0) # (N, head_dim)
if self.bias is not None:
h = h + self.bias
if return_attention:
return h, attn
else:
return h
class GAT(nn.Module):
"""Multi-layer GAT for node classification."""
def __init__(self, in_features, hidden_features, out_features, num_heads=8, num_layers=2):
super().__init__()
self.layers = nn.ModuleList()
# Input layer
self.layers.append(GATLayer(in_features, hidden_features, num_heads, concat=True))
# Hidden layers
for _ in range(num_layers - 2):
self.layers.append(GATLayer(hidden_features * num_heads, hidden_features, num_heads, concat=True))
# Output layer (no concatenation)
self.layers.append(GATLayer(hidden_features * num_heads, out_features, num_heads=1, concat=False))
def forward(self, x, adj):
attention_weights = []
for i, layer in enumerate(self.layers[:-1]):
x, attn = layer(x, adj, return_attention=True)
x = F.elu(x)
attention_weights.append(attn)
x = self.layers[-1](x, adj)
return x, attention_weights
# Example: GAT
gat = GAT(in_features=128, hidden_features=64, out_features=10, num_heads=8)
x = torch.randn(100, 128)
adj = torch.randn(100, 100)
adj = (adj > 0.5).float()
adj = (adj + adj.T) / 2
adj += torch.eye(100)
logits, attention_weights = gat(x, adj)
print(f"Output shape: {logits.shape}")
print(f"Attention shape: {attention_weights[0].shape}")
print(f"Parameters: {sum(p.numel() for p in gat.parameters()):,}")
7.3 GraphSAGE Implementation
class GraphSAGELayer(nn.Module):
"""GraphSAGE layer with sampling."""
def __init__(self, in_features, out_features, aggregator='mean'):
super().__init__()
self.aggregator = aggregator
self.linear = nn.Linear(in_features + out_features, out_features)
def forward(self, x, adj, num_samples=10):
"""
Forward pass with neighbourhood sampling.
Parameters
----------
x : Tensor (N, in_features)
adj : Tensor (N, N) - adjacency matrix
num_samples : int - number of neighbours to sample
Returns
-------
h : Tensor (N, out_features)
"""
N = x.size(0)
# Sample neighbours for each node
neighbours = []
for v in range(N):
# Get neighbours
neighbours_v = adj[v].nonzero().squeeze()
# Sample if too many
if neighbours_v.numel() > num_samples:
idx = torch.randperm(neighbours_v.numel())[:num_samples]
neighbours_v = neighbours_v[idx]
neighbours.append(neighbours_v)
# Aggregate
aggregated = []
for v, nbrs in enumerate(neighbours):
if nbrs.numel() == 0:
# No neighbours, use self
agg = x[v:v+1]
elif self.aggregator == 'mean':
agg = x[nbrs].mean(dim=0, keepdim=True)
elif self.aggregator == 'max':
agg = x[nbrs].max(dim=0, keepdim=True).values
elif self.aggregator == 'lstm':
# LSTM requires ordered input
perm = torch.randperm(nbrs.numel())
agg = x[nbrs[perm]].unsqueeze(0).sum(dim=0, keepdim=True)
aggregated.append(agg)
aggregated = torch.cat(aggregated, dim=0)
# Concatenate self features and transform
h = torch.cat([x, aggregated], dim=1)
h = F.relu(self.linear(h))
return h
class GraphSAGE(nn.Module):
"""Multi-layer GraphSAGE."""
def __init__(self, in_features, hidden_features, out_features,
num_layers=2, aggregator='mean'):
super().__init__()
self.layers = nn.ModuleList()
self.layers.append(GraphSAGELayer(in_features, hidden_features, aggregator))
for _ in range(num_layers - 2):
self.layers.append(GraphSAGELayer(hidden_features, hidden_features, aggregator))
self.layers.append(GraphSAGELayer(hidden_features, out_features, aggregator))
def forward(self, x, adj):
for layer in self.layers[:-1]:
x = layer(x, adj)
x = self.layers[-1](x, adj)
return x
# Example: GraphSAGE
sage = GraphSAGE(in_features=128, hidden_features=64, out_features=10)
x = torch.randn(50, 128)
adj = torch.randn(50, 50)
adj = (adj > 0.5).float()
adj = (adj + adj.T) / 2
adj += torch.eye(50)
logits = sage(x, adj)
print(f"Output shape: {logits.shape}")
print(f"Parameters: {sum(p.numel() for p in sage.parameters()):,}")
8. Summary
-
Message passing is the universal framework for GNNs, aggregating information from local neighbourhoods.
-
GCN approximates spectral graph convolutions with a first-order polynomial, enabling efficient training.
-
GraphSAGE enables inductive learning through neighbourhood sampling and multiple aggregation functions.
-
GAT learns attention weights for edges, enabling adaptive neighbourhood aggregation.
-
Over-squashing is a fundamental limitation; solutions include graph rewiring and graph transformers.
References
- Kipf, T. N., & Welling, M. (2017). Semi-supervised classification with graph convolutional networks. ICLR.
- Hamilton, W. L., Ying, Z., & Leskovec, J. (2017). Inductive representation learning on large graphs. NeurIPS.
- Veličković, P., et al. (2018). Graph attention networks. ICLR.
- Gilmer, J., et al. (2017). Neural message passing for quantum chemistry. ICML.
- Alon, U., & Yahav, E. (2021). On the bottleneck of graph neural networks and its practical implications. ICLR.
- Ying, C., et al. (2021). Do transformers really perform bad for graph representation? NeurIPS.
- Rampášek, L., et al. (2022). Recipe for a general, powerful, scalable graph transformer. NeurIPS.
- Bronstein, M. M., et al. (2017). Geometric deep learning: Going beyond Euclidean data. IEEE Signal Processing Magazine.