Learning Hyperbolic
Representations for Unsupervised
3D Segmentation
Joy Hsu@joycjhsu
joycj@stanford.edu

Segmenting coronavirus

Overview

Applications in cryo-EM
Hyperbolic representation
learningSegmentingcoronavirus

Problem introduction
What do you do with an unknown and unexplored input?
How do you discover features of interest from it?

Object discovery

Object discovery
Unsupervised 3D segmentation
Voxel-level class prediction
Focus on biomedical data

Object discovery
Unsupervised 3D segmentation
Voxel-level class prediction
Focus on biomedical dataGranularity of segmentation
Many “correct” segmentations at diﬀerent levels of semantic hierarchy
Diﬃcult to do unsupervised segmentation without base granularity

Coronavirus scans
Slice of 3D tomogram
How do we learn more about spikes oncoronavirus particles?

Automated particle picking
RetinaNet to identify particles

Spike discovery

Our ModelCount spikes

Applications inCryo-EM

CryoEM
Cryogenic electron microscopy
Allows for observation of biological specimens in their native environment at
cryogenic temperatures
Produces high-resolution three-dimensional views of samples
Each voxel represents number of electrons
Abundance of undiscovered information

Imaging
Samples are imaged as they are tilted
Angle restricted by thickness
Series of 2D images are combined to produce a 3D
reconstruction

Volumes

CryoEM process
Imaging
Alignment
Reconstruction
DenoisingParticle identiﬁcation
Model generation

Problem
For Cryo-EM
Identify molecular components within the crowded cellular environment
Large variation within each protein structure
Identiﬁed particles can be used to generate protein structure models

Can we 3d segment features of interests at diﬀerent levels of hierarchy with nosupervision?

Initial approach
VAE based probabilistic approach
Learn semantically relevant representation of sampled 3D volumes

VAE
Encoder retrieves feature representation
Downsamples input

VAE
Sample latent vector z
Learn to approximate true p(z|x) with encoder q(z|x)

VAE
Decoder reconstructs input from latent vector
Reconstruction ensures semantic features kept

Initial approach
Learn downsampled representation of input

Probabilistic VAE pros
Can learn robust and generalizable representations
Can qualitatively visualize representations
Can create potential segmentations

Segmentation results
Use qualitative visualizations to measure performance
Utilize plotly to showcase sampled 3D visualizations after clusteringRepresentation space visualizations
Representation 𝛍 generated of dimension 64
Apply dimension reduction through PCA & uniform manifold approximation
Representation of nucleus & mitochondria & microtubules

Component analysis
Representation 𝛍 of dimension 3Sliding window inference
Component should represent higher levelfeature information

Latent space reconstructions

Apoferritin
Clustering on representation space
High sensitivity for apoferritin
Noise around detected instances
Incorrectly identiﬁed gold ﬁducials ‘X’s

Neuronal cells
Identiﬁed protein aggregate
Can visualize structureNeuronal cells
Learning interesting structureDiﬃcult to capture semantic information

Microtubules
Noise within segmentation
Requires post-processingProcess
Representation space
Component understanding
Segmentation for objects of interest

Need to learn more natural representation

Representation learning
How do we learn the most eﬀective and natural form of representation 𝛍?

Hyperbolic
representationlearningPreliminaries

Introduction to manifolds
A manifold M of dimension d is an object where theneighborhood of every point locally resembles
d-dimensional Euclidean space (R^d)
Deﬁne a tangent space TzM for each point z on M
Think of a manifold as the generalization of a 2D
surface to n dimensions

Riemannian manifolds
A Riemannian manifold is a manifold M
with a metric g at every point x
The metric g allows us to locally deﬁneangles, distance, & area on M

Riemannian manifolds
Deﬁne the inner product at point z

Deﬁne a normRiemannian manifolds
From the matrix representation G deﬁne areas

The length of a curveis deﬁned as

Deﬁne the global distance on M as

The shortest path connecting two points on a manifold is called a geodesic

Riemannian manifolds
Deﬁne the exponential map asRiemannian Manifolds
Deﬁne the logarithm map as the inverse of the exponential map

Warning: the logarithm map doesn’t always exist!

Poincare ball model
The Poincare ball model is one of ﬁve equivalentmodels of constant curvature hyperbolic space
The Poincare ball has closed form equations for all
quantities of interest
Formally the Poincare ball is deﬁned as
The second term is the metric tensor

Gyrovector operations
Mobius addition

Exponential map

Logarithm map

MethodRepresentation learning
Hyperbolic representations eﬀectively learn hierarchical structure
Can embed tree-like structure in hyperbolic spaces with arbitrarily low error

Hyperbolic VAE
Unsupervised representation learningHyperbolic VAE
𝛍 = mean = representation
𝝈 = standard deviationz = latent vector

Methods
Euclidean encoder with 3D convolutional layers to learn featuresDense layers output 𝛍 and 𝝈Methods
Map 𝛍 into hyperbolic space through the exponential map
Apply softplus activation to 𝝈Sampling𝛍 and 𝝈 parameterize the hyperbolic variational posterior
Need to sample with hyperbolic 𝛍 and 𝝈 to get latent vector z
Wrapped normal and Riemannian normal as prior and posterior distribution

Riemannian normal
Distribution maximizing entropy given expectation and varianceWrapped normal
Push forward measure mapping normal distribution along the exponential map

Distributions on the Poincare ballBoth normal distributions are similar
Riemannian normal has slightly larger mode

Methods
Decoder uses gyroplane layer to map hyperbolic representation to Euclidean
On the Poincare ball a gyroplane layer is deﬁned asWith closed form distance equation asGyroplane convolutional layer
Encoder output 𝛍 is produced by 3D convolutional layers
Hence both 𝛍 and sampled vector z can be thought of a 3D volume
Each voxel is represented by a vector in hyperbolic spaceGyroplane convolutional layer
Deﬁne gyroplane 3D convolutional layer as the Euclidean 3D convolutional layer
Euclidean linear operation replaced by a gyroplane operation
Gyroplane layer is a hyperbolic aﬃne transformation from Poincare ball toEuclidean

Methods
Decoder includes 3D up-convolutional layers back to input size

Hyperbolic VAE
Unsupervised representation learning

VAE Euclidean Objective
Goal is to learn true posterior of latent vector z given input x
Minimize KL divergence (information loss) between approximate posterior fromencoder q(z|x) and true posterior p(z|x)
Maximizing ELBO (evidence lower bound) minimizes KL divergence without p(x)Hyperbolic ELBO ObjectiveELBOReparameterization through hyperbolic polar change of coordinates

Details in Appendix of Continuous Hierarchical Representations with Poincaré
Variational Auto-Encoders by Mathieu et al.

Hierarchical objective
Enforce relations in diﬀerent levels of hierarchy
Hierarchical triplet loss samples anchor parent, positive child, and negative childmotivated by the compositional hierarchy of 3D volumes

Hierarchical objective
Encourages the anchor patch (parent) and a sub-patch (positive child) to have
similar representations
In hyperbolic space this has the interpretation of belonging to the same hierarchy
Anchor patch and a distant patch (negative child) to have dissimilar representation
Belonging to the diﬀerent hierarchies

Hierarchical triplet loss

LossExperiments

Clustering
Retrieve representation of sampled patches from input volume through a sliding
window
Perform hyperbolic kmeans
Use Hungarian algorithm for assignment

Evaluation

Biologically-inspired toy dataset

Prior approaches

Ablations

BRATS dataset
3D MRI scans of the brain
Annotated with tumors
Output tumor masks on volume

BRATS resultsPrior approaches

Hierarchical segmentations
Can retrieve diﬀerent levels of segmentations given one modelCryoEMJoy Hsu joycj@stanford.edu
Jeff Gu jeffgu@stanford.edu
Michael Zhang mzhang20@stanford.edu
Serena Yeung syyeung@stanford.edu
Thanks to Gong Her Wu, Jesus Montoya, David Chmielewski,
Michael Schmid, Wah Chiu from SLAC

Questions?