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

Original: 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 static snapshot samples taken at sporadic timepoints.

Original: MIOFlow combines dynamic models, manifold learning, and optimal transport by training neural ordinary differentialequations (Neural ODE) to interpolate between static population snapshots as penalized by optimal transport with manifold ground distance.
Revised: MIOFlow combines dynamic models, manifold learning, and optimal transport by training neural ordinary differential equations (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 that the flow follows the geometry by operating in the latent space of an autoencoder that we call a geodesic autoencoder (GAE).

Original: In GAE the latent space distancebetween points is regularized to match a novel multiscale geodesic distance on thedata manifold that we define.
Revised: In GAE the latent space distance between points is regularized to match a novel multiscale geodesic distance on the data 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 normalizing flows, Schrödinger bridges and other generative models that are designed to flow from noise to data in terms of interpolating between populations.

Original: Theoretically, welink these trajectories with dynamic optimal transport.
Revised: Theoretically, we link 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 method on simulated data with bifurcations and merges, as well as scRNA-seq data from embryoid 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 data manifold.

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 low dimensional 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 manifold hypothesis, this assumption has led to many successful models of biological, chemical, and physical systems.

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 single cell data the entire transcriptomic profile of the cell is measured with each mRNA or gene species a 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 intrinsic dimensionality of the data is low dimensional.

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

Original: Laplacian eigenmaps [2] and the wider graph signal processing literature, has focused on learningthis structure via data graphs [28, 29] and autoencoders [27, 11].
Revised: Laplacian eigenmaps [2], and the wider graph signal processing literature, has focused on learning this structure via data graphs [31, 32] and autoencoders [30, 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 [38, 39], diffusion models [19], Schrödinger bridges [30, 10], and continuousnormalizing flows (CNF) 
Revised: Examples include score-based generative matching [41, 42], diffusion models [21], Schrödinger bridges [33, 10], and continuous normalizing flows (CNF) [7, 17].

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 such systems 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) 3 , a new framework for flows based on dynamical 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 differential equations (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 penalized to 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 [41], dynamics modeled by a recurrent neuralnetwork [18], and Waddington-OT [34].
Revised: Works tackling a similar problem include TrajectoryNet [44], dynamics modeled by a recurrent neural network [20], and Waddington-OT [38].

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

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

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

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

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

Original: Hence, to model the intrinsic stochasticity in biology, we have to forcechaotic dynamics starting from slight noise added to any initial distributions.
Revised: Hence, to model the intrinsic stochasticity in biology, we have to force chaotic 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, continuous normalizing 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 flows on 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 manifold distance called diffusion geodesic distance.

Original: Second, we penalize the transport using a manifoldoptimal transport method [31, 46].
Revised: Second, we penalize the transport using a manifold optimal transport method [34, 50].

Original: These two steps enforce flows on the manifold, whereas in earlierworks, such as TrajectoryNet, the KL-divergence is used to match distributions, and in [24] access tothe metric tensor of a Riemannian manifold is required.
Revised: These two steps enforce flows on the manifold, whereas in earlier works, such as TrajectoryNet, the KL-divergence is used to match distributions, and in [27] access to the 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 results showing 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 on single-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 time t ∈ 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: For each time t , we observe a sample X t ∼ µ t of size n t .

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

Original: We also want to characterize the evolution of the support of µ t .
Revised: We also want to characterize the 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 define a 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 T samples { X i } T − 1 i =1 , and initial condition X 0 .

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

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

Original: We  the equivalence  two distances d 1 ≃ d 
Revised: We note the equivalence between two distances d 1 ≃ d 2 .

Original:  Stochastic Differential Equations ) admit a solution [14] proofs are presented in the supplementary material
Revised: We assume that all Stochastic Differential Equations (SDE) admit a solution [15]. All proofs are presented in the supplementary material.

Original: Optimal Transport In this section, we provide a brief overview of optimal transport [31, 46],our primary approach for interpolating between distributions using a neural ODE.
Revised: Optimal Transport In this section, we provide a brief overview of optimal transport [34, 50], 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 two distributions µ and ν defined on X and Y , and Π( µ, ν ) the set of joint distributions on X × Y with marginals µ and ν , i.e. π ( dx, Y ) = µ ( dx ) and π ( X , dy ) =

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

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

Original: This formulation gives rise to the p -Wasserstein distance
Revised: This 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 a fix 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 transport plan is replaced by a time-evolving distribution ρ t , such that ρ 0 = µ and ρ 1 = ν and that satisfies the 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 movement of 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 of a 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 [25]).
Revised: ∼ ν [Equation] where the infimum is over all absolutely continuous stochastic path X t (see [28]).

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 in the paths or trajectories.

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

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 this formulation, one can also retrieve the Benamou-Brenier optimal transport.

Original: Indeed, Mikami [25]established the convergence to the W 2 when the diffusion term goes to zero.
Revised: Indeed, Mikami [28] 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 can phrase the problem using an SDE

Original: [Equation]
Revised: [Equation]

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

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

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 data originates from an intrinsic low dimensional manifold that is mapped via nonlinear functions to observable 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 of k ≫ 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 dimensional ambient 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 is collected.

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 [-=
-] ∈ M ofthe underlying manifold via x i = 
Revised: = { z i } ni [-=
-]=1 ∈ M of the underlying manifold via x i = f

Original: f ( z i ) , i = 1 , . . . , n , and aim to learn a low dimensional intrinsicrepresentation that approximates the manifold geometry of M .
Revised: ( z i ) , i = 1 , . . . , n , and aim to learn a low dimensional intrinsic representation 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 [9] proposed to use an anisotropic kernel k ϵ,β ( x, y ) :
Revised: To learn a manifold geometry from collected data that is robust to sampling density variations, 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 ϵ ( [Equation] , · ) ∥ β 1 ∥ k ϵ ( y, · ) ∥ β 1 , where 0 ≤ β ≤ 1 controls the separation of geometry from density,with β = 0 yielding the isotropic kernel, and β = 1 completely removing density and providinga geometric equivalent to uniform sampling of the underlying manifold.
Revised: k ϵ ( x, · ) ∥ β 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 providing a 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 ϵ,β is normalized 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 := p ϵ,β ( x i , x j ) that describes a Markovian diffusion process over the intrinsic geometry of the data.
Revised: k ϵ,β ( x, · ) ∥ 1 and an n × n row stochastic matrix ( P ϵ,β ) ij := p ϵ,β ( x i , x j ) that describes a Markovian diffusion process over the intrinsic geometry of the data.

Original: Finally, a diffusion map [9] Φ 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 eigenvalues and 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, we consider a multiscale diffusion distance that relates to the geodesic distance on the manifold.

Original: ManifoldGeodesics
Revised: Manifold

Original:  The dynamic optimal transport formulations (1) and (2) are valid when theground distance is Euclidean.
Revised: Geodesics The dynamic optimal transport formulations (1) and (2) are valid when the ground 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 lower dimensional 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 underlying geometry of the data.

Original: Static optimal transport with a geodesic distance is known to perform well ongeometric domains [37, 42], or when the data lie on a manifold [43].
Revised: Static optimal transport with a geodesic distance is known to perform well on geometric domains [40, 45], or when the data lie on a manifold [46].

Original: Here, we extend the use of aground geodesic distance to the dynamic formulation.
Revised: Here, we extend the use of a ground 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 step that 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 equivalent to 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 theoretical results 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 representing the 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 point x ∈ M , the heat kernel h t ( x, · ) induces a measure, i.e. how the heat has propagated on the manifold at time t given an arbitrary initial distribution.

Original: The diffusion ground distance between x, y ∈ M isbased on the L 1 norm between the measures induced by the heat kernel given the initial conditionsδ x and δ y .
Revised: The diffusion ground distance between x, y ∈ M is based 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: ) || 1 , for α ∈ (0 , 1 / 2) , the scale parameter k ≥ 0 , and h t the heat kernel on M .

Original: Next, we state an important result from Leeb and Coifman (Thm. 2 in [22]).
Revised: Next, we state an important result from Leeb and Coifman (Thm. 2 in [25]).

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

Original: Theorem 1 ([22] Thm. 2) .
Revised: Theorem 1 ([25] 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 [9] provide a natural approximation.
Revised: However, convergence results 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 [9]).
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 paper we 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 which preserves the diffusion geodesic distance—which we define using our Geodesic Autoencoder.

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

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

Original: Neural ODE [7], allowing us to learn non-linear paths and interpolate between timepoints.
Revised: Neural ODE [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) from a 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 autoencoder regularized to match this distance.

Original: Similar ideas are explored in [27, 11] where the latent space of theautoencoder is regularized to match a manifold learning embedding, either PHATE or diffusion map.
Revised: Similar ideas are explored in [30, 11] where the latent space of the autoencoder 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 graph with 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 normalized matrix M ϵ : =

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

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

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

Original: =
Revised: =

Original: D − 1 M ϵ , where D ii := (cid:80) nj =1 ( M ϵ ) ij is a diagonal matrix.
Revised: D − 1 M ϵ , where D ii := (cid:80) nj =1 ( M ϵ ) ij is a diagonal matrix.

Original: The stationarydistribution associated to P ϵ is π , where π i = D ii / (cid:80) j D jj , since P ϵ is π -reversible.
Revised: The stationary distribution 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, we note that the matrix P ϵ follows a similar construction as the kernel p ϵ ; the integrals are approximated by sums.

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

Original: Until now, we have defined a matrix P ϵ that approximates the operator P ϵ , which in turnconverges to the heat operator.
Revised: Until now, we have defined a matrix P ϵ that approximates the operator P ϵ , which in turn converges 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 the transition 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 scales given two different initial states x i , x j .

Original: In [43], the authors use the diffusion operator P ϵ to define adistance equivalent to the Wasserstein with ground distance D α .
Revised: In [46], the authors use the diffusion operator P ϵ to define a distance equivalent to the Wasserstein with ground distance D α .

Original: Their method comes from the staticformulation of optimal transport.
Revised: Their method comes from the static formulation 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 approximation of 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 encoder outputting ϕ : 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 (MSE)

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 since we use it to compute distances between predicted and ground truth observations.

Original: Moreover, theencoder can be trained to denoise, hence becoming robust to predicted values that are not close tothe manifold.
Revised: Moreover, the encoder can be trained to denoise, hence becoming robust to predicted values that are not close to the 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 diffusion matrix.

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 N to have few observations in most regions of the manifold, thus making the computation of G α very efficient, 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 .
Revised: Assuming X is sampled from a closed Riemannian manifold M with geodesic d M .

Original: Then, for all α ∈ (0 , 1 / 2) , sufficiently large K, N and small ϵ > 0 , we have with high probability
Revised: Then, for all α ∈ (0 , 1 / 2) , sufficiently large K, N and small ϵ > 0 , we have with high probability

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

Original: ( x i ) − ϕ ( x j ) ∥ 2 ≃d 2 α M ( x i , x j ) for all x i , x j ∈ X .
Revised: ( x i ) − ϕ ( x j ) ∥ 2 ≃ d 2 α M ( x i , x j ) for all x i , x j ∈ X .

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

Original: The law or large numbers guarantees the convergence.
Revised: The law or large numbers guarantees the convergence.

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

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

Original: From that embedding Z , wecan also train a decoder ϕ − 1 : Z → R k , with a reconstruction loss L r := (cid:80) x
Revised: From that embedding Z , we can 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 can learn the trajectories in the lowerdimensional 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 lower dimensional space Z , then decode them in the ambient space.

Original: For example, we use the decoder in
Revised: For example, we use the decoder in

Original: Sec. 4 to infer cellular trajectories in the gene space, enabling us to understand how specific genesevolved over time.
Revised: Sec. 4 to infer cellular trajectories in the gene space, enabling us to understand how specific genes evolved over time.

Revised: We describe the training procedure of the GAE in algorithm 2 in the supplementary material.

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 ∼ µ i for 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 [41], to consider the path space, and other types of dissimilaritiesbetween distributions.
Revised: We adapt the main theorem from [44], to consider the path space, and other types of dissimilarities between 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 ( µ, ν )

Original: {+= 
 iff µ =+} ν .
Revised: {+= 
 iff µ =+} ν .

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 with geodesic 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: In practice, it justifies the method of matching the marginals and adding a regularization on the vector field f ( x, t ) , that way the optimal ( ρ t , f ) corresponds to the one from the W 2 .

Original: If in addition we modelthe trajectories in the embed space Z , then the transport is equivalent to the one on the manifold.
Revised: If in addition we model the 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: Denote by ψ θ : R k × T → R d |T | the function that represents the Neural ODE, where T is a set of time indices.

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 [-
-] it returns the approximation of (4) for all t ∈ T .
Revised: , T − 1 } ) ; given an initial set X [-
-]0 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 initial condition, we predict (cid:98) X t +1 =

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

Original: Whereas for the global, we use the entire trajectoriesgiven the initial condition X 0 .
Revised: Whereas for the global, we use the entire trajectories given 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, representing the 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 integral can be approximate with the forward pass of the ODE solver, and λ e ≥ 0 is a hyperparameter.

Original: [Equation]
Revised: [Equation]

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

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-likelihood approach in CNF methods which requires evaluation of the instantaneous change of variables at every integration timestep at O ( k 2 ) additional cost per function evaluation [7].

Original: We add a final density loss L d inspired by [41] to encourage the trajectories to stay on the manifold.
Revised: We add a final density loss L d inspired by [44] 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 t given a lower bound h >[-
-]

Original: [Equation]
Revised: [Equation]

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

Original: = k (cid:88)

Original: i = max(0 , min

Original: [Equation]
Revised: [Equation]

Original: k ( {∥ x [Equation] y
Revised: (cid:88) i =

Original: [Equation]
Revised: [Equation]

Original: y [Equation] X
Revised: We describe the overall training procedure of MIOFlow in algorithm 1.

Original: [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
Revised: Modeling Diffusion We add a diffusion term in the trajectories, resulting in the SDE dX t = f ( X t , t ) dt + √ σ t dB t

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

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

Original: It has the advantage of mappingthe same cell to different trajectories stochastically as in many naturalistic systems.
Revised: It has the advantage of mapping the 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 the training 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: As training continues, we empirically observed that σ → 0 .

Original: Hence, as shown in [25], converging tothe W 2 formulation.
Revised: Hence, as shown in [28], converging to the 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 to the 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, thus drastically increasing the computational cost. Algorithm

Revised: Input: Datasets X 0 , . . .

Revised: , X T − 1 , graph kernel K ϵ , maximum scale K , noise scale ξ , batch size N , maximum iteration n max , initialized encoder ϕ and decoder ϕ − 1 , maximum local iteration n local , maximum global iteration n global , initialized neural ODE ψ θ .

Original: Results
Revised: Output: Trained neural ODE ψ θ . GAE: ϕ, ϕ − 1 ← GAE ( X , K ϵ , K, ξ, N, n max , ϕ, ϕ − 1 ) ▷ See algorithm 2 in SM if use GAE then X 0 , . . . , X T − 1 ← ϕ

Revised: ( X 0 ) , . . .

Revised: , ϕ ( X T − 1 ) end if for i=1 to n local do ▷ Local training for t=0 to T − 2 do (cid:98)

Revised: X t +1 ← ψ θ ( X t , t + 1) end for

Revised: L ← L m + L e + L d update ψ θ with gradient descent w.r.t.

Original: [Figure]
Revised: the loss L end for for i=1 to n global do ▷ Global training (cid:98) X 1 , . . . , (cid:98) X T − 1 ← ψ θ ( X 0 , { 1 , . . . , T − 1 } ) L ← L m + L e + L d update ψ θ with gradient descent w.r.t.

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: the loss L end for  Results Here, we validate our method empirically on both synthetic test cases and single-cell RNA sequencing datasets chosen to test the capability of MIOFlow to follow divergent trajectories from population snapshot 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 qualitatively follow the data manifold efficiently, we also quantitatively assess their accuracy using the Wasserstein

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

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

Original: We compare our results to those other methods thatpreform population flows including TrajectoryNet [41] which is based on a CNF that is regularized toachieve efficient paths, and Diffusion Schrödinger’s Bridge (DSB) [10] which is an optimal transportframework for generative modeling.
Revised: We compare our results to those other methods that preform population flows including TrajectoryNet [44] which is based on a CNF that is regularized to achieve efficient paths, and Diffusion Schrödinger’s Bridge (DSB) [10] which is an optimal transport framework 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 corresponds to 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: Additional experiments, 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 cellular differentiation, 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) and evolves 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 dataset for 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 keeping trajectories 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 more accurate entropic interpolation for both W 1 and MMD(G), while requiring much less computational time.

Revised: In the supplementary material, we present additional results for DSB on the petal dataset.

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 reductionmethod PHATE [29].
Revised: We first embed the observations in five dimensions using the non-linear dimensionality reduction method PHATE [32].

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

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

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

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

Original: Moreover, the bifurcation is asymmetric; the curve in the lower branch is lesspronounce than the upper one, making it harder to learn a continuous bifurcation.
Revised: Moreover, the bifurcation is asymmetric; the curve in the lower branch is less pronounce than 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: 3, we show the trajectories for TrajectoryNet, DSB, and our method.

Original: MIOFlow is more accurate as it learnscontinuous bifurcating trajectories along the two main branches.
Revised: MIOFlow is more accurate as it learns continuous 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 far from 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: 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 course of human embryoid body (EB) differentiation over a period of 27 days, measured in 3 timepoints between days 0-3, 6-9, 12-15, 18- [Table] and 24-27.

Original: We reduce the dimensionality of this dataset to 200dimensions using PCA.
Revised: We reduce the dimensionality of this dataset todimensions 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 steps and 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, we can decode the trajectories back into the gene space (17846 dimensions), thus describing the evolution of 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 neuronal cells 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 divergence in trends artificially showing higher heterogeneity in gene trends then exist in neuronal specification processes.

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

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

Original: Another example is ONECUT2 which is consistently known to be higher in cells that divergetowards the neuronal lineage [45].
Revised: Another example is ONECUT2 which is consistently known to be higher in cells that diverge towards the neuronal lineage [48].

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

Original: ONECUT biologically
Revised: we compare the prediction accuracy of MIOFlow whether we use the GAE embedding or no embedding. For the GAE we show results for two kernels; the Gaussian kernel (scale 0 . 5 ), and the α -decay kernel [32] (based on a 30 nearest neighbors graph). For almost all metrics, the prediction accuracy is greater with the GAE. We present additional results for the choice of kernel 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: [Equation] MMD( [Table] AML Cancer Data with Treatment-based Progression We also assessed whether MIOFlow can give insights on treatment resistant cancer cell populations, rather than differentiation.

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

Original: [12] in which AF9 mutant cancer cells implanted into thebone marrow of mice as a model of Acute Myeloid Leukemia (AML).
Revised: [12] in which AF9 mutant cancer cells implanted into the bone 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 treated for 7 days with intensive chemotherapy with single cell measurements taken at day 2, day 5, and day 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 the trajectories, the majority of cells surviving to the last timepoint showing an increase in the markers

Original: SLPI, MYB, and B2M—all members of the LSC pathway [12].
Revised: SLPI, MYB, and B2M—all members of the LSC pathway [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 start with this signature but are made to transition to this signature to resist the effects of chemotherapy.  Related Work

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

Revised: Static methods such as pseudotime assign an abstract measure of time for all cells depending on an initial state [19, 47].

Original: Conclusion
Revised: In can be combined to infer the cell lineage, i.e. the principal curves or trajectories [43]. When estimates of cell growth rates are available, Zhang et al. [52] proposed a method based on optimal transport. Another line of research is on the direct rate of differentiation of cell or RNA velocity [23, 4].

Revised: This method does not provide trajectories, but it can be combined to define a transition matrix to infer trajectories [53].

Revised: Saelens et al.

Revised: [37] provide a rigorous comparison of 45 static methods.

Revised: Similar methods used the dynamic formulation of optimal transport [5, 13, 20, 35, 36, 38, 44] to infer trajectories.

Revised: The dynamics are either modeled with a recurrent neural network [20], a generative adversarial networks [35], flows defined by the gradient of an input convex neural network [5], or a CNF [44].

Revised: Here, we circumvent the need for a CNF by modeling the trajectories directly on the support of the distributions, hence maintaining the computational cost low even in high dimension.

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

Revised: Other methods use the Schrödinger bridge (SB) [10, 22, 49, 51] problem to define trajectories between distributions.

Revised: For example [10], solve the SB between two distributions using a series of Sinkhorn iterations.

Revised: However, it does not explicitly learn dynamics, making their interpolations inaccurate.

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 multiple cross-sectional samples of a time-evolving population.

Original: Our method includes a novel manifoldlearning framework, where we learn an embedding preserving the geodesic distance.
Revised: Our method includes a novel manifold learning 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 show theoretically how MIOFlow is related to dynamical optimal transport on an underlying manifold, and it’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 continuous trajectories 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-seq data, we leveraged an autoencoder with the geodesic embedding to learn the trajectories in a lower dimensional 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 stiff through 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 the time-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 be alleviated by using unbalanced optimal transport, which allows transporting distributions with unequal masses.

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.

Revised: Acknowledgments and Disclosure of Funding

Original: {+This work was partially funded and supported by CIFAR AI Chair [G.W.], NSERC Discovery grant 03267 [G.W.], NIH grants 1F30AI157270-01 [M.K.], R01GM135929 [G.W., S.K.], R01HD100035, and R01GM130847 [G.W., S.K.], NSF Career grant 2047856 [S.K.], the Chan-Zuckerberg Initiative grants CZF2019-182702 and CZF2019-
+}
Revised: {+This work was partially funded and supported by CIFAR AI Chair [G.W.], NSERC Discovery grant 03267 [G.W.], NIH grants 1F30AI157270-01 [M.K.], R01GM135929 [G.W., S.K.], R01HD100035, and R01GM130847 [G.W., S.K.], NSF Career grant 2047856 [S.K.], the Chan-Zuckerberg Initiative grants CZF2019-182702 and CZF2019-
+}

Revised: [S.K.], and the Sloan Fellowship FG-2021- [S.K.].

Revised: The content provided here is solely the responsibility of the authors and does not necessarily represent the official views of the funding agencies.

Revised: The funders had no role in study design, data collection & analysis, decision to publish, or preparation of the manuscript.

Original:  A.1 Proof of Thm.
Revised: A Theory and Algorithm A.1 Geodesic Autoencoder Algorithm

Revised: Using the definitions from Sec. 3, we present the algorithm to train the geodesic autoencoder. Algorithm Input: Dataset X of size n , graph kernel K ϵ , maximum scale K , noise scale ξ , batch size N , maximum iteration n max , initialized encoder ϕ and decoder ϕ − 1 .

Revised: Output: Trained encoder ϕ and decoder ϕ − 1 .

Revised: for i=1 to n max do

Revised: Sample batch { x 1 , . . .

Revised: , x N } ⊆ X of size N

Revised: Q ← diag ( (cid:80) j ( K ϵ ) ij )

Revised: M ϵ ← Q − 1 K ϵ Q −

Revised: D ← diag ( (cid:80) j ( M ϵ ) ij )

Revised: P ϵ ← D − 1 M ϵ ▷ N × N diffusion matrix

Revised: G α ← 2 − Kα pairwise_distances (( P ϵ ) 1: , . . .

Revised: , ( P ϵ ) N : ) for k=1 to K do

Revised: G α ← G α + 2 − ( K − k )

Revised: α pairwise_distance (( P 2 k ϵ ) 1: , . . .

Revised: , ( P 2 k ϵ ) N : ) end for ˜ x 1 , . . . ,

Revised: ˜ x N ← x 1 + ξz 1 , . . .

Revised: , x N + ξz N , where { z 1 , . . .

Revised: , z N } ∼ N (0 , 1)

Revised: L ← MSE (cid:0) pairwise_distances ( ϕ (˜ x 1 ) , . . .

Revised: , ϕ (˜ x N )) , G α (cid:1)

Revised: L r ← (cid:80) i ∥ ϕ − 1 ◦ ϕ

Revised: (˜ x i ) − x ∥

Revised: Update ϕ, ϕ − 1 with gradient descent w.r.t.

Revised: L ( ϕ ) , L r ( ϕ, ϕ − 1 ) end for

Revised: A.2 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 [43], it can be summarized in two parts.
Revised: We adapt the convergence proof from [46], 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, the equivalence 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 convergence results P t/ϵϵ

Original: → H t as ϵ → 0 from Coifman and Lafon [9].
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) [Equation]

Original: ||
Revised: ||

Original: A 2 − k ( δ x − δ y )
Revised: [Equation] k [Equation] − δ 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 show that 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 [9], 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 ( δ

Original: ∥ 1 < γ ∥
Revised: [Equation]

Original: δ x − δ y ∥ 1 = 2 γfor x ̸ 
Revised: 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 reverse triangle 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 [54] we can lower bound the heat kernel, and thus the distance D α,H t ( x, y ) > C forsome C > 0 .
Revised: According to [59] we can lower bound the heat kernel, and thus the distance D α,H t ( x, y ) > C for some 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

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

Original: Since we willlet K → ∞ , we ignore the second term in G α
Revised: Since we will let 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 [9] (Sec. 5), we have P → P as n → ∞ , using Monte-Carlo integration (approximationof integral with summation by the law of large numbers).
Revised: Similar to Coifman and Lafon [9] (Sec. 5), we have P → P as n → ∞ , using Monte-Carlo integration (approximation of integral with summation by the law of large numbers).

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

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

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]

Revised: Hence

Original: Hence, if we take ϵ := 2 − K , we have for sufficiently large n and K (implying small ϵ > 0 ),
Revised: [Equation] if we take ϵ := 2 − K [Equation] 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. 
Revised: A.3 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 ( µ, ν )

Original: {+= 
 iff µ =+} ν .
Revised: {+= 
 iff µ =+} ν .

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

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: if X t defined on the embedded space Z [Equation] then W 2 is equivalent to the Wasserstein with geodesic 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: [41] (Thm. 4.1), showed that, for large λ > 0 and ρ t satisfying a) , this minimizationproblem is equivalent to
Revised: [44] (Thm. 4.1), showed that, for large λ > 0 and ρ t satisfying a) , this minimization problem 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 integral as 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 Corr.

Original: 1 for α ∈ (0 , 1 / 2) , we have∥ ϕ
Revised: 1 for α ∈ (0 , 1 / 2) , 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
Revised: Then, for all π ∈ Π( µ, ν ) , we have

Original: [Equation]
Revised: [Equation]

Original: 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.4 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
Revised: C ([0 , 1] , R k ) , with initial and final distribution µ and ν , and N the Wiener measure on C , then the Schrödinger bridge (SB) problem is to find the time-evolving distribution ρ t such that

Original: [Equation]
Revised: [Equation]

Original: subject toρ t ∈ D ( µ, ν )
Revised: to

Original: [Equation]
Revised: [Equation]

Original: where KL is the Kullback–Leibler divergence.
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. minimizing with 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 and final 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 statesunder the Wiener measure N .
Revised: KL ( π || π N ) + KL ( ρ yx || N yx ) , where π N is the joint distribution between the initial and final states under 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ödinger bridge is min KL

Original: [Equation]
Revised: [Equation]

Original:  subject to π ∈ Π( µ, ν ) .
Revised: || π N ) subject to π ∈ Π( µ, ν ) .

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

Original: [Equation]
Revised: [Equation]

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

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

Original: We compare our method with Diffusion Schrödinger Bridge(DSB) [10], which solve the SB problem with an approximation of the iterative proportional fittingalgorithm.
Revised: We compare our method with Diffusion Schrödinger Bridge (DSB) [10], which solve the SB problem with an approximation of the iterative proportional fitting algorithm.

Original: B Experiment details
Revised: B Experiment details

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

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 dimensions to 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 activation functions 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 given sample size without replacement.

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

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, and it requires more careful parametrization (see Sec.

Original: C).
Revised: C).

Original: For all experiments, we used TrajectoryNetwith 1000 iterations and whiten .
Revised: For all experiments, we used TrajectoryNet with 1000 iterations and whiten .

Original: We used the authors’ implementation of DSB with the basic modeland default parameters.
Revised: We used the authors’ implementation of DSB with the basic model and 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 and 32 , 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, with hidden 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 sample size 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 . 15 , − 0 .
Revised: The diffusion parameters are generally small while training, for example when trained on the entire dataset σ t ∈ { 0 . 18 , 0 . 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 maximum time 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, we used λ 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 to hold-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: EB data The EB data are publicly available 3 .
Revised: EB data The EB data are publicly available 4 .

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 geodesic autoencoder 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: The encoder 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 agood scale as with the RBF kernel.
Revised: This kernel is more expensive to 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.

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 CELU activation 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 ofper 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 5 .

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

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 foriterations 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 three hidden layers of size 16 , 32 , 16 , with LeakyReLU in between layers, trained without noise forlocal epochs with λ d = 20 .

Original: C Additional ResultsEnergy
Revised: C Additional Results

Original:  loss In Fig.
Revised: Energy 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 function evaluation 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 the accuracy 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, and thus 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 as training 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 not add 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, one could 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 use the W 1 , the MMD with a Gaussian kernel, and the MMD with identity map ( L 2 norm between the sample 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 the average 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: The parametrization 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 one timepoint 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 ablation on 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, and the 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. 6, we present the 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-out timepoints 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. 3, 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 4, we compare the prediction 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 consider the choice of two kernels when computing the distance G ; the Gaussian kernel for various bandwidth parameters, and the α -decay kernel for different K -nearest neighbors.

Original: Using the GAE improves theinterpolation accuracy for all metrics.
Revised: Using the GAE improves the interpolation 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 with layers 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 by holding out timepoint t for the GAE and the neural ODE.

Original: For the MMD with the Gaussian kernel, weset the scale to 1500 .DSB.
Revised: For the MMD with the Gaussian kernel, we set the scale to 1500 .

Original: Lower is better.

Original: [Table]
Revised: [Table]

Original: Nearest neighbor distance To assist in evaluating how well the generated points land on the knownmanifold (ground truth points), we employ 1-NN distance from the generated points to the groundtruth points.
Revised: Nearest neighbor distance To assist in evaluating how well the generated points land on the known manifold (ground truth points), we employ 1-NN distance from the generated points to the ground truth points.

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

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

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

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 to outliers, 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 [51, 50], and compared its results to MIOFlow (Fig.
Revised: Neural SDE Additionally, we employed a Neural SDE from the PyTorch implementation of differentiable SDE Solvers [56, 55], 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 trained with a fixed diffusion function, but due to extensive training time (approximately two hours forepochs for each dataset) we were unable to perform extensive hyperparameter optimization.

Original: Static methods such as pseudotime assign an abstract measure of time for all cells depending onan initial state [17, 44]. In can be combined to infer the cell lineage, i [Table]
Revised: DSB on petal In Fig.

Original: e. the principal curves ortrajectories [40]. When estimates of cell growth rates are available, Zhang et al. [47] proposed a [Table]
Revised: 14, we show prediction of the DSB algorithm for different ranges of noise scale γ .

Original: method based on optimal transport.

Original: 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 [48]. Saelens et al. [33] provide a rigorous comparisonof 45 static methods. The task of recovering dynamics from multiple snapshots has gained a lot of attention, especially inbiology with scRNA-seq data.
Revised: With too little noise, the prediction appear to collapse on a subset of the data, while the larger scales are inaccurate for this dataset.

Original: 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 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 target canalso be used to recover trajectories.

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

Original: E Software Versions
Revised: D 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==
Revised: In Tab. 7, we present the GPU used during this work, and the following software version were used: cudatoolkit==11.1.1, graphtools==1.5.2, numpy==1.19.5, phate==1.0. pytorch==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.

Original: [Table]
Revised: [Table]

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

Original: Supplement References
Revised: [Figure]