Graph Classification With Graph Neural Networks

Node classification gives every node a label and the graph stays fixed. Graph classification gives the entire graph one label — is this molecule toxic, is this binary malicious, is this scene a kitchen — and every example is a different graph with a different number of nodes. That shape difference is what the pooling layer exists to solve, and the pooling choice bounds what the model can ever learn.

One label for a variable-size object

A message-passing GNN layer updates each node from its neighbours. After K layers every node carries a vector summarising its K-hop neighbourhood. You now have a matrix with one row per node, of a height that changes from example to example, and you need a fixed-length vector to feed a classifier.

The batching also works differently and it is worth knowing before you read any implementation. Graphs of different sizes cannot be stacked into a tensor, so libraries build a batch by placing the graphs into one big disconnected graph — block-diagonal adjacency — and carrying a vector that records which graph each node belongs to. Message passing then runs once over the whole batch, and because there are no edges between the blocks, no information crosses between examples. The readout is a scatter operation keyed by that vector.

The function that does that — the READOUT — has one hard requirement: it must be permutation invariant. Node 3 in one graph and node 3 in another mean nothing in common, and relabelling the nodes of a graph must not change its prediction. That rules out anything positional (concatenation, an RNN over the rows, a dense layer on a flattened matrix) and leaves aggregation over the multiset of node vectors: sum, mean, max, or something learned that is itself invariant.

The readout, worked

Take two graphs whose node vectors after the final GNN layer are two-dimensional, and where the second graph is simply two disjoint copies of the first:

Graph G1  (4 nodes)          Graph G2  (8 nodes)
  h1 = [1, 0]                  four nodes at [1, 0]
  h2 = [1, 0]                  four nodes at [0, 1]
  h3 = [0, 1]
  h4 = [0, 1]

MEAN readout
  G1: [ (1+1+0+0)/4 , (0+0+1+1)/4 ] = [0.50, 0.50]
  G2: [ 4/8 , 4/8 ]                  = [0.50, 0.50]     ← identical

SUM readout
  G1: [2, 2]
  G2: [4, 4]                                            ← distinguishable

MAX readout
  G1: [1, 1]
  G2: [1, 1]                                            ← identical

Mean and max cannot tell a graph from two copies of it. Sum can, because sum preserves multiplicity rather than proportion. That is not a quirk of this example; it is the general result in How Powerful are Graph Neural Networks? (Xu et al., ICLR 2019), which proves sum is injective over multisets while mean captures only the distribution and max only the underlying set. If your classes differ in size or in how many of a substructure they contain — count the benzene rings, count the loops — mean pooling has thrown that away before your classifier ever sees it.

Sum is not free. Its magnitude grows with the number of nodes, so a dataset containing 10-node and 10,000-node graphs feeds the classifier activations three orders of magnitude apart and training becomes unstable. The practical resolution is usually sum pooling with normalisation, or concatenating sum and mean so both the count and the proportion survive. The GIN paper additionally concatenates the readout from every layer rather than only the last, so shallow structural information is not smeared away by the deep layers.

The ceiling: Weisfeiler-Lehman

The same paper proves the harder result: any GNN of this aggregate-and-update form is at most as powerful at telling graphs apart as the 1-dimensional Weisfeiler-Lehman colour-refinement test. If WL cannot distinguish two graphs, no amount of width, depth or training will let a message-passing GNN distinguish them either.

Here is the standard pair that makes this concrete:

Graph A: one 6-cycle          Graph B: two disjoint triangles
  1-2-3-4-5-6-1                   1-2-3-1   and   4-5-6-4

Every node in both graphs has degree 2.
Every node's neighbours have degree 2.
Refine as many rounds as you like: all twelve nodes keep the
same colour. WL cannot separate them, so no message-passing GNN
can either — the readouts are identical whatever the weights.

These two graphs are not isomorphic and the difference is obvious to a human — one is connected, the other is not. If your task depends on counting cycles or on connectivity, a plain GNN is the wrong tool and the fix is to give it information it cannot derive: add cycle counts or other structural descriptors as node features, add random node identifiers, or use a higher-order architecture. GIN itself reaches the WL bound (with sum aggregation and an injective update, using the (1 + eps) self-weight so a node’s own vector is not confusable with its neighbours’) but does not exceed it.

Hierarchical pooling

Flat readout throws away the whole hierarchy in one step. Hierarchical pooling coarsens the graph in stages instead, in imitation of the striding a CNN does over an image. Two families are in common use:

• Soft cluster assignment. DiffPool (Ying et al., NeurIPS 2018) learns a soft assignment matrix mapping N nodes to a smaller fixed number of clusters at each layer, then contracts the graph. It is differentiable end to end, and it is dense — the assignment matrix is N × clusters, which is quadratic-ish in memory and the reason it does not scale to large graphs.
• Node dropping. Top-k pooling and self-attention pooling score every node, keep the highest-scoring fraction, and discard the rest. Sparse and cheap, but genuinely lossy: a dropped node’s information is gone, and dropping can disconnect the graph so later layers pass no messages across the gap.

Why reported gains often do not reproduce

The standard graph-classification benchmarks are small — many of the widely used molecular and social datasets hold a few hundred to a few thousand graphs. At that size, the gap between a careful evaluation and a careless one is larger than the gap between architectures. A Fair Comparison of Graph Neural Networks for Graph Classification (Errica et al., ICLR 2020) re-ran a set of published models under a single protocol and found that structure-agnostic baselines — models that ignore the edges and only aggregate node features — were competitive with, and on some datasets better than, the GNNs they were compared against.

The practical reading of that is not that GNNs do not work. It is that on a small benchmark you must fix the split protocol before you look at results, use nested cross-validation rather than tuning on the test fold, report variance across seeds, and always run the edges-removed baseline. If removing the graph does not hurt, the graph was not carrying the signal and everything on this page is beside the point for your dataset.

Depth is the other recurring trap. Stacking layers to widen the receptive field makes node vectors converge toward each other, which removes exactly the differences the readout needs — see GNN oversmoothing. Most published graph classifiers are two to five layers deep for that reason, not from a lack of ambition.