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

Original: Abstract
Revised: Abstract

Original: 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.
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 differential equations (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 that the flow follows the geometry by operating in the latent space of an autoencoder that 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 distance between points is regularized to match a novel multiscale geodesic distance on the data 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 normalizing flows, Schrödinger bridges and other generative models that are designed to flow from 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, we link these trajectories with dynamic optimal transport.
Revised: Theoretically, we link these trajectories with dynamic optimal transport.

Original: 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.
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 data manifold.
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 low dimensional 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 manifold hypothesis, this assumption has led to many successful models of biological, chemical, and physical systems.
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 single cell data the entire transcriptomic profile of the cell is measured with each mRNA or gene species a 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 intrinsic dimensionality 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 good model 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 learning this structure via data graphs [31, 32] and autoencoders [30, 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 [41, 42], diffusion models [21], Schrödinger bridges [33, 10], and continuous normalizing flows (CNF) [7, 17].
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 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.
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) 3 , a new framework for flows based on dynamical 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 differential equations (Neural ODE)
Revised: MIOFlow uses a neural ordinary differential equations (Neural ODE)

Original: [7] to transport a sampling of high dimensional data points between timepoints such that 1.
Revised: [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 penalized to 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 [44], dynamics modeled by a recurrent neural network [20], and Waddington-OT [38].
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 flow where 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, this approach suffers from several drawbacks.
Revised: However, this approach suffers from several drawbacks.

Original: The first drawback is the requirement of starting from a 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 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.
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 normalizing flows are deterministic.
Revised: Second, continuous normalizing flows are deterministic.

Original: Hence, to model the intrinsic stochasticity in biology, we have to force chaotic 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, continuous normalizing 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 flows on 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 manifold distance 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 manifold optimal transport method [34, 50].
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 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.
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.
Revised: • The Geodesic Autoencoder to create a manifold embedding in the latent space of the autoencoder.

Original: • 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.
Revised: • 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 on single-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 time t ∈ 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: For each 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 size n := (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 define a 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 T samples { 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 note the equivalence between two distances d 1 ≃ d 2 .
Revised: We note the equivalence between two distances d 1 ≃ d 2 .

Original: We assume that all Stochastic Differential Equations (SDE) admit a solution [15].
Revised: We assume that all Stochastic Differential Equations (SDE) admit a solution [15].

Original: All proofs are presented in the supplementary material.
Revised: All proofs are presented in the supplementary material.

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

Original: = 0 , where ∇· is the divergence operator.
Revised: = 0 , where ∇· is the divergence operator.

Original: The movement of 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 of a 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: ∼ ν [Equation] where the infimum is over all absolutely continuous stochastic path X t (see [28]).
Revised: 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 in the 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, 29, 16], and mean field games [24].
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 this formulation, 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 [28] 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 can phrase 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 [33, 28].
Revised: The two previous formulations admit the same entropic interpolation ρ t [33, 28].

Original: In this paper, we utilize such a diffusion term, with the knowledge that it still converges to the transport between distributions.
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 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.
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 of k
Revised: Formally, let M be a hidden m dimensional manifold that is only observable via a collection of k

Original: ≫ m nonlinear functions f 1 , . . .
Revised: ≫ m nonlinear functions f 1 , . . .

Original: , f k : M → R that enable its immersion in a high dimensional ambient 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 is collected.
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 =1 ∈ M of the underlying manifold via x i = f
Revised: = { z i } ni =1 ∈ M of the underlying manifold via x i = f

Original: ( z i ) , i = 1 , . . . , n , and aim to learn a low dimensional intrinsic representation 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 density variations, 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 ϵ ( 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.
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 ϵ,β is normalized 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, · )
Revised: k ϵ,β ( x, · )

Original: ∥ 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.
Revised: ∥ 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 eigenvalues and 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, we consider 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: Manifold
Revised: Manifold

Original: Geodesics The dynamic optimal transport formulations (1) and (2) are valid when the ground 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 lower dimensional 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 underlying geometry 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 on geometric domains [40, 45], or when the data lie on a manifold [46].
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 a ground 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 step that 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 equivalent to 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 theoretical results 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 representing the 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 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.
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 is based 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: ) || 1 , 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 [25]).
Revised: Next, we state an important result from Leeb and Coifman (Thm. 2 in [25]).

Original: This theorem links the diffusion 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 ([25] 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, convergence results 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 paper we 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 which preserves 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) from a 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 autoencoder regularized to match this distance.
Revised: We then train an autoencoder regularized to match this distance.

Original: 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.
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 graph with 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 normalized matrix M ϵ :
Revised: k ϵ ( x i , x j ) , and we consider the density normalized matrix M ϵ :

Original: =
Revised: =

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, we choose k ϵ to be a Gaussian kernel or the α -decay kernel [32].
Revised: In practice, we choose k ϵ to be a Gaussian kernel or the α -decay kernel [32].

Original: Lastly, a Markov diffusion operator is 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 stationary distribution 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, we note that the matrix P ϵ follows a similar construction as the kernel p ϵ ; the integrals are approximated by 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 the operator 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 turn converges 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 the transition 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 scales given 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 [46], the authors use the diffusion operator P ϵ to define a distance 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 static formulation 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 approximation of 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 encoder outputting ϕ : 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 (MSE)
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 since we 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, the encoder can be trained to denoise, hence becoming robust to predicted values that are not close to the 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 diffusion matrix.
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 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.
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
Revised: G α ( x i , x j ) ≃ d 2 α M ( x i

Original: , x j ) for all x i , x j ∈ X
Revised: , 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 ) ∥
Revised: ( x i ) − ϕ ( x j ) ∥

Original: 2 ≃ d 2 α M ( x i , x j ) for all x i , x j ∈ X .
Revised: 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 approximation error bounds with high probability, see for example [8, 39].
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 so that 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 , we can 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 lower dimensional space Z , then decode them in the ambient space.
Revised: This is particularly interesting in high dimensions, since we can learn the trajectories in the 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 genes evolved over time.
Revised: Sec. 4 to infer cellular trajectories in the gene space, enabling us to understand how specific genes evolved over time.

Original: We describe the training procedure of the GAE in algorithm 2 in the supplementary material.
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 ∼ µ i for 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 [44], to consider the path space, and other types of dissimilarities between 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 ( µ, ν )
Revised: ( X t , t ) dt with density ρ t , and a dissimilarity between distributions such that D ( µ, ν )

Original: = 

Revised: =

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 with geodesic 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: In practice, it justifies the method of matching the marginals and adding a regularization on the vector field f
Revised: In practice, it justifies the method of matching the marginals and adding a regularization on the vector field f

Original: ( x, t ) , that way the optimal ( ρ t , f ) corresponds to the one from the W 2 .
Revised: ( x, t ) , that way the optimal ( ρ t , f ) corresponds to the one from the W 2 .

Original: If in addition we model the 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: Denote by ψ θ : R k × T → R d |T | the function that represents the Neural ODE, where T is a set of time indices.
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 initial set X 0 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 initial condition, 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 trajectories given 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, representing the 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 integral can 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 implementation from the library Python Optimal Transport [14].
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-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].
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 [44] 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 t given 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 > L d [Equation] λ d

Revised: [Equation]

Revised: (cid:88) x ∈

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: We describe the overall training procedure of MIOFlow in algorithm 1.
Revised: We describe the overall training procedure of MIOFlow in algorithm 1.

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

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

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

Original: It has the advantage of mapping the 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 the training 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: As training continues, we empirically observed that σ → 0 .
Revised: As training continues, we empirically observed that σ → 0 .

Original: Hence, as shown in [28], converging to the 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 to the 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, thus drastically increasing the computational cost.
Revised: To model complex diffusion, one would need an SDE solver, thus drastically increasing the computational cost.

Original: Algorithm
Revised: Algorithm

Original: Input: Datasets X 0 , . . .
Revised: Input: Datasets X 0 , . . .

Original: , 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 ψ θ .
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: Output: Trained neural ODE ψ θ .
Revised: Output: Trained neural ODE ψ θ .

Original: GAE: ϕ, ϕ − 1 ← GAE ( X , K ϵ , K, ξ, N, n max , ϕ, ϕ − 1 )
Revised: GAE: ϕ, ϕ − 1 ← GAE ( X , K ϵ , K, ξ, N, n max , ϕ, ϕ − 1 )

Original: ▷ See algorithm 2 in SM if use GAE then X 0 , . . .
Revised: ▷ See algorithm 2 in SM if use GAE then X 0 , . . .

Original: , X T − 1 ← ϕ
Revised: , X T − 1 ← ϕ

Original: ( X 0 ) , . . .
Revised: ( X 0 ) , . . .

Original: , ϕ ( 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 ) end if for i=1 to n local do ▷ Local training for t=0 to T − 2 do (cid:98)

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

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

Original: the loss L end for for i=1 to n global do ▷ Global training (cid:98) X 1 , . . .
Revised: the loss L end for for i=1 to n global do ▷ Global training (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 } )
Revised: , T − 1 } )

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

Original: the loss L end for
Revised: the loss L end for

Original:  Results
Revised:  Results

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

Original: The MMD is evaluated with 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 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.
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 corresponds to 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: Additional experiments, 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 cellular differentiation, 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) and evolves 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 dataset for 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 keeping trajectories 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 more accurate entropic interpolation for both W 1 and MMD(G), while requiring much less computational time.
Revised: MIOFlow results in a more accurate entropic interpolation for both W 1 and MMD(G), while requiring much less computational time.

Original: In the supplementary material, we present additional results for DSB on the petal dataset.
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 reduction method PHATE [32].
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 differentiation stage.
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 petal dataset.
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 to the 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 less pronounce 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: 3, we show the 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 learns continuous 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 far from 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: [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 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-
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-

Original: [Table]
Revised: [Table]

Original: and 24-27.
Revised: and 24-27.

Original: We reduce the dimensionality of this dataset todimensions using PCA.
Revised: We reduce the dimensionality of this dataset todimensions using PCA.

Original: See the supplementary material for further details on the preprocessing steps and 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, we can decode the trajectories back into the gene space (17846 dimensions), thus describing the evolution of 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 neuronal cells 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 divergence in trends artificially showing higher heterogeneity in gene trends then exist in neuronal specification processes.
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 expression during development from embryonic stem cells to neural early progenitors and then decreasing as the lineages mature [26].
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 more chaotic.
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 diverge towards the neuronal lineage [48].
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.

Original: In Tab.
Revised: In Tab.

Original: we compare the prediction accuracy of MIOFlow whether we use the GAE embedding or no embedding.
Revised: we compare the prediction accuracy of MIOFlow whether we use the GAE embedding or no embedding.

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

Original: For almost all metrics, the prediction accuracy is greater with the GAE.
Revised: For almost all metrics, the prediction accuracy is greater with the GAE.

Original: We present additional results for the choice of kernel in the supplementary material.
Revised: We present additional results for the choice of kernel in the supplementary material.

Original: [Equation]
Revised: [Equation]

Original: MMD(
Revised: MMD(

Original: [Table]
Revised: [Table]

Original: AML Cancer Data with Treatment-based Progression We also assessed whether MIOFlow can give insights on treatment resistant cancer cell populations, rather than differentiation.
Revised: 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 this purpose, 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 the bone 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 treated for 7 days with intensive chemotherapy with single cell measurements taken at day 2, day 5, and day 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 the trajectories, 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 start with 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.

Original:  Related Work
Revised:  Related Work

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

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

Original: In can be combined to infer the cell lineage, i.e. the principal curves or trajectories [43].
Revised: In can be combined to infer the cell lineage, i.e. the principal curves or trajectories [43].

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

Original: [52] proposed a method based on optimal transport.
Revised: [52] proposed a method based on optimal transport.

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

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

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

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

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

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

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

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

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

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

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

Original:  Conclusion
Revised:  Conclusion

Original: We proposed MIOFlow a computationally efficient method to recover trajectories from multiple cross-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 manifold learning 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 show theoretically how MIOFlow is related to dynamical optimal transport on an underlying manifold, and it’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 continuous trajectories 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-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.
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 stiff through 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 the time-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 be alleviated by using unbalanced optimal transport, which allows transporting distributions with unequal masses.
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.

Original: Acknowledgments and Disclosure of Funding
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-
-]

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

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

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

Original: A Theory and Algorithm A.1 Geodesic Autoencoder Algorithm

Original: Using the definitions from Sec. 3, we present the algorithm to train the geodesic autoencoder.

Original: Algorithm

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

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

Original: for i=1 to n max do

Original: Sample batch { x 1 , . . .

Original: , x N } ⊆ X of size N

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: (˜ x i ) − x ∥

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

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

Original: A.2 Proof of Thm.

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

Original: Proof.

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

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

Original: This relies on the convergence results P t/ϵϵ

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

Original: Second, the approximation of the operator

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

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

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

Original: [Equation]

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

Original: =

Original: (cid:80)

Original: [Equation]

Original: ||

Original: [Equation]

Original: k

Original: [Equation]

Original: − δ y )

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

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

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

Original: ≃ D α,H t , where

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

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

Original: H t ∥

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

Original: ∥ P t/ϵϵ − H t

Original: ∥ L 1 ( M ) → 0 . Now let Γ t := P t/ϵϵ − H t , and for γ > 0 choose ϵ > 0 such that ∥ Γ t ( δ [Equation]
Revised: We thank Tom Marty for helping to prototype the GAE.

Original: for x ̸ =

Original: y , we have

Original: [Equation]

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

Original: From the reverse triangle inequality, we have

Original: [Equation]

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

Original: [Equation]

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

Original: [Equation]

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

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

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

Original: Now we define

Original: [Equation]

Original: 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 will let K → ∞ , we ignore the second term in G α

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

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

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

Original: min( K, n ) , and therefore

Original: [Equation]

Original: Hence

Original: [Equation]

Original: if we take ϵ := 2 − K

Original: [Equation]

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

Original: A.3 Proof of Thm.

Original: Theorem 3.

Original: We consider a time-varying vector field f

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

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

Original: [-= 
-]
Revised: [-= 
-]

Original: iff µ = ν .

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

Original: [Equation]

Original: Moreover

Original: [Equation]

Original: if X t defined on the embedded space Z

Original: [Equation]

Original: then W 2 is equivalent to the Wasserstein with geodesic distance W 2 ( µ, ν ) ≃ W d 2 α M ( µ, ν ) .

Original: Proof.

Original: We recall that

Original: [Equation]

Original: is equivalent to

Original: [Equation]

Original: with the three constraints

Original: [Equation]

Original: Tong et al.

Original: [44] (Thm. 4.1), showed that, for large λ > 0 and ρ t satisfying a) , this minimization problem is equivalent to

Original: [Equation]

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

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

Original: [Equation]

Original: Assuming the encoder ϕ

Original: achieves a loss of zero, then from Corr.

Original: 1 for α ∈ (0 , 1 / 2) , we have ∥ ϕ

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

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

Original: ∥ ϕ

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

Original: Then, for all π ∈ Π( µ, ν ) , we have

Original: [Equation]

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

Original: A.4 Schrödinger bridge

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

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

Original: to

Original: [Equation]

Original: where KL is the Kullback–Leibler divergence.

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

Original: 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 states under the Wiener measure N .

Original: By choosing ρ yx = N yx , the static formulation of the Schrödinger bridge is min KL

Original: [Equation]

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

Original: When the Wiener measure has variance σ 2 , the static Schrödinger bridge is equivalent to minimizing

Original: [Equation]

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

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

Original: B Experiment details

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

Original: 01 , and 5 nearest neighbors.

Original: 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 activation functions in between layers.

Original: 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 [57] with default parameters.

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

Original: For all experiments, we used

Original: TrajectoryNet with 1000 iterations and whiten .

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

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

Original: 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, with hidden layers size 16 , 32 , 16 , and initial scales σ t = 0 .

Original: We trained for 40 local epoch with sample size 60 .

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

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

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

Original: 5 , and maximum time scale 2 5 .

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

Original: We used

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

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

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

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

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

Original: 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 CELU activation in between layers.

Original: We trained for 20 local and 10 global epochs with sample size ofper time.

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

Original: AML The AML dataset is publicly available 5 .

Original: 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 foriterations with the RBF kernel (scale 0.5).

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

Original: Energy loss In Fig.

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

Original: Sec.

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

Original: 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 the accuracy of the trajectories is very sensitive to this loss.

Original: 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 as training continues, it collapses in the center (right, ten epochs).

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

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

Original: 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 the average for the scales 0 . 1 and 0 .

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

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

Original: In Fig.

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

Original: In Fig.

Original: 10, we reproduce the same ablation on Dyngen data.

Original: 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. 6, we present the average of the KNN distance computed on the top 5% .

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

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

Original: TrajectoryNet, and DSB.

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

Original: In Tab 4, we compare the prediction accuracy on the EB data, depending on using the GAE or not.

Original: 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 the interpolation accuracy for all metrics.

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

Original: 2 , linear layers of size 16 , 32 , 16 , 50 with CELU activation in between layers.

Original: We train by holding out timepoint t for the GAE and the neural ODE.

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

Original: [Table]

Original: 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 5 shows the 1-NN distance for several methods and datasets.

Original: Of note is the aggregate column.

Original: Where aggregate is specified as mean , we compute the mean 1-NN distance from the predicted point to the 1-NN at the same time index. 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 . 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. 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.
Revised: This work was partially funded and supported by CIFAR AI Chair [G.W.], NSERC Discovery grant 03267 [G.W.], NIH grants 1F30AI157270- [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- [S.K.], and the Sloan Fellowship FG-2021-15883 [S.K.].

Original: DSB on petal In Fig.

Original: 14, we show prediction of the DSB algorithm for different ranges of noise scale γ . With too little noise, the prediction appear to collapse on a subset of the data, while the larger scales are inaccurate for this dataset. D Software Versions In Tab. 7, we present the GPU used during this work, and  the following software version were used:
Revised: The content provided here is solely the responsibility of the authors and does not necessarily represent the official views of the funding agencies.

Original: 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. [Table] [Table]
Revised: The funders had no role in study design, data collection & analysis, decision to publish, or preparation of the manuscript.

Original: [Table]

Original: [Figure]