Graph neural networks

Problem

A convolution slides one small filter over a grid because every pixel has the same neighbours in the same positions: up, down, left, right. A transformer works on a sequence, where position is an index. Many datasets have neither. Atoms in a molecule have between one and four bonds. A paper cites any number of others. A user in a social network has three friends or three thousand. There is no "left neighbour", and renumbering the nodes changes nothing about the graph, so a model that depends on the numbering is wrong.

The questions asked of such data are of three kinds: about a node (what topic is this paper, is this account a bot), about an edge (will these two users connect, do these two proteins interact) and about a whole graph (is this molecule toxic). A graph neural network learns a vector for every node from its own features and from its neighbourhood, using the same weights for every node whatever its degree, and in a way that does not depend on how the nodes are numbered.

Intuition

Each node starts with a feature vector: an atom's element, a paper's bag of words. In one layer, every node collects its neighbours' vectors, combines them with its own and passes the result through a small learned function. After one layer a node's vector summarizes its immediate neighbours; after two, its neighbours' neighbours; after k layers, everything within k hops. A paper about graph learning that cites three other graph papers ends up with a vector close to theirs.

The combining step has to treat the neighbours as a set, not a list, so it uses an order-free reduction: a sum, a mean, a maximum, or a weighted sum whose weights are themselves computed from the vectors. Everything else is a linear layer applied to every node with shared weights. The variants differ mostly in how they weight the neighbours, and the cost of the whole idea shows up when the stack gets deep: averaging over and over makes every node look like every other.

Mechanics

A graph as two matrices

Number the n nodes 0 to n − 1. The adjacency matrix A is n × n with A_ij = 1 when there is an edge between i and j and 0 otherwise; for an undirected graph it is symmetric. The feature matrix X is n × f, one row per node. The degree d_i = Σ_j A_ij counts node i's neighbours, and D is the diagonal matrix of degrees. Renumbering the nodes with a permutation matrix P turns A into PAPᵀ and X into PX; a layer is permutation equivariant when its output turns into P times the old output, which is the property every layer on this page has.

Real graphs are sparse: a citation graph with millions of nodes has a few dozen edges per node, so A is stored as an edge index, a 2 × E tensor of (source, target) pairs, with each undirected edge listed in both directions. That is a COO sparse matrix without the values, and PyTorch's torch.sparse_coo_tensor takes exactly that layout.

Message passing

The general layer, from the message passing neural network framework, has three parts. For each edge j → i a message m_ji = φ(h_i, h_j, e_ji) is computed from the two endpoints and any edge features. Each node aggregates its incoming messages with a permutation-invariant reduction ⊕ (sum, mean, max). Then it updates:

On an edge index this is three tensor operations: gather the endpoint rows with x[src] (an index_select), compute one message per edge, and scatter the messages back into their target rows with index_add_ or scatter_reduce_. The cost is O(E · f) instead of the O(n² · f) of a dense matrix product, which is what makes graphs with millions of nodes trainable.

GCN: normalized propagation, derived

The simplest choice is to sum the neighbours: AX. It has two defects. A node's own features are dropped, since A_ii = 0, and the scale of the result grows with the degree, so a hub's vector is hundreds of times larger than a leaf's. Adding self-loops fixes the first: à = A + I, with degrees d̃_i = d_i + 1. Dividing by the degree fixes the second: D̃⁻¹ÃX is the mean over the closed neighbourhood. That mean is not symmetric, though: it weights the edge i ← j by 1/d̃_i and the edge j ← i by 1/d̃_j. The symmetric compromise splits the normalization between both ends:

Kipf and Welling reached the same matrix from spectral graph theory. The normalized Laplacian L = I − D^−½AD^−½ is symmetric with eigenvalues in [0, 2], and its eigenvectors play the role of a Fourier basis on the graph; a spectral filter multiplies each frequency by g_θ(λ). Approximating the filter by a first-order Chebyshev polynomial, with λ_max ≈ 2, gives g ⋆ x ≈ θ₀x + θ₁(L − I)x = θ₀x − θ₁D^−½AD^−½x, a filter that only looks one hop away. Tying the two parameters, θ = θ₀ = −θ₁, leaves θ(I + D^−½AD^−½)x. That operator has eigenvalues in [0, 2], so stacking it can blow activations up; the renormalization trick replaces it with D̃^−½ÃD̃^−½, which is  above and whose eigenvalues lie in (−1, 1]. A filter with one parameter per pair of input and output features, applied to all f inputs at once, is the matrix W. The paper's two-layer model for node classification is Z = softmax( ReLU(ÂXW⁽⁰⁾) W⁽¹⁾), trained with cross-entropy on the few labelled nodes.

Read as message passing, GCN's message is the transformed neighbour W^ᵀh_j scaled by the fixed coefficient 1/√(d̃_id̃_j), the aggregation is a sum, and the update is the nonlinearity. The coefficients depend only on the graph, so they can be computed once.

GraphSAGE: sample, aggregate, concatenate

GCN as written needs the whole graph in every step, and it mixes a node's own vector into the same average as its neighbours. GraphSAGE changes both. Its layer keeps the self and the neighbourhood apart and concatenates them:

so the model can learn how much to trust the node versus its context. The aggregator can be a mean, an elementwise max after a small MLP ("pooling"), or an LSTM over a random ordering. The second change is training on samples: each layer uses a fixed number of randomly sampled neighbours (the paper uses 25 at the first hop and 10 at the second), so the cost of a minibatch of target nodes is bounded however large the graph is, and the learned function applies to nodes and graphs never seen in training. This is the inductive setting, where GCN as originally trained is transductive.

GAT: attention over the neighbourhood

GCN gives every neighbour a weight fixed by degrees. A graph attention network learns the weights from the features, in the spirit of attention but restricted to the edges. With z_i = W h_i and a learned vector a of length 2f′:

The softmax runs only over the node's own neighbourhood including itself, so a degree-2 node and a degree-200 node both get weights that sum to 1. The paper uses a LeakyReLU slope of 0.2 (PyTorch's F.leaky_relu defaults to 0.01), and K independent heads whose outputs are concatenated in hidden layers and averaged in the output layer, exactly as multi-head attention splits a transformer layer. Because a splits into a part that sees z_i and a part that sees z_j, the score is a monotone function of a₂ᵀz_j plus a constant per node, so every node ranks shared neighbours in the same order; GATv2 applies the nonlinearity before a to make the attention dynamic.

Over-smoothing

Each layer multiplies by Â, and  is a smoothing operator: its largest eigenvalue is 1, with eigenvector proportional to (√d̃₁, …, √d̃_n) on a connected graph, and with self-loops every other eigenvalue has magnitude below 1. Applying it k times shrinks every other component by its eigenvalue to the k-th power, so Â^kX tends to a rank-one matrix: row i becomes √d̃_i times one shared vector. Nodes then differ only by their degree, and the features that distinguished them are gone. Li, Han and Wu showed this is Laplacian smoothing, the thing that makes shallow GCNs work, taken too far. Weights and nonlinearities between the layers slow it but do not stop it in practice, which is why most GNNs use two to four layers.

The remedies keep some of each node's own signal alive: residual connections or a jumping-knowledge concatenation of every layer's output, normalization schemes that re-spread the rows, randomly dropping edges during training, and decoupling the depth of propagation from the depth of the network. A related problem, over-squashing, is the opposite: information from an exponentially growing number of nodes k hops away has to pass through a fixed-size vector, so long-range dependencies are hard to learn even before smoothing sets in.

Batching graphs

For graph-level tasks the data is many small graphs of different sizes, so they cannot be stacked into one tensor. The standard trick is to treat a minibatch as one big disconnected graph. Concatenate the feature matrices, add to each graph's edge index the number of nodes that come before it, and keep a batch vector giving each node's graph. The adjacency of the result is block-diagonal, so no message ever crosses between graphs and every layer above runs unchanged. At the end a readout pools each graph's nodes into one vector with a sum, mean or max over the nodes sharing a batch index; that pooling is one more index_add_ or scatter_reduce_. Sum pooling keeps the count of nodes, which mean pooling loses; the GIN paper shows sum aggregation is what lets a GNN be as discriminating as the Weisfeiler-Lehman graph isomorphism test.

PyTorch v2.15.0 has no graph layers of its own; libraries such as PyTorch Geometric and DGL build them from exactly these index and scatter operations.

Worked example

Five nodes and five undirected edges: A–B, A–C, B–C, C–D, D–E, a triangle with a tail. Each node has two features, chosen so that A and B differ.

Implementation

A GCN layer, a neighbourhood softmax for GAT and mean pooling for a batch, in plain PyTorch on an edge index. The example graph's tensors reproduce the worked numbers above; nn.Linear computes x·Wᵀ, so the page's W goes in transposed.

import torch
import torch.nn as nn

def gcn_layer(x, edge_index, lin):
    n = x.size(0)
    loops = torch.arange(n)
    src = torch.cat([edge_index[0], edge_index[1], loops])   # both directions plus self-loops
    dst = torch.cat([edge_index[1], edge_index[0], loops])
    deg = torch.zeros(n).index_add_(0, dst, torch.ones(dst.numel()))
    coef = deg[src].rsqrt() * deg[dst].rsqrt()                # 1 / sqrt(d_s * d_t)
    h = lin(x)
    return torch.zeros_like(h).index_add_(0, dst, coef.unsqueeze(1) * h[src])

def neighbourhood_softmax(scores, dst, n):
    top = torch.full((n,), float("-inf")).scatter_reduce(0, dst, scores, reduce="amax")
    e = (scores - top[dst]).exp()                             # subtract each node's max for stability
    return e / torch.zeros(n).index_add_(0, dst, e)[dst]

def mean_pool(h, batch, num_graphs):
    total = torch.zeros(num_graphs, h.size(1)).index_add_(0, batch, h)
    count = torch.zeros(num_graphs).index_add_(0, batch, torch.ones(batch.numel()))
    return total / count.unsqueeze(1)

x = torch.tensor([[1., 0.], [0., 1.], [1., 1.], [0., 2.], [2., 0.]])
edge_index = torch.tensor([[0, 0, 1, 2, 3], [1, 2, 2, 3, 4]])
lin = nn.Linear(2, 2, bias=False)
with torch.no_grad():
    lin.weight.copy_(torch.tensor([[0.5, -1.0], [1.0, 0.5]]).T)
print(gcn_layer(x, edge_index, lin))       # row 2 is C: [1.3854, 0.0193]

The docstrings of index_add_ and scatter_reduce_ warn that they may behave nondeterministically on a CUDA device, so sums can differ in the last bits from run to run; torch.use_deterministic_algorithms(True) makes PyTorch pick a deterministic implementation or raise an error. For a fixed graph the coefficients can be computed once and the whole propagation stored as one sparse matrix, after which each layer is a single torch.sparse.mm. Libraries such as PyTorch Geometric and DGL package these layers, neighbour samplers and batching collators; the operations underneath are the ones above.

Tradeoffs

• GCN vs GraphSAGE vs GAT. GCN is the cheapest, with fixed coefficients and one matrix per layer, and is a strong baseline on citation-style graphs where neighbours tend to share labels. GraphSAGE separates the self from the neighbourhood and scales through sampling. GAT learns neighbour weights at the cost of one score per edge per head, which helps when some neighbours matter more than others.
• Depth. More layers see further but over-smooth and over-squash; two to four layers are typical, with residual connections or jumping knowledge to go deeper.
• Full graph vs sampling. Full-batch training is exact but needs the whole graph and all activations in memory. Neighbour sampling bounds memory per minibatch but adds variance, and the number of sampled nodes still grows exponentially with depth.
• Expressiveness. A message-passing GNN cannot tell apart nodes that the 1-WL test cannot, like A and B in the worked example. Sum aggregation (GIN) reaches that bound; mean and max fall below it. Going further needs positional encodings, subgraph features or higher-order methods.
• Homophily. Averaging neighbours assumes they resemble the node. On heterophilous graphs, where connected nodes tend to differ (fraudsters and their victims), plain GCN can do worse than an MLP that ignores the edges; keeping the self apart helps.
• Graph transformers. Full attention over all nodes avoids over-squashing but costs O(n²) and needs positional encodings to know the structure at all; most such models still mix in local message passing.

Connections

References

• Kipf and Welling. Semi-Supervised Classification with Graph Convolutional Networks. ICLR 2017. arXiv:1609.02907
• Veličković, Cucurull, Casanova, Romero, Liò and Bengio. Graph Attention Networks. ICLR 2018. arXiv:1710.10903
• Hamilton, Ying and Leskovec. Inductive Representation Learning on Large Graphs. NeurIPS 2017. arXiv:1706.02216
• Gilmer, Schoenholz, Riley, Vinyals and Dahl. Neural Message Passing for Quantum Chemistry. ICML 2017. arXiv:1704.01212
• Li, Han and Wu. Deeper Insights into Graph Convolutional Networks for Semi-Supervised Learning. AAAI 2018. arXiv:1801.07606
• Brody, Alon and Yahav. How Attentive are Graph Attention Networks? ICLR 2022. arXiv:2105.14491
• Xu, Hu, Leskovec and Jegelka. How Powerful are Graph Neural Networks? ICLR 2019. arXiv:1810.00826
• PyTorch v2.15.0: Tensor.index_add_, Tensor.scatter_reduce_, torch.sparse_coo_tensor, Tensor.to_sparse_coo.