---
title: Torch Geometric Pool (TGP) Overview
url: https://www.emergentmind.com/topics/torch-geometric-pool-tgp
type: topic
---

# Torch Geometric Pool (TGP) Overview

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 [2512.12642].

## 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 [2512.12642, 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):

```bash
pip install torch-geometric-pool
# Latest (development) version:
pip install git+https://github.com/filippoMB/torch-geometric-pool.git
```

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

```python
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:

```python
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:

```python
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 [2512.12642, 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 [2512.12642].

Source: https://www.emergentmind.com/topics/torch-geometric-pool-tgp