Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI ๐Ÿ‘จ๐Ÿปโ€๐Ÿ’ป๐Ÿค–

DiffPool: Learning Hierarchical Graph Pooling

6 minute read

Published:

TL;DR: DiffPool (Ying et al., 2018) uses a second GNN to learn soft cluster assignments โ€” each node is assigned fractionally to each cluster. The coarsened graph is then passed to another GNN, building a hierarchy from fine-grained nodes to coarse super-nodes. It's end-to-end differentiable and captures multi-scale structure that flat pooling misses.
DiffPool hierarchical pooling
DiffPool: differentiable graph pooling with soft cluster assignments (Ying et al., 2018)

Intuition First: Hierarchical Pooling as a Convolutional Pyramid

In image CNNs, max-pooling after each conv layer progressively coarsens the spatial grid: 32ร—32 โ†’ 16ร—16 โ†’ 8ร—8 โ†’ 1ร—1. This captures multi-scale features โ€” edges at fine scale, shapes at coarser scale, objects at the coarsest. DiffPool brings this pyramid idea to graphs: nodes โ†’ clusters โ†’ super-clusters โ†’ graph. The key challenge is that unlike pixels, graph nodes have no fixed spatial ordering, so the pooling must be learned and permutation-invariant.

Key Insight: DiffPool's assignment matrix S is the graph analogue of a pooling kernel. But instead of a fixed stride-2 window, S is learned โ€” the model discovers which nodes should be clustered together based on both features and graph structure, without any external supervision on the clustering.
Level 0 (8 nodes) Sโฝโฐโพ Level 1 (4 clusters) Sโฝยนโพ Level 2 (2) global pool โ†’ h_G
DiffPool hierarchically coarsens the graph: 8 nodes โ†’ 4 clusters โ†’ 2 super-clusters โ†’ graph embedding. Each level learns which nodes to merge via a soft assignment matrix S.

The Limitation of Flat Global Pooling

Global mean/sum/max pooling jumps directly from N node embeddings to a single graph embedding. For graphs with hierarchical structure โ€” like molecules (atoms โ†’ functional groups โ†’ whole molecule) or social networks (people โ†’ communities โ†’ factions) โ€” this skips all intermediate scales.

CNNs solve this with hierarchical pooling (max-pool after each conv layer). DiffPool brings the same idea to graphs, with the key challenge: graph pooling must be permutation-invariant and must handle variable numbers of nodes.

The DiffPool Architecture

DiffPool processes each pooling level l with two GNNs:

1. Embedding GNN โ€” computes node embeddings:

Z^{(l)} = GNN_{embed}^{(l)}( A^{(l)}, H^{(l)} )

2. Pooling GNN โ€” computes cluster assignments:

S^{(l)} = softmax( GNN_{pool}^{(l)}( A^{(l)}, H^{(l)} ) )

S^{(l)} โˆˆ โ„^{N_l ร— k_l} is the soft assignment matrix: S[i,j] is the probability that node i belongs to cluster j. softmax is applied row-wise.

Coarsening: compute the next levelโ€™s adjacency and embeddings:

H^{(l+1)} = S^{(l)T} Z^{(l)} โˆˆ โ„^{k_l ร— d}
A^{(l+1)} = S^{(l)T} A^{(l)} S^{(l)} โˆˆ โ„^{k_l ร— k_l}

The new adjacency A^{(l+1)} is the cluster-to-cluster connectivity โ€” how much two clusters share edges through their constituent nodes.

Why Soft Assignment?

Hard assignment (each node assigned to exactly one cluster) would require argmax โ€” not differentiable. Soft assignment (each node fractionally assigned to all clusters) allows end-to-end gradient flow.

The assignment S^{(l)} is learned jointly with the rest of the network. The model discovers which nodes should be clustered together โ€” without any external supervision on the clustering.

The analogy to attention: DiffPool's soft assignment is conceptually similar to attention in Transformers โ€” it learns a soft selection over items. Here, each cluster "attends" to nodes, and the resulting cluster embedding is a weighted sum of node embeddings (S^T Z). The difference: DiffPool reduces the number of tokens (nodes coarsened to clusters), while attention preserves the token count.

Auxiliary Loss Terms

DiffPool adds two regularisation losses to guide clustering quality:

Link prediction loss: encourages nodes connected by an edge to be assigned to the same cluster:

L_{LP} = ||A^{(l)} - S^{(l)} S^{(l)T}||_F

If S^{(l)} S^{(l)T} approximates A^{(l)}, clusters are graph-connected.

Entropy loss: encourages each nodeโ€™s assignment to be concentrated (not uniformly spread):

L_E = (1/N) ฮฃ_i H( S_i )

Where H is entropy. Low entropy โ†’ sharp cluster assignment โ†’ more interpretable clusters.

Concrete Worked Example: One DiffPool Step

Suppose we have 4 nodes and want to pool to 2 clusters. After GNN_embed, node embeddings Z โˆˆ โ„^{4ร—3}:

Z = [[1, 0, 1],   # node 0
     [1, 0, 0],   # node 1 โ€” similar to node 0
     [0, 1, 1],   # node 2
     [0, 1, 0]]   # node 3 โ€” similar to node 2

GNN_pool outputs raw logits, softmax โ†’ assignment S โˆˆ โ„^{4ร—2}:

S = [[0.9, 0.1],   # node 0 โ†’ mostly cluster A
     [0.8, 0.2],   # node 1 โ†’ mostly cluster A
     [0.1, 0.9],   # node 2 โ†’ mostly cluster B
     [0.2, 0.8]]   # node 3 โ†’ mostly cluster B

New cluster embeddings: Hโ€™ = S^T Z โˆˆ โ„^{2ร—3}:

H'[A] = 0.9ร—[1,0,1] + 0.8ร—[1,0,0] + 0.1ร—[0,1,1] + 0.2ร—[0,1,0]
       โ‰ˆ [1.7, 0.3, 0.9]   (cluster A = nodes 0+1)
H'[B] = 0.1ร—[1,0,1] + 0.2ร—[1,0,0] + 0.9ร—[0,1,1] + 0.8ร—[0,1,0]
       โ‰ˆ [0.3, 1.7, 1.0]   (cluster B = nodes 2+3)

The soft assignment correctly merged structurally similar nodes (0,1) into cluster A and (2,3) into cluster B โ€” learned purely from features and graph structure.

Computational Cost

DiffPool runs two GNNs at each level. For a graph with N nodes coarsened to k clusters:

  • Forward pass: O(Nยฒ d) for the dense assignment matrix
  • Memory: O(Nยฒ) โ€” DiffPool operates on dense adjacency matrices

This makes DiffPool quadratic in N โ€” practical for graphs with hundreds of nodes (molecules), but not for social networks or knowledge graphs with millions of nodes.

When DiffPool Helps

DiffPool outperforms flat pooling when:

  1. The task requires multi-scale understanding: molecular property prediction benefits from atom-level and functional-group-level representations simultaneously
  2. Graphs have natural hierarchical structure: trees, clustered communities, hierarchical molecules
  3. The graph is small enough: typically โ‰ค 1000 nodes

It is a standard strong baseline on the TUDataset graph classification benchmarks (PROTEINS, MUTAG, IMDB, RDT).

Limitations

  1. Quadratic memory: no scaling to large graphs
  2. Fixed number of clusters: must specify k^{(l)} per level before training
  3. No guarantee of meaningful clusters: the auxiliary losses help but do not force semantically meaningful groupings
  4. Sensitive to the number of pooling levels: too many levels โ†’ over-compression; too few โ†’ flat pooling

Summary

ComponentRole
GNN_{embed}Compute node representations at this scale
GNN_{pool}Learn soft cluster assignments
S^T ZAggregate node embeddings into cluster embeddings
S^T A SCoarsen adjacency to cluster graph
L_{LP} + L_EEncourage connectivity-aligned, sharp clusters

DiffPool introduced the idea of learned hierarchical pooling for graphs โ€” differentiable, end-to-end, and structure-aware. Its quadratic complexity limits scale, but for small-graph tasks (molecules, proteins), it remains a reference architecture.

References

You May Also Enjoy

Modern Activation Functions: GELU, SiLU, Mish, and Smooth Gating

11 minute read

Published:

Once ReLU became the default, researchers started asking a better question: can we keep the easy optimization while making the activation smoother, softer, and more expressive? This chapter covers the modern answers.

Activation Functions in Neural Networks: Why Non-Linearity Matters

11 minute read

Published:

Activation functions are the reason neural networks can model curved decision boundaries instead of collapsing into one giant linear map. This chapter builds the intuition first, then walks through the classical functions that shaped deep learning.

FoPE: Fourier Position Embedding for Length Generalization

6 minute read

Published:

FoPE rethinks long-context positional encoding from a frequency-domain perspective. Instead of only stretching RoPE heuristically, it explicitly improves attentionโ€™s periodic extension so Transformers generalize more gracefully to longer sequences.