Total Variation Graph Neural Networks

TL;DR

This paper introduces a novel GNN model that optimizes a tighter total variation-based relaxation for vertex clustering, resulting in sharper cluster boundaries and improved performance over existing methods.

Contribution

The paper proposes a new GNN approach that directly minimizes graph total variation for better vertex clustering and graph pooling, surpassing spectral clustering relaxations.

Findings

Outperforms existing GNNs in vertex clustering tasks

Produces sharper and more accurate cluster boundaries

Enhances graph classification performance

Abstract

Recently proposed Graph Neural Networks (GNNs) for vertex clustering are trained with an unsupervised minimum cut objective, approximated by a Spectral Clustering (SC) relaxation. However, the SC relaxation is loose and, while it offers a closed-form solution, it also yields overly smooth cluster assignments that poorly separate the vertices. In this paper, we propose a GNN model that computes cluster assignments by optimizing a tighter relaxation of the minimum cut based on graph total variation (GTV). The cluster assignments can be used directly to perform vertex clustering or to implement graph pooling in a graph classification framework. Our model consists of two core components: i) a message-passing layer that minimizes the $\ell_1$ distance in the features of adjacent vertices, which is key to achieving sharp transitions between clusters; ii) an unsupervised loss function that minimizes the GTV of the cluster assignments while ensuring balanced partitions. Experimental results show that our model outperforms other GNNs for vertex clustering and graph classification.