Modelling single-cell RNA-seq trajectories on a flat statistical manifold

Published: 28 Oct 2023, Last Modified: 09 Nov 2023NeurIPS2023-AI4Science OralEveryoneRevisionsBibTeX
Keywords: scRNAseq, Information geometry, Optimal Transport
TL;DR: We perform optimal transport on a flat statistical manifold to infer single-cell trajectories through time.
Abstract: Optimal transport has demonstrated remarkable potential in the field of single-cell biology, addressing relevant tasks such as trajectory modelling and perturbation effect prediction. However, the standard formulation of optimal transport assumes Euclidean geometry in the representation space, which may not hold in traditional single-cell embedding methods based on Variational Autoencoders. In this study, we introduce a novel approach for matching the latent dynamics learnt by Euclidean optimal transport with geodesic trajectories in the decoded space. We achieve this by implementing a "flattening" regularisation derived from the pullback metric of a Negative Binomial statistical manifold. The method ensures alignment between the latent space of a discrete Variational Autoencoder modelling single-cell data and Euclidean space, thereby improving compatibility with optimal transport. Our results in four biological settings demonstrate that these constraints enhance the reconstruction of cellular trajectories and velocity fields. We believe that our versatile approach holds promise for advancing single-cell representation learning and temporal modelling.
Submission Track: Original Research
Submission Number: 152