BN-Pool: Bayesian Nonparametric Pooling for Graphs

TL;DR

BN-Pool introduces a Bayesian nonparametric clustering pooling method for Graph Neural Networks that adaptively determines the optimal number of supernodes, balancing graph reconstruction and compactness.

Contribution

It is the first clustering-based pooling method for GNNs that automatically learns the number of clusters using a Bayesian nonparametric model.

Findings

BN-Pool effectively discovers optimal coarsening levels for graphs.

It maintains high performance while reducing redundancy in pooled graphs.

Code is publicly available at the provided GitHub link.

Abstract

We introduce BN-Pool, the first clustering-based pooling method for Graph Neural Networks that adaptively determines the number of supernodes in a coarsened graph. BN-Pool leverages a generative model based on a Bayesian nonparametric framework for partitioning graph nodes into an unbounded number of clusters. During training, the node-to-cluster assignments are learned by combining the supervised loss of the downstream task with an unsupervised auxiliary term, which encourages the reconstruction of the original graph topology while penalizing unnecessary proliferation of clusters. By automatically discovering the optimal coarsening level for each graph, BN-Pool preserves the performance of soft-clustering pooling methods while avoiding their typical redundancy by learning compact pooled graphs. The code is available at https://github.com/NGMLGroup/Bayesian-Nonparametric-Graph-Pooling.