#  $\Gamma$-VAE: Curvature regularized variational autoencoders for uncovering
emergent low dimensional geometric structure in high dimensional data

Jason Z. Kim<EMAIL_ADDRESS>Department of Physics, Cornell University,
Ithaca, NY 14853, USA  Nicolas Perrin-Gilbert  Sorbonne Université, CNRS,
Institut des Systèmes Intelligents et de Robotique, ISIR, F-75005 Paris,
France  Erkan Narmanli  Inserm U830, Institut Curie Research Center, PSL
University, Paris, France  Department of Translational Research, Institut
Curie Research Center, PSL University, Paris, France  Paul Klein  Inserm U830,
Institut Curie Research Center, PSL University, Paris, France  Department of
Translational Research, Institut Curie Research Center, PSL University, Paris,
France  Christopher R. Myers  Department of Physics, Cornell University,
Ithaca, NY 14853, USA  Center for Advanced Computing, Cornell University,
Ithaca, NY 14853, USA  Itai Cohen  Department of Physics, Cornell University,
Ithaca, NY 14853, USA  Kavli Institute at Cornell for Nanoscale Science,
Cornell University, Ithaca, NY 14853  Joshua J. Waterfall  Inserm U830,
Institut Curie Research Center, PSL University, Paris, France  Department of
Translational Research, Institut Curie Research Center, PSL University, Paris,
France  James P. Sethna<EMAIL_ADDRESS>Department of Physics, Cornell
University, Ithaca, NY 14853, USA

###### Abstract

Natural systems with emergent behaviors often organize along low-dimensional
subsets of high-dimensional spaces. For example, despite the tens of thousands
of genes in the human genome, the principled study of genomics is fruitful
because biological processes rely on coordinated organization that results in
lower dimensional phenotypes. To uncover this organization, many nonlinear
dimensionality reduction techniques have successfully embedded high-
dimensional data into low-dimensional spaces by preserving local similarities
between data points. However, the nonlinearities in these methods allow for
too much curvature to preserve general trends across multiple non-neighboring
data clusters, thereby limiting their interpretability and generalizability to
out-of-distribution data. Here, we address both of these limitations by
regularizing the curvature of manifolds generated by variational autoencoders,
a process we coin “$\Gamma$-VAE”. We demonstrate its utility using two example
data sets: bulk RNA-seq from the The Cancer Genome Atlas (TCGA) and the
Genotype Tissue Expression (GTEx); and single cell RNA-seq from a lineage
tracing experiment in hematopoietic stem cell differentiation. We find that
the resulting regularized manifolds identify mesoscale structure associated
with different cancer cell types, and accurately re-embed tissues from
completely unseen, out-of distribution cancers as if they were originally
trained on them. Finally, we show that preserving long-range relationships to
differentiated cells separates undifferentiated cells—which have not yet
specialized—according to their eventual fate. Broadly, we anticipate that
regularizing the curvature of generative models will enable more consistent,
predictive, and generalizable models in any high-dimensional system with
emergent low-dimensional behavior.

Many natural systems with high dimensional states give rise to emergent
behaviors that can be described using low dimensional models: the average
behavior of $\sim 10^{23}$ gas molecules can be described using the ideal gas
law, and the expression of $\sim 10^{4}$ genes organize along regulatory
networks with a relatively sparse number of phenotypes [1]. Successful models
of such systems often explain the data using only a few interpretable
dimensions [2, 3], make accurate predictions about new conditions outside of
the training data that have not yet been tested, and elucidate fundamental
governing principles of the phenotypes that emerge in the low dimensional
structure.

The challenge of finding low-dimensional models of high-dimensional data has
spanned over a century [4]. Linear projection methods such as PCA provide
perfectly consistent embedding dimensions as linear subspaces that span the
entire dataset. However, PCA often requires many dimensions to capture the
majority of the data variance, making the dimensions difficult to interpret.
More modern techniques such as UMAP and variational autoencoders (VAEs) learn
low dimensional nonlinear embeddings by using methods that preserve the
similarity between data points [5, 6, 7, 8, 9, 10, 11, 12, 13, 14]. While many
of these state-of-the-art methods excel in preserving the topology and local
structure of the data, they significantly distort the longer-range data
trends, especially in regions where there are no training samples. These
distortions limit the interpretability of the embedding dimensions and make it
difficult to form accurate predictions on out-of distribution samples, both of
which moderate their utility in elucidating the principles underlying emergent
behaviors [15, 16].

Excitingly, recent work has begun to use tools from differential geometry to
more accurately reflect the geometry of the data in the latent representation
by measuring and controlling the curvature of the manifold [17, 18, 19]. Of
particular relevance for the problem of reducing manifold non-linearity is the
development of techniques for regularizing estimates of the Ricci and
extrinsic curvatures to improve reconstruction accuracy [20]. Their pioneering
advance was to stay faithful to the nonlinear geometry of the data rather than
relying on the inductive bias of the neural network architecture to model data
manifolds [21].

Here we build on this approach in two fundamental ways. First, we formulate
the extrinsic curvature as a projection of the second derivative vector rather
than estimating the full projection operator—resulting in computing an
$N$-dimensional vector rather than estimating an $N\times N$-dimensional
matrix—where $N$ is the large dimensionality of the data. This formulation
allows us to scale to datasets with thousands of features using an exact
calculation of curvature rather than an estimate via sampling. Second, instead
of the Ricci curvature, we regularize the parameter-effects curvature [22],
which explicitly accounts for the distortion along every direction on the
manifold. The parameter-effects curvature is also scalable, exactly
calculable, and yields an important theoretical limit: at zero extrinsic and
parameter-effects curvature, the second derivative of the manifold in all
directions is 0, such that the embedding is analagous to PCA.

To demonstrate its utility, we apply this method to both bulk and single-cell
RNA-seq gene expression data. For bulk RNA-seq, we use a joint dataset of both
TCGA and GTEx [23, 24, 25], and demonstrate the ability to capture the
mesoscale organization of healthy and cancer tissues along interpretable axes,
and the accurate and fully out-of-distribution prediction of cancer
phenotypes. Further, we demonstrate the utility of our approach in single cell
RNA sequencing (scRNA-seq) by accurately predicting cell fates from
multipotent cells in a lineage-tracing experiment of hematopoietic stem cell
differentiation [26]. Taken together, our method provides a model manifold
through complex and multi-scale data that preserves interpretable dimensions,
makes accurate out-of-distribution predictions, and uncovers biologically
relevant geometry.

##  I Curvature Regularized VAEs

The standard formulation of VAEs models the data as being generated from a
probability distribution over the data, $\bm{x}_{i}\in\mathcal{R}^{N}$, and a
set of unobserved latent variables $\bm{z}_{j}\in\mathcal{R}^{m}$. The idea is
to find a probability distribution, $p(\bm{x},\bm{z})$, that best matches the
distribution of the measured data after marginalizing over the unobserved
latent variables,

$\displaystyle p(\bm{x})=\int p(\bm{x},\bm{z})\mathrm{d}\bm{z}.$ (1)

However, the direct optimization of Eq. 1 is challenging because it is usually
difficult to calculate.

VAEs use deep neural networks (DNNs) to define an encoder
$q_{\bm{\phi}}(\bm{z}|\bm{x})$ and a decoder $\bm{f}_{\bm{\theta}}(\bm{z})$,
and optimize a bound on Eq. 1, the Evidence Lower Bound (ELBO):

$\displaystyle-\underbrace{\mathbb{E}_{q_{\bm{\phi}}(\bm{z}|\bm{x})}||\bm{x}-\bm{f}_{\bm{\theta}}(\bm{z})||_{2}^{2}}_{\mathrm{reconstruction~{}error}}-\beta
D_{KL}(q_{\bm{\phi}}(\bm{z}|\bm{x})||\underbrace{p(\bm{z})}_{\mathrm{prior}}),$
(2)

where $\bm{\phi}$ and $\bm{\theta}$ are the DNN weights, and the specific form
of the reconstruction error is for a Gaussian decoder. The first term is the
reconstruction error between the true data point, $\bm{x}_{i}$, and decoder’s
reconstruction of that point, $\bm{f}_{\bm{\theta}}(\bm{z}_{i})$, from a
sample in latent space $\bm{z}_{i}\sim q_{\bm{\phi}}(\bm{z}|\bm{x}_{i})$. The
second term is the Kullback-Leibler divergence between the encoded data in
latent space, $q_{\bm{\phi}}(\bm{z}|\bm{x}_{i})$, and the prior, $p(\bm{z})$,
which is usually taken to be a standard normal distribution.

By maximizing Eq. 2 (or minimizing the negative of the ELBO as the loss,
$\bar{\mathcal{L}}$), the encoder makes the latent-space embeddings similar to
a standard normal distribution, and the decoder learns a continuous manifold
from the latent space to data space that best reconstructs the data. Due to
the nonlinearity of the decoder, VAEs often yield sharp variations of this
manifold [27], making the latent representation difficult to interpret. The
hyperparmeter $\beta$ serves to regularize the interpretability of the latent
space (termed $\beta$-VAE [28]), yet does not fully cure the distortion,
leading to the dependence of representations on inductive biases and random
seeds for initialization [29].

To ameliorate this distortion, we desire two properties in our model manifold.
First, we want a unit distance anywhere in latent space to correspond roughly
to the same distance on the manifold. Hence, we regularize the parameter-
effects curvature at each sampled point on the manifold, given by

$\displaystyle\mathcal{L}_{PE,i}=\sum_{a}g^{\mu\bar{\mu}}g^{\nu\bar{\nu}}\left(\Gamma_{\mu\nu}^{\kappa}\frac{\partial\bm{f}^{a}_{\bm{\theta}}}{\partial
x^{\kappa}}\right)\left(\Gamma_{\bar{\mu}\bar{\nu}}^{\bar{\kappa}}\frac{\partial\bm{f}^{a}_{\bm{\theta}}}{\partial
x^{\bar{\kappa}}}\right),$ (3)

where $g^{\mu\nu}$ is the inverse of the metric tensor,
$\Gamma_{\mu\nu}^{\kappa}$ are the Christoffel symbols, and Eq. 3 is written
in Einstein summation notation (repeated Greek indices are summed). At a point
on the model manifold, $\bm{f}_{\bm{\theta}}(\bm{z}_{i})$, Eq. 3 is the norm
of the component of the second derivative of $\bm{f}_{\bm{\theta}}$ that lies
tangent to the manifold at that point. Intuitively, Eq. 3 measures by how much
a regular grid of points in latent space is distorted by the decoder in data
space [22].

Second, we want straight lines in latent space to correspond roughly to
straight lines in data space. Drawing inspiration from prior work, [20, 22],
we also regularize the extrinsic curvature at each point on the manifold,
given by

$\displaystyle\begin{split}\mathcal{L}_{EX,i}=\sum_{a}g^{\mu\bar{\mu}}g^{\nu\bar{\nu}}\left(\frac{\partial^{2}\bm{f}^{a}_{\bm{\theta}}}{\partial
x^{\mu}\partial
x^{\nu}}-\Gamma_{\mu\nu}^{\kappa}\frac{\partial\bm{f}^{a}_{\bm{\theta}}}{\partial
x^{\kappa}}\right)\\\
\left(\frac{\partial^{2}\bm{f}^{a}_{\bm{\theta}}}{\partial
x^{\bar{\mu}}\partial
x^{\bar{\nu}}}-\Gamma_{\bar{\mu}\bar{\nu}}^{\bar{\kappa}}\frac{\partial\bm{f}^{a}_{\bm{\theta}}}{\partial
x^{\bar{\kappa}}}\right),\end{split}$ (4)

which is the norm of the component of the second derivative of
$\bm{f}_{\bm{\theta}}$ that lies perpendicular to the manifold at that point.
Intuitively, Eq. 4 measures how much the manifold bends in data space.
Importantly, we can directly compute and backpropagate on Eq. 3 and Eq. 4
instead of estimating them via sampling, even if the data live in thousands of
dimensions. The final loss that we minimize then is given by

$\displaystyle\mathcal{L}=\bar{\mathcal{L}}+\frac{1}{M}\sum_{i=1}^{M}(\gamma\mathcal{L}_{PE,i}+\delta\mathcal{L}_{EX,i}),$
(5)

where $i$ indexes points that are uniformly sampled in the latent space based
on the covariance of the embedded points. We call this method $\Gamma$-VAE due
to the role of the Christoffel symbols in computing curvature.

Figure 1: Explicit regularization of the model manifold curvature. (a)
Schematic of tissues (colored points) connected to nearest neighbors (lines)
in the high-dimensional space of gene expression. (b) 2-D UMAP embedding of
the joint TCGA + GTEX datasets, with a sampled grid in the center. (c) PCA of
the sampled grid projected back into gene space. (d) Schematic of a continuous
and differentiable manifold through the tissue samples. (e) Embedding of the
dataset using $\beta$-VAE, where the embedding is colored by the maximum
parameter-effects curvature (copper, left half), and the maximum extrinsic
curvature (grayscale, right), with visibly extreme distortions. (f) PCA of a
decoded grid shows a sharply deformed manifold. (g) Schematic of a manifold
through tissue samples with less curvature. (h) Embedding of the dataset using
$\Gamma$-VAE, demonstrating significantly less parameter-effects and extrinsic
curvature, which can be further seen in (i) the PCA of a decoded portion of
the latent space. (j). Schematic of a PCA through the datapoints, with (k) a
corresponding PCA of the linear embedding and (l) a subset of points.

##  II Geometry of model manifolds

To provide some intuition about the information contained in different low-
dimensional embeddings, we consider the manifolds learned on our joint
TCGA+GTEx data by four different methods: UMAP, $\beta$-VAE, $\Gamma$-VAE, and
PCA. For $n$ data points in $N$-dimensional space, the UMAP algorithm (and
related nearest-neighbor embedding methods) constructs a weighted
dissimilarity graph between each point and its $k$-nearest neighbors (Fig.
1a). Then, it uses a force-directed layout to embed the data in a
$d$-dimensional space (Fig. 1b). We observe that UMAP achieves phenomenal
clustering between different tissue types. To test whether the embedding
contains interpretable geometric information in gene-space, we sample a grid
of points in the embedding, decode the sample back into gene-space, plot the
PCA of the decoded sample (fig. 1c), and find that the manifold is quite
jagged and heterogeneous.

Next, we train a $\beta$-VAE on the same data, which learns a continuous
manifold through the data in gene-space (Fig. 1d). We use the Softplus
activation function to guarantee that the manifold is smooth and
differentiable. The $\beta$-VAE embedding also clearly separates the different
tissue types (Fig. 1e), but when we decode a sampled grid on the manifold, we
find that the manifold geometry is distorted (Fig. 1f). We measure this
distortion using the norms of the largest parameter-effects curvature (Fig.
1e, copper, left) and extrinsic curvature (Fig. 1e, grayscale, right) in any
direction along the manifold. As is expected, $\beta$-VAE significantly
distorts the regular grid of latent-space points on the manifold, and sharply
curves the manifold in orthogonal directions between neighboring tissues (Fig.
1f).

The idea behind $\Gamma$-VAE is to regularize the parameter-effects curvature
(which reduces distortion within the manifold) and the extrinsic curvature
(which reduces the bending of the manifold in gene space) to better capture
long-range covariance in the data (Fig. 1g). The corresponding embedding of
the TCGA+GTEx data shows the points as being distributed much more uniformly
throughout the manifold, with much less parameter-effects and extrinsic
curvature (Fig. 1h). A PCA of a decoded grid sample—containing several of the
same tissues as in Fig. 1e—shows a gently curving manifold that moves roughly
in the same direction between multiple tissue types (Fig. 1i). At the extreme
limit of very high curvature regularization, we get PCA that is a linear
manifold through the data (Fig. 1j) with no extrinsic or parameter-effects
curvature (Fig. 1k), where the tissues are not well separated (Fig. 1l).
Hence, $\Gamma$-VAE provides both regularly-spaced and smoothly varying
manifolds in gene space with good data separation.

Figure 2: A geometric 3 dimensional atlas of human tissue and cancer gene
expression. (a) Embedding of the joint TCGA + GTEX dataset into a 3
dimensional latent space of a highly regularized VAE, with a green plane
spanning liver and muscle, and a red plane spanning the blood to the brain.
(b) Projection of the decoded red and green planes onto a linear PCA of the
data, showing the significant yet regularized curvature of the VAE manifold.
(c) Angles between the tangent space (in gene space) at the origin versus the
tangent space radially away from the origin. (d) Plot of the VAE embedding
colored by the gene signature of adaptive immune response—which defines the
axis from blood to brain—(e) p53 pathway—which defines an axis from bone
marrow to the majority of cancer tissues—(f) and epithelial mesenchymal
transition—which shows complex spatial gradients in the cancer clusters. (g)
Plot of a subset of nine healthy tissues with arrows drawn to their
corresponding adenocarcinomas, forming a distinct and uniform cancer axis. The
embedding also shows a distinct orthogonal separation of squamous cell
carcinomas. (h) Decoded gene trajectories from healthy colon to colon
adenocarcinoma for select genes, showing curved, nonlinear pahtways through
gene space.

##  III A curved 3-dimensional atlas of human gene expression

With the intuitions from the previous section, we now construct an
interpretable, curved 3-dimensional atlas of human gene expression
(TCGA+GTEx). We train a $\Gamma$-VAE with a 3-dimensional latent space, and
large curvature regularization (Fig. 2a). We first observe two distinct axes
of biological function: one between the liver and muscle (along the green
plane), and the second between blood and brain (along the red plane). By
“interpretable,” we mean that the geometry of this embedding is only mildly
distorted in gene space. To demonstrate this distortion, we decode the grid of
points in the red and green planes into gene space, and project the decoded
planes onto a 3-dimensional PCA of the data (Fig. 2b), showing a gentle
curving of the planes. To further quantify this distortion, we compute the
three angles formed by the 3-dimensional tangent spaces of the manifolds
between the origin of the embedding and at randomly sampled points (Fig. 2c).
The manifolds do not become orthogonal in any direction until the edges of the
latent embedding, demonstrating long-range coherence.

We find that these axes are not only consistent, but also biologically
meaningful. We color each tissue in the embedding by the gene signature—the
sum of the expression for a subset of genes—associated with adaptive immune
response, and find a striking gradient from the blood to the brain (Fig. 2d).
A gene set enrichment analysis (GSEA [30]) of the tangent vector along this
axis from the origin is enriched for processes such as inflammatory response,
allograft rejection, and TNFA signaling via NFKB. We find a second global axis
from healthy tissues towards the squamous cell carcinomas with a monotonic
gradient in the p53 pathway gene signature (Fig. 2e). The tangent vector along
this axis is further enriched in early and late estrogen response, epithelial-
mesenchymal transition (EMT), and KRAS signaling. Beyond monotonic gradients,
our method also organizes points along spatially heterogeneous patterns, where
we find both an EMT axis from healthy tissues to the carcinomas (squamous cell
and adenocarcinomas), as well as a gradient for each branch of the carcinomas
(Fig. 2f).

Figure 3: Out-of-distribution generalizability of reguarlized embeddings to
unseen cancers. (a) A $\Gamma$-VAE embedding of the joint TCGA + GTEX dataset,
with a subset of tissues (breast carcinoma, BRCA) colored in a copper gradient
based on their distance to healthy breast tissue. (b) a zoomed-in plot of the
joint breast and BRCA tissues in the VAE space, where the triple-negative
breast carcinomas (TNBC) are plotted as tetrahedra. (c) A regularized
embedding trained while leaving out all BRCA and BRCA juxta-tumor samples,
with the BRCA tissues re-embedded after training. (d) A zoomed-in plot of the
joint breast and re-embedded BRCA tissues that successfully separates TNBC
from non-TNBC tissues, where the BRCA tissues are colored the same as in panel
(b). (e) Pairwise-distance correlations for UMAP, unregularized VAE, and
regularized VAE of each cancer tissue and all other tissues between embeddings
where the cancer tissue was included or excluded from the training. (f)
Column-normalized density plots of all pairwise distances for each cancer
tissue and all other tissues between the original embedding including the
tissue, and the out-of-sample embedding that excluded the cancer tissue from
training.

To better understand the organization of healthy versus cancer tissues, we
look at a subset of carcinomas and their corresponding healthy tissues (Fig.
2g). We find that the atlas clearly and geometrically separates the squamous
cell carcinomas (tetrahedra) from the adenocarcinomas (cubes) and healthy
tissues (spheres). Further, we report a consistent geometric vector from 9
healthy tisues (breast, colon, esophagus, lung, ovary, prostate, stomach,
uterus, and thyroid) to their corresponding adenocarcinomas. Our atlas also
separates squamous cell carcinomas from adenocarcinomas for cancers
originating from the same tissue (e.g. lung squamous cell versus lung
adenocarcinoma), and that for TCGA categories that contain a mixture of both
(e.g. esophageal carcinoma), our atlas correctly embeds the two types as
determined by the histology.

Finally, due to VAEs learning an explicit geometric manifold in gene-space, we
can study the trajectory of gene expression along axes. In Fig. 2h, we plot
the trajectory of gene expression from healthy colon tissues to colon
adenocarcinoma for a subset of genes. Notably, because straight lines in our
latent space correspond to gently curved lines in gene space, these gene
trajectories have the potential to yield staging information during the
progression of tumorogenesis. We note that in this dataset, there is no
concrete evidence that these specific trajectories correspond to the process
of tumorogenesis. Rather we demonstrate the potential of our method for
reconstructing the nonlinear sequence of gene expression [31].

##  IV Out-of-distribution prediction of cancer phenotypes

While part of the advantage of constructing geometric manifolds in data space
is improved and interpretable visualizations, the manifolds themselves are
also a model of the data. We explored the model capabilities by studying the
enrichment of tangent vectors (Fig. 2d-f) and trajectories (Fig. 2h) of the
model manifold. Here, we demonstrate the power of the model to accurately
predict out of distribution.

Figure 4: Cell fate prediction using curved embeddings. (a) $\Gamma$-VAE
embedding of a lineage tracing experiment on hematopoietic stem cell
differentiation. (b) Subset of multipotent LSK cells from a PCA of the data,
yielding a 49% classification accuracy of the eventual fate of their daughter
cells. (c) The same subset of cells from the cf-VAE of the data, yielding a
60% classification accuracy of the eventual fate of their daughter cells. (d)
re-embedding of all data on the $\Gamma$-VAE manifold trained only on day 2
cells. (e) Column-normalized density plot of pairwise distances between day 4
and 6 cells in the embedding, with the pairwise distances of a $\Gamma$-VAE
trained on all data on the x-axis, and the pairwise distances of a
$\Gamma$-VAE trained on only day 2 data on the y-axis. (f) Re-embedding of day
4, 6 cells on the PCA projections defined on only the day 2 cells.

As a concrete example, we train a $\Gamma$-VAE on all of the TCGA + GTEx data
at a lower regularization than in Fig. 2, which clearly separates most tissues
(Fig. 3a), and color the samples from breast carcinoma (BRCA) in a copper
gradient based on their distance from healthy breast tisuses. We choose BRCA
because there is significant known within-cancer heterogeneity within BRCA,
i.e. triple negative (or basal, TNBC) versus HER2+ and hormone receptor
positive (or luminal) subtypes [32]. Zooming into the regions of the embedding
near BRCA (Fig. 3b), we see a clear and expected separation between the non-
TNBC BRCA samples (copper spheres) and the TNBC BRCA samples (copper
tetrahedra). When we train another $\Gamma$-VAE using a different random
number seed after removing all BRCA samples—along with their corresponding
juxta-tumor tissues—and re-embed the BRCA samples after training, we find that
the geometry of the embedding of all tissues remain relatively fixed, and that
the BRCA samples re-embed in the same global location (Fig. 3c). Zooming into
the same region surrounding the re-embedded BRCA samples, we find that the
relative distance from the healthy breast tissues (the copper color gradient)
is preserved, as is the embedding positions of the non-TNBC and TNBC samples.

We observed remarkable re-embedding consistency with the out-of-distribution
re-embedding of the BRCA samples. To demonstrate the utility of the method
more generally, we perform the same procedure of removing each cancer type
(along with its corresponding juxta-tumor samples), training a separate
$\Gamma$-VAE on the reduced data, and re-embedding the cancer tissues. We
quantify re-embedding consistency by first computing the pairwise distance
between the out-of-distribution tissues and all other tissues for a
$\Gamma$-VAE trained on all of the data, and the pairwise distance for a
$\Gamma$-VAE trained on the reduced data (initialized with a different random
number seed). We then compute the Spearman rank correlation coefficient
between these two sets of pairwise distances for each of the 33 cancer types
in TCGA. We find the re-embedding consistency to be significantly higher for
$\Gamma$-VAE ($\rho=0.9$) than for either $\beta$-VAE ($\rho=0.67$) or UMAP
($\rho=0.58$) using a non-parmetric Wilcoxon signed rank test ($p<10^{-6}$ for
both) (Fig. 3e). Column-normalized density plots of the pairwise distances
across all tissues show the striking ability of $\Gamma$-VAE to preserve local
and global embedding distances on out-of-distribution data (Fig. 3f).

##  V Uncovering and predicting cell fates on curved manifolds

To demonstrate the utility of $\Gamma$-VAE in another biological context and
to extend to single cell RNA sequencing (scRNAseq) data, we look at a lineage
tracing experiment for hematopoietic stem cell differentiation [26]. In this
work, the authors uniquely barcoded a population of hematopoietic stem cells
on day 0, collected and sequenced subsets of the population on days 2, 4, and
6, and used the barcodes to trace their differentiation trajectories.
Embedding this data using $\Gamma$-VAE, we find that the differentiated cell
types are well separated (Fig. 4a).

The first advantage of embedding on a curved manifold is the ability to
separate out the eventual fates of uncommitted cells. A subset of cells
collected on day 2 were multipotent LSK cells whose future fates were obtained
from their barcode-labelled daughter cells on days 4 and 6. A plot of this LSK
subset from a PCA of the data shows poor separation between LSK cells of
different eventual fates (Fig. 4b), with a 49% classification accuracy using a
discriminant classifier on the 3 PCA dimensions. Plotting the same LSK cells
in our $\Gamma$-VAE shows significantly better separation between their
eventual fates, with a 60% classification accuracy using a discriminant
classifier on the 3 $\Gamma$-VAE dimensions: the same accuracy reported by the
original authors using a deep neural network classifier on 447 differentially
expressed genes [26].

The second advantage of embedding on a curved manifold is the ability to
accurately re-embed out of distribution. We train a second $\Gamma$-VAE on
only day 2 cells, and re-embed the day 4 and day 6 cells (Fig. 4d). Visually,
we notice a striking similarity with the embedding trained on all data (Fig.
4a). We quantify this similarity using the pairwise distances between the day
4 and 6 embedded points, and find a strong correlation in the pairwise
distances between the embedding trained on all data, and the embedding trained
on only day 2 data (Fig. 4e) with a Spearman rank correlation coefficient of
0.92. By comparison, the equivalent PCA re-embedding of day 4 and 6 cells on a
PCA of only day 2 cells yields significantly worse separation between cell
types, and in particular between the differentiated and undifferentiated cells
(Fig. 4f).

##  VI Discussion

The search for interpretable, low-dimensional models has spanned an incredible
breadth of time [4, 33] subject areas [34, 35], methodologies [11], and
mathematical objects such as graphs [36], mathematical models [1], and
probability distributions [3, 2]. One prevailing paradigm is to search for an
appropriate kernel similarity function that serves to untangle the majority of
the variance of the data along a few interpretable axes [11] and is the basis
for many manifold learning methods today [5, 6, 7, 8, 9]. In a very real and
practical sense, this is also the prevailing mindset in deep learning, where
learning in deep neural networks (DNNs) at a theoretical limit can be thought
of as evolution according a kernel function that is determined by the network
architecture: the Neural Tangent Kernel (NTK) [21, 37]. The hope is that, at
the limit of the ever-increasing quantities of experimental data and
conditions [38, 39] alongside the correct sophisticated DNN architectures
[40], we can obtain powerful universal representations across a wide variety
of data [41].

In this work, we provide a complementary approach based on a simple idea: that
the geometry of the data in natural coordinates is meaningful and worth
preserving. This idea is particularly fruitful in heterogeneous atlas data
with highly dissimilar clusters where the data topology do not reflect the
geometry [42, 43], and where statistical modeling methods fail to correctly
capture global and meso-scale structure that span between clusters. We find
that regularizing the manifold curvature leads to more uniformly distributed
embeddings where the genes co-vary across multiple tissue types (Fig. 1). We
also find that regularizing the manifold curvature reveals both global and
mesoscale axes of biological function, clear separation of distinct cancer
phenotypes, and homogeneous curved gene expression pathways from healthy to
cancerous tissues (Fig. 2). Crucially, our method can accurately predict
expression phenotypes not only out of sample, but out of distribution, with
far superior consistency than UMAP or $\beta$-VAE (Fig. 3). Finally, our
method can predict future cell fates both by curving the projection of
multipotent cells towards their future fate with state-of-the-art
classification accuracy in just 3 dimensions, and accurately predict future
cell expression from only day 2 data (Fig. 4). In sum, many biological and
emergent processes contain low-dimensional and geometric organization in high-
dimensional spaces, and by carefully preserving that geometry, our method
constructs more interpretable and predictive low-dimensional models.

##  VII Acknowledgements

We gratefully acknowledge conversations and feedback from Alexandre Lanau,
Phillipe Martin, Jason Cosgrove, Dr. Melody Lim, and Dr. Erich Mueller. JZK
was supported by the Bethe/KIC/Wilkins post-doctoral fellowship. This work
received financial support from ITMO Cancer of Aviesan within the framework of
the 2021-2030 Cancer Control Strategy, on funds administered by Inserm.

##  VIII References

## References

  * Transtrum and Qiu [2014] M. K. Transtrum and P. Qiu, Model reduction by manifold boundaries, Physical review letters 113, 098701 (2014).
  * Teoh  _et al._ [2020] H. K. Teoh, K. N. Quinn, J. Kent-Dobias, C. B. Clement, Q. Xu, and J. P. Sethna, Visualizing probabilistic models in minkowski space with intensive symmetrized kullback-leibler embedding, Physical Review Research 2, 033221 (2020).
  * Quinn  _et al._ [2019] K. N. Quinn, C. B. Clement, F. De Bernardis, M. D. Niemack, and J. P. Sethna, Visualizing probabilistic models and data with intensive principal component analysis, Proceedings of the National Academy of Sciences 116, 13762 (2019).
  * Pearson [1901] K. Pearson, Liii. on lines and planes of closest fit to systems of points in space, The London, Edinburgh, and Dublin philosophical magazine and journal of science 2, 559 (1901).
  * McInnes  _et al._ [2018] L. McInnes, J. Healy, and J. Melville, Umap: Uniform manifold approximation and projection for dimension reduction, arXiv preprint arXiv:1802.03426 (2018).
  * Becht  _et al._ [2019] E. Becht, L. McInnes, J. Healy, C.-A. Dutertre, I. W. Kwok, L. G. Ng, F. Ginhoux, and E. W. Newell, Dimensionality reduction for visualizing single-cell data using umap, Nature biotechnology 37, 38 (2019).
  * Van der Maaten and Hinton [2008] L. Van der Maaten and G. Hinton, Visualizing data using t-sne., Journal of machine learning research 9 (2008).
  * Coifman and Lafon [2006] R. R. Coifman and S. Lafon, Diffusion maps, Applied and computational harmonic analysis 21, 5 (2006).
  * Belkin and Niyogi [2003] M. Belkin and P. Niyogi, Laplacian eigenmaps for dimensionality reduction and data representation, Neural computation 15, 1373 (2003).
  * Moon  _et al._ [2017] K. R. Moon, D. van Dijk, Z. Wang, W. Chen, M. J. Hirn, R. R. Coifman, N. B. Ivanova, G. Wolf, and S. Krishnaswamy, Phate: a dimensionality reduction method for visualizing trajectory structures in high-dimensional biological data, BioRxiv 120378 (2017).
  * Schölkopf  _et al._ [1997] B. Schölkopf, A. Smola, and K.-R. Müller, Kernel principal component analysis, in  _International conference on artificial neural networks_ (Springer, 1997) pp. 583–588.
  * Williams [2000] C. Williams, On a connection between kernel pca and metric multidimensional scaling, Advances in neural information processing systems 13 (2000).
  * Lopez  _et al._ [2018] R. Lopez, J. Regier, M. B. Cole, M. I. Jordan, and N. Yosef, Deep generative modeling for single-cell transcriptomics, Nature methods 15, 1053 (2018).
  * Kingma  _et al._ [2019] D. P. Kingma, M. Welling, _et al._ , An introduction to variational autoencoders, Foundations and Trends® in Machine Learning 12, 307 (2019).
  * Kobak and Linderman [2021] D. Kobak and G. C. Linderman, Initialization is critical for preserving global data structure in both t-sne and umap, Nature biotechnology 39, 156 (2021).
  * Arvanitidis  _et al._ [2017] G. Arvanitidis, L. K. Hansen, and S. Hauberg, Latent space oddity: on the curvature of deep generative models, arXiv preprint arXiv:1710.11379 (2017).
  * Acosta  _et al._ [2023] F. Acosta, S. Sanborn, K. D. Duc, M. Madhav, and N. Miolane, Quantifying extrinsic curvature in neural manifolds, in  _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ (2023) pp. 610–619.
  * Davidson  _et al._ [2018] T. R. Davidson, L. Falorsi, N. De Cao, T. Kipf, and J. M. Tomczak, Hyperspherical variational auto-encoders, arXiv preprint arXiv:1804.00891 (2018).
  * Bose  _et al._ [2020] J. Bose, A. Smofsky, R. Liao, P. Panangaden, and W. Hamilton, Latent variable modelling with hyperbolic normalizing flows, in  _International Conference on Machine Learning_ (PMLR, 2020) pp. 1045–1055.
  * Lee and Park [2023] Y. Lee and F. C. Park, On explicit curvature regularization in deep generative models, in  _Topological, Algebraic and Geometric Learning Workshops 2023_ (PMLR, 2023) pp. 505–518.
  * Jacot  _et al._ [2018] A. Jacot, F. Gabriel, and C. Hongler, Neural tangent kernel: Convergence and generalization in neural networks, Advances in neural information processing systems 31 (2018).
  * Transtrum  _et al._ [2011] M. K. Transtrum, B. B. Machta, and J. P. Sethna, Geometry of nonlinear least squares with applications to sloppy models and optimization, Physical Review E 83, 036701 (2011).
  * Weinstein  _et al._ [2013] J. N. Weinstein, E. A. Collisson, G. B. Mills, K. R. Shaw, B. A. Ozenberger, K. Ellrott, I. Shmulevich, C. Sander, and J. M. Stuart, The cancer genome atlas pan-cancer analysis project, Nature genetics 45, 1113 (2013).
  * Lonsdale  _et al._ [2013] J. Lonsdale, J. Thomas, M. Salvatore, R. Phillips, E. Lo, S. Shad, R. Hasz, G. Walters, F. Garcia, N. Young, _et al._ , The genotype-tissue expression (gtex) project, Nature genetics 45, 580 (2013).
  * Wilks  _et al._ [2021] C. Wilks, S. C. Zheng, F. Y. Chen, R. Charles, B. Solomon, J. P. Ling, E. L. Imada, D. Zhang, L. Joseph, J. T. Leek, _et al._ , recount3: summaries and queries for large-scale rna-seq expression and splicing, Genome biology 22, 1 (2021).
  * Weinreb  _et al._ [2020] C. Weinreb, A. Rodriguez-Fraticelli, F. D. Camargo, and A. M. Klein, Lineage tracing on transcriptional landscapes links state to fate during differentiation, Science 367, eaaw3381 (2020).
  * Higgins  _et al._ [2016] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner, beta-vae: Learning basic visual concepts with a constrained variational framework, in  _International conference on learning representations_ (2016).
  * Fil  _et al._ [2021] M. Fil, M. Mesinovic, M. Morris, and J. Wildberger, beta-vae reproducibility: Challenges and extensions, arXiv preprint arXiv:2112.14278 (2021).
  * Locatello  _et al._ [2020] F. Locatello, S. Bauer, M. Lucic, G. Rätsch, S. Gelly, B. Schölkopf, and O. Bachem, A sober look at the unsupervised learning of disentangled representations and their evaluation, The Journal of Machine Learning Research 21, 8629 (2020).
  * Subramanian  _et al._ [2005] A. Subramanian, P. Tamayo, V. K. Mootha, S. Mukherjee, B. L. Ebert, M. A. Gillette, A. Paulovich, S. L. Pomeroy, T. R. Golub, E. S. Lander, _et al._ , Gene set enrichment analysis: a knowledge-based approach for interpreting genome-wide expression profiles, Proceedings of the National Academy of Sciences 102, 15545 (2005).
  * Fearon and Vogelstein [1990] E. R. Fearon and B. Vogelstein, A genetic model for colorectal tumorigenesis, cell 61, 759 (1990).
  * Chiu  _et al._ [2018] A. M. Chiu, M. Mitra, L. Boymoushakian, and H. A. Coller, Integrative analysis of the inter-tumoral heterogeneity of triple-negative breast cancer, Scientific reports 8, 11807 (2018).
  * Hastie and Stuetzle [1989] T. Hastie and W. Stuetzle, Principal curves, Journal of the American Statistical Association 84, 502 (1989).
  * Cunningham and Yu [2014] J. P. Cunningham and B. M. Yu, Dimensionality reduction for large-scale neural recordings, Nature neuroscience 17, 1500 (2014).
  * Bruce  _et al._ [2002] L. M. Bruce, C. H. Koger, and J. Li, Dimensionality reduction of hyperspectral data using discrete wavelet transform feature extraction, IEEE Transactions on geoscience and remote sensing 40, 2331 (2002).
  * Hamilton  _et al._ [2017] W. L. Hamilton, R. Ying, and J. Leskovec, Representation learning on graphs: Methods and applications, arXiv preprint arXiv:1709.05584 (2017).
  * Alemohammad  _et al._ [2020] S. Alemohammad, Z. Wang, R. Balestriero, and R. Baraniuk, The recurrent neural tangent kernel, arXiv preprint arXiv:2006.10246 (2020).
  * Replogle  _et al._ [2022] J. M. Replogle, R. A. Saunders, A. N. Pogson, J. A. Hussmann, A. Lenail, A. Guna, L. Mascibroda, E. J. Wagner, K. Adelman, G. Lithwick-Yanai, _et al._ , Mapping information-rich genotype-phenotype landscapes with genome-scale perturb-seq, Cell 185, 2559 (2022).
  * Ståhl  _et al._ [2016] P. L. Ståhl, F. Salmén, S. Vickovic, A. Lundmark, J. F. Navarro, J. Magnusson, S. Giacomello, M. Asp, J. O. Westholm, M. Huss, _et al._ , Visualization and analysis of gene expression in tissue sections by spatial transcriptomics, Science 353, 78 (2016).
  * Vaswani  _et al._ [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, Attention is all you need, Advances in neural information processing systems 30 (2017).
  * Rosen  _et al._ [2023] Y. Rosen, Y. Roohani, A. Agrawal, L. Samotorcan, T. S. Consortium, S. R. Quake, and J. Leskovec, Universal cell embeddings: A foundation model for cell biology, bioRxiv , 2023 (2023).
  * Lu  _et al._ [2023] T.-C. Lu, M. Brbić, Y.-J. Park, T. Jackson, J. Chen, S. S. Kolluru, Y. Qi, N. S. Katheder, X. T. Cai, S. Lee, _et al._ , Aging fly cell atlas identifies exhaustive aging features at cellular resolution, Science 380, eadg0934 (2023).
  * Comitani  _et al._ [2023] F. Comitani, J. O. Nash, S. Cohen-Gogo, A. I. Chang, T. T. Wen, A. Maheshwari, B. Goyal, E. S. Tio, K. Tabatabaei, C. Mayoh, _et al._ , Diagnostic classification of childhood cancer using multiscale transcriptomics, Nature Medicine 29, 656 (2023).