Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

1 Introduction

• MIOFlow targets continuous dynamics of probability distributions on data manifolds, motivated by high-dimensional measurements with low intrinsic dimensionality.The introduction uses single-cell transcriptomic profiles as an example of high-dimensional measurements with informational redundancy.Source: §1 · p. 1
• Static cellular snapshots leave interpolation at unmeasured times and inference of individual manifold-following trajectories as key tasks.Measurements are destructive, so current technology cannot follow a cell through time.Source: §1 · pp. 1–2
• MIOFlow uses Neural ODEs to transport sampled high-dimensional data between timepoints, with manifold-constrained transport, Wasserstein penalties against measured populations, and stochastic trajectories.The framework is designed to interpolate populations and infer individual trajectories from static snapshots.Source: §1 · p. 2
• Prior trajectory methods include TrajectoryNet, recurrent neural networks, and Waddington-OT; the section contrasts MIOFlow with generative flows from noise to data.The authors identify Gaussian initialization, deterministic trajectories, and Jacobian-trace computation as drawbacks of continuous normalizing flows.Source: §1 · p. 2
• The Geodesic Autoencoder embeds data so latent distances match a multiscale manifold distance, while manifold optimal transport penalizes the learned flow.The two steps enforce flows on the manifold, unlike approaches that operate in ambient dimensions.Source: §1 · p. 2
• The stated contributions include a diffusion geodesic distance with a convergence result and empirical modeling of divergent trajectories on toy and single-cell data.The distance is described as converging to a manifold geodesic as the number of points tends to infinity.Source: §1 · p. 2

Problem Formulation and Notation

• The problem observes samples from a finite set of time-indexed distributions and seeks trajectories from an initial sample that match distributions at observed times.The formulation also aims to characterize the evolving support of the population.Source: §2 · p. 2
• The assumptions include absolute continuity with respect to Lebesgue measure, equivalence of the compared distances, and existence of solutions for the SDEs.Proofs are presented in the supplementary material.Source: §2 · p. 2

Optimal Transport

• The Benamou–Brenier dynamic formulation replaces a transport plan with an evolving distribution and vector field satisfying the continuity equation.The formulation connects endpoint distributions through the time evolution of density and mass movement.Source: §2 · p. 3

Adding Diffusion

• The stochastic formulation adds Brownian diffusion and uses the corresponding Fokker–Planck equation; as diffusion vanishes, its infimum converges to W2.The paper uses this diffusion term while retaining convergence to transport between distributions.Source: §2 · p. 3

Manifold Geodesics

• Dynamic optimal transport with an ambient Euclidean ground distance can ignore data-manifold geometry; an embedding is sought whose Euclidean distances correspond to manifold geodesics.The section uses this embedding to make dynamic optimal transport equivalent to Wasserstein transport with a geodesic ground distance.Source: §2 · p. 4
• The diffusion ground distance compares heat-kernel-induced measures, while a multiscale diffusion distance relates to manifold geodesic distance.For β = 1, the diffusion operator is described as converging to the heat operator as ϵ approaches zero.Source: §2 · pp. 3–4

3 Manifold Interpolating Optimal-Transport Flow

• MIOFlow uses Neural ODEs to learn continuous trajectories between snapshots, matching population marginals with an optimal-transport loss and regularizing the vector field.Training can predict only the next snapshot locally or the full trajectory globally from the initial condition.Source: §3 · p. 4, §3.2 · p. 6

3.1 Geodesic Autoencoder Embedding

• The Geodesic Autoencoder regularizes latent distances to approximate a sample-based diffusion geodesic that compares random-walk transition probabilities across scales.The diffusion geodesic distance is used to perform dynamic optimal transport in the learned space.Source: §3.1 · p. 5

3.2 Inferring Trajectories

• In the embedded space, Euclidean Wasserstein matching is equivalent to transport using the manifold geodesic distance under the stated result.The theorem motivates matching marginals and regularizing the vector field to recover the dynamic W2 formulation.Source: §3.2 · p. 6
• A density loss penalizes predicted points for being distant from their k-nearest neighbors in the observed timepoint data.The loss uses a lower bound h > 0.Source: §3.2 · p. 6
• Adding diffusion yields stochastic trajectories that empirically promote early bifurcation, while the learned diffusion scale approaches zero as training continues.Modeling complex diffusion would require an SDE solver and substantially increase computational cost.Source: §3.2 · p. 7

4 Results

• The evaluation compares qualitative trajectory fit and held-out-timepoint Wasserstein W1 and MMD, using TrajectoryNet and DSB as population-flow baselines.The reported baseline measure is the average distance between the previous and next timepoints for a given time.Source: §4 · p. 7

4.1 Artificial Data

• MIOFlow is more accurate than TrajectoryNet and DSB on held-out interpolation for both Petal and Dyngen, while requiring much less computational time.On Petal it improves W1 and MMD(G); on Dyngen it is more accurate for all metrics and learns continuous trajectories along both main branches.Source: Table 1 · p. 8, §4.1 · p. 8

4.2 Single-Cell Data

• MIOFlow captures HAND2’s non-monotonic expression trend—rising from embryonic stem cells to neural early progenitors, then falling as lineages mature—while TrajectoryNet shows more chaotic gene trends and a tendency toward early divergence.ONECUT2 is described as higher in cells diverging toward the neuronal lineage.Source: §4.2 · p. 9
• For almost all EB prediction metrics, accuracy is greater with GAE than without embedding.The comparison includes Gaussian and α-decay kernels; lower values are better in Table 2.Source: §4.2 · p. 9, Table 2 · p. 9
• In AML treatment data, most cells surviving to the last timepoint increase SLPI, MYB, and B2M, markers of the leukemic stem-cell pathway.The authors interpret this as cells transitioning to the signature associated with resistance to chemotherapy.Source: §4.2 · p. 9, Figure 4 · p. 9

• Static approaches include pseudotime and lineage inference, optimal transport using cell-growth estimates, and RNA velocity combined with transition matrices; a separate study compares 45 static methods.Source: §5 · p. 10
• Dynamic optimal-transport approaches model trajectories with recurrent networks, generative adversarial networks, input-convex neural-network gradients, or continuous normalizing flows.Source: §5 · p. 10
• Schrödinger-bridge methods interpolate between distributions, but the cited example does not explicitly learn dynamics and its interpolations are described as inaccurate.Source: §5 · p. 10

• MIOFlow is presented as a computationally efficient method for recovering trajectories from multiple cross-sectional samples of a time-evolving population.Source: §6 · p. 10
• Its manifold-learning framework learns an embedding that preserves geodesic distance, and the paper theoretically relates MIOFlow to dynamical optimal transport and stochastic trajectories.Source: §6 · p. 10
• On artificial datasets, MIOFlow produced trajectories more in line with the manifold than DSB and TrajectoryNet, while scRNA-seq trajectories were decoded to visualize gene-expression evolution.Source: §6 · p. 10

7 Limitations and Broader Impact

• Generic ODE solvers and adjoint dynamics can fail under stiff dynamics, and training may produce stiffness that decreases model accuracy.Source: §7 · p. 10
• MIOFlow is sensitive to time bins and worked best with relatively uniform sampling; unbalanced optimal transport is suggested as a possible remedy.Source: §7 · p. 10
• The authors state that the work can help understand dynamical systems, particularly cellular differentiation, and presents no potential negative societal impacts to their knowledge.Source: §7 · p. 10

A.1 Geodesic Autoencoder Algorithm

• The GAE algorithm takes a dataset and graph-kernel and scale settings, samples noisy batches, computes diffusion-based pairwise distances, and updates encoder and decoder parameters.Source: §A.1 · p. 15

A.2 Proof of Thm.2

• For samples from a closed Riemannian manifold, the theorem states that Gα approximates squared geodesic distance when K and N are sufficiently large and ϵ is small.Source: §A.2 · pp. 15–16
• The proof has two parts: convergence of the powered diffusion operator to the heat operator, followed by approximation of the operator using a finite point set.Source: §A.2 · p. 15
• For sufficiently large n and K with ϵ = 2^-K, the argument relates Gα to the heat-operator distance and then to squared manifold geodesic distance.Source: §A.2 · pp. 15–16

A.3 Proof of Thm.3

• The theorem assumes a time-varying vector field generating trajectories and a distribution dissimilarity that is zero only for identical distributions.The proof also specifies continuity and endpoint constraints for the density path.Source: §A.3 · p. 16
• For trajectories in the embedded space, W2 is equivalent to Wasserstein distance with a geodesic ground distance.The proof derives this equivalence under its embedding assumptions.Source: §A.3 · pp. 16–17

A.4 Schrödinger bridge

• The Schrödinger bridge seeks evolving distributions between endpoint measures relative to Wiener measure, and its static form minimizes KL divergence over endpoint couplings.With Brownian variance σ2, the formulation is entropically regularized optimal transport and can be solved with Sinkhorn’s algorithm.Source: §A.4 · p. 17

B Experiment details

• Across experiments, the Neural ODE used RK4, three layers, two appended input dimensions and time, while the encoder used three ReLU layers.Training used AdamW and 20 batches per epoch; energy loss was omitted in practice because trajectories were already smooth and it required careful parametrization.Source: B Experiment details · p. 17
• Petal experiments used 8- and 32-unit encoder layers, a 16-32-16 ODE, and 40 local epochs with density-loss weight λd = 35.The hold-out setup used 30 local and 15 global epochs with λd = 1; diffusion-geodesic kernels used maximum time scale 25.Source: B Experiment details · pp. 17–18
• Dyngen used an RBF kernel, an 8- and 32-unit encoder, and a 16-32-16 ODE trained for 50 global epochs with λd = 5.For hold-out time t = 3, the adjacent intervals were both withheld and training used 50 global epochs; EB and AML used separate encoder, ODE, and training configurations.Source: B Experiment details · p. 18

Energy loss

• With energy loss λe = 1, Dyngen trajectories become straight and cannot bifurcate, while Petal trajectories collapse toward the center during training.The authors omit this loss by default because their trajectories were already smooth; careful parametrization may yield more accurate trajectories.Source: C Additional Results · p. 18

Ablation and hold-out

• The density loss and geodesic embedding both appear important for accurate trajectories; on Dyngen, they support manifold adherence and bifurcation, respectively.Using the GAE improves EB interpolation accuracy across all reported metrics.Source: C Additional Results · p. 19
• Holding out Dyngen timepoint t = 3 appears hardest to learn, and the hold-out evaluation compares predictions with ground truth using W1 and two MMD variants.The Petal and Dyngen Gaussian MMD averages scales 0.1 and 0.5.Source: C Additional Results · p. 19, Figure 9 · p. 20
• MIOFlow is more accurate on W1 than TrajectoryNet and DSB across Petal and Dyngen leave-one-out predictions.Table 3 reports W1 and Gaussian- or mean-kernel MMD between predicted and ground-truth points; lower is better.Source: Table 3 · p. 19, C Additional Results · p. 19

DSB on petal

• On Petal, DSB predictions collapse onto a subset with too little noise, while larger noise scales are inaccurate.The comparison evaluates different noise scales γ.Source: C Additional Results · p. 20

D Software Versions

• The listed software environment includes cudatoolkit 11.1.1, PyTorch 1.10.2, and torchdiffeq 0.2.2.The appendix also lists versions for graphtools, NumPy, PHATE, POT, scikit-learn, SciPy, scprep, and torchsde.Source: D Software Versions · p. 20
• The figures cover density-loss and GAE comparisons, held-out-time trajectories, Neural SDE versus MIOFlow on Petal and Dyngen, and DSB across noise scales.The captions identify comparisons but do not report their outcomes.Source: Figure 10 · p. 21, Figure 11 · p. 21, Figure 12 · p. 22
• Tables 4–6 describe EB interpolation metrics by GAE and kernel, 1-NN distances across methods and datasets, and Dyngen 10-NN distances.Table 7 lists GPU specifications and notes that results used multiple GPU generations and models.Source: Table 4 · p. 21, Table 5 · p. 22, Table 6 · p. 22