Manifold Interpolating Optimal-Transport Flows for Trajectory Inference
Authors: Guillaume Huguet
Daniel Sumner Magruder
Alexander Tong
Oluwadamilola Fasina
Manik Kuchroo
Guy Wolf
Smita Krishnaswamy
Venue: NeurIPS 2022 Accept (2022)

Abstract: 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. 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. 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). In GAE the latent space distance between points is regularized to match a novel multiscale geodesic distance on the data manifold that we define. We show that this method is superior to normalizing flows, Schr\"odinger bridges and other generative models that are designed to flow from noise to data in terms of interpolating between populations. Theoretically, we link these trajectories with dynamic optimal transport. 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. 

Manifold Interpolating Optimal-Transport Flows for Trajectory Inference
Guillaume Huguet
D.S. Magruder
Alexander Tong!*
Oluwadamilola Fasina²
Manik Kuchroo²
Guy Wolf
Smita Krishnaswamy
Université de Montréal; Mila Quebec Al Institute
Yale University
Abstract
We present method called Manifold Interpolating (MIOFlow) that learns stochastic, continuous population dynamics from static snapshot samples taken sporadic timepoints. MIOFlow combines dynamic mod- els, manifold learning, and optimal transport by training neural ordinary differential equations (Neural ODE) interpolate between static population snapshots as pe- nalized by optimal transport with manifold ground distance. Further, we ensure that the flow follows the geometry operating in the latent space of autoencoder that we call geodesic autoencoder (GAE). In GAE the latent space distance between points is regularized to match novel multiscale geodesic distance on the data manifold that we define. 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. Theoretically, we link these trajectories with dynamic optimal transport. We evaluate our method on simulated data with bifurcations and merges, well as scRNA-seq data from embryoid body differentiation, and acute myeloid leukemia treatment.
Flow
Introduction
Here, tackle the problem of continuous dynamics of probability distributions defined on data manifold. Data from naturalistic systems are often modeled as generated from an underlying low dimensional manifold embedded in high dimensional measurement space. Termed the manifold hypothesis, this assumption has led to many successful models of biological, chemical, and physical systems. Mcasurements in such systems are increasingly high dimensional. For instance, single cell data the entire transcriptomic profile of the cell is measured with each mRNA or gene species a dimension. However, because of informational redundancy between these genes, the intrinsic dimensionality of the data is low dimensional. Mathematically. Riemannian manifold is good model for such a system, and much of manifold learning literature including, diffusion maps (9). Laplacian eigenmaps [2), and the wider graph signal processing literature, has focused on learning this structure via data graphs 31 and autoencoders [30]
Building on this literature, we consider snapshot measurements of cellular populations over time. Current technology cannot follow cell in time as the measurements are destructive. We frame this as learning population flow on data manifold. Recently several neural networks that perform flows, or transports between probability distributions, have been proposed in the literature. However, most of these works have focused on generative modeling: i.e., flowing from noise distribution such as Gaussian distribution to a data distribution in order to generate data. Examples include score-based
Equaal contribution
Corresponding authors: guy . wolf folfmutunrr....... . ca and smita. www.a...rissnavaaar.. www..isnaaaayyyall..ccomm
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

generative matching A1 42) diffusion models 22], Schrödinger bridges B3 10 and continuous normalizing flows (CNF)
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.
To continuously interpolate populations over time on data with manifold structure, we propose Manifold Interpolation Optimal- -ranspoot Flow (MIOFlow a new framework for flows based on dynamical optimal transport in manifold embedding. MIOFlow uses a neural ordinary differential equations (Neural ODE) 7 to transport a sampling of high dimensional data points between time- points such that 1. the transport occurs on manifold defined by samples, 2. the transport is penalized to agree with measured timepoints using Wasserstein, the transport inherently stochastic.
Works tackling a similar problem include Trajectory 44 dynamics modeled by a recurrent neural network [20), and Waddington-OT [38] TrajectoryNet is based on continuous normalizing flow where the magnitude of the derivative is penalized to create a more optimal flow. However, this approach suffers from several drawbacks. The first drawback is the requirement of starting from a Gaussian distribution. 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. Second, continuous normalizing flows are deterministic. Hence, to model the intrinsic stochasticity in biology, we have to force chaotic dynamics starting from slight noise added to any initial distributions. Third, continuous normalizing flow models in & dimensions require calculating the trace of the Jacobian, which requires O(k²) operations to compute 7 making our method i times faster per function evaluation.
Additionally, unlike normalizing flows which operate in ambient data dimensions, focus on flows on a data manifold. We feature a two-pronged approach with which to enforce this. First, propose the Geodesic Autoencoder to embed the data such that distances in the latent space match new manifold distance called diffusion geodesic distance. Second, we penalize the transport using a manifold
optimal transport method [34 50]. 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 (2 access to the metric tensor of a Riemannian manifold required.
The main contributions of our work include:
The MIOFlow framework for efficiently learning continuous stochastic dynamics of static snap- shot data based on Neural ODEs that implements optimal transport flows on a data manifold.
The Geodesic Autoencoder to create manifold embedding in the latent space of the autoencoder.
A new multiscale manifold distance called diffusion geodesic distance, and theoretical results showing its convergence to geodesic on the manifold in the limit of infinitely many points. Empirical results showing that our flows can model divergent trajectories on toy data and on single-cell developmental and treatment response data.
2 Preliminaries and Background
Problem Formulation and Notation We consider a distribution μ4, over TR evolving over time t E R. from which we only observe samples from finite set of T distributions {ur |T-1 For each time t, we observe sample X, μ4 of size n4. We note the set of all observations X of size := E. ny. We also want to characterize the evolution of the support of 14, We aim to define trajectory from an initial points Xo No to XT-1 MT-1, given the intermediate conditions X2 ~ u2,..... .., XT-1 ~ pT--1. We are thus interested in matching distributions given a set of T samples (x,) and initial condition Xo-
In the following, we assume that the distributions are absolutely continuous with respect to the Lebesgue measure, and we use the same notation for the distribution and its density. We note the equivalence between two distances dj ~ d2. We assume that all Stochastic Differential Equations (SDE) admit solution [5 All proofs are presented in the supplementary material.
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, We consider two distributions and defined on X' and y, and 11(,, v) the set of joint distributions on X' x y with
Code is available here: httpp////gthuu...c/////sssaaaass

---

marginals je and V, i.e. ((zz, )) = p(dx) and T(X', dy) = v(dy) for т E [(,, v). The transport plan и moves the mass from p to v. where the cost of moving a unit mass from the initial X to the final y is d(ar, y). This formulation gives rise to the p- Wassersteii distance
J d(z, y) (dx,dy).
Wp(, := inf nEll(H.)) JxYY
where p E oo). Suppose X' == = R², and d(x, y) == || yl12 then Benamou and Brenier 31 provide dynamic formulation of the optimal transport problem. For simplicity, assume fix time interval [0, 1]. but note that the following holds for any time interval [to, f1]. The transport plan is replaced by a time-evolving distribution Pt. such that Pa = μ and P1 / and that satisfies the continuity equation pep: + V (p.v) = O, where V is the divergence operator. The movement of mass is described by time-evolving vector field v(r, t). When Pt satisfies these conditions, Benamou and Brenier 3 show that
W2(u, inf [ [/eL [lu(x, (1l2p((d)ddl
(pr.0)
This formulation arises from the field of fluid mechanics; the optimal vector field is the divergence of a pressureless potential flow. We can also view the problem on the path space of R
W2(n, u)² = inf [[, [IS(X,, (1Baa s.t. dx, = f(Xe, f)dl, No ~ 14, X1 ~M
(2)
Xt
where the infimum is over all absolutely continuous stochastic path X, (see (28)
Adding Diffusion For various applications, it is interesting to incorporate a diffusion term in the paths or trajectories. These types of flows are often used in control dynamics [1 29 16], and mean field games [24]. We consider an SDE d.X, fXx,, t) + VodB, where B, is standard Brownian motion. To model this new dynamic, one can replace the continuity equation with the
Fokker-Planck equation of this SDE with diffusion, i.e. 0,p: + V (p.v) Aoop. From this formulation, one can also retrieve the Anmamou-Brenied optimal transport. Indeed, Mikami [28
established the convergence to the W2 when the diffusion term goes to zero. Similar to ( we can phrase the problem using an SDE
inf EE [. [X, 112dt s.t. dx,; = f(x,, t)dt VodB,, Xo M4 X1 ~ M. y
(3)
The two previous formulations admit the same entropic interpolation Pt 33 28 In this paper, we utilize such diffusion term, with the knowledge that it still converges to the transport between distributions. Moreover, as the diffusion goes to zero, the infimum converges the W2.
Manifold Learning A useful assumption representation learning is that high dimensional data originates from intrinsic low dimensional manifold that is mapped via nonlinear functions to observable high dimensional measurements; this commonly referred to as the manifold assumption. Formally, let M be a hidden m dimensional manifold that is only observable via a collection of k > m nonlinear functions f1 fi....,Je Su M - R that enable its immersion in a high dimensional ambient space F(M) (f(=) collected. Conversely, given dataset X (II, ]n) C R of high dimensional observations, manifold learning methods assume data points originate from a sampling 2 = } E M of the underlying manifold via x; f(4), =1 1..... n, and aim to learn low dimensional intrinsic representation that approximates the manifold geometry of M.
(f)( ((()))...,((()))
E M) R from which data is
To learn manifold geometry from collected data that is robust to sampling density variations, Coifman and Lafon proposed to use an anisotropic kernel ke,9 (x. y)
(x, w)/lkk Ilke III where 0 < 3 < 1 controls the separation of geometry from density, with 8 yielding the isotropic kernel. and B
completely removing density and providing a geometric equivalent to uniform sampling of the underlying manifold. Next, the kernel ke,e is normalized define transition probabilities Pc. y) == (x, y) [/ -)11 and an n n row stochastic matrix (P (P.,a);j hjj == Pc,,(x,,,0) that describes a Markovian diffusion process over the intrinsic geometry of the data. Finally, diffusion map 9 (xx)) ) is defined by the eigenvalues and eigenvectors of the matrix P c..- Most notably, this embedding preserves the diffusion distance between 2,, X, where t is time-scale parameter. Next,
II'',,(x,,)) :) ( 3//1 (-)]l2 consider multiscale diffusion distance that relates to the geodesic distance on the manifold.
3

---

Neural Network
0100000
x((.)
5386
QDE Solver
Loss
0101010101110
AutdErooder
Figure I: Overview of the MIOFlow pipeline. The Geodesic Autoencoder learns a latent space that preserves the diffusion geodesic distance. neural network predicts the derivative trajectories with respect to time. For an initial sample at lo. the ODE Solver produces the trajectories xx41) ) At population level, predictions of points at time t are penalized by the Wasserstein distance during training.
Manifold Geodesics The dynamic optimal transport formulations ① and (2) valid when the ground distance is Euclidean. This is restrictive as we might have observations sampled from a lower dimensional manifold. Thus, using the Euclidean distance in the ambient space ignores the underlying geometry of the data. Static optimal transport with a geodesic distance is known to perform well on geometric domains 440 451, when the data lie manifold [66.. Here, extend the of ground geodesic distance to the dynamic formulation. To do so, we define a manifold learning step that precedes the trajectory inference. The goal is to learn embedding Z such that the Euclidean in Z is equivalent to the geodesic distance. In that case, dynamic optimal transport in Z is equivalent the Wasserstein with a geodesic ground distance. In the following section, we present theoretical results that will justify our approximation of a geodesic distance on a closed Riemannian manifold.
We consider closed Riemannian manifold (M, d.). where d.M is the geodesic distance representing the shortest path between two points on the manifold. We note h2 the heat kernel on M. For a point M, the heat kernel (x, induces measure, i.c. how the heat has propagated on the manifold at time given an arbitrary initial distribution. The diffusion ground distance between æ, E . is based on the L norm between the measures induced by the heat kernel given the initial conditions 5, and by-
Definition 1. The diffusion ground distance between 2, y M is
Da (x ka (1) - h2-* (11
R20
for E (0, 1/2), the scale parameter A O, and he the heat kernel M.
Next, we state an important result from Leeb and Coifman (Thm. in (25)) This theorem links the diffusion ground distance and the geodesic on a Riemannian manifold.
Theorem 1 (25) Thm. 2). Let (M, a closed Riemannian manifold, with geodesic d for (0, 1/2), the distance D is equivalent to dag.
In practice, we cannot always evaluate the heat kernel for arbitrary manifold. However, convergence results from Coifman and Lafon [9] provide a natural approximation. Indeed, for B = 1, the diffusion operator pt/s converges to the heat operator
goes to zero (Prop. 3 ( For the rest of the paper we assume 3 = 1, we define P. = Pe.1 and the matrix P. = P.,1.
Manifold Interpolating Optimal-Transpon Flow
To learn individual trajectories from multiple cross-sectional samples of population. we propose MIOFlow shown in Fig. [ Our method consists two steps. first learn an embedding which preserves the diffusion geodesic distance-whit
we define using our Geodesic Autoencoder. Then, we learn continuous trajectories based on an optimal transport loss. We model the trajectories with Neural ODE 77). allowing to learn non-linear paths and interpolate between timepoints.

---

3.1 Geodesic Autoencoder Embedding
In light of Thm. I we define an approximation of the diffusion ground distance Da (Def. [T) from a collection of samples X that we call diffusson geodesic distance. We then train an autoencoder regularized match this distance. Similar ideas are explored [T] where the latent space of the autoencoder is regularized to match manifold learning embedding, either PHATE or diffusion map.
Our approximation relies the diffusion matrix Pc. with the affinity matrix defined by (K.) == (z,, and we consider the density normalized matrix M.- == Q Q where Q is a diagonal matrix with Q == Σ, (Ke).. In practice, we choose ke to be a Gaussian kernel or the a-decay kernel [32]. Lastly, a Markov diffusion operator is defined by P.
We first recall its construction. We build graph
D 'M., where D ΣØ= (M.) by is a diagonal matrix. The stationary distribution associated to P. π, where æ, D&/ E, Djj. since P. is r-reversible. Finally, we note that the matrix P. follows similar construction as the kernel p.: the integrals approximated by sums. We refer to [9] for more details about the convergence of the matrix operator P. to the operator P.. Until now, have defined matrix P. that approximates the operator P., which in turn converges to the heat operator. Now, we define an approximation of the diffusion ground distance, based on the matrix P.. We use the notation (P.). : represent the i-th row of P it represents the transition probabilities of t-steps random walk started at zi..
Definition 2. We define the diffusion geodesic distance between T4, Iy E X as
Ga(Di,uz)= = 2 (K 1(P (K--)))0[(PP))) (P.) 11 2 (K+ -Tyl1-
The diffusion geodesic compares the transitions probabilities of a random walk at various scales given two different initial states z;, Tzj. In [46), the authors use the diffusion operator P. to define distance equivalent to the Wasserstein with ground distance D.. Their method comes from the static formulation of optimal transport. Here we propose to learn space Z that preserves an approximation of the distance D in order to do dynamic optimal transport in 2.
To use the diffusion geodesic distance, we train a Geodesic Autoencoder, with encoder outputting Ø IR* + Z, such that [lØ(x;) - p(x,)); x Go GT),, We draw subsample of size N, evaluate G. and minimize the Mean Square Error (MSE)
Training
N
2 ( = N
Σ
Σ S>
(I(b(,)) - (11))12 - Ga(ss, x;))²
Learning the embedding 2 over using Ga as the advantage of being inductive, is useful since to compute distances between predicted and ground truth observations. Moreover, the encoder can be trained to denoise, hence becoming robust to predicted values that are not close to the manifold. Computing G., on the entire dataset is inefficient due to the powering of the diffusion matrix. We circumvent this difficulty with the encoder, since we train on subsamples. We choose N to have few observations in most regions of the manifold, thus making the computation of Ga very efficient, which allows us to consider more scales.
Theorem 2. Assuming X is sampled from a closed Riemannian manifold M with geodesic dy- Then, for all E (0, 1/2), sufficiently large K, N and small > 0. we have with high probability G. (Zi, xz) ~ d3a (I;, for all Z;, J E X.
Corollary 1. If the encoder such that L(Ø) 0, then with high probability [|ő((.)) ((x,))l2= die mry) for all Fú, xy X.
The stochasticity in the two previous results arises from the discrete approximation of the operator Pa- The law or large numbers guarantees the convergence. For fix sample size, there exist approximation error bounds with high probability, see for example g 33. In practice, we choose close to 1/2 so that the diffusion distance is equivalent to the geodesic on the manifold. From that embedding 2, we can also train decoder Φ-1 Z + R* with reconstruction loss L,- := Σ | - Φ(x) x1|2- 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. For example, we use the decoder in
5

---

Sec. to infer cellular trajectories in the gene space, enabling us to understand how specific genes evolved over time. We describe the training procedure of the GAE in algorithm 2 in the supplementary material.
3.2 Inferring Trajectories
Given T distributions {0 we want model the trajectories respecting the conditions X; ~ f4 for fix timepoints i E { T 1}. Formally, we want to learn parametrized function fo( t) such that
X, = Xo + / fo(Xur u)du, with Xo ~ ,o,..., XT-1 ~ PT-1-
(4)
We adapt the main theorem from 44 to consider the path space, and other types of dissimilarities between distributions.
Theorem 3.
We consider a time-varying vector field f(x,t) defining the trajectories dX, f(X,, with density pt, and dissimilarity between distributions such that D(,, v) iff v. Given these assumptions, there exist sufficiently large X > O such that
[ ((xx,,))133 + AD((I. S.I. No ~ n..
inf =i3 X.
(5)
Moreover, if X, is defined on the embedded space 2, then W2 is equivalent to the Wasserstein with geodesic distance W2(u, v) Wazz (H, v).
This theorem enables us to add the second marginal constraint in the optimization problem. In practice, it justifies the method of matching the marginals and adding regularization on the vector field [ax 4), that way the optimal (pt. J) corresponds to the one from the W2. If in addition we model the trajectories in the embed space Z, then the transport is equivalent the
the manifold.
Training We observe discrete distributions μ4 == (1/n:) dn, for x; X,, and we approxi- mate (4) with a Neural ODE 77]. where so is modeled by a neural network with parameters 0. Denote by tho R x T + 47 the function that represents the Neural ODE, where T is a set of time indices. We define the predicted sets X..... X = the (Xo, {1...., T 1}); given an initial The resulting discrete distributions are i = (L/n:) Σ5 62, for xj X, To match the marginals, we used two approaches of training; local or global. For the local method, we only predict the next sample, hence given X, as initial condition, we predict X++1 = the (X,, t+ 1). Whereas for the global, we use the entire trajectories given the initial condition Xg- For both cases, we formulate loss I'm the marginals, representing the second term in (5). To take into account the first term in (5). we add the loss Le. where the integral can be approximate with the forward pass of the ODE solver, and A Z 0 is a hyperparameter.
set Xo returns the approximation of (4) for all
T-1 Lm = Σ Wa(apa)
T-1 E Le de Σ
(6)
In practice, to compute the Wasserstein between discrete distributions, we use the implementation from the library Python Optimal Transport [14). This is in contrasts with the maximum-lkeelibodd approach in CNF methods which requires evaluation of the instantancous change of variables every
integration timestep at O(k2) additional cost per function evaluation (7].
We add final density loss La inspired by 44 encourage the trajectories to stay the manifold. Intuitively, for a predicted point x Xx. we minimize the distance to its k-nearest neighbors in X given a lower bound h > 0
T-1
k
L= Aa
Σ
Σ
ls(z,t), t), where lx(x,t) = max(0, min-k((lle= yll : X,)) - h).
1=1
We describe the overall training procedure of MIOFlow in algorithm

---

Modeling Diffusion We add a diffusion term in the trajectories, resulting in the SDE dXx;, fX,, t) + VoidB, In practice, we learn I parameters {or) TT-1 It has the advantage of mapping the same cell to different trajectories stochastically as in many naturalistic systems. Early in the training process, we notice that it promotes bifurcation, as it helps to spread the initial Xp. As training continues, we empirically observed that a - O. Hence, as shown in [28], converging to the W2 formulation. For small we train with a regular ODE solver, essentially corresponding to Maruyama method. To model complex diffusion, one would need an SDE solver, thus drastically increasing the computational cost.
Algorithm 1 MIOFlow
Xy- 1, graph kernel K,, maximum scale K, noise scale batch size N, maximum iteration Remoz, initialized encoder & and decoder Φ maximmm local iteration Moccll maximum global iteration Nggobbat initialized neural ODE 1h0-
Input: Datasets Xp,-..,Xyy
Output: Trained neural ODE 4/g.
GAE: o, 20_ GAE(X, K., K, K, N,,,N,,ma Rmaz 0,0 21) if use GAE then , t end if
See algorithm 2
in SM
((XT
for i=1 to nocco do
Local training
for t=0 to T 2 do
X(+1 + tha (X2, t
end for
L - Imm + The + Ltd
update the with gradient descent w.r.f. the loss end for
for i=1 to Aglobal do +/o (1...., - 1})
P Global training
L Im Tee + La
update the with gradient descent w.r.t. the loss L
end for
4 Results
Here, we validate our method empirically on both synthetic test cases and single-cell RNA sequencing datasets chosen to test the capability of MIOFlow follow divergent trajectories from population snapshot data. Using these test cases we assess whether MIOFlow-des follow the data manifold efficiently, we also quantitatively assess their accuracy using the Wasserstein W1, and Maximum Mean Discrepancy (MMD) [18] on held-out timepoints. The MMD is evaluated with the Gaussian (G) kernel and the mean (M). We compare our results to those other methods that preform population flows including TrajectoryNet [4]] which is based on CNF that is regularized to achieve efficient paths, and Diffusion Schrödinger's 's Bridge (DSB) [00 which is an optimal transport framework for generative modeling. The baseline measure in the quantitative results corresponds to the average distance between the previous and next timepoints for a given time I. Additional experiments, such as ablation studies, and details are provided in the supplementary material.
trajectories qualitatively
TTggectryy
MEDE toe Joun!
Figure 2: Comparisons between predicted flows of (a) TrajectoryNet, (b) DSB, and (c) MIOFlow on a toy Petal dataset. MIOFlow trajectories both match the data and transition along the data manifold whereas TrajectoryNet takes shortcut paths between the the curved paths and Schroedinger's bridge does not follow the data. Time indicated by color. Ground truth points are designated by translucent "X"s in the background, whereas predicted points are opaque circles, whilst trajectories are solid lines.
7

---

4.1
Artificial Data
Petal The petal dataset is simple yet challenging, as it mimics natural dynamics arising in cellular differentiation, including bifurcations and merges. The dynamics start in the center (purple) and evolves through the end of the leaf (yellow). In Fig. 2 we present the trajectories on the petal dataset for TrajectoryNet, DSB, and MIOFlow. Only our method is able to learn the dynamic, and keeping trajectories on the manifold. We then hold-out timepoint 2 corresponding to the initial bifurcation, and, in Tab. [ we compare the interpolation between the three methods. MIOFlow results in more accurate entropic interpolation for both W and MMD(G), while requiring much less computational time. In the supplementary material, we present additional results for DSB on the petal dataset.
Dyngen We use Dyngen 6 to simulate scRNA -sqq dataset from dynamical cellular process. We first embed the observations in five dimensions using the non-linear dimensionality reduction method PHATE 32]. We discretize the space in five bins, each representing a differentiation stage. This dataset includes one bifurcation. However, it proves more challenging than the petal dataset. The number of samples per bins not equal, which tends to give more influence to the lower branch. 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. In Fig. 3 we show the trajectories for TrajectoryNet. DSB, and our method. MIOFlow is more accurate it learns continuous bifurcating trajectories along the two main branches. In Tab. I we hold-out timepoint 2, and show the interpolation results against the ground truth. MIOFlow is more accurate for all metrics, while requiring much less computational time. We see that in timepoint I, TrajectoryNet deviates far from the trajectories particularly the branching point where it struggles with induced heterogeneity (see Fig. 3p. We have further quantification reflecting this in the supplementary material.
voe
Figure 3: Comparisons between predicted flows of (a) TrajectoryNet, (b) DSB, and (c) MIOFlow on Dyngen 66 simulated dataset. MIOFlow trajectories both match the data and transition along the data manifold throughout the timepoints. Time indicated by color. Ground truth points are designated by translucent "X"s in the background, whereas predicted points are opaque circles, whilst trajectories are solid lines.
Table 1: Leave-one-out Wi and MMD with (G)aussian (M)an kernel between the predicted and ground truth points as well as training time in (m)inutes and (s)econds on the Petal and Dyngen datasets of TrajectoryNet, DSB, and MIOFlow. Lower is better.
Petal Wi
Dyngen Wi
MMD(G)
MMD(M)
Runtime
MMD(G)
MMD(M)
Runtime
MIOFlow (ours)
0.090
0.029
0.005
284.60 (s)
0.783
0.199
0.509
95.63 (s)
TrajectoryNet
0.181
0.136
0.002
64 (m)
1.499
0.303
1.806
62 (m)
DSB
0.199
0.157
0.008
86 (m)
2.051
0.932
2.367
81 (m)
Baseline
0.221
0.196
<0.001
N/A
1.388
1.008
0.926
N/A
4.2 Single-Cell Data
Embryoid Body Here we evaluate our method on scRNA-seq data from the dynamic time course
of human embryoid body (EB) differrnntiatinn over period of 27 days, measured in 3 timepoints between days 0-3, 6-9, 12-15, 18-21, and 24-27. We reduce the dimensionality of this dataset to 200 dimensions using PCA. See the supplementary material for further details the preprocessing steps and the parameters.
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 cell. In Fig. 5 we plot the trajectories, at a gene level of 20 neuronal

---

cells against 20 random one. We generally observe distinct differences between TrajectoryNet and MIOFlow. TrajectoryNet (which is shown the bottom) shows tendency have early divergence in trends artificially showing higher heterogeneity in gene trends then exist in neuronal specification processes. Examples include HAND22, which has 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). MIOFlow clearly shows this trend whereas TrajectoryNet trends are more chaotic. Another example is ONECUT2 which is consistently known to be higher in cells that diverge towards the neuronal lineage (48] Other examples are discussed the supplementary material.
gone
Figure 4: MIOFlow on the AML dataset. a) PCA Embedding comparing ground truth to MIOFlow predicted points and trajectories. b) Cellular gene expression trajectories for SLPI, B2M, and MYB for MIOFlow.
Nousinal
TranctoryMet
Figure Neuronal gene trajectories vs. random trajectories for MIOFlow (top) vs. TrajectoryNet [46 (bottom). TrajectoryNet consistently shows tendency for early divergence in trajectories, examples include ONECUT2 which is biologically known to be higher expressed 33 and HAND2 which is known to peak in neural early progenitors (NEPs).
In Tab. we compare the prediction aceuracy of MIOFlow whether we use the GAE embedding or no embedding. For the GAE, we show results for two kernels; the Gaussian kernel (scale 0.5), and the a-decay kernel 32 (based on a 30 nearest neighbors graph). For almost all metrics, the prediction accuracy is greater with the GAF. We present additional results for the choice of kernel in the supplementary material.
Table Leave-one-out (t) W1 and MMD with (G)aussian and (M)ean kernel between the predicted and ground truth point on the EB datasets for MIOFlow with the Gaussian, a-decay kernel, or without GAE. Lower is better.
t=
t == 3
W1
MMD(G)
MMD(M)
W1
MMD(G)
MMD(M)
GAE Gaussian
26.318
0.059
101.969
32.963
0.138
223.679
GAE a-decay
25.744
0.061
99.747
32.227
0.135
213.465
No GAE 29.243
0.063
126.608
35.709
0.119
246.609
Baseline 33.415
0.103
227.279
35.319
0.095
213.492
AML Cancer Data with Treatment-baat Progression We also assessed whether MIOFlow can give insights on treatment resistant cancer cell populations, rather than differentiation. For this purpose, we used data from Fennell al. [22 in which AF9 mutant cancer cells implanted into the bone marrow of mice as a model of Acute Myeloid Leukemia (AML). The mice are then treated for 7 days with intensive chemotherapy with single cell measurements taken day 2, day and day 7. 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. We see heterogeneity trajectories, the majority of cells surviving to the last timepoint showing an increase in the markers SLPI, MYB, and B2M all members of the LSC pathway [2]. This indicates that not all cells start with this signature but made transition to this signature to resist the effects of chemotherapy.
the

---

5
Related Work
The task of recovering dynamics from multiple snapshots has gained a lot of attention, especially in biology with scRNA-seq data, Static methods such as pseudotime assign an abstract measure of time for all cells depending on an initial state
[19 47 can be combined to infer the cell lineage, i.e. the principal curves or trajectories [43] When estimates of cell growth rates are available, Zhang
et al. [52) proposed a method based on optimal transport. Another line of research is on the direct rate of differentiation of cell or RNA velocity [23] 4.. This method does not provide trajectories, but
it can be combined to define a transition matrix to infer trajectories [53]. Saelens et al. [37
provide a
rigorous comparison of 45 static methods.
44 infer trajectories. The dynamics are either modeled with a recurrent neural network 200, a generative
Similar methods used the dynamic formulation of optimal transport [3 20
35
36
38
adversarial networks 355, flows defined by the gradient of an input convex neural network [5], or a CNF [44. Here, we circumvent the need for CNF by modeling the trajectories directly on the support of the distributions, hence maintaining the computational cost low even in high dimension. Some generative methods that interpolate between a prior to a target can also be used
to recover trajectories. Other methods use the Schrödinger bridge (SB) [10 22 49 51 problem to
define trajectories between distributions. For example 10), solve the SB between two distributions using a series of Sinkhorn iterations. However, it does not explicitly learn dynamics, making their interpolations inaccurate.
6 Conclusion
We proposed MIOFlow computationally efficient method to recover trajectories from multiple cross-sectional samples of time-evolving population. Our method includes novel manifold learning framework, where we learn an embedding preserving the geodesic distance. We show theoretically how MIOFlow is related to dynamical optimal transport on an underlying manifold. and it's 's capacity to model stochastic trajectories. On artificial datasets, MIOFlow produced continuous trajectories more in line with the manifold compared to DSB and TrajectoryNet. 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. This way we could visualize the evolution of the gene expression of a cell.
Limitations and Broader Impact
Our method does have some limitations. 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. Like other methods that learn dynamics is possible that the learned dynamics may become stiff through training leading to decreased model accuracy. In addition, our work is sensitive to the time-bins, and worked best for relatively uniform sampling. This second issue could potentially be alleviated by using unbalanced optimal transport, which allows transporting distributions with unequal masses. 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.
Acknowledgments and Disclosure of Funding
We thank Tom Marty for helping to prototype the GAE. This work was partially funded and supported
by CIFAR AI Chair [G. W.J, NSERC Discovery grant 03267 [.. NIH grants 1F30AI157270-M
[M.K.]. 201GM135929 [G.W., S.K.J, R01HD100035, and R01GM130841 [G. W., S.K.J. NSF Career grant 2047856 [S.K.], the Chan-Zuckerberr Initiative grants CZF2019- 18009--188 182702 and CZF2019-002440 [S.K.]. and the Sloan Fellowship (G-2021-15888 [S.K.]. The content provided here solely the responsibility of the authors and does not necessarily represent the official views of the funding agencies. The funders had no role in study design, data collection & analysis, decision to publish, or preparation of the manuscript.
10

---

References
[1] Andrei Agrachev and Paul WY Lee. Continuity of optimal control costs and its application weak kam theory. Calculus of Variations and Partial Differential Equations, 39(1):1132322 2010.
[2] Mikhail Belkin and Partha Niyogi. Semi-Superviseer Learning on Riemannian Manifolds. Machine Learning, 56(1-3)2009 239, 2004.
[3] Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
[4] Volker Bergen, Ruslan A. Soldatov, Peter V. Kharchenko, and Fabian J. Theis, RNA velocity- current challenges and future perspectives. Molecular Systems Biology, 1788)x10282, 2021.
[5] Charlotte Bunne, Laetitia Lengg-Papaxanthoss Andreas Krause, and Marco Cuturi. Proximal Optimal Transport Modeling of Population Dynamics. AISTATS, 2022.
[6] Robrecht Cannoodt, Wouter Saelens, Louise Deconinck, and Yvan Saeys. Spearheading future omics analyses using dyngen, a multi-modal simulator of single cells. Nature Communications, 12(1):1-9, 2021.
[7] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018.
[8] Xiuyuan Cheng and Nan Wu. Eigen-convergence of gaussian kernelized graph laplacian by manifold heat interpolation. Applied and Computational Harmonic Analysis, 61:132-190, 2022.
[9] Ronald R. Coifman and Stéphane Lafon. Diffusion maps. Applied and Computational Harmonic Analysis, 21(1):5-30, 2006.
[10] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances Neural Information Processing Systems, 34, 2021.
[11] Andrés F Duque, Sacha Morin, Guy Wolf, and Kevin Moon. Extendable and invertible manifold learning with geometry regularized autocncoders. In 2020 IEEE International Conference on Big Data (Big Data), pages 5027-5036. IEEE, 2020.
[12] Katie A Fennell, Dane Vassiliadis, Enid YN Lam, Luciano G Martelotto, Jesse J Balic, Sebastian Hollizeck, Tom Weber, Timothy Semple, Qing Wang. Denise C Miles, al. Non-genetic determinants of malignant clonal fitness at single-cell resolution. Nature, 601((891): 2022.
[13] Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. In International conference on machine learning, pages 3154-3164. PMLR, 2020.
[14] Rémi Flamary, Nicolas Courty. Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
[15] Crispin W Gardiner et al. Handbook of stochastic methods, volume 3. springer Berlin, 1985.
[16] Nassif Ghoussouh, Young-Heon Kim, and Aaron Zeff Palmer. A solution to the monge transport problem for brownian martingales. the Annals of Probobility, 49(2)=877-907, 2021.
[17] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: Free- form Continuous Dynamies for Scalable Reversible Generative Models. In 7th International Conference on Learning Representations, 2019.
[18] Arthur Gretton, Karsten Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. kernel method for the two-sample problem. Journal of Machine Learning Research, 1:1-10, 2008.
[19] Laleh Haghverdi, Maren Büttner, F. Alexander Wolf. Florian Buettner, and Fabian J. Theis. Difffusion pscudotime robustly reconstructs lincage branching. Nature Methods, 13(10)845-8844 2016.
[20] Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-leve diffusions with generative mns. In International Conference on Machine Learning, pages 2417-2426 PMLR, 2016.
11

---

[21] Jomathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models, Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[22] Takeshi Koshizuka and Issei Sato. Neural lagrangian schrödinger bridge. arXiv preprint arX1w:2204.04853. 2022.
[23] Gioele La Manno, Ruslan Soldatov, Amit Zeisel, Emelie Braun, Hannah Hochgerner, Viktor Petukhov, Katja Lidschreiber, Maria E. Kastriti, Peter Lönnerberg, Alessandro Furlan, Jean Fan, Lars E. Borm, Zehua Liu, David van Bruggen, Jimin Guo, Xiaoling He, Roger Barker, Erik Sundström, Gonçalo Castelo- Branco, Patrick Cramer, Igor Adameyko, Sten Linnarsson, and Peter V. Kharchenko. RNA velocity of single cells. Natwre, (6000700))49444499950 2018.
[24] Jean-Micha Lasry and Pierre-Loui Lions. Mean field games. Japanese journal of mathematics, 2(1): 229-260, 2007.
[25] William Lech and Ronald Coifman. Höldeer-Lipschitz Norms and Their Duals on Spaces with Semigroups. with Applications to Earth Mover's Distance. J Fourier Anal Appl, 22(4))910-953, 2016.
[26] Jun Lei and Marthe Howard. Targeted deletion hand2 enteric neural precursor cells affects its func- tions in neurogenesis, neurotransmitter specification and gangliogenesis, causing functional aganglionosis. Development, 138(21) 4870--0078, 2011.
[27] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neurul Information Processing Syxtems, 33:2503-2515 2020.
[28] Toshio Mikami. Monge's problem with a quadratic cost by the zero-noise limit of h-path processes. Probability theory and related fields, 129(2):245-28) 260. 2004.
[29] Toshio Mikami and Michèle Thieullen. Duality theorem for the stochastic optimal control problem. Stochastic processes and their applications, (16(12)1881--188) 2006.
[30] Gal Mishne, Uri Shaham, Alexander Cloninger, and Israel Cohen. Diffusion nets. Applied and Computa- tional Harmonic Analysis, 47(2):259-283 2019.
[31] Kevin R. Moon, Jay S. Stanley, Daniel Burkhardt, David van Dijk, Guy Wolf, and Smita Krishnaswann Manifold learning- based methods for analyzing single-cell RNA RNA-sequnn data. Current Opinion in Systems Biology, 7:36 2018
[32] Kevin R. Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel B. Burkhardt, William S. Chen, Kristina Yim. Antonia van den Elzen, Matthew J. Hirn, Ronald R. Coifman, Natalia B. Ivanova, Guy Wolf, and Smita Krishnaswamy, Visualizing structure and transitions in high-dimensional biological data. Nat Biotechnol, 37(12):1482-144 2019.
[33] Michele Pavon, Giulio Trigila, and Esteban Tabak. The data-driven schrödinger bridge. Communications on Pure and Applied Mathematics, 74(7):1545-) 1573, 2021.
[34] Gabriel Peyré and Marco Cuturi. Computational Optimal Transport.
2020.
[35] Neha Prasad, Karren Yang, and Caroline Uhler. Optimal transport using gans for lineage tracing. arXiv preprint arXiv:2007. 12098, 2020.
[36] Jack Richter-Powell, Yaron Lipman, and Ricky TQ Chen. Neural conservation laws: A divergence-free perspective. arXiv preprint 10.07:2200007741, 2022.
[37] Wouter Saelens, Robrecht Cannoodt, Helena Todorov, and Yvan Saeys. A comparison of single-cell trajectory inference methods. Nature biotechnology, 37(5)55475554, 2019.
[38] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Arych Solomon, Siyan Liu, Stacie Lin, Peter Berube, Lia Lee, Jenny Chen, Justin Brumbaugh, Philippe Rigollet, Konrad Hochedlinger, Rudolf Jaenisch, Aviv Regev, and Eric S. Lander. Reconstruction of developmental land- scapes by optimal- analysis of single-cell gene expression sheds light on cellular reprogramming. Cell, 2019.
[39] A. Singer. From graph to manifold Laplacian: The convergence rate. Applied and Computational Harmonic Analysis, (211::1281133 2006.
[40] Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions Graphics (ToG), 34(4):1-11, 2015.
12

---

[41] Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. Advances in Neural Information Processing Systems, pages 11918-11930, 2019.
[42] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
[43] Kelly Street, Davide Risso, Russell B Fletcher, Diya Das, John Ngai, Nir Yosef, Elizabeth Purdom, and Sandrine Dudoit. Slingshot: cell lineage and pseudotime inference for single-cell transcriptomics. BMC genomics, 19(1):116) 2018.
[44] Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. TrajectoryNet. A Dynamic Optimal Transport Network for Modeling Cellular Dynamics, In Proceedings of the 37th International Conference on Machine Learning, pages 9526-9536. PMLR, 2020.
[45] Alexander Tong. Guillaume Huguet, Dennis Shung. Amine Natik, Manik Kuchroo, Guillaume Lajoie, Guy Wolf. and Smita Krishnaswamy, Embedding signals on graphs with unbalanced diffusion earth mover's distance. In ICASSP 2022-2022 IEEF International Conference Acoustics, Speech and Signal Processing (ICASSP), pages 5647-5651, 2022.
[46] Alexander Y Tong, Guillaume Huguet, Amine Natik, Kincaid MacDonald, Manik Kuchroo, Ronald Coif- man, Guy Wolf, and Smita Krishnaswamy. Diffusion earth mover's s distance distribution embeddings. In International Conference on Machine Learning, pages 10336-10346. PMLR, 2021.
[47] Cole Trapnell, Davide Cacchiarelli, Jonna Grimsby, Prapti Pokharel, Shuqiang Li, Michael Morse, Niall Lennon, Kenneth Livak, Tarjei S Mikkelsen, and John Rinn. Pscudo-temporal ordering of individual cells reveals dynamics and regulators of cell fate decisions. Nature biotechnology, 32(4):381, 2014.
[48] Jori van der Raadt, Sebastianus HC van Gestel, Nael Nadif Kasri, and Cornelis A Albers. Onecut transcription factors induce neuronal characteristics and remodel chromatin accessibility. Nucleic acids research, 47511)55587-562,, 2019.
[49] Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schrödinger Bridges via Maximum Likelihood. Entropy, 23(9): 1134, 2021.
[50] Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
[51] Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International Conference on Machine Learning pages 10794-10804. PMIRR, 2021.
[52] Stephen Zhang, Anton Afanassicv, Laura Greenstreet, Tetsuya Matsumoto, and Geoffrey Schiebinger. Optimal transport analysis reveals trajectories in steady-state systems. Technical report, 2021.
[53] Ziqi Zhang and Xiuwei Zhang. Inference of high-resolution trajectories in single-cell RNA-seq data by using RNA velocity. Cell Reports Methods, (6)) 100095, 2021.
13

---

1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In the limitations and broader impact section.
(c) Did you discuss any potential negative societal impacts of your work?
(Yes] In the limitations and broader impact section. We do not foresee any immediate negative societal impacts of this work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] We stated all assumptions the main body of the text and provide additional details in the supplementary materials.
(b) Did you include complete proofs of all theoretical results? [Ys]] Complete proofs are available in the supplement.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material [Yes] Code and links to publicly available data are provided in the supplemental material.
URL)?
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] All training details are provided in the supplement.
(c) Did you report eтror bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Compute resources are discussed in the supplement.
4. If you are using existing assets code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes] See supplement
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you't 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[Yes] Our data does not contain any personally identifiable information or offensive content. The Embryoid body data is cell culture, and the AML Cancer dataset is patient derived but contains no identifiable information.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "diZ7y-V0EE", "reviewer": "Reviewer_JuWZ", "summary": "The authors propose a new architecture, called Manifold Interpolating Optimal-Transport Flow (MIOFlow), to learn the dynamics of snapshot data like single-cell data. It consists of two major parts: 1. An autoencoder that embeds the data into a latent space, conserving a geodesic distance. 2. The MIOFlow that is based on neural ordinary differential equations (nODE), which is trained to model the stochastic trajectories of snapshot data. To this end the embedding of the autoencoder is combined with noise as the input to the dynamics function of the nODE.  \nThey evaluate their model on artificial data as well as cell real word single-cell data, showing the advantages of their method. \n", "questions": "Minor comment: Some sentences have too many/wrong words. Examples include:\n- line 18\n- Caption of Figure one", "limitations": "Yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- To my knowledge, the MIOFlow architecture is a novel idea. \n- The manuscript is clearly and well written.\n- They compare MIOFlow with other models and show superior performance on several datasets including single-cell data. \n\nWeaknesses:\n- Most results are qualitative rather than quantitative", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "TywpUOghXFI", "reviewer": "Reviewer_CfAn", "summary": "The authors consider a problem of modeling distribution evolution given samples for different discrete time moments.  \n\nThe paper consists of several parts.\n\nFirst, the authors describe some concepts on how to construct an autoencoder that preserves distances over a Riemannian manifold on which the data is located. They define the diffusion ground distance on the manifold. And then describe an approximation of the diffusion geodesic distance. Based on this approximation, we can estimate the diffusion ground distance and train an autoencoder such that the L2 distance in the latent space approximates these estimates.\n\nSecond, the authors introduce a definition of OT based on a time-evolving vector field and on a diffusion formulation (see (3) and (4)). In Th. 3 they repeat a theorem from [35] that OT between the measures mu and nu can be represented as an optimization w.r.t. the vector field of an integral term plus some regularization.\n\nThird, the authors model the vector field with a neural ODE and train the parameters of a neural ODE so that between each consecutive step the neural ODE solves the OT problem, formulated based on the diffusion formulation.\n\nFinally, the authors applied the proposed algorithm to model some trajectories of cells.", "questions": "- 18: “However, most have these works have focused on generative modeling” - bad wording\n\n- TrajectoryNet, and other papers, mentioned in the introduction, are not the only approaches to model stochastic flows. For example, what about the approaches from papers\n1) Large-Scale Wasserstein Gradient Flows\nPetr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Evgeny Burnaev\nNeurips, 2021\n2) Charlie Frogner and Tomaso Poggio. Approximate inference with Wasserstein gradient flows. AISTATS, 2020\n\n- The authors use a geodesic autoencoder. Any particular motivation why for using such an autoencoder? Why not some standard autoencoder? Any comparison between using the geodesic autoencoder and the standard one? How does this choice influences the experimental results?\n\n- 66: “an initial state [15, 38] in can be combined to infer the cell lineage…” - in can be? Looks like a misprint.\n\n- 91: Here, d1 and d2 are distances. Earlier, d was used as a dimension. Moreover, again in 104 d(x,y) denotes a distance. I propose to use some consistent notations. \n\n- Why consider a Riemannian manifold? What if the manifold is not a Riemannian one? What is the motivation? How does this choice influences the experimental results?\n\n- The authors use the heat kernel and estimate the diffusion ground distance using the Gaussian kernel function. What if we use some other kernel function? Will the choice significantly influence the accuracy of the approximation and final results? How does this choice influences the experimental results?\n\n- How does a value of alpha influence the experimental results?\n\n- It is unclear to what extent the diffusion geodesic distance approximates anything accurately. Any comments on the accuracy of approximation in Th. 2? Also, the approximation here is stochastic. So the formulation should be either w.r.t. probability or a.s. Any comments about this?\n\n- The geodesic distance is introduced in 139, but there is no exact definition.\n\n- 153: mu is an actual distribution on M. Is it connected somehow with mu in (3)?\n\n- How does the choice of beta in 153 influence the experimental results?\n\n- Why did the authors use Neural ODE? Why not standard RNN? Any rationale behind the choice?\n\n- Fig. 1: “With the combination of noise, a neural network learns to predict the derivative x′(t0), and an ODE Solver produces the predicted x(t1).” - x(t) is a stochastic process. Which derivative do the authors mean? Also, nothing was told about using the derivatives w.r.t the trajectories in the paper's main text.\n\n- Fig. 1: “Geodesic Autoencoder learns the latent space of representation of data is learnt such that is preserves a diffusion geodesic distance. ” - is this crucial for the considered applications?\n\n- 182: “With a slight abuse of notation, a similar distance between distributions (rather than points) is constructed in [37] where it was applied to approximate the Earth Mover’s Distance with a geodesic ground cost.” What is the novelty of this approach?\n\n- 186: instead of = we should place \\approx here.\n\n- 189: “we use it to compute distances between predicted and ground truth observations.” Which ground truth observations?\n\n- 191: “Empirically, we observed that G fails to capture meaningful distances for a large number\n of observations. ” - if it fails, then what is the point of using it?\n\n- 202: “to understand how specific genes evolved over time, by decoding in the gene space.” - I do not understand the paper's main topic then. Do the authors propose a new method? Or is this about the modeling of genes? or what?\n\n- (6): W_d -> W_2\n\n- Th. 3. Which D should we consider? Is any D OK?\n\n- 224: too many indexes i are used in the same formula. I guess a summation should be done w.r.t. to some other index.\n\n- (7): How do the authors differentiate L_m w.r.t. parameters of the Neural ODE?\n\n- 238: T -> $T$\n\n- 260: “trajectories on the manifold”. Why are they located on any manifold? How should the reader analyze Fig. 2 and conclude that the results for the proposed method are better? What is the ground truth? Any accuracy characteristics to compare methods?\n\n- Table 1: Are the numbers in the table big or small? We should compare this with some simple baseline approaches. Possibly all methods provide some bad results?\n\n- 293: “[39], Other …” -> “[39]. Other …”\n\n- Appendix, 519: “Replacing the continuity equation by the path dynamic, we have…” Why can we make such a replacement? Any justification that their proof is valid for any dissimilarity?\n\n- Appendix, 520-521: Here, the authors wrote that the autoencoder approximates the geodesic distance for sample points. First, this is not completely true, as in Th. 2 we also have a parameter alpha. Also, if these quantities are equal for sample points, it does not follow that for any pair of points from some distribution, the autoencoder will have good generalization ability so that inequalities in 522 are valid.", "limitations": "The potential negative societal impact of the work is not presented in the paper, although the authors did not explicitly state this.\n\nSome limitations of the paper are addressed.", "rating": 4, "confidence": 4, "soundness": 1, "presentation": 1, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The topic of the paper is about modeling stochastic dynamics, which is important\n\nWeaknesses\n- The text of the paper is of low quality and clarity, and contains errors and misprints, see comments below.\n- The paper is a combination of several moving pieces, and it is unclear why exactly such choices were made.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "6PAC5vgVwgh", "reviewer": "Reviewer_JVYo", "summary": "The authors propose a geodesic auto encoder to map points on a manifold to latent space whilst preserving geodesic distance on the data manifold in Euclidean distance in the latent space.\n\nThe geodesic auto encoder is then used to simulate trajectories in latent space, using diffusions and Neural ODEs. One may decode back to the data space. The trajectories are trained through an OT loss to samples. \n\nThe authors show a number of theoretical results to justify the preservation of geodesic loss in latent space, and how Wasserstein in latent space with L2 ground cost is equivalent to Wasserstein on the decoded space with geodesic ground cost.\n\nExperiments are carried out on biological data to justify the approach, and method seem to work well and results are convincing.", "questions": "- In Table 1, what does runtime refer to? Is it train time or runtime at deployment?\n- DSB interpolates via reverse diffusions and approximates the OT cost between initial and terminal distributions (no intermediary distributions) in Euclidean space where the ground cost is determined by the initial forward diffusion (reference diffusion). Which reference diffusion is used here?\n- It is stated in the appendices that DSB is used with \"default\" parameters, what are these default parameters? What network was used? How many diffusion steps?\n- The diagrams are not clear to me. Are the colours points labelled \"samples\" samples from the training data or the model? How come Figure 2 a and c have grey lines but Figure 2b does not? The legend states ground truth is labelled by a cross and the predicted by a filled in dot, then how come the ground truth varies so much by the method used to interpolate?\n- What layer widths were used for the Petal and Dyngen experiments and what was the latent dimension?", "limitations": "The authors acknowledge training ODE may be challenging in some settings and that they use <10 time steps.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- Training a geodesic auto encoder appears novel, has a lot of applications, and allows standard machine learning tools in Euclidean space to be applied to data living on a manifold. The method presented is simple and easy to understand.\n- The geodesic auto encoder + trajectory with Neural ODE/ SDE approach presented seems quite flexible and able to be modified to intermediate distributions, unlike the Schordinger bridge, which has practical importance\n- The method seem to work well and results are convincing\n- Interesting theoretical results connecting intermediate trajectories to OT\n\nWeaknesses\n- It is not clear how well the decoder generates data close to the manifold in question. Comparison to manifolds with analytic form such as a sphere/ hypersphere may be helpful to compare to a ground truth.\n- [1] appears highly relevant and should perhaps be used as a baseline, especially given the similarities using W2 loss to map cross sectional samples of a time evolving population\n- The diagrams are not clear to me. The coloured points are labelled \"samples\", are these samples from the training data or the model? How come Figure 2a and Figure2c have grey lines but Figure 2b does not? The legend states ground truth is labelled by a cross and the predicted by a filled in dot, then how come the ground truth varies so much by method used to interpolate?\n- The link to code does not work: https://anonymous.4open.science/r/bananarama-FADD/README.md\n- It is not clear why the geodesic auto encoder does not work well for a large N, as mentioned in line 191\n- A written algorithm would be helpful to understand the training procedure\n\n- Experiments\n    - It appears prediction at time 3,4 in the Appendices Table 2 for leave one out prediction are quite different to the results in Table 1 in the main text. It may give a clearer, more balanced view to include all the results. Also why was W1 chosen over W2?\n    - Number of time steps is relatively small\n  \n\nMinor:\n- \"sinkhorn\" is a name and should be capitalised in line 79 \n- line 310, should be \"compared\"\n\n[1] Proximal Optimal Transport Modeling of Population Dynamics, Bunne, 2021", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "-YDB3iG0G2m", "reviewer": "Reviewer_j3Kk", "summary": "The paper proposes a computationally efficient dynamic optimal transport framework to track samples overtime.\n\n\n ", "questions": "- in Theorem 1: what is $d^{2\\alpha}_{\\mathcal{M}}$ ? \n- what is the meaning of the sentence preceding Theorem 1 mentioning \"Thm.2\" ?\n\n", "limitations": "The limitations are well addressed. The proposed method works best with uniform sampling in the time dimension and also consider balanced optimal transport (equal mass of density functions whatever time t).  ", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall the paper is well written.\n\nSome elements are unclear however\ne.g.\n- Fig 1 caption \"Geodesic Autoencoder learns the latent space of representation of data is learnt such that is preserves a diffusion geodesic distance.\"  (?)\n\n- the notation \"d\" is used for dimension and also for distances.\n\n- some element of section 3 are well known and may not need to be written (e.g. wasserstein distance equation), \n\n- section 3.3 is unclear. \n\n- the connection between section 3  and section 4 is not clear.\n \n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept