Torch Geometric Pool (TGP) Overview

• Torch Geometric Pool (TGP) is a modular PyTorch-based library for hierarchical pooling in Graph Neural Networks, featuring both trainable and non-trainable operators.
• Its unified API, built on the SRCL paradigm with selector, reducer, connector, and lifter components, enables rapid prototyping and component-level experimentation.
• Precomputed pooling and caching mechanisms achieve 10–100× speedups, making TGP ideal for systematic benchmarking and efficient GNN pipeline design.

Torch Geometric Pool (tgp) is a PyTorch-based library designed for hierarchical pooling in Graph Neural Networks (GNNs), extending the PyTorch Geometric (PyG) ecosystem. It provides a broad array of pooling operators—both sparse and dense, trainable and non-trainable—integrated under a unified and modular API. tgp emphasizes usability, extensibility, and efficiency, supporting features such as precomputed pooling to greatly accelerate training with certain methods. The toolkit enables systematic benchmarking and fast prototyping for various GNN pooling approaches and downstream tasks (Bianchi et al., 14 Dec 2025).

1. Architectural Design and Unified API

tgp is built as a lightweight extension on PyG, centering on the abstract base class torch_geometric_pool.pooling.SRCPooling. This class implements the Select-Reduce-Connect-Lift (SRCL) paradigm, decomposing every pooling operator into four modular and interchangeable submodules:

• selector: maps original nodes $\{1,\dots,N\}$ to $K$ supernodes, represented as an assignment matrix $S\in\mathbb{R}^{N\times K}$.
• reducer: aggregates and transforms node features from selected nodes to supernodes.
• connector: constructs the coarsened adjacency matrix (edge index) among supernodes.
• lifter: optionally propagates coarsened features back to the original node set.

Input and output for all poolers are standardized through two data structures:

• SelectOutput: encapsulates the assignment matrix $S$, selected node indices, and utility methods.
• PoolingOutput: returns pooled node features $X'$, pooled edge_index, updated batch vector, the original SelectOutput, and auxiliary loss terms (when present).

API conventions:

• Poolers are instantiated either via get_pooler(name, **kwargs) or by direct class construction, with typical arguments including in_channels, ratio (or K), and seed.
• Core attributes include is_dense (indicator for dense outputs), has_loss (whether the pooler provides auxiliary losses), and is_precoarsenable (whether caching/pre-coarsening is supported).
• The preprocessing(data, use_cache=False) call formats the input (e.g., padding for dense poolers).
• The main pooling operation is executed by forward.
• Global graph-level summaries are extracted via global_pool(x, batch).

This structural unification enables interchangeability and modular experimentation at the submodule or operator level.

2. Catalog of Implemented Pooling Operators

tgp natively supports 17 pooling operators, partitioned by sparsity, trainability, and clustering/assignment paradigm. The major operators and their defining operations are:

Operator	Category	Core Assignment/Formula
TopK [Gao & Ji]	Sparse, Trainable	$s = X W_s$; select top-$k$; $S_{i,j}=1$ if selected; $X' = S^T X$; $A' = S^T A S$
SAGPool [Lee et al.]	Sparse, Trainable	$a = \sigma(\mathrm{GNN}(X,A)v)$; select top-$k$ nodes
ASAP [Ranjan et al.]	Sparse, Trainable	GNN + motif-based local structure, thresholding
ECPool [Diehl, Landolfi]	Sparse, Trainable	Edge contraction via learned scores
PAN [Ma et al.]	Sparse, Trainable	$\Phi = \sum_{t=0}^T \alpha_t A^t X W$; cluster by assignment
GraClus [Dhillon et al.]	Sparse, Non-trainable	Multilevel weighted matching of adjacency
LaPool [Noutahi et al.]	Sparse, Non-trainable	Clustering on top-$k$ Laplacian eigenvectors
KMIS [Bacciu et al.]	Sparse, Non-trainable	Repeated maximum independent set
NDP [Bianchi et al.]	Sparse, Non-trainable	Node decimation by degree criterion
DiffPool [Ying et al.]	Dense, Trainable	$S = \mathrm{softmax}(X W_S)$; $X' = S^T X$, $A' = S^T A S$; link/entropy aux. losses
MinCutPool [Bianchi et al.]	Dense, Trainable	$S=\mathrm{softmax}(X W)$; min-cut, entropy aux. losses
DMoN [Tsitsulin et al.]	Dense, Trainable	Modularity-based clustering
ACC [Hansen & Bianchi]	Dense, Trainable	Total-variation clustering
HOSC [Duval & Malliaros]	Dense, Trainable	Higher-order spectral motif Laplacian
JBPool [Bianchi]	Dense, Trainable	Joint-balance loss
NMF [Bacciu & di Sotto]	Dense, Non-trainable	$A \approx S H S^T$ via offline NMF

Each operator is further annotated with respect to sparsity (output format), trainability, and whether it supports auxiliary loss terms. The spectrum of implemented methods covers local heuristics, learned attention, spectral and modularity approaches, edge contraction, and clustering-based strategies. The architecture encourages reuse and extension of individual components for rapid ablation or new pooler creation.

3. Precomputed Pooling for Computational Efficiency

Non-trainable poolers (e.g., GraClus, NDP, KMIS, NMF) rely exclusively on graph topology and can be significantly accelerated by precomputing expensive sub-steps. tgp supports two mechanisms:

• In-memory caching: For single-graph tasks, setting cached=True stores results of Select and Connect on the first pass, reducing all subsequent forwards to just the Reduce phase.
• Pre-coarsening transforms: For datasets of multiple graphs, the PoolTransform(pooler, **kwargs) is applied as a PyG pre_transform, executing Select and Connect offline and embedding the assignment and coarsened adjacency into each sample. PoolDataLoader collates these into pooled batches, so only the linear Reduce remains at runtime.

The reduction in forward computational complexity is as follows:

Phase	Standard	Cached / Pre-coarsened
Select	$O(\text{score‐GNN or matching})$	$O(1)$
Connect	$O(M)$	$O(1)$
Reduce	$O(NF\times K)$	$O(NF\times K)$

Empirically, this yields 10–100$\times$ speedup and reduces pre-coarsened forward computation to a single matrix multiply [(Bianchi et al., 14 Dec 2025), Table 4]. This approach is appropriate exclusively for operators whose selection is independent of node features, i.e., non-trainable poolers.

4. Usage and Customization: Example Workflows

tgp is distributed via PyPI and supports PyTorch (≥1.11) and PyG (≥2.0):

pip install torch-geometric-pool pip install git+https://github.com/filippoMB/torch-geometric-pool.git

A typical use-case with pre-coarsening enabled is as follows:

from torch_geometric.loader import DataLoader from tgp import get_pooler from torch_geometric.datasets import TUDataset dataset = TUDataset( root='data/NCI1', name='NCI1', pre_transform=lambda d: get_pooler('graclus', ratio=0.5).pre_transform(d)) loader = DataLoader( dataset, batch_size=32, shuffle=True, collate_fn=lambda batch: PoolDataLoader.collate(batch))

Model instantiation and training proceeds by inserting any pooling operator (selected by name or class) into a PyG-compatible GNN pipeline, with downstream layers conditional on the pooler's density:

self.pool = get_pooler('diffpool', in_channels=32, K=10) self.conv1 = GINConv(...) self.conv2 = DenseGINConv(...) if self.pool.is_dense else GINConv(...)

Customizing individual SRCL components is also supported:

from tgp.connect import KronConnect pool = get_pooler('topk', in_channels=64, ratio=0.5) pool.connector = KronConnect()

This modular structure and strict input-output typing enables rapid operator ablation and experimentation.

5. Benchmarking Results

An extensive benchmarking study in tgp evaluates 17 operators across four major tasks:

• Unsupervised Node Clustering (Cora, Citeseer, PubMed, DBLP, Community): Normalized Mutual Information (NMI) metric. ACC, DMoN, and MinCut achieve the best results on homophilic citation networks; DiffPool underperforms without supervised guidance.
• Node Classification in Heterophilic Graphs (Amazon-ratings, Minesweeper, Roman-empire, Questions, Tolokers): Accuracy/AUROC metrics. MaxCutPool, ASAP, and NDP are consistently superior.
• Graph-level Classification/Regression (NCI1, Reddit-b, EXPWL1, Multipartite, GCB-H, ogbg-molhiv, peptides-func, peptides-struct): A variety of metrics (Acc, ROC-AUC, AP, MAE) are used. No single pooler consistently dominates; e.g., ACC, MinCut, and HOSC excel on peptides, while MaxCutPool outperforms on Multipartite, and JBPool, KMIS, ACC lead on GCB-H.
• Pooling Efficiency (NDP, KMIS, NMF): Pre-coarsening and caching provide 10–1,000$\times$ speedups, confirming the computational advantage for topologically-invariant poolers.

The choice of optimal pooling operator is dataset- and task-dependent, supporting the need for broad coverage and consistent comparison facilities [(Bianchi et al., 14 Dec 2025), Tables 1–4].

6. Practical Recommendations

• For large graphs ($N\gg 1,000$) where coarse granularity suffices, sparse, non-trainable poolers (GraClus, KMIS, NDP) with pre-coarsening yield maximal throughput.
• For community detection and unsupervised node clustering in homophilic settings, dense, spectral poolers (ACC, DMoN, MinCut) are most effective.
• For supervised tasks, trainable poolers (ASAP, TopK, SAGPool, DiffPool variants) offer greater adaptivity; auxiliary loss term weighting should be controlled, though defaults are typically robust.
• The pooling ratio or cluster count ($K$) dictates trade-off between information retention and compression, with $r\in[0.1,0.5]$ as a standard initial range.
• The SRCL modularity enables efficient component-level ablation (e.g., swapping connectors) for analytical studies.
• Caching or dataset-level pre-coarsening should always be leveraged for fixed-topology poolers to achieve substantial speed improvements.

tgp's consistent, modular design, extensive pooling repertoire, and computational optimizations provide a practical and rigorous toolchain for hierarchical GNN research and prototyping (Bianchi et al., 14 Dec 2025).