Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

Guillaume HuguetDaniel Sumner MagruderAlexander TongOluwadamilola FasinaManik KuchrooGuy WolfSmita Krishnaswamy

article2022NeurIPS126 citationsBest Paper Award

Develops MIOFlow, a framework combining neural ordinary differential equations and optimal transport to reconstruct continuous biological trajectories from static snapshot measurements while respecting intrinsic data geometry.

Listen

Modern biomedical technologies, such as single-cell RNA sequencing, capture rich biological snapshots but destroy cells during measurement, making it impossible to observe continuous individual cell trajectories directly over time. Reconstructing continuous, dynamic cellular processes from these discrete, static snapshots is critical for understanding developmental biology, disease progression, and treatment response mechanisms.

The article develops and evaluates a computational framework called Manifold Interpolating Optimal-Transport Flow (MIOFlow) to learn continuous, stochastic population dynamics and infer individual trajectories from static snapshot data. The objective is to demonstrate that MIOFlow accurately interpolates unmeasured intermediate states while constraining dynamic paths to the natural geometry of the data manifold.

The researchers designed a two-part computational approach combining manifold learning and neural ordinary differential equations regularized by optimal transport. First, they constructed a Geodesic Autoencoder that maps high-dimensional data into a low-dimensional latent space while preserving intrinsic manifold distances using a newly defined diffusion geodesic metric. Second, they trained neural differential equations with an added diffusion term to model population trajectories directly in this latent space, penalizing deviations between predicted and observed distributions using optimal transport. The method was evaluated on synthetic branching datasets and two real-world biological datasets: human embryoid body differentiation over 27 days and acute myeloid leukemia cells undergoing a seven-day chemotherapy regimen.

The evaluation revealed several key findings in order of importance. First, MIOFlow accurately reconstructed continuous paths along complex bifurcations and merging branches without taking artificial shortcuts across empty space, substantially outperforming alternative models like TrajectoryNet and Diffusion Schrödinger Bridges. Second, in held-out timepoint evaluations, MIOFlow reduced interpolation error metrics by roughly 50% to over 60% compared to baseline and competing methods on synthetic benchmarks. Third, MIOFlow operated dramatically faster than continuous normalizing flow baselines, reducing training times from over an hour to under five minutes on benchmark tasks because it avoids high-order Jacobian trace calculations. Fourth, on biological single-cell data, the framework accurately captured non-monotonic gene expression patterns during neuronal development and revealed how leukemia cells transition toward drug-resistant stem cell signatures during chemotherapy.

These findings demonstrate that embedding dynamic flows within an intrinsic manifold geometry provides a scalable, computationally efficient solution for modeling complex biological trajectories. By operating in latent space and eliminating the need for rigid Gaussian priors or heavy derivative computations, the approach lowers computational costs and minimizes modeling artifacts. This allows researchers to reliably map out cellular drug-resistance pathways, reducing technical risk when identifying therapeutic targets and biomarkers.

Organizations analyzing dynamic population snapshots can adopt this framework to reconstruct lineage trajectories and study time-dependent interventions. To support future implementation, practitioners should evaluate unbalanced optimal transport formulations when addressing severe cell proliferation or death, and conduct preliminary validation on sampling density before scaling up analysis.

The primary limitations include sensitivity to irregular timepoint spacing and the potential for standard differential equation solvers to struggle with numerically stiff dynamics. Confidence in the empirical results is high given rigorous benchmarking across both synthetic datasets and real-world genomic datasets, though caution is recommended when applying the model to populations with extreme, unmeasured growth-rate imbalances.

Mila – Québec Artificial Intelligence Institute Université de Montréal Yale University autoencoders generative model geometry Manifold Interpolating Optimal-Transport Flows trajectory inference optimal transport

arXiv: 2206.14928

Paper: Neural Ordinary Differential Equations, Ricky T. Q. Chen et al. (2018). Introduces Neural Ordinary Differential Equations and continuous-time normalizing flows, providing the fundamental differential equation framework that MIOFlow adapts for latent trajectory learning.
Paper: Score-Based Generative Modeling through Stochastic Differential Equations, Yang Song et al. (2021). Establishes the continuous stochastic differential equation formulation for generative modeling that underlies continuous population dynamics and diffusion bridge approaches.
Paper: Supervised Training of Conditional Monge Maps, Charlotte Bunne et al. (2022). Develops methods for learning continuous optimal transport maps between unpaired single-cell populations across biological conditions, framing the core snapshot-matching problem addressed by MIOFlow.
Paper: Variational Inference with Normalizing Flows, Danilo Jimenez Rezende et al. (2015). Provides the foundational principles of continuous density transformations and normalizing flows used to construct invertible latent representations.
Paper: Denoising Diffusion Probabilistic Models, Jonathan Ho et al. (2020). Introduces the standard formulation of denoising diffusion models, establishing the stochastic path mechanics that MIOFlow constrains along data manifolds.

Paper: Matching Normalizing Flows and Probability Paths on Manifolds, Heli Ben-Hamu et al. (2022). Generalizes flow matching directly to Riemannian manifolds without requiring numerical ODE solvers during training, extending continuous probability path learning on non-Euclidean geometries.
Paper: Flow Matching for Generative Modeling, Yaron Lipman et al. (2023). Introduces simulation-free Flow Matching with optimal transport paths, offering an alternative fast framework for learning continuous flows between distributions.
Paper: Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, Xingchao Liu et al. (2023). Develops rectified flow to straighten ODE trajectories between arbitrary distributions, continuing the effort to build computationally efficient, straight-path dynamical generative models.
Paper: SE(3)-Stochastic Flow Matching for Protein Backbone Generation, Avishek Joey Bose et al. (2024). Applies optimal transport-guided continuous flows and stochastic bridges to non-Euclidean SE(3) manifold geometries for generative macromolecular modeling.
Paper: Stochastic Interpolants: A Unifying Framework for Flows and Diffusions, Michael S. Albergo et al. (2025). Unifies deterministic flows and stochastic diffusions bridging arbitrary boundary distributions in finite time through stochastic interpolants.

Cover for Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

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.

Topics

autoencoders generative model geometry inference ode stochastic data manifolds probabilistic inference

Keywords

Manifold Interpolating Optimal-Transport Flows trajectory inference optimal transport continuous normalizing flows diffusion maps geodesic autoencoder entropic interpolation neural ordinary differential equations manifold learning

1 Introduction
2 Preliminaries and Background
3 Manifold Interpolating Optimal-Transport Flow
3.1 Geodesic Autoencoder Embedding
3.2 Inferring Trajectories
4 Results
4.1 Artificial Data
4.2 Single-Cell Data
5 Related Work
6 Conclusion
7 Limitations and Broader Impact
Acknowledgments and Disclosure of Funding
References

Knowls

Knowl 1 — Diffusion Geodesic Distance
definition
Let $\mathcal{X} = \{x_1, \dots, x_n\} \subset \mathbb{R}^k$ be a discrete dataset sampled from an underlying Riemannian manifold $\mathcal{M}$ . Let $K_\epsilon \in \mathbb{R}^{n \times n}$ be an affinity matrix with entries $(K_\epsilon)_{ij} = k_\epsilon(x_i, x_j)$ , where $k_\epsilon$ is a kernel function (such as a Gaussian kernel or an $\alpha$ -decay kernel with bandwidth $\epsilon$ ). Let $Q$ be a diagonal matrix with $Q_{ii} = \sum_j (K_\epsilon)_{ij}$ , and let $M_\epsilon = Q^{-1} K_\epsilon Q^{-1}$ denote the density-normalized kernel matrix. The row-stochastic Markov diffusion operator is defined as $P_\epsilon = D^{-1} M_\epsilon$ , where $D$ is a diagonal matrix with $D_{ii} = \sum_{j=1}^n (M_\epsilon)_{ij}$ . The stationary distribution of the Markov chain is $\pi$ , with $\pi_i = D_{ii} / \sum_j D_{jj}$ .

Let $(P_\epsilon)_{i:}^{2^k}$ denote the $i$ -th row of the $2^k$ -step transition matrix $(P_\epsilon)^{2^k}$ , representing the transition probability distribution after $2^k$ random walk steps starting at $x_i$ .

For $\alpha \in (0, 1/2)$ and a maximum scale integer $K \ge 0$ , the diffusion geodesic distance $G_\alpha(x_i, x_j)$ between two points $x_i, x_j \in \mathcal{X}$ is defined as:

$G_\alpha(x_i, x_j) := \sum_{k=0}^K 2^{-(K-k)\alpha} \| (P_\epsilon)_{i:}^{2^k} - (P_\epsilon)_{j:}^{2^k} \|_1 + 2^{-(K+1)/2} \|\pi_i - \pi_j\|_1$

This multiscale metric accumulates $L^1$ differences between diffusion distributions across dyadic time scales $2^0, 2^1, \dots, 2^K$ , discretely approximating the continuous heat kernel diffusion ground distance on $\mathcal{M}$ .
Knowl 2 — Geodesic Autoencoder Architecture and Objective
model/method
The Geodesic Autoencoder (GAE) embeds high-dimensional data $\mathcal{X} \subset \mathbb{R}^k$ into a latent space $\mathcal{Z} = \mathbb{R}^d$ ( $d \ll k$ ) such that the Euclidean distance between latent representations preserves the diffusion geodesic distance $G_\alpha$ .

The network comprises an encoder $\phi : \mathbb{R}^k \to \mathcal{Z}$ and an optional decoder $\phi^{-1} : \mathcal{Z} \to \mathbb{R}^k$ . Given a batch of $N$ subsampled points $\{x_1, \dots, x_N\}$ , the encoder is trained to minimize the mean squared error (MSE) distance matching loss:

$L(\phi) := \frac{2}{N} \sum_{i=1}^N \sum_{j > i} \left( \|\phi(x_i) - \phi(x_j)\|_2 - G_\alpha(x_i, x_j) \right)^2$

When ambient space reconstruction is required (e.g., to trace individual gene expression over time), the autoencoder is trained with an additional reconstruction loss:

$L_r := \sum_{x \in \mathcal{X}} \|\phi^{-1}(\phi(x)) - x\|_2$

Because $G_\alpha$ is calculated on subsamples of size $N$ , the computational overhead of powering the full $n \times n$ diffusion matrix is eliminated, allowing evaluation across multiple dyadic diffusion scales while maintaining inductive generalization and denoising capability.
Knowl 3 — Convergence of Diffusion Geodesic Distance to Riemannian Manifold Geodesic
theoretical result
Let $(\mathcal{M}, d_\mathcal{M})$ be a closed $m$ -dimensional Riemannian manifold with geodesic distance $d_\mathcal{M}$ , and let $\mathcal{X} = \{x_1, \dots, x_n\}$ be samples drawn from $\mathcal{M}$ .

For any $\alpha \in (0, 1/2)$ , sufficiently large dyadic scale limit $K$ , sample size $N$ , and kernel bandwidth $\epsilon \to 0$ , the discrete diffusion geodesic distance $G_\alpha(x_i, x_j)$ approximates the manifold geodesic distance with high probability:

$G_\alpha(x_i, x_j) \simeq d_\mathcal{M}^{2\alpha}(x_i, x_j) \quad \forall x_i, x_j \in \mathcal{X}$

Consequently, if an encoder $\phi : \mathbb{R}^k \to \mathcal{Z}$ achieves zero distance matching loss $L(\phi) = 0$ , the Euclidean distance in the latent space satisfies:

$\|\phi(x_i) - \phi(x_j)\|_2 \simeq d_\mathcal{M}^{2\alpha}(x_i, x_j) \quad \forall x_i, x_j \in \mathcal{X}$

When $\alpha$ is chosen close to $1/2$ , the Euclidean metric $\|\phi(x_i) - \phi(x_j)\|_2$ in the latent embedding space $\mathcal{Z}$ is equivalent to the Riemannian geodesic distance $d_\mathcal{M}(x_i, x_j)$ on the data manifold.
Knowl 4 — Dynamic Optimal Transport on Latent Manifold Path Space
theoretical result
Let $\mu$ and $\nu$ be two probability distributions on a data manifold, and let $f(x, t)$ be a time-varying vector field defining trajectories via the ordinary differential equation $dX_t = f(X_t, t)dt$ with time-evolving density $\rho_t$ . Let $D(\mu, \nu)$ be a distribution dissimilarity measure satisfying $D(\mu, \nu) = 0$ if and only if $\mu = \nu$ .

There exists a sufficiently large regularization weight $\lambda > 0$ such that the 2-Wasserstein distance $W_2(\mu, \nu)^2$ equals the penalized path-space energy functional:

$W_2(\mu, \nu)^2 = \inf_{X_t} \mathbb{E}\left[ \int_0^1 \|f(X_t, t)\|_2^2 \, dt \right] + \lambda D(\rho_1, \nu) \quad \text{subject to } X_0 \sim \mu$

Furthermore, when the trajectories $X_t$ are defined within the Geodesic Autoencoder latent space $\mathcal{Z}$ whose Euclidean metric preserves the diffusion geodesic distance $G_\alpha \simeq d_\mathcal{M}^{2\alpha}$ , the dynamic 2-Wasserstein transport in $\mathcal{Z}$ is equivalent to dynamic optimal transport with respect to the manifold geodesic ground distance:

$W_2(\mu, \nu) \simeq W_{d_\mathcal{M}^{2\alpha}}(\mu, \nu)$
Knowl 5 — MIOFlow Trajectory Inference Loss Functions
equation
Given $T$ discrete empirical distributions $\mu_i = \frac{1}{n_i} \sum_{x \in X_i} \delta_x$ observed at timepoints $i \in \{0, 1, \dots, T-1\}$ , a Neural ODE parameterized by vector field $f_\theta(x, t)$ produces transported empirical distributions $\hat{\mu}_i = \frac{1}{n_i} \sum_{x \in \hat{X}_i} \delta_x$ , where $\hat{X}_i = \psi_\theta(X_0, i)$ .

MIOFlow trains the Neural ODE parameters $\theta$ by minimizing the composite loss $L = L_m + L_e + L_d$ , comprising:
1. Marginal matching loss (enforcing endpoint distribution consistency via discrete 2-Wasserstein distance): $L_m := \sum_{i=1}^{T-1} W_2(\hat{\mu}_i, \mu_i)$
2. Kinetic energy regularization loss (minimizing trajectory path energy): $L_e := \lambda_e \sum_{i=1}^{T-1} \int_{i-1}^i \|f_\theta(x_t, t)\|_2^2 \, dt$
3. Manifold density loss (penalizing predicted points that drift away from the data manifold support): $L_d := \lambda_d \sum_{t=1}^{T-1} \sum_{x \in \hat{X}_t} \ell_d(x, t)$ $\ell_d(x, t) := \sum_{i=1}^k \max\left(0, \text{min-k}(\{\|x - y\| : y \in X_t\}) - h\right)$ where $\text{min-k}$ denotes the distance to the $k$ -nearest neighbors in the observed sample $X_t$ , $h > 0$ is a distance threshold, and $\lambda_e, \lambda_d > 0$ are regularization hyperparameters.
Knowl 6 — Stochastic Population Dynamics via Learnable Diffusion Annealing
model/method
To account for stochasticity in biological differentiation trajectories (such as cell fate bifurcations), MIOFlow models particle paths via a stochastic differential equation (SDE):

$dX_t = f_\theta(X_t, t)dt + \sqrt{\sigma_t} \, dB_t$

where $B_t$ is standard Brownian motion and $\{\sigma_t\}_{t=0}^{T-1}$ is a set of $T$ learnable time-dependent variance parameters.

Early in training, non-zero diffusion $\sigma_t > 0$ disperses initial samples $X_0$ , promoting exploration across divergent paths and bifurcating branches. As optimization progresses, the learned diffusion parameters empirically converge toward zero ( $\sigma_t \to 0$ ). In the zero-diffusion limit, the trajectory paths converge to deterministic dynamic 2-Wasserstein optimal transport geodesics. For small $\sigma_t$ , the flow is numerically integrated using standard ODE solvers (analogous to the Euler–Maruyama scheme), avoiding the severe computational cost of general SDE solvers.
Knowl 7 — MIOFlow Training Algorithm
algorithm
MIOFlow infers continuous population trajectories by first training a Geodesic Autoencoder (GAE) to project samples onto a geodesic-preserving latent manifold, followed by two-stage local and global Neural ODE training.
Input: Datasets $X_0, \dots, X_{T-1}$ , graph kernel $K_\epsilon$ , maximum scale $K$ , noise scale $\xi$ , batch size $N$ , max GAE iterations $n_{\max}$ , initialized encoder $\phi$ and decoder $\phi^{-1}$ , local iterations $n_{\text{local}}$ , global iterations $n_{\text{global}}$ , initialized neural ODE $\psi_\theta$ , loss weights $\lambda_e, \lambda_d$ , lower bound $h$ .
Output: Trained Neural ODE $\psi_\theta$ .
$\phi, \phi^{-1} \leftarrow \text{GAE}(X, K_\epsilon, K, \xi, N, n_{\max}, \phi, \phi^{-1})$
if using GAE embedding then
    for $t = 0$ to $T-1$ do
         $X_t \leftarrow \phi(X_t)$
    end for
end if
for $i = 1$ to $n_{\text{local}}$ do
    for $t = 0$ to $T - 2$ do
         $\hat{X}_{t+1} \leftarrow \psi_\theta(X_t, t + 1)$
    end for
     $L_m \leftarrow \sum_{t=0}^{T-2} W_2(\hat{\mu}_{t+1}, \mu_{t+1})$
     $L_e \leftarrow \lambda_e \sum_{t=0}^{T-2} \int_t^{t+1} \|f_\theta(x_u, u)\|_2^2 \, du$
     $L_d \leftarrow \lambda_d \sum_{t=0}^{T-2} \sum_{x \in \hat{X}_{t+1}} \ell_d(x, t+1)$
     $L \leftarrow L_m + L_e + L_d$
    Update $\theta$ via gradient descent on $L$
end for
for $i = 1$ to $n_{\text{global}}$ do
     $\hat{X}_1, \dots, \hat{X}_{T-1} \leftarrow \psi_\theta(X_0, \{1, \dots, T-1\})$
     $L_m \leftarrow \sum_{t=1}^{T-1} W_2(\hat{\mu}_t, \mu_t)$
     $L_e \leftarrow \lambda_e \sum_{t=1}^{T-1} \int_{t-1}^t \|f_\theta(x_u, u)\|_2^2 \, du$
     $L_d \leftarrow \lambda_d \sum_{t=1}^{T-1} \sum_{x \in \hat{X}_t} \ell_d(x, t)$
     $L \leftarrow L_m + L_e + L_d$
    Update $\theta$ via gradient descent on $L$
end for
return $\psi_\theta$

Knowl 8 — Quantitative Trajectory Interpolation Performance on Synthetic and Simulated Datasets

data/table

The trajectory interpolation accuracy of MIOFlow was benchmarked against TrajectoryNet (continuous normalizing flow) and Diffusion Schrödinger's Bridge (DSB) on two synthetic datasets with complex bifurcations: the 2D Petal dataset and the 5D Dyngen single-cell simulator dataset. Models were evaluated by holding out timepoint $t=2$ and computing 1-Wasserstein distance ( $W_1$ ), Gaussian-kernel Maximum Mean Discrepancy (MMD(G)), linear-kernel Maximum Mean Discrepancy (MMD(M)), and total training runtime.

	Petal				Dyngen
Method	$W_1$	MMD(G)	MMD(M)	Runtime	$W_1$	MMD(G)	MMD(M)	Runtime
MIOFlow (ours)	0.090	0.029	0.005	284.60 s	0.783	0.199	0.509	95.63 s
TrajectoryNet	0.181	0.136	0.002	64 m	1.499	0.303	1.806	62 m
DSB	0.199	0.157	0.008	86 m	2.051	0.932	2.367	81 m
Baseline	0.221	0.196	$<$ 0.001	N/A	1.388	1.008	0.926	N/A

MIOFlow achieved lower $W_1$ and MMD(G) errors on both datasets while running in under 5 minutes, compared to over an hour for TrajectoryNet and DSB. The efficiency gain stems from avoiding the $O(k^2)$ Jacobian trace evaluations required by continuous normalizing flows and computing optimal transport directly on the manifold support.

Knowl 9 — Ablation of Geodesic Autoencoder Embeddings on Embryoid Body scRNA-seq Data

data/table

The impact of the Geodesic Autoencoder (GAE) manifold embedding was evaluated on single-cell RNA sequencing data of human embryoid body (EB) differentiation over 27 days (reduced to 200 PCA components). Leave-one-out interpolation accuracy was assessed at timepoints $t=2$ (days 12-15) and $t=3$ (days 18-21) using $W_1$ , MMD(G), and MMD(M) metrics comparing GAE with a Gaussian kernel (scale 0.5), GAE with an $\alpha$ -decay kernel (30-nearest neighbor graph), and training without GAE embedding (No GAE).

	$t = 2$			$t = 3$
Configuration	$W_1$	MMD(G)	MMD(M)	$W_1$	MMD(G)	MMD(M)
GAE Gaussian	26.318	0.059	101.969	32.963	0.138	223.679
GAE $\alpha$ -decay	25.744	0.061	99.747	32.227	0.135	213.465
No GAE	29.243	0.063	126.608	35.709	0.119	246.609
Baseline	33.415	0.103	227.279	35.319	0.095	213.492

Embedding the data with GAE prior to flow learning consistently reduced $W_1$ distance and linear MMD across held-out timepoints compared to operating without GAE, with the $\alpha$ -decay kernel yielding the best overall $W_1$ accuracy ( $25.744$ at $t=2$ and $32.227$ at $t=3$ ).

Knowl 10 — Limitations of MIOFlow
limitation
MIOFlow has two main limitations:
1. Numerical stiffness in ODE solvers: The method relies on generic ODE numerical solvers and adjoint sensitivity methods; if the learned velocity field $f_\theta(x, t)$ exhibits stiff dynamics during training, ODE integration can fail or suffer from diminished prediction accuracy.
2. Sensitivity to time-bin uniformity: The optimal transport interpolation assumes relatively uniform temporal sampling across snapshot timepoints. Highly irregular sampling or variable population growth/death rates can skew the transport flow (a limitation that could potentially be addressed by replacing balanced optimal transport with unbalanced optimal transport).

Coverage note — None was omitted; all key theoretical definitions, convergence theorems, model formulations, optimization algorithms, empirical benchmarking tables, and stated limitations are included.

References

1.Andrei Agrachev and Paul WY Lee. Continuity of optimal control costs and its application to weak kam theory. Calculus of Variations and Partial Differential Equations, 39(1):213–232, 2010.
2.Mikhail Belkin and Partha Niyogi. Semi-Supervised Learning on Riemannian Manifolds. Machine Learning, 56(1-3):209–239, 2004.
3.Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375–393, 2000.
4.Volker Bergen, Ruslan A. Soldatov, Peter V. Kharchenko, and Fabian J. Theis. RNA velocity—current challenges and future perspectives. Molecular Systems Biology, 17(8):e10282, 2021.
5.Charlotte Bunne, Laetitia Meng-Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal Optimal Transport Modeling of Population Dynamics. AISTATS, 2022.
6.Robrecht Cannoodt, Wouter Saelens, Louise Deconinck, and Yvan Saeys. Spearheading future omics analyses using dyngen, a multi-modal simulator of single cells. Nature Communications, 12(1):1–9, 2021.
7.Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018.
8.Xiuyuan Cheng and Nan Wu. Eigen-convergence of gaussian kernelized graph laplacian by manifold heat interpolation. Applied and Computational Harmonic Analysis, 61:132–190, 2022.
9.Ronald R. Coifman and Stéphane Lafon. Diffusion maps. Applied and Computational Harmonic Analysis, 21(1):5–30, 2006.
10.Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34, 2021.
11.Andrés F Duque, Sacha Morin, Guy Wolf, and Kevin Moon. Extendable and invertible manifold learning with geometry regularized autoencoders. In 2020 IEEE International Conference on Big Data (Big Data), pages 5027–5036. IEEE, 2020.
12.Katie A Fennell, Dane Vassiliadis, Enid YN Lam, Luciano G Martelotto, Jesse J Balic, Sebastian Hollizeck, Tom S Weber, Timothy Semple, Qing Wang, Denise C Miles, et al. Non-genetic determinants of malignant clonal fitness at single-cell resolution. Nature, 601(7891):125–131, 2022.
13.Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. In International conference on machine learning, pages 3154–3164. PMLR, 2020.
14.Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1–8, 2021.
15.Crispin W Gardiner et al. Handbook of stochastic methods, volume 3. springer Berlin, 1985.
16.Nassif Ghoussoub, Young-Heon Kim, and Aaron Zeff Palmer. A solution to the monge transport problem for brownian martingales. the Annals of Probability, 49(2):877–907, 2021.
17.Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models. In 7th International Conference on Learning Representations, 2019.
18.Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel method for the two-sample problem. Journal of Machine Learning Research, 1:1–10, 2008.
19.Laleh Haghverdi, Maren Büttner, F. Alexander Wolf, Florian Buettner, and Fabian J. Theis. Diffusion pseudotime robustly reconstructs lineage branching. Nature Methods, 13(10):845–848, 2016.
20.Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative rnns. In International Conference on Machine Learning, pages 2417–2426. PMLR, 2016.
21.Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
22.Takeshi Koshizuka and Issei Sato. Neural lagrangian schrödinger bridge. arXiv preprint arXiv:2204.04853, 2022.
23.Gioele La Manno, Ruslan Soldatov, Amit Zeisel, Emelie Braun, Hannah Hochgerner, Viktor Petukhov, Katja Lidschreiber, Maria E. Kastriti, Peter Lönnerberg, Alessandro Furlan, Jean Fan, Lars E. Borm, Zehua Liu, David van Bruggen, Jimin Guo, Xiaoling He, Roger Barker, Erik Sundström, Gonçalo Castelo-Branco, Patrick Cramer, Igor Adameyko, Sten Linnarsson, and Peter V. Kharchenko. RNA velocity of single cells. Nature, 560(7719):494–498, 2018.
24.Jean-Michel Lasry and Pierre-Louis Lions. Mean field games. Japanese journal of mathematics, 2(1): 229–260, 2007.
25.William Leeb and Ronald Coifman. Hölder–Lipschitz Norms and Their Duals on Spaces with Semigroups, with Applications to Earth Mover’s Distance. J Fourier Anal Appl, 22(4):910–953, 2016.
26.Jun Lei and Marthe J Howard. Targeted deletion of hand2 in enteric neural precursor cells affects its functions in neurogenesis, neurotransmitter specification and gangliogenesis, causing functional aganglionosis. Development, 138(21):4789–4800, 2011.
27.Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503–2515, 2020.
28.Toshio Mikami. Monge’s problem with a quadratic cost by the zero-noise limit of h-path processes. Probability theory and related fields, 129(2):245–260, 2004.
29.Toshio Mikami and Michèle Thieullen. Duality theorem for the stochastic optimal control problem. Stochastic processes and their applications, 116(12):1815–1835, 2006.
30.Gal Mishne, Uri Shaham, Alexander Cloninger, and Israel Cohen. Diffusion nets. Applied and Computational Harmonic Analysis, 47(2):259–285, 2019.
31.Kevin R. Moon, Jay S. Stanley, Daniel Burkhardt, David van Dijk, Guy Wolf, and Smita Krishnaswamy. Manifold learning-based methods for analyzing single-cell RNA-sequencing data. Current Opinion in Systems Biology, 7:36–46, 2018.
32.Kevin R. Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel B. Burkhardt, William S. Chen, Kristina Yim, Antonia van den Elzen, Matthew J. Hirn, Ronald R. Coifman, Natalia B. Ivanova, Guy Wolf, and Smita Krishnaswamy. Visualizing structure and transitions in high-dimensional biological data. Nat Biotechnol, 37(12):1482–1492, 2019.
33.Michele Pavon, Giulio Trigila, and Esteban G Tabak. The data-driven schrödinger bridge. Communications on Pure and Applied Mathematics, 74(7):1545–1573, 2021.
34.Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. arXiv:1803.00567, 2020.
35.Neha Prasad, Karren Yang, and Caroline Uhler. Optimal transport using gans for lineage tracing. arXiv preprint arXiv:2007.12098, 2020.
36.Jack Richter-Powell, Yaron Lipman, and Ricky TQ Chen. Neural conservation laws: A divergence-free perspective. arXiv preprint arXiv:2210.01741, 2022.
37.Wouter Saelens, Robrecht Cannoodt, Helena Todorov, and Yvan Saeys. A comparison of single-cell trajectory inference methods. Nature biotechnology, 37(5):547–554, 2019.
38.Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Siyan Liu, Stacie Lin, Peter Berube, Lia Lee, Jenny Chen, Justin Brumbaugh, Philippe Rigollet, Konrad Hochedlinger, Rudolf Jaenisch, Aviv Regev, and Eric S. Lander. Reconstruction of developmental landscapes by optimal-transport analysis of single-cell gene expression sheds light on cellular reprogramming. Cell, 2019.
39.A. Singer. From graph to manifold Laplacian: The convergence rate. Applied and Computational Harmonic Analysis, 21(1):128–134, 2006.
40.Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1–11, 2015.
41.Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. Advances in Neural Information Processing Systems, pages 11918–11930, 2019.
42.Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
43.Kelly Street, Davide Risso, Russell B Fletcher, Diya Das, John Ngai, Nir Yosef, Elizabeth Purdom, and Sandrine Dudoit. Slingshot: cell lineage and pseudotime inference for single-cell transcriptomics. BMC genomics, 19(1):1–16, 2018.
44.Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics. In Proceedings of the 37th International Conference on Machine Learning, pages 9526–9536. PMLR, 2020.
45.Alexander Tong, Guillaume Huguet, Dennis Shung, Amine Natik, Manik Kuchroo, Guillaume Lajoie, Guy Wolf, and Smita Krishnaswamy. Embedding signals on graphs with unbalanced diffusion earth mover’s distance. In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 5647–5651, 2022.
46.Alexander Y Tong, Guillaume Huguet, Amine Natik, Kincaid MacDonald, Manik Kuchroo, Ronald Coifman, Guy Wolf, and Smita Krishnaswamy. Diffusion earth mover’s distance and distribution embeddings. In International Conference on Machine Learning, pages 10336–10346. PMLR, 2021.
47.Cole Trapnell, Davide Cacchiarelli, Jonna Grimsby, Prapti Pokharel, Shuqiang Li, Michael Morse, Niall J Lennon, Kenneth J Livak, Tarjei S Mikkelsen, and John L Rinn. Pseudo-temporal ordering of individual cells reveals dynamics and regulators of cell fate decisions. Nature biotechnology, 32(4):381, 2014.
48.Jori van der Raadt, Sebastianus HC van Gestel, Nael Nadif Kasri, and Cornelis A Albers. Onecut transcription factors induce neuronal characteristics and remodel chromatin accessibility. Nucleic acids research, 47(11):5587–5602, 2019.
49.Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schrödinger Bridges via Maximum Likelihood. Entropy, 23(9):1134, 2021.
50.Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
51.Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International Conference on Machine Learning, pages 10794–10804. PMLR, 2021.
52.Stephen Zhang, Anton Afanassiev, Laura Greenstreet, Tetsuya Matsumoto, and Geoffrey Schiebinger. Optimal transport analysis reveals trajectories in steady-state systems. Technical report, 2021.
53.Ziqi Zhang and Xiuwei Zhang. Inference of high-resolution trajectories in single-cell RNA-seq data by using RNA velocity. Cell Reports Methods, 1(6):100095, 2021.

Citation

MLA

Huguet, G., et al. “Manifold Interpolating Optimal-Transport Flows for Trajectory Inference”. Advances in Neural Information Processing Systems, vol. 35, 2022, pp. 29705–18, https://proceedings.neurips.cc/paper_files/paper/2022/file/bfc03f077688d8885c0a9389d77616d0-Paper-Conference.pdf.

APA

Huguet, G., Magruder, D. S., Tong, A., Fasina, O., Kuchroo, M., Wolf, G., & Krishnaswamy, S. (2022). Manifold Interpolating Optimal-Transport Flows for Trajectory Inference. Advances in Neural Information Processing Systems, 35, 29705–29718. https://proceedings.neurips.cc/paper_files/paper/2022/file/bfc03f077688d8885c0a9389d77616d0-Paper-Conference.pdf

Chicago

Huguet, G., D. S. Magruder, A. Tong, et al. 2022. “Manifold Interpolating Optimal-Transport Flows for Trajectory Inference”. Advances in Neural Information Processing Systems 35: 29705–18. https://proceedings.neurips.cc/paper_files/paper/2022/file/bfc03f077688d8885c0a9389d77616d0-Paper-Conference.pdf.

Harvard

Huguet, G. et al. (2022) “Manifold Interpolating Optimal-Transport Flows for Trajectory Inference”, Advances in Neural Information Processing Systems. Curran Associates, Inc., pp. 29705–29718. Available at: https://proceedings.neurips.cc/paper_files/paper/2022/file/bfc03f077688d8885c0a9389d77616d0-Paper-Conference.pdf.

Vancouver

1. Huguet G, Magruder DS, Tong A, Fasina O, Kuchroo M, Wolf G, Krishnaswamy S (2022) Manifold Interpolating Optimal-Transport Flows for Trajectory Inference. In: Advances in Neural Information Processing Systems. Curran Associates, Inc., pp 29705–29718

BibTeX

@inproceedings{huguet2022manifold,
  title = {Manifold Interpolating Optimal-Transport Flows for Trajectory Inference},
  author = {Huguet, Guillaume and Magruder, Daniel Sumner and Tong, Alexander and Fasina, Oluwadamilola and Kuchroo, Manik and Wolf, Guy and Krishnaswamy, Smita},
  year = {2022},
  booktitle = {Advances in Neural Information Processing Systems},
  publisher = {Curran Associates, Inc.},
  volume = {35},
  pages = {29705-29718},
  url = {https://proceedings.neurips.cc/paper_files/paper/2022/file/bfc03f077688d8885c0a9389d77616d0-Paper-Conference.pdf}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Authors

Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Diffusion Geodesic Distance

Knowl 2 — Geodesic Autoencoder Architecture and Objective

Knowl 3 — Convergence of Diffusion Geodesic Distance to Riemannian Manifold Geodesic

Knowl 4 — Dynamic Optimal Transport on Latent Manifold Path Space

Knowl 5 — MIOFlow Trajectory Inference Loss Functions

Knowl 6 — Stochastic Population Dynamics via Learnable Diffusion Annealing

Knowl 7 — MIOFlow Training Algorithm