TopKPool and SAGPool: Sparse Graph Pooling
Published:
Intuition First: Selecting the Most Important Witnesses
Imagine summarising a long meeting by selecting the 5 most informative speakers and ignoring the rest. TopKPool does exactly this for graphs: it learns a score for each node (how informative is this node for the prediction?) and keeps only the top-k scoring nodes. The key question is how to score nodes — by their own features alone (TopKPool) or by how important they are in the context of their neighbourhood (SAGPool).
The Motivation for Sparse Pooling
DiffPool’s soft assignment is expressive but quadratic in graph size. For large graphs, this is prohibitive. A simpler approach: select a subset of nodes (the “important” ones) and form the induced subgraph.
This hard selection is naturally sparse (the selected nodes inherit only edges between them) and avoids dense matrix computation. The challenge: how to define “importance” and how to make selection differentiable.
TopKPool (gPool)
Score computation: learn a single projection vector \(p \in \mathbb{R}^{d}\), shared by every node. Each node’s importance score is the length of its projection onto \(p\):
Note what is not in this expression: the adjacency \(A\). A node’s score depends only on its own row of \(H\). Message passing in earlier layers has of course already mixed neighbourhood information into \(H\), but the scoring function itself is blind to structure.
Selection: rank the scores and keep the \(k\) largest, \(\mathrm{idx} = \operatorname{top-}k(y)\). Then
where \(\sigma\) is the logistic sigmoid and the gate \(\sigma(y_{\mathrm{idx}}) \in \mathbb{R}^{k}\) is broadcast across the \(d\) feature columns. That elementwise multiplication is what makes \(p\) trainable: \(y\) appears in the output, not only in the ranking.
Subgraph: \(A' = A_{\mathrm{idx},\mathrm{idx}}\) is the adjacency restricted to selected nodes. It stays sparse if the original graph was sparse.
Complexity: \(O(Nd)\) for the scores, \(O(N \log N)\) (or \(O(N)\) with a selection algorithm) for the ranking, and \(O(E)\) to extract the induced subgraph from a sparse adjacency. No dense \(N \times N\) object is ever formed.
SAGPool: Self-Attention Graph Pooling
SAGPool (Lee et al., 2019) changes exactly one thing about TopKPool: how the score is produced. Instead of a projection onto a learned vector, the score comes from a graph convolution with a single output channel:
where \(\hat{A} = \tilde{D}^{-1/2}\tilde{A}\tilde{D}^{-1/2}\) is the usual GCN propagation matrix. The key difference from TopKPool is the presence of \(\hat{A}\): node \(i\)’s score is a function of its neighbours’ features as well as its own, computed at pooling time rather than inherited from earlier layers.
Intuition: a node should be selected as important if both it and its neighbourhood are informative for the task. A node that sits on an important information flow scores highly even when its own features are unremarkable.
Selection and subgraph formation then follow TopKPool exactly:
Note that the gate here is \(y\) itself, already squashed by \(\tanh\) inside the score — SAGPool uses \(\tanh\) where gPool uses a sigmoid. The consequence is that a SAGPool gate can be negative, flipping the sign of a retained node’s features, whereas a gPool gate only ever attenuates.
Differentiability via Score Gating
Hard top-\(k\) selection is not differentiable: the ranking is a piecewise-constant function of the scores, so \(\partial\,\mathrm{idx}/\partial y = 0\) almost everywhere. If the scores entered the model only through the ranking, \(p\) (or \(\Theta_{\mathrm{att}}\)) would receive zero gradient and never train.
The fix in both methods is to multiply the retained features by their own gated score:
Now \(y\) appears in the forward output as a smooth multiplicative factor, so \(\partial H'_i/\partial y_{\mathrm{idx}(i)} = H_{\mathrm{idx}(i),:}\, g'(y_{\mathrm{idx}(i)})\) is nonzero and the scorer learns. What still does not receive gradient is the selection: a node that was dropped contributes nothing to the loss and therefore gets no signal about whether it should have been kept. Training can only refine the ranking of nodes it already keeps, which is why these methods are sensitive to initialisation and why an unlucky early ranking can persist.
This is the same device attention uses to avoid one-hot selection — soften the discrete choice into a multiplication so gradients have somewhere to flow.
Hierarchical Pooling with TopK/SAGPool
Both methods are designed for stacking:
Layer 1: N nodes → GNN → TopKPool → k₁ nodes
Layer 2: k₁ nodes → GNN → TopKPool → k₂ nodes
Layer 3: k₂ nodes → GNN → Global pool → graph embedding
At each level, the graph shrinks. The final global pooling (mean/sum/max) operates on a small set of “important” nodes — the hierarchically selected representatives.
Comparison with DiffPool
| Property | DiffPool | TopKPool | SAGPool |
|---|---|---|---|
| Assignment | Soft (continuous) | Hard (top-\(k\)) | Hard (top-\(k\)) |
| Scoring input | \(\mathrm{GNN}(A, X)\) | Projection \(Hp/\lVert p\rVert\) | \(\mathrm{GNN}(A, H)\) |
| Memory | \(O(N^2)\) | \(O(N + E)\) | \(O(N + E)\) |
| Scales to large graphs | No | Yes | Yes |
| Neighbourhood-aware scores | Yes | No | Yes |
| Nodes discarded | None — every node contributes to every cluster | Unselected nodes dropped | Unselected nodes dropped |
| Coarsened graph | Dense | Sparse (induced subgraph) | Sparse (induced subgraph) |
| Differentiability | Full | Score gating only; selection is not | Score gating only; selection is not |
Practical Notes
Ratio \(k/N\): typically set to 0.5 or 0.25 per level — halving or quartering the graph at each pooling step. Because the ratio is relative, the absolute \(k\) adapts to each graph’s size, which is one advantage over DiffPool’s fixed cluster count. Too aggressive → information loss. Too gentle → insufficient compression.
Edge dropping: nodes dropped at level \(l\) take their edges with them. If two retained nodes were connected only through dropped nodes, they become disconnected in \(A'\) — the induced subgraph does not reconnect them. Stacked over several levels this can fragment the graph into isolated nodes, at which point further message passing does nothing and only the final global readout still carries signal.
Batch handling: when training on graphs of different sizes, pooling ratios produce different absolute node counts. PyTorch Geometric handles this with batch indexing.
Summary
TopKPool and SAGPool trade DiffPool’s expressiveness for scalability: by selecting a sparse subset of nodes rather than soft-assigning all nodes to all clusters, they pool in time and memory linear in the graph, at the cost of discarding unselected nodes entirely. The two differ only in how a node is scored — a projection of its own features versus a one-layer GNN over its neighbourhood — and in both cases it is the multiplication of retained features by their gated score, not the ranking, that makes the scorer trainable at all.
References
- Gao, H., & Ji, S. (2019). Graph U-Nets. ICML 2019 (TopKPool / gPool).
- Lee, J., Lee, I., & Kang, J. (2019). Self-Attention Graph Pooling. ICML 2019 (SAGPool).
- Ying, R., You, J., Morris, C., Ren, X., Hamilton, W. L., & Leskovec, J. (2018). Hierarchical Graph Representation Learning with Differentiable Pooling. NeurIPS 2018 (DiffPool — the alternative approach).