Enforcing Latent Euclidean Geometry in VAEs for Statistical Manifold Interpolation

We thank the reviewer for their detailed feedback and suggestions, which improved the quality of the manuscript. Below, we provide detailed answers to the raised questions and concerns.

Confusion about manifold

Thank you for your feedback. We acknowledge that referring to the manifold of interest as a "complex data manifold" is a coarse and under-elaborated description. Here, we take the opportunity to clarify what we aim to convey instead.

One important aspect of our model is its application to experimental data, which often includes noise, non-continuous domains, and lacks a clear association with a geometry that has closed-form geodesics. This is the primary reason we leverage the latent space of a VAE for manifold learning and regularise it towards a simple geometric structure to facilitate interpolation. In essence, the main desideratum of our model is to provide a framework for efficiently learning how to interpolate the experimental data manifold. Unlike well-defined geometries, such as that of a sphere where geodesic paths are straightforward to compute, experimental data manifolds typically do not allow for direct geodesic traversal. This data-oriented approach has been considered quite often recently due to its relevance in applied fields [1,2].

Since our primary focus is single-cell RNA-seq, we specifically design our method around the concept of a statistical manifold—i.e., a manifold of probability distributions from a particular family. This concept is well-suited to data regimes where an accurate noise model can be assumed and captured by the decoder of a VAE, such as the negative binomial distribution for scRNA-seq. In this context, we leverage information geometry to connect the geometry of the latent space manifold with that of the decoded space, governed by the Fisher Information Metric.

We have revised the text to better reflect our purpose, replacing the term complex manifold with the idea of experimental data manifold or, in more specific passes, with a specific reference to the statistical manifold spanned by the VAE deocder. Examples of such clarifications are in the abstact (L15) and in the caption of Fig.1. Additionally, we updated Fig.1 by removing the hole to avoid potential misconceptions and to provide a more general visualisation of the concept of a data manifold. We also updated the introduction (see lines 54–64) to more accurately convey our modelling objectives and to explain how our framework optimises towards them.

We hope the reviewer finds our elaboration and the changes made to the text insightful.

Evaluation of the flattening loss in a lighter setting

We thank the reviewer for their relevant suggestion. Inspired by the toy example in [3], where the authors consider noisy circular data, we designed an experiment simulating data points sampled from a 2D circular manifold embedded within a 3D Gaussian statistical manifold with constant variance. The 3D manifold is equipped with the Fisher Information metric for the Gaussian distribution, as described in [3].

In this experiment, we trained a 2D VAE with regularisation strengths $\lambda \in \{0, 0.5, 1\}$, as suggested by the reviewer. After training, we computed linear paths between random pairs of points in the latent space and plotted the resulting geodesic interpolations on the 2D manifold.

The results are shown in Fig. 10 (Section I.2), where the red paths between point pairs represent the decoded means of the Gaussian likelihood. Analysing the results, we observed several interesting patterns that support the working principles of our model. Without regularisation ($\lambda = 0$), decoded geodesic paths deviate from the manifold and fail to take the shortest route along it between source and target points. With $\lambda = 0.5$, the adherence to the manifold improves, but interpolations still follow longer routes between points. Only with $\lambda = 1$ does the regularisation ensure that geodesic paths remain on the manifold and correspond to the shortest path between points.

These results highlight the effectiveness of our approach, with the empirical findings providing support of our model assumptions and design.

We thank the reviewer for their thorough and constructive feedback. We appreciate the positive recognition of our technical approach, as well as the highlighting of our model's potential for reliable data interpretations in domains such as the life sciences. Below, we provide detailed answers to the raised questions and concerns.

Better communication towards our contribution

• We thank the reviewer for raising this point. In the updated version of our manuscript, we tried to delineate better between known results and our contribution towards aligning latent space interpolations with geodesics on statistical manifolds. Our primary focus was the introduction, where the reviewer detected ambiguity in the aspects that distinguish our framework from prior work. We reformulate lines 54 to 64 by listing a series of desiderata for an ideal latent manifold interpolation model, and explaining why existing work does not exhaustively fulfill them all. Right after this paragraph, FlatVI is introduced as a method incorporating all desirable aspects. We believe the reviewer's suggestion significantly helped us improve the flow of our introduction.
• Furthermore, we discuss our contribution in detail in our response to reviewer tvQk.
• We also updated our manuscript to reflect that our methods regularises towards / approximates a Euclidean geometry. See, for example, L22, L83, L418, or L529 (all such changes are highlighted in blue in the text). We appreciate the reviewer's feedback and believe that these changes contributed towards the clarity of our work.

Computational complexity

• Expectedly, introducing an additional term to regularise the pullback metric tensor causes an increase in the model's runtime, as it involves the evaluation of the Jacobian of the decoder ($\mathbb{J}_h$ in the paper, where $h$ is the decoder function).
• In practice, this work focuses on tabular data, which we embed into the latent space using MLP models with one to three layers of non linearities. Although the regularisation introduces a computational overhead, a single forward step is completed in a reasonable time (see Tab. 6). Moreover, in the considered setting, our Jacobian-based regularised model scales better than GAE, which requires the computation of pairwise distances between input observations to estimate the diffusion geodesic distance in the data domain.
• To compute the pullback metric more efficiently, we use the Jacobian-vector product (JVP) rather than constructing the full Jacobian. For the full metric tensor, we apply JVP to compute directional derivatives along each standard basis vector, dynamically assembling the necessary components of the decoder's Jacobian. This approach avoids the computational and memory overhead of explicitly forming the full Jacobian while providing the required terms for the metric computation.

Q: Is a KL-approximation also possible in our setting?

This is a very interesting question. To clarify, given a latent code $\mathbf{z}$ and an observation $\mathbf{x}$, the authors in [2] define an approximation for the elements of the metric tensor $\mathbf{M}$ as a function of the KL divergence between the likelihoods $p(\mathbf{x}|h(\mathbf{z}))$ and $p(\mathbf{x}|h(\mathbf{z}+\delta \mathbf{z}))$, where $\delta$ represents a perturbation of $\mathbf{z}$ and $h$ is a decoder function. Initially, we considered applying this scheme for computational efficiency. However, it's important to recall our likelihood model:

$\mathbf{x} \sim \mathrm{NB}(\mu, \theta)$

$\mu = l \cdot \mathrm{softmax}(h_{\phi}(\mathbf{z}))$

where $l = \sum_g \mathbf{x_g}$ and $h_\phi$ is a neural network that decodes gene proportions (indicated as $\rho_\phi$ in the manuscript).

When implementing the approximate metric method, we encountered an issue: Decoding perturbations of $\mathbf{z}$ introduced some ambiguity in how to handle the size factor. Since $l$ is not a function of $\mathbf{z}$, it was unclear what value to use for this variable when computing the perturbed likelihood $p(\mathbf{x}|h(\mathbf{z} + \delta \mathbf{z}))$, which is required for the metric approximation. As a result, we decided to adhere to the standard version of the metric tensor, which provides a closed-form expression in the case of the negative binomial distribution.

Q: Will the code be made available

Yes, we will make the code available as part of the final submission. The software release will follow the structure of scVI tools [1] to promote a faster introduction of the model to the community.

Q: Is there a connection to model uncertainty?

Unlike FlatVI, the authors in [2] do not regularise the latent manifold to a simple geometry. Consequently, to perform geodesic interpolations between latent codes $\mathbf{z}_1$ and $\mathbf{z}_2$, they must learn and optimise a cubic spline that connects the two points based on energy minimisation. To ensure the fitted curve does not model paths that leave the data's support, they introduce an uncertainty regularisation term. This term extrapolates to high uncertainty if the curve represents a trajectory incompatible with the data support.

In contrast, FlatVI does not require learning parameterised latent curves via cubic splines as in [2]. This is because we regularise the latent pullback geometry to approximate a Euclidean manifold, where geodesics between $\mathbf{z}_1$ and $\mathbf{z}_2$ are simple linear paths that map to data geodesics via the decoder. Ideally, when Euclidean geometry is well-approximated and geodesics are accurately decoded along the data manifold, our model remains uninfluenced by uncertainty.

Nevertheless, implementing uncertainty regularisation in our setting is still feasible, particularly in scenarios where FlatVI is underregularised and the latent space is not fully flattened. As the reviewer suggested, our loss pushes the model towards a latent Euclidean geometry but does not necessarily guarantee it, especially depending on the regularisation strength. If regularisation is incomplete, one might opt to optimise geodesic curves instead of relying on linear paths to trace interpolations between $\mathbf{z}_1$ and $\mathbf{z}_2$. In such cases, it would be necessary to define a notion of maximum uncertainty for the negative binomial likelihood spanned by the decoder under mean/inverse dispersion regularisation. For instance, we could set the inverse dispersion parameter of the distribution to infinity when modelling transitions outside the data's support, as the variance of the negative binomial distribution is given by $\mu + \theta \mu^2$, where $\theta$ is the inverse dispersion parameter. This strategy could help enforce meaningful transitions while preserving compatibility with the underlying data structure. Adding such an extension to FlatVI would be an interesting avenue for the future.

We once again thank the reviewer for their insightful suggestions and hope that we were able to address all remaining concerns satisfactorily. We look forward to further discussion.

[1] Gayoso, Adam, et al. "A Python library for probabilistic analysis of single-cell omics data." Nature biotechnology 40.2 (2022): 163-166.

[2] Arvanitidis, Georgios, et al. "Pulling back information geometry." arXiv preprint arXiv:2106.05367 (2021).

We thank the reviewer for taking the time to review our work and for providing the opportunity to elaborate further on our contribution. Below, we highlight our contribution in greater detail.

The biological contribution

• scRNA-seq profiles high-dimensional states, with each feature representing a gene's expression as a discrete count of RNA molecules. The data is inherently sparse and skewed due to high dropout rates and overdispersion caused by experimental inaccuracies and biological variability.

• These characteristics necessitate modeling noise with appropriate discrete models. VAEs are ideal, as their decoder maps the latent space $\mathbf{z}$ to the parameters of an arbitrary likelihood model. For scRNA-seq, a negative binomial decoder is most suitable, as it models counts while capturing the mean-variance trend beyond Poisson assumptions. Crucially, decoding $\mathbf{z}$ represents a cell as the parameters of a negative binomial distribution, effectively making each decoded cell a probability distribution.

• In the VAE literature, the decoder is often assumed to span a statistical manifold of probability distributions from a specific family. Similarly, we can interpret decoded single cells as points on such a manifold. To the best of our knowledge, no prior work has approached manifold learning assuming that single-cell data lie on the statistical manifold of negative binomial distributions defined by the NB-VAE decoder.

• Why take this approach? Many scRNA-seq tasks rely on interpolations in a latent or reduced-dimensional space. Our work explores whether we can push such interpolations to meaningful paths on the statistical manifold spanned by the decoder. Since the negative binomial model is the most empirically accurate for scRNA-seq noise, we aim for latent interpolations to align with the geometry of this statistical manifold.

The machine learning contribution

• To connect with the application setting, we develop an approach to approximate geodesic walks on the decoded manifold as trajectories on a simpler latent manifold geometry. Although geodesic interpolation methods exist, we aim for the following properties:
  • Interpolation of empirical data manifolds: The approach should interpolate an empirical data manifold where no closed-form geodesic distance between two points exists. This distinguishes our framework from methods studying interpolations on standard manifolds with tractable geodesics (e.g., Riemannian Flow Matching [1] or Fisher Flow Matching [2]).
  • Cost-efficient latent manifold interpolation: Interpolations along the latent manifold should be computationally inexpensive. Without regularising the latent manifold to a simpler geometry, interpolations are typically parameterised geodesic curves requiring costly simulations. For instance, Arvaniditis et al. [3] learn cubic splines that minimise the geodesic distance between two latent points using the pullback metric $M(\mathbf{z})$ (see Eq. 5 in our paper). Computing such interpolations involves iterative updates to the spline's weights, resulting in an expensive optimisation problem.
  • Suitability for complex, non-continuous data domains: The strategy must handle challenging domains, such as scRNA-seq data. While methods like Metric Flow Matching [4] address some of these challenges, their manifold regularisation technique relies on a continuous data space. Indeed, the authors primarily consider low-dimensional approximations of gene expression vectors, which is often insufficient for capturing biological variability.
  • Real-world applicability: In the spirit of the Applications to physical sciences track, the framework should connect to existing tools to encourage its adoption in practical scenarios.

• FlatVI achieves all of these goals:
  • It builds on the pullback geometry of VAEs to enable latent manifold learning with real high-dimensional data.
  • It introduces a regularisation strategy that encourages linear latent interpolations to approximate geodesic paths on the statistical manifold defined by the decoder. Notably, no prior method has explored flattening a VAE's latent space using the Fisher-based pullback metric.
  • It is applied to the complex scRNA-seq domain, formalising for the first time the connection between the geometry of latent cellular representations and the stochastic negative binomial manifold of the decoder. This framework is also generalisable to other count-based distributions, such as the Poisson.

Dear Authors

I appreciate your reply. The explanation helps me better understand the paper. I am so sorry I still think a regularization term is incremental as the main contribution.

However, the authors have nice experiments in scRNA-seq. Some theoretical analyses when this tool is applied to negative binomial distributions are also provided. I am sorry those are topics I am not familiar with. I will increase my rating, but keep my confidence low.

I deeply appreciate the authors' paper and the AC's effort.

Best wishes reviewer

Dear Reviewer tvQk,

We sincerely thank you for the time invested in reviewing our responses and for highlighting the positive aspects of our work. We also appreciate your honesty and openness in considering an increased score despite the low confidence.

Regarding novelty concerns, our regularisation strategy uniquely imposes geometric constraints on a VAE, introducing a novel, unexplored approach. Specifically, we are the first to induce correspondence between geodesics on an (approximately) Euclidean latent space and the statistical manifold spanned by a decoder with arbitrary likelihood. This creates a robust representation space for manifold interpolations across data types with known noise models, an important step forward in scientific domains. This contribution advances trajectory modelling and interpolations in experimental data, addressing a combination of desiderata unmet by current methods like Flow Matching on manifolds or geometry-aware VAEs (see desiderata bullet point list in our rebuttal).

We also encourage you to review the new experimental result directed at 7m6g (Section I2, Fig. 10). It illustrates how linear latent space paths map to geodesic paths on a 2D circular manifold within a 3D statistical manifold, enabled only by our flattening approach. This example uses a Gaussian likelihood, highlighting the framework’s flexibility beyond the negative binomial case.

We hope these clarifications and additional insights emphasise the novelty and impact of our work, addressing your concerns and encouraging a more favourable evaluation. Thank you again for your thoughtful feedback.

Sincerely, The Authors

We thank the reviewer for their detailed feedback on our work. We appreciate the positive comments regarding the manuscript's content and presentation, and we value their constructive criticism as an opportunity to enhance the quality of our contribution. Below, we provide our responses to the questions and concerns raised during the review process.

Addition of Spearman correlation and neighbourhood overlap metrics.

We acknowledge the reviewer's comment and have implemented the suggested changes. The updated results are presented in Tab. 4, section I.1.3. To compute the metrics of interest, we follow these steps:

• Sampling: For each regularisation strength, we randomly sample 50 synthetic observations across 5 repetitions.
• Distance Computation: For all pairs of observations within the sampled batch, we compute both Euclidean and pullback geodesic distances, resulting in two distinct distance matrices. The pullback geodesic approximation is derived using Eq. 38 from the manuscript.
• Neighborhood Evaluation: The two distance matrices induce neighbourhood structures. We evaluate two metrics:
  1. Average Spearman correlation: We compute the Spearman correlation between Euclidean and geodesic distance vectors for each observation. A high correlation indicates alignment in the ordering of distances from a given observation $i$ across the two metrics.
  2. Neighborhood overlap score: For each observation, we calculate the $k$-nearest neighbours based on both Euclidean and geodesic distance matrices, then compute the average proportion of shared neighbours across all sampled observations.

We report the neighbourhood overlap metric for $k = \{3, 5\}$, reflecting the preservation of local metric properties. Conversely, the Spearman correlation captures the global structure by operating on the entire vector of distances for each observation.

The results in Tab. 4 demonstrate that regularisation improves both Spearman correlation and neighbourhood overlap values. Specifically, the Spearman correlation starts at relatively high values even before regularisation. However, the neighbourhood overlap score shows more significant improvement with increasing $\lambda$, particularly for $k = 3$, where a regularisation strength of $\lambda$ = 10 yields an average neighbourhood overlap proportion of 0.80. These findings validate our model's principles, where regularisation effectively aligns the pullback geometry closer to Euclidean geometry.

Notes on geodesic computation

1. We use batches of 50 observations for neighbourhood computations to manage the high computational complexity of estimating geodesics with Eq. 38.
2. Eq. 38 provides an approximation of the true geodesic via cubic spline optimisation [1], as the exact geodesic distance in the pullback metric setting is intractable.

Rename the 'Model' column to 'Lambda'.

We applied the recommended change in the updated version of the document.

Add CN and VoR features to a Euclidean space so that the reader can visually compare the different representations.

Please find the suggested plot in Fig. 6 of the updated document. The bottom row illustrates the expected behaviour of the VoR and CN metrics in a perfectly Euclidean space. Specifically:

• The VoR metric is equal to 0.
• The CN metric is equal to 1.

On the right-hand side, we visualise the geodesic distances between pairs of observations in the Euclidean space. These distances are represented as straight lines connecting the points, similar to the regularised latent space in FlatVI.