Learning Hyperbolic Representations for Unsupervised 3D Segmentation
Authors: Joy Hsu
Jeffrey Gu
Gong Her Wu
Wah Chiu
Serena Yeung
Venue:  (2021)
Keywords: unsupervised learning
representation learning
segmentation
biocomputing

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.


Under review as a conference paper at ICLR 2021
LEARNING HYPERBOLIC REPRESENTATIONS FOR UN- SUPERVISED 3D SEGMENTATION
Anonymous authors Paper under double-bling review
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 hi- erarchical, we propose learning effective representations of 3D patches for unsu- pervised 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 hier- archical triplet loss and multi-scale patch sampling scheme to embed relationships across varying levels of granularity. We demonstrate the effectiveness of our hy- perbolic representations for unsupervised 3D segmentation on a hierarchical toy dataset, BraTS whole tumor dataset, and cryogenic electron microscopy data.
INTRODUCTION
Recent advances in technology have greatly increased both the availability of 3D data, as well as the need process and learn from 3D data. In particular, technologies such magnetic resonance imaging and cryogenic electron microscopy (cryo-EM) have led to greater availability of 3D voxel data. Deep learning a promising technique to do so, but producing annotations for 3D data can be extremely expensive, especially for richer tasks such as segmentation in dense voxel grids. In some cases, labels may also be impossible to produce due to the limitations of current knowledge, or may introduce bias if we want to conduct scientific discovery. Unsupervised learning, which does not require annotations, is promising approach for overcoming these limitations.
In this work, we tackle the challenging problem of unsupervised segmentation on complex 3D voxel data by addressing the essential challenge of representation learning. We expand from prior literature in the hyperbolic domain that conducts classification in simple data to the task of segmentation in 3D images, which requires significantly more representation discriminability. In order to learn effective representations, need to capture the structure of our input data. We observe that 3D images often have inherent hicrarchical structure: as a biomedical example, a cryo-EM tomogram of a cell has hierarchy that at the highest level comprises the entire cell; such as the mitochondria and nucleus; and at an even finer level comprises sub-structures such as the nucleolus of nucleus or proteins within organelles. For downstream analysis, we are typically interested in the unsupervised discovery and segmentation of structures spanning multiple levels of hierarchy. However, prior work on representation learning for unsupervised 3D segmentation does not explicitly model hierarchical structure between different regions of a 3D image. We argue that this hampers the ability to leverage hierarchical relationships to improve segmentation in complex 3D images.
finer level comprises organelles
Our key insight is that we can utilize hyperbolic embedding space to learn effective hierarchical representations of voxel regions in 3D images. Hyperbolic representations have been proposed as a continuous way to represent hierarchical data, as trees can be embedded in hyperbolic space with arbitrarily low error (Sarkar 2011). These methods have shown promise for modeling data types such as natural language word taxonomies (Niceel & Kiela 2017 2018), graphs (Nickel & Kiela 2017 Mathieu et al. 2019 Ovinnikov 2019; Chami et al. (201). well simple MNIST (LeCun ct al. 2010) image data for classification (Mathieu et al. 2019). To the best of our knowledge, our work is the first to introduce learning hyperbolic representations to capture hierarchical structure

---

Under review as a conference paper at ICLR 2021
among subregions of complex 3D images, and to utilize the learned hyperbolic representations to perform a complex computer vision task such as segmentation.
Our approach for learning hyperbolic representations of 3D voxel grid data is based several key innovations. First, to handle larger and more complex 3D data such as biomedical images, propose hyperbolic 3D convolutional VAE along with a new gyroplane convolutional layer that respects hyperbolic geometry. Second, we enhance our VAE training objective with a novel self-supervised hierarchical triplet loss that helps our model learn hierarchical structure within the VAE's hyperbolic latent space. Finally, since our goal in segmentation is to learn hierarchy within voxel regions of 3D input, we present multi-scale sampling scheme such that our 3D VAE can simultaneously embed hierarchical relationships across varying levels of granularity.
In summary, our key contributions are follows:
We introduce hyperbolic 3D convolutional VAE with novel gyroplane convolutional layer that scales the learning of hyperbolic representations to complex 3D data.
We propose a multi-scale sampling scheme and hierarchical triplet loss in order to encode hierarchical structure in the latent space and perform 3D unsupervised segmentation.
We demonstrate the effectiveness of our approach through experiments on a synthetic 3D toy dataset, the Brain Tumor Segmentation (BraTS) dataset (Menze al. 2014 Bakas et al 2017 2018), and cryo-EM data.
2 RELATED WORK
Segmentation on 3D voxel data Since 3D voxel grids are dense, computer vision tasks such as supervised segmentation are commonly performed using deep learning architectures with 3D con- volutional layers (Chen et al.) 2016; Dou et al. 2017, Hesamian et al., 2019; [Zheng et al.) 2019). 3D, there has been sig- nificant effort in finding semi-supervised approaches, including using labels only from several fully annotated 2D slices of an input volume (Çięek et al. (2016). using a smaller set of segmentations with joint segmentation and registration (Xu & Niethammer 2019), and using one segmented input in conjunction with other unlabelled data (Zhao et al. 2019).
However, due to the challenges of obtaining voxel-level segmentations
Unsupervised approaches for 3D segmentation are useful not only for further reducing the manual annotation effort required, but also for scientific discovery tasks where we lack the sufficient knowl- edge provide representative training examples for structures of interest. |Moryya et al. (2018)
tends to 3D data an iterative approach of feature learning followed by clustering (Yang et al.) | (016)). Nalepa al. (2020) uses 3D convolutional autoencoder architecture and performs clustering the latent representations. Another approach, (Dalca et al.) 20188), uses network pre-trained on man- segmentations from separate dataset perform unsupervised segmentation of 3D biomedical images. However, this limits applicability to areas where we already have dataset with manual annotations and makes unsuitable for unbiased unsupervised discovery. Gur et al. (2019) and Kitrungrotsakul et al. (2019) developed unsupervised methods for 3D segmentation of vessel struc- tures, but these are specialized and do not generalize to the segmentation of other structures. Beyond unsupervised 3D segmentation, there has been work such as Ji et al. (2019) that performs unsuper- vised 2D segmentation based on a mutual information objective, and Caron et al.] (2018), which proposes using the clustered output of encoder pseudo-labels. While these methods can be applied to 2D slices of a 3D volume to perform 3D segmentation, they generally suffer limitations due to insufficient modeling of the 3D spatial information. None of the aforementioned approaches explicitly model hierarchical structure, which is the main focus of our work.
Hyperbolic representations recent line of work has employed hyperbolic space model hi- erarchical structure, with the intuition that tree structures can be naturally embedded into continu- ous hyperbolic space (Nickel Kiela 2017). Several works have proposed hyperbolic variational autoencoders (VAEs)
unsupervised method to learn hyperbolic representations. Ovinnikov (2019) proposes a Wasserstein autoencoder on the Poincaré ball model of hyperbolic geometry. Nagano al. (2019) proposes VAE on the hyperboloid model of hyperbolic geometry where the last layer of the encoder is an exponential map, and derives a reparametrisable sampling scheme for the wrapped normal distribution, which they for the prior and posterior. Mathieu al. (2019)
2

---

Under review as a conference paper at ICLR 2021
proposes VAE on the Poincaré ball model of hyperbolic geometry. In addition to having the last layer of the encoder be an exponential map, Mathieu et al. (2019) also proposes to have the first layer of the decoder be the gyroplane layer proposed by Ganea et al. (2018) in order to better handle the geometry of the hyperbolic latent space, and applies their model to MNIST image classification. Our work differs by introducing approach for learning hyperbolic representations that models the hierarchy between sub-volumes of complex 3D images, and uses novel hierarchical triplet loss and sampling scheme to capture relationships among multiple levels of granularity in a given input.
In addition, a related field of study has sought to generalize traditional Euclidean neural networks or their components to non- non-Euulidean spaces. Ganea et al. (2018) proposes hyperbolic feed-forward and recurrent architectures based on the theory of gyrovector spaces. Building on this work. CChami al. (2019) propose a hyperbolic graph convolutional network. Other works such as Bachhann et al. 2019); : (Beiigneul & Ganea (2019); Gu et (2019) have also proposed learning with product space of manifolds. Our work generalizes a layer of Ganea et al. (2018) in order to create and use new hyperbolic convolutional layer, which we call the gyroplane convolutional layer.
PRELIMINARIES
Hyperbolic space is non-Euclidean space with constant negative curvature. Curvature measure of the deviation of the geometry from a flat plane (Chami et al. 2019).
Hyperbolic Space
There are five equivalent models of hyperbolic geometry. Following previous work (Mathieu et al., 2019 Ganea et al., 2018 Lou et al., 2020), we use the Poincaré ball model. Hyperbolic space can be considered the continuous version of trees (Nickel & Kiela 2017), making natural choice for embedding hierarchical data. Trees can be embedded in the Poincaré ball with arbitrarily low error (Sarkar 2011), and like the leaves of tree, the area of disc in the Poincaré ball increases exponentially with the radius. Unlike trees, hyperbolic space is smooth, permitting deep learning.
Poincaré ball model of hyperbolic geometry The Poincaré ball (of curvature C = 1) is the open ball of radius 1 centered at the origin equipped with the metric tensor 9p = (Az)² ge. where the conformal factor A, Tー and ge is Euclidean metric tensor (i.e., the usual dot product). Formally, this makes the Poincaré ball a Riemannian manifold. The distance dp between points on the Poincaré ball is given by:
dp(x, v) = =osh (1
2
V v12
(1 +
|10|| (1 1)
The exponential and logarithm maps are a useful way to map from Euclidean space to the Poincaré ball and vice versa (in general, to map from a tangent space to a Riemannian manifold and vice versa). On the Poincaré ball, the exponential and logarithm maps have the closed forms
>xe|vul> 2
V |w||
2 tanh -1I 1 vl1)
+ y
(u) = + tanh
logg (v)
z @ y/
where denotes Mobius addition, which was first introduced by Ungar (0001) as way to define vector operations on hyperbolic space (see Appendix).
4 METHODS
In this section, we describe our approach for learning hyperbolic representations of subvolumes (3I) patches) from 3D voxel grid data. We propose model that comprises 3D convolutional variational autoencoder (VAE) with hyperbolic representation space and a new gyroplane convolutional layer, along with novel hierarchical triplet loss and multi-scale sampling scheme that facilitates learning hierarchical structure within the hyperbolic latent space. To produce segmentations, we cluster the leamned hyperbolic representations. In Section 4.1 we deseribe VAE framework as well as our proposed gyroplane convolutional layer and hiearcchicll triplet loss. In Section 4.2 our approach of hyperholic clustering for segmentation.
we introduce
3

---

Under review as a conference paper at ICLR 2021
4.1 UNSUPERVISED HYPERBOLIC REPRESENTATION LEARNING
3D Hyperbolic VAE framework The VAE framework (Kingma & Welling 2013 Rezende et al., 2014) is widely used for unsupervised representation learning, but requires new innovations to lean effective hierarchical representations in 3D image data. Our proposed hyperbolic VAE consists of 3D convolutional encoder which maps sampled 3D patches of the input volume into hyperbolic space and produces the parameters of the variational posterior, as well as a 3D convolutional decoder which reconstructs the patch from sampled latent hyperbolic representations. The last layer of the encoder is an exponential map that ensures that the output is in hyperbolic space, and the first layer of the decoder is our proposed gyroplane convolutional layer which maps hyperbolic space to Euclidean space. This ensures that both the encoder and decoder respect the hyperbolic geometry of the latent space. We use the wrapped normal distribution as our prior and posterior distribution (see Appendix). Figure iillustrates overview of this VAE framework.
Our variational autoencoder takes as imput a patch of fixed size m X m X m. This allows our model to learn representations of subvolumes that can subsequently be used to perform voxel-level segmentation of the whole 3D volume. To learn hierarchical structure in the 3D scene of each input, generate training examples using multi-scale sampling scheme that samples patches of size r x r X r, where r is randomly sampled. We use two sampling schemes, one for input of smaller sizes and one for input of larger sizes. In both schemes, for a given 3D volume, we sample i patch centers uniformly.
In the sampling scheme for smaller inputs, we sample U(Cnin: maxx) where min: max are hyperparameters. The patch is then upsampled or downsampled to size mx Xm. For larger inputs, we observe that semantic changes tend to occur on a logarithmic scale, so we instead first sample e ~ [f fmaxx and then set r = 2". This sampling scheme is motivated by the intuition that for larger patches, a small change in r less likely correspond to significant semantic difference.
Gyroplane convolutional layer Since R- = =..... R. nigh-dimensional Euclidean spaces can be decomposed into a product of low-dimensional Euclidean spaces. An equivalent decomposition does not hold for arbitrary Riemannian manifolds, making it difficult generalize the usual (Eu- clidean) convolutional layer to arbitrary Riemannian manifolds. For manifolds that are products of manifolds, we can generalize the usual convolution by replacing the Euclidean affine transformation with an affine transformation on the manifold. For the Poincaré ball, one analogue of the Euclidean affine transformation the gyroplane operator fa.p (see Appendix). The details are as follows: for simplicity, suppose x is 4D tensor containing elements of the Poincaré ball and our kernel size is k le x k; with an odd ie value. Our gyroplane convolutional layer is defined as:
++ 1+12/2
s+ L+/2/
1+/2/2
Σ
Σ
Σ
Dr,stt
[A/2]
[/2]
14/2]
Our gyroplane convolutional layer can be extended in the same way as Euclidean convolutional layers to incorporate even kernel size k, input and output channels, padding, stride, and dilation. Our model's 's encoder mean output (p in Figure [ can be interpreted as a product of Poincaré balls, justifying our definition and use of the gyroplane convolutional layer.
Hierarchical triplet loss As our model is trained on patches of the whole 3D volume, the hi- erarchical structure of the volume is not readily apparent from the individual inputs. To help the model infer hierarchical structure, we provide self oef-suuervvisio in the form of a hierarchical triplet loss where positive examples are sub- aub-patches of an anchor patch and negative examples are patches that do not overlap with the anchor patch.
To sample 3D patches for the triplet loss, we first generate an anchor patch centered at voxel with size r x r x r according to one of the above sampling schemes. A positive child patch is generated as a smaller sub-patch of the anchor patch as follows: the positive child patch centered at with size Tchild X Tehild X rchild, where Fehild ~ ((Emin, r -g and Tgap is hyperparameter representing the gap in size between the anchor size and the child size. A negative child is patch of size rehild X rehild X Tehild centered at Uneg, where Uneg is sampled uniformly from the set of voxels w such that a patch of size Tehild X Tchild X Tebild centered at w does not overlap with the anchor patch.
4

---

evidence lower bound Under review as a conference paper at ICLR 2021
Sampled anchor, positive child, and negative child from input volume
Hierarchical triplet loss
Uoos
pos
Euclidean Hyperbolic
neg
Gyroplane convolutional
Input volume
Sampled 3D patch
W Wrapped normal distribution
layer
Exponential map
3D
X
Conv
Z
Conv Decoder
X"
Encoder
Softplus
Our method learns hyperbolic representations of subvolumes of 3D voxel grid data, through a 3D hyperbolic VAE with gyroplane convolutional layer. We enhance the VAE training objective with a self-superrviee hierarchical triplet loss that facilitates learning hierarchical structure within the VAE's hyperbolic latent space.
Figure 1:
Our choice of positive and negative patches motivated by the compositional hierarchy of 3D volumes. Our hierarchical triplet loss encourages the anchor patch and sub-patch (positive child) to have similar representations, while encouraging the anchor patch and a distant patch (negative child) have dissimilar representations. In hyperbolic space, this has the interpretation of belonging to the same hierarchy and belonging to different hierarchies, respectively. We learn hierarchy within 3D image through this triplet loss.
The hierarchical triplet loss can be formulated with any dissimilarity measure between the encoder outputs μ (see Figure T of the anchor Up. positive child Hposs and negative child Pneg- For our model, we take d to be the Poincaré ball distance dp and define our triplet loss with margin a as:
Ltripiet (/4p Pposs /ueg) == max(0, dp /Mpos) dp (Ap- Mnog) + a)
(4)
This formulation can be extended to any metric space by taking the dissimilarity measure to be the space's 's metric. In particular, for our ablations using an Euclidean latent space we take the dissimilarity measure d to be the Euclidean distance.
Optimization We optimize a loss function that be decomposed (ELBO) loss and our new hierarchical triplet loss that encourages the learning of hierarchical struc- ture in the latent representations. The total loss be formulated as Leotal LELBO + BLeriplett where / is hyperparameter that controls the strength of the triplet loss.
4.2 SEGMENTATION BY CLUSTERING REPRESENTATIONS
Hyperbolic clustering In 3D segmentation, we seek assign each voxel v segmentation label sy {1....,n), n., where is the number of segmentation classes. We perform segmentation by clus- tering the representations of patches centered each voxel. We first generate latent representations My for each voxel by running our trained VAE upsampled downsampled encoder input size m x m m if necessary. We then cluster the fh into clusters, and produce segmentation by assigning each the cluster label of μ4, Clustering done using a k-means algorithm that respects hyperbolic geometry, which we derive by replacing the Euclidean centroid and distance computations of classical k-means with their appropriate coun- terparts in Riemannian geometry, the Fréchet mean and manifold distance. We calculate the Fréchet mean using the algorithm of Lou al. (2020).
patch of fixed size p X p centered at v,
EXPERIMENTS
Though method is general to any 3D voxelized grid data, evaluate several biomedical datasets due to the availability of annotated 3D voxel data in the field. We evaluate our method quantitatively both synthetic 3D toy dataset simulating biological image data, well as the
5

---

Under review as a conference paper at ICLR 2021
Figure 2: Sampled 2D slices from 3D volume in our biologically-inspire. toy dataset. The top row showcases the raw input data, and the bottom row showcases the ground truth segmentation.
BraTS tumor segmentation dataset. Our biologically-inspire toy dataset allows quantitative eval- uation of segmentation at multiple levels of hierarchy, while the BraTS dataset is a well-known benchmark for 3D MRI segmentation. We also demonstrate the use of unsupervised segmentation for discovering new biological features in real-world cryo-EM data.
For all models, the encoder of variational autoencoder is comprised of four 3D convolutional layers with kernel size 5 of increasing filter depth {16, 32, 64, 128). The decoder is of the same structure, except with decreasing filter depth and gyroplane convolutional layer the initial layer. We use / = 1e3 as the weighting factor between LELBO and Lnripler and a = 0.2 as the triplet margin, and train the model using the Adam optimizer (Kingma & Ba 2014). We fix the representa- tion dimension to be d = 2. For training on the toy dataset, we sample 3D volume sizes uniformly, and for BraTS and the cryo-EM dataset we sample using an exponential scale (see Section 4.1). For inference, we obtain the latent representations of 5 x 5 x 5 patches densely across the full volume, and then perform hyperbolic k-means clustering, where the number of clusters eter that controls the granularity of the segmentation. For quantitative evaluation, we then use the Hungarian algorithm (Kuhn 1955) to match each predicted segmentation class with corresponding ground truth label.
a hyperparam-
5.1 BIOLOGICALLY TOY DATASET
Since most 3D image datasets not annotated multiple levels of hierarchy, first generate a hierarchical toy dataset to enable more thorough evaluation of the effectiveness of our model for unsupervised 3D segmentation. We note that datasets such ShapeNet (Chang al.) 2015) unsuitable since they contain 3D shape models instead of 3D voxel grid data, which is the focus of our work. Our toy dataset is inspired by cryo-EM tomograms of cells. Each volume in our toy dataset contains multiple levels of hierarchy with objects at each level determined by texture and size. Figure Saaure2ssooee shows an example input volume with sampled slices shown.
Each 3D image of our toy dataset consists of a background and large sphere which represents a cell, which we will refer as Level of the image's 's hierarchy. The large sphere contains medium- size sphere and cube meant to represent cellular substructures such as vesicles, which we will refer as Level 2. In turn, each of these shapes contains two smaller objects of the same shape in Level 3. The color, size, and location of each shape vary randomly. We also apply biologically realistic noise in the form of pink noise. More details can found in the Appendix.
To measure the ability of our model to capture the hierarchy of the toy dataset, we separately evaluate on the three levels of hierarchy defined above and use the average class DICE score to compare segmentation performance. Since our model is unsupervised, segmentation classes are assigned to ground truth labels using the Hungarian algorithm. See results Table and Table S
Comparison with prior approaches Table [ssows quantitative comparison of our method with prior state-of-the-an 3D unsupervised, 2D unsupervised (which we extend to 3D), and semi- supervised models. As unsupervised 3D segmentation a relatively unexplored field, we provide these baselines with different levels of supervision for additional reference. ÇÇçękk et al. (2016) was trained with 2% of the ground truth slices each of the zy, yz. and x² planes, and Zhao et al. 2019) trained with one fully annotated atlas. Di et al. (2019) was implemented using the thors" original code and extrapolated to 3D. For Nalepa et al. (2020) and Morrya et al. (2018), we re-implemented their methods as the original code unavailable. Our model performs signifi-

---

Under review as a conference paper at ICLR 2021
Table 1: Comparison with prior approaches toy dataset
Dice Level I
Dice Level 2
Dice Level 3
Supervision type
Ciçek et al.
(2016)
0.96843
0.82867
0.66771
3D Semi-supervised
Zhao et al.
2019)
0.98892
0.65478
0.35651
3D Semi-supervised
Nalepa ct al.
(2020)
0.53017
0.27612
0.11997
3D Unsupervised
al. (2019
0.58866
0,29086
0.14999
2D to 3D Unsupervised
Morıy et al.
2018)
0,62777
0.31120
0,14131
3D Unsupervised
Ours
0.95211
0.54065
0.21623
3D Unsupervised
Table 2: Ablation studies on toy dataset
Latent Space
Configuration
Dice Level Dice Level
Dice Level 3
Euclidean
Base
0.78370
0.32170
0.10890
Triplet
0.76111
0.34202
0.15349
Hyperbolic
Base
0.83231
0.35185
0.13528
GyroConv
0.90468
0.47297
0.20363
Triplet
0,94522
0.53392
0.22217
GyroConv Triplet
0.95211
0.54065
0.21623
cantly better at all levels of hierarchy compared to its unsupervised counterparts, and comparably to the semi-supervised approach of Zhao al. (2019).
Ablation Table 8 presents ablation studies on the hierarchical toy dataset comparing our contribu- tions: Euclidean vs. hyperbolic representations. the addition our gyroplane convolutional layer, and the addition of our hierarchical triplet loss. The Base Euclidean configuration is the 3D convo- lutional VAE with Euclidean latent space, no gyroplane convolutional layer, and trained with just the ELBO loss. The Triplet Euclidean configuration adds the hierarchical triplet loss to the base Euclidean configuration. The Base Hyperbolic configuration is the same as the Base Euclidean configuration except with hyperbolic latent space. The Triplet configuration is the hyperbolic ana- logue of the Euclidean Triplet configuration, and GyroConv configurations have the addition of the gyroplane convolutional layer.
Hyperbolic representations outperform their Euclidean counterparts in all experiments. We attribute this the more efficient and better organization of hyperbolic representations. When we intro- duce the hierarchical triplet loss, performance improves significantly for our hyperbolic models, but performance for our Euclidean model does not improve as much, likely due information loss in representing hierarchical input. Introducing the gyroplane convolutional layer shows clear improve- ment over our Base Hyperbolic model, which shows the benefit of having layer that respects the geometry of the latent space. The combination of the triplet loss and gyroplame convolutional layer exhibits the most gain over the Base Hyperbolic model, but only small gains over the model with just the added triplet loss. This shows the importance of the our triplet loss for learning effective hierarchical representations.
5.2 BRAIN TUMOR SEGMENTATION CHALLENGE DATASET
The BraTS 2019 dataset is a public, well-established benchmark dataset containing 3D MRI scans of brain tumors along with voxel ground truth annotations of tumor segmentation masks. The seans are of dimension 200 200 x 155 and have four modalities; we use the FLAIR modality, which is the most commonly used one-modality input. We use the same evaluation metric as in the BraTS challenge, and compare DICE score on whole tumor (WT) segmentation, which detectable solely from FLAIR. There are 259 high grade glioma (HGG) labelled training examples, which we split into 180 train examples, 39 validation examples, and 40 test examples. We do not use the official validation available. Table shows the comparison of our results against prior work; we train all baselines on the specified data split for fair comparison. The only exception is the current state-of-the oftthe-art fully-
test sets because the ground truth annotations for these sets not publicly

---

Under review as a conference paper at ICLR 2021
supervised result (Jiang et al., 2019) in Table which also uses all 4 modalities. We show this for reference as an upper bound; the reported number is trained on the full train set and evaluated on the BraTS test set.
Our best model performs significantly better than the unsupervised basclines, and in addition outper- forms one 3D semi-supervised model. This illustrates the ability of our hyperbolic latent represen- tations to effectively capture the hierarchical structure individual brain seans. We use a granular segmentation with three clusters for quantitative evaluation in order capture the tumor, brain, and background, then use the Hungarian algorithm for assignment. In addition, we also show qualitative results for our model (see Figure 3). which include byproduct segmentations from the same model with different numbers of clusters specified, showcasing additionally discovered features in the scan that could also clinically useful.
Table 3: Table shows comparison BraTS 2019 dataset. Figure shows qualitative example where top left image is a slice from a 3D test volume, and the three other images show segmentations with 2,3,4 4 numbers of clustering centroids respectively, illustrating multiple levels of hierarchy learned.
BraTS dataset
Dice WT
Algorithm type
SOTA Jiang et
2019
0.88796
3D Fully-supervised
Zhao et al. (2019)
0.64826
3D Semi-supervised
Çiçek et al.
2016)
0.75965
3D Semi-supervised
et al. (2019)
0,21076
2D to 3D Unsupervised
Moriya et al.
2018
0.42515
3D Unsupervised
Nalepa al.
2020
0.49503
3D Unsupervised
Ours
0.68391
3D Unsupervised
5.3 CRYOGENIC ELECTRON MICROSCOPY TOMOGRAMS
Finally, we show an example of unsupervised 3D segmentation in real-world scenario where un- supervised discovery important. Cryogenic electron microscopy is a technique that images cells at cryogenic temperatures with a beam of electrons. The value of each voxel is the electron density at that location, and is created through reconstruction from tilt slices 70 degrees from electron microscopy. Cryo-EM tomograms features that are unknown or unexplored. We train our model on three 512 x 512 x 250 cryo-EOM tomograms of cells collected from a research laboratory, and run inference Figure 3 shows segmentations produced by our model on a mitochondria from the evaluation to- mogram, using the proposed hyperbolic embedding space vs. Euclidean embedding space, and at a coarse and finer level of granularity. Unlike the Euclidean approach, the hyperbolic approach discovers a fine-gras fine-grained class corresponding to small features on the mitochondria, which may be macromolecular aggregates. As an example of performing unsupervised discovery with our model, the discovered features can now be investigated for their chemical identities and functions.
rich source of biological data, capturing many subcellular
fourth tomogram.
Input
Ours Hyperbolic
Ours Enclidean
Figure 3: Leftmost image is partial slice from a 3D cryo-EM image. The middle box shows segmentation from our best hyperbolic model, the rightmost box shows segmentation from our best Euclidean model. The segmentations in each box correspond to clustering using 2 vs. 4 classes.
8

---

Under review as a conference paper at ICLR 2021
6 CONCLUSION
We propose method for learning hyperbolic representations of subvolumes in 3D voxel grid data, that is based on hyperbolic 3D convolutional VAE with new gyroplane convolutional layer that respects hyperbolic geometry. We enhance the VAE training objective with self-supervised hier- archical triplet loss that facilitates learning hierarchical structure within the VAE's hyperbolic latent space, and a multi-scale sampling scheme. We demonstrate that hyperbolic clustering of learned voxel-level representations can be used to achieve state-of-the-aa unsupervised 3D segmentation,, on hierarchical toy dataset and the BraTS dataset. We also illustrate the promise of using model for unsupervised scientific discovery on an example of cryogenic electron microscopy data.
REFERENCES
Gregor Bachmann, Gary Becigneul, and Octavian-Eugen Ganea. Constant curvature graph convo- lutional networks. arXiv preprint arXiv:1911.0500 1,0666, 2019.
Spyridon Bakas, Hamed Akbari, Aristeidis Sotiras, Michel Bilello, Martin Rozycki, Justin Kirby, John B Freymann, Keyvan Farahani, and Christos Davatzikos, Advancing the cancer genome atlas glioma mri collections with expert segmentation labels and radiomic features. Scientific data, 4: 1701 17, 2017.
Spyridon Bakas, Mauricio Reyes, Andras Jakab, Stefan Bauer, Markus Rempfler, Alessandro Crimi, Russell Takeshi Shinohara, Christoph Berger, Sung Min Ha, Martin Rozycki, et al. Identifying the best machine learning algorithms for brain tumor segmentation, progression assessment, and overall survival prediction in the brats challenge. arXiv preprint arXiv:1811.0262) 2018.
Gary Becigneul and Octavian-Eugen Ganea. Riemannian adaptive optimization methods. In Inter- national Conference on Learning Representations, 2019.
Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsu- pervised learning of visual features. In Proceedings of the European Conference on Computer Vision (ECCV), pp.
2018
Ines Chami, Rex Ying, Christopher Re, and Jure Leskovic. Hyperbolic graph convolutional neural networks. Advances in Neural Information Processing Systems 32 (NIPS 2019), 2019.
Angel X. Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, Jianxiong Xiao, Li Yi, and Fisher Yu. ShapeNet: An Information-Rie 3D Model Repository. Technical Report 1cXiv:1512.03011 [cs.GR], Stanford University - Princeton University - Toyota Technological Institute at Chicago, 2015.
Jianxu Chen, Lin Yang, Yizhe Zhang, Mark Alber, and Danny Z Chen. Combining fully convo- lutional and recurrent neural networks for 3d biomedical image segmentation. In Advances in neural information processing systems, pp. 3036-3044, 2016.
Özgün Çiçek, Ahmed Abdulkadir, Soeren S Lienkamp, Thomas Brox, and Olaf Ronneberger. 3d u- net: learning dense volumetric segmentation from sparse annotation. In International conference on medical image computing and computer-assisted intervention, pp. 424-432, Springer, 2016.
Adrian Dalca, John Guttag, and Mert Sabuncu. Anatomical priors in convolutional networks for unsupervised biomedical segmentation. P Proceedings of the IEEE Conference on Computer Vision and Paltern Recognition, pp. 9290-9299, 2018.
Qi Dou, Lequan Yu, Hao Chen, Yueming Jin, Xin Yang, Jing Qin, and Pheng-Am Ann Heng. 3d deeply supervised network for automated segmentation of volumetric medical images. Medicul image analysis, 41:40-54, 2017.
Octavian- -uggn Ganea, Gary Becigneul, and Thomas Hoffmann.. Hyperbolic neural networks. Ad- vances in Neural Information Processing Systems 3/ (NIPS 2018), 2018

---

Under review as a conference paper at ICLR 2021
Albert Gu, Frederic Sala, Beliz Gunel, and Christopher Re. Learning mixed-curvatuie represen- tations in products of model spaces. In International Conference on Learning Representations, 2019.
Shir Gur, Lior Wolf, Lior Golgher, and Pablo Blinder. Unsupervised microvascular image segmenta- tion using active contours mimicking neural network. In Proceedings of the IEEE International Conference on Computer Vision, pp. 10722-10731, 2019. 2019.
Mohammad Hesam Hesamian, Wenjing Jia, Xiangjian He, and Paul Kennedy. Deep learning tech- niques for medical image segmentation: Achievements and challenges. Journal of digital imag- ing, 32(4):582-59)) 2019.
Xu Ji, João Henriques, and Andrea Vedaldi. Invariant information clustering for unsupervised image classification and segmentation. In Proceedings of the IEEE International Conference Computer Vision, pp. 9865-9874, 2019.
Zeyu Jiang, Changxing Ding, Minfeng Liu, and Dacheng Tao. Two-stage cascaded u-net: 1st place solution to brats challenge 2019 segmentation task. In International MICCAI Brainlesion Work- shop, pp. 231-241, Springer, 2019.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv: /4/2. 2014.
arXiv preprint
Diederik Kingma and Max Welling. Auto-encoding arXiv:1312.6114, 2013.
variational bayes.
arXiv preprint
Titinunt Kitrungrotsakul, Xian-Hua Han, Yutaro Iwamoto, Lanfen Lin, Amir Hossein Foruzan, Wei Xiong, and Yen-Wei Chen. Vesselnet: A deep convolutional neural network with multi pathways for robust hepatic vessel segmentation. Computerized Medical Imaging and Graphics, 75:74-83, 2019.
Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):83-97, 1955.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yamnllecmn.c wewpu..com//xxd//misstttt 2, 2010.
Aaron Lou, Isay Katsman, Qingxuan Jiang, Serge Belongie, Ser-Nam Lim, and Christopher De Sa. Differentiating through the fr). 'echet mean. arXiv preprint arXhv:2003.00335, 2020.
Emile Mathieu, Charline Le Lan, Chris Maddison, and Yee Whye Tee Ryota Tomioka. Contin- uous hierarchical representations with poincaré variational auto-encoders. Advances in Neural Information Processing Systems 32 (NIPS 2019), 2019.
Bjoern H Menze, Andras Jakab, Stefan Bauer, Jayashree Kalpathy-Cramer Keyvan Farahani, Justin Kirby, Yuliya Burren, Nicole Porz, Johannes Slotboom, Roland Wiest, et al. The multimodal brain tumor image segmentation benchmark (brats). IEEE transactions medical imaging, 34 (10):1993-2024 2014.
Takayasu Moriya, Holger R Roth, Shota Nakamura, Hirohisa Oda, Kai Nagara, Masahiro Oda, and Kensaku Mori. Unsupervised segmentation of 3d medical images based on clustering and deep representation learning. In Medical Imaging 2018: Biomedical Applications in Molecular; Structural, and Functional Imaging, volume 10578, pp. 1057820. International Society for Optics and Photonics, 2018.
Yoshihiro Nagano, Shoichiro Yamaguchi, Yasuhiro Fujita, and Masanori Koyama. A wrapped nor- mal distribution hyperbolic space for gradient-based based learning. Proceedings of Machine Learn- ing Research, 97, 2019.
Jakub Nalepa, Michal Myller, Yasuteru Imai, Ken ichi Honda, Tomomi Takeda, and Marek Anto- niak. Unsupervised segmentation of hyperspectral images using 3d convolutional autoencoders. IEFE Geoscience and Remote Sensing Letters, pp. 1-5, 2020.
10

---

Under review as a conference paper at ICLR 2021
Maximilian Nickel and Douwe Kiela. Learning continuous hierarchies in the lorentz model of hyperbolic geometry. arXiv preprint arXiv:180.033117, 2018.
Maximillian Nickel and Douwe Kiela. Poincaré embeddings for learning hierarchical representa- tions. In Advances in neural information processing systems, pp. 6338-63477 2017.
Ivan Ovinnikov. Poincar' 'e wasserstein autoeneoder. arXiv preprint arXiv: 19t1::0012000 2019.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. arXiv preprint arXiv:1401.4082, 2014.
Rik Sarkar. Low distortion delaunay embedding of trees in hyperbolic plane. In International Symposium on Graph Drawing, pp. 355-366. Springer, 2011.
Abraham A Ungar. Hyperbolic trigonometry and its application in the poincaré ball model of hy-
perbolic geometry. Computers & Mathematics with Applications, 1377 2001.
Abraham Albert Ungar. A gyrovector space approach to hyperbolic geometry. Synthesis Lectures Mathematics and Statisties, 1(1):1-194, 2008
Zhenlin Xu and Mare Niethammer. Deepatlas: Joint semi- sumi-supervised learning of image registration and segmentation. In International Conference on Medical Image Computing and Computer- Assisted Intervention, pp. 429. Springer, 2019.
Jianwei Yang, Devi Parikh, and Dhruv Batra. Joint unsupervised learning of deep representa- tions and image clusters. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 5147-5156, 2016.
Amy Zhao, Guha Balakrishnan, Fredo Durand, John V Guttag, and Adrian Dalca. Data augmen- tation using learned transformations for one-shot medical image segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3543-5553, 2019.
Hao Zheng, Yizhe Zhang, Lin Yang, Peixian Liang, Zhuo Zhao, Chaoli Wang, and Danny Z Chen. A ensemble learning framework for 3d biomedical image segmentation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5909-5916, 2019.

---

Under review as a conference paper at ICLR 2021
A APPENDIX
A. RIEMANNIAN MANIFOLDS
In this section we give more complete introduction to Riemannian manifolds, of which hyperbolic space is an example. Riemannian manifolds are spaces that locally resemble Euclidean space. To define this mathematically, we first introduce a manifold as a set of points M that locally resembles the Euclidean space IR. Associated with each point E M is a vector space called the tangent spuce at x, denoted TxM, which the space of all directions a curve on the manifold .M can tangentially pass through point x. A metric tensor g defines inner product 8x on every tangent space, and a Riemannian manifold is a manifold . together with a metric tensor g. Distance on Riemannian manifold
defined as the following, Let Y [a + M be curve the manifold M. The length of y is defined to be (o)) dt and denoted L(y)- The distance between any two points x, y on the manifold defined as dx y) inf I(7)). where the inf is taken over all curves yy that begin at x and end at y. This distance makes M a metric space.
The exponential map (xpx(u) T&M + .M a useful way to map vectors from the (Euclidean) tangent space to the manifold. The exponential map is defined as expx(v) 7[), where is the unique geodesic, the shortest possible curve between two points, starting at x with starting direction Intuitively, one can think of the exponential map as telling how travel one step starting from a point x on the manifold in the e direction. The logarithmic map log,( :M + TxM is the inverse of the exponential map, and maps vectors back to Euclidean space.
A.2 GYROVECTOR OPERATIONS IN THE POINCARÉ BALL
Gyrovector operations were first introduced by Ungar 2008) to generalize the Euclidean theory of vector spaces to hyperbolic space. Mobius addition the Poincaré ball analogue of vector addition in Euclidean spaces. The closed-form expression for Mobius addition on the Poincaré ball with negative curvature is Mathieu et al. (2019):
tr y +1+30 20(z, y) + clyy )) + -||==|////
(5)
1 20(z, 3) + c² 1
we recover Euclidean vector addition. Additionally, the analogue of Euclidean vector subtraction is Mobius subtraction, which defined as x Oc y = @ (-v), and the analogue of Euclidean scalar multiplication is Mobius scalar multiplication, which can be defined for a scalar r as (Ganea et al. 2018):
As one might hope anticipate, when
tanh tanh - '(V[l=))) x Ve
1
[x]|
where we also recover Euclidean scalar multiplication when e 0. this paper, we only con- sider the Poincaré ball with fixed constant negative curvature which allows to drop the dependence on c.
A.3 WRAPPED NORMAL DISTRIBUTION
The importance of the normal distribution in Euclidean space has led many attempts to generalize the normal distribution to Riemannian manifolds. The wrapped normal distribution is one popular way to do this (Mahhieu al.) 2019 Nagano et al., 2019). The wrapped normal distribution can be defined on an arbitrary Riemannian manifold as the push-forward measure obtained by mapping the normal distribution Euclidean space along the manifold's exponential map. The probability density function of the wrapped normal with mean μ and covariance Σ is:
(p, 2)
N E) = =
(7)
sinh(dd
where the subscripts P, E indicate whether the distribution is over the Poincaré ball or Euclidean space, respectively. To use the wrapped normal in VAE, we require both way to sample from the wrapped normal as well as a way to train its parameters. Mathieu et al. (2019) provides a reparametrization and sampling scheme for the wrapped normal on the Poincaré ball.
12

---

Under review as a conference paper at ICLR 2021
A.4 GYROPLANE OPERATOR
The gyroplane layer can be thought of as a hyperbolic affine transformation, and is motivated by the fact we can express Euclidean affine transformation as (a, p-p) sign (a. 2-p. ||a (z, Ha.r) (Ganea et al. 2018), where dg is Euclidean distance and H.pp = {= E IR' Ia - p) = =0.. H
called the decision hyperplane. Ganea et al. (2018) defined the gyroplane operator fa.p from this formulation by replacing each component with its hyperbolic equivalent:
fa.p(2) (2) sign (fa, (=g) [al Ha.p)
(8)
where Ha.p is the hyperbolic decision boundary Ha.p - f E Ba, lo%p(())) = 0}. and the distance to the hyperbolic decision boundary dp Ho.p) is
p 2, a)
dp(2, Ha.p) sinh-1
(1- - + 2 |3]||a|
A.5 TOY DATASET
Our ielogicaallyinssiiee toy dataset has 120 total volumes, which we split into 80 train examples, 20 validation examples, and 20 test examples. Each toy volume in our dataset is 50 x 50 x 50 and contains multiple levels of hierarchy.
The first level of hierarchy (Level 1) is a an outer sphere centered in the volume of radius (2,5,1)) 1). Using cell analogy, this represents the outer cell. The second level (Level 2) consists of spheres ("vesicles") and cuboids ("mitochondriaa), both of which are textured, hollow, and completely contained within the outer cell wall. The positions are randomly sampled with ra- dius of r ~ (88, 0.5) and with side length of 8 ~ 2 -W 0.5). In the third level (Level 3) we introduce small spheres and cuboids ("proteins")) in the vesicle spheres and mitochondria cubloids respectively. The Level 3 proteins are randomly positioned with radius of r N22 0.2) and with side length of 8 (33, 0.15).
Each instance of a shape with a particular size is also given its own unique texture to mimic the different organelles of the cell. The color of each object chosen randomly, according to a standard normal distribution. We also apply pink noise with magnitude m 0.25 to the volume as it is commonly seen in biological data.
In addition, we have also added biologically-inspieed toy dataset with irregular shapes for evaluat- ing datasets with different characteristics. This dataset was created through applying smooth noise to the boundaries of each shape. Specifically, we generate smooth noise by first sampling random points in our voxel grid and random values according to a Gaussian distribution, and interpolate to get a smooth noise. We then use this smooth noise function to perturb the points that fall within the interior of the three largest shapes. See an example of the dataset in Figure 4
We demonstrate our method's performance in comparison prior work on the aforementioned irregular dataset in Table and an ablation study applied on the same irregular dataset in Table both with error bars over four independent runs.
We note that in Table 4 our proposed method outperforms prior work significantly on the irregular dataset, following our initial observations from Ta show state-of-the-an performance. We can see that while all methods show slight decrease in performance, our method is still able to maintain the lead in performance as compared to prior work across all levels.
For ablations on the irregular toy dataset in Table we find that our best models with hyperbolic latent space still outperform models with Euclidean latent space, as with our original toy dataset. We also demonstrate that the gyroplane convolutional layer and hicrarchical triplet loss are both effective compared to the base hyperbolic configuration. However, despite being effective compared to the base hyperbolic configuration, models with hyperbolic hierarchical triplet loss performed less well across the board as compared to the original toy dataset. We hypothesize that this is due to the specific challenges that the irregular dataset brings, for example, needing to recognize noisy instances of irregular shape as the same class in Levels 2 and 3. Therefore, our proposed gyroplane convolutional layer by itself is able to add more effective learning capacity, and shows significant improvement. The added hierarchical triplet loss performs less well on the irregular dataset than in our original toy dataset because in our multi-patch sampling method, each patch is sampled at
13

---

Under review as a conference paper at ICLR 2021
Figure 4: Sampled 2D slices from two examples of 3D volumes in our irregular biologically toy dataset, showing large variance in shapes across input. For each 3D volume example, the top row showcases the raw input data, and the bottom row showcases the ground truth segmentation.
Table 4: Comparison with prior approaches on irregllar toy dataset
Dice Level
Dice Level
Dice Level 3
Supervision type
Çiçek et al.
2016
0,93459
0.76226
0.65480
3D Semi-supervised
7hho et al.
2019)
0.97308
0.63275
0.30164
3D Semi-supervised
Nalepa al. (2020)
0.51008
0.24476
0.09324
3D Unsupervised
J1 et al. (2019)
0.53370
0.24476
0.18865
2D to 3D Unsupervised
Moriya et al.
2018
0.54804
0,23930
0.13181
3D Unsupervised
Ours
0.950 +0.006
0.474 +0.063
0.202 +0.026
3D Unsupervised
random capturing parts of the 3D input. Since the boundary of the shape changes in every image, with random sampling learning more difficult for our hierarchical triplet loss. We don't 't see the same phenomenon for Level since background/foreground segmentation is an easier task. We conclude that with the level of irregularity added to our dataset (see examples Figure 4 the gyroplane convolutional layer with the hyperbolic latent space provides more effectiveness than the triplet loss.
We also note that in real-world datasets, such as in work in cryogenic electron microscopy, the overall shapes of each class of object is similar, and do not contain such dramatic irregularity. For example, vesicles almost-cireular ellipses with only slight eccentricity (deformations with slight stretch), but without distinctive irregularities and protrusions in our irregular dataset. Over- all, our experiments demonstrate that different components of our method are useful for different scenarios, and that our method overall robustly outperforms prior work across data with different characteristics. All hyperbolic configurations of our method seen in Table 4 outperform past unsu- pervised methods, and our approach of learning hyperbolic representations of complex 3D data for segmentation more effective than methods with canonical Euclidean representations.
Last, for runtime on the toy datasets, our implementations of the proposed models take between five to eight hours to train on single Titan RTX GPU for both Euclidean and Hyperbolic variants. We note that for our current implementation, hyperbolic k-means clustering takes on the order of few hours versus minutes for Euclidean k-means. However, this is because we are using our own unoptimized implementation based on recent research in fast Frechet mean algorithms, and standard packages such as scikit-learn do not include hyperbolic k-means algorithms. The Euclidean k k-means algorithms in these packages heavily optimized with parallelization. We anticipate that such optimization would bring the hyperbolic k-means's 's runtime to the order of the Euclidean k- k-means, the computational complexity of the algorithms similar in practice.
14

---

Under review as a conference paper at ICLR 2021
Table 5: Ablation studies irregular toy dataset
Latent Space
Configuration
Dice Level I
Dice Level
Dice Level 3
Euclidean
Base
0.705 = 0.161
0.349 + 0.102
0.148 = 0.066
Triplet
0.906 + 0.037
0,506 0.078
0.201 = 0.061
Hyperbolic
Base
0,845 = 0.123
0,413 + 0.129
0.180 = 0.052
GyroConv
0,947 = 0.006
0.530 + 0.008
0.220 = 0.007
Triplet
0,951 + 0.002
0.471 + 0.011
0.192 = 0.005
GyroConv & Triplet
0.950 + 0.006
0.474 + 0.063
0.202 = 0.026
A.6 BRATS DATASET
We also conduct ablation study on the BraTS dataset with each of our added components with emoor bars over independent runs. Results are shown in Table We can see that our best Hyperbolic model outperforms our best Euclidean model significantly. The addition of the triplet loss improved both Euclidean and Hyperbolic models, while the Hyperbolic models see more improvement due to ability to encode hierarchy in hyperbolic latent space. Our gyroplane convolutional layer also improves performance, while both our additions jointly improve upon our Hyperbolic baseline, showing the benefit of these added components to learning effective representations.
Table 6: Ablation study for BraTS dataset. We report the mean and standard deviation of DICE scores for independent runs.
Latent Space
Contiguration
Dice
Euclidean
Base
0.388 + 0.022 0.517 + 0.050
Triplet
Hyperbolic
Base
0.414 + 0.017
GyroConv
0.539 0.014
Triplet
0.610 + 0.028
GyroConv Triplet
0.692 0.009
We include the average and 95 percentile Hausdorff distance as complementary evaluation met- method compared other
rics the BraTS dataset. See Table We show performance unsupervised baselines, our model outperforms all prior methods on both metrics.
Table Comparison of our method against prior unsupervised work in Hausdorff distance.
Average Hausdorff
95% Hausdorff
Moriya al. (2018)
118.1439
170.434
J1 et al. (2019
96.865
1 114.400
Nalepa et (202
87.704
110.803
Ours
77.940
97.641
A.7 EVALUATION
We DICE score quantitatively evaluate segmentation performance. The DICE score is defined as the following:
2TP
DICE 2TP + FN++FN+F
(10)
where TP is the number of true positives, FN is the number of false negatives, and FP is the number of false positives. For our toy dataset, we first assign predicted classes to ground truth labels using the Hungarian algorithm Kuhn (1955), then evaluate using the average class DICE score. For the BraTS dataset [Menze et al. (2014); Bakas et al. (2017) 2018), we evaluate DICE of the whole tumor segmentation following official evaluation guidelines.
15

---

Under review as a conference paper at ICLR 2021
We also use Hausdorff distance to evaluate the worst-case performance of model. For two sets of points A, B, the directed Hausdorff distance from A to is defined as
((A, B) max aεA
min d(a, b) bEB
(11)
where is any distance function. distance is then defined to be
We will take d to be the Euclidean distance. The Hausdorff
H (AA, B) = max (h(A, (B,
(12)
The official BraTS evaluation uses 95 percentile Hausdorff distance as measure of model robustness (Bakas et al. 2018
A.8 HYPERPARAMETERS
We single set of hyperparameters all of our evaluation datasets, and these hyperparameters are not tuned on any of the evaluation datasets. In order to obtain reasonable set of hyperparam- eters, created completely separate synthetic dataset on which we trained models and tuned hyperparameters. This synthetic dataset was created in similar manner to our toy dataset; how- ever, we designed it to have different and fewer objects, simpler nesting structure, no noise, and fewer textures. The application of this single set of hyperparameters to our evaluation datasets our toy dataset, the BraTS dataset, and the cryogenic electron microscopy dataset, demonstrates robustness of our approach.
With the synthesis dataset, we tuned over a range of hyperparameter values using its validation
set. This includes weight of triplet loss 3 {10 2 1,1001,1 10 102 103 10 100} with the final weight used as 3 103, The patch size for inference was tuned with range = (5, 10, 15, 20, 40} with chosen size as 5 5 x 5. The number of epochs was tuned with range {3, 5, 8, 10, 12, 15) with final epoch number 8.
The BraTS 2019 Menze al. (2014); Bakas et al. 2018)) can be downloaded following directions from https:/ // www. med. upenn.. edw/ cbii ca wtta/baass219//reissttttooo html We will release toy dataset with the final code release.
A.9 MULTI-PATCH SAMPLING
Our method is designed to model the compositional hierarchy of 3D data, where we often find visual substructures contained within other structures. Based on this idea, we sample triplets of 3D volume patches that capture this notion of hierarchical structure. Triplets are sampled through the following process: First, we sample 3D patch of data to be the anchor element, and consider this to be the parent in the triplet. Second, we sample smaller patch of data that is completely contained within the parent patch, and consider this to be the positive child patch. Then, we sample a smaller patch of data that does not overlap with the anchor patch, and consider this to be the negative child patch. See Section 4.1 for further details on sampling procedure. We input the (parent, positive child, negative child) tuples into our hierarchical triplet loss, where the loss encourages the anchor parent and positive child to have closer representations relative to the anchor and the negative child. See Figure for an overview.
Input volume
Positive child
Hierarchieal triplet loss
Upos
3D Conv Encoder
Sampled anchor patch
Up
X
Uneg
Negative child
Figure 5: Example of multi-patch sampling procedure with sampled anchor patch, positive child, and negative child.
16

---

Under review as a conference paper at ICLR 2021
A.10 LATENT DIMENSION ABLATION
In Section 5.1, Section 5.2 and Section 5.3 our experiments were all run with latent dimension of two. To show the effect of higher latent space dimensions, report ablation study for both hy- perbolic and Euclidean representations. As expected, the performance increases with dimension for our model with Euclidean latent space, but our model with hyperbolic latent space still outperforms the Euclidean model at all tested dimensions.
Table 8: Ablation study of latent space dimension for Euclidean and Hyperbolic models on the toy dataset. Dice scores for all three levels are reported.
Latent Space
Dice Level
d=2
d=3
d=5
d=8
d=16
Euclidean
Level Level 2
0.95211
0.95943
0.95574
0,94159
0.95350
0.54065
0.53827
0.54959
0.52889
0.54097
Level 3
0.21623
0.21283
0.21850
0.22612
0.22791
Hyperbolic
Level
0.76111
0.83793
0.84664
0.87080
0.87210
Level 2
0.34202
0.36218
0.37751
0,48133
0.49511
Level 3
0.15349
0.17568
0.16543
0.22521
0.22767
17


Reviews:
[{"review_id": "r6TFhYrfZJ", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Overall, the unsupervised 3D segmentation task tackles one major pain point of the medical imaging field, where the professional annotation is hard and expensive to obtain. This work adopts 3D VAE to map 3D patches into a hyperbolic latent space, then apply clustering followed by the Hungarian algorithm to finish the segmentation task. Though with merits on the introduced hyperbolic space and the gyroplane conv, I have the following concerns and questions. \n\nAlthough I totally appreciate the exploration of the fully unsupervised method to do the 3D segmentation to tackle the lacking of annotation problem, my major concern is that a reliable medical imaging system cannot be obtained through this way and won't be clinically meaningful. Results in Table 1 and Table 3 shows that the 3D semi-supervised by training a vanilla 3D UNet only on 2% of the ground-truth data can perform much better than the proposed unsupervised method. I believe even this simple semi-supervised way is closer to the practical scenario, not to mention that a medical AI system needs to be as reliable as possible for doctors to use.\n\nMy questions:\n1) What is the running time comparison between the between the Hyperbolic space and the Euclidean space?\n2) Can the ablation studies done on the BraTS data? I would personally think this way can be more serious and convincing rather than on a synthesized toy dataset.", "extra_scores": {}, "extra_text": {}}, {"review_id": "pGnEAh8XnEj", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper considers learning hyperbolic representations for unsupervised 3D segmentation. Since the general task of producing annotations for 3D data can be expensive (e.g. for segmentation in dense voxel grids), this is an important problem. The paper proposes to learn hierarchical data structures (e.g. 3D biomedical images) with a hyperbolic variational autoencoder. The paper adapts different metric learning approaches, such as triplet loss and computing a Frechet mean on Riemannian manifolds for clustering. \n\nThe paper is easy to read and interesting to the computer vision and biomedical imaging communities. In terms of machine learning theory, there is not much novelty since the paper combines different existing approaches into a single framework. However, the paper tackles a difficult problem which is the representation of hierarchical 3D data with a simple solution.\n\n- I think that the self-supervised aspect of the approach in the hierarchical triplet loss paragraph is important. However, there does not seem to be much emphasis on it, and I did not fully understand how the positive and negative child patches are generated. Could you please clarify that part and maybe illustrate it in a figure?\n\n- Since the approach is used in the unsupervised setup, how are hyperparameters tuned? There is a discussion about hyperparameter tuning in Section A.7 but the setup is not clear to me. Do you use annotations/labels for the validation set? What is the stopping criterion? How do you pick the best set of hyperparameters? This needs clarification.\n\n- It is nice to see an ablation study of different components of the Euclidean vs hyperbolic approach in Table 2. However, it is known that hyperbolic representations perform better than Euclidean representations to representation trees in low-dimensional space. In high-dimensional space, Euclidean and hyperbolic representations generally perform similarly. \nIf I understand correctly, the dimensionality of the wrapped normal distribution is 2 for the toy dataset. Could we have an ablation study of the impact of higher dimensions for both hyperbolic and Euclidean representations? This would be particularly interesting for real-world applications.\nWhat is the dimensionality of the latent space in Section 4.2 and 4.3?\n\nIn conclusion, I think that the paper is interesting, but needs some clarifications.\n\nMinor comment: representation dimensionality d and distance d use the same font.", "extra_scores": {}, "extra_text": {}}, {"review_id": "FjfWacR0LHE", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes to train a variational autoencoder with a hyperbolic hidden layer augmenting the ELBO loss with a hierarchical triplet loss between nested subsets of data suitable for spatial data structured hierarchically. This process results in representations that can be clustered and used for unsupervised segmentation.\n\n# Strengths\n\nObtaining segmentation labels for 3D data is notoriously difficult and error-prone as even the experts are unable to stay focused for a very long time required to label voxels in a large MRI volume, as an example. Practical advances in unsupervised labeling are a welcome addition to application areas relying on 3D segmentation.\n\nThis paper is able to exploit the recently developed approach to hyperbolic convolutions in a variational autoencoder training with a simple addition of a hierarchical triplet loss. Notably, this simple idea leads to performance improvements over other existing unsupervised or self-supervised methods.\n\nDemonstrations on synthetic data and a couple of real datasets, together with comparison to existing methods is a plus that helps to understand the value of the proposed approach.\n\n# Weaknesses and comments\n\nThe limited technical novelty of the contribution, as most of the innovations are although recent but already published elsewhere.\n\nThe application component is not as rigorously evaluated as one would need to assess the real practical applicability of the approach. For example, in the case of the tumor, the size may bear a large meaning, and volume error in addition to the DICE coefficient would provide much-needed information. For electron microscopy, tomograms or even for the tumor to be treated with radiation therapy the Hausdorff distance may be an informative and potentially more appropriate metric of comparison. Basically, it is very difficult to assess experiments in any meaningful way besides - \"this might be interesting\".\n\nNeither of the measurements provides error bars (all tables), so it is impossible to assess the uncertainty (variability of the approach) and for values that are not too far apart estimate if the prediction is even statistically different.\n\nIt would be good to discuss the training process, hyperparameter search, and computational complexity of the proposed method.\n\nCould the encoder be trained directly using the triplet loss without invoking the VA at all?", "extra_scores": {}, "extra_text": {}}, {"review_id": "EzoneBSA-WJ", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors of this manuscript propose an unsupervised learning framework for 3D segmentation of biomedical images. Specifically, the proposed method learns effective representations for 3D patches using variational autoencoder (VAE) with a hyperbolic latent space. Its main contribution lies at that it introduces a new unsupervised learning framework including hyperbolic convolutional VAE and hierarchical triplet loss. This work conducts experiments on toy dataset, the Brain Tumor Segmentation dataset, and cryo-EM data. The experiment demonstrates competitive performance of the proposed method.\n\nMajor Strengths of The Paper\n1. The idea of hyperbolic 3D convolutional VAE is interesting.\n2. The experiment shows that this idea performs favorably against the state-of-the-art.\n\nMajor Weaknesses of The paper \n1. The manuscript contains some grammatical errors, and the writing could also be improved. For example, there are some grammar mistakes like “our work is a first to introduce…”. The paper will be more clear if the authors could further polish the manuscript.\n2. The proposed model is not novel enough. In my personal opinion, the proposed model generalizes the gyroplane layer and hyperbolic representations in previous literature to 3D biomedical images. The multi-scale sampling and triplet-loss are also commonly used mechanisms. It would be better if the authors could elaborate the novelty of the proposed model.\n3. The biologically-inspired toy dataset does not look close to a biomedical dataset. Although we totally understand there should be differences between simulated dataset and real dataset, the toy dataset (i.e. composed of regular shapes) oversimplifies the biomedical image, and makes the experiment results less convincing. It will be better if the authors can simulate images close to real biomedical images, at least with irregular shapes or fuzzy boundaries. \n4. For the baseline selection, the authors used the model in paper “Unsupervised Segmentation of Hyperspectral Images Using 3D Convolutional Autoencoders,” which is published in IEEE Geoscience and Remote Sensing Letters. There should be a domain gap between remote sensing images and biomedical images. I personally think it may not be a fair comparison for the baseline methods. It will be better if the authors could select some of the state-of-the-art unsupervised 3D segmentation methods as baselines.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This application paper applies hyperbolic convolutions in VAE learning to perform unsupervised 3D segmentation. 
Addition of these components enables performance improvements in the unsupervised segmentation task.
Overall, the paper is borderline and the reviewers mention the limited novelty of the approach, which largely uses components that have been developed before. Even though the paper presents an application of these methods to a relevant and significantly more challenging task than prior work, I recommend rejection from ICLR due to concerns about novelty.
Decision: Reject