Original: Manifold Interpolating Optimal-Transport Flows for Trajectory Inference
Revised: Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Revised: AffiliationAddress

Original: Abstract
Revised: Abstract

Original: We present a method called Manifold Interpolating Optimal-Transport Flow(MIOFlow) that learns stochastic, continuous population dynamics from staticsnapshot samples taken at sporadic timepoints.
Revised: Here, 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.

Original: MIOFlow combines dynamic models, manifold learning, and optimal transport by training neural ordinary differentialequations (Neural ODE) to interpolate between static population snapshots as penalized by optimal transport with manifold ground distance.
Revised: MIOFlow combines dynamic models,manifold learning, and optimal transport by training neural ordinary differentialequations (Neural ODE) to interpolate between static population snapshots aspenalized by optimal transport with manifold ground distance.

Original: Further, we ensure thatthe flow follows the geometry by operating in the latent space of an autoencoderthat we call a geodesic autoencoder (GAE).
Revised: Further, we ensurethat the flow follows the geometry by operating in the latent space of an autoencoder that we call a geodesic autoencoder (GAE).

Original: In GAE the latent space distancebetween points is regularized to match a novel multiscale geodesic distance on thedata manifold that we define.
Revised: In GAE the latent space distancebetween points is regularized to match a novel multiscale geodesic distance on thedata manifold that we define.

Original: We show that this method is superior to normalizingflows, Schrödinger bridges and other generative models that are designed to flowfrom noise to data in terms of interpolating between populations.
Revised: We show that this method is superior to normalizingflows, Schrödinger bridges and other generative models that are designed to flowfrom noise to data in terms of interpolating between populations.

Original: Theoretically, welink these trajectories with dynamic optimal transport.
Revised: Theoretically, welink these trajectories with dynamic optimal transport.

Original: We evaluate our methodon simulated data with bifurcations and merges, as well as scRNA-seq data fromembryoid body differentiation, and acute myeloid leukemia treatment.
Revised: We evaluate our methodon simulated data with bifurcations and merges, as well as scRNA-seq data fromembryoid body differentiation, and acute myeloid leukemia treatment.

Original:  Introduction
Revised:  Introduction

Original: Here, we tackle the problem of continuous dynamics of probability distributions defined on a datamanifold.

Original: Data from naturalistic systems are often modeled as generated from an underlying lowdimensional manifold embedded in a high dimensional measurement space.

Original: Termed the manifoldhypothesis, this assumption has led to many successful models of biological, chemical and physicalsystems.

Original: Measurements in such systems are increasingly high dimensional.

Original: For instance, in singlecell data the entire transcriptomic profile of the cell is measured with each mRNA or gene speciesa dimension.

Original: However, because of informational redundancy between these genes, the intrinsicdimensionality of the data is low dimensional.

Original: Mathematically, a Riemannian manifold is a goodmodel for such a system and much of manifold learning literature including, diffusion maps [9],

Original: Laplacian eigenmaps [2] and the wider graph signal processing literature, has focused on learningthis structure via data graphs [28, 29] and autoencoders [27, 11].

Original: Building on this literature, we consider snapshot measurements of cellular populations over time.

Original: Current technology cannot follow a cell in time as the measurements are destructive.

Original: We frame this as learning a population flow on a data manifold. Recently several neural networks that perform flows, or transports between probability distributions, have been proposed in the literature.
Revised: Recently several neural networks that perform flows, or transports between probability distributions,have been proposed in the literature.

Original: However, most of these works have focused on generative modeling: i.e., flowing from a noise distribution such as a Gaussian distribution to a data distribution in order to generate data.
Revised: However, most have these works have focused on generativemodeling: i e., flowing from a noise distribution such as a Gaussian distribution to a data distributionin order to generate data.

Original: Examples include score-based generative matching [38, 39], diffusion models [19], Schrödinger bridges [30, 10], and continuousnormalizing flows (CNF) 
Revised: Examples include score-based generative matching [32, 33], diffusionmodels [17], Schrödinger bridges [26, 9] and continuous normalizing flows (CNF) [6, 14].

Original: However, here we focus on learning continuous dynamics of suchsystems using static snapshot measurements, with a key task being interpolation in unmeasured times,as well as inference of individual trajectories that follow the manifold structure.
Revised: However,here we focus on a different problem—learning the continuous dynamics of probability distributionsthat are measured as static snapshots over time.

Original: To continuously interpolate populations over  time on data with manifold structure, we propose
Revised: Such distributions occur in naturalistic systems wherepopulations are measured at sporadic intervals but individuals are not (or cannot be followed) overtime. An example in modern biology is single-cell data that is measured over a time course of weeksmeasuring cellular development, where the measurements are destructive to the cell at each timepointbut the entire population is followed over time.

Original:  Manifold Interpolation Optimal-transport Flow (MIOFlow) 1 , a new framework for flows based ondynamical optimal transport  in a manifold embedding.
Revised: To continuously interpolate such populations, wepropose Manifold Interpolation Optimal-transport Flow (MIOFlow) 1 , a new framework for flowsbased on optimal transport penalties applied to dynamics learned in a manifold embedding.

Original: MIOFlow uses a neural ordinary differentialequations (Neural ODE) [7] to transport a sampling of high dimensional data points between timepoints such that 1.
Revised: MIOFlowis a neural ordinary differential equations (Neural ODE) [6] network to transport a sampling of highdimensional data points between timepoints such that 1.

Original: the transport occurs on a manifold defined by samples, 2.
Revised: the transport occurs on a manifold defined by samples, 2.

Original: the transport is penalizedto agree with measured timepoints using Wasserstein, 3.
Revised: the transport is penalized to agree with measured timepoints using Wasserstein, 3.

Original: the transport is inherently stochastic.
Revised: thetransport is inherently stochastic.

Original: Works tackling a similar problem include TrajectoryNet [41], dynamics modeled by a recurrent neuralnetwork [18], and Waddington-OT [34].
Revised: The key works tackling a similar problem include TrajectoryNet [35], dynamics modeled by arecurrent neural network [16], and Waddington-OT [30].

Original: TrajectoryNet is based on a continuous normalizing flowwhere the magnitude of the derivative is penalized to create a more optimal flow.
Revised: TrajectoryNet is based on a continuousnormalizing flow where the magnitude of the derivative is penalized to create a more optimal flow.

Original: However, thisapproach suffers from several drawbacks.
Revised: However, this approach suffers from several drawbacks.

Original: The first drawback is the requirement of starting froma Gaussian distribution.
Revised: The first drawback is the requirement ofstarting from a Gaussian distribution.

Original: It can be difficult to match distributions from the real world to Gaussiandistributions in an interpolating sense, all intermediate distributions must be penalized to start froma Gaussian rather than being trained to flow from one to the next.
Revised: It can be difficult to match distributions from the real worldto Gaussian distributions in an interpolating sense, all intermediate distributions must be penalizedto start from a Gaussian rather than being trained to flow from one to the next.

Original: Second, continuous normalizingflows are deterministic.
Revised: Second, continuousnormalizing flows are deterministic.

Original: Hence, to model the intrinsic stochasticity in biology, we have to forcechaotic dynamics starting from slight noise added to any initial distributions.
Revised: Hence, to model the intrinsic stochasticity in biology, wehave to force chaotic dynamics starting from slight noise added to any initial distributions.

Original: Third, continuousnormalizing flow modelsin k dimensions require calculating the trace of the Jacobian, which requires
Revised: Third,continuous normalizing flow models require calculating the trace of the Jacobian, which requires

Original: O ( k 2 ) operations to compute [7], making our method k times faster per function evaluation.
Revised: O ( d 2 ) operations to compute [6], making our method d times faster per function evaluation [6].

Original: Additionally, unlike normalizing flows which operate in ambient data dimensions, we focus on flowson a data manifold.
Revised: Additionally, unlike normalizing flows which operate in ambient data dimensions, we focus on flowson a data manifold.

Original: We feature a two-pronged approach with which to enforce this.
Revised: We feature a two-pronged approach with which to enforce this.

Original: First, propose the
Revised: First, propose the

Original: Geodesic Autoencoder to embed the data such that distances in the latent space match a new manifolddistance called diffusion geodesic distance.
Revised: Geodesic Autoencoder to embed the data such that distances in the latent space match a new manifolddistance called diffusion geodesic distance.

Original: Second, we penalize the transport using a manifoldoptimal transport method [31, 46].
Revised: Second, we penalize the transport using a manifoldoptimal transport method [27, 40].

Original: These two steps enforce flows on the manifold, whereas in earlierworks, such as TrajectoryNet, the KL-divergence is used to match distributions, and in [24] access tothe metric tensor of a Riemannian manifold is required.
Revised: These two steps enforce flows on the manifold, whereas in earlierworks such as TrajectoryNet KL-divergence is used to match distributions, and in [22] access to themetric tensor of a Riemannian manifold is required.

Original: The main contributions of our work include:
Revised: The main contributions of our work include:

Original: • The MIOFlow framework for efficiently learning continuous stochastic dynamics of static snapshot data based on Neural ODEs that implements optimal transport flows on a data manifold.•
Revised: • The MIOFlow framework for efficiently learning continuous stochastic dynamics of staticsnapshot data based on Neural ODEs that implement optimal transport flows on a datamanifold.

Original:  The Geodesic Autoencoder to create a manifold embedding  in the latent space of the autoencoder.• A new multiscale manifold distance called diffusion geodesic distance, and theoretical resultsshowing its convergence to a geodesic on the manifold in the limit of infinitely many points.
Revised: • The Geodesic Autoencoder to create a manifold embedding of high dimensional data in thelatent space of the autoencoder.• A new multiscale manifold distance called diffusion geodesic distance, and theoreticalresults showing its convergence to a geodesic on the manifold in the limit of infinitely manypoints.• Theoretical results relating MIOFlow trajectories to dynamic optimal transport.• Empirical results showing that our flows can model divergent trajectories on toy data.• Empirical results on single-cell developmental and treatment response data.  Related Work Static methods such as pseudotime assign an abstract measure of time for all cells depending onan initial state [15, 38] in can be combined to infer the cell lineage, i.e. the principal curves ortrajectories [34]. When estimates of cell growth rates are available, Zhang et al.

Revised: [41] proposed amethod based on optimal transport.

Original: • Empirical results showing that our flows can model divergent trajectories on toy data and onsingle-cell developmental and treatment response data.
Revised: Another line of research is on the direct rate of differentiation ofcell or RNA velocity [18, 3]. This method does not provide trajectories, but it can be combined todefine a transition matrix to infer trajectories [42]. Saelens et al. [29] provide a rigorous comparisonof 45 static methods.

Revised: The task of recovering dynamics from multiple snapshots has gained a lot of attention, especially inbiology with scRNA-seq data.

Original: Preliminaries and Background
Revised: Similar methods used optimal transport [30] either with a recurrentneural network [16], generative adversarial networks [28], flows defined by the gradient of an inputconvex neural network [4], or a CNF [35]. Here, we circumvent the need for a CNF by modelingthe trajectories directly on the support of the distributions, hence maintaining the computational costlow even in high dimension. Some generative methods that interpolate between a prior to a targetcan also be used to recover trajectories. For example [9], solve the Schrödinger bridge (SB) between two distributions.

Revised: SB methods however use a series of sinkhorn iterations and do not explicitly learndynamics, making their interpolations inaccurate.

Original: Problem Formulation and Notation We consider a distribution µ t over R k evolving over timet ∈ R , from which we only observe samples from a finite set of T distributions { µ i } T − 1 i =0 .
Revised: We consider a distribution µ t over R d evolving over time t ∈ R , from which we only observe samplesfrom a finite set of T distributions { µ i } T − 1 i =0 .

Original: Foreach time t , we observe a sample X t ∼ µ t of size n t .
Revised: For each time t , we observe a sample X t ∼ µ t of size n t .

Original: We note the set of all observations X of sizen := (cid:80) i n i .
Revised: We note the set of all observations X of size n := (cid:80) i n i .

Original: We also want to characterize the evolution of the support of µ t .
Revised: We also want to characterize the evolutionof the support of µ t .

Original: We aim to definea trajectory from an initial points X 0 ∼ µ 0 to X T − 1 ∼ µ T − 1 , given the intermediate conditions X 2 ∼ µ 2 , . . .
Revised: We aim to define a trajectory from an initial points X 0 ∼ µ 0 to X T − 1 ∼ µ T − 1 ,given the intermediate conditions X 2 ∼ µ 2 , . . .

Original: , X T − 1 ∼ µ T − 1 .
Revised: , X T − 1 ∼ µ T − 1 .

Original: We are thus interested in matching distributions given a set of Tsamples { X i } T − 1 i =1 , and initial condition X 0 .
Revised: We are thus interested in matchingdistributions given a set of T samples { X i } T − 1 i =1 , and initial condition X 0 .

Original: In the following, we assume that the distributions are absolutely continuous with respect to the
Revised: In the following, we assume that the distributions are absolutely continuous with respect to the

Original: Lebesgue measure, and we use the same notation for the distribution and its density.
Revised: Lebesgue measure, and we use the same notation for the distribution and its density.

Original: We the equivalence two distances d 1 ≃ d 
Revised: We note theequivalence between two distances d 1 ≃ d 2 , meaning that there exist. We assume that all Stochastic

Original: Stochastic Differential Equations ) admit a solution [14] proofs are presented in the supplementary material
Revised: Differential Equations (SDE) admit a solution [12]. All proofs are presented in the supplementarymaterial.

Original: Optimal Transport In this section, we provide a brief overview of optimal transport [31, 46],our primary approach for interpolating between distributions using a neural ODE.
Revised: In this section, we provide a brief overview of optimal transport, our primary approach for interpolating between distributions using a neural ODE.

Revised: We also review its dynamic formulation based ontwo monographs [27, 40].

Revised: The Monge problem, the initial optimal transport formulation, is to find atransport plan T : X → Y that minimizes the total cost (cid:80) d ( x, T ( x )) of moving elements in X to Yfor a given cost function d : X × Y → [0 , ∞ ) .

Revised: With this formulation, it is not possible to divide themass of an element in X to multiple elements in Y .

Revised: The Kantorovich formulation solves this issue byformulating the problem with distributions.

Original: We consider twodistributions µ and ν defined on X and Y , and Π( µ, ν ) the set of joint distributions on X × Y withmarginals µ and ν , i.e. π ( dx, Y ) = µ ( dx ) and π ( X , dy ) =
Revised: We consider two distributions µ and ν defined on X and Y , and Π( µ, ν ) the set of joint distributions on X × Y with marginals µ and ν , i.e. π ( dx, Y ) = µ ( dx )and π ( X , dy ) =

Original: ν ( dy ) for π ∈ Π( µ, ν ) .
Revised: ν ( dy ) for π ∈ Π( µ, ν ) .

Original: The transportplan π moves the mass from µ to ν , where the cost of moving a unit mass from the initial x ∈ X tothe final y ∈ Y is d ( x, y ) .
Revised: The transport plan π moves the mass from µ to ν , wherethe cost of moving a unit mass from the initial x ∈ X to the final y ∈ Y is d ( x, y ) .

Original: This formulation gives rise to the p -Wasserstein distance
Revised: This formulationgives rise to the p -Wasserstein distance

Original: [Equation]
Revised: [Equation]

Original: where p ∈ [1 , ∞ ) .
Revised: where p ∈ [1 , ∞ ) .

Original: Suppose X = Y = R k , and d ( x, y ) :
Revised: For the rest of the paper, we suppose X = Y = R d , and d ( x, y ) :

Original: = || x − y ||
Revised: = || x − y ||

Original: then Benamou and Brenier[3] provide a dynamic formulation of the optimal transport problem.
Revised: 2 , the Euclidean distance. With these assumptions, Benamou and Brenier [2] provide a dynamic formulationof the optimal transport problem.

Original: For simplicity, we assume afix time interval [0 , 1] , but note that the following holds for any time interval [ t 0 , t 1 ] .
Revised: For simplicity, we assume a fix time interval [0 , 1] , but note that thefollowing holds for any time interval [ t 0 , t 1 ] , and we let C :

Original:  The transportplan is replaced by a time-evolving distribution ρ t , such that ρ 0{+thatρ 
+} = µ and ρ 1 = ν and that satisfiesthe continuity equation ∂ t ρ t + ∇ · ( ρ t v ) = 0 , where ∇· is the divergence operator.
Revised: = C ([0 , 1] , R d ) the space of R d valuedcontinuous function. The transport plan is replaced by a time-evolving distribution ρ t , such {+thatρ 
+} = µ and ρ 1 = ν .

Original: The movementof mass is described by a time-evolving vector field v ( x, t ) .
Revised: The movement of mass is described by a time-evolving vector field v ( x, t ) .

Original: When ρ t satisfies these conditions,
Revised: It isalso required that the total mass of ρ t must be the same for all t ∈ [0 , 1] (see [7] for the unbalancedcase). That is ρ t must satisfy the conservation of mass (or continuity equation) [Equation] where ∇ is the divergence operator.

Original:  Benamou and Brenier [3] show that
Revised: When ρ t satisfies these two conditions, Benamou and Brenier[2] show that

Original: [Equation]
Revised: [Equation]

Original: This formulation arises from the field of fluid mechanics; the optimal vector field is the divergence ofa pressureless potential flow.
Revised: This formulation arises from the field of fluid mechanics; the optimal vector field is the divergence ofa pressureless potential flow.

Original: We can also view the problem on the path space of R k
Revised: It can also be formulated on the path space of R d

Original: [Equation]
Revised: [Equation]

Original: where the infimum is over all absolutely continuous stochastic path X t  (see [25]).
Revised: where the infimum is over all absolutely continuous stochastic path X t on C (see [23]).

Original: Adding Diffusion For various applications, it is interesting to incorporate a diffusion term inthe paths or trajectories.
Revised: Adding Diffusion For various applications, it can be interesting to incorporate a diffusion termin the paths or trajectories.

Original: These types of flows are often used in control dynamics [1, 26, 15], andmean field games [21].
Revised: These types of flows are often used in control dynamics [1, 24, 13], andmean field games [19].

Original: We consider an SDE dX t = f
Revised: We consider an SDE dX t = f

Original: ( X t , t ) dt + √ σdB t , where B t is a standard
Revised: ( X t , t ) dt + √ σdB t , where B t is a standard

Original: Brownian motion.
Revised: Brownian motion.

Original: To model this new dynamic, one can replace the continuity equation with the Fokker–Planck equation of this SDE with diffusion, i.e. ∂ t ρ t + ∇ · ( ρ t v ) =
Revised: To model this new dynamic, one can replace (1) with the Fokker–Planck equationof this SDE with diffusion, i.e. ∂ t ρ t + ∇ · ( ρ t v ) =

Original: ∆( σρ t / 2) .
Revised: ∆( σρ t / 2) .

Original: From thisformulation, one can also retrieve the Benamou-Brenier optimal transport.
Revised: From this formulation, one can alsoretrieve the Benamou-Brenier optimal transport.

Original: Indeed, Mikami [25]established the convergence to the W 2 when the diffusion term goes to zero.
Revised: Indeed, Mikami [23] established the convergenceto the W 2 when the diffusion term goes to zero.

Original: Similar to (2), we canphrase the problem using an SDE
Revised: Similar to (3), we can phrase the problem using an SDE

Original: [Equation]
Revised: [Equation]

Original: The two previous formulations admit the same entropic interpolation ρ t [30, 25].
Revised: The two previous formulations admit the same entropic interpolation ρ t [26, 23].

Original: In this paper, weutilize such a diffusion term, with the knowledge that it still converges to the transport betweendistributions.
Revised: In this paper, weutilize such a diffusion term, with the knowledge that it still converges to the transport betweendistributions.

Original: Moreover, as the diffusion goes to zero, the infimum converges to the W 2 .
Revised: Moreover, as the diffusion goes to zero, the infimum converges to the W 2 .

Original: Manifold Learning A useful assumption in representation learning is that high dimensional dataoriginates from an intrinsic low dimensional manifold that is mapped via nonlinear functions toobservable high dimensional measurements; this is commonly referred to as the manifold assumption.

Original: Formally, let M be a hidden m dimensional manifold that is only observable via a collection ofk ≫ m nonlinear functions f 1 , . . .

Original: , f k : M → R that enable its immersion in a high dimensionalambient space as F ( M ) = { f ( z ) = ( f 1 ( z ) , . . . , f k ( z )) T : z ∈ M} ⊆
Revised: The dynamic formulations (2) and (3) are equivalent with W 2 holds when the ground distance is the

Original: R k from which data iscollected.

Original: [-Conversely, given a dataset X = { x 1 , . . . , x n } ⊂ R k of high dimensional observations,manifold learning methods assume data points originate from a sampling Z = { z i } ni =
-]
Revised: [-Conversely, given a dataset X = { x 1 , . . . , x n } ⊂ R k of high dimensional observations,manifold learning methods assume data points originate from a sampling Z = { z i } ni =
-]Euclidean.

Original: ∈ M ofthe underlying manifold via x i = f ( z i ) , i = 1 , . . . , n , and aim to learn a low dimensional intrinsicrepresentation that approximates the manifold geometry of M . To learn a manifold geometry from collected data that is robust to sampling densityvariations, Coifman and Lafon [9] proposed to use an anisotropic kernel k ϵ,β ( x, y ) :
Revised: In practice, we might have observations sampled from or generated by a lower dimensionalmanifold.

Original: =

Original: k ϵ ( x, y ) /

Original: ∥

Original: k ϵ ( [Equation] , · ) ∥ β 1 ∥ k ϵ ( y, · ) ∥ β 1 , where 0 ≤ β ≤ 1 controls the separation of geometry from density,with β = 0 yielding the isotropic kernel, and β = 1 completely removing density and providinga geometric equivalent to uniform sampling of the underlying manifold. Next, the kernel k ϵ,β isnormalized to define transition probabilities p ϵ,β ( x, y ) :
Revised: Thus using the Euclidean distance in the ambient space ignores the underlying geometry.

Original: =

Original: k ϵ,β ( x, y ) /

Original: ∥ k ϵ,β ( x, · ) ∥ 1 and an n × nrow stochastic matrix ( P ϵ,β ) ij := p ϵ,β ( x i , x j ) that describes a Markovian diffusion process over the intrinsic geometry ofthe data. Finally, a diffusion map [9] Φ t ( x i ) is defined by the eigenvaluesand eigenvectors of the matrix P tϵ,β .
Revised: Inthe following section, we present theoretical results that will justify our approximation of a distanceequivalent to the geodesic distance on a closed Riemannian manifold.

Original: Most notably, this embedding preserves the diffusion distance∥ p tϵ,β ( x i , · ) − p tϵ,β ( x j , · ) /ϕ 1 ( · ) ∥ 2 between x i , x j ∈ X , where t is a time-scale parameter. Next, weconsider a multiscale diffusion distance that relates to the geodesic distance on the manifold.
Revised: 3.3 Manifold Geodesics

Original: The dynamic optimal transport formulations (1) and (2) are valid when theground distance is Euclidean. This isrestrictive as we might have observations sampled from a lowerdimensional manifold. Thus, using the Euclidean distance in  the ambient space ignores the underlyinggeometry of the data.
Revised: Part of MIOFlow is a Geodesic Autoencoder to learn a representation Z , such that the L 2 norm in Zis equivalent to the geodesic on a manifold M .

Original: Static optimal transport with a geodesic distance is known to perform well ongeometric domains [37, 42], or when the data lie on a manifold [43]. Here, we extend the use of aground geodesic distance to the dynamic formulation. To do so, we define a manifold learning stepthat precedes the trajectory inference. The goal is to learn an embedding Z such that the Euclidean in
Revised: To motivate this, we use results from Riemanniangeometry.

Original: Z is equivalent to the geodesic distance. In that case, dynamic optimal transport in Z is equivalentto the Wasserstein with a geodesic ground distance. In the following section, we present theoreticalresults that will justify our approximation of a geodesic distance on a closed Riemannian manifold. We consider a closed Riemannian manifold ( M , d M ) , where d M is the geodesic distance representingthe shortest path between two points on the manifold.
Revised: We consider a closed Riemannian manifold ( M , d M ) , where d M is the geodesic distance, and wenote h t the heat kernel on M .

Original: We note h t the heat kernel on M . For a pointx ∈ M , the heat kernel h t ( x, · ) induces a measure, i.e. how the heat has propagated on the manifoldat time t given an arbitrary initial distribution.
Revised: For a point x ∈ M , the heat kernel h t ( x, · ) induces a measure, i.e.how the heat has propagated on the manifold at time t given an arbitrary initial distribution.

Original: The diffusion ground distance between x, y ∈ M isbased on the L 1 norm between the measures induced by the heat kernel given the initial conditionsδ x and δ y .
Revised: Thediffusion ground distance between x, y ∈ M is based on the L 1 norm between the measures inducedby the heat kernel given the initial conditions δ x and δ y .

Original: Definition 1.
Revised: Definition 1.

Original: The diffusion ground distance between x, y ∈ M is
Revised: The diffusion ground distance between x, y ∈ M is

Original: [Equation]
Revised: [Equation]

Original: for α ∈ (0 , 1 / 2) , the scale parameter k ≥ 0 , and h t the heat kernel on M .
Revised: for α ∈ (0 , 1 / 2) , the scale parameter k ≥ 0 , and h t the heat kernel on M .

Original: Next, we state an important result from Leeb and Coifman  (Thm. 2 in [22]). This theorem links thediffusion ground distance and the geodesic on a Riemannian manifold.
Revised: Next we state an important result from Leeb and Coifman [20] (Thm. 2), this theorem links thediffusion ground distance and the geodesic on a Riemannian manifold.

Original: Theorem 1 ([22] Thm. 2) .
Revised: Theorem 1.

Original: Let ( M , d M ) a closed Riemannian manifold, with geodesic d M , forα ∈ (0 , 1 / 2) , the distance D α is equivalent to d 2 α M .
Revised: Let ( M , d M ) a closed Riemannian manifold, with geodesic d M , for α ∈ (0 , 1 / 2) , thedistance D α is equivalent to d 2 α M .

Original: In practice, we cannot always evaluate the heat kernel for an arbitrary manifold.
Revised: In practice, we cannot always evaluate the heat kernel for an arbitrary manifold.

Original: However, convergenceresults from Coifman and Lafon [9] provide a natural approximation.
Revised: However, convergence results from Coifman and Lafon [8] provide a natural approximation.

Revised: For a rotationinvariant kernel k ϵ : M × M → R + , where ϵ > 0 is a bandwidth parameter, we definek ϵ,β ( x, y ) :

Revised: =

Original: Indeed, for β = 1 ,  the diffusion operator P t/ϵϵ,β converges to the heat operator as ϵ goes to zero (Prop. 3 [9]).
Revised: k ϵ ( x, y ) /q β ( x ) q β ( y ) where q ( x ) := (cid:82) k ϵ ( x, y ) dµ ( y ) , and µ is the underlying distribution on M .

Revised: The kernel k ϵ,β aims to capture the local geometry of the manifold.

Original: For the rest of the paperwe assume β = 1 , we define P ϵ :
Revised: To mimicthe propagation of the heat, we define the transition kernel p ϵ,β ( x, y ) : = k ϵ,β ( x, y ) /q β ( x ) whereq β ( x ) :

Original: =  P ϵ, 1 and the matrix P ϵ :=
Revised: = (cid:82) k ϵ,β ( x, y ) dµ ( y ) . We note its resulting operator P ϵ,β . For β = 1 , it is shown that theoperator P t/ϵϵ, 1 will converge to the heat operator as ϵ goes to zero (Prop. 3 [8]). For the rest of thepaper we assume β = 1 , and we define P ϵ :

Revised: =

Original: P ϵ, 1 .
Revised: P ϵ, 1 .

Original:  Manifold Interpolating Optimal-Transport Flow
Revised:  Manifold Interpolating Optimal-Transport Flow

Original: To learn individual trajectories from multiple cross-sectional samples of a population, we propose
Revised: To learn individual trajectories from multiple cross-sectional samples of a population, we propose

Original: MIOFlow shown in Fig.
Revised: MIOFlow shown in Fig.

Original: Our method consists of two steps.
Revised: Our method consists of two steps.

Original: We first learn an embedding Z whichpreserves the diffusion geodesic distance—which we define using our Geodesic Autoencoder.
Revised: We first learn an embedding Z which preserves the diffusion geodesic distance—which we define using our Geodesic Autoencoder.

Original: Then,we learn continuous trajectories based on an optimal transport loss.
Revised: Then,we learn continuous trajectories based on an optimal transport loss.

Original: We model the trajectories with a
Revised: We model the trajectories with a

Original: Neural ODE [7], allowing us to learn non-linear paths and interpolate between timepoints.
Revised: Neural ODE [6], allowing us to learn non-linear paths and interpolate between timepoints.

Original: 3.1 Geodesic Autoencoder Embedding
Revised: 4.1 Geodesic Autoencoder Embedding

Original: In light of Thm. 1,we define an approximation of the diffusion ground distance D α (Def. 1) froma collection of samples X that we call diffusion geodesic distance .We then train an autoencoderregularized to match this distance.
Revised: In light of Thm. 1, define an approximation of the diffusion ground distance D α (Def. 1) from acollection of samples X that we call diffusion geodesic distance .

Original: Similar ideas are explored in [27, 11] where the latent space of theautoencoder is regularized to match a manifold learning embedding, either PHATE or diffusion map.

Original: Our approximation relies on  the diffusion matrix P ϵ .
Revised: Our approximation relies on a discrete version of the operator P ϵ .

Original: Wefirst recall its construction. We build a graphwith the affinity matrix defined by ( K ϵ ) ij :
Revised: We build a graph with theaffinity matrix defined by ( K ϵ ) ij :

Original: =
Revised: =

Original: k ϵ ( x i , x j ) , and we consider the density normalizedmatrix M ϵ :=
Revised: k ϵ ( x i , x j ) , and we consider the density normalized matrix M ϵ :=

Original: Q − 1 K ϵ Q − 1 , where Q is a diagonal matrix with Q ii := (cid:80) j ( K ϵ ) ij .
Revised: Q − 1 K ϵ Q − 1 , where Q is a diagonal matrix with Q ii := (cid:80) j ( K ϵ ) ij .

Original: In practice, wechoose k ϵ to be a Gaussian kernel or the α -decay kernel [29].
Revised: In practice, we choosek ϵ to be a Gaussian kernel or the α -decay kernel [25].

Original: Lastly, a Markov diffusion operatoris defined by P ϵ:
Revised: Lastly, a Markov diffusion operator is definedby P ϵ

Original: =
Revised: =

Original: D − 1 M ϵ , where D ii :=
Revised: D − 1 M ϵ , where D ii :=

Original: (cid:80) nj
Revised: (cid:80) nj

Original: =1 ( M ϵ ) ij is a diagonal matrix.
Revised: =1 ( M ϵ ) ij is a diagonal matrix.

Original: The stationarydistribution associated to P ϵ is π , where π i = D ii / (cid:80) j D jj , since P ϵ is π -reversible.
Revised: The stationary distributionassociated to P ϵ is π , where π i = D ii / (cid:80) j D jj , since P ϵ is π -reversible.

Original: Finally, wenote that the matrix P ϵ follows a similar construction as the kernel p ϵ ; the integrals are approximatedby sums.
Revised: Finally, we note that thematrix P ϵ follows a similar construction as the kernel p ϵ ; the integrals are approximated by sums.

Original: We refer to [9] for more details about the convergence of the matrix operator P ϵ to theoperator P ϵ .
Revised: We refer to [8] for more details about the convergence of the matrix operator P ϵ to the operator P ϵ .

Original: Until now, we have defined a matrix P ϵ that approximates the operator P ϵ , which in turnconverges to the heat operator.
Revised: Until now, we have defined a matrix P ϵ that approximates the operator P ϵ , which in turn convergesto the heat operator.

Original: Now, we define an approximation of the diffusion ground distance,based on the matrix P ϵ .
Revised: Now, we define an approximation of the diffusion ground distance, based onthe matrix P ϵ .

Original: We use the notation ( P ϵ ) ti : to represent the i-th row of P tϵ , it represents thetransition probabilities of a t-steps random walk started at x i .
Revised: We use the notation ( P ϵ ) ti : to represent the i-th row of P tϵ , it represents the transitionprobabilities of a t-steps random walk started at x i .

Original: Definition 2.
Revised: Definition 2.

Original: We define the diffusion geodesic distance between x i , x j ∈ X as
Revised: We define the diffusion geodesic distance between x i , x j ∈ X as

Original: [Equation]
Revised: [Equation]

Original: The diffusion geodesic compares the transitions probabilities of a random walk at various scalesgiven two different initial states x i , x j .
Revised: The diffusion geodesic compares the transitions probabilities of a random walk at various scalesgiven two different initial states x i , x j .

Original: In [43], the authors use the diffusion operator P ϵ to define adistance equivalent to the Wasserstein with ground distance D α . Their method comes from the staticformulation of optimal transport.
Revised: With a slight abuse of notation, a similar distance betweendistributions (rather than points) is constructed in [37] where it was applied to approximate the Earth

Original: Here we propose to learn a space Z that preserves an approximationof the distance D α , in order to do dynamic optimal transport in Z . Training To use the diffusion geodesic distance, we train a Geodesic Autoencoder , with encoderoutputting ϕ : R k → Z , such that ∥ ϕ ( x i ) − ϕ ( x j ) ∥
Revised: Mover’s Distance with a geodesic ground cost. Training To use the diffusion geodesic distance, we train a Geodesic Autoencoder , with encoderoutputing ϕ : R d → Z , such that ∥ ϕ ( x i ) − ϕ ( x j ) ∥

Original: 22 ≈ Gα ( x i , x j ) .
Revised: 22 = G ( x i , x j ) .

Original: We draw a subsample of size N ,evaluate Gα and minimize the mean square error
Revised: We draw a subsample of size N ,evaluate G and minimize the mean square error

Original: [Equation]
Revised: [Equation]

Original: Learning the embedding Z over using Gα as the advantage of being inductive, it is useful sincewe use it to compute distances between predicted and ground truth observations.
Revised: Learning the embedding Z over using G as the advantage of being inductive, it is useful sincewe use it to compute distances between predicted and ground truth observations.

Original: Moreover, theencoder can be trained to denoise, hence becoming robust to predicted values that are not close tothe manifold.
Revised: Moreover, the encoder can be trained to denoise, hence becoming robust to predicted values that are not close to themanifold.

Original: Computing G α on the entire dataset is inefficient due to the powering of the diffusionmatrix.
Revised: Empirically, we observed that G fails to capture meaningful distances for a large numberof observations.

Original: We circumvent this difficulty with the encoder, since wetrain on subsamples. We choose Nto have few observations in most regions of the manifold, thus making the computation of G α veryefficient, which allows us to consider more scales.
Revised: In practice, we choose N to have few observations in most regions of the manifold,which also has the benefit of reducing the size of the matrix P ϵ , thus making the computation of Gvery efficient.

Original: Theorem 2.
Revised: Theorem 2.

Original: Assuming X is sampled from a closed Riemannian manifold M with geodesic d M. Then, for all α ∈ (0 , 1 / 2) ,  sufficiently large K, N and small ϵ > 0 , we havewith high probability Gα ( x i , x j ) ≃ d 2 α M ( x i , x j )  for all x i , x j ∈ X .
Revised: Assuming X is sampled from a closed Riemannian manifold M with geodesic d M , thenfor a sufficiently large K, N and small ϵ > 0 , we have G ( x i , x j ) ≃ d 2 α M ( x i , x j ) , for all α ∈ (0 , 1 / 2) .

Original: Corollary 1.
Revised: Corollary 1.

Original: If the encoder is such that L ( ϕ ) = 0 , thenwith high probability ∥ ϕ ( x i ) − ϕ ( x j ) ∥ 2 ≃d 2 α M ( x i , x j ) for all x i , x j ∈ X .The stochasticity in the two previous results arises from the discrete approximation of the operator P ϵ .
Revised: If the encoder is such that L ( ϕ ) = 0 , then ∥ ϕ ( x i ) − ϕ ( x j ) ∥ 2 ≃ d 2 α M ( x i , x j ) for allx i , x j ∈ X .

Original: The law or large numbers guarantees the convergence.

Original: For fix sample size, there exist approximationerror bounds with high probability, see for example [8, 35].

Original: In practice, we choose α close to 1 / 2 sothat the diffusion distance is equivalent to the geodesic on the manifold.

Original: From that embedding Z , wecan also train a decoder ϕ − 1 : Z → R k , with a reconstruction loss L r := (cid:80) x
Revised: From that embedding, we can also train a decoder ϕ − 1 : Z → R d , with a reconstruction loss L r := (cid:80) x

Original: ∥ ϕ − 1 ◦ ϕ
Revised: ∥ ϕ − 1 ◦ ϕ

Original: ( x ) − x ∥ 2 .
Revised: ( x ) − x ∥ 2 .

Original: This is particularly interesting in high dimensions, since we can learn the trajectories inthe lowerdimensional space Z , then decode them in the ambient space.
Revised: This is particularly interesting, since we can learn the trajectories in Z , then decode them in the ambient space.

Original: Forexample, we use the decoder in Sec. 4 to infer cellular trajectories in the gene space, enabling us to understand how specific genesevolved over time.
Revised: For the task of inferring cells trajectories, this enables usto understand how specific genes evolved over time, by decoding in the gene space.

Original: 3.2 Inferring Trajectories
Revised: 4.2 Inferring Trajectories

Original: Given T distributions { µ i } T − 1 i =0 , we want to model the trajectories respecting the conditions X i ∼ µ ifor fix timepoints i ∈ { 0 , . . .
Revised: Given T distributions { µ i } T − 1 i =0 , we want to model the trajectories respecting the conditions X i ∼ µ ifor fix timepoints i ∈ { 0 , . . .

Original: , T − 1 } .
Revised: , T − 1 } .

Original: Formally, we want to learn a parametrized function f θ ( x, t )such that 
Revised: Formally, we want to learn a parametrized function f θ ( x, t )such that [Equation] u [Equation]

Revised: u

Original: [Equation]
Revised: [Equation]

Revised: du, with X 0 ∼ µ 0 ,

Original:  We adapt the main theorem from [41], to consider the path space, and other types of dissimilaritiesbetween distributions.
Revised: [Equation] , X [Equation] We adapt the main theorem from [35], to consider the path space, and other types of dissimilaritiesbetween distributions.

Original: Theorem 3.
Revised: Theorem 3.

Original: We consider a time-varying vector field f
Revised: We consider a time-varying vector field f

Original: ( x, t ) defining the trajectories dX t =f
Revised: ( x, t ) defining the trajectories dX t =f

Original: ( X t , t ) dt with density ρ t , and a dissimilarity between distributions such that D ( µ, ν ) = 0 iffµ =
Revised: ( X t , t ) dt with density ρ t , and a dissimilarity between distributions such that D ( µ, ν ) = 0 iffµ =

Original: ν .
Revised: ν .

Original: Given these assumptions, there exist a sufficiently large λ > 0 such that
Revised: Given these assumptions, there exist a sufficiently large λ > 0 such that

Original: [Equation]
Revised: [Equation]

Original: Moreover, if X tis defined on the embedded space Z , then W 2 is equivalent to the Wasserstein withgeodesic distance W 2 ( µ, ν ) ≃ W d 2 α M( µ, ν ) .
Revised: Moreover, if X t defined on the embedded space Z , then W 2 is equivalent to the Wasserstein withgeodesic distance W 2 ( µ, ν ) ≃ W d 2 α M .

Original: This theorem enables us to add the second marginal constraint in the optimization problem.
Revised: The previous theorem enables us to add the second marginal constraint in the optimization problem.

Original: Inpractice, it justifies the method of matching the marginals and adding a regularization on the vectorfield f ( x, t ) , that way the optimal ( ρ t , f ) corresponds to the one from the W 2 .
Revised: In practice, it justifies the method of matching the marginals and adding a regularization on thevector field f ( x, t ) , that way the optimal ( ρ t , f ) corresponds to the one from the W 2 .

Original: If in addition we modelthe trajectories in the embed space Z , then the transport is equivalent to the one on the manifold.
Revised: If in additionwe model the trajectories in the embed space Z , then the transport is equivalent to the one on themanifold.

Original: Training We observe discrete distributions µ i :=
Revised: Training We observe discrete distributions µ i :=

Original: (1 /n i ) (cid:80) i δ x i for x i ∈
Revised: (1 /n i ) (cid:80) i δ x i for x i ∈

Original: X i , and we approximate (4) with a Neural ODE [7], where f θ is modeled by a neural network with parameters θ .
Revised: X i , and we approximate (5) with a Neural ODE [6], where f θ is modeled by a neural network with parameters θ .

Original: Denoteby ψ θ : R k × T → R d |T | the function that represents the Neural ODE, where T is a set of timeindices.
Revised: Denoteby ψ θ : R d × T → R d |T | the function that represents the Neural ODE, where T is a set of timeindices.

Original: We define the predicted sets (cid:98) X 1 , . . .
Revised: We define the predicted sets (cid:98) X 1 , . . .

Original: , (cid:98) X T − 1 = ψ θ ( X 0 , { 1 , . . .
Revised: , (cid:98) X T − 1 = ψ θ ( X 0 , { 1 , . . .

Original: , T − 1 } ) ; given an initialset X 

Revised: , T − 1 } ) ; given an initialset X

Original: it returns the approximation of (4) for all t ∈ T .
Revised: it returns the approximation of (5) for all t ∈ T .

Original: The resulting discrete distributions areˆ µ i :=
Revised: The resulting discrete distributions areˆ µ i :=

Original: (1 /n i ) (cid:80) j δ x j for x j ∈ (cid:98) X i .
Revised: (1 /n i ) (cid:80) i δ x i for x i ∈ (cid:98) X i .

Original: To match the marginals, we used two approaches of training;local or global .
Revised: To match the marginals, we used two approaches of training; localor global .

Original: For the local method, we only predict the next sample, hence given X t as initialcondition, we predict (cid:98) X t +1 =
Revised: For the local method, we only predict the next sample, hence given X t as initial condition,we predict (cid:98) X t +1 =

Original: ψ θ ( X t , t + 1) .
Revised: ψ θ ( X t , t + 1) .

Original: Whereas for the global, we use the entire trajectoriesgiven the initial condition X 0 .
Revised: Whereas for the global, we use the entire trajectories given theinitial condition X 0 .

Original: For both cases, we formulate a loss L m on the marginals, representingthe second term in (5).
Revised: For both cases, we formulate a loss L m on the marginals, representing thesecond term in (6).

Original: To take into account the first term in (5), we add the loss L e , where the integralcan be approximate with the forward pass of the ODE solver, and λ e ≥ 0 is a hyperparameter.
Revised: To take into account the first term in (6), we add the loss L e , where the integralcan be approximate with the forward pass of the ODE solver, and λ e ≥ 0 is a hyperparameter.

Original: [Equation]
Revised: [Equation]

Original: In practice, to compute the Wasserstein between discrete distributions, we use the implementationfrom the library Python Optimal Transport [13].
Revised: In practice, to compute the Wasserstein between discrete distributions, we use the implementationfrom the library Python Optimal Transport [11].

Original: This is in contrasts with the maximum-likelihoodapproach in CNF methods which requires evaluation of the instantaneous change of variables at everyintegration timestep at O ( k 2 ) additional cost per function evaluation [7].
Revised: This is in contrasts with the maximum-likelihoodapproach in CNF methods which requires evaluation of the instantaneous change of variables at everyintegration timestep at O ( d 2 ) additional cost per function evaluation [6].

Original: We add a final density loss L d inspired by [41] to encourage the trajectories to stay on the manifold.
Revised: We add a final density loss L d inspired by [35] to encourage the trajectories to stay on the manifold.

Original: Intuitively, for a predicted point x ∈ (cid:98) X t , we minimize the distance to its k-nearest neighbors in X tgiven a lower bound h > 

Revised: Intuitively, for a predicted point x ∈ (cid:98) X t , we minimize the distance to its k-nearest neighbors in X tgiven a lower bound h >

Original: [Equation]
Revised: [Equation]

Original: ℓ d( x, t ) , where ℓ d ( x, t ) : = k (cid:88) i = max(0 , min
Revised: ℓ d

Original: [Equation]
Revised: [Equation]

Original: k ( {∥ x [Equation]
Revised: x, t ) , where ℓ d ( x, t ) :=

Original: y
Revised: k

Original: [Equation]
Revised: [Equation]

Original: y
Revised: i = max [Equation]

Original: [Equation]
Revised: [Equation]

Original: X
Revised: min

Original: [Equation]
Revised: [Equation]

Original: h )
Revised: k

Original: [Equation]
Revised: [Equation]

Original: Modeling Diffusion We add a diffusion term in the trajectories, resulting in the SDE dX t =f ( X t , t ) dt + √ σ t dB t
Revised: Modeling Diffusion We add a diffusion term in the trajectories, resulting in the SDE dX t =f ( X t , t ) dt + √ σ t dB t

Original: In practice, we learn T parameters { σ t } T − 1 t =

Revised: In practice, we learn T parameters { σ t } T − 1 t =

Original: It has the advantage of mappingthe same cell to different trajectories stochastically as in many naturalistic systems.
Revised: It has the advantage of mappingthe same cell to different trajectories stochastically as in many naturalistic systems.

Original: Early in thetraining process, we notice that it promotes bifurcation, as it helps to spread the initial X 0 .
Revised: Early in thetraining process, we notice that it promotes bifurcation, as it helps to spread the initial X 0 .

Original: Astraining continues, we empirically observed that σ → 0 .
Revised: Astraining continues, we empirically observed that σ → 0 .

Original: Hence, as shown in [25], converging tothe W 2 formulation.
Revised: Hence, as shown in [23], converging tothe W 2 formulation.

Original: For small σ t , we train with a regular ODE solver, essentially corresponding tothe Euler–Maruyama method.
Revised: For small σ t , we train with a regular ODE solver, essentially corresponding tothe Euler–Maruyama method.

Original: To model complex diffusion, one would need an SDE solver, thusdrastically increasing the computational cost.
Revised: To model complex diffusion, one would need an SDE solver, thusdrastically increasing the computational cost.

Original:  Results
Revised:  Results

Original: [Figure] Here, we validate our method empirically on both synthetic test cases and single-cell RNA sequencingdatasets chosen to test the capability of MIOFlow to follow divergent trajectories from populationsnapshot data.
Revised: Here, we validate our method empirically on both synthetic test cases and single-cell RNA sequencingdatasets chosen to test the capability of MIOFlow to follow divergent trajectories from populationsnapshot data.

Original: Using these test cases we assess whether MIOFlow-derived trajectories qualitativelyfollow the data manifold efficiently, we also quantitatively assess their accuracy using the Wasserstein
Revised: Using these test cases we assess whether MIOFlow-derived trajectories qualitativelyfollow the data manifold efficiently, we also quantitatively assess their accuracy using the Wasserstein

Original: W 1 , and Maximum Mean Discrepancy (MMD) [36] on held-out timepoints.
Revised: W 1 , and Maximum Mean Discrepancy (MMD) [31] on held-out timepoints.

Original: The MMD is evaluatedwith the Gaussian (G) kernel and the mean (M).

Original: We compare our results to those other methods thatpreform population flows including TrajectoryNet [41] which is based on a CNF that is regularized toachieve efficient paths, and Diffusion Schrödinger’s Bridge (DSB) [10] which is an optimal transportframework for generative modeling.The baseline measure in the quantitative results correspondsto the average distance between the previous and next timepoints for a given time t .
Revised: We compare our resultsto those other methods that preform population flows including TrajectoryNet [35] which is based ona CNF that is regularized to achieve efficient paths, and Diffusion Schrödinger’s Bridge (DSB) [9]which is an optimal transport framework for generative modeling.

Original: Additionalexperiments, such as ablation studies, and details are provided in the supplementary material.
Revised: Additional experiments such asablation studies, and details are provided in the supplementary material.

Original: 4.1 Artificial Data
Revised: 5.1 Artificial Data

Original: Petal The petal dataset is a simple yet challenging , as it mimics natural dynamics arising in cellulardifferentiation, including bifurcations and merges.
Revised: Petal The petal dataset is a simple yet challenging , as it mimics natural dynamics arising in cellulardifferentiation, including bifurcations and merges.

Original: The dynamics start in the center (purple) andevolves through the end of the leaf (yellow).
Revised: The dynamics start in the center (purple) andevolves through the end of the leaf (yellow).

Original: In Fig.
Revised: In Fig.

Original: 2, we present the trajectories on the petal datasetfor TrajectoryNet, DSB, and MIOFlow.
Revised: 2, we present the trajectories on the petal datasetfor TrajectoryNet, DSB, and MIOFlow.

Original: Only our method is able to learn the dynamic, and keepingtrajectories on the manifold.
Revised: Only our method is able to learn the dynamic, and keepingtrajectories on the manifold.

Original: We then hold-out timepoint 2 corresponding to the initial bifurcation, and, in Tab. 1, we compare the interpolation between the three methods.
Revised: We then hold-out timepoint 2 corresponding to the initial bifurcation,and in Tab. 5.1 we compare the interpolation between the three methods.

Original: MIOFlow results in a moreaccurate entropic interpolation for both W 1 and MMD(G), while requiring much less computationaltime.
Revised: MIOFlow results in a moreaccurate entropic interpolation for both metrics, while requiring much less computational time.

Original: Dyngen We use Dyngen [6] to simulate a scRNA-seq dataset from a dynamical cellular process.
Revised: Dyngen We use Dyngen [5] to simulate a scRNA-seq dataset from a dynamical cellular process.

Original: We first embed the observations in five dimensions using the non-linear dimensionality reductionmethod PHATE [29].
Revised: We first embed the observations in five dimensions using the non-linear dimensionality reductionmethod PHATE [25].

Original: We discretize the space in five bins, each representing a differentiationstage.
Revised: We discretize the space in five bins, each representing a differentiation stage.

Original: This dataset includes one bifurcation.
Revised: This dataset includes one bifurcation.

Original: However, it proves more challenging than the petaldataset.
Revised: However, it proves more challenging than the petal dataset.

Original: The number of samples per bins is not equal, which tends to give more influence tothe lower branch.
Revised: The number of samples per bins is not equal, which tends to give more influence to the lower branch.

Original: Moreover, the bifurcation is asymmetric; the curve in the lower branch is lesspronounce than the upper one, making it harder to learn a continuous bifurcation.
Revised: Moreover, the bifurcation is asymmetric; the curve in the lower branch is less pronounce than theupper one, making it harder to learn a continuous bifurcation.

Original: In Fig.
Revised: In Fig.

Original:  we showthe trajectories for TrajectoryNet, DSB, and our method.
Revised: 3, we show the trajectories for TrajectoryNet, DSB, and our method.

Original: MIOFlow is more accurate as it learnscontinuous bifurcating trajectories along the two main branches.
Revised: MIOFlow is more accurate as it learns continuous bifurcatingtrajectories along the two main branches.

Original: In Tab. 1, we hold-out timepoint 2,and show the interpolation results against the ground truth.
Revised: In Tab. 5.1, we hold-out timepoint 2, and show theinterpolation results against the ground truth.

Original: MIOFlow is more accurate  for all metrics,while requiring much less computational time.
Revised: MIOFlow is more accurate in terms of EMD, and issimilar to TrajectoryNet for the MMD, but requiring much less computational time.

Original: We see that in timepoint 1, TrajectoryNet deviates farfrom the trajectories particularly at the branching point where it struggles with induced heterogeneity(see Fig. 3).
Revised: However, we seethat in timepoint 1, TrajectoryNet deviates far from the trajectoriesm particularly at the branchingpoint where it struggles with induced heterogeneity (see Figure 3).

Original: We have further quantifications reflecting this in the supplementary material. 
Revised: We have further quantificationsreflecting this in the supplementary material. Petal Dyngen W

Original: ground truth points as well as training time
Revised: [Table] MMD Runtime W

Revised: [Table]

Revised: MMD

Revised: Runtime

Revised: TrajectoryNet 1.

Original: [Table]
Revised: [Table]

Original: 4.2 Single-Cell Data
Revised: 5.2 Single-Cell Data

Original: Embryoid Body Here we evaluate our method on scRNA-seq data from the dynamic time courseof human embryoid body (EB) differentiation over a period of 27 days, measured in 3 timepointsbetween days 0-3, 6-9, 12-15, 18-21, and 24-27.
Revised: Embryoid Body Here we evaluate our method on scRNA-seq data from the dynamic time courseof human embryoid body (EB) differentiation over a period of 27 days, measured in 4 timepointsbetween days 0-3, 6-9, 12-15, 18-21, and 24-27.

Original: We reduce the dimensionality of this dataset to 200dimensions using PCA.
Revised: We reduce the dimensionality of this dataset to 200dimensions using PCA.

Original: See the supplementary material for further details on the preprocessing stepsand the parameters.
Revised: See the supplementary material for further details on the preprocessing stepsand the parameters.

Original: By combining PCA and the autoencoder with diffusion geodesic distance described in section 3.1, wecan decode the trajectories back into the gene space (17846 dimensions), thus describing the evolutionof the gene expression of a cell.
Revised: By combining PCA and the autoencoder with diffusion geodesic distance described in section 4.1, wecan decode the trajectories back into the gene space (17846 dimensions), thus describing the evolutionof the gene expression of a cell.

Original: In Fig.
Revised: In Fig.

Original: 5 we plot the trajectories, at a gene level of 20 neuronalcells against 20 random one.
Revised: 5 we plot the trajectories, at a gene level of 20 neuronalcells against 20 random one.

Original: We generally observe distinct differences between TrajectoryNet and
Revised: We generally observe distinct differences between TrajectoryNet and

Original: MIOFlow.
Revised: MIOFlow.

Original: TrajectoryNet (which is shown in the bottom) shows a tendency to have early divergencein trends artificially showing higher heterogeneity in gene trends then exist in neuronal specificationprocesses.
Revised: TrajectoryNet (which is shown in the bottom) shows a tendency to have early divergencein trends artificially showing higher heterogeneity in gene trends then exist in neuronal specificationprocesses.

Original: Examples include HAND2, which has a non-monotonic trend, first increasing in expressionduring development from embryonic stem cells to neural early progenitors and then decreasing as the lineages mature [23].
Revised: Examples include HAND2, which has a non-monotonic trend, first increasing in expressionduring development from embryonic stem cells to neural early progenitors and then decreasing asthe lineages mature [21].

Original: MIOFlow clearly shows this trend whereas TrajectoryNet trends are morechaotic.
Revised: MIOFlow clearly shows this trend whereas TrajectoryNet trends are morechaotic.

Original: Another example is ONECUT2 which is consistently known to be higher in cells that divergetowards the neuronal lineage [45]. Other examples are discussed in the supplementary material.
Revised: Another example is ONECUT2 which is consistently known to be higher in cells that divergetowards the neuronal lineage [39], Other examples are discussed in the supplementary material.

Original: ONECUT biologically 
Revised: bottom). TrajectoryNet consistently shows a tendency for early divergence in trajectories, examples include ONECUT2 which is biologically known to be higher expressed [25] and HAND2 which is known to peak in neural early progenitors (NEPs).

Original: AML Cancer Data with Treatment-based Progression We also assessed whether MIOFlowcan give insights on treatment resistant cancer cell populations, rather than differentiation.
Revised: AML Cancer Data with Treatment-based Progression We also assessed whether MIOFlowcan give insights on treatment resistant cancer cell populations, rather than differentiation.

Original: For thispurpose, we used data from Fennell et al.
Revised: For thispurpose, we used data from Fennell et al.

Original: [12] in which AF9 mutant cancer cells implanted into thebone marrow of mice as a model of Acute Myeloid Leukemia (AML).
Revised: [10] in which AF9 mutant cancer cells implanted into thebone marrow of mice as a model of Acute Myeloid Leukemia (AML).

Original: The mice are then treatedfor 7 days with intensive chemotherapy with single cell measurements taken at day 2, day 5, andday 7.
Revised: The mice are then treatedfor 7 days with intensive chemotherapy with single cell measurements taken at day 2, day 5, andday 7.

Original: Here we show MIOFlow on this dataset along with signatures of leukemic stem cells (LSCs)which are hypothesized to be dominant in chemotherapy resistant cells.
Revised: Here we show MIOFlow on this dataset along with signatures of leukemic stem cells (LSCs)which are hypothesized to be dominant in chemotherapy resistant cells.

Original: We see heterogeneity in thetrajectories, the majority of cells surviving to the last timepoint showing an increase in the markers
Revised: We see heterogeneity in thetrajectories, the majority of cells surviving to the last timepoint showing an increase in the markers

Original: SLPI, MYB, and B2M—all members of the LSC pathway [12].
Revised: SLPI, MYB, and B2M—all members of the LSC pathway [10].

Original: This indicates that not all cells startwith this signature but are made to transition to this signature to resist the effects of chemotherapy.
Revised: This indicates that not all cells startwith this signature but are made to transition to this signature to resist the effects of chemotherapy.

Original:  Conclusion
Revised:  Conclusion

Original: We proposed MIOFlow a computationally efficient method to recover trajectories from multiplecross-sectional samples of a time-evolving population.
Revised: We proposed MIOFlow a computationally efficient method to recover trajectories from multiplecross-sectional samples of a time-evolving population.

Original: Our method includes a novel manifoldlearning framework, where we learn an embedding preserving the geodesic distance.
Revised: Our method includes a novel manifold learningframework, where we learn an embedding preserving the geodesic distance.

Original: We showtheoretically how MIOFlow is related to dynamical optimal transport on an underlying manifold, andit’s capacity to model stochastic trajectories.
Revised: We show theoreticallyhow MIOFlow is related to dynamical optimal transport on an underlying manifold, and it’s capacityto model stochastic trajectories.

Original: On artificial datasets, MIOFlow produced continuoustrajectories more in line with the manifold compared to DSB and TrajectoryNet.
Revised: On artificial datasets, MIOFlow produced continuous trajectoriesmore in line with the manifold compare to DSB and TrajectoryNet.

Original: On scRNA-seqdata, we leveraged an autoencoder with the geodesic embedding to learn the trajectories in a lowerdimensional space while preserving meaningful distances which can then be decoded in gene space.
Revised: On scRNA-seq data, we leveragedan autoencoder with the geodesic embedding to learn the trajectories in a lower dimensional spacewhile preserving meaningful distances which can then be decoded in gene space.

Original: This way we could visualize the evolution of the gene expression of a cell.
Revised: This way we couldvisualize the evolution of the gene expression of a cell.

Original:  Limitations and Broader Impact
Revised:  Limitations and Broader Impact

Original: Our method does have some limitations.
Revised: Our method does have some limitations.

Original: Similarly to other works based on continuous time networks,our work makes use of generic ODE solvers and adjoint dynamics which can fail under stiff dynamics.
Revised: Similarly to other works based on continuous time networks,our work makes use of generic ODE solvers and adjoint dynamics which can fail under stiff dynamics.

Original: Like other methods that learn dynamics it is possible that the learned dynamics may become stiffthrough training leading to decreased model accuracy.
Revised: Like other methods that learn dynamics it is possible that the learned dynamics may become stiffthrough training leading to decreased model accuracy.

Original: In addition, our work is sensitive to thetime-bins, and it worked best for relatively uniform sampling.
Revised: In addition, our work is sensitive to thetime-bins, and it worked best for relatively uniform sampling.

Original: This second issue could potentially bealleviated by using unbalanced optimal transport, which allows transporting distributions with unequalmasses.
Revised: This second issue could potentially bealleviated by using unbalanced optimal transport, which allows transporting distributions with unequalmasses.

Original: Our work can help understand dynamical systems, and in particular cellular differentiation,and thus to the best of our knowledge presents no potential negative societal impacts.A.1 Proof of Thm.  Theorem 2.
Revised: Our work can help understand dynamical systems, and in particular cellular differentiation,and thus to the best of our knowledge presents no potential negative societal impacts.

Original: Assuming X is sampled from a closed Riemannian manifold M with geodesic d M ,then for a sufficiently large K, N and small ϵ > 0 , we have G α ( x i , x j ) ≃ d 2 α M ( x i , x j ) , for allα ∈ (0 , 1 / 2) . Proof. We adapt the convergence proof from [43], it can be summarized in two parts. First, theequivalence of the distances defined by P t/ϵϵ and the heat operator H t .
Revised: Checklist

Original: This relies on the convergenceresults P t/ϵϵ → H t as ϵ → 0 from Coifman and Lafon [9]. Second, the approximation of the operator P ϵ based on a finite set of points.
Revised: For all authors...

Original: Throughout the proof we use the same notation δ x for the delta Dirac function, and δ x i for a row vector with i -th element 1 , and 0 otherwise. For a family of operator ( A t ) t ∈ R we define
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Equation]
Revised: [Yes]

Original: and equivalently D α,A ( x, y ) = (cid:80) 2 − kα ||
Revised: (b) Did you describe the limitations of your work?

Original: A 2 − k ( δ x − δ y ) || 1 , where Aδ x := (cid:82) A ( u, · ) δ x ( u ) du . This definition is the same as Def. 1, but for an arbitrary operator. Our final goal is to showthat G α ( x i , x j ) ≃ D α,H t ( x i , x j ) , and then conclude with Thm. 1 which gives the equivalence
Revised: [Yes] In the limitations and broader

Original: D α,H t ( x i , x j ) ≃ d 2 α M ( x i , x j ) . We first want to show that, for a sufficiently small ϵ > 0 , we have D α,P t/ϵϵ ( x, y ) ≃ D α,H t , where P is the anisotropic operator defined in Sec. 2.
Revised: impact section.

Original: From Coifman and Lafon [9], we have ∥ P t/ϵϵ − H t ∥ L 2 ( M ) → 0 as ϵ → 0 , and also ∥ P t/ϵϵ − H t
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: ∥
Revised: [Yes] In the limitations and broader impact section.

Original: L 1 ( M ) → 0 . Now let Γ t := P t/ϵϵ − H t , and for γ > 0 choose ϵ > 0 such that ∥ Γ t ( δ x − δ y ) ∥ 1 < γ ∥
Revised: We do not foresee any immediate

Original: δ x − δ y ∥ 1 = 2 γfor x ̸ = y , we have
Revised: negative societal impacts of this work.

Original: [Equation] Thus, for all t > 0 , and all γ > 0 there exists ϵ > 0 such that D α, Γ t ( x, y ) < γ . From the reversetriangle inequality, we have [Equation]
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: using D α, Γ t ( x, y ) < γ , we obtain [Equation] According to [54] we can lower bound the heat kernel, and thus the distance D α,H t ( x, y ) > C forsome C > 0 . For γ < C/ 2 , and a sufficiently small ϵ > 0 , we have
Revised: them?

Original: [Equation]
Revised: [Yes]

Original: This proves our first claim, that D α,P t/ϵϵ ( x, y ) ≃ D α,H t for small ϵ > 0 . Next, we consider the approximation of the operator P t/ϵϵ with a finite set of points. Now we define
Revised: If you are including theoretical results...

Original: G [Equation] if ϵ = 2 − K , then the first term of G α ( x i , x j ) is equal to G K,ϵ,α ( x i , x j ) as α → 1 / 2 .
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: Since we willlet K → ∞ , we ignore the second term inG α ( x i , x j ) , as it converges to zero. Similar to Coifmanand Lafon [9] (Sec. 5), we have P → P as n → ∞ , using Monte-Carlo integration (approximationof integral with summation by the law of large numbers). By the strong law oflarge numbers, the convergence is with probability one.
Revised: [Yes] We stated all assumptions in the main body of the text and provide additional

Original: For a finite number of samples, we have a high probabilitybound on the convergence, see for example [8, 35]. Now we let N := min( K, n ) , and therefore [Equation]
Revised: details in the supplementary materials.

Original: Hence, if we take ϵ := 2 − K , we have for sufficiently large n and K (implying small ϵ > 0 ), G α ( x i , x j ) ≃ D α,H t ≃ d 2 α M ( x i , x j ) from Thm. 1. A.2 Proof of Thm.  Theorem 3.
Revised: (b) Did you include complete proofs of all theoretical results?

Original: We consider a time-varying vector field f ( x, t ) defining the trajectories dX t =f ( X t , t ) dt with density ρ t , and a dissimilarity between distributions such that D ( µ, ν ) = 0 iffµ = ν .
Revised: [Yes] Complete proofs are available in the supplement.

Original: Given these assumptions, there exist a sufficiently large λ > 0 such that [Equation] Moreover, if X t defined on the embedded space Z , then W 2 is equivalent to the Wasserstein withgeodesic distance W 2 ( µ, ν ) ≃ W d 2 α M ( µ, ν ) .
Revised: If you ran experiments...

Original: Proof. We recall that [Equation] is equivalent to 
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: [Equation] with the three constraints [Equation] Tong et al.
Revised: mental results (either in the supplemental material or as a URL)?

Original: [41] (Thm. 4.1), showed that, for large λ > 0 and ρ t satisfying a) , this minimizationproblem is equivalent to [Equation] where KL is the Kullback–Leibler divergence, we note that their proof is valid for any dissimilarity D ( ρ 1 || ν ) respecting the identity of indiscernibles.
Revised: [Yes] Code and links to publicly available data will be provided in the supplemental

Original: Using the path formulation, by writing the integralas an expectation and taking the infimum over all absolutely continuous path, we have

Original: [Equation]
Revised: material.

Original: Assuming the encoder ϕ achieves a loss of zero, then from Cor. 1 for α ∈ (0 , 1 / 2) , we have∥ ϕ
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: ( x i ) − ϕ ( x j ) ∥ 2 ≃ d 2 α M ( x i , x j ) for all x i , x j ∈ X ⊆ M , and sufficiently large n and K . That is,there exist c, C > 0 such that c d M ( x i , x j ) ≤ ∥ ϕ
Revised: how they

Original: ( x i ) − ϕ ( x j ) ∥ 2 ≤ Cd M ( x i , x j ) for all x i , x j ∈ X . Then, for all π ∈ Π( µ, ν ) , we have [Equation] and taking the infimum with respect to π ∈ Π( µ, ν ) yields the desired result W 2 ( µ, ν ) ≃ W d 2 α M ( µ, ν ) .
Revised: were chosen)?

Original: A.3 Schrödinger bridge

Original: Let D ( µ, ν ) the space of probability distributions on C := C ([0 , 1] , R k ) , with initial and finaldistribution µ and ν , and N the Wiener measure on C , then the Schrödinger bridge (SB) problem isto find the time-evolving distribution ρ t such thatmin KL [Equation] subject to ρ t ∈ D ( µ, ν )
Revised: [Yes] All training details are provided in the supplement.

Original: [Equation] where KL is the Kullback–Leibler divergence. It also admits a static formulation, i.e. minimizingwith respect toπ ∈ Π( µ, ν ) . Using the measure N yx over a Brownian bridge with initial andfinal condition x and y , and ρ yx defined similarly for ρ ∈ D ( µ, ν ) , we can write KL ( ρ t || N ) =
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: KL ( π || π N ) + KL ( ρ yx || N yx ) , where π N is the joint distribution between the initial and final statesunder the Wiener measure N . By choosing ρ yx = N yx , the static formulation of the Schrödingerbridge is [Equation] subject to π ∈ Π( µ, ν ) .
Revised: ments multiple times)?

Original: When the Wiener measure has variance σ 2 , the static Schrödinger bridge is equivalent to minimizinginf π ∈ Π( µ,ν ) (cid:90)

Original: [Equation]
Revised: [N/A]

Original: where H is the entropy (see [30] and references therein).

Original: That is an optimal transport problemwith entropic regularization, the same formulation as the Sinkhorn divergence [49], which can besolved with Sinkhorn’s algorithm [53].

Original: We compare our method with Diffusion Schrödinger Bridge(DSB) [10], which solve the SB problem with an approximation of  the iterative proportional fittingalgorithm. B Experiment details An anonymized version of our repository 2 is available for reviewing purpose. For all experiments,we used the Runge-Kutta RK4 ODE solver, for the density loss we used h = 0 .
Revised: (d) Did you include the total amount of compute and the type of resources used (e

Original: 01 , and 5 nearestneighbors.

Original: The ODE network consists of three layers, and we concatenate two extra dimensionsto the input as well as the time index.

Original: The encoder network is three layers with ReLU activationfunctions in between layers.

Original: For each epoch (local or global), we sampled 20 batches of a givensample size without replacement.

Original: For both networks, we optimize with AdamW [52] with defaultparameters. In practice, we did not use the energy loss, as our trajectories were already smooth, andit requires more careful parametrization (see Sec. C). For all experiments, we used TrajectoryNetwith 1000 iterations and whiten .
Revised: g., type

Original: We used the authors’ implementation of DSB with the basic modeland default parameters.
Revised: of GPUs, internal cluster, or cloud provider)?

Original: Petal For diffusion geodesic, we experiment with an RBF kernel with scale 0 .
Revised: [Yes] Compute resources are discussed in the supplement.

Original: 1 and α -decay kernel(knn=5), both with maximum time scale 2 5 . The encoder consists of two linear layers of size 8 and32 , trained for 1000 iterations. We used LeakyReLU in between layers of the ODE network, withhidden layers size 16 , 32 , 16 , and initial scales σ t = 0 . We trained for 40 local epoch with samplesize 60 .
Revised: If you are using existing assets (e

Original: We used λ d = 35 to weight the density loss and λ e = 0 . 001 . For the hold-out experiments,we trained for 30 local and 15 global epoch with λ d = 1 . The diffusion parameters are generally small while training, for example when trained onthe entire dataset σ t ∈ { 0 . 18 , 0 .
Revised: g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: 15 , − 0 .
Revised: [Yes]

Original: 08 } , and for the hold-out time three σ t ∈{ 0 . 08 , 0 . 02 , 0 . Dyngen The diffusion geodesic was evaluated with an RBF kernel with scale 0 .
Revised: (b) Did you mention the license of the assets? [Yes] See supplement

Original: 5 , and maximumtime scale 2 5 . The encoder has two linear layers of size 8 and 32 , trained for 1000 iterations.
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: We used
Revised: [Yes]

Original: CELU in between layers of the ODE network, with hidden layers size 16 , 32 , 16 , and initial scalesσ t = 0 . We trained for 50 global epochs of size 60 , with λ d = 5 . For the hold-out experiments, we used λ d = 5 , and for t = 5 we trained on 5 local and 10 global epochs.
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: For t = 3 , since we have tohold-out the steps (2 , 3) and (3 , 4) , we only trained using 50 global epochs. Similar to the petal dataset, the diffusion parameters is empirically low, on the entire dataset it isσ t ∈ { 0 . 18 , 0 . 20 , 0 . 16 } , and for the hold-out three σ t ∈ { 0 . 02 , 0 . 08 , 0 . 24 , 0 . 22 , 0 . 24 } . EB data The EB data are publicly available 3 .
Revised: using/curating?

Original: For the genes trajectories we used the full geodesicautoencoder because we want to decode in PCA space, and then project back into gene space. Theencoder consists of three layers of size 200 , 100 , 100 , and the decoder is 100 , 100 , 200 , both with ReLU activation between layers. We trained for 1000 epochs with sample size of 100 per time,with the α -decay kernel ( 5 nearest neighbors, and max scale of 2 6 ).
Revised: [N/A]

Original: This kernel is more expensiveto compute since it relies on the approximation of the knn, but it circumvents the need to find agood scale as with the RBF kernel.

Original: The ODE network is three hidden layers of size 64 , with CELUactivation in between layers. We trained for 20 local and 10 global epochs with sample size of 400per time. We used the density loss with weight λ d = 20 . AML The AML dataset is publicly available 4 .
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: We first embed the data in 50 dimensions PCAspace. We use an encoder of size 50 , 8 , 8 with ReLU activations in between layers, trained for 1000iterations with the RBF kernel (scale 0.5). For the trajectories we used an ODE network with threehidden layers of size 16 , 32 , 16 , with LeakyReLU in between layers, trained without noise for 50local epochs with λ d = 20 . C Additional ResultsEnergy
Revised: information or offensive content?

Original: loss In Fig. 6, we present the trajectories using the same parametrization as detailed in Sec. B with the energy loss ( λ e = 1 ).
Revised: [Yes] Our data does not contain any personally identifiable information or offensive

Original: This loss requires much more memory since for each functionevaluation we need to save the norm of the derivative.

Original: It should encourage smooth trajectory, but theaccuracy of the trajectories is very sensitive to this loss.

Original: On Dyngen (left), it learns a straight line, andthus cannot bifurcate. On the petal, it initially learns a straight trajectory (middle, one epoch), and astraining continues, it collapses in the center (right, ten epochs). Unless otherwise stated, we did notadd it to our model since our trajectories were smooth. However, with careful parametrization, onecould learn more accurate trajectories trajectories (see Fig. 7).
Revised: content.

Original: Ablation and hold-out To evaluate the prediction accuracy, we train by holding-out one timepoint,and we compute the distance between the held-out sample (ground truth) and the prediction.

Original: We usethe W 1 , the MMD with a Gaussian kernel, and the MMD with identity map ( L 2 norm between thesample means).

Original: For the Petal and Dyngen datasets, we compute the Gaussian MMD by taking theaverage for the scales 0 . 1 and 0 .

Original: For the EB data, we set the scale to 1500 . In Fig. 8, we present an ablation study for the density loss and the geodesic embedding. Theparametrization was found for the model with density loss and embedding, both seem important to learn accurate trajectories. In Fig.
Revised: The Embryoid body data is a cell culture

Original: 9, we present the trajectories when we trained by holding onetimepoint out both with the embedding and density loss.

Original: In Fig. 10, we reproduce the same ablationon Dyngen data. Here, the density loss isimportant for the trajectories to stay on the manifold, andthe embedding appears to help with the bifurcation. To complement this figure, in Tab. 5, we presentthe average of the KNN distance computed on the top 5% .
Revised: and the AML Cancer dataset is

Original: On Dyngen, the trajectories with hold-outtimepoints are presented in Fig. Holding out the timepoint t = 3 seems to be the hardest to learn. In Tab. 2, we show the W 1 and MMD between the predicted values and ground truth, for MIOFlow, TrajectoryNet, and DSB.
Revised: patient derived but contains no identifiable information.

Original: For the W 1 , MIOFlow is always more accurate. In Tab 3, we compare theprediction accuracy on the EB data, depending on using the GAE or not. For the GAE, we considerthe choice of two kernels when computing the distance G ; the Gaussian kernel for various bandwidthparameters, and the α -decay kernel for different K -nearest neighbors. Using the GAE improves theinterpolation accuracy for all metrics.
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: We embed the data in 200 PCA, and train the autoencoder withlayers 200 , 100 , 100 , 50 for 1000 epochs.

Original: We trained the Neural ODE for 80 epochs with λ d = 20 ,σ t = 0 .

Original: 2 , linear layers of size 16 , 32 , 16 , 50 with CELU activation in between layers. We train byholding out timepoint t for the GAE and the neural ODE. For the MMD with the Gaussian kernel, weset the scale to 1500 . DSB.
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: Lower is better.

Original: [Table]
Revised: [Table]

Original: Nearest neighbor distance To assist in evaluating how well the generated points land on the knownmanifold (ground truth points), we employ 1-NN distance from the generated points to the groundtruth points. Table 4 shows the 1-NN distance for several methods and datasets. Of note is theaggregate column. Where aggregate is specified as mean , we compute the mean 1-NN distance fromthe predicted point to the 1-NN at the same time index.
Revised: applicable?

Original: Additionally, to help reduce sensitivity due tooutliers, we computed these distances and report only the mean of the worst (highest 1-NN distance)quartile, where aggregate is specified as quartile . Neural SDE Additionally, we employed a Neural SDE from the PyTorch implementation ofdifferentiable SDE Solvers [51, 50], and compared its results to MIOFlow (Fig. We trainedwith a fixed diffusion function, but due to extensive training time (approximately two hours for 20epochs for each dataset) we were unable to perform extensive hyperparameter optimization.
Revised: [N/A]

Original: Static methods such as pseudotime assign an abstract measure of time for all cells depending onan initial state [17, 44]. In can be combined to infer the cell lineage, i
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: [Table]
Revised: Board

Original: e. the principal curves ortrajectories [40]. When estimates of cell growth rates are available, Zhang et al. [47] proposed a
Revised: ) approvals, if applicable?

Original: [Table]
Revised: [N/A]

Original: method based on optimal transport.

Original: Another line of research is on the direct rate of differentiation ofcell or RNA velocity [20, 4]. This method does not provide trajectories, but it can be combined todefine a transition matrix to infer trajectories [48]. Saelens et al. [33] provide a rigorous comparisonof 45 static methods.
Revised: ) Did you include the estimated hourly wage paid to participants and the total amount

Original: The task of recovering dynamics from multiple snapshots has gained a lot of attention, especially inbiology with scRNA-seq data. Similar methods used optimal transport [34] either with a recurrentneural network [18], generative adversarial networks [32], flows defined by the gradient of an inputconvex neural network [5], or a CNF [41]. Here, we circumvent the need for a CNF by modeling thetrajectories directly on the support of the distributions, hence maintaining the computational cost loweven in high dimension. Some generative methods that interpolate between a prior to a target canalso be used to recover trajectories.
Revised: [Table]

Original: For example [10], solve the Schrödinger bridge (SB) betweentwo distributions. SB methods however use a series of Sinkhorn iterations and do not explicitly learndynamics, making their interpolations inaccurate. E Software Versions In Tab. 6, we present the GPU used during this work, and the following software version wereused: cudatoolkit==11.1.1, graphtools==1.5.2, numpy==1.19.5, phate==1.0.7pytorch==1.10.2, pot==0.7.0, scikit-learn==0.24.2, scipy==1.5.3,scprep==1.1.0, torchdiffeq==0.2.2, torchsde==
Revised: spent on participant compensation

Original: [Table] [Table] [Table]
Revised: [N/A] [Table]

Original: Supplement References