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

TL;DR

This paper introduces FlatVI, a training framework for single-cell variational autoencoders that enforces Euclidean geometry in the latent space, improving manifold interpolation and trajectory reconstruction in single-cell RNA sequencing data.

Contribution

FlatVI is a novel method that regularizes the latent space of VAEs towards Euclidean geometry, specifically designed for single-cell count data, enhancing interpolation accuracy.

Findings

Supports theoretical soundness with synthetic data

Improves trajectory reconstruction in real single-cell data

Enhances compatibility with Euclidean-based downstream methods

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.