Enforcing Latent Euclidean Geometry in Single-Cell VAEs for Manifold Interpolation

We sincerely thank zkLA for their thorough review and positive assessment of our work, and we welcome the opportunity to address their remaining questions.

A1

$\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.

A2

In 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.

A3

In 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.

A4

Our 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.

A5

The 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.

A6

In 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.

A7

Many 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).

A8

To 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.

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.

Global regularisation

Similar to GAE, we propose to:

• Sample a couple of cells and encode them.
• Approximate the latent geodesic distance between representations using the pullback. One could fit an energy-minimising cubic spline weighing local shifts by the pullback.
• Minimise the difference between the latent Euclidean and the geodesic distances.

This approach is expensive because it contains an unstable inner optimisation loop requiring decoding and gradient computation.

Local vs global

While 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.

Contextualisation

Apologies for the missing citations; we will add them upon revision.

• 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.

• 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:

  • Only consider low-dimensional and continuous input data (e.g., PCA).
  • Do not enforce a tractable latent manifold, hence, they have to learn geodesics between latent observations with a NeuralODE.

These aspects may hinder scaling to the gene space, as sparse expression vectors are unsuitable for Euclidean distances, and geodesic NeuralODEs are unstable (see below).

Manifold assumptions

There 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].

FlatVI's benefits

Kindly refer to A1. to NJwg and A7. to zkLA.

Benchmark

We 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.

We 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).

Since the FM algorithm is not the novelty of our work, we did not compare it with MFM.

Consistency

Kindly refer to A2 to NJwg.

PCA - MEF and EB

We 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.

[1] de Kruiff et al., Pullback Flow Matching on Data Manifolds

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.

Interpretability of Fig. 2

The 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:

• 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.

• 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.

In 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.

Despite 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.

Moving tables 4 and 8 into the main body.

We 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.

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.

A1. Limited use/comparison with other algorithms.

While our specific examples focus on OT-CFM, the benefits of FlatVI extend to a broader range of single-cell analysis techniques.

• 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.
• 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.
• 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.
• The latent Euclidean assumption is also key in batch integration [1] and perturbation prediction [2], broadening the use of our regularisation.
• 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.

A2. Velocity consistency

Importantly, 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.

A3. Contribution

Our work presents a novel methodological formulation that, to the best of our knowledge, has not been explored before:

• 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).
• 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.
• 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.
• 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).

Finally, 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.

[1] Lücken et al., Benchmarking atlas-level data integration in single-cell genomics

[2] Eyring et al., Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation

[3] Gayoso et al., Deep generative modelling of transcriptional dynamics for RNA velocity analysis in single cells

[4] Huguet et al., Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

[5] Sun et al., Geometry-Aware Generative Autoencoders for Warped Riemannian Metric Learning and Generative Modeling on Data Manifolds