Learning Hierarchical Graph Neural Networks for Image Clustering

Yifan Xing, Tong He, Tianjun Xiao, Yongxin Wang, Yuanjun Xiong, Wei Xia, David Wipf, Zheng Zhang, Stefano Soatto; Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021, pp. 3467-3477


We propose a hierarchical graph neural network (GNN) model that learns how to cluster a set of images into an unknown number of identities using a training set of images annotated with labels belonging to a disjoint set of identities. Our hierarchical GNN uses a novel approach to merge connected components predicted at each level of the hierarchy to form a new graph at the next level. Unlike fully unsupervised hierarchical clustering, the choice of grouping and complexity criteria stems naturally from supervision in the training set. The resulting method, Hi-LANDER, achieves an average of 49% improvement in F-score and 7% increase in Normalized Mutual Information (NMI) relative to current GNN-based clustering algorithms. Additionally, state-of-the-art GNN-based methods rely on separate models to predict linkage probabilities and node densities as intermediate steps of the clustering process. In contrast, our unified framework achieves a three-fold decrease in computational cost. Our training and inference code are released.

Related Material

[pdf] [supp] [arXiv]
@InProceedings{Xing_2021_ICCV, author = {Xing, Yifan and He, Tong and Xiao, Tianjun and Wang, Yongxin and Xiong, Yuanjun and Xia, Wei and Wipf, David and Zhang, Zheng and Soatto, Stefano}, title = {Learning Hierarchical Graph Neural Networks for Image Clustering}, booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)}, month = {October}, year = {2021}, pages = {3467-3477} }