Seg-HGNN: Unsupervised and Light-Weight Image Segmentation with Hyperbolic Graph Neural Networks

TL;DR

Seg-HGNN develops a lightweight, unsupervised image segmentation framework that operates in hyperbolic space to capture latent hierarchical structure with minimal parameters. By projecting patch-level features from a frozen transformer into the Lorentz model and applying a single hyperbolic graph convolutional layer, the method achieves competitive localization and segmentation results on standard benchmarks while using less than 7.5k trainable parameters and enabling fast inference on common GPUs. The approach combines a Lorentz linear transform, Einstein midpoint aggregation, and a relaxed normalized-cut loss to cluster patches without supervision. This work demonstrates that hyperbolic representations offer compact, scalable, and effective image analysis, with practical implications for edge-friendly vision systems and real-time applications.

Abstract

Image analysis in the euclidean space through linear hyperspaces is well studied. However, in the quest for more effective image representations, we turn to hyperbolic manifolds. They provide a compelling alternative to capture complex hierarchical relationships in images with remarkably small dimensionality. To demonstrate hyperbolic embeddings' competence, we introduce a light-weight hyperbolic graph neural network for image segmentation, encompassing patch-level features in a very small embedding size. Our solution, Seg-HGNN, surpasses the current best unsupervised method by 2.5\%, 4\% on VOC-07, VOC-12 for localization, and by 0.8\%, 1.3\% on CUB-200, ECSSD for segmentation, respectively. With less than 7.5k trainable parameters, Seg-HGNN delivers effective and fast ($\approx 2$ images/second) results on very standard GPUs like the GTX1650. This empirical evaluation presents compelling evidence of the efficacy and potential of hyperbolic representations for vision tasks.

Figures (4)

• Figure 1: Method overview. We extract patch-level features from a frozen model $\boldsymbol{T}$, flattening them to $\boldsymbol{f}$. Edge weights are obtained from the Gram matrix(eqn.(\ref{['eq:normalization']})) of $\boldsymbol{f}$. We project $\boldsymbol{f}$ to the lorentz manifold to obtain initial node embeddings. We optimize an unsupervised graph partitioning loss to obtain $k$ clusters.
• Figure 2: A few samples comparing the quality of segmentation achieved by Seg-HGNN. Here, the second row shows the predicted masks and the third row has the ground-truth.
• Figure 3: Samples showing the quality of semantic segmentation. The first image shows that even for multiple objects, the underlying semantic behind clustering remains intact. For example, the heads of both the cats are clustered together in red. The projected hyperbolic embeddings also show the quality of clustering.
• Figure 4: Low-dimensional performance of hyperbolic embeddings.