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).$$

Clustering of graph vertices

Tensorflow

The official implementation of the JustBalancePool layer is on Spektral, the Tensorflow/Keras library for Graph Neural Networks.

Run main.py to perform node clustering.

Pytorch

just_balance_pyg.py provides a Pytorch implementation based on Pytorch Geometric. Run main_pyg.py to perform node clustering in Pytorch.

Note: the results on the paper are based on the Tensorflow implementation.

Poster

The poster presentation of the paper Simplifying Clustering with Graph Neural Networks can be downloaded here.

Citation

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