Software implementation and code to reproduce the results of the Just Balance GNN (JBGNN) model for graph clustering as presented in the paper Simplifying Clustering with Graph Neural Networks.

The JBGNN architecture consists of:

• a GCN layer operating on the connectivity matrix: $\mathbf{I} - \delta ( \mathbf{I} - \mathbf{D}^{-\frac{1}{2}} \mathbf{A} \mathbf{D}^{-\frac{1}{2}} )$;
• a pooling layer that computes a cluster assignment matrix $\mathbf{S} \in \mathbb{R}^{N \times K}$ as

$$ \mathbf{S} = \texttt{softmax} \left( \texttt{MLP} \left( \mathbf{\bar X}, \boldsymbol{\Theta}_\text{MLP} \right) \right) $$

       where $\mathbf{\bar X}$ are the node features returned by a stack of GCN layers.

Each pooling layer is associated with an unsupervised loss that balances the size of the clusters and prevents degenerate partitions

$$\mathcal{L} = - \text{Tr}\left( \sqrt{\mathbf{S}^T\mathbf{S} } \right).$$

Node clustering and graph classification

Tensorflow

A TF/Keras implementation of the JustBalancePool layer is on Spektral.

Run example_clustering_tf.py to perform node clustering.

Pytorch

just_balance.py provides a Pytorch implementation based on Pytorch Geometric.

Run example_clustering.py to perform node clustering in Pytorch.

Run example_classification.py to perform graph classification in Pytorch.

Important

The results on the paper were obtained using the Tensorflow/Spektral implementation.

Citation

@misc{bianchi2022simplifying,
  doi = {10.48550/ARXIV.2207.08779},
  author = {Bianchi, Filippo Maria},
  title = {Simplifying Clustering with Graph Neural Networks},
  publisher = {arXiv},
  year = {2022},
}