BN-Pool: a Bayesian Nonparametric Pooling for Graphs

TMLR Paper6703 Authors

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 non-parametric 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://anonymous.4open.science/r/BN-Pool.

Submission Type: Long submission (more than 12 pages of main content)

Assigned Action Editor: ~Roman_Garnett1

Submission Number: 6703