BN-Pool: a Bayesian Nonparametric Approach to Graph Pooling

TL;DR

BN-Pool tackles graph pooling by removing the need to specify a pooling ratio. It introduces a Bayesian non-parametric pooling operator based on a Dirichlet process with stick-breaking to infer an unbounded number of clusters per graph, approximated via variational inference. Training combines a graph-reconstruction objective with KL-based priors to regulate cluster usage and connectivity, yielding an adaptive coarsening mechanism. Empirically, BN-Pool achieves competitive or superior graph classification results and yields meaningful, graph-tailored clusterings, including notable gains on Colors-3 and Enzymes datasets.

Abstract

We introduce BN-Pool, the first clustering-based pooling method for Graph Neural Networks (GNNs) that adaptively determines the number of supernodes in a coarsened graph. By leveraging a Bayesian non-parametric framework, BN-Pool employs a generative model capable of partitioning graph nodes into an unbounded number of clusters. During training, we learn the node-to-cluster assignments 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. This adaptive strategy allows BN-Pool to automatically discover an optimal coarsening level, offering enhanced flexibility and removing the need to specify sensitive pooling ratios. We show that BN-Pool achieves superior performance across diverse benchmarks.

Figures (12)

• Figure 1: Graphical representation of in plate notation. The orange dashed arrow represents a deterministic computation.
• Figure 2: The SEL operation of and the components of the auxiliary loss.
• Figure 3: Clusters found on a graph with five communities.
• Figure 4: Adjacency matrix of Cora, class labels visualization, and adjacency matrix reconstruction by and .
• Figure 5: Original node features and assignments $\mathbf{S}$ of