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)

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: We present a method called Manifold Interpolating Optimal-Transport Flow(MIOFlow) that learns stochastic, continuous population dynamics from staticsnapshot 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 as penalized 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 ensure thatthe flow follows the geometry by operating in the latent space of an autoencoderthat 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.
Revised: 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.
Revised: 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.
Revised: 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.
Revised: 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.
Revised: 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.
Revised: 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 [8],
Revised: 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 [27, 28] and autoencoders [26, 10].
Revised: 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.
Revised: 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.
Revised: 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.
Revised: We frame this as learning a population flow on a data manifold.

Original: 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
Revised: However, most of these works have focused on generative modeling: i.e., flowing from a noise distribution such as a

Original: Gaussian distribution to a data distribution in order to generate data.
Revised: Gaussian distribution to a data distribution in order to generate data.

Original: Examples include score-based generative matching [36, 37], diffusion models [18], Schrödinger bridges [29, 9], and continuousnormalizing flows (CNF)
Revised: Examples include score-based generative matching [38, 39], diffusion models [19], Schrödinger bridges [30, 10], and continuousnormalizing flows (CNF)

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 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.

Original: To continuously interpolate populations over time on data with manifold structure, we propose
Revised: To continuously interpolate populations over time on data with manifold structure, we propose

Original: Manifold Interpolation Optimal-transport Flow (MIOFlow) 1 , a new framework for flows based ondynamical optimal transport in a manifold embedding.
Revised: Manifold Interpolation Optimal-transport Flow (MIOFlow) 1 , a new framework for flows based ondynamical optimal transport 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: MIOFlow uses a neural ordinary differentialequations (Neural ODE) [7] to transport a sampling of high dimensional 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 penalizedto agree with measured timepoints using Wasserstein, 3.

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

Original: Works tackling a similar problem include TrajectoryNet [39], dynamics modeled by a recurrent neuralnetwork [17], and Waddington-OT [33].
Revised: Works tackling a similar problem include TrajectoryNet [41], dynamics modeled by a recurrent neuralnetwork [18], and Waddington-OT [34].

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 continuous normalizing flowwhere the magnitude of the derivative is penalized to create a more optimal flow.

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

Original: The first drawback is the requirement of starting froma Gaussian distribution.
Revised: The first drawback is the requirement of starting froma 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 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.

Original: Second, continuous normalizingflows are deterministic.
Revised: Second, continuous normalizingflows 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, we have to forcechaotic dynamics starting from slight noise added to any initial distributions.

Original: Third, continuousnormalizing flow models in k dimensions require calculating the trace of the Jacobian, which requires
Revised: Third, continuousnormalizing flow models in k dimensions 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 ( k 2 ) operations to compute [7], making our method k times faster per function evaluation.

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 [30, 44].
Revised: Second, we penalize the transport using a manifoldoptimal transport method [31, 46].

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 [23] 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, the KL-divergence is used to match distributions, and in [24] access tothe metric 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 static snapshot data based on Neural ODEs that implements optimal transport flows on a data manifold.•

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 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.

Original: • Empirical results showing that our flows can model divergent trajectories on toy data and onsingle-cell developmental and treatment response data.
Revised: • Empirical results showing that our flows can model divergent trajectories on toy data and onsingle-cell developmental and treatment response data.

Original:  Preliminaries and Background
Revised:  Preliminaries and Background

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: 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 .

Original: Foreach time t , we observe a sample X t ∼ µ t of size n t .
Revised: Foreach 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 sizen := (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 evolution of 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
Revised: We aim to definea trajectory from an initial points X 0 ∼ µ 0 to X T − 1 ∼ µ T − 1 , given the intermediate conditions

Original: X 2 ∼ µ 2 , . . .
Revised: 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 matching distributions given a set of Tsamples { 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 equivalencebetween two distances d 1 ≃ dall Stochastic Differential Equations
Revised: We the equivalence two distances d 1 ≃ d Stochastic Differential Equations

Original: ) admit a solution [13] All proofs are presented in the supplementary material
Revised: ) admit a solution [14] proofs are presented in the supplementary material

Original: Optimal Transport In this section, we provide a brief overview of optimal transport [30, 44],our primary approach for interpolating between distributions using a neural ODE.
Revised: 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.

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 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 ) =

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 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 ) .

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

Original: [Equation]
Revised: [Equation]

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

Original: Suppose X =
Revised: Suppose X =

Original: Y =
Revised: Y =

Original: R k , and d ( x, y ) :
Revised: R k , 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: then Benamou and Brenier[3] provide a dynamic formulation of 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 afix time interval [0 , 1] , but note that the following holds for any time interval [ t 0 , t 1 ] .

Original: The transportplan is replaced by a time-evolving distribution ρ t , such that ρ 0 = µ and ρ 1 = ν and that satisfiesthe continuity equation ∂ t ρ t + ∇ · ( ρ t v ) = 0 , where ∇· is the divergence operator.
Revised: The transportplan is replaced by a time-evolving distribution ρ t , such that ρ 0 = µ and ρ 1 = ν and that satisfiesthe continuity equation ∂ t ρ t + ∇ · ( ρ t v ) = 0 , where ∇· is the divergence operator.

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

Original: When ρ t satisfies these conditions,
Revised: When ρ t satisfies these conditions,

Original: Benamou and Brenier [3] show that
Revised: Benamou and Brenier [3] 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: We can also view the problem on the path space of R k

Original: [Equation]
Revised: [Equation]

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

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 is interesting to incorporate a diffusion term inthe paths or trajectories.

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

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
Revised: To model this new dynamic, one can replace the continuity equation with the

Original: Fokker–Planck equation of this SDE with diffusion, i.e. ∂ t ρ t + ∇ · ( ρ t v ) =
Revised: Fokker–Planck equation of 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 thisformulation, one can also retrieve the Benamou-Brenier optimal transport.

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

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

Original: [Equation]
Revised: [Equation]

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

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.
Revised: 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 , . . .
Revised: 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 ) =
Revised: , f k : M → R that enable its immersion in a high dimensionalambient space as F ( M ) =

Original: { f ( z ) =
Revised: { f ( z ) =

Original: ( f 1 ( z ) , . . .
Revised: ( f 1 ( z ) , . . .

Original: , f k ( z )) T : z ∈ M} ⊆
Revised: , f k ( z )) T : z ∈ M} ⊆

Original: R k from which data iscollected.
Revised: R k from which data iscollected.

Original: Conversely, given a dataset X = { x 1 , . . .
Revised: Conversely, given a dataset X = { x 1 , . . .

Original: , x n } ⊂ R k of high dimensional observations,manifold learning methods assume data points originate from a sampling Z
Revised: , x n } ⊂ R k of high dimensional observations,manifold learning methods assume data points originate from a sampling Z

Original: = { z i } ni
Revised: = { z i } ni

Original: =

Revised: =

Original: ∈
Revised: ∈

Original: M ofthe underlying manifold via x i =
Revised: M ofthe underlying manifold via x i =

Original: f ( z i ) , i = 1 , . . . , n , and aim to learn a low dimensional intrinsicrepresentation that approximates the manifold geometry of M .
Revised: f ( z i ) , i = 1 , . . . , n , and aim to learn a low dimensional intrinsicrepresentation that approximates the manifold geometry of M .

Original: To learn a manifold geometry from collected data that is robust to sampling densityvariations, Coifman and Lafon [8] proposed to use an anisotropic kernel k ϵ,β ( x, y ) :
Revised: 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 ) :

Original: =
Revised: =

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

Original: ∥
Revised: ∥

Original: k ϵ 
Revised: k ϵ (

Original: [Equation]
Revised: [Equation]

Original: , · ) ∥ β 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.
Revised: , · ) ∥ β 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.

Original: Next, the kernel k ϵ,β isnormalized to define transition probabilities p ϵ,β ( x, y ) :
Revised: Next, the kernel k ϵ,β isnormalized to define transition probabilities p ϵ,β ( x, y ) :

Original: =
Revised: =

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

Original: ∥
Revised: ∥

Original: k ϵ,β ( x, · ) ∥ 1 and an n × nrow stochastic matrix ( P ϵ,β ) ij :=
Revised: k ϵ,β ( x, · ) ∥ 1 and an n × nrow stochastic matrix ( P ϵ,β ) ij :=

Original: p ϵ,β ( x i , x j ) that describes a Markovian diffusion process over the intrinsic geometry of the data.
Revised: p ϵ,β ( x i , x j ) that describes a Markovian diffusion process over the intrinsic geometry of the data.

Original: Finally, a diffusion map [8] Φ t ( x i ) is defined by the eigenvaluesand eigenvectors of the matrix P tϵ,β .
Revised: Finally, a diffusion map [9] Φ t ( x i ) is defined by the eigenvaluesand eigenvectors of the matrix P tϵ,β .

Original: Most notably, this embedding preserves the diffusion distance∥ p tϵ,β
Revised: Most notably, this embedding preserves the diffusion distance∥ p tϵ,β

Original: ( x i , · ) − p tϵ,β
Revised: ( x i , · ) − p tϵ,β

Original: ( x j , · ) /ϕ 1 ( · ) ∥ 2 between x i , x j ∈ X , where t is a time-scale parameter.
Revised: ( x j , · ) /ϕ 1 ( · ) ∥ 2 between x i , x j ∈ X , where t is a time-scale parameter.

Original: Next, weconsider a multiscale diffusion distance that relates to the geodesic distance on the manifold.
Revised: Next, weconsider a multiscale diffusion distance that relates to the geodesic distance on the manifold.

Original: Manifold Geodesics
Revised: Manifold Geodesics

Original: The dynamic optimal transport formulations (1) and (2) are valid when theground distance is Euclidean.
Revised: The dynamic optimal transport formulations (1) and (2) are valid when theground distance is Euclidean.

Original: This is restrictive as we might have observations sampled from a lowerdimensional manifold.
Revised: This is restrictive as we might have observations sampled from a lowerdimensional manifold.

Original: Thus, using the Euclidean distance in the ambient space ignores the underlyinggeometry of the data.
Revised: Thus, using the Euclidean distance in the ambient space ignores the underlyinggeometry of the data.

Original: Static optimal transport with a geodesic distance is known to perform well ongeometric domains [35, 40], or when the data lie on a manifold [41].
Revised: 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].

Original: Here, we extend the use of aground geodesic distance to the dynamic formulation.
Revised: Here, we extend the use of aground geodesic distance to the dynamic formulation.

Original: To do so, we define a manifold learning stepthat precedes the trajectory inference.
Revised: To do so, we define a manifold learning stepthat precedes the trajectory inference.

Original: The goal is to learn an embedding Z such that the Euclidean in
Revised: The goal is to learn an embedding Z such that the Euclidean in

Original: Z is equivalent to the geodesic distance.
Revised: Z is equivalent to the geodesic distance.

Original: In that case, dynamic optimal transport in Z is equivalentto the Wasserstein with a geodesic ground distance.
Revised: In that case, dynamic optimal transport in Z is equivalentto the Wasserstein with a geodesic ground distance.

Original: In the following section, we present theoreticalresults that will justify our approximation of a geodesic distance on a closed Riemannian manifold.
Revised: In the following section, we present theoreticalresults that will justify our approximation of a geodesic distance on a closed Riemannian manifold.

Original: 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 representingthe shortest path between two points on the manifold.

Original: We note h t the heat kernel on M .
Revised: We note h t the heat kernel on M .

Original: 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 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.

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: 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 .

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 [21]).
Revised: Next, we state an important result from Leeb and Coifman (Thm. 2 in [22]).

Original: This theorem links thediffusion ground distance and the geodesic on a Riemannian manifold.
Revised: This theorem links thediffusion ground distance and the geodesic on a Riemannian manifold.

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

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

Original: α is equivalent to d 2 α M .
Revised: α 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 [8] provide a natural approximation.
Revised: However, convergenceresults from Coifman and Lafon [9] provide a natural approximation.

Original: Indeed, for β = 1 , the diffusion operator P t/ϵϵ,β converges to the heat operator as ϵ goes to zero (Prop. 3 [8]).
Revised: Indeed, for β = 1 , the diffusion operator P t/ϵϵ,β converges to the heat operator as ϵ goes to zero (Prop. 3 [9]).

Original: For the rest of the paperwe assume β = 1 , we define P ϵ :
Revised: For the rest of the paperwe assume β = 1 , we define P ϵ :

Original: = P ϵ, 1 and the matrix P ϵ :=
Revised: = P ϵ, 1 and the matrix P ϵ :=

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 whichpreserves 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 [7], allowing us to learn non-linear paths and interpolate between timepoints.

Original: 3.1 Geodesic Autoencoder Embedding
Revised: 3.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 .
Revised: 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 .

Original: We then train an autoencoderregularized to match this distance.
Revised: We then train an autoencoderregularized to match this distance.

Original: Similar ideas are explored in [26, 10] where the latent space of theautoencoder is regularized to match a manifold learning embedding, either PHATE or diffusion map.
Revised: 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 the diffusion matrix P ϵ .

Original: We first recall its construction.
Revised: We first recall its construction.

Original: We build a graphwith the affinity matrix defined by ( K ϵ ) ij :
Revised: We build a graphwith the affinity 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 normalizedmatrix 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 [28].
Revised: In practice, wechoose k ϵ to be a Gaussian kernel or the α -decay kernel [29].

Original: Lastly, a Markov diffusion operatoris defined by P ϵ :
Revised: Lastly, a Markov diffusion operatoris defined by 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 stationarydistribution associated 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, wenote that the matrix P ϵ follows a similar construction as the kernel p ϵ ; the integrals are approximatedby sums.

Original: We refer to [8] for more details about the convergence of the matrix operator P ϵ to theoperator P ϵ .
Revised: We refer to [9] for more details about the convergence of the matrix operator P ϵ to theoperator 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 turnconverges to 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 on the 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 thetransition probabilities 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 [41], the authors use the diffusion operator P ϵ to define adistance equivalent to the Wasserstein with ground distance D α .
Revised: In [43], the authors use the diffusion operator P ϵ to define adistance equivalent to the Wasserstein with ground distance D α .

Original: Their method comes from the staticformulation of optimal transport.
Revised: Their method comes from the staticformulation of optimal transport.

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 .
Revised: Here we propose to learn a space Z that preserves an approximationof the distance D α , in order to do dynamic optimal transport in Z .

Original: Training To use the diffusion geodesic distance, we train a Geodesic Autoencoder , with encoderoutputting ϕ : R k
Revised: Training To use the diffusion geodesic distance, we train a Geodesic Autoencoder , with encoderoutputting ϕ : R k

Original: → Z , such that ∥ ϕ ( x i ) − ϕ ( x j ) ∥
Revised: → 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, theencoder can be trained to denoise, hence becoming robust to predicted values that are not close tothe manifold.

Original: Computing G α on the entire dataset is inefficient due to the powering of the diffusionmatrix.
Revised: Computing G α on the entire dataset is inefficient due to the powering of the diffusionmatrix.

Original: We circumvent this difficulty with the encoder, since we train on subsamples.
Revised: We circumvent this difficulty with the encoder, since we train on subsamples.

Original: 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: 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.

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 have  G α ( x i , x j ) ≃ d 2 α M ( x i , x j ) forall x i , x j ∈ X .
Revised: 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 have with high probability G α ( x i , x j ) ≃ d 2 α M ( x i , x j ) for all x i , x j ∈ X .

Original: Corollary 1.
Revised: Corollary 1.

Original: 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 .
Revised: If the encoder is such that L ( ϕ ) = 0 , then with high probability ∥ ϕ ( x i ) − ϕ ( x j ) ∥ 2 ≃d 2 α M ( x i , x j ) for all x i , x j ∈ X .

Revised: The stochasticity in the two previous results arises from the discrete approximation of the operator P ϵ .

Revised: The law or large numbers guarantees the convergence.

Revised: 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 so that the diffusion distance is equivalent to the geodesic on themanifold.
Revised: 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 , we can also train a decoder ϕ − 1 : Z → R k , with a reconstructionloss L r := (cid:80) x
Revised: From that embedding Z , wecan also train a decoder ϕ − 1 : Z → R k , 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 canlearn the trajectories in the lower dimensional space Z , then decode them in the ambient space.
Revised: This is particularly interesting in high dimensions, since we can learn the trajectories in the lowerdimensional space 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 tounderstand how specific genes evolved over time.
Revised: For example, we use the decoder in Sec. 4 to infer cellular trajectories in the gene space, enabling us to understand how specific genesevolved over time.

Original: 3.2 Inferring Trajectories
Revised: 3.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

Original: [Equation]
Revised: [Equation]

Original: We adapt the main theorem from [39], to consider the path space, and other types of dissimilaritiesbetween distributions.
Revised: We adapt the main theorem from [41], 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 t is 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 is 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: This 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: 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 .

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 addition we modelthe trajectories in the embed space Z , then the transport is equivalent to the one on the manifold.

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 (4) with a Neural ODE [7], 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 k × 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 (4) for all t ∈ T .

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

Original: (1 /n i ) (cid:80)
Revised: (1 /n i ) (cid:80)

Original: j δ x j for x j ∈ (cid:98) X i .
Revised: j δ x j for x j ∈ (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;local or 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 initialcondition, 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 trajectoriesgiven the initial 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, representingthe second term in (5).

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 (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.

Original: [Equation]
Revised: [Equation]

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

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 ( k 2 ) additional cost per function evaluation [7].

Original: We add a final density loss L d inspired by [39] to encourage the trajectories to stay on the manifold.
Revised: We add a final density loss L d inspired by [41] 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: L d := λ d T − 1 (cid:88) t = (cid:88) x ∈ (cid:98)
Revised: [Equation]

Original: X t ℓ d ( x, t ) , where ℓ d ( x, t ) :
Revised: ℓ d ( x, t ) , where ℓ d ( x, t ) :

Original: = k (cid:88)
Revised: = k (cid:88)

Original: i = max(0 , min-k ( {∥ x − y ∥ :
Revised: i = max(0 , min [Equation] k ( {∥ x [Equation]

Revised: y

Original:  y ∈ Xt } ) − h ) .
Revised: [Equation] y [Equation] X

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: [Equation] h ) [Equation] 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 [24], converging tothe W 2 formulation.
Revised: Hence, as shown in [25], 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:  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: [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.

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) [34] on held-out timepoints.
Revised: W 1 , and Maximum Mean Discrepancy (MMD) [36] on held-out timepoints.

Original: The MMD is evaluatedwith the Gaussian (G) kernel and the mean (M).
Revised: 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 [39] which is based on a CNF that is regularized toachieve efficient paths, and Diffusion Schrödinger’s Bridge (DSB) [9] which is an optimal transportframework for generative modeling.
Revised: 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.

Original: The baseline measure in the quantitative results correspondsto the average distance between the previous and next timepoints for a given time t .
Revised: The baseline measure in the quantitative results correspondsto the average distance between the previous and next timepoints for a given time t .

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

Original: 4.1 Artificial Data
Revised: 4.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. 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 W 1 and MMD(G), while requiring much less computationaltime.

Original: Dyngen We use Dyngen [6] to simulate a scRNA-seq dataset from a dynamical cellular process.
Revised: Dyngen We use Dyngen [6] 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 reduction method PHATE [28].
Revised: We first embed the observations in five dimensions using the non-linear dimensionality reductionmethod PHATE [29].

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

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 petaldataset.

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 tothe 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 lesspronounce than the upper 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: we showthe 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 learnscontinuous bifurcating trajectories 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. 1, we hold-out timepoint 2,and show the interpolation results against the ground truth.

Original: MIOFlow is more accurate for all metrics,while requiring much less computational time.
Revised: MIOFlow is more accurate for all metrics,while 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: 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).

Original: We have further quantifications reflecting this in the supplementary material.
Revised: We have further quantifications reflecting this in the supplementary material.

Original: [Figure]
Revised: ground truth points as well as training time

Original: [Table]
Revised: [Table]

Original: 4.2 Single-Cell Data
Revised: 4.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 3 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 3.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 asthe lineages mature [22].
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 as the lineages mature [23].

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 [43].
Revised: Another example is ONECUT2 which is consistently known to be higher in cells that divergetowards the neuronal lineage [45].

Original: Other examples are discussed in the supplementary material.
Revised: Other examples are discussed in the supplementary material.

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: ONECUT biologically 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 this purpose, we used data from Fennell et al.
Revised: For thispurpose, we used data from Fennell et al.

Original: [11] in which AF9 mutant cancer cells implanted into thebone marrow of mice as a model of Acute Myeloid Leukemia (AML).
Revised: [12] 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 [11].
Revised: SLPI, MYB, and B2M—all members of the LSC pathway [12].

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 manifoldlearning framework, 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 showtheoretically how MIOFlow is related to dynamical optimal transport on an underlying manifold, andit’s capacity to 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 continuoustrajectories more in line with the manifold compared 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-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.

Original: This way we could visualize the evolution of the gene expression of a cell.
Revised: This way we could visualize 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.
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: A.1 Proof of Thm. 
Revised: A.1 Proof of Thm.

Original: Theorem 2.
Revised: Theorem 2.

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) .
Revised: 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) .

Original: Proof.
Revised: Proof.

Original: We adapt the convergence proof from [41], it can be summarized in two parts.
Revised: We adapt the convergence proof from [43], it can be summarized in two parts.

Original: First, theequivalence of the distances defined by P t/ϵϵ and the heat operator H t .
Revised: First, theequivalence of the distances defined by P t/ϵϵ and the heat operator H t .

Original: This relies on the convergenceresults P t/ϵϵ
Revised: This relies on the convergenceresults P t/ϵϵ

Original: → H t as ϵ → 0 from Coifman and Lafon [8].
Revised: → H t as ϵ → 0 from Coifman and Lafon [9].

Original: Second, the approximation of the operator
Revised: Second, the approximation of the operator

Original: P ϵ based on a finite set of points.
Revised: P ϵ based on a finite set of points.

Original: Throughout the proof we use the same notation δ x for the delta
Revised: Throughout the proof we use the same notation δ x for the delta

Original: Dirac function, and δ x i for a row vector with i -th element 1 , and 0 otherwise.
Revised: Dirac function, and δ x i for a row vector with i -th element 1 , and 0 otherwise.

Original: For a family of operator ( A t ) t ∈ R we define
Revised: For a family of operator ( A t ) t ∈ R we define

Original: [Equation]
Revised: [Equation]

Original: and equivalently D α,A ( x, y )
Revised: and equivalently D α,A ( x, y )

Original: = (cid:80) 2 − kα
Revised: = (cid:80) 2 − kα

Original: ||
Revised: ||

Original: A 2 − k ( δ x − δ y )
Revised: A 2 − k ( δ x − δ y )

Original: || 1 , where Aδ x := (cid:82) A ( u, · ) δ x ( u ) du .
Revised: || 1 , where Aδ x := (cid:82) A ( u, · ) δ x ( u ) du .

Original: This definition is the same as Def. 1, but for an arbitrary operator.
Revised: This definition is the same as Def. 1, but for an arbitrary operator.

Original: 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: 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

Original: D α,H t ( x i , x j ) ≃ d 2 α M ( x i , x j ) .
Revised: D α,H t ( x i , x j ) ≃ d 2 α M ( x i , x j ) .

Original: We first want to show that, for a sufficiently small ϵ > 0 , we have D α,P t/ϵϵ ( x, y )
Revised: We first want to show that, for a sufficiently small ϵ > 0 , we have D α,P t/ϵϵ ( x, y )

Original: ≃ D α,H t , where
Revised: ≃ D α,H t , where

Original: P is the anisotropic operator defined in Sec. 2.
Revised: P is the anisotropic operator defined in Sec. 2.

Original: From Coifman and Lafon [8], we have ∥ P t/ϵϵ −
Revised: From Coifman and Lafon [9], we have ∥ P t/ϵϵ −

Original: H t ∥
Revised: H t ∥

Original: L 2 ( M ) → 0 as ϵ → 0 , and also
Revised: L 2 ( M ) → 0 as ϵ → 0 , and also

Original: ∥ P t/ϵϵ − H t
Revised: ∥ P t/ϵϵ − H t

Original: ∥
Revised: ∥

Original: L 1 ( M ) → 0 .
Revised: L 1 ( M ) → 0 .

Original: Now let Γ t := P t/ϵϵ − H t , and for γ > 0 choose ϵ > 0 such that ∥ Γ t ( δ x − δ y )
Revised: Now let Γ t := P t/ϵϵ − H t , and for γ > 0 choose ϵ > 0 such that ∥ Γ t ( δ x − δ y )

Original: ∥ 1 < γ ∥
Revised: ∥ 1 < γ ∥

Original: δ x − δ y ∥ 1 = 2 γfor x ̸
Revised: δ x − δ y ∥ 1 = 2 γfor x ̸

Original: = y , we have
Revised: = y , we have

Original: [Equation]
Revised: [Equation]

Original: Thus, for all t > 0 , and all γ > 0 there exists ϵ > 0 such that D α, Γ t ( x, y ) < γ .
Revised: Thus, for all t > 0 , and all γ > 0 there exists ϵ > 0 such that D α, Γ t ( x, y ) < γ .

Original: From the reversetriangle inequality, we have
Revised: From the reversetriangle inequality, we have

Original: [Equation]
Revised: [Equation]

Original: using D α, Γ t ( x, y ) < γ , we obtain
Revised: using D α, Γ t ( x, y ) < γ , we obtain

Original: [Equation]
Revised: [Equation]

Original: According to [52] we can lower bound the heat kernel, and thus the distance D α,H t ( x, y ) > C forsome C > 0 .
Revised: According to [54] we can lower bound the heat kernel, and thus the distance D α,H t ( x, y ) > C forsome C > 0 .

Original: For γ < C/ 2 , and a sufficiently small ϵ > 0 , we have
Revised: For γ < C/ 2 , and a sufficiently small ϵ > 0 , we have

Original: [Equation]
Revised: [Equation]

Original: This proves our first claim, that D α,P t/ϵϵ ( x, y ) ≃
Revised: This proves our first claim, that D α,P t/ϵϵ ( x, y ) ≃

Original: D α,H t for small ϵ > 0 .
Revised: D α,H t for small ϵ > 0 .

Original: Next, we consider the approximation of the operator P t/ϵϵ with a finite set of points.
Revised: Next, we consider the approximation of the operator P t/ϵϵ with a finite set of points.

Original: Now we define
Revised: Now we define

Revised: G

Original: [Equation]
Revised: [Equation]

Original: 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: 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 .

Original: Since we willlet K → ∞ , we ignore the second term in G α
Revised: Since we willlet K → ∞ , we ignore the second term in G α

Original: ( x i , x j ) , as it converges to zero.
Revised: ( x i , x j ) , as it converges to zero.

Original: Similar to Coifmanand Lafon [8] (Sec. 5), we have P → P as n → ∞ , using Monte-Carlo integration (approximation ofintegral with summation by the law of large numbers).
Revised: 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).

Revised: By the strong law of large numbers, the convergence is with probability one.

Revised: For a finite number of samples, we have a high probabilitybound on the convergence, see for example [8, 35].

Original: Now we let N :=
Revised: Now we let N :=

Original: min( K, n ) , and therefore
Revised: min( K, n ) , and therefore

Original: [Equation]
Revised: [Equation]

Original: Hence, if we take ϵ := 2 − K , we have for sufficiently large n and K (implying small ϵ > 0 ),
Revised: Hence, if we take ϵ := 2 − K , we have for sufficiently large n and K (implying small ϵ > 0 ),

Original: G α ( x i , x j ) ≃ D α,H t ≃ d 2 α M ( x i , x j ) from Thm. 1.
Revised: G α ( x i , x j ) ≃ D α,H t ≃ d 2 α M ( x i , x j ) from Thm. 1.

Original: A.2 Proof of Thm.3 
Revised: A.2 Proof of Thm.

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 t 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: Proof.
Revised: Proof.

Original: We recall that
Revised: We recall that

Original: [Equation]
Revised: [Equation]

Original: is equivalent to
Revised: is equivalent to

Original: [Equation]
Revised: [Equation]

Original: with the three constraints
Revised: with the three constraints

Original: [Equation]
Revised: [Equation]

Original: Tong et al.
Revised: Tong et al.

Original: [39] (Thm. 4.1), showed that, for large λ > 0 and ρ t satisfying a) , this minimizationproblem is equivalent to
Revised: [41] (Thm. 4.1), showed that, for large λ > 0 and ρ t satisfying a) , this minimizationproblem is equivalent to

Original: [Equation]
Revised: [Equation]

Original: where KL is the Kullback–Leibler divergence, we note that their proof is valid for any dissimilarity
Revised: where KL is the Kullback–Leibler divergence, we note that their proof is valid for any dissimilarity

Original: D ( ρ 1 || ν ) respecting the identity of indiscernibles.
Revised: D ( ρ 1 || ν ) respecting the identity of indiscernibles.

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

Original: [Equation]
Revised: [Equation]

Original: Assuming the encoder ϕ
Revised: Assuming the encoder ϕ

Original: achieves a loss of zero, then from Cor.
Revised: achieves a loss of zero, then from Cor.

Original: 1 for α ∈ (0 , 1 / 2) ,
Revised: 1 for α ∈ (0 , 1 / 2) ,

Original: we have∥ ϕ
Revised: we have∥ ϕ

Original: ( x i ) − ϕ ( x j ) ∥
Revised: ( x i ) − ϕ ( x j ) ∥

Original: 2 ≃ d 2 α M ( x i , x j ) for all x i , x j ∈ X ⊆ M , and sufficiently large n and K .
Revised: 2 ≃ d 2 α M ( x i , x j ) for all x i , x j ∈ X ⊆ M , and sufficiently large n and K .

Original: That is,there exist c, C > 0 such that c d M ( x i , x j ) ≤
Revised: That is,there exist c, C > 0 such that c d M ( x i , x j ) ≤

Original: ∥ ϕ
Revised: ∥ ϕ

Original: ( x i ) − ϕ ( x j ) ∥ 2 ≤ Cd M ( x i , x j ) for all x i , x j ∈ X .
Revised: ( x i ) − ϕ ( x j ) ∥ 2 ≤ Cd M ( x i , x j ) for all x i , x j ∈ X .

Original: Then, for all π ∈ Π( µ, ν ) , we have[Equation] π
Revised: Then, for all π ∈ Π( µ, ν ) , we have

Original: [Equation]
Revised: [Equation]

Original: dx

Original: [Equation] dy [Equation] and taking the infimum with respect to π ∈ Π( µ, ν ) yields the desired result W 2 ( µ, ν ) ≃ W d 2 α M ( µ, ν ) .
Revised: and taking the infimum with respect to π ∈ Π( µ, ν ) yields the desired result W 2 ( µ, ν ) ≃ W d 2 α M ( µ, ν ) .

Original: A.3 Schrödinger bridge
Revised: A.3 Schrödinger bridge

Original: Let D ( µ, ν ) the space of probability distributions on C :=
Revised: Let D ( µ, ν ) the space of probability distributions on C :=

Original: 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 ( ρ t || N ) subject to ρ t ∈ D ( µ, ν ) , where KL is the Kullback–Leibler divergence.
Revised: 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 ( µ, ν ) [Equation]

Revised: where KL is the Kullback–Leibler divergence.

Original: It also admits a static formulation, i.e. minimizingwith respect to π ∈ Π( µ, ν ) .
Revised: It also admits a static formulation, i.e. minimizingwith respect to π ∈ Π( µ, ν ) .

Original: 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: 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 ) =

Original: KL ( π || π N ) + KL ( ρ yx || N yx ) , where π N is the joint distribution between the initial and final states under the Wiener measure N .
Revised: KL ( π || π N ) + KL ( ρ yx || N yx ) , where π N is the joint distribution between the initial and final statesunder the Wiener measure N .

Original: By choosing ρ yx = N yx , the static formulation of the Schrödingerbridge is
Revised: By choosing ρ yx = N yx , the static formulation of the Schrödingerbridge is

Original: [Equation]
Revised: [Equation]

Original: subject to π [Equation]
Revised: subject to π ∈ Π( µ, ν ) .

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

Original: [Equation]
Revised: [Equation]

Original: where H is the entropy (see [29] and references therein).
Revised: 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 [47], which can besolved with Sinkhorn’s algorithm [51].
Revised: 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) [9], which solve the SB problem with an approximation of the iterative proportional fittingalgorithm.
Revised: We compare our method with Diffusion Schrödinger Bridge(DSB) [10], which solve the SB problem with an approximation of the iterative proportional fittingalgorithm.

Original: B Experiment details
Revised: B Experiment details

Original: An anonymized version of our repository 2 is available for reviewing purpose.
Revised: An anonymized version of our repository 2 is available for reviewing purpose.

Original: For all experiments,we used the Runge-Kutta RK4 ODE solver, for the density loss we used h = 0 .
Revised: For all experiments,we used the Runge-Kutta RK4 ODE solver, for the density loss we used h = 0 .

Original: 01 , and 5 nearestneighbors.
Revised: 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.
Revised: 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.
Revised: 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.
Revised: For each epoch (local or global), we sampled 20 batches of a givensample size without replacement.

Original: For both networks, we optimize with AdamW [50] with defaultparameters.
Revised: For both networks, we optimize with AdamW [52] with defaultparameters.

Original: In practice, we did not use the energy loss, as our trajectories were already smooth, andit requires more careful parametrization (see Sec.
Revised: In practice, we did not use the energy loss, as our trajectories were already smooth, andit requires more careful parametrization (see Sec.

Original: C).
Revised: C).

Original: For all experiments, we used TrajectoryNetwith 1000 iterations and whiten , for the  DSB we used the  default parameters.
Revised: For all experiments, we used TrajectoryNetwith 1000 iterations and whiten . We used the authors’ implementation of DSB with the basic modeland default parameters.

Original: Petal For diffusion geodesic, we experiment with an RBF kernel with scale 0 .
Revised: Petal For diffusion geodesic, we experiment with an RBF kernel with scale 0 .

Original: 1 and α -decay kernel(knn=5), both with maximum time scale 2 5 .
Revised: 1 and α -decay kernel(knn=5), both with maximum time scale 2 5 .

Original: The encoder consists of two linear layers of size 8 and32 , trained for 1000 iterations.
Revised: The encoder consists of two linear layers of size 8 and32 , trained for 1000 iterations.

Original: We used LeakyReLU in between layers of the ODE network, withhidden layers size 16 , 32 , 16 , and initial scales σ t = 0 .
Revised: We used LeakyReLU in between layers of the ODE network, withhidden layers size 16 , 32 , 16 , and initial scales σ t = 0 .

Original: We trained for 40 local epoch with samplesize 60 .
Revised: We trained for 40 local epoch with samplesize 60 .

Original: We used λ d = 35 to weight the density loss and λ e = 0 . 001 .
Revised: We used λ d = 35 to weight the density loss and λ e = 0 . 001 .

Original: For the hold-out experiments,we trained for 30 local and 15 global epoch with λ d = 1 .
Revised: For the hold-out experiments,we trained for 30 local and 15 global epoch with λ d = 1 .

Original: The diffusion parameters are generally small while training, for example when trained onthe entire dataset σ t ∈ { 0 . 18 , 0 .
Revised: The diffusion parameters are generally small while training, for example when trained onthe entire dataset σ t ∈ { 0 . 18 , 0 .

Original: 15 , − 0 .
Revised: 15 , − 0 .

Original: 08 } , and for the hold-out time three σ t ∈{ 0 . 08 , 0 . 02 , 0 .
Revised: 08 } , and for the hold-out time three σ t ∈{ 0 . 08 , 0 . 02 , 0 .

Original: Dyngen The diffusion geodesic was evaluated with an RBF kernel with scale 0 .
Revised: Dyngen The diffusion geodesic was evaluated with an RBF kernel with scale 0 .

Original: 5 , and maximumtime scale 2 5 .
Revised: 5 , and maximumtime scale 2 5 .

Original: The encoder has two linear layers of size 8 and 32 , trained for 1000 iterations.
Revised: The encoder has two linear layers of size 8 and 32 , trained for 1000 iterations.

Original: We used
Revised: We used

Original: CELU in between layers of the ODE network, with hidden layers size 16 , 32 , 16 , and initial scalesσ t = 0 .
Revised: CELU in between layers of the ODE network, with hidden layers size 16 , 32 , 16 , and initial scalesσ t = 0 .

Original: We trained for 50 global epochs of size 60 , with λ d = 5 .
Revised: We trained for 50 global epochs of size 60 , with λ d = 5 .

Original: For the hold-out experiments, weused λ d = 5 , and for t = 5 we trained on 5 local and 10 global epochs.
Revised: For the hold-out experiments, we used λ d = 5 , and for t = 5 we trained on 5 local and 10 global epochs.

Original: For t = 3 , since we have tohold-out the steps (2 , 3) and (3 , 4) , we only trained using 50 global epochs.
Revised: For t = 3 , since we have tohold-out the steps (2 , 3) and (3 , 4) , we only trained using 50 global epochs.

Original: Similar to the petal dataset, the diffusion parameters is empirically low, on the entire dataset it isσ t ∈ { 0 . 18 , 0 . 20 , 0 .
Revised: Similar to the petal dataset, the diffusion parameters is empirically low, on the entire dataset it isσ t ∈ { 0 . 18 , 0 . 20 , 0 .

Original: 16 } , and for the hold-out three σ t ∈ { 0 . 02 , 0 . 08 , 0 . 24 , 0 . 22 , 0 . 24 } .
Revised: 16 } , and for the hold-out three σ t ∈ { 0 . 02 , 0 . 08 , 0 . 24 , 0 . 22 , 0 . 24 } .

Original:  data The EB data are publicly available 3 .
Revised: EB data The EB data are publicly available 3 .

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.
Revised: For the genes trajectories we used the full geodesicautoencoder because we want to decode in PCA space, and then project back into gene space.

Original: Theencoder consists of three layers of size 200 , 100 , 100 , and the decoder is 100 , 100 , 200 , both with
Revised: Theencoder consists of three layers of size 200 , 100 , 100 , and the decoder is 100 , 100 , 200 , both with

Original: ReLU activation between layers.
Revised: ReLU activation between layers.

Original: 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: 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 ).

Original: This kernel is more expensiveto compute since it relies on the approximation of the knn, but it circumvents the need to find a good scale as with the RBF kernel.
Revised: 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.
Revised: The ODE network is three hidden layers of size 64 , with CELUactivation in between layers.

Original: We trained for 20 local and 10 global epochs with sample size of 400per time.
Revised: We trained for 20 local and 10 global epochs with sample size of 400per time.

Original: We used the density loss with weight λ d = 20 .
Revised: We used the density loss with weight λ d = 20 .

Original: AML The AML dataset is publicly available 4 .
Revised: AML The AML dataset is publicly available 4 .

Original: We first embed the data in 50 dimensions PCAspace.
Revised: We first embed the data in 50 dimensions PCAspace.

Original: 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).
Revised: 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).

Original: 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 .
Revised: 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 .

Original: C Additional Results
Revised: C Additional ResultsEnergy

Original: Energy loss In Fig.
Revised: loss In Fig.

Original: 6, we present the trajectories using the same parametrization as detailed in
Revised: 6, we present the trajectories using the same parametrization as detailed in

Original: Sec.
Revised: Sec.

Original: B with the energy loss ( λ e = 1 ).
Revised: B with the energy loss ( λ e = 1 ).

Original: This loss requires much more memory since for each functionevaluation we need to save the norm of the derivative.
Revised: 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.
Revised: 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.
Revised: On Dyngen (left), it learns a straight line, andthus cannot bifurcate.

Original: On the petal, it initially learns a straight trajectory (middle, one epoch), and astraining continues, it collapses in the center (right, ten epochs).
Revised: On the petal, it initially learns a straight trajectory (middle, one epoch), and astraining continues, it collapses in the center (right, ten epochs).

Original: Unless otherwise stated, we did notadd it to our model since our trajectories were smooth.
Revised: Unless otherwise stated, we did notadd it to our model since our trajectories were smooth.

Original: However, with careful parametrization, onecould learn more accurate trajectories trajectories (see Fig. 7).
Revised: However, with careful parametrization, onecould learn more accurate trajectories trajectories (see Fig. 7).

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.
Revised: 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).
Revised: 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 .
Revised: 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 .
Revised: For the EB data, we set the scale to 1500 .

Original: In Fig. 8, we present an ablation study for the density loss and the geodesic embedding.
Revised: In Fig. 8, we present an ablation study for the density loss and the geodesic embedding.

Original: Theparametrization was found for the model with density loss and embedding, both seem important to learn accurate trajectories.
Revised: Theparametrization was found for the model with density loss and embedding, both seem important to learn accurate trajectories.

Original: In Fig.
Revised: In Fig.

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

Original: In Fig.
Revised: In Fig.

Original: 10, we reproduce the same ablationon Dyngen data.
Revised: 10, we reproduce the same ablationon Dyngen data.

Original: Here, the density loss is important for the trajectories to stay on the manifold, andthe embedding appears to help with the bifurcation.
Revised: Here, the density loss is important for the trajectories to stay on the manifold, andthe embedding appears to help with the bifurcation.

Original: To complement this figure, in Tab. 5, we presentthe average of the KNN distance computed on the top 5% .
Revised: To complement this figure, in Tab. 5, we presentthe average of the KNN distance computed on the top 5% .

Original: On Dyngen, the trajectories with hold-outtimepoints are presented in Fig.
Revised: On Dyngen, the trajectories with hold-outtimepoints are presented in Fig.

Original: Holding out the timepoint t = 3 seems to be the hardest to learn.
Revised: Holding out the timepoint t = 3 seems to be the hardest to learn.

Original: In Tab. 2, we show the W 1 and MMD between the predicted values and ground truth, for MIOFlow,
Revised: In Tab. 2, we show the W 1 and MMD between the predicted values and ground truth, for MIOFlow,

Original: TrajectoryNet, and DSB.
Revised: TrajectoryNet, and DSB.

Original: For the W 1 , MIOFlow is always more accurate.
Revised: For the W 1 , MIOFlow is always more accurate.

Original: In Tab 3, we compare theprediction accuracy on the EB data, depending on using the GAE or not.
Revised: In Tab 3, we compare theprediction accuracy on the EB data, depending on using the GAE or not.

Original: 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.
Revised: 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.

Original: Using the GAE improves theinterpolation accuracy for all metrics.
Revised: Using the GAE improves theinterpolation accuracy for all metrics.

Original: We embed the data in 200 PCA, and train the autoencoder withlayers 200 , 100 , 100 , 50 for 1000 epochs.
Revised: 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 .
Revised: 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.
Revised: 2 , linear layers of size 16 , 32 , 16 , 50 with CELU activation in between layers.

Original: We train byholding out timepoint t for the GAE and the neural ODE.
Revised: We train byholding out timepoint t for the GAE and the neural ODE.

Original: For the MMD with the Gaussian kernel, weset the scale to 1500 .t [Equation] t
Revised: For the MMD with the Gaussian kernel, weset the scale to 1500 .

Original: [Equation]

Original: t

Original: [Equation] W MMD(G) MMD(M) W
Revised: DSB.

Original: MMD(G) MMD(M)

Original: W MMD(G) MMD(M) Petal MIOFlow(ours) 0.
Revised: Lower is better.

Original: 0.009 TrajectoryNet 0.

Original: 0.004 DSB 0.

Original: < 0. 0.081 Baseline 0. Dyngen MIOFlow(ours) 0.
Revised: [Table]

Original: 1.719 TrajectoryNet 1. 1.913 DSB 0. 5.546 Baseline 1. 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.
Revised: 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.

Original: Table 4 shows the 1-NN distance for several methods and datasets.
Revised: Table 4 shows the 1-NN distance for several methods and datasets.

Original: Of note is theaggregate column.
Revised: Of note is theaggregate column.

Original: Where aggregate is specified as mean , we compute the mean 1-NN distance from the predicted point to the 1-NN at the same time index.
Revised: 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.

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 .
Revised: 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 .

Original: Neural SDE Additionally, we employed a Neural SDE from the PyTorch implementation ofdifferentiable SDE Solvers [49, 48], and compared its results to MIOFlow (Fig.
Revised: Neural SDE Additionally, we employed a Neural SDE from the PyTorch implementation ofdifferentiable SDE Solvers [51, 50], and compared its results to MIOFlow (Fig.

Original: 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: 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.

Original: D Related Work
Revised: D Related Work

Original: Static methods such as pseudotime assign an abstract measure of time for all cells depending onan initial state [16, 42].
Revised: Static methods such as pseudotime assign an abstract measure of time for all cells depending onan initial state [17, 44].

Original: In can be combined to infer the cell lineage, i.e. the principal curves ortrajectories [38].
Revised: In can be combined to infer the cell lineage, i [Table] e. the principal curves ortrajectories [40].

Original: When estimates of cell growth rates are available, Zhang et al.
Revised: When estimates of cell growth rates are available, Zhang et al.

Original: [45] proposed amethod based on optimal transport.
Revised: [47] proposed a [Table] method based on optimal transport.

Original: Another line of research is on the direct rate of differentiation ofcell or RNA velocity [19, 4].
Revised: Another line of research is on the direct rate of differentiation ofcell or RNA velocity [20, 4].

Original: This method does not provide trajectories, but it can be combined todefine a transition matrix to infer trajectories [46].
Revised: This method does not provide trajectories, but it can be combined todefine a transition matrix to infer trajectories [48].

Original: Saelens et al.
Revised: Saelens et al.

Original: [32] provide a rigorous comparisonof 45 static methods.
Revised: [33] provide a rigorous comparisonof 45 static methods.

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

Original: Similar methods used optimal transport [33] either with a recurrent [Table] neural network [17], generative adversarial networks [31], flows defined by the gradient of an inputconvex neural network [5], or a CNF [39].
Revised: 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].

Original: 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.
Revised: 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.

Original: Some generative methods that interpolate between a prior to a targetcan also be used to recover trajectories.
Revised: Some generative methods that interpolate between a prior to a target canalso be used to recover trajectories.

Original: For example [9], solve the Schrödinger bridge (SB) betweentwo distributions.
Revised: For example [10], solve the Schrödinger bridge (SB) betweentwo distributions.

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

Original: E Software Versions
Revised: E Software Versions

Original: 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==0.2.
-]
Revised: 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==0.2.
-]torchsde==

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Supplement References
Revised: Supplement References