Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Authors: Alessandro Palma
Sergei Rybakov
Leon Hetzel
Stephan Günnemann
Fabian J Theis
Venue: ICML 2025 spotlightposter (2025)
Keywords: scRNA-seq
Riemannian geometry
representation learning
trajectory inference
VAEs
statistical manifolds
Primary area: applications->health_medicine
TL;DR: We enforce a local Euclidean geometry in the latent manifold of scRNA-seq VAEs to improve downstream applications.

Abstract: Latent space interpolations are a powerful tool for navigating deep generative models in applied settings. An example is single-cell RNA sequencing, where existing methods model cellular state transitions as latent space interpolations with variational autoencoders, often assuming linear shifts and Euclidean geometry. However, unless explicitly enforced, linear interpolations in the latent space may not correspond to geodesic paths on the data manifold, limiting methods that assume Euclidean geometry in the data representations. We introduce FlatVI, a novel training framework that regularises the latent manifold of discrete-likelihood variational autoencoders towards Euclidean geometry, specifically tailored for modelling single-cell count data. By encouraging straight lines in the latent space to approximate geodesic interpolations on the decoded single-cell manifold, FlatVI enhances compatibility with downstream approaches that assume Euclidean latent geometry. Experiments on synthetic data support the theoretical soundness of our approach, while applications to time-resolved single-cell RNA sequencing data demonstrate improved trajectory reconstruction and manifold interpolation.

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Alessandro Palma" *
Sergei Rybakov 123
Leon Hetzel * 124 Stephan Günnemann' 24
Fabian J. Theis 2
Abstract
2023; Lopez et al., 2018; Griffiths & Hernández-Lobato,
2020). One reason is that real-world experimental data often poses significant modelling challenges as is inherently noisy, high-dimensional, and complex (Sarker, 2021). As a solution, learning compressed and dense latent representation of the data has gained traction in applied machine learning- For example, interpolations in well-behaved, low-dimensional data embeddings useful for modelling the dynamics of complex systems, allowing insights into sample evolution over time (Džeroski & Todorovski, 2003; Bunne et al., 2022).
Latent space interpolations are a powerful tool for navigating deep generative models in applied settings. An example is single-cell RNA sequenc- ing, where existing methods model cellular state transitions as latent space interpolations with vari- ational autoencoders, often assuming linear shifts and Euclidean geometry. However, unless explic- itly enforced, linear interpolations in the latent space may not correspond to geodesic paths on the data manifold, limiting methods that assume Euclidean geometry in the data representations. We introduce FlatVI, novel training framework that regularises the latent manifold of discrete- likelihood variational autoencoders towards Euclidean geometry, specifically tailored for modelling single-cell count data. By encouraging straight lines in the latent space to approximate geodesic interpolations the decoded single-cell manifold, FlatVI enhances compatibility with downstream approaches that assume Euclidean latent geometry. Experiments on synthetic data support the theoretical soundness of our approuch, while applications to time-resolved single-cell RNA sequencing data demonstrate improved tra- jectory reconstruction and manifold interpolation.
In particular, VAEs have shown great promise repre- senting both continuous and discrete data, as the decoder parameterises flexible likelihood model. This flexibil- ity has demonstrated unprecedented potential in cellular data (Lopez et al., 2018 8), particularly in gene expression, which is measured in counts that reflect the number of RNA molecules produced by individual genes and collected through single-cell RNA sequencing (scRNA-seq) (Haque et al., 2017). Such technique allows the measurement of thousands of genes in parallel, and the resulting vector describes the state of cell across diverse biological settings (Regev et al., 2017). Leveraging latent interpolations within single-cell VAEs can provide insights into cellular state tran- sitions, capturing dynamic changes in gene expression that reveal underlying biological processes.
The representation learnt by VABs is tightly connected the latent space parametrisation of low-dimensional manifold (Arvanitidis et al., 2021). When modelling latent cellular dynamics on single-cell data, popular approaches still rely on assuming Fuclidean geometry in the representation space, for example by modelling cellular transitions through linear-cost Opti- mal Transport (OT) (Peyré et al., 2019; Klein al., 2025; Tong al., 2020). However, building linear latent trajec- tories using the Euclidean assumption is sub-optimal when the duta lies on non-Euclidean manifold, as straight la- tent lines do not necessarily reflect geodesic paths on the manifold induced by the decoder.
1. Introduction
Riemannian geometry, one
Generative models for representation learning, such Variational Autoencoders (VARs), have influenced compu- tational sciences across multiple fields (Zhong& & Meidani,
"Equal contribution Inssittte of Computational Biology, Helmholtz Munich, Munich, Germany "School of Computation Information and Technology, Technical University of Munich, Germany "Lamin Labs MMunich Data Science Institute, Techni- cal University of Munich, Germany 'TUM School of Life Sci- ences Weihenstephan, Technical University of Munich, Germany. Correspondence to: Fabian Theis <ffaiinn theis @ helmholtz- munich.de>.
learn effective interpolations on a single-cell manifold, we establish the following desiderata: (i) Approximate trajectories intractable data manifolds via interpolations
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Latent space
Parameter space
Linear interpolation
VAE
Pullback metric
FlatVl
FlatVI
minimise
M(z)
En, (slxe) (M(x) alal
(H, O)
Z
Decoder h(z)
Decoded
parameters
Figure 1. Visual conceptualisation of the FlatVI approach. The decoder of a single-cell VAE maps to the parameter space of negative binomial statistical manifold probability distributions. In standard VAE settings, straight latent paths are not guaranteed to map to meaningful statistical manifold interpolations through the decoder. By regularising the pullback metric of the stochastic decoder, Flat induces correspondence between straight paths the latent space and geodesic interpolations along the manifold the decoded space.
on a simpler latent manifold with a tractable geometry. (ii) Design decoding scheme that encourages straight paths in the latent manifold to map to approximate geodesics in the decoded space. (iii) Formalise the geodesic matching framework in way that supports a flexible choice of the de- coder's likelihood. To achieve (i) and (ii), existing methods regularise the latent representation of Gaussian AEs using Euclidean geometry (Chen et al., 2020; Yonghyeon et al., 2021), but limit their application to continuous data by neglecting the decoder's general likelihood model support. Other works explore the connection between stochastic de- coders' geometry and the latent space manifold (Arvanitidis et al., 2022), even for discrete data, but do not address regu- larising the latent manifold to simple, traversable geometry while preserving geodesic paths on the decoded manifold.
Euclidean geometry, while enabling likelihood parameter reconstruction on par with standard VAEs. Our method finds direct applications to single-cell representation learn- ing and trajectory inference, which we demonstrate across multiple biological settings by providing an improved data representation for OT-based modelling of cellular popula- tion dynamics and latent interpolation. In summary, we make the following contributions:
We introduce a regularisation technique for discrete- likelihood VAEs encourage straight latent interpola- tions to approximate geodesic paths on the statistical manifold induced by the decoder.
We provide an explicit formulation of the flattening loss for the negative binomial case, which directly impacts modelling -igh-dimessional scRNA-seq data.
In this work, we close this gap and introduce FlatVI-Flaa Variational Inference-a theoretically principled approach pushing straight paths in the latent space of VAEs to approximate geodesic paths along the manifold of the decoded data. Our focus is on statistical manifolds, whose points are probability distributions of pre-defined family. This enables us to draw connections to the theory of VAËS and information geometry. When trained as a likelihood model, a VAE's decoder image maps latent codes to a statistical manifold's parameter space. This formulation finds direct application in scRNA-seq, where individual gene counts are assumed to follow negative binomial distribution, reflecting relevant data properties such as discreteness and overdispersion (Zhou et al., 2011).
We empirically validate our model on simulated data and latent geodesic interpolations.
We show that our method offers a better representation space for existing OT-based trajectory inference tools than existing VAE-based approaches on real data.
2. Related Work
Geometry and AEs. Prior work by Arvanitidis et al. (2021) introduced optimal latent paths reflecting observation space geometry in deterministic and Gaussian stochastic decoders, extended by Arvanitidis et al. (2022) to VAEs with arbi- trary likelihoods. Meanwhile, Chen al. (2020) explored representation learning benefits by modelling latent spaces of deterministic AEs as flat manifolds, while other studies incorporate data geometry via isometric (Yonghyeon et al., 2021) and Jacobian (Nazari et al., 2023) regularisations.
Crucially, FI regularises the latent space through a flat- tening loss that pushes the pullback metric from a stochastic VAE decoder towards a spatially-unifome scaled identity matrix, thereby regularising towards a locally Fuclidean latent geometry (Figure 1). In controlled simulation set- ting, we demonstrate that our regularisation successfully constrains the latent manifold to exhibit an approximate
Geometry in single-cell representations. Latent variable models for ScRNA-seq data are established (Lopez et al., 2018: Eraslan et al., 2019), and geometric regularisations for continuous approximation of high-dimensional cellular
2

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
data have been proposed before for deterministic AEs (Duque al., 2020; Sun al., 2025). Combining single-cell representations and geometry, diffusion-based manifold learning (Moon et al., 2019; Huguet et al., 2024; Fasina et al., 2023) offers insights into geometry-aware low-dimensional representations. Investigating single-cell geometry extends to gene expression data (Korem et al., 2015; Qiu et al., 2022) and dynamic settings (Rifkin & Kim, 2002).
termed decoder. VAEs additionally include an encoder network fe X' - Z optimized jointly with ha through the Evidence Lower Bound Objective (ELBO) (Kingma & Welling, 2014). Overall, as long as one can select a parametric family of distributions as a reasonable noise model for the dataset properties, the likelihood of the data can be modelled by the decoder of a VAE.
In the field of scRNA-seq, biological and technical vari- ability causes sparsity and overdispersion properties in the expression counts, making the negative binomial likelihood a natural choice for modelling gene expression. Sparsity arises from technical limitations in detecting gene transcripts or from unexpressed genes in specific conditions. Overdis- persion refers to genes having higher variance than the mean, deviating from a Poisson model. This is influenced by tech- nical factors and modelled by the inverse dispersion pa- rameter of a negative binomial distribution (Heumos et al., 2023). Thus, we assume that genes follow a negative bino- mial noise model NB(Ung, 0g), where Ung and 0, represent the del--gee--sseiiee mean and the gene-specific inverse dispersion parameters, respectively. In the VAE setting, given gene-expression vector x, we define the following parameterizations (Lopez al., 2018):
Modelling single-cell state transitions in oiw-einesiional spaces. Reconstructing cellular state transitions in a biolog- ically meaningful low-dimensional space is a key challenge in single-cell transcriptomics. Various approaches address this, including applications to drug perturbation prediction (Bunne et al., 2023; Lotfollahi et al., 2023; Hetzel et al., 2022) and multi-modal trajectory inference (Klein et al., 2025). Several works focus on learning continuous gene pression trajectories in latent spaces for modelling cellular dynamics. Our work is tightly linked to Huguet et al. (2022), where the authors employ Geodesic Autoencoder (GAE) where distances in the latent space approximate geodesic distances in single-cell data, while Haviv al. (2024) intro- duce a framework regularising autoencoders' latent spaces to approximate Wasserstein distances, with applications in spatial transcriptomics. Flow Matching (Lipman et al., 2023; Albergo & Vanden-Eijnden, 2023), generative mod- elling approach, has also been explored for reconstruct- ing manifold-an -aware cellular trajectories in low-dimensional spaces (Kapusniak et al., 2024), with OT-based formulations showing promise in this context (Tong et al., 2024a).
z - (x)), p h (z. I) softmax(po(z)). (2)
where Ps : IR. - R models expression proportions of individual genes and I is the observed cell-specific size factor directly derived from the data as cell's total number of counts ( = -= Zg.. The encoder f already takes into account the reparametrisation trick (Kingma & Welling, 2014). Assuming global, gene-specific technical effects, Og is treated as free parameter independent of the cell's 's state.
3. Background
3.1. Discrete VAEs for Single-Cell RNA-seq
In this work, we deal with discrete count data, formally collected in a high-dimensional matrix X E Ng NNG where N represents the number of observations and G the number of features. We assume that individual sample features Wng are independent realisations of a discrete random variable Xng (P(- (ggg) with observatioonsppecfic real parameters Png. Letx X Nº be single realisation vector.
3.2. The Geometry of Autoencoders
Continuous deterministic AEs. A possible assumption is that continuous data lies near low-dimensional Riemannian manifold MMx with X = R², associated with a d-dimensional latent space Z. The decoder h of a deterministic AE model can be seen as an immersion h Rd + R of the latent space Z into the embedded Riemannian manifold Mx equipped with a metric tensor M and defined as follows:
We consider a joint latent variable model describing the probability of an observation x and its associated latent vari- able z. The model factorizes as (x, z) = P& (x)z) where z E 2 Rdd is a d-dimensional latent variable with d < G. ~ p(z), and p(z) = .V Ld). Here, La is the squared identity matrix with dimension d.
Definition 3.1. A Riemannian manifold is a smooth man- ifold Mx endowed with a Riemannian metric M(x) for x E Mxx M(x) is a positive-definite matrix that changes smoothly and defines a local inner product on the tangent space Mx as (u, M) u²M(x)vv, with v,u TxMx (Do Carmo & Flaherty Francis, 1992).
The factor Ps defines likelihood model following a dis- crete distribution from a pre-defined family, with parameters expressed as a function of the latent variable z as:
From Definition 3.1, it follows that a Riemannian manifold in the decoded space has Euclidean geometry if M(() = Ia everywhere, as the dot product between tangent vectors reduces to linear product.
(x|z) = P(x(ha(z))
(1)
In deep latent variable models, he is deep neural network
3

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
In this setting, one can define a Riemannian manifold in the latent space, called Mz, whose geometry is directly linked to the geometry of the decoded manifold by the pullback metric M(z):
where M(x) E Rdxd, Note that the calculation of the FIM is specific for the chosen likelihood type and depends on initial assumptions on the data distribution (see in Appendix C.2),
((z) =
(3)
4. The FI Model
where Jn (2) is the Jacobian matrix of h(z). Here, the decoder h is assumed to be a diffeomorphism between the latent space and its image, such that JA (z) is full rank for all z. The relationship between latent and decoded metric tensors in Equation (3) enables us to connect distances on the latent manifold with quantities measured in the data space. For example, Equation (3) allows to define the shortest curve (t) connecting pairs of latent codes z and 22 as the one minimising the distance between their images )
4.1. Latent Euclidean Assumption in Single-Cell Biology
Modelling high-dimensional cellular processes from discrete count data poses significant challenges. common approach is to study variations in cell states as latent interpo- lations using the continuous latent representation learned by negative binomial VAEs. In applications like perturbation modelling (Hetzel et al., 2022; Lotfollahi et al., 2023) or gene expression distance quantification (Luecken et al., 2021), it is a common assumption to model state transitions as linear shifts in the latent space. This also applies to trajec- tory inference for modelling continuous population dynam- ics, where the trajectory of single cells is learnt by matching subsequent cellular snapshots collected over time using dy- namic OT in Euclidean spaces (Tong et al., 2020; Koshizuka & Sato, 2023; Tong et al., 2024b; Neklyudov et al., 2023). In all the above cases, the standard assumption is that linear interpolations of the latent manifold reflect optimal trajec- tories on the decoded single-cell manifold. However, the standard negative binomial VAE formulation does not nat- urally enforce such correspondence, violating modelling assumptions and potentially leading to sub-optimal decoded cell-state trajectories and latent distance estimations.
(Z)) and l(2)) on Mx. More formally:
( [I(()(t)]ldtt
diatem (2) inf
(4)
>(1)
inf L(" *(t)
V(O)M((())d
(5)
where 7(0) = 21: 7(1) = 22
Here, 7(t) : R + Z is a curve in the latent space with boundary conditions 7()) = 21 and r(1) = Zy, and 5(1) its derivative along the manifold (more details in Appendix C.I). Importantly, when the metric tensor satisfies M(z) Il for all z, the curve 7 (t) minimising Equation (4) is the straight line between latent codes, and the geodesics coincide with Euclidean lines in latent space.
To address this issue and complement existing methods with representations meeting their assumptions, we propose to regularise the latent space of single-cell VAE in such a way that the latent manifold Mz has an approximately Eu- clidean geometry. In other words, our goal is to encourage correspondence between straight paths in the latent space and geodesic interpolations along the statistical manifold. When this condition is satisfied, decoded trajectories generated using linear interpolations in the latent space of the VAE respect the geometry of the data manifold.
Stochastic decoders.
While in AES one deals with deter- ministic manifolds, in VAEs the decoder function h maps a latent code z E Z to the parameter configuration yp E H of the data likelihood. If the likelihood has continuous parame- ters, H. = R represents the parameter space. As such, the image of the decoder lies statistical manifold, which is smooth manifold of probability distributions. Such mani- folds have a natural metric tensor called Fisher Information Metric (FIM) (Nielsen, 2020; Arvanitidis et al., 2022). The FIM defines the local geometry of the statistical manifold and can be used to build the pullback metric for arbitrary decoders. For statistical manifold My with parameters &p E H. the FIM is formulated as
4.2. Assumptions
Before describing our regularisation approach, we state two assumptions about the geometry of the single-cell manifold induced by the negative binomial decoder:
M((p) = Ep(x(p) [V log ( logp(x((p)"] , (6)
1. Geodesic convexity: The manifold is assumed to be geodesically convex; that is, any two points on the manifold are connected by a unique geodesic.
where 4 h(z) and the metric tensor M(p)) E ROx Analogous to deterministic AEs, one can combine Equa- tion (6) and Equation (3) to formulate the pullback metric for an arbitrary statistical manifold, with the difference that the metric tensor is defined based on the parameter space H. Thus, the latent space of a VAE is endowed with the pullback metric for a statistical manifold.
2. Local-to-gloo to-gobal approximation: The manifold is suf- ficiently sampled such that enforcing local geometric constraints via the pullback metric at observed points yields a good approximation of the global geometry.
Assumption (1) is appropriate for acyclic biological pro-
((x) = (M)x))(()))TT((
4

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
cesses, such differentiation or perturbation responses, where cellular transitions follow smooth and directed pro- gressions. Assumption (2) aligns with standard practice in single-cell manifold learning, where local neighbourhood structure is leveraged to infer global geometry under the assumption of smooth state transitions.
for each gene g independently:
PND (N 0g)= - C
(0,
0g
Hg
(10)
0g + 1/y
9 + M4y
I(O,
Fg)
where
1. >
Ig !1 (0g)
with Tig E No and H4g = (z). Notably, since the decoder h produces cell-specific means, each cell is deemed as an individual probability distribution. Consequently, we as- sume that the single-cell data lies on a statistical manifold parameterised by the decoder in the space of negative bino- mial distributions. According to Equation (7), we pull back the FIM of the statistical manifold of the negative binomial probability distribution the latent manifold Ma.
4.3. Flattening Loss
To ensure that straight latent paths approximate geodesics on the decoded statistical manifold, we introduce a regular- isation term in the VAE objective that introduces a locally Euclidean geometry in the latent manifold MMz
As shown in Section 3.2, the local geometry of VAE'S latent space is determined by the metric in Equation (7). This metric is function of the Fisher information of the decoder's s likelihood, which depends on the decoded pa- rameters y E H. From Equation (4) we also know that if M(z) = la, then the geodesic distance between each pair of latent points is given by the straight line between them. Therefore, regularising the product Ju (z))M( (z) to- wards la forces a VAE to model locally Euclidean latent geometry. Crucially, the non-lineat decoder is still trained to reconstruct the original data space under the likelihood opti- misation task in the ELBO, preserving the local geometry of the decoded statistical manifold described by M(sp).
Proposition 4.1. The pullback metric at latent point 2 of the statistical manifold of negative binomial distributions, parameterised by a decoder h and fixed inverse dispersion 0, is given by:
6 M(z) = Σ hg(2)(h)(2)) 0g) )
Vali(z)) & Vahg(2),
(11)
where @ is the outer product of vectors, g indexes individual decoded dimensions, and hy(z) denotes the decoded mean for gene 9.
In summary, we introduce a flattening loss, Last, which encourages locally Euclidean latent geometry in VAEs with flexible decoders. This loss is combined with the BLBO to form the full VI objective:
We provide the derivation of Proposition 4.1 in Appendix B. Note that we only take the gradient of the mean decoder h since the inverse dispersion parameter is not function of the latent space in single-cell VAEs (see Section 3.1). This expression for M(z) is then used in the flattening loss (Equa- tion (8)) to train a geometry-regularised single-cell VAF.
Laa w. a) = (xlx) [M(()) aTd
8
Here, e and to are the VAE'S parameters, q the approximate posterior on the latent space learnt by the encoder, and M(z) is calculated by Equation (7). Meanwhile, a is trainable parameter offering some flexibility on the scale of the diagonal constraint while preserving straight geodesics. The Frobenius norm encourages each local pullback metric to be close to a scaled identity. In VAEs, the loss of FI combined with the ELBO:
5. Experiments
We evaluate FlatVI both simulated and real single-cell scenarios. We begin in Section 5.1 by demonstrating that our regularisation improves the approximation of constant Euclidean geometry in the latent manifold while preserving the reconstruction of likelihood parameters on synthetic data. For real-world validation, we focus on modelling single-cell population dynamics using linear dynamic OT, We investigate whether enforcing Euclidean geometry in the latent space improves both latent and gene-wise trajectory reconstruction. These results are presented in Section 5.2 and Section 5.3. Finally, in Section 5,4 and Section 5.5, we explore the representations learned by FIatVI and evaluate the effectiveness of linear latent interpolations for modelling developmental processes.
(ylavi(d, ai) LOLBO +) XCna a),
(9)
where À controls the strength of the flattening regularisation. We summarise the procedure used to train FlatVI in Algorithm 1.
4.4. VI on a Negative Binomial Single-Cell Manifold
In this work, we model cellular trajectories to study the evolution of biological processes through interpolations on a statistically grounded manifold. As outlined in Section 3.1, single-cell counts are modelled with a negative binomial decoder with the following univariate point mass function
5.1. Simulated Data
Task and datasets. To assess the effect of FlatVI on latent geometry, we evaluate its performance on a synthetic dataset
5

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
VoR NB VAE
CN NB VAE
Geodesics NB VAE
Table Comparison between VI and an unregularised NB-
80
VAE (A 0) in terms of MSE for and 0, and the 3-NN overlap
between Euclidean and geodesic neighbourhoods.
Latent
30
Latent
60
Latent
Reg. strength
MSE (Jt) (Q)
MSE (0) (
3-NN overlap (f)
À = 0
15.52 0.34
3.10
0,66 0.00
10
20
A
16.34
0.85
0.63 1.00
Latent
Latent
Latent 1
À 16.35 1 0.53
4.39
0.77 2 1.000
A= ==
14.75 a 0.12
3.20
0.67 a 0.01
VoR FlatVi A=7
CN Flatvi A=7
Geodesics FlatVi x=7
50
== 7
15.47 0.20
3.38 0.09
0.72 0.01
80
400
15.41
3.08 0.13
0.80
Latent
30
Latent
60
Latent
20
40
-10
20
generated from a multivariate negative binomial distribution. We aim to induce Euclidean latent geometry in 2D space while maintaining faithful data reconstruction.
Latent
Latent 1
Latent 1
Figure 2. Comparison of latent geometries for NB-VAE (top) and FIatVT with À 7 (hottom) using: (a) Variance of the Riemannian metric (VoR), (b) Condition Number (CN), and (c) Straightness of geodesics between high VoR regions.
We simulate 1000 observations from a 10-dimensional nega- tive binomial distribution with known mean (J4) and inverse dispersion (0). Each observation represents a simulated cell, uniformly assigned to one of three categories mimicking biological cell types (see Figure 6). For each type, the mean vectors are sampled from normal distributions centred at -1, O, and 1, respectively, with standard deviation of 1, and exponentiated to ensure positivity. Gene-spe speci-cci in- verse dispersion parameters are sampled from a Gamma distribution (concentration rate 1) and made positive. To comply with assumptions in real single-cell data, all data points share the same dispersion values. We provide further details on the simulation setup in Appendix 1.2.
that our regularisation forces geodesic distances under the pullback metric to better approximate Euclidean topology in the latent space compared to an unregularised Negative Bino- mial VAE (NB-VAE) with À = 0. In other words, increasing À improves the correspondence between the neighbourhood structures induced by pullback geodesic and Euclidean dis- tances (see Appendix 1.3.3 for additional metrics). Mean- while, the capabilities of our model to reconstruct the mean (Ju) and inverse dispersion (σr parameters do not degrade when the regularisation strength is increased.
For À E {0, 1, 5, 7, 10), we assess: (i) The MSE in reconstructing the true μl and 0 from decoded latent variables; and (ii) the 3-Nearest-Neige Neiihbor (3-NN) overlap between neighbourhoods defined by Euclidean and pullback geodesic distances in the latent space (see Appendix 1.3 for results using more neighbours). Geodesics are approxi- mated by parameterised cubic splines minimising the length under the pullback metric defined in Equation (7) (see Equa- tion (4)). In our setting, a successful regularisation recovers the true parameters of the data-generating process while imposing a similar neighbourhood structure between latent Euclidean and pullback geodesic distances.
Evaluation.
The plots in Figure 2 serve as additional proof of the flat- tening mechanism. Inducing Euclidean geometry into the latent space ensures a more uniform local geometry, as the latent manifold of FlatVI does not exhibit as many regions of systematically high VoR or CN as in the standard NB- VAE setting (see Figure 2a-b). Despite the flattening, some limited regions with high CN and VoR remain in the VI embedding, We qualitatively investigate the cause for high VoR and CN values in Appendix L...4.
In Figure 2c we sample couples of points from re- gions of high VoR in the NB-VAE latent space and plot geodesie paths approximated according to Equation (4) on both FI and the unregularised model. FlatVI achieves straight paths, while pullback-based geodesic interpolations in the standard NB-VAE bottleneck show a curvature (see Figure 7 for comparison with Euclidean manifold). These findings support our objective of ensuring that pullback- driven geodesies exhibit linear behaviour in the latent space.
We also visually assess Riemannian characteristics of the latent space: The variance of the Riemannian metric (VoR)
and the condition number (CN), following Chen et al. (2020)
and Yonghyeon et al. (2021). VoR quantifies how much the metric deviates from its spatial average M Eg-pp [M)) VoR of 0 implies a globally uniform metric, and we esti- mate its value using batches of 256 latent encodings. CN, defined as the ratio of the largest to the smallest eigenvalue of (z), approaches 1 when the metric resembles the iden- tity. Low VoR and CN values indicate well-flattened latent space. See Appendix I.I for definitions.
5.2. Reconstruction of scRNA-seq Trajectories
Task and dataset. Our core hypothesis is that FlatVT's latent space provides a better embedding for cell-state
Results. In our simulation setting, results in Table show
6

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Table Comparison of cellular trajectory reconstruction on held-out time points using different representation models. Latent trajectories are learnt with OT-CFM, leaving out intermediate time points and using them as ground truth for evaluating the interpolation of the cellular dynamies. Distribution matching metrics are evaluated to compare real held-out points and reconstructions thereof in both the decoded and latent space across three seeds.
EB
MEF
Latent
Decoded
Latent
Decoded
2- Wasserstein ())
I.2 (L)
2- Wasseesttei (1)
MMD (4)
2- 2-Wasserstein (L)
1.2 (L)
2-Wasserstein (4)
MMD (L)
GAF
2.16.014
0.40 0.06
70.29 0.00
0.14+004
2.49 a.2
0.57 0.07
106.83 000
0.38 0.01
NB-VAE
2.07 2007
0.30
43,36 0.39
0.09
2.07 2012
0.40 0.08
103.29
0.19
FlatVI
1.54 6.00
0.27 0.83
41.99 0.04
0.07 a 0.01
1.64
0.36 0.45
97.12
0.16 6.01
interpolation methods that assume Euclidean geometry, as our flattening loss encourages the VAE approximate local Euclidean geometry in the latent manifold. We demonstrate the benefits of using FlatVT's representation space in combination with Euclidean, continuous OT to map single-cell trajectories over time. As a dynamic OT algorithm, we use the OT Conditional Flow Matching (OT-CFM) model (Tong et al., 2024a), which leverages straight-line interpolation between samples to learn velocity field transporting cells across time (Appendix D).
sion quality using linear-kernel Maximum Mean Discrep- ancy (MMD) (Borgwardt et al., 2006) and 2-Wasserstein distance.
We set FlatVT's regularisation strength to x=1 for the EB dataset and X=0.1 for the MEF dataset. We tune the hyper- parameter based on the value that leads to the best repre- sentation for OT-based trajectory reconstruction on training data (see Appendix H for more details).
Results. Table 2 reports reconstruction metrics between true and interpolated latent cells. Across all datasets and metries, trajectories in FlatVI's Euclidean latent space result in better time point reconstruction compared to the baseline models, highlighting the effectiveness of our regularisation. Furthermore, FlatVI improves the quality of decoded gene expression trajectories, as reflected in lower MMD and Wasserstein distances. Table 6 additionally provides biological validation via improved reconstruction of lineage marker trajectories using FlatVI.
We evaluate using two real-world datasets: (i) The Embryoid Body (EB) dataset (Moon et al., 2019), comprising 18,203 differentiating human embryoid cells over five time points and spanning four lineages; and (ii) the MEF reprogram- ming dataset (Schiebinger et al., 2019), containing 165,892 cells across 39 time points, tracing the reprogramming of mouse embryonic fibroblasts into induced pluripotent stem cells. Full dataset details are provided in Appendix J.1.
Baselines. We compare FlatVI with a standard NB-VAE trained without regularisation (Lopez al., 2018) as a representation model for continuous OT. Additionally, we evaluate latent OT on embeddings produced by the GAE model from Huguet et al. (2022), described in Section 2. This model is trained on log-normalised gene expression to compensate for the absence of a discrete probabilistic decoder. Differences between FIatVI and GAE are discussed in Appendix 1.2. All three approaches are used to generate latent embeddings of time-resolved gene expression datasets. These embeddings are then used to train an OT-CFM model that learns latent trajectories from unpaired observations at consecutive time points.
5.3. Latent Vector Field and Lineage Mapping
Task and dataset. We evaluate the capacity of continuous OT to identify biologically meaningful cell velocity field using the representation spaces computed by FlatVI, the unregularised NB-VAE and the GAE model. We hypoth- esise that applying dynamic OT with Euclidean cost to a flat representation space is beneficial. As a dataset for the analysis, we employ the panereatic endocrinogenesis (here shortly denoted as Pancreas) by Bastidas which measures 16,206 cells and spans embryonic days 14.5 to 15.5, revealing multipotent cell differentiation into endocrine and non- non-endocrine lineages. More specifically, we train the compared representation learning frameworks on the dataset and learn separate vector fields for all models' embeddings, matching days 14.5 to 15.5 with OT-CFM. The learnt vector field represents the directionality of the observations on the cellular development manifold.
et al. (2019),
Evaluation. Following Tong et al. (2020), we leave out intermediate time points during training and assess the model's s ability reconstruct them via OT. The accuracy of reconstructing an unseen time point from t-1 reflects the model's 's interpolation ability along the data manifold. We use this paradigm to compare different representation spaces. For quantitative evaluation, we compute the 2- Wasserstein and mean L distances between real and reconstructed latent cells at each time point. We also assess decoded gene expres-
Evaluation. Using the CellRank model (Lange et al., 2022; Weiler al., 2024), we build random walks on a cell graph based the directionality of latent velocities learnt by OT-CFM in the different representation spaces.
7

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
a
GAZ
NB NAE
Fatvi
EB
MEF
Pancreas
Plakk is WeEk
1.00
PCZ -205%
PC2 18.07%
Initial Intermadlate
Velocity Corsistency
0.95
NB-VAE
0.90
Flatin lgs ling
0.85
PQ1- 1618%
18.25%
PCT- 8883%
0.80
Aciiar Cactal
0.75
10.22%
10.0%
Aphh
0.70
GAE
Pound tamiinal
PCD
X Mixsed taminall ctis
Epalin
Latenn Dimension
PCT 10.60%
101-11.22%
FC1- 8.9%
Figure 3. Learning terminal states from OT-CFM's cell velocities in the Pancreas datasct. (a) Terminal states found by CellRank using different representation models. (b) Latent velocity consis- tency computed for cells across different latent space sizes.
2751K
VIX 2008
Poz 33.04%
32.41%
PCI 60100%
PCI 201-4880 85%
Walks converge to macrostates representing the endpoints of the biological process if the learnt velocity field points to biologically meaningful directions. We quantify the quality of vector fields learnt by OT in different latent spaces based on (i) the number of macrostates identified by random walks, and (ii) the velocity consistency, measured as the correlation of the latent velocity field of single datapoints with that of the neighbouring cells (Gayoso et al., 2024). Higher consistency indicates smoother transitions in the vector field, suggesting that the representation space facilitates more coherent and biologically meaningful dynamics, making it suitable space for learning trajectories (see Appendix L.I..
Figure 4. 2D PCA plots of the latent spaces computed by GAE, BBBVAE and FlatVI. Marked are initial, intermediate and terminal cell states along the biological trajectory.
between initial and terminal cell types. In Table 8 we show that such a separation is more pronounced than competing models, also on the Pancreatic dataset based on quantitative clustering metrics. Moreover, the higher variance explained by individual PCs in FlatVI's latent space suggests that our model captures the main sources of variation (the biological trajectories) more efficiently, reducing latent space dimen- sionality and enhancing information compression while pre- serving even improving biological fidelity. Consequently, FIatVI is well-suited for smaller latent spaces, making it promising input for trajectory inference methods.
Results. Figure 3a summarises the number of terminal cell states identified by following the velocity graph. From prior biological knowledge, it is known that the dataset contains six terminal states, which are all identified on the represen- tation computed by our FIatVI (A = 1). In contrast, on the GAE and NB-VAE's representations, CellRank only captures four and five terminal states, respectively. In Fig- ure 3b, we further evaluate the velocity consistency within neighbourhoods of cells as a function of latent dimensional- ity. In line with previous results, OT on the approximately Fuclidean latent space computed by FIatVI yields more consistent velocity field across latent space sizes.
5.5. Decoded Geodesic Interpolations
Task and evaluation. We qualitatively evaluate whether linear interpolations in the latent space of FlatVI yield bio- logically meaningful trajectories when decoded into gene expression space. Given two cells at different stages along a developmental lineage, we linearly interpolate between their latent representations and decode the intermediate points, We then inspect whether the expression of known marker genes along these decoded paths evolves consistently with expected biological progression. If so, this suggests that linear paths in FlatVI's latent space approximate geodesics aligned with the underlying developmental manifold.
5.4. Single-Cell Data Representations
We visualise single-cell latent representations on the previ- ously introduced datasets computed using the FIatVI, NB- VAE, and GAE models. For FI the value of is set to 1 for EB and Pancreas and 0.1 for the MEF dataset, in line with previous settings. In Figure 4, we compare the Princi- pal Component (PC) embeddings of FlatVT's latent space with competing models, highlighting initial and terminal cellular states. Despite the regularisation, FlatVI represents the biological structure in the latent space better or on par with the baselines, as illustrated by the separation between initial and terminal states. This is particularly evident in the MEF dataset, where FlatVI provides clearer division
Baseline. As baseline, we use GAGA (Sun al., 2025). which explicitly enforces alignment between the latent ge- ometry and data strueture using neighbour-basee regulari- sation. Unlike Fl GAGA does not assume tractable latent manifold. Instead, it learns data-driven representa- tion where geodesics are approximated via a neural ODE. In contrast, VI enables efficient approximations of latent geodesics through simple linear interpolation.
8

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
a
Lineage: Multipotent - Beta/ / Marker: Cpe
Lineage: Multipotent Beta/ Marker: Nnat
Lineage: Multipotent Delta / Marker: Cpe
Flatvi
Expression
GAGA
GAGA
3
GAGA
2
Multipotent
Ngna low EP
Ngn3 High early
Ngn3 High late
Fov-+ Beta
Beta
Multipotent
Ngná low EP
Ngn3 High early
Non3 High late
Fav-+ Beta
Beta
Muttipotent
Ngn3 low EP
Ngn3 High early
Ngn3 High late
Fevt Delta
Detra
Lineage-
Lineage
Lineage
Figure 5. Decoded marker trajectories from latent interpolations on the Pancreas datasct, comparing FlatVI and GAGA. We sample 100 pairs of multipotent and mature cells from each terminal state, encode their gene expression profiles, and perform latent interpolations (linearly in VIatVI and using neural ODE in GAGA). Intermediate states along these trajectories are decoded, and the resulting marker gene expression visualised alongside the real expression along the lineage (dark grey points). Each trajectory is shown as low-opacity line, while the solid line denotes the mean trajectory. We show three examples (a, b, and e) from the Beta and Delta lineages.
Dataset. We use the pancreas endocrinogenesis dataset described in Section 5.3, focusing the endocrine devel- opmental trajectory. We randomly sample batches of 100 multipotent progenitor cells and 100 terminal-stage cells from each of the endocrine branch lineages. Pairs of mul- tipotent and mature cells are randomly matched, and their decoded interpolation paths are analysed for marker gene expression dynamics.
cellular manifolds. Furthermore, linear interpolations between latent cellular states offer interpretable insights into the progression of cell states, providing straightforward approach to exploring dynamical biological processes. Col- lectively, these improvements enhance core tasks in cellular development, such as fate mapping and the reconstruction of differentiation pathways, establishing FlatVI as a useful tool for trajectory inference in single-cell transcriptomics.
Results. While both approaches produce reasonable results on average (see Figure 13), in Figure 5 we highlight some recurrent sub-optimal patterns in GAGA's performance that do not arise in VI Specifically, in Figure we present examples of unstable optimisation (Figure 5a), underestima- tion of the marker expression (Figure 5b) and overestimation of intermediate expression patterns (Figure 5c). In contrast, FlatVT's latent space consistently produces decoded marker dynamics that faithfully recapitulate the expected fate- specific trajectories. Reasonably, interpolations computed with FIatVI are less computationally expensive (Figure 14).
Limitations and future work. To improve FlatVI's applicability to real-world datasets, we aim to enhance its robustness to reconstruction loss on biological data, reducing potential trade-offs between flattening and reconstruction likelihood. As discussed in Section 4, enforcing a locally Euclidean latent geometry imposes strong assumptions that may not hold for all datasets (e.g.. those dominated by cyclic processes such as the cell cycle). Future work will investigate alternative latent geometries to better capture diverse biological structures. We also plan to extend FlatVl to a broader class of statistical manifolds and single-cell tasks, including Poisson-based modelling of chromatin accessibility, evaluation in batch correction settings, and OT-mediated perturbation modelling.
These findings highlight the effectiveness of FlatVI as straightforward and reliable method for exploring cellular manifolds, without requiring a parameterised interpolant.
Impact Statement
6. Conclusion
The presented work deals with fundamental characteristics of scRNA-seq data and studies how efficient representa- tions of complex high-dimensional cellular data can help to address key biological questions. We envision the release of FIatVI as user-friendly, open-source tool to enable its widespread use as option for single-cell analysis. Deal- ing with biological data, FlatVI could be used in sensitive settings involving clinical information and patient data.
We addressed the problem of modelling cellular trajectories in seRNA-seq data by introducing FlatVI, a VAE training strategy that enforces locally Euclidean geometry in the la- tent space by regularising the pullback metric of the stochas- tic decoder. This regularisation encourages straight latent paths to approximate geodesic interpolations in the decoded data space. Experiments on synthetic data demonstrate that VIatVI successfully induces a latent Euclidean geometry while preserving accurate parameter reconstruction. When combined with dynamic OT, FI improves trajectory pre- diction performance and yields more consistent vector fields
9

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Acknowledgments
Borgwardt, K. M., Gretton, A., Rasch, M. J., Kriegel, H.-P., Schölkopf, B.,, and Smola, A. J. Integrating structured biological data by kernel maximum mean discrepancy. Bioinformaties, 22(14):449-c) 2000
The authors give special thanks to Sören Becker for pro- viding valuable feedback on the manuscript. Additionally, A.P. and L.H. are supported by the Helmholtz Association under the joint research school Munich School for Data Science (MUDS). A.P., S.G. and F.J.T. also acknowledge support from the German Federal Ministry of Education and Research (BMBF) through grant numbers 031L0289A, 031L0289C, OIIS DIIS18053, 18053A and 03 03100269C. F.J.T. acknowl- edges support from the Helmholtz Association's 's Initiative and Networking Fund via the CausalCellDynamiis project (grant number interlabs-0029) Finally, F.J.T. acknowledges support from the European Union (ERC, DeepCell grant number 101054957 and ERA PerMed-NET Transcan-3- grant agreement number O1KT2308).
Bunne, C., Papaxanthos, L.,, Krause, A., and Cuturi, M. Proximal optimal transport modeling of population dy- namics. In International Conference on Artificial Intelli- gence and Statistics, pp. 651 6511-6528. PMLR, 2022.
Bunne, C., Stark, S. G., Gut, G., Del Castillo, J. S., Levesque, M., Lehmann, K.-V., Pelkmans, L., Krause, A., and Rätsch, G. Learning single-cell perturbation re- sponses using neural optimal transport. Nature methods, 20(11):1759-177 2023.
Chen, N., Klushyn, A., Ferroni, F., Bayer, J., and Van Der Smagt, P. Learning Flat Latent Manifolds with VABs, In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of o Proceedings of Machine Learning Research, pp. 1587-1596. PMLR, 13-18 Jul 2020.
References
Albergo, M. S. and Vanden-Eijnden E. Building Normaliz- ing Flows with Stochastic Interpolants. In The Eleventh International Conference on Learning Representations, 2023.
Do Carmo, M. and Flaherty Francis, J. Riemannian geometry, volume 2. Springer, 1992.
Arvanitidis, G., Hauberg, S., and Schölkopf, B. Geometri- cally Enriched Latent Spaces. In Banerjee, A. and Fuku- mizu, K. (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, vol- ume 130 of Proceedings of Machine Learning Research, pp. 631-6339 PMLR, 13-15 Apr 2021.
Duque, A. F. Morin, S., Wolf, G., and Moon, K. Extendable and invertible manifold learning with geometry regular- ized autoencoders. In 2020 IEEE International Con- ference on Big Data (Big Data), pp. 5027-5036. IEEE, 2020.
Džeroski, S. and Todorovski, L. Learning population dynam- ics models from data and domain knowledge. Ecological Modelling, 170(2-3): 129-)11211100000 140, 2003.
Arvanitidis, G., González-Duque, M., Pouplin, A., Kalatzis, D., and Hauberg, S. Pulling back information geometry In Camps-Vall G., Ruiz, F. J. R., and Valera, (eds.), Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Pro- ceedings of Machine Learning Research, pp. 4872-4894. PMLR, 28-30 Mar 2022.
Eraslan, G., Simon, L. M., Mircea, M., Mueller, N. S., and RNA-seq
Theis, F. J. Single-cell RNA- seq denoising using a deep count autoencoder. Nature communications, 10(1):390, 2019.
Eyring, L., Klein, D., Palla, G., Becker, S., Weiler, P., Kil- bertus, N., and Theis, F. J. Modeling Single-Cell Dynam- ies Using Unbalanced Parameterized Monge Maps. In NeurIPS 2022 Workshop on Learning Meaningful Repre- sentations of Life. 2022.
Bastidas- Ponce, A., Tritschler, S., Dony, L., Scheibner, K., Tarquis-Meding M., Salinno, C., Schirge, S., Burtscher, I., Böttcher, A., Theis, F. J., et al. Comprehensive single cell mRNA profiling reveals detailed roadmap for pancreatic endocrinogenesis. Development, 146(12); dev dev173849, 2019.
Fasina, O., Huguet, G., Tong, A., Zhang, Y., Wolf, G., Nickel, M., Adelstein, I., and Krishnaswamy, S. Neural FIM for learning Fisher information metrics from point cloud data. In International Conference on Machine Learning, pp. 9814-9826. PMLR, 2023.
Benamou, J. D. and Brenier, Y. A computational fluid me- chanics solution to the Monte-Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, Jan- uary 2000. doi: 10.1007/s0021 10050002.
Gayoso, A., Weiler, P., Lotfollahi, M., Klein, D., Hong. J., Streets, A., Theis, F. J., and Yosef, N. Deep generative modeling of transcriptional dynamics for RNA velocity analysis in single cells. Nature methods, 21(1):50-59, 2024.
Bergen, V., Lange M., Peidli. S., Wolf, F., and Theis, F. Generalizing RNA velocity to transient cell states through dynamical modeling. Nature Biotechnology, 38( 1408 1414, August 2020. doi: 10. 10111//411177020--0000
10

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Griffiths, R.-R. and Hernández-Lobato, M. Constrained Bayesian optimization for automatic chemical design us- ing variational autoencoders. Chemical science, 11(2): 577-586, 2020.
M., and Theis, J.J. Mapping cells through time and space with moscot. Nature, January 2025. ISSN 1476-4687. doi: 10. 10088641156-002-88445...
Korem, Y., Szekely, P., Hart, Y., Sheftel, H., Hausser, J., Mayo, A., Rothenberg, M. E., Kalisky, T., and Alon, U. Geometry of the gene expression space of individual cells. 2015.
Haque, A., Engel, J., Teichmann, S. A., and Lönnberg, T. A practical guide to single-cell RNA- -equeennug for biomedical research and clinical applications. Genome medicine, 9:1-12, 2017.
PLoS computational biology,
Koshizuka, T. and Sato, I. Neural Lagrangian Schrödinger Bridge: Diffusion Modeling for Population Dynamics. In The Eleventh International Conference on Learning Representations, 2023.
Haviv, D., Kunes, R. Z., Dougherty, T., Burdziak, C., Nawy, T., Gilbert, A., and Pe'Er, D. Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformer. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.). Pro- ceedings of the 4 st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learn- ing Research, pp. 17697-17718 PMLR, 21-27 Jul 2024.
Lange, M., Bergen, V., Klein, M., Setty, M., Reuter, B., Bakhti, M., Lickert, H., Ansari, M., Schniering, J., Schiller, H. B., et al. CellRank for directed single-cell fate mapping. Nature methods, 19(2):19--770, 2022.
Lehmann, E. L. and Casella, G. Theory of point estimation. Springer Science & Business Media, 2006.
Hetzel, L., Boehm, S., Kilbertus, N., Günnemann, S., Theis, E., et al. Predicting cellular responses novel drug per- turbations at a single-cell resolution. Advances in Neural Information Processing Systems, 35:2671 11-2222, 2022.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow Matching for Generative Modeling. In The Eleventh International Conference on Learning Representations, 2023.
Heumos, L., Schaar, A. C., Lance, C., Litinetskaya, A., Drost, F., Zappia, L., Lücken, M. D., Strobl, D. C., Henao, J., Curion, F., et al. Best practices for single-cell analysis across modalities. Nature Reviews Genetics, 24(8):550- 572, 2023.
Lopez, R., Regier, J., Cole, M. B., Jordan, M. L,, and Yosef, N. Deep generative modeling for single-cell transcrip- tomics. Nature methods, 15(12):1053-1058 2018.
Huguet, G., Magruder, D. S., Tong, A., Fasina, O., Kuchroo, M.. Wolf, G., and Krishnaswamy, S. Manifold inter- polating aptimal-transpont flows for trajectory inference. Advances in neural information processing systems, 35: 29705-29718, 2022.
Lotfollahi, M., Klimovskaia Susmelj. A., De Donno, C., Het- zel, L., Ji, Y., Ibarra, I. L., Srivatsan, S. R., Naghipourfar, M., Daza, R. M., Martin, B., Shendure, J., McFaline- Figueroa, J. L., Boyeau, P., Wolf, F. A., Yakubova, N., Günnemann, S., Trapnell, C., Lopez-Paz, D., and Theis, F. J. Predicting cellular responses to complex perturbations in high- gighthhroughpt screens. Molecular Sys- lems Biology, 19(6), May 2023. ISSN 1744-4292. doi: (0.15252//sb.2222115111
Huguet, G., Tong, A., De Brouwer, E., Zhang, Y., Wolf, G., Adelstein, I., and Krishnaswamy, S. A Heat Diffu- sion Perspective on Geodesic Preserving Dimensionality Reduction. Advances in Neural Information Processing Systems, 36, 2024.
Luecken, M. D., Büttner, M., Chaichoompu, K., Danese, A., Interlandi, M.. Mueller, M. F., Strobl, D. C., Zap- pia, L., Dugas, M., Colooéé-Tatché, M., and Theis, F. J. Benchmarking atlas-level data integration in single-cell genomics. Nature Methods, 19(1):41-50, December 2021. ISSN 1548-7105. doi: 10.1033/641592-021-111
Kapusniak, K., Potaptchik, P., Reu, T., Zhang. L., Tong, A., Bronstein, M., Bose, J., and Di Giovanni, F. Metric flow matching for smooth interpolations on the data manifold. Advances in Neural Information Processing Systems, 37: 13501 2024.
Moon, K. R., Van Dijk, D., Wang, Z., Gigante, S., Burkhardt, D. B., Chen, W. S., Yim, K., Elzen, A. v. d., Hir,, M. J., Coifman, R. R., et al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechnol- ogy, 37(11): 1482-142, 2019.
Kingma, D. P. and Welling, M. Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
Klein, D., Palla, G., Lange, M., Klein, M., Piran, Z., Gander, M., Paang--Paxaattoos L., Sterr, M., Saber, L., Jing, C., Bastidas-Ponee A., Cota, P., Tarquis-Medina, M., Parikh, S., Gold, I., Lickert, H., Bakhti, M., Nitzan, M., Cuturi,
Nazari, P., Damrich, S., and Hamprecht, F. A. Geomet- ric Autoencoders - What You See is What You Decode. In Krause, A., Brunskill, E.,, Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp- 25834-25857. PMLR, 23-29 Jul 2023.
Sun, X., Liao, D., MacDonald, K., Zhang, Y., Huguet, G., Wolf, G., Adelstein, I., Rudner, T. G. J., and Krish- naswamy, S. Geometry-Aware Generative Autoencoder for Warped Riemannian Metric Learning and Generative Modeling on Data Manifolds. In The 28th International Conference on Artificial Intelligence and Statistics, 2025.
Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. Action matching: Learning stochastic dynamics from samples. In International conference on machine learn- ing. pp. 25858 258889 PMLR, 2023.
Tong, A., Huang, J., Wolf, G., Van Dijk, D., and Krish- naswamy, S. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9526-9536. PMLR, 2020.
Nielsen, F. An elementary introduction to information ge- ometry. Entropy, 22(10): 2020.
Pennec, X., Fillard, P.. and Ayache, N. A Riemannian Framework for Tensor Computing. International Journal of Computer Vision, 66(1) 66(1):41-4)) 1166666666, January 2006. doi: 10. 10.1007/s1 1233-005-3222-
Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. Improv- ing and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024a.
Peyré, G., Cuturi, M., al. Computational optimal transport: With applications to data science. Technical Report 5-6, 2019.
Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulation- Free Schrödinger Bridges via Score and Flow Matching. Dasgupta, S., Mandt, S., and Li, (eds.), Proceed- ings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pp. 1287. PMLR, 02 04 May 2024b.
Pooladian, A.-A., Ben-Hamu, H., Domingo-Enrich, C., Amos, B., Lipman. Y., and Chen, R. T. Multisample Straightening Flows with Minibatch Cou- plings. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, vol- ume 202 of Proceedings of Machine Learning Research, pp. 28 18100-220 28127. PMLR, 23-29 Jul 2023.
Flow Matching:
Weiler, P., Lange, M., Klein, M., Pe'er, D., and Theis, F, CellRank 2: unified fate mapping in multiview single-cell data. Nature Methods, 21(7): 196-1205 2024.
Qiu, X., Zhang, Y., Martin Rufinn, J. D., Weng, C., Hos- seinzadeh, S., Yang, D., Pogson, A. N., Hein, M. Y., Min, K. H. J., Wang, L., et al. Mapping transcriptomic vector fields of single cells. Cell, 185(4))690-711, 2022.
Wolf, F. A., Angerer, P., and Theis, SCANPY: large- scale single-cell gene expression data analysis. Genome biology, 19:1-5, 2018
Regev, A., Teichmann, A., Lander, S., Amit, I., Benoist, C., Birney, E., Bodenmiller, B., Campbell, P., Carninci, P., Clatworthy, M., et al. The human cell atlas. elife, 6: e27041, 2017.
Yonghyeon, L., Yoon, S., Son, M., and Park, F. C. Regular- ized autoencoders for isometric representation learning. In International Conference on Learning Representations, 2021.
Reuter, B., Fackeldey, K., and Weber, M. Generalized Markov modeling of nonreversible molecular kinetics. The Journal of Chemical Physics, 150(17):174103,, 2019. doi: 10. 00063/155064530
Zamir, R. A proof of the Fisher information inequality via a data processing argument. IEEE Transactions on Information Theory, 44(3):1246-1255 1250, 1998.
Zhong, W. and Meidani, H. Pi-vae: Physics-informed varia- tional auto-encoder for stochastic differential equations. Computer Methods in Applied Mechanics and Engineer- ing. 403:1 15664, 2023.
Rifkin, S. A. and Kim, J. Geometry of gene expression dynamics. Bioinformatics, 1809):117-11888 2002.
Sarker, I. H. Machine learning: Algorithms, real-world ap- plications and research directions. SN computer science, 2(3):160, 2021.
Zhou, Y.H.,, Xia, K., and Wright, F. A. A powerful and flexible approach to the analysis of RNA sequence count data. Bioinformatics, 27(19) 2011.
Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subrama- nian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in repro- gramming. Cell, 176(4):928 943, 2019.
12

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
A. Code and Datasets
have made our code publicly available at https: / www/tteesssab//llaaattooeooe All datasets used this study are open source, and their associated publications are cited in the manuscript.
B. Derivation of the Fisher Information Metric for the Negative Binomial Distribution
We first show that the Fisher information of univariate Negative Binomial (NB) distribution parameterised by the mean and inverse dispersion with respect to yu is
M(U) μlf + +)
(12)
We then move on with the derivation of the pullback metric in Proposition 4.1.
Fisher information of the NB distribution. The univariate NB probability distribution parameterised by mean μ and inverse dispersion 0 is
r r(0+x) + x) x0T(0)
H
(N 1,,0)
(13)
1 + u
1 +
The Fisher information the distribution can be computed with respect to ja as:
0²
M(²) = p)
log PNB( 1,,0)
(14)
Oμ12
where
PN 14, 0) = 0 (log(0) - log(0 + [4] + (log(u) log(O [4]
(15)
with C = log(T(( + z)) log(z!) lo)( Then, it can be shown that
0²
+
log 1)
(16)
0p²
(A p)²
2²
Using the fact that the parameterisation involving the mean y and inverse dispersion implies that
Ep(elu.0) [x]] = u
(17)
we can expand Equation (14) as follows
0 ++x
M(²) E Ep(elu,)) ()
(0 1u)2 1x2
1
Ep(a m.0) 10 + x] +
w ppt [x]
(18)
plu + 0)
Derivation of the Fisher information metric. We here consider the NB-VAE case, where the likelihood is parameterised
by P' = (2) and 0g independently for each gene g.
When h is a continuously differentiable function of the pullback metric M, of the output g w.r.t by the reparameteri- sation property (Lehmann & Casella, 2006) is
Mz(z) Vehy (z)
0
V (z) (zllg())&Vah((e
(19)
(2) 0g)
where & is the outer product of vectors, and the gradients are column vectors.
13

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
By the chain rule, the joint Fisher information of independent random variables equals the sum of the Fisher information values of each variable (Zamir, 1998). As all Vy are independent given in the NB-VAE, the resulting Fisher Information Metric (FIM) is
M(z)= = Σ (g(z)
0g
h (z) (2) Vehy(z) & Vzhg(z).
(20)
C. The Geometry of AEs
We deal with the assumption that the observed data lies near Riemannian manifold Mx embedded in the ambient space X' = =R The manifold M, is defined as follows:
Definition C.1. A Riemannian manifold is smooth manifold M, endowed with a Riemannian metric M(x) for x E Mxx M(x) changes smoothly and identifies an inner product on the tangent space Mx at point E MMx as (u, v) u² v,)v,, with v, E Mx.
For an embedded manifold Mx with intrinsic dimension d, we can assume the existence of an invertible global chart map {: M:M + Ed mapping the manifold Mxx to its intrinsic coordinates. A vector Vx E Mx on the tangent space of Mxx can be expressed as pushforward Vx = Jg-1 (z)v2 of a tangent vector V2 E R at z = f(x), where I indicates the Jacobian. Therefore, Jg-- maps vectors Vz E Rd into the tangent space of the embedded manifold Mx. The ambient metric M(x) can be related the metric M(z) defined in terms of intrinsic coordinates via:
M(z) = Jg 1 (z) (x)-((2))J (z).
(21)
In other words, we can use the metric M(z) to compute quantities on the manifold, such as geodesic paths. However, for an embedded manifold Mx, the chart map & is usually not known. A workaround is to define the geometry of MM on another Riemannian manifold Mı with a trivial chart map E(z) z for z E M2, which can be mapped to MMx via a smooth immersion In the next section, we elaborate on the connection between manifold learning and autoencoders.
C.1. Deterministic AEs
We assume the decoder h: = X RRG a deterministic autoencoder is an immersion of a latent manifold with trivial chart map into Riemannian manifold Mx embedded in X' and with metric M. This valid one also assumes that is the intrinsic dimension of Mx. As explained before, the Jacobian of the decoder maps tangent vectors V2 E T2 to tangent vectors Vx=h(z) Tx.xx The decoder induces metric into the latent space following Equation (21) as
M(z) (z)² M(h(2))Jx.(z))
(22)
called pullback metric. The pullback metric defines the geometry of the latent manifold Mz compared that of the manifold Mx. The metric tensor M(z) regulates the inner product of tangent vectors u, and v. on the tangent space TzM2z
(uy, Vz)Me uu M(zV(
(23)
To enhance latent representation learning, distances in the latent space 2 can be optimised according quantities of interest in the observation space X,, following the geometry of Mx. For instance, we can define the length of curve 7 : [0. Z in the latent space by measuring its length on the manifold Mx:
L(7) [ |
= 1''' V((0)M((()))())....
(24)
where the equality is derived by applying the chain rule of differentiation.
14

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
C.2. Pulling Back the Information Geometry
In machine learning, exploring latent spaces is crucial. particularly in generative models such as VAEs, One challenge is defining meaningful distances in the latent space 2, which often depends the properties of stochastic decoders and their alignment with the observation space. Injecting the geometry of the decoded space of VAE into the latent space requires different theoretical framework, where the data is assumed to lie near a statistical manifold.
VAES can model various data types by utilising the decoder function as a non-linear likelihood parameter estimation model. We consider the decoder's output space as a parameter space H. for a probability density function. Depending on the data type, we express a likelihood function p(x (p) with parameters & E H. reformulated as p(x z) through a mapping h: 2 + H. We aim to define a natural distance measure in Z for infinitesimally close points 2 and z2 2i òz when seen from H. One can show that such distance corresponds to the Kullback-Leibber (KI.) divergence:
dist (z1. z2) (L(((x 21). p(x z2))
(25)
To define the geometry of the statistical manifold, one can resort to information geometry, which studies probabilistic densities represented by parameters &p E H. In this framework, H becomes a statistical manifold equipped with a FIM:
M((p) V logp(x e) log 40)" p( (0) dx.
(26)
The FIM locally approximates the KL divergence. For a univariate density p. parameterised by it is known that
1 60 M(())) + + o(6x²).
KL(y(> p(r I p + oΦ))
(27)
In the VAE setting, we view the decoder not as a mapping to the observation space X' but as transformation to the parameter space H. This perspective allows us to naturally incorporate the FIM into the latent space 2. Consequently, the VAE's decoder can be seen as spanning a manifold Mr in H. with Mæ inheriting the metric in Equation (26) via the Riemannian pullback. Based on this, we define statistical manifold.
Definition C.2. A statistical manifold represented by parameter space of a distribution p(x (p) and is endowed with the FIM as the Riemannian metric.
The Riemannian pullback metric derived as in Equation (22). Having defined the Riemannian pullback metric for VAEs with arbitrary likelihoods, one can extend the measurement of curve lengths in Z when mapped to H through h as displayed by Equation (24). This approach allows flexibility in the choice of the decoder, as long as the FIM of the chosen distribution type is tractable.
D. Learning Population Dynamics with Optimal Transport
The complexity of learning trajectories in high-dimensional data can be prevented by interpolating latent representations and decoding intermediate results the data space for inspection. Here, we deal with learning population dynamics, which consists of modelling the temporal evolution of dynamical system from unpaired samples of observations through time. As such, the task is naturally formulated as distribution matching problem, and dynamic Optimal Transport (OT) has been a popular avenue for population dynamics.
Let the data be defined on a continuous space X' = Rdd OT computes the most efficient mapping for transporting mass from one measure y to another 7), defined on X. Relevant to dynamical systems, Benamou & Brenier (2000) introduced a continuous formulation of the OT problem. this setting, let pr time-varying density over IRd constrained by po = v and P1 = 17. Dynamic OT learns a time-dependent marginal vector field u [o, 1] x Rd + Ed. where ux(x) u(t, x). Such field is associated with an Ordinary Differential Equation (ODE), dx u, (x)df, whose solution matches the source with the target distribution. Therefore, one can use dynamic OT to learn a system's 's dynamics from snapshots of data collected over time.
An efficient simulation- free formulation of dynamic OT comes from the OT Conditional Flow Matching (OT-CFM) model by Tong et al. (2024a) and Pooladian et al. (2023), who demonstrated that the time-resolved marginal vector field u, (x()) has the same minimiser as the data-conditioned vector field u, (x) Xn, x1), where (Xo, x1) ~ a are tuples of points sampled
15

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
from the static OT coupling between source and target batches. Assuming Gaussian marginals Pt and Xo and x1 be connected by Gaussian flows, both pt (x X1) and uy X Xo, x1) become tractable:
Pc Xo; xi) = ((x) | + (1 - x)xo, o²)
(28)
ux(x Xo: Xi) = x1 x0
(29)
where the value of o² is small pre-defined constant. Accordingly. the OT-CFM loss is
(5)C =
(t, x) - 4. (x|xx,, x1) x1)
with t U(, 1).
(30)
Here, ve((, x) is neural network approximating the marginal vector field (x). Given this formulation of dynamic OT, we highlight three aspects:
1. Dynamic OT only applies to continuous spaces.
2. OT-CFM benefits from low-dimensional representations since the OTTcoupling is optimised from distances in the state space.
3. Based on Equation (28), OT-CFM uses straight lines to optimise the conditional vector field, thus assuming Euclidean geometry.
In the presence of discrete data like scRNA-seq counts modelled with VAEs, one can tackle (1) and (2) by learning dynamics in low-dimensional representation of the state space, the latent space of VAE with a however, that (3) still a shortcoming, since straight lines the latent space of a VAE do not reflect geodesic paths on the decoded data manifold unless enforced otherwise, see Figure 1. In this work, we address this remaining limitation by regularising the VAB's latent space induce locally Euclidean geometry, ensuring that straight lines in latent space better approximate geodesics the decoded data manifold.
decoder. Note,
E. Baseline Description
E.1. GAE
E.1.1. THE MODEL
Here, we describe the Geodesic Autocncoder (GAE) from Huguet et al. (2022). For more details on the theoretical framework, we refer to the original publication. The GAE works by matching Euclidean distances between latent codes with the diffusion geodesic distance, which is an approximation the diffusion ground distance in the observation space.
Briefly, the authors compute graph with an affinity matrix based on distances between observations i andd j using Gaussian kernel as:
(K.ig = ke (x,,
(31)
with scale parameter e, where X, X' and X' the observation space. The affinity is then density- normalised by M. -- T-KQQ-- where Q is diagonal matrix such that Q.; E3 (K.).. To compute the diffusion geodesic distance, the authors additionally calculate the diffusion matrix P. = D M,, with D = Σド= ""_((M,, (M.)sy and stationary distribution π, = D,a/ Σ, D,y- The diffusion geodesic distance between observations x; and x, is
K
Ga(xi,Xg) = Σ (K -)K--)ajjn (P (P.) 11 2 1
(32)
k=
with a E (0, 1/2). The running value of &s in Equation (32) defines the scales which similarity between the random walks starting at x; and Xj computed.
Given the diffusion geodesic distance G., defined in Equation (32), the GAE model is trained such that the pairwise Euclidean distances between latent codes approximate the diffusion geodesic distances in the observation space X,, in a batch of size B. Given an encoder fø I > IRd, the reconstruction loss is optimised alongside geodesic loss
Cooodesic (w)
2 B
Σ
(x))) fo(x2)112 (x,,x,)))
(33)
i=
16

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
1.1. ADDITIONAL COMPARISON BETWEEN FLATVI AND GAE
Although related in scope, VI significantly differs from the Geodesic Autoencoder (GAB) proposed by Huguet et al. (2022). Firstly, GAE is deterministic autoencoder optimised for reconstruction based Mean Squared Error (MSE) loss. As such, the model is not tailored to simulate gene counts. On the contrary, VlatVT's decoder parameterises Negative Binomial likelihood, allowing realistic generation of count data through sampling. This aspect has two advantages. By focusing on learning continuous parameters of a discrete likelihood, FlatVI explicitly models distributional properties of single-cell transcriptomics data, such as overdispersion, sparsity and discreteness. In contrast, a fully connected Gaussian decoder produces dense and continuous cells, failing to preserve the characteristics of the data. Moreover, the GAE regularises the latent space by approximating geodesic distances via a k-nearest neighbou graph constructed in the observation space. This method requires the computation of pairwise Euclidean distances in the observation space. As suggested previously, gene expression is nigh-dimensional and, therefore, deceiving due to the curse of dimensionality, On the contrary, leveraging only the Jacobian and the output of the decoder to enforce latent space Euclideanicity, FlatVI is more suitable for larger datasets and eludes computing distances in high dimensions.
E.2. Geometry-Aware Generative Autoencoder (GAGA)
E.2.1. THE MODEL
GAGA (Sun et al., 2025) is an AE model trained such that latent distances approximate those on the data manifold, as estimated via PHATE (Moon et al., 2019). This estimation in data space enables the imposition of pullback Riemannian metric in the latent space via the encoder, thereby aligning the geometries of the latent space and the data manifold. In addition, GAGA introduces a warping approach, wherein large distances are assigned to points outside the data manifold. The off-manifold status determined using an auxiliary dimension in the latent space, which is trained adversarially.
The defined latent geometry supports various tasks, including (i) uniform sampling of the manifold, (ii) interpolation, and (ii) generation. In our experiments, compare GAGA with FIat on task (ii). Specifically, on-manifold interpolations are achieved by training neural ODE that minimises the curve length connecting two points on the manifold, based on the encoder's pullback metric. In other words, GAGA requires parameterised neural network to define latent interpolations constrained to the manifold.
E.2.2. COMPARISON BETWEEN FLATVI AND GAGA
GAGA and FIatVI share a common objective: To learn a latent geometry amenable to interpolation and manifold exploration. GAGA does so by learning distances on the manifold through local neighbourhood approach, whereas FlatVI performs local manifold regularisation by assuming a metric defined on the statistical manifold spanned by the decoder of a discrete single-cell VAF. Consequently, FlatVI avoids the computation of pairwise distances for metric regularisation, relying on the assumption that aligning the pullback metric with an Euclidean geometry locally induces globally consistent manifold flattening, especially under dense sampling conditions.
In this context, FlatVT is explicitly designed for latent manifold regularisation in high-dimensional, over-dispersed count data, while GAGA assumes a continuous approximation of the gene expression space. Moreover, FlatVI encourages a simple and tractable latent geometry, whereas GAGA combines manifold regularisation with parameterised neural network that approximates the shortest curve between pairs of points.
F. Additional Notes on the Novelty of the Contribution
In VAE-based approaches, the decoder often parameterises a statistical manifold over the space of probability distributions. In our work, we explicitly model decoded single-cell profiles as points on the statistical manifold of negative NB distributions defined by the NB- VAB decoder. To the best of our knowledge, this is the first systematic approach manifold learning in single-cell analysis that leverages the geometry of this specific statistical family.
This perspective is particularly impactful for scRNA-seq tasks that rely on interpolations in a latent or reduced space, such as trajectory inference and cellular fate mapping. Instead of interpolating in an arbitrary Fuclidean latent space, we ask whether these transitions can be aligned with the intrinsic geometry of the NB statistical manifold. Since the NB distribution widely accepted as the most accurate noise model for scRNA-seq data, enforcing such geometric consistency has the potential to produce more biologically meaningful trajectories.
17

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
G. Complexity and Implementation
G.1. Computational Complexity of the FIM Computation
We provide a breakdown the computational complexity involved in evaluating the pullback FIM, as defined in Equa- tion (20).
We begin by analyzing the complexity of Equation (7), which generalizes Equation (20):
M(z) JI. (z) M( (z),
where (z) is the Jacobian of the decoder with respect to the latent variable z.
Let G be the number genes and the latent dimensionalityy, with G > d. The decoded mean parameter E IRG yields one value per gene, and the Jacobian JA (z) E R@xdd
The Fisher information matrix M({)) E RGxG encodes the second-order structure of the negative binomial likelihood with respect to the mean parameters.
The product J&. (z) M((+) costs O(dG²) and yields d x G matrix. Multiplying this with [z) gives final cost of O) Since G >> d, the total complexity is dominated by (ddG)).
An equivalent result is obtained if one directly evaluates the sum of outer product matrices from Equation (20), both forms represent the pullback metric.
Note that this analysis only reflects the cost of evaluating the pullback metric itself. In practice, computing (z) requires evaluating the decoder h, which may be expensive.
G.2. Implementation via the Jacobian- Vector Product
To compute the pullback metric efficiently. we use the Jacobian-vector product (JVP) instead of forming the full Jacobian. JVP is applied along each standard basis vector, dynamically assembling the necessary components while reducing memory and computational overhead. Our models, MLPs with one to three nonlinear layers, remain efficient despite the added cost. Notably, this approach scales better than GAE (Table 7), which requires pairwise distance computations for geodesic estimation.
H. Model Setup
Experimental details for Autoencoder models. The Geodesic AE, NB-VA VAE, and FlatVI models are trained using shallow
2-layer neural networks with hidden dimensions [256, 10]. Batch normalisation is applied between layers, as we observed that it improves reconstruction loss. Non-linearities are introduced using the ELU activation function. All models are trained for 1000 epochs with early stopping based on the VAE loss and patience value of 20 epochs. The default learning rate is set to 1e-3. For VAE-based models, we linearly anneal the KL divergence weight from 0 to over the course of training.
NB-VAE and FlatVI models use batch size of 32, while Geodesic AE is trained with batch size of 256, selected after evaluating [ (64, 100, 256} based on validation loss. Notably, Geodesic AE is trained to reconstruct log normalised counts, in contrast to NB-VAE and FIatVI, which model raw counts via a negative binomial decoder. To ensure training stability, all encoders receive inputs transformed via log(1 + x).
summary of hyperparameter sweeps for FIatVI, along with selected values based on validation loss, shown in Table 3.
Table Hyperparameter sweeps for training FlatVI. The Hidden dims column excludes final latent layer, which fixed at dimension 10. Selected values used for main results are shown in bold.
Batch size
Hidden dims
EB
32, 20,25,,555 512
[1024, 512, 256), [512, 256). [256]
0.001, 0.01, 1. 10
Pancreas
32, 256. 512
11024, 512, 256). [512, 256). [256]
0.001, 0.01, 0.1. 10
MEF
32, 256, 512
[1024, 512, 256], [512, 256], [256]
0.001, 0.01, 0.1, I, 10
18

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Experimental details for OT-CFM. To parameterise the velocity field in OT-CFM, units per layer and SELJ activations. The learning rate is fixed 1 1e-3. The network receives concatenation of the latent vector and scalar time value as imput, as required for conditional velocity estimation. Following recommendations from the official OT-CFM repository', we sample batches that include cells from all time points in each epoch. The variance hyperparameter a is set to 0.1 by default.
a 3-layer MLP with 64 hidden
The choice of the hyperparameter A. The hyperparameter À controls the strength of latent flatness regularisation in FlatVI (Table 1). While increasing À flattens the latent geometry (i.e., reduces curvature and variation of reconstruction), an overly large value can harm data fidelity and lower the model likelihood. A higher leads to more uniform (lower VoR) and flatter (lower CN) latent space.
We select À by evaluating the quality of trajectories inferred by OT-CFM on top of FIatVI embeddings. In practice, we increase À until trajectory reconstruction quality no longer improves. For most real datasets, increasing À from 0.1 to I enhances reconstruction performance. However, further increasing it to 10 offers no additional benefit. For the more complex MEF dataset, performance was best at À = 0.1, so we retained that setting.
I. Evaluation
L.. Metric Description
Condition number. Given metric tensor M(z), let Smin and Smax be its lowest and highest eigenvalues, respectively,
The condition number (CN) is defined as the ratio
Smax
CN(M(2)) =
(34)
Smin
Notably, an identity matrix has a CN equal to 1. The CN is an indicator of the stability of the metric tensor. A well- conditioned metric with a CN close to I suggests that the lengths and angles induced by the metric are stable. A large condition number means that the distances are more stretched in some directions than others. On Euclidean manifold with scaled diagonal metric tensor, distances are preserved in all directions.
Variance of the Riemannian metric. In assessing the Riemannian metric, we introduce key evaluation called the Variance of the Riemannian Metric (VoR) (Pennec et al., 2006). VoR is defined as the mean square distance between the Riemannian metric M(z) and its mean M = Ez-pp [M(z)]. As suggested in Yonghyeon et al. (2021), we compute the VoR employing an affine-invariant Riemannian distance metric d, expressed as:
m
( B) =
Σ
(log (B-1
(35)
where () indicates the 7th eigenvalue of the matrix B- VoR provides insights into how much the Riemannian metric varies spatially across different z values. When VoR is close to zero, it indicates that the metric remains constant throughout. This evaluation procedure focuses solely on the spatial variability of the Riemannian metric and is an essential aspect of assessing the manifolds. Note that the expected value in Equation (35) is estimated using batches of latent observations with size 256.
Velocity Consistency. (Gayoso al., 2024) This metric quantifies the average Pearson correlation between the velocity v(x,) of a reference cell xj and the velocities of its neighbouring cells within the k-nearest-neighhour graph. It is mathematically expressed as:
Σ
corr(u(x;), ux)).
(36)
k
(x))
Here, Cj represents the velocity consistency, i denotes the number of nearest neighbours considered in the k-nearest- neighbour graph, x, the reference cell, (x,)) represents the set of neighbouring cells. The value corr(v(x)). v(x)) is the Pearson correlation between the velocity of the reference cell v(x;) and the velocity of each neighbouring cell v(x). Higher values of ej indicate greater local consistency in velocity across the cell manifold.
hhtpps:/// :/tp:s::///iae thub. www camm//ttoogo//oooodio tional 1- f1 flox-matchiin
19

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
1.2. Experiment Description
Simulation details. We simulate 10-dimensional negative binomial data from three distinct categories parameterised by means following distinct distributions and the same inverse dispersion. The negative binomial means μa are drawn from 10-dimensional Gaussian distributions with category-specd -sregifc- means -1, 0 and 1. The inverse dispersion parameters 0 are again random and drawn from the same distribution across the different classes, namely a Gamma distribution with concentration equal to 2 and rate equal to 1. We exponentiate the means and take the absolute value of inverse dispersions to make them strictly positive. Note that we do not use size factors in the simulation experiment. Overall, we simulate 1000 observations drawn uniformly from different categories.
Estimating latent pullback geodesies to generate Table and Table 4. We evaluate FlatVI trained with varying strengths of flattening regularisation, controlled via the hyperparameter A. To assess the model's reconstruction performance, we consider how accurately FlatVI recovers the mean (4) and inverse dispersion (0) parameters used to simulate individual cells.
As an additional evaluation, we assess how well Euclidean distances in the latent space approximate geodesic distances on the latent manifold. To evaluate this similarity, we proceed in the following way:
We sample pairs of simulated cells, encode them using FlatVI (under different À values) and obtain their latent representations zj and zg.
For each pair, we compute:
1. The Euclidean distance between ž and in latent space.
2. The pullback geodesic distance diatent (Z1,z)) on the statistical manifold M2.
Geodesic distances are computed using the StochMan software?, which approximates the shortest curve connecting two points on a manifold defined by a metric tensor. Specifically, we define the geodesic length via the KL divergence between infinitesimally close conditional distributions, leveraging the link between the Fisher information metric and the KL divergence of nearby decoded points (Fquation (14), Equation (27)). The geodesic distance is thus formulated as:
disteut (31. 22) inf ( KL (p(x (((())),, p(x h(r(t+ + dt)))) dt, 7(1)
(37)
where 7(0) = Z1 7(1) = 72-
Here, h()(()) denotes the decoder output interpolated latent point (1), and p(x (((((())) represents negative binomial likelihood conditioned on those decoded parameters. In practice, parameterise (t) as cubic spline (t) over 100 interpolation steps and optimise it to minimise the objective in Equation (37).
This procedure yields two vectors of pairwise distances (Euclidean and geodesic), which we compare in terms of their induced neighbourhoods and MSE.
Spearman correlation between Euclidean and geodesic distances, and neighbourhood overlap metrics (Table 4). We assess how closely the Euclidean and latent geodesic distances agree in terms of both correlation and local neighbourhood structure. Specifically, for VAEs trained with varying levels of regularisation, perform the following procedure over 10 random repetitions:
In each repetition, we randomly sample 50 simulated observations and encode them into the latent space.
For all pairs of encoded points, compute:
The geodesic distances using the KL-based formulation in Equation (37).
2. The Euclidean distances in the latent space.
Due to the computational cost evaluating Equation (37), we limit the analysis to 50 observations per repetition. Once both pairwise distance matrices are obtained, we compute the following metrics:
Spearman correlation: For each data point, calculate the Spearman rank correlation between its Euclidean and geodesic distances to all other points and report the average across all data points and repetitions.
Neighbourhood overlap: Using the Euclidean and geodesic distance matrices, we extract the neneersst neighbours (k: = 3 and l 5) for each point and compute the average proportion of neighbours shared across both metrics. A
Antpps:/ //g thub. com Mh net.anni feSci wew.eiiencee/tttcccaaaeeeeeeeeee
bocccck-------- random- reee-rr-g-gmeeetrr
20

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
higher overlap indicates greater agreement between the local structures induced by Euclidean and pullback geodesic distances.
Euclidean matrix distance: To quantify the overall similarity between the geodesic and Euclidean distances, we compute the MSE between the two full pairwise distance matrices. Lower values indicate better global agreement between the metrics.
Riemannian metrics and path visualisation. Riemannian metrics only take the metric tensor at individual latent codes as input (see Appendix I.I.. The metric tensor for individual observations is calculated following Equation (1I). The geodesic paths in Figure 2c are again approximated by cubic splines optimising the task in Equation (37). We compute such paths for 10 randomly drawn pairs in the region exhibiting high VoR in the unregularised VAE model for demonstration purposes.
Trajectory reconstruction experiments. In Table 2. we explore the performance of OT on different embeddings based on the reconstruction of held-out time points. For the EB dataset, we evaluate the leave-out performance on all intermediate time points. Conversely, on the MEF reprogramming dataset (Schiebinger et al., 2019) we conduct our evaluation holding out time points 2, 5, 10, 15, 20, 25 and 30 to limit the computational burden of the experiment. Note that the Pancreas dataset used for Figure 3 could not be used for this analysis, since it only has two time points: An initial and a terminal one.
After training OT-CFM excluding the hold-out time point t, we collect the latent representations of cells at t --tt-1 and simulate their trajectory until time i, where we compare the generated cells with the ground truth via distribution matching metrics both in the latent and in the decoded space (mean I,2 and 2- Wasserstein distance the latent space, 2- Wasserstein distance and MMD in the decoded space). Note that replaced the L,2 distance with the MMD in the decoded space, as the former behaved uniformly across models in higher dimensions. For the latent reconstruction quantification, generated latent cell distributions at time t are standardised with the mean and standard deviation of the latent codes of real cells at time to make the results comparable across embedding models. Results in Table 2 are averaged across three seeds.
Fate mapping with CellRank. We first train representation learning models on the Pancreas dataset. Then, following the setting proposed by Eyring al. (2022), we leam velocity field over the latent representations of cells by matching time points 14.5 15.5 with OT-CFM and input the velocities to CellRank (Lange et al., 2022: Weiler et al., 2024). Using the function g. compute macrostates (n states, cluster key) of the GPPCA estimator for macrostate identification (Reuter et al., 2019), we look for 10 to 20 macrostates. If OT-CFM cannot find one of the 6 terminal states within 20 macrostates for certain representation, we mark the terminal state as missed (see Figure 3). Terminal states are computed with the function compute_tert na sstates hmethod, n_states). Velocity consistency is estimated using the se Velo package (Bergen et al., 2020) through the function scv.tl tl. www..eeoottt.coooiiiooooo fidence. The value is then averaged across cells.
Latent interpolation comparison with GAGA. We train both GAGA and FlatVI on the pancreas dataset. For FIat we use the same parameters in Table 3. For GAGA, choose hyperparameters to reflect the FlatVI setting, but make the neural network deeper, as it produced better results. Below we list the final array of hyperparameters used to train the GAGA model:
batch.size: 256.
latent.space.dim: 10.
act.ivation.funetiooo ReLU.
learning.rate: 0.001.
dist_reconstr.weighhe [0.9, 0.1, 0.0).
Once we optimise the models, we draw 100 pairs of multipotent cells and mature cells (a batch for each lineage). Pairing is done randomly. Then, we encode both multipotent and mature cells using both FlatVI and GAGA and interpolate the latent spaces of paired cells. For FIat we interpolate linearly. For GAGA, we perform geodesic interpolations by training an energy-minimistie neural ODE as explained in Sun et al. (2025) and Appendix E.2 over 200 epochs with a learning rate of 0.01. Intermediate results for both models are decoded with the respective decoder network, and the predicted expression is compared with the real gene expression distribution of dineage-specifer markers.
21

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
J. Data
J.1. Data Description
Embryoid Body (EB). Moon et al. (2019) measured the expression of 18,203 diffcrentiating human embryoid single cells across time points. From an initial population of stem cells, approximately four lineages emerged, including Neural Crest, Mesoderm, Neuroectoderm and Endoderm cells. Here, resort reduced feature space of 1241 highly variable genes. OT has been readily applied the embryoid body datasets in multiple scenarios (Tong al., 2020; 2024a), making it solid benchmark for time-resolved single-cell trajectory inference. The data is split into 80% training and 20% test sets.
Pancreatic Endocrinogenesis (Pancreas). We consider 16,206 cells from Bastidas- Ponce et al. (2019) measured across 2 time points corresponding to embryonic days 14.5 and 13.5. In the dataset, multipotent cells differentiate branching into endocrine and non-endocrine lineages until reaching 6 terminal states. Challenges concerning such dataset include bifurcation and unbalancedness of cell state distributions across time points (Eyring et al., 2022). The data is split into 80% training and 20% test sets.
Reprogramming Dataset (MEF). We consider the dataset introduced in Schiebinger et al. (2019), which studies the reprogramming of Mouse Embryonic Fibroblasts (MEF) into induced Pluripotent Stem Cells (iPSC). The dataset consists of 165,892 cells profiled across 39 time points and 7 cell states. For this dataset, we keep 1479 highly variable genes. Due to its number of cells, such a dataset is the most complicated to model among the considered. The data was split into 80% training and 20% test sets.
J.2. Data Preprocessing
We use the Scanpy (Wolf et al., 2018) package for single-cell data preprocessing. The general pipeline involves normalisa- log-transforrmaion via sc. pp- loglp and highly-variable gene selection using se. sc.pp. highly gy variable genes. 50-dimensional embeddings are then computed via PCA through sc pp pea. We then the PCA representation compute the 30-nearest-neighhouur graphs around single observations and use them for learning 2D UMAP embeddings of the data. For the latter steps, employ the Scanpy functions sc. sc.pp.... pp..pp.neeee neighbuurs and sc. tl umap. Raw counts are preserved in adata. layers "x_counts) to train FlatVI.
tion via sc -P www.onooo.ettto..ooooooo
J.3. Details about Computational Resources
Our model is implemented in Python 3.10, and for deep learning models, used Py Torch 2.0. For the implementation of neural-ODE-based simulations, we use the torchdyn package. Our experiments ran on different GPU servers with varying specifications: GPU: 16x Tesla V100 GPUs (32GB RAM per card) / GPU: 2x Tesla V100 GPUs (16GB RAM per card) / GPU: 8x A A100-SXM4 GPUs (40GB RAM per card).
22

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
K. Algorithms
Train FI Below, we provide a training algorithm for FlatVI.
Algorithm Train FlatVI
Require: Data matrix E No NG batch size B, maximum iterations Rmma encoder f decoder hss flatness loss scale
Ensure: Trained encoder fø, decoder he and inverse dispersion parameter 0
Randomly initialize gene-wise inverse dispersion f
Randomly initialize the identity matrix scale
trainable parameter
for i 1 to Rmax do Sample batch X < -x1, XB) from X
1 +- ccompute_sieee wic.sooorrrf.c
Z - f+(1+ log x*)
μ +- (Zb, 14)
CKL f compute
Lrecon & compute.nb.1kkeehhoooxx.
M(Z)) + Equation (20)
Liat MSE(M al.)
C = Lrcon - CKL + ACaat
Update parameters via gradient descent
end for
Train OT-CFM on latent representations. For time-resolved scRNA-seq data, cells are collected in unpaired distribu- tions {u. Individual time points correspond to separate snapshot datasets {X,)?= each with N, observations. Every snapshot mapped tuples {(Z,, 1)) of latent representations Ze Nxdd and size factors Nç following the setting described in Section 3.1. We wish to learn the dynamics of the system through a parameterised function in the latent space 2 of a VAE, taking advantage of its continuity and lower dimensionality properties.
Algorithm 2 Train latent OT-CFM with FlatVI
Require: Datasets
{X,]==>> variance a, batch size B, initial velocity function Uy. maximum iterations Tamax trained
encoder S
Ensure: Trained velocity function U/
(Zr) + f+ log oo)
(l,)20 < hompptt...o..szee.facaee..
(S <- t-imewiss concatenate (Z (4 1T ()
for i 1 to Timax do
Initialize empty aray of velocity predictions V Initialise empty array of velocity ground truth U for 4raj = to
do
Randomly sample batches with B observations S St s
Ferui S +1
(Sh Keaa Sh
Fryn 1ron++1) +1)
((,
<N(1
S) b
t
+1. Ie)
Append (t + heraj, 2t) to
Append (St 1imjt11 8b
to
end for
2 Update parameters via gradient descent
LOT-CFM IV
end for
23

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Size factor treatment. Since the size factors 1, required for decoding, are observed variables derived from single-cell counts in the dataset, their values not available when simulating novel cell trajectories from t 0, hindering the use of the decoder to recover individual gene evolution. Assuming that the size factor is real number and related to the cell state, we include log 4 in the latent dynamics and infer its trajectory together with the latent state representation Z,. The log taken for training stability. Therefore, we use OT-CFM to learn a velocity field vyg [0, 1] x Rd+1 + R4+1 on the concatenated state S; [z, log4]. The time-resolved vector field 0g is modelled by matching subsequent pairs cell distributions.
Gene expression trajectories. If the latent space Z can be mapped to the parameter manifold H. trajectories in Z correspond to walks across the continuous parameter space via the stochastic decoder h. The temporal trajectory of the likelihood parameter vector μr is given by
Me h 80 f/' ug(t', sv) dt')
(38)
where M4 h(s)) and express the decoder function hizz, 42) from Equation (2) h(sa) for simplicity. Then, assuming gene-wise constant inverse dispersion 0g. discrete trajectories of gene counts follow the noise model x, ~ NB(,,,
24

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
L. Additional Results
L.1. Simulation Data
L.I.I. SIMULATED DATA VISUALISATION
PCA representation of simu ated data
Figure The dimensionality reduction plot of the simulated data coloured by category.
1.2. COMPARISON WITH EUCLIDEAN SPACE METRICS
In Figure 7 provide an extension Figure where we add how the VoR, CN and geodesic paths should appear in the Euclidean space. Notably, both CN and VoR are uniform, equating to 1 and O, respectively. Furthermore, while some points in the simulated data exhibit high values for VoR and CN, the representation from FlatVI is more compatible with the expected one under Euclidean geometry than a normal NB- BAE. Additionally, path straightness is better preserved in FlatVT's latent geodesics compared to the counterpart, validating the purpose of our model.
VoR NB-VAE
CN NB-VAE
Geodesics VA-VAE
80
Latent
Latent
Latent 2
20
40
20
Latent
Latent
Latent 1
FlatVI X=7
CN FlatVi
Geodesics FlatVI
50
80
40
Latent
30
Latent
Latent
20
40
10
Latent
Latent
Latent 1
VoR Euclidean space,
to
CN Euclidean space
Euclidean space
Latent 2
Latent 2
Latent
Latent 1
Latent
Latent
Figure 7. Comparison between the latent geometries of the NB-VAE (top row), FlatVI trained with 7 (middle row) and example
Euclidean space of 1000 points evaluated in terms of (a) Variance of the Riemannian metric (VoR), (b) Condition Number (CN) and (c) straightness of the geodesic paths connecting pairs of latent points. The Euclidean panels are simulated uniformly sampled points on regular grid.
25

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
L.3. ADDITIONAL METRICS
In addition to the metrics reported in Table 1, in Table 4 we provide further results that justify the usage of our flattening loss in the simulated data setting. The metries are described in Appendix I.2. Such values represent how much the latent pullback geodesic and Euclidean distances and the neighbourhood structures deriving from them correspond. As can be inferred from the table, adding the regularisation promotes an overall improvement of the metrics. While Spearman correlation does not drastically separate results, adding the flattening regularisation marks an improvement in neighbourhood preservation metrics of up to 18% when considering 5 neighbours and 14% when using 3-point neighbourhoods. In other words, the pullback geodesics are better reflected by the Euclidean distance when applying our regularisation, which provides evidence of the working principles of our flattening loss.
Table Comparison between FlatVI and the unregularised NB- VAE (A = 0). Spcarman (Geo-Euc) represents the Spearman correlation between latent Euclidean and pullback geodesic distances. Neighbourhood metrics measure the proportion of shared nearest neighbours (5-NN and 3-NN) across distance types. MSE (Geo-Euc) quantifies the absolute discrepancy between Euclidean and geodesic distances
across 1000 pairs.
A
Spearman (Geo-Euc) (1)
5-NN overlap (1)
3-NN overlap (1)
MSE (Geo-Euc) (+)
A 0
0.94
0.50 0.01
0.66 0.00
47.75 + 2.80
0.95 0.00
0.57 +.00
0.63
46.34 = 6.45
À = 3
0.96 0.01
0,68 0.03
0.77 10000
16.74 a
0.97 a 0.01
0.58 1 0.01
0.67 2001
8.65 1.68
À = 7
0.97 0.01
0.60 1.01
0.72 2001
8.02 0.67
A 10
0.97
0.68
0.80 0.03
11.80 1.04
L.1.4. ANALYSTS OF SUB-OPTIMALLY FLATTENED REGIONS
In Figure 2, we show that introducing our flattening loss component in the VAE model training ensures lower and more uniform Riemannian metric throughout the space, as well as lower CN. However, some points of our simulation dataset still display high decoding distortion and variance in the Riemannian metric as signals of insufficient flattening. In Figure 8, we show that latent paths between points of high VoR and their neighbours do display some curvature, indicating regions the manifold with sub-optimal Fuclideanisation. We investigated what causes points to exhibit a high VoR in both FlatVT (A=7) and the regular NB-VAE. First, we found that regions of high VoR and CN in FIatVI tend to overlap, while in NB-VAE they less correlated (see Figure 9a). Hence, while for most of the observations, flattening applies, the pullback metric from the decoder violates uniformity and preservation of angles and distances in some portion of the space.
We check the label annotation of the insufficiently flattened regions by overlaying their VoR and CN values onto the UMAP plot of the real data (see Figure 9b-c). By this analysis, we note that the high VoR and CN data points are concentrated at the inter-class boundaries FlatVT's latent space (see Figure 10). In other words, observations in regions of the manifold enriched by different classes representing state transitions are more likely to fail flatten. The fact that high VoR and CN are concentrated in regions of class heterogeneity may suggest that FlatVI fails unfold some fast-changing manifold areas at the intersection between classes, and the decoder needs to violate the correspondence between the Euclidean latent space and the statistical manifold to ensure proper reconstruction.
26

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
-6.73
-10
-1.00
-2.4
VoR = 51.96
-149
-12
-114
0.1
65
34
5430
145
05
VoR 51.63
35
-30
34
94
5.0 3.0 M0
-05
150
-52
15
-0.500
-2
0.300
-255
-35
-0.
00
022
VoR - 48.10.....0
32-
-10
01
62
955
18
18
90
05
0.0
05
-2.00
-30
VoR 40.03
-19
VoR 41,23 -25
-63
D%
-12
-12
Figure 8. Geodesic paths between regions with high VoR and CN in the FlatVT emheddings with Every row represents point with high VoR and CN. The columns are five randomly sampled neighbours in the dataset. Every plot represents the geodesic path between the point with high VoR and CN and the associated neighbour. representation of high VoR and CN, we select points with the VoR value larger than 30.
300
250
r- 0.21
0.58
250
700
200
CN FaMi
200
50
100
250
20
40
VoR Flatvi
CN NB VAEE
CN Flat X-7
150
Label
Voa TlacV Am
20
Figure (a) The scatter plot and correlation values between VoR and CN in the embeddings computed by NB-VAE and FlatVI. (b-c) The UMAP plots of the real data coloured by class and CN and VoR values from the NB- VAE and FlatVI embeddings.
27

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
CN FlatVi class 0
FlatVl class
CN FlatVI class 2
18
UMAP2
16
UMAP2
UMAP22
14
14
12
12
10
17.5
20.0
22.5
15.0
17.3
20.0
22.5
15.0
2255
UMAP1
MMAP1
UMAP1
Figure 10. The UMAP plots computed the simulated count data coloured by the CN from the FlatVI embeddings. Highlighted regions with high VoR and CN. External boxes represent the label compositions of the highlighted regions (for a label-specific colour legend, Figure 9b).
L.2. Velocity Estimation
We compared our whole pipeline involving the combination FlatVI and OT-CFM with the seVelo (Bergen et al., 2020) and veloVI (Gayoso al., 2024) models for RNA velocity analysis. Figure 11 and Table show examples how our approach favourably compares with velocity estimation methods, namely inferring a proper velocity field in Acinar cells and detecting all terminal states with high velocity consistency in the Pancreas dataset.
Pancress Velocity field Acinar branch
FlatVOOT-OFM
Figure 1.. Comparison between the vector field learnt on the Acinar branch of the Panereas dataset by using OT-CFM in combination with VI and standard RNA velocity algorithms.
Table Number of terminal states computed by CellRank and velocity consistency using the representations and velocities learnt by Fla+VI+TT-CM and standard RNA velocity algorithms.
Method
Terminal states
Consistency
FlatVI + OT-CEM
0.94
veloVI
0.92
sc Velo
4
0.80
28

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
L.3. Trajectory Visualisation of Real Data
MD VAE CFM
Ane Predicted
OT-CFM
Brarch Neural crest Cardicc
Faai + OT-GFM
WD 161
Cardiac
HANDI
OYPC
Roal
Flatdi OT-CFM
46% unespressed 45% Lnnpprmannd
155 unscpressed
92% unecressed 10% unexpressed
Figure Prediction scRNA-seq time. (a) Overlap between real and simulated latent samples the dataset. WD indicates the 2- 2-Wasserstein distance between real and generated latent cell representations. (b) 2D UMAP plots of real and predicted cell counts from the cardiac and neural crest lineages of the EB dataset, comparing FlatVI GAE as representations for OT-CFM. Colours indicate the predicted log gene expression of the reported lineage drivers GYPC and HAND1. Under each UMAP plot, calculate the percentage of unexpressed marker instances along the trajectory.
L.4. Additional Tables
Table Univariate 2-Wasserstein distance between simulated and real marker gene expression for different lineage branches of the EB dataset models. A lower value indicates that the model better approximates marker gene trajectories along the branch.
Wasserstein real-simulated markers
Cardiac
Neural Crest
Endoderm
Neuronal
GATA6
HAND1
TNN2
NGFR
GYPC
PDGFRB
SOX17
GATA3
CDX2
LMXIA
ISL1
CXCR4
GAF NB-VAE
0.17
0.28
0.23
0.05
0.24
0.07
0.24
0.14
0.11
0.04
0.04
0.18
0.03
0.24
0.07
0.07
0.09
0.04
0.08
0.07
0.02
0.02
0.06
0.07
FlaiVI
0.02
0.09
0.03
0.02
0.05
0.03
0.08
0.02
0.02
0.01
0.02
0.03
Table 7. Runtime, in seconds, evaluated over a single forward pass considering diffcrent batch sizes for cach compared model. The runtime is tested over 10 repetitions using random inputs with 2k dimensions.
Runtime (s)
Batch size
GAE
NB VAE
FlatVI
8
0.119
0.000
0.009
16
0.095 0.012
0,001 0.0000
0.007 0.002
32
0.115=00r
0,001 0.000
0,004 0.003
64
0.099
0.001
0.005
128
0.116 = 0.077
0.002 0.000
0.009 0.004
256
0.218 0.177
0,002 0.000
0,009 0.002
512
0.513 0.021
0.002
0.015 0.004
1024
1.522 0013
0.003 0.000
0.015 0.001
Table 8. Separation between initial and terminal lineage states evaluated in terms of clustering metrics in the latent spaces of the distinct models, Different representation spaces are compared on how well they unroll developmental trajectories.
Silhouette Score (f)
Calinski- Calinski-Harabasz (f)
Bouils- (1)
EB
Pancreas
MEF
EB
Pancrcas
MEF
EB
Pancreas
MEF
GAE
0.28
0.15
0.09
1608.56
1723.13
11232.84
1.03
1,50
2.99
NB-VAE
0.19
0.26
0.21
940.87
2191.48
19440.38
1.28
1.56
2.35
0.21
0.50
0.31
983.41
6986.31
45372.75
1.18
0.73
29

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
L.5. Additional comparison with GAGA
Lineage: Multipotent Alpha/ Marken: Ope
Lineage: Multiputent Alpha Merker: Gness
Lineage: Mullipotent Alphte Marker: Gpe
GAGA
GAGA
GAGA
Multipetent
Ngra EP
8pr33 High exrty
Ngn3 Hight lake
Fev Alpha
Alpha
Multipetent
agna High early
Ngnz High late
Pev Alpha
Alpha
Multipotent
Hgn3 low EP
Nona High earl
Nomz High lare
Peve Alpha
Alpha
Hpna low
Lineage
Unesge
Lineage
Lineage: Mullipotent Alphe/ / Marker:: PcskIn
Lineage: Multipotant a Alpha/ Marker: Pcsk2
Lineage: Multipotent - Alpha/ Marker: Scg3
FlatVi
Flatvi
FlatVi
GAGA
GAGA
GAGA
Multipotent
Bpnd High early
Ngn3 High late
Pev Alpha
Alpha
Multspatent
Ngra tow EP
agn3 High eanty
Ngrz High tate
Few Alpha
Alpha
Multipotent
Ngn low EP
Agn3 Hight early
Npna High late
Fast Alpha
Aipha
Ngn3 low
Lineage-
Lineage
Unnege
Lineage: Multipocent
Lineage: Multipotent
Lineage:
Poplrla
FlatVi
Flatvi
FlatVi
GAGA
2
Multipatent
Ngr3 EP
Ngr33 High earl
Ngn3 High late
Far+ Beta
Multipstntt
Noră lowr EP
Ngn3 High early
Hon3 High tate
Fev Beta
Bela
Multipatnt
age3 EP
Hgn3 High early
Ngnă High late
Fev Beta
Beta
Lineage
Lineage
Lineage
Lineage: Multipotent - Delta/ Marker: isll
Lineage: Multipotent - Delta Marker: Raplb
Lineage: Multipotent Delta Marker: Stt
FlatVi
Flatil
FlatVi
Expression
GAGAA
GAGA
0
Multipotent
3gg High earty
Nor3 High late
Fairt Data
Delta
Muttipotent
Ngróà High early
Hona High late
Delta
Multipatent
Ngn3 low ER
Nor' High early
Agni High late
Fest Delta
Dota
Ngn3
Non3
For+
Lineage-..
Lineage
Lineage
Lineage: Multipctent Epsiion
Lineage: Multipotent Epsilon / Marker:
Lineage: Multipotent - Epsilon/ Marker: Rap1b
FlatVi
4-
GAGA
GAGA
GAGA
2
Multipstett
Non3 High early
Agm3 High Inte
For+ Epoilon
Epaton
Multiestett
Ngn3 High early
Non3 High Late
Fev Epuionn
Epoioon
Multieohent eett
Ngn3 High early
Ngn3 High late
Fov Ejsiton
Epsiton
Ngn3 low
Non3 low
Won' low
Lineage-
Lineage-
Lineage-
Figure 13. Decoded marker trajectories from latent interpolations on the pancreas dataset across multiple lineages, comparing FlatVI and GAGA. We draw 100 pairs of multipotent and mature cells from terminal encode their gene expression and perform latent interpolation (lincarly for VIatVI and with a neural ODE in GAGA). The intermodiate states are decoded, and their marker gene expression visualised together with the real marker expression distribution along the lincage (dark grey points). The individual trajectorics are represented as low-opacity lines, while the centre solid line is the mean trajectory.
30

---

Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation
Interpolation Time Comparison
Execution Time (s)
10a
100
10-1
GAGA FlatVI
10
10a
102
10
Number of Interpolated Samples
Figure 14. Scaling behaviour function of different numbers of interpolation samples for GAGA and FlatVI. On the y-axis, the execution time seconds is presented on log scale.
31


Reviews:
[{"review_id": "x9uFGXYy7P", "reviewer": "Reviewer_NJwg", "summary": "The paper proposes FlatVI, a novel training framework for variational autoencoders (VAEs) applied to single-cell RNA-seq (scRNA-seq) data. Its central goal is to enforce Euclidean geometry in the latent space of discrete-likelihood VAEs—specifically, negative binomial VAEs commonly used for modeling gene expression counts—so that straight lines in the latent space correspond more closely to geodesics on the decoded statistical manifold.\n\nThe authors develop a flattening loss based on the pullback of the Fisher Information Metric (FIM) from the decoder, which they regularize toward a scaled identity matrix. The method is designed to retain data reconstruction performance while improving manifold geometry.\n\nEmpirical evaluations include:\n\n- Simulations showing improved alignment between geodesic and Euclidean distances in latent space.\n- Application with OT-CFM for scRNA-seq cellular trajectory reconstruction\n- Application with CellRank for latent vector field and lineage mapping\n- Some qualitative PCA-2D plots\n\n## update after rebuttal\n\nI appreciate the authors’ rebuttal, and the rebuttal Figure 1 serves as a new case study to show that the changes of three genes can be better estimated by FlatVI on the pancreas dataset. Therefore, I raised my score by 1 from broadline reject to broadline accept.\n\nNonetheless, I wish to clarify my stance (1) I agree this is a very interesting work extending VAE with a regularization to enhance euclidean metric (2) I think table 2 servers as a good benchmark to show the usefulness of this method on two real dataset (3) I am okay to use qualitative results in 5.4 to show the goodness of FlatVI representation at the end of main text to further enhance the work.\n\nThe main shortcoming from my perspective is the limited heuristic results for Section 5.3. Given that the authors only give one OT method as the basis for the analysis. Would the velocity consistency be enhanced at different levels of complexity of different datasets (maybe some datasets with some simple induced differentiation process and some complex organism developmental dataset like Mouse or C. elegans)? I feel like Figure 3 is a case study, but not comprehensive enough as the evaluation of a task. To this end, if the full section 5 consists of only one simulation, one task evaluation with 2 datasets, one case study, and one qualitative representation. I feel like the overall workload is under my expectations. \n\nMoreover, I think rebuttal Figure 1 is a very nice addition to the results section. But again, to demonstrate the advantage of a method, we may consider a more systematic shows the usefulness among all the genes (or a subset of genes that are identified to be time-variant by a set of time analysis methods) in the dataset instead of curating three genes with better fitted trends.\n\nIn the end, I raised my score by one based on the rebuttal and the comments from zkLA and 5Fan. But I think ideally, Figure 3 and rebuttal Figure 1 should be expanded into two task evaluations to not only give a comprehensive heuristic evaluation but also better define this task to facilitate future work to build upon.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We thank NJwg for thoroughly reviewing our paper and providing constructive criticism to improve our submission. We kindly point the reviewer to the anonymous link https://figshare.com/s/74ca822781c60b2a85f2 containing the figures we refer to in some of our answers.\n\n> A1. Limited use/comparison with other algorithms.\n\nWhile our specific examples focus on OT-CFM, the benefits of FlatVI extend to a broader range of single-cell analysis techniques.\n\n* Many of the trajectory inference methods mentioned by the reviewer still fundamentally use distance calculations between representations of cells to define relationships and biological state proximity. For example, algorithms that build k-NN graphs (including pseudotime and diffusion) require a distance metric, and the Euclidean distance is often the default or a common choice. \n* By better aligning latent cell state Euclidean distances with biological distances in discrete gene expression counts (see our simulations), FlatVI's representation offers stronger guarantees for k-NN graph construction, which can be used by diffusion-based methods to assign cellular ordering and infer biological processes.\n*  In other words, FlatVI provides a representation space where biological distances are easy to compute and suitable for building temporal graphs and inferring cell ordering. \n*  The latent Euclidean assumption is also key in batch integration [1] and perturbation prediction [2], broadening the use of our regularisation.\n*  In Rebuttal Fig.1b and 2b-c (addressing oZ7a), we extend our comparisons with GAGA [5], a geodesic autoencoder, on the additional task of gene trajectory reconstruction across lineages. Here, we show that our simple Euclideanisation approach is faster, more stable, and more accurate than the most recent geodesic interpolation approach in capturing the dynamics of gene features on cellular manifolds.\n\n\n> A2. Velocity consistency\n\nImportantly, we evaluate our model with ground truth terminal state labels in Fig. 3a, showing that FlatVI’s representations guide Markov paths computed by CellRank to the correct terminal states, unlike competing methods. Consistency, introduced by VeloVI [3], measures vector field smoothness in a cellular representation space. Higher consistency means nearby latent profiles have correlated velocities, resulting in a smoother vector field. When estimating trajectories, related cellular states should have similar gradients, with gradual changes in vector field direction along the manifold. Regularising geometry before vector field estimation, as suggested in [3], enables smoother state transitions with Euclidean latent geometry.\n\n> A3. Contribution \n\nOur work presents a novel methodological formulation that, to the best of our knowledge, has not been explored before:  \n\n- **Manifold learning:** Most manifold learning methods rely on continuous approximations of single-cell data. We are the first to define the cellular manifold as the negative binomial manifold spanned by the decoder of a discrete VAE. This means our geometric regularisation is directly informed by the statistical properties of single-cell data rather than assuming an arbitrary continuous structure.  Our simple approach leads to more biologically meaningful interpolations (Rebuttal figures).\n- **Pullback-based regularisation:** Unlike most existing methods that rely on k-NN neighbourhood estimation, we introduce pullback-based local geometric regularisation from the decoder of a VAE. This is a new approach to single-cell machine learning.  \n- **Geometry-aware VAEs:** Our work extends existing geometry-aware VAEs by connecting latent Euclidean interpolations with geodesic paths on discrete statistical manifolds with tractable likelihoods. Similar geometric regularisation techniques have primarily been explored in deterministic autoencoders, but not in the probabilistic single-cell setting.  \n- **Beyond trajectory inference:** While we focus on trajectory inference, the improved correspondence between latent Euclidean distances and statistical manifold geodesics has broader implications (see our first answer).\n\nFinally, we highlight that existing standalone ML conference papers in single-cell geometry [4, 5] often introduce new simulation-driven insights while targeting specific downstream tasks. FlatVI combines novel representation approaches with modelling assumptions in single-cell analysis. \n\n[1] Lücken et al., Benchmarking atlas-level data integration in single-cell genomics\n\n[2] Eyring et al., Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation\n\n[3] Gayoso et al., Deep generative modelling of transcriptional dynamics for RNA velocity analysis in single cells\n\n[4] Huguet et al., Manifold Interpolating Optimal-Transport Flows for Trajectory Inference\n\n[5] Sun et al., Geometry-Aware Generative Autoencoders for Warped Riemannian Metric Learning and Generative Modeling on Data Manifolds", "claims_and_evidence": "The paper makes three central claims:\n1. Latent geometry regularization via flattening loss\n2. Improved trajectory inference when combined with OT-CFM\n3. Improved velocity field consistency in the latent space\n\nI find the evidence for Claim 1 to be strong and clearly demonstrated through simulation experiments. The use of the pullback metric and flattening loss is well-motivated and effectively shown to induce Euclidean-like geometry in the latent space.\n\nClaim 2 and Claim 3 are supported by results in Table 2 and Figure 3, respectively. However, I believe these claims highlight improvements that are largely specific to a narrow set of use cases. The improvements shown for trajectory inference rely on OT-CFM, which assumes Euclidean latent geometry. Many commonly used trajectory inference methods (e.g., pseudotime algorithms, diffusion-based or graph-based methods) operate in non-Euclidean or PCA-based spaces, where the benefits of FlatVI may not hold or may be less impactful. As such, the broader significance of the method is unclear when considering the wide variety of available alternatives in the single-cell analysis space.\n\nFor Claim 3, while Figure 3(b) reports improvements in “velocity consistency,” this metric reflects self-consistency of the learned vector field rather than ground-truth accuracy. It would strengthen the paper if the authors provided a clearer discussion about the biological interpretability and utility of velocity consistency in practical single-cell analyses.\n\nOverall, I appreciate the authors’ effort and find the core idea to be well-presented. The work is technically sound and the writing is clear. That said, the methodological contribution is somewhat incremental, and the scope of demonstrated utility is limited to scenarios where specific assumptions (e.g., Euclidean latent geometry) hold. From my perspective, this work could be a useful enhancement to tools like scVI or CellRank, but may fall short of the bar for a standalone ICML paper in its current form.", "methods_and_evaluation_criteria": "I am in general agree with the metric and criteria. Moreover, I think Table 2 is a good choice of benchmark showing OT-CFM got improved by FlatVI. It would be a good reference for later works to do trajectory reconstruction task.", "theoretical_claims": "The proofs for theoretical claims look good to me.", "experimental_designs_or_analyses": "As discussed in Claims And Evidence.", "relation_to_broader_scientific_literature": ".", "essential_references_not_discussed": ".", "other_strengths_and_weaknesses": ".", "other_comments_or_suggestions": ".", "questions_for_authors": "- I would love to hear the authors thoughts about the significance of the FlatVI in trajectory inference task considering those alternative methods. \n- I would love to the authors explain/discuss a bit more about their thoughts about the usefulness of velocity consistency in single cell analysis.", "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": "", "extra_scores": {}}, {"review_id": "NW79MdOv87", "reviewer": "Reviewer_oZ7a", "summary": "The paper proposes FlatVI, a VAE training strategy that enforces local Euclidean geometry in the encoder’s latent space by regularizing the pullback metric of the decoder to approximate local Euclidean geometry in latent space. It validated the proposed method on simulated data and the single-cell trajectory inference task. It also demonstrated two other use cases of the locally Euclidean geometry in latent space, including latent vector fields and single-cell latent space visualization.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "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": "We sincerely thank oZ7a for their elaborate review. Their informed criticism offers us an opportunity to improve our submission. We will refer to two additional rebuttal figures stored at https://figshare.com/s/74ca822781c60b2a85f2. \n\n> Global regularisation\n\nSimilar to GAE, we propose to:\n- Sample a couple of cells and encode them.\n- Approximate the latent geodesic distance between representations using the pullback. One could fit an energy-minimising cubic spline weighing local shifts by the pullback.\n- Minimise the difference between the latent Euclidean and the geodesic distances. \n\nThis approach is expensive because it contains an unstable inner optimisation loop requiring decoding and gradient computation. \n\n> Local vs global\n\nWhile our flattening loss enforces local Euclidean geometry in the latent manifold, we emphasise that this regularisation is applied specifically in the regions where the data lies. Under the assumption that the latent manifold is sufficiently densely sampled, these local constraints can collectively approximate a globally Euclidean structure for the data-supported manifold. This assumption aligns with common practices in single-cell manifold estimation, where smooth transitions between neighbouring states define the manifold structure. We will add a clarifying assumption statement. \n\n> Contextualisation\n\nApologies for the missing citations; we will add them upon revision.\n\n* We contextualize our model in App. E (ref. in l.239, col.2). Methodologically, FlatVI is a representation learning model, creating a space where tractable interpolations map to meaningful data-space paths. In contrast, MFM is an interpolation model, operating in a fixed PCA space and focusing on FM regularisation to align with a neighbourhood-based manifold. Unlike FlatVI, MFM does not account for probabilistic aspects of high-dimensional single-cell data.\n\n* More related are GAGA and NeuralFIM. Both methods perform latent geodesic regularisation to reflect neighbourhood structures in single-cell datasets, using k-NN graph-based approaches. Different from FlatVI, both models:\n\t* Only consider low-dimensional and continuous input data (e.g., PCA). \n\t* Do not enforce a tractable latent manifold, hence, they have to learn geodesics between latent observations with a NeuralODE.\n\nThese aspects may hinder scaling to the gene space, as sparse expression vectors are unsuitable for Euclidean distances, and geodesic NeuralODEs are unstable (see below).\n\n> Manifold assumptions \n\nThere is a distinction between the goals of the cited studies and our paper. For instance, PHATE produces a geometrically sound embedding of the data, exploiting neighbouring structures to **discover** patterns in the dataset. We implement a strategy that connects distances on a latent manifold to the geometry of high-dimensional and noisy discrete data to replace the standard VAE logic in combination with distance-based tools. FlatVI operates through the decoder function. When the decoder is expressive and assuming geodesic convexity, it reconstructs the data manifold's geometry and aligns our latent structure to the statistical manifold. Interpolating simpler proxy manifolds is established [1]. \n\n> FlatVI's benefits\n\nKindly refer to A1. to NJwg and A7. to zkLA.\n\n> Benchmark\n\nWe focus on FlatVI’s representation aspect and how our guarantees improve downstream tasks. Our baselines, NB-VAE and GAE, reflect this. Comparing to NB-VAE serves as ablation to highlight our regularisation’s effects. GAE, like FlatVI, enforces Euclidean geometry via a k-NN-based approach, making the comparison fair.\n\nWe validate our claims above by comparing FlatVI with GAGA (concurrent work under ICML guidelines). Both methods encode single-cell data into a latent space, where we interpolate between Multipotent and mature cells—using geodesic interpolations for GAGA and linear for FlatVI. Decoding the results, we assess marker gene reconstruction along lineages (Reb. Fig. 1a). FlatVI’s interpolations align better with true gene trends, whereas GAGA shows unnatural oscillations (Reb. Fig. 1a). GAGA’s simulations are also more unstable (Reb. Fig. 2a) and slower (Reb. Fig. 2b). Due to memory issues, we couldn’t run geodesic interpolations with NeuralFIM, but its training was slower than FlatVI (Reb. Fig. 2c).\n \nSince the FM algorithm is not the novelty of our work, we did not compare it with MFM.\n\n> Consistency\n\nKindly refer to A2 to NJwg.\n\n> PCA - MEF and EB \n\nWe agree that the sentence may sound ambiguous and overstated. In the text, we only mention the MEF and Pancreas datasets, where FlatVI better unrolls the dynamics. We quantified the separation between initial and terminal states in latent spaces using clustering metrics in Tab. 8. While our method outperforms competitors on MEF and Pancreas, the reviewer correctly noted that it does not show a clear advantage on EB. We will clarify this in the text.\n\n[1] de Kruiff et al., Pullback Flow Matching on Data Manifolds", "claims_and_evidence": "The authors make claims in the paper about straight paths and enforcing locally Euclidean geometry. This is only the case for globally euclidean data. Indeed curved manifolds can be locally euclidean and a lot of work has gone into showing how to drive geodesics there (see Neural FIM, Fasina et al. ICML 2023 ), (GAGA Xu et al. AISTATS 20250 , (Metric Flow Matching Kapusniak, Kacper, et al. 2024),.", "methods_and_evaluation_criteria": "I think the idea of using the decoder’s pullback metric is very interesting, but I am unsure about the benefits of enforcing Euclidean geometry in latent space, especially for single-cell applications. \n\nTo my best knowledge, single cells are often hypothesized to lie on a manifold (manifold hypothesis), and the geometric and topological analysis of single-cell data have been proven very useful (Rizvi, A., Camara, P., Kandror, E. et al. Single-cell topological RNA-seq analysis reveals insights into cellular differentiation and development, Moon, K.R., van Dijk, D., Wang, Z. et al. Visualizing structure and transitions in high-dimensional biological data). Enforcing Euclidean structure might not be ideal for capturing the intrinsic non-Euclidean geometry often present in single-cell data. \n\nThe paper used OT Conditional Flow (OT-CFM) in their experiments, and it makes sense since OT-CFM does straight path interpolation, but there have been a lot of works recently leveraging the non-Euclidean geometry of the underlying manifold for trajectory inference tasks such as Metric Flow Matching (Kapusniak, Kacper, et al.), Geometry-Aware Generative Autoencoder (Sun, Xingzhi, et al.), WLF-SB (Neklyudov et al., 2024) etc.", "theoretical_claims": "The paper derives the Fisher Information Metric for the negative binomial distribution. I scanned the proof and it seems reasonable..", "experimental_designs_or_analyses": "The paper showcased four experiments: 1) sanity check on simulated data, 2) single-cell trajectory inference, 3) latent vector field, and 4) single-cell representations.  The experiments showed that locally euclidean geometry is preserved in the latent space of FlatVI, but they do not lend strong support for why FlatVI is better for modeling single-cell data. Overall, the experiments are interesting but not compelling enough. Indeed they did not compare to Neural FIM, GAGA or metric flow matching. \n\nThe experiment on simulated data is more like a sanity check for whether the latent space is regularized as expected. The experiment on single-cell trajectory shows that FlatVI performs better than VAE without the flatness regularization and GAE. However, the interpolation method used in the experiment is OT Conditional Flow (OT-CFM), which assumes straight path interpolation, so the OT-CFM could compound the experiment results, making FlatVI more favorable. The paper can also benefit from benchmarking with more methods than just two methods.\n\nThe same can be said for the latent vector field experiment. Here, OT-CFM was used to compute the velocity fields, and then the vector fields were used to learn the terminal states of the Pancreas data. The experiment showed that FlatVI has more consistent velocity fields compared to the other two methods, but it’s not clear to me why consistent velocity is more ideal for single-cell data. \n\nFor the data representation, the paper claims that “FlatVI effectively represents the biological structure in the latent space as illustrated by the separation between initial and terminal states.”. However, in Figure 4, I do not see a clear advantage of FlatVI compared to the other methods, especially on EB and MEF.", "relation_to_broader_scientific_literature": "The key contributions of this paper are about regularizing the latent space of VAE to approximate a locally Euclidean geometry through the pullback metric of the decoder. However, I am not convinced of the benefits of enforcing Euclidean geometry on the latent space, especially for single-cell data, considering there have been a lot of works showing the benefits of leveraging the non-Euclidean characteristics of the single-cell data (Rizvi, A., Camara, P., Kandror, E. et al. Single-cell topological RNA-seq analysis reveals insights into cellular differentiation and development, Moon, K.R., van Dijk, D., Wang, Z. et al. Visualizing structure and transitions in high-dimensional biological data) . They need to cite and discuss these as well as the other papers I mentioned above.", "essential_references_not_discussed": "This paper uses the pullback metric of the decoder to regularization the latent space but lacks citations on previous works in geometrically regularizing the latent space, such as Geometry-Aware Generative Autoencoder (Sun, Xingzhi, et al.) that regularizes the local latent space to match geodesic distances in data space, Geometry regularized autoencoders (\nDuque, Andrés F., et al.), and Gromov-Monge Embedding (Lee et al., 2024), Neural FIM (Fasina et al ICML 20230.", "other_strengths_and_weaknesses": "Overall, I think deriving a metric through the decoder pullback is very interesting, but the paper did not provide a compelling case for enforcing latent Euclidean geometry for single-cell data. The empirical experiments could benefit a lot from benchmarking with more interesting methods rather than the VAE baseline and GAE.", "other_comments_or_suggestions": "Please put the contributions in relation to related work, and offer more insight on the global euclidean assumption.", "questions_for_authors": "Are there regularizations need to make the geometry globally euclidean ? For example can L_flat be modified?", "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": "", "extra_scores": {}}, {"review_id": "UxKr7iGfSX", "reviewer": "Reviewer_5Fan", "summary": "The manuscript focuses on learning a specific representation intended to understanding dynamics, specifically arising from single-cell RNA sequencing (scRNA-seq). The authors regularize the latent space of a Variational Autoencoder (VAE) using the geodesics of a statistical manifold. This approach ensures that straight lines in the latent space correspond to geodesics on the statistical manifold. The authors demonstrate the advantages of their method on toy datasets with known geodesics, as well as on real-world scRNA-seq datasets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "We would like to thank Reviewer 5Fan for taking the time to review our work and for highlighting the aspects they found positive. We also appreciate their constructive suggestions regarding structural improvements and areas that could benefit from further clarification. \n\n> Interpretability of Fig. 2\n\nThe first two columns of Fig. 2 illustrate how the values of the Condition Number (CN) and the Variance of the Riemannian Metric (VoR) are distributed across latent cell observations. We offer a brief, intuitive description of these metrics:\n\n* VoR: Measures how the Riemannian metric at a given latent point varies relative to its immediate neighbours. In a perfectly Euclidean latent space, this metric would remain constant throughout, as Euclidean manifolds exhibit uniform geometry. Therefore, we expect this metric to take low values in the regularised space.\n\n* CN: Quantifies how uniformly distances are scaled in different directions within the latent space geometry. Specifically, it represents the ratio between the largest and smallest eigenvalues of the metric tensor. In a perfectly Euclidean space, this ratio would be 1, meaning distances are preserved equally in all directions. A higher CN indicates greater distortion, with distances being stretched more in some directions than others.\n\nIn the first two columns of Fig. 2, we show that the latent manifold produced by FlatVI consistently exhibits lower VoR and CN values compared to its unregularised counterpart, suggesting a closer approximation to Euclidean manifold properties. These results should be considered alongside Tab. 1, which demonstrates that stronger regularisation does not compromise reconstruction accuracy. In other words, FlatVI enforces the Euclidean constraint for downstream tasks while enabling precise trajectory reconstruction.\n\nDespite these improvements, some regions with elevated VoR and CN values remain in FlatVI’s output. We discuss these regions in Section K.1.4, where we hypothesise that they correspond to rapidly changing manifold areas located at the intersections between different cell-type categories.\n\n> Moving tables 4 and 8 into the main body. \n\nWe agree with the reviewer and thank them for their suggestion. In case we get the chance to present a revised version with an extra page, we will add Tab.4 to Section 5.1 and Tab.8 to Section 5.4.", "claims_and_evidence": "The claims are supported by theoretical and experimental results.", "methods_and_evaluation_criteria": "The method is theoretically grounded and is explained clearly. It has been rigorously tested on synthetic datasets with known geodesic properties, as well as on scRNA-seq datasets that exhibit temporal dynamics.", "theoretical_claims": "I checked the derivation of proposition 4.1 and it seems correct.", "experimental_designs_or_analyses": "I reviewed all the experiments and found them to be reasonable. However, I find Figure 2 difficult to interpret, and the conclusions drawn from the first two columns are unclear. However, I appreciate the quantitative experiments presented in Table 2.", "relation_to_broader_scientific_literature": "The key contributions focus on enhancing the interpretability of the latent spaces in Variational Autoencoders (VAEs). Other studies have also attempted to regularize these latent spaces. The authors reference these related works and provide a comparison with one of them.", "essential_references_not_discussed": "I believe there are no essential reference missing.", "other_strengths_and_weaknesses": "The manuscript is well written, and the method is clearly explained. I find most of the experiments to be of high quality; however, I find Figure 2 difficult to interpret.\n\nI believe the experimental section could benefit from slight revisions. Specifically, I suggest moving some tables from the supplementary material, such as Table 4 and Table 8, into the main body of the text, as they present compelling qualitative results.", "other_comments_or_suggestions": "I don't have any other comments or suggestions.", "questions_for_authors": "I don't have any other questions.", "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": "", "extra_scores": {}}, {"review_id": "FyDeq7zAhO", "reviewer": "Reviewer_zkLA", "summary": "This paper presents a novel training framework for VAEs for count data (i.e., negative binomial likelihood), where it enforces straight lines in latent space to map to geodesic paths on the data manifold, an assumption often made in VAE models for scRNA data but not explicitly enforced (or checked). In doing so, and combined with optimal transport-based modeling of cell fate transitions, FlatVI achieves good data reconstruction while better guiding the mapping between straight latent trajectories and cellular transitions.\n\n## update after rebuttal\nI thank the authors for their clarifications, and maintain my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "We sincerely thank zkLA for their thorough review and positive assessment of our work, and we welcome the opportunity to address their remaining questions.\n\n> A1\n\n$\\alpha$ relaxes the strictness of Euclidean regularisation. When the metric tensor $\\mathrm{M}(\\mathbf{z})$ is identity matrix $\\mathbb{I}_d$, the latent space is strictly Euclidean, and the inner product reduces to the linear product. Our goal is to enforce a constant local geometry in the latent space. On a Riemannian manifold, distances are computed as in Eq.5 in the manuscript. If $\\mathrm{M}(\\mathbf{z}) = \\mathbb{I}_d$ everywhere, the shortest path is a straight line. By setting $\\mathrm{M}(\\mathbf{z}) = \\alpha \\mathbb{I}_d$, we preserve straight geodesics while uniformly scaling tangent vectors. The trainable $\\alpha$ provides flexibility in latent metric scaling, improving reconstruction; freezing it restricts optimal geometry learning to an assumed scale.  \n\n> A2\n\nIn single-cell analysis, $\\theta$ is often modelled as a gene-specific parameter, capturing overdispersion from technical or biological effects. These are systematic gene-level patterns reflecting detection rates or expression variability. Thus, $\\theta$ is learned as a free parameter by the likelihood model, not as a function of latent space. Gradients from the likelihood loss update both the VAE and gene-wise $\\theta$. Though $\\theta$ is independent of latent states, its estimation reflects reconstruction quality by interacting directly with the likelihood loss. Hence, we include it in Tab.1.\n\n> A3\n\nIn Tab.1, we measure how closely straight paths align with geodesics by computing their MSE. Since geodesics on the latent manifold (equipped with the pullback metric) lack closed-form solutions without perfect regularisation, we approximate them numerically. Specifically, we use cubic splines that minimise the total Riemannian displacement between points, with local distances weighted by each model’s learned pullback metric (see l. 995-1023). When $\\alpha=1$, the pullback metric reduces to the identity and the minimum MSE achievable is 0. In practice, our spline-based method introduces numerical errors, so the MSE rarely reaches 0. This metric remains robust for comparing how models deviate from ideal Euclidean behaviour.\n\n> A4\n\nOur simulations focus on understanding model behaviour and validating our claims. On real data, where selection matters, we recommend choosing $\\lambda$ based on the downstream task. For trajectory inference, we tested $\\lambda$ and selected, for each dataset, the highest possible $\\lambda$ that improved trajectory reconstruction. For a detailed discussion, see App.G.  \n\n> A5\n\nThe mismatch occurs because Fig.9 displays a UMAP projection of the same latent encodings shown directly in Fig.2. While Fig. 2 presents the raw 2D latent space, we applied UMAP in Sec. K.1.4 to better visualize how unflattened regions are enriched in decision boundaries. Both figures represent identical data, the UMAP simply offers an alternative view for better visualisation. We will clarify this in the Appendix.  \n\n> A6\n\nIn Fig.4, we present data point embeddings rather than trajectories. We believe the non-linear structure observed, particularly for the Pancreas dataset, arises because PCA is a linear method that captures the global variance but does not preserve the local geometry of the data. Although our regularisation aims to approximate Euclidean space locally, some residual non-Euclidean structures may remain since we trained the model with $\\lambda=1$. PCA, being a global method, emphasises this non-linearity when reducing the data to 2D. This is less apparent in the MEF dataset, where flattening appears to better unfold the natural temporal structure (Fig.4b). To better preserve the local structure, non-linear methods such as Isomap could be considered. We chose PCA for its interpretability and to analyse dimensionality across latent spaces.\n\n> A7\n\nMany single-cell tools assume that Euclidean distances in the representation space reflect biological relationships. FlatVI explicitly links the VAE's latent manifold to a locally Euclidean latent space, ensuring these distances better match biological proximity in terms of single-cell counts. This creates a principled approach for downstream analysis (kindly refer to the 1st answer to NJwg).\n\n> A8\n\nTo extend FlatVI to $\\beta-$VAEs with hyperparameters controlling the prior spread, we can incorporate the flattening loss alongside the modified KL term. This allows us to regularise both the latent trajectory straightness and posterior variance. Crucially, this extension does not disrupt the core guarantees of FlatVI, as the FlatVI loss only influences the latent space geometry shaped by the decoder model (Eq.7), while the KL term balances the posterior with the prior. This ensures that geometry regularisation and posterior variance are controlled without compromising the underlying structure of the latent space.", "claims_and_evidence": "The main claim is that, by adding the novel flattening loss to NB-VAE training, straight latent trajectories that is often used for interpolation would better correspond to minimal paths on the data manifold, e.g., by going through an intermediate cell state. As far as I can tell, the authors successfully demonstrate this on synthetic and multiple real datasets.", "methods_and_evaluation_criteria": "The method is elegant and simple, adding just one additional loss term to achieve the desired result, though requiring an additional regularization hyperparameter (see Questions). The experiments and evaluation criteria chosen appear to be appropriate, and demonstrate FlatVI’s improved performance over existing baselines.", "theoretical_claims": "I did not go through the loss proof in the amount of detail required to reconstruct it, but superficially everything checks out, aside from some questions in the end.", "experimental_designs_or_analyses": "The experimental design and results seem comprehensive (checked all), and overall quite performant.", "relation_to_broader_scientific_literature": "This paper is a nice addition to the VAE, and more broadly, generative modeling literature. Having the correspondence between straight latent trajectories and geodesic trajectories on the data space would be useful in general for interpretability, and the specific negative binomial likelihood VAE would be useful in other applications with (low) count data, like neural recordings with spike counts.", "essential_references_not_discussed": "None I’m aware of.", "other_strengths_and_weaknesses": "- The paper is clearly and concisely written, and overall quite informative.\n- The exposition from deterministic to stochastic AE was helpful to guide the reader.\n- Really nice results overall and great visualization.", "other_comments_or_suggestions": "None", "questions_for_authors": "- can you explain in more detail why alpha (in Eq 8) is a trainable parameter, and how varying a frozen alpha affects the results?\n- could you clarify why the inverse dispersion parameter is not a function of the latent space (line 232-234), and whether it is estimated through the decoder (e.g., Eq. 2)? The evaluation metric considers reconstruction error on dispersion, so it seems that it is given by the decoder?\n- What is the scale we should expect for the MSE between geodesic and euclidean distances? Is 0 the theoretical minimum if everything worked correctly?\n- How should one choose lambda in practice, especially since the optimal for the three metrics do not necessarily coincide in Table 1 (though the differences in mu and theta are minimal)?\n- Why are there these “streaks” of very high condition number for FlatVI in Fig 2b? While K.1.4 and Figure 7-9 touches on this point, I don’t quite see how the boundary effect highlighted in Fig 9 translate to the high CN regions in Figure 2?\n- Naive question: in figure 4, imagining paths from initial to terminal state seems quite non-linear given the curved manifold (e.g., for Pancreas). Given that the trajectories should be straight and PCA is a linear reduction, wouldn’t we expect more linear manifolds and “straighter” paths from early to terminal? Instead of PCA and taking the first two PCs, is there a way to do dimensionality reduction but preserving the linear paths? Or am I fundamentally misunderstanding something?\n- Another naive question, more broadly, why is it assumed (or desirable) that straight paths in latent space map to geodesic paths in data space through a cell’s transition?\n- how would this scheme extend to work with, e.g., beta VAEs with an additional hyperparameter (on the KL loss) controlling the prior spread? It would be interesting to be able to regularize both the latent trajectory straightness and posterior variance?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This is a methodologically novel and experimentally solid contribution at the intersection of machine learning, geometry, and computational biology. All reviewers appreciated the conceptual clarity and technical novelty, though some were hesitant due to limited comparative scope or concern over assumed geometry. However, the rebuttal mitigated these reservations, showing that FlatVI is not only a tractable regularizer but also a foundation for future geometric learning in probabilistic settings. The technique is widely applicable beyond scRNA-seq, such as in other count-based domains (e.g., spike train analysis, molecular measurements).
Decision: Accept (spotlight poster)