Capturing implicit hierarchical structure in 3D biomedical images with self-supervised hyperbolic representations

1 Introduction . Advances in biomedical imaging techniques such as cryogenic electron tomography ( cryo-ET ) and magnetic resonance imaging ( MRI ) have resulted in an ever-increasing amount of 3D biomedical image data . In these data domains , a growing body of work shows that , when provided with labels , machine learning models achieve good performance on many tasks [ Çiçek et al. , 2016 , Milletari et al. , 2017 , Dou et al. , 2017 , Falk et al. , 2019 ] . However , these labels , especially for segmentation , are very costly as they often have to be provided by experts in the appropriate field . Consequently , supervised learning and even semi-supervised learning remain limited in this setting as ( 1 ) tasks and domains are very diverse , making it intractable for experts to provide labelled data for all problems ; and ( 2 ) experts can only label objects they already know , restricting the potential for scientific discovery using machine learning methods . In this work , we tackle the task of unsupervised segmentation in 3D biomedical image data . Our key insight is that 3D biomedical images have inherent hierarchical structure . For example , in the cryo-ET domain , an image of a cell has a hierarchy that at the highest level comprises the entire cell ; at a finer level comprises organelles such as the mitochondria and nucleus ; and at an even ∗These authors contributed equally . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . finer level comprises sub-structures such as the nucleolus of a nucleus or protein machineries within organelles . Such types of hierarchies are present in many types of biomedical images ( e.g. , nested anatomical structures within MRI images ) . We hypothesize that in the unsupervised setting , models that are able to encode this internal hierarchical structure will provide better data representations for downstream tasks . To that end , we propose learning representations based on embedding subvolumes of 3D images in hyperbolic space . In contrast to traditional Euclidean embeddings , hyperbolic embeddings better preserve hierarchical relationships present in the data . Hyperbolic representations have been proposed as a continuous way to represent hierarchical data , due to their ability to embed trees with arbitrarily low error [ Sarkar , 2011 ] . A recent line of work utilizes hyperbolic representations to model hierarchical data across domains ranging from natural language word taxonomies [ Nickel and Kiela , 2017 , 2018 ] and graphs [ Nickel and Kiela , 2017 , Mathieu et al. , 2019 , Ovinnikov , 2019 , Chami et al. , 2019 ] , to image classification [ Mathieu et al. , 2019 ] . In these settings , the objects , and in most cases their relationships , are explicitly encoded in the data . However , 3D biomedical images consist of subvolumes that represent parts of an implicit hierarchical structure . In our case , for any single 3D voxel-grid , we wish to embed and infer the implicit relationships between all of its subvolumes without any supervision . To embed our 3D images in hyperbolic space , we use a 3D hyperbolic variational autoencoder ( VAE ) . For the decoder , we propose a gyroplane convolutional layer which maps from the latent space back to 3D images while respecting hyperbolic geometry . In addition to the VAE loss , we propose an essential self-supervised loss to capture the hierarchical structure present in the data . More specifically , we consider reconstruction of implicit hierarchies as a pretext task . Concretely , we add a triplet loss which encourages a child subvolume to be mapped close to its parent subvolume in hyperbolic space . To capture hierarchical relationships of varying granularity , we train on subvolumes sampled at multiple scales . Finally , for a specified scale , we cluster the subvolumes in latent space and produce a segmentation map . We evaluate our model on datasets with different domains : synthetic datasets and a medical image dataset . We construct synthetic datasets where we generate structures at various scales and show that our model segments objects at multiple levels of hierarchy better than all prior unsupervised segmentation methods . We demonstrate performance gains ranging from 7 % for the smallest objects to 32 % for the largest objects . On the real-world medical image dataset ( BraTS Brain Tumor Segmentation Challenge ) [ Menze et al. , 2014 , Bakas et al. , 2017 , 2018 ] , we show that our method outperforms prior works by 19 % , and even achieves comparable performance to semi-supervised methods although we do not use any labels . 2 Related Work . Segmentation on 3D voxel data Many diverse biomedical images , ranging from MRI and CT scans to fluorescence microscopy , come in the form of 3D voxel-grids . Since 3D voxel-grids are dense , computer vision tasks such as supervised segmentation are commonly performed using deep learning architectures with 3D convolutional layers [ Chen et al. , 2016 , Dou et al. , 2017 , Hesamian et al. , 2019 , Zheng et al. , 2019 ] . However , due to the challenges of obtaining voxel-level annotations in 3D , there has been significant 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 and Niethammer , 2019 ] , and using one segmented input in conjunction with other unlabelled data [ Zhao et al. , 2019 ] . Unsupervised approaches for 3D segmentation are useful not only for further reducing the manual annotation effort required , which is especially expensive for segmentation , but also for scientific discovery tasks where we would like to identify previously unknown structures for which annotations are impossible to produce . Moriya et al . [ 2018 ] extends to 3D data an iterative approach of feature learning followed by clustering [ Yang et al. , 2016 ] . Nalepa et al . [ 2020 ] uses a 3D convolutional autoencoder architecture and performs clustering of the latent representations . Another approach , [ Dalca et al. , 2018 ] , uses a network pre-trained on manual segmentations from a separate dataset to perform unsupervised segmentation of 3D biomedical images . However , this limits applicability to areas where we already have a dataset with manual annotations and makes it unsuitable for unbiased scientific discovery . Gur et al . [ 2019 ] and Kitrungrotsakul et al . [ 2019 ] develop unsupervised methods for 3D segmentation of vessel structures , but these are specialized and do not generalize to the segmentation of other structures , [ Uzunova et al. , 2019 ] utilizes knowledge of background patches with no patholoogy , and Baur et al . [ 2018 ] uses deep autoencoding models for unsupervised anomaly detection . Another line of work performs unsupervised 2D segmentation , such as Ji et al . [ 2019 ] which proposes a mutual information objective for clustering , and Caron et al . [ 2018 ] , which uses the clustered output of an encoder as 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 learn hierarchical structure of the data , which is the main focus of our work . Hyperbolic representations A recent line of work employs hyperbolic space to represent hierarchical structures , with the intuition that tree structures can be naturally embedded into continuous hyperbolic space [ Nickel and Kiela , 2017 ] . These works utilize hyperbolic representations for a variety of tasks , including MNIST classification [ Mathieu et al. , 2019 , Nagano et al. , 2019 , Ovinnikov , 2019 ] , natural language processsing tasks such as embedding word taxonomies and entailment [ Nickel and Kiela , 2017 , Ganea et al. , 2018 ] , link prediction and node classification [ Chami et al. , 2019 ] , and game playing [ Nagano et al. , 2019 ] . In most of these works , hierarchical structure is explicitly encoded in data . In contrast , we seek to capture implicit hierarchical structure arising from composition within 3D images . Several architectures have been proposed in order to learn hyperbolic representations , including hyperbolic VAEs [ Mathieu et al. , 2019 , Nagano et al. , 2019 , Ovinnikov , 2019 ] , feed-forward and recurrent hyperbolic neural networks architectures [ Ganea et al. , 2018 ] , and hyperbolic graph convolutional networks [ Chami et al. , 2019 ] . We extend the hyperbolic VAE framework to the task of learning hyperbolic representations from subvolumes of complex 3D images , and use this to perform unsupervised segmentation . Self-supervision Providing self-supervision by solving pretext tasks is one common approach for learning unsupervised visual representations . Pretext tasks leverage properties of the input data or prior knowledge as supervisory signals in order to learn better representations . Examples of pretext tasks include finding the relative position of two patches sampled from an image [ Doersch et al. , 2015 ] , solving jigsaw puzzles [ Noroozi and Favaro , 2016 ] , and predicting pixel movements of videos in subsequent frames [ Pathak et al. , 2017 ] . In contrast , we propose the pretext task of reconstructing implicit hierarchy in 3D voxel-grid images , to learn effective hyperbolic representations for downstream segmentation . 3 Methods . In this section , we describe our approach for learning hyperbolic representations of subvolumes ( 3D patches ) from 3D voxel-grid data . We propose a model that comprises a 3D convolutional variational autoencoder ( VAE ) with hyperbolic representation space and a gyroplane convolutional layer . We train our model with self-supervision through a novel hierarchical triplet loss and multipatch sampling scheme . Then , we cluster the learned representations using hyperbolic k-means to produce 3D segmentations . In Section 3.1 , we provide an overview of hyperbolic space . In Section 3.2 , we describe our VAE framework with the proposed gyroplane convolutional layer and self-supervised hierarchical triplet loss . Finally , in Section 3.3 , we discuss our approach of hyperbolic clustering for segmentation . 3.1 Hyperbolic formulation . Hyperbolic space We embed subvolumes of 3D voxel-grid data in hyperbolic space , a nonEuclidean space with constant negative curvature . In negative curvature spaces , the area of a disc increases exponentially with the radius . We can think of this growth as analogous to the exponential increase of leaves at each level of a tree . Hence hyperbolic space can encode trees with arbitrarily low error [ Sarkar , 2011 ] and can be considered as the continuous version of hierarchical structures . Unlike trees , hyperbolic space is smooth , permitting our use of deep learning on representations . For additional background on geometry and hyperbolic space , see the Appendix . Poincaré ball model of hyperbolic geometry In this work we use the Poincaré ball as our 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 gp = ( λx ) 2ge , where the conformal factor λx = 21−||x||2 and ge is the Euclidean metric tensor ( i.e. , the Euclidean dot product ) . Formally , this makes the Poincaré ball a Riemannian manifold . For an introduction to Riemannian manifolds , see the Appendix . The distance dp between points on the Poincaré ball is given by : dp ( x , y ) = cosh −1 ( 1 + 2 ||x− y||2 ( 1− ||x||2 ) ( 1− ||y||2 ) ) ( 1 ) We use the exponential and logarithm maps to map from Euclidean space to the Poincaré ball and vice versa . On the Poincaré ball , we note that the exponential and logarithm maps have the closed form expressions expz ( v ) = z ⊕ ( tanh ( λz||v|| 2 ) v ||v|| ) ( 2 ) logz ( y ) = 2 λz tanh−1 ( || − z ⊕ y|| ) −z ⊕ y || − z ⊕ y|| ( 3 ) where ⊕ denotes Mobius addition , which was introduced by Ungar [ 2001 ] as a way to define vector operations on hyperbolic space . Wrapped Normal Distribution The importance of the normal distribution in Euclidean space has led to many generalizations of the normal distribution to Riemannian manifolds . We use the wrapped normal distribution [ Mathieu et al. , 2019 , Nagano et al. , 2019 ] , which can be defined on an arbitrary Riemannian manifold as the push-forward measure obtained by mapping the normal distribution in Euclidean space along the manifold ’ s exponential map . On the Poincaré ball , the probability density function of the wrapped normal with mean µ and covariance Σ is : NP ( z|µ , Σ ) = NE ( λµ ( z ) |0 , Σ ) ( dp ( µ , z ) sinh ( dp ( µ , z ) ) ) ( 4 ) where the subscripts P , E indicate distributions over the Poincaré ball and Euclidean space , respectively . We use the sampling and reparametrization scheme of Mathieu et al . [ 2019 ] in order to sample and train our VAE using the wrapped normal distribution .

Summary: The authors present a novel representation for learning hierarchical relations within samples (3D image pixels) in an embedding space. A component that that enables this mapping is a novel self-supervised loss function. Results on real data indicate advantageous properties.