[Paper Review] Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

Guillaume Huguet, D S Magruder+5 more|arXiv (Cornell University)|Jun 29, 2022

Traffic Prediction and Management TechniquesEngineering19 citations

TL;DR

MIOFlow proposes a neural ODE-based framework that learns stochastic, continuous population dynamics from static, time-sparse snapshots by combining manifold learning with optimal transport. It uses a geodesic autoencoder to embed data in a latent space where distances reflect a novel multiscale diffusion geodesic distance, enabling accurate interpolation and trajectory inference on complex manifolds, outperforming normalizing flows and Schrödinger bridges in modeling bifurcations and biological dynamics.

ABSTRACT

We present a method called Manifold Interpolating Optimal-Transport Flow (MIOFlow) that learns stochastic, continuous population dynamics from static snapshot samples taken at sporadic timepoints. MIOFlow combines dynamic models, manifold learning, and optimal transport by training neural ordinary differential equations (Neural ODE) to interpolate between static population snapshots as penalized by optimal transport with manifold ground distance. Further, we ensure that the flow follows the geometry by operating in the latent space of an autoencoder that we call a geodesic autoencoder (GAE). In GAE the latent space distance between points is regularized to match a novel multiscale geodesic distance on the data manifold that we define. We show that this method is superior to normalizing flows, Schrödinger bridges and other generative models that are designed to flow from noise to data in terms of interpolating between populations. Theoretically, we link these trajectories with dynamic optimal transport. We evaluate our method on simulated data with bifurcations and merges, as well as scRNA-seq data from embryoid body differentiation, and acute myeloid leukemia treatment.

Motivation & Objective

• To address the challenge of inferring continuous, stochastic population dynamics from static, time-ordered snapshots of high-dimensional data.
• To model trajectory inference in systems where direct longitudinal tracking is impossible, such as in single-cell RNA sequencing.
• To ensure that inferred trajectories respect the intrinsic geometry of the underlying data manifold, rather than the ambient high-dimensional space.
• To develop a computationally efficient method that avoids the need for Gaussian priors and deterministic flows, which are common limitations in existing generative models.
• To enable accurate interpolation and reconstruction of unobserved timepoints and individual cell trajectories in complex biological systems.

Proposed method

• MIOFlow employs neural ordinary differential equations (Neural ODEs) to model continuous flows between static population snapshots at different timepoints.
• It uses a geodesic autoencoder (GAE) that regularizes the latent space to preserve a novel multiscale diffusion geodesic distance on the data manifold.
• The diffusion geodesic distance is defined via a kernel-based graph Laplacian and converges to the true geodesic distance on the manifold in the limit of infinite samples.
• The flow is penalized via optimal transport with a manifold ground distance, ensuring that the transport cost between timepoints aligns with the intrinsic geometry of the data.
• The method operates in the latent space of the GAE, reducing computational cost and improving generalization by learning dynamics on the low-dimensional manifold.
• Stochasticity is introduced naturally through the initial sampling of points in the latent space, avoiding the need for explicit noise injection as in deterministic flows.

Experimental results

Research questions

• RQ1 Can we learn continuous, stochastic population dynamics from static, time-sparse snapshots without assuming a Gaussian prior?
• RQ2 How can we ensure that the inferred trajectories respect the intrinsic geometric structure of the underlying data manifold?
• RQ3 Can we construct a data-driven, multiscale distance metric in the latent space that converges to the true geodesic distance on the manifold?
• RQ4 How does MIOFlow compare to existing methods like normalizing flows, Schrödinger bridges, and TrajectoryNet in modeling complex dynamics such as bifurcations and merges?
• RQ5 Can MIOFlow accurately reconstruct individual cell trajectories and capture dynamic changes in gene expression during biological processes like differentiation and treatment response?

Key findings

• MIOFlow outperforms normalizing flows and Schrödinger bridges in interpolating between populations on synthetic data with bifurcations and merges, producing smoother and more geodesically accurate trajectories.
• On simulated data, MIOFlow successfully models divergent trajectories, with the flow accurately capturing the branching structure of the underlying manifold.
• In scRNA-seq data from embryoid body differentiation, MIOFlow reconstructs biologically plausible lineage trajectories that align with known developmental markers.
• On acute myeloid leukemia (AML) treatment data, MIOFlow reveals that surviving cells transition toward a leukemic stem cell (LSC) signature, with increased expression of SLPI, MYB, and B2M at later timepoints.
• The geodesic autoencoder with diffusion geodesic distance enables accurate manifold reconstruction, and the latent-space flows preserve meaningful biological distances that decode to biologically interpretable gene expression patterns.
• MIOFlow achieves faster inference per function evaluation than normalizing flows due to avoiding Jacobian trace computation, scaling as O(k) instead of O(k^2) in ambient dimension k.