[Paper Review] Learning Hyperbolic Representations for Unsupervised 3D Segmentation

Joy Hsu, Jeffrey Gu+3 more|arXiv (Cornell University)|May 4, 2021

3D Shape Modeling and Analysis25 references4 citations

TL;DR

This paper proposes a hyperbolic variational autoencoder with gyroplane convolutional layers and hierarchical triplet loss to learn hierarchical 3D representations for unsupervised segmentation. By modeling the intrinsic hierarchy in volumetric data using hyperbolic geometry, the method achieves improved unsupervised segmentation on BraTS, cryo-EM, and synthetic hierarchical datasets.

ABSTRACT

There exists a need for unsupervised 3D segmentation on complex volumetric data, particularly when annotation ability is limited or discovery of new categories is desired. Using the observation that much of 3D volumetric data is innately hierarchical, we propose learning effective representations of 3D patches for unsupervised segmentation through a variational autoencoder (VAE) with a hyperbolic latent space and a proposed gyroplane convolutional layer, which better models the underlying hierarchical structure within a 3D image. We also introduce a hierarchical triplet loss and multi-scale patch sampling scheme to embed relationships across varying levels of granularity. We demonstrate the effectiveness of our hyperbolic representations for unsupervised 3D segmentation on a hierarchical toy dataset, BraTS whole tumor dataset, and cryogenic electron microscopy data.

Motivation & Objective

• To address the challenge of unsupervised 3D segmentation in complex volumetric data where annotations are scarce or new categories need discovery.
• To leverage the inherent hierarchical structure present in 3D volumetric data for improved representation learning.
• To develop a deep learning framework that models hierarchical relationships more effectively than Euclidean space in 3D segmentation tasks.
• To enable discovery of semantic structures without category labels by learning disentangled, hierarchical representations in hyperbolic space.

Proposed method

• Utilizes a variational autoencoder (VAE) with a hyperbolic latent space to better capture the hierarchical nature of 3D volumetric data.
• Introduces a novel gyroplane convolutional layer that operates in hyperbolic space, enabling equivariant feature learning on hyperbolic manifolds.
• Employs a hierarchical triplet loss to enforce relative ordering of patches across multiple levels of granularity, promoting structured representation learning.
• Applies a multi-scale patch sampling scheme to capture relationships at varying levels of detail during training.
• Uses hyperbolic geometry to model tree-like or hierarchical structures more naturally than Euclidean space, improving representation quality.

Experimental results

Research questions

• RQ1 Can hyperbolic representations better capture the hierarchical structure of 3D volumetric data compared to Euclidean representations for unsupervised segmentation?
• RQ2 How does the proposed gyroplane convolutional layer improve feature learning in hyperbolic space for 3D data?
• RQ3 To what extent does the hierarchical triplet loss enhance the quality of learned representations across multiple scales?
• RQ4 Does the multi-scale patch sampling scheme improve generalization and hierarchical consistency in representation learning?
• RQ5 Can the proposed method achieve competitive unsupervised segmentation performance on real-world 3D datasets like BraTS and cryo-EM without annotations?

Key findings

• The proposed method achieves state-of-the-art unsupervised 3D segmentation performance on the hierarchical toy dataset, demonstrating superior ability to recover hierarchical structures.
• On the BraTS whole tumor dataset, the model learns meaningful, semantically coherent segmentations without any category labels, indicating effective discovery of anatomical structures.
• The method shows improved representation learning on cryo-EM data, capturing structural hierarchies inherent in macromolecular complexes.
• The use of hyperbolic space leads to more compact and structured latent representations compared to Euclidean baselines.
• The hierarchical triplet loss and multi-scale sampling significantly improve the consistency and granularity of learned representations across levels of abstraction.