Torch Geometric Pool: the Pytorch library for pooling in Graph Neural Networks

2025-12-14

AI Summary

Hierarchical pooling methods for Graph Neural Networks (GNNs) suffer from fragmentation, poor integrability, and the absence of a unified evaluation benchmark. Method: This paper introduces the first PyTorch-native open-source library dedicated to GNN pooling, built upon PyTorch Geometric. It innovatively incorporates pre-computed pooling to accelerate inference and features a standardized API alongside a modular operator architecture, enabling flexible extension and plug-and-play deployment. Contribution/Results: Through systematic benchmarking across diverse downstream tasks—including graph classification and node classification—the study demonstrates that pooling performance is highly task- and data-distribution-dependent, establishing a “task-aware selection” paradigm. The library provides a reproducible, comparable, and efficient unified infrastructure for GNN pooling research, significantly lowering barriers to method development, integration, and empirical evaluation.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Data Mining & Knowledge Management: Graph Mining, Social Network Analysis & Community
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web evaluation methodologies and metrics

Abstract

We introduce Torch Geometric Pool (tgp), a library for hierarchical pooling in Graph Neural Networks. Built upon Pytorch Geometric, Torch Geometric Pool (tgp) provides a wide variety of pooling operators, unified under a consistent API and a modular design. The library emphasizes usability and extensibility, and includes features like precomputed pooling, which significantly accelerate training for a class of operators. In this paper, we present tgp's structure and present an extensive benchmark. The latter showcases the library's features and systematically compares the performance of the implemented graph-pooling methods in different downstream tasks. The results, showing that the choice of the optimal pooling operator depends on tasks and data at hand, support the need for a library that enables fast prototyping.

Problem

• Introduces a library for hierarchical pooling in Graph Neural Networks
• Provides various pooling operators with a consistent API and modular design
• Enables fast prototyping by comparing pooling methods across different tasks

Innovation

• Hierarchical pooling library for Graph Neural Networks
• Unified API with modular design for pooling operators
• Precomputed pooling accelerates training for specific operators