Inside the Graph Neural Network

Molecules are atoms joined by bonds. A social network is people joined by friendships. Papers cite papers, proteins bind proteins, roads join crossings. All of these are graphs: a set of nodes, and a set of edges that connect pairs of nodes.

This graph has 8 nodes, numbered 0 to 7, and 10 edges. The edges have no direction: if 0 is linked to 1, then 1 is linked to 0.

Every node also carries a vector of features. In a molecule they might describe the element of each atom. In a citation network, the words in each paper.

Here every node has two features, so each node can be drawn as a point in the plane on the right. The number next to a point is its node.

A graph holds two kinds of information: what each node is, and what it is connected to.

Say the nodes belong to two classes, blue and rose, and the task is to predict the class of every node. From the features alone, that is hard.

Node 1 is blue, but its point sits among the rose ones. Node 6 is rose, but sits next to the blue ones. No straight line in the plane puts all blue points on one side and all rose points on the other.

The edges tell a different story. Node 1 is linked to nodes 0 and 3, and both are blue. Nine of the ten edges join two nodes of the same class. Only the bridge between 3 and 4 crosses over.

Many real graphs work this way. Friends share interests, and papers cite papers on the same topic. The pattern is called homophily, and it is what the simplest graph neural networks rely on.

A graph neural network decides about a node using its own features and the features of its neighbors.

Stack the feature vectors as rows and you get the feature matrix $X$. Row $i$ belongs to node $i$. With $n$ nodes and $d$ features its shape is (n, d), here (8, 2).

The edges go into the adjacency matrix $A$, of shape (n, n). Entry $A_{ij}$ is 1 when nodes $i$ and $j$ share an edge, and 0 otherwise.

Row 3 has ones in columns 0, 1, 2, 4. Those are the neighbors of node 3, written $\mathcal{N}(3)$.

Edges have no direction, so $A$ is symmetric: row 3 and column 3 hold the same entries, and $A = A^{\top}$.

Adding up a row counts the neighbors of that node. This is its degree $d_v$. Node 3 has degree 4, node 7 only 1.

Real graphs are sparse. The Cora citation graph, a standard benchmark, has 2,708 papers. Its adjacency matrix would have 7,333,264 entries, and only 10,556 of them are ones: 0.14%.

So libraries store only the edges. edge_index has shape (2, E): a row of source nodes and a row of target nodes, with every undirected edge listed in both directions. It holds the same information as $A$.

Nothing in the graph says which node is number 0. The numbers are names we gave the nodes so we could write them into matrices. A model that gives different answers for different names has learned something that is not there.

Renumber the nodes: node 3 becomes 0, node 0 becomes 5, and so on. The graph is the same, but the rows of $X$ move, and the rows and columns of $A$ move with them. The amber row and column belong to the ringed node.

A permutation matrix $P$ describes the renumbering. It has a single 1 in every row and column. Multiplying by $P$ on the left reorders rows, and multiplying by $P^{\top}$ on the right reorders columns. That is why the rows of $A$ move first, and then its columns.

This rules out the obvious approach of flattening $A$ into one long vector and feeding it to an ordinary network. After the renumbering, 20 of its 64 entries are different, so the network would see a different input for the same graph.

There are $n!$ ways to number $n$ nodes, 40,320 for this small graph. An ordinary network would have to learn from examples that all of them mean the same thing, and the count explodes as graphs grow.

So the requirement is built in. A layer that produces one vector per node must be permutation equivariant: renumber the input, and the output rows are renumbered in the same way, with nothing else changed.

A model that produces one answer for the whole graph, such as whether a molecule is toxic, must be invariant: the answer does not change at all.

Every layer in this essay is built to satisfy this rule.

Take node 1. Its own features say rose, but its neighbors, 0 and 3, are both blue. The idea behind every graph neural network is to let each node collect the vectors of its neighbors and update its own vector with what it finds.

The collecting step is called aggregation. It has one hard rule from the previous chapter: it may not depend on the order in which the neighbors are numbered.

The simplest aggregation is a sum. Geometrically, the neighbor vectors are laid tip to tail, and the sum is the arrow from the start of the first to the end of the last.

Addition does not care about order: $x_0 + x_3$ equals $x_3 + x_0$. That is exactly the property we need.

A sum grows with the number of neighbors. A node with fifty neighbors gets a vector about twenty-five times longer than a node with two, even if the neighbors look alike.

Dividing by the count gives the mean. In the plane, the mean of two points is the midpoint between them. The mean of many points is their centroid, the center of mass they would balance on.

Averaging only the neighbors would throw away the node's own features. So the node counts itself as one of its neighbors. The mean is now the centroid of the shaded triangle.

Node 1 has moved from [0.90, −0.40] into the blue region. Its neighbors outvoted its own features.

Now every node does the same, all at once, each using the old vectors of its neighbors. Every point moves to the centroid of its own neighborhood.

Node 6 crosses over too, from [−0.70, −0.20] to [0.70, −0.17]. A straight line now separates the two classes, which no line could do before.

Averaging over neighbors smooths the features along the edges. When neighbors tend to share a class, smoothing removes noise.

Row $v$ of $A$ has a 1 for every neighbor of $v$ and a 0 everywhere else. So row $v$ of the product $AX$ adds up exactly the feature vectors of the neighbors of $v$.

To let every node count itself, add the identity matrix. This puts a 1 on the diagonal, a self-loop on every node. Papers write it with a tilde.

For the mean, divide every row by its number of ones. Collect those counts in a diagonal matrix $\tilde{D}$, with $\tilde{D}_{vv} = \sum_u \tilde{A}_{vu}$. Multiplying by its inverse divides each row by its count.

Row 1 of $H$is the centroid from before. Every other row is the centroid of that node's neighborhood.

Graph libraries usually skip $\tilde{A}$ altogether. Each edge copies the vector of its source node, a message, to its target node, and each node reduces the messages it receives. The cost grows with the number of edges, not with $n^2$. This is why the whole family is called message passing.

Equivariance now comes for free. Renumbering turns $\tilde{A}X$ into $(P\tilde{A}P^{\top})(PX)$, and because $P^{\top}P = I$ that is $P(\tilde{A}X)$: the same rows in the new order.

Averaging has no parameters, so there is nothing to learn. A graph convolutional network, or GCN, adds the two ingredients of an ordinary neural network layer: a weight matrix and a nonlinearity. It was introduced by Kipf and Welling in 2017 and is still the usual starting point.

Start again from the raw features. The grey arrows mark the two unit vectors along the axes, so we can see what happens to the plane itself.

First, aggregate. A GCN uses a slightly different average. Instead of dividing by the degree of the receiving node, it divides each message by the square root of both degrees.

A node with many neighbors now sends weaker messages, because its features are shared with many nodes. The weights in a row no longer add up to exactly 1, but the result is close to the mean, and the matrix stays symmetric.

Look at the grid: it does not move. Aggregation is not a function of the plane. Where a point goes depends on its neighbors, not on where it is.

Then, transform. Every row is multiplied by the same weight matrix $W$, the same for every node.

This is a linear map of the whole plane. Grid lines stay straight, parallel and evenly spaced. The two grey arrows land on the rows of $W$, and every point moves along with the grid.

Here $W$ is 2 by 2 so we can draw it. In practice it also changes the width. On Cora the first layer maps 1433 features to 16.

Finally, bend. ReLU replaces every negative coordinate with zero, which folds the plane onto the positive quadrant. Points left of the vertical axis land on it, points below the horizontal axis land on that, and the bottom-left quarter collapses onto the origin.

The fold sorts the nodes. The blue points end up on or right next to the vertical axis, the rose points on or next to the horizontal one. Node 6 started at [−0.70, −0.20] and ends at [0.60, 0.00].

Together, that is one GCN layer.

Papers often write the same layer for a single node. The sum runs over the node and its neighbors.

The name comes from images. A convolution applies the same small filter at every pixel and its neighbors. A GCN applies the same $W$ at every node and its neighbors. An image is a graph too: a grid where each pixel links to the pixels around it.

Aggregation moves information between nodes. The weights and the nonlinearity transform each node on its own, with the same parameters everywhere.

The layer first builds $\hat{A}$. d[:, None] * d[None, :] is an (n, n) matrix whose entry $(v, u)$ is $\tilde{d}_v \tilde{d}_u$, so one division normalizes every entry at once.

nn.Linear without a bias is the multiplication by $W$. It works on every row at once and turns shape (n, d) into (n, d′).

Matrix products can be grouped either way: $\hat{A}(HW) = (\hat{A}H)W$. Applying $W$ first is cheaper when it shrinks the vectors, as with the 1433 to 16 layer on Cora.

Multiplying by A_norm does the aggregation. The number of rows stays $n$, one vector per node, and the rows are in the same order as the input. The layer is equivariant.

A typical GCN stacks two of these layers with a ReLU between them. The last layer outputs one score per class for every node, and training uses cross-entropy on the nodes whose labels are known.

Libraries such as PyTorch Geometric provide this layer as GCNConv, working from edge_index instead of a dense matrix, with a bias added.

One layer mixes each node with its direct neighbors. After one layer, node 0 holds information from nodes 1, 2, 3 and from itself, and from nothing else.

In the second layer, node 0 aggregates again. But node 3 already took in node 4 during the first layer, so information from node 4 now reaches node 0 through node 3.

After $k$ layers, a node has heard from every node at most $k$edges away. This set is the node's receptive field, the same idea as in a convolutional network for images.

The matrix view says the same. Entry $(v, u)$ of $\tilde{A}^k$ counts the walks of length $k$ from $v$ to $u$, where a self-loop lets a walk pause. It is nonzero exactly when $u$ is at most $k$ hops away.

Node 7 is 4 hops from node 0, so it takes 4 layers before node 0 can depend on it at all. Now every entry in row 0 is filled.

Change the number of layers and watch both views grow together.

Most GCNs use two or three layers. That sounds like very few, but real graphs tend to be small worlds: friends of friends of friends already cover a large part of a social network.

So why not use many layers? Repeat the mean aggregation of chapter 2 without any weights and watch what happens. The crosses mark the mean of each class.

One round cuts the average distance of a point to its class mean from 1.21 to 0.45, while the gap between the two class means stays at 1.85. This is the useful part: each class pulls together.

After four rounds the classes are tighter still. But the gap between them has started to shrink too, down to 1.36. The bridge nodes 3 and 4 drift toward each other and drag their classes along.

After 32 rounds, every node sits at nearly the same point. This is over-smoothing, and it is not a coincidence of this graph.

Each row of $\tilde{D}^{-1}\tilde{A}$ sums to 1, so it describes a random walk that moves to a random neighbor or stays put. Run the walk long enough and where it ends no longer depends on where it started. Every node ends up with the same mixture of all features, weighted by degree.

Drag the slider. The differences between nodes never vanish completely, but after the first round they shrink by about 10 percent every round, until they are smaller than the noise and the bias terms that real layers add.

Learned weights and nonlinearities between the layers do not prevent this in general. Deeper models use residual connections, so each layer only adds a change, $h^{(l+1)} = h^{(l)} + f(h^{(l)})$, or combine the outputs of all layers, or normalize the vectors to keep them apart.

Depth widens what a node can see, and blurs what it sees.

Node 4 sits on the bridge between the two groups. Its neighbors are 3, 5, 6, and node 3 belongs to the other class.

The mean, and the GCN, give every neighbor a weight that depends only on the graph. With itself included, node 4 takes a quarter from each, and its new vector [0.42, 0.25] is pulled toward node 3.

A graph attention network, or GAT, lets the model choose the weights. For every neighbor $u$ it computes a score from the two vectors involved.

$\Vert$ means concatenation: the two vectors are placed end to end. $a$ is a learned vector that turns the pair into one number. Here $W$ is the identity and $a = [0.3,\ 0.3,\ 0.6,\ -0.6]$, picked by hand. Training would learn both.

Softmax over the neighbors turns the scores into weights that are positive and sum to 1. Node 5 gets 0.64, node 3 only 0.08. The edge widths show the weights.

One detail matters when you read later papers. Split $a$ into $a_1$, which multiplies $W h_v$, and $a_2$, which multiplies $W h_u$. LeakyReLU and softmax never change the order of their inputs, so neighbors are ranked by $a_2^{\top} W h_u$ alone: a neighbor that ranks high for one node ranks high for every node. GATv2 applies $a$ after the nonlinearity, so the ranking can depend on who is asking.

The new vector is the weighted sum of the neighbors. It moves from the mean to [1.24, −0.33], deep among its own class.

Because the weights are positive and sum to 1, this is a convex combination, and the result always lies inside the shaded region, the convex hull of the neighbors. The mean always lands on the plain average of the points. Attention can choose any point in the hull, but never one outside it.

Attention picks a point inside the hull of the neighbors. The mean always picks their average.

If this looks familiar, it should. A Transformer does the same with queries and keys instead of concatenation, and lets every token attend to every other token. That is attention on a complete graph, the dashed edges.

Set the scores of non-neighbors to $-\infty$ before the softmax, and Transformer attention only flows along the edges of the graph, as in GAT. The borrowing goes both ways: GAT usually runs several attention heads side by side and concatenates their outputs, just like a Transformer.

An aggregation receives a multiset: a collection in which order does not matter but the same vector may appear more than once. Ideally, different multisets give different results.

Compare a node with neighbors $u$ and $w$ to a node with two copies of each. The mean lands on the same centroid. The element-wise maximum, which takes the largest value in every coordinate, lands on the same corner. Only the sum notices: it doubles.

Now add a small vector $c$. It is below the maximum in every coordinate, so the maximum does not move. The mean does move, and so does the sum.

The pattern is general. The maximum ignores how often a vector occurs, and can miss a vector entirely, like $c$. The mean sees the proportions of the neighbors but not how many there are. The sum keeps the count.

The sum remembers how many neighbors there are. The mean and the maximum forget it.

The graph isomorphism network, or GIN, builds on this. It sums the neighbors and passes the result through a small MLP instead of a single linear map.

The factor $1 + \epsilon$keeps a node's own vector from blending in with its neighbors. Xu and colleagues proved in 2019 that when node features come from a countable set, such as categories, a suitable transformation followed by a sum can tell apart any two different multisets. No mean or maximum can do that. With enough capacity, GIN is as powerful as message passing gets.

That is about what a model can tell apart, not what it learns best. When the number of neighbors is mostly noise, a mean often works better in practice.

How much can message passing tell apart at best? A graph algorithm from 1968 gives the answer. It is called color refinement, or the Weisfeiler-Leman test, WL for short.

Every node starts with the same color, as if all features were equal. Only the structure of the graph is left.

In each round, a node gets a new color determined by its own color and the multiset of its neighbors' colors. Nodes with the same old color and the same neighbor colors get the same new color.

After one round the colors simply count neighbors: 4 colors, one for each degree.

In the second round, nodes with the same degree split up if their neighbors have different degrees. There are now 7 colors.

Nodes 1 and 2 still share a color. They are mirror images: swap them and the graph does not change, so nothing could tell them apart.

A third round changes nothing, so the colors are stable.

Compare the code with a GNN layer. The new color depends on the node's own color and the multiset of its neighbors' colors, exactly like aggregate and update. The hash is a perfect aggregation that never confuses two multisets.

It follows that if WL, started from the node features, gives two nodes the same color, every message-passing GNN gives them the same vector, whatever its weights. Xu and colleagues and Morris and colleagues proved this independently in 2019.

Here are two graphs that are clearly different. One is a single ring of six nodes. The other is two separate triangles. Every node in both has exactly two neighbors.

Run the refinement and nothing ever splits. Every node sees two neighbors with its own color, in every round. The color histograms of the two graphs are identical, so any sum or mean over the nodes is too.

No message-passing GNN can tell these graphs apart, however deep or wide, as long as the node features are equal. In particular, it cannot count triangles.

Stronger models break the symmetry. Some add features that describe a node's position in the graph, others pass messages between pairs of nodes or count small subgraphs.

Message passing sees a graph only through what each node can learn from its neighbors.