Trajectory Inference via Mean-field Langevin in Path Space
Authors: Lénaïc Chizat
Stephen Zhang
Matthieu Heitz
Geoffrey Schiebinger
Venue: NeurIPS 2022 Accept (2022)
Keywords: trajectory inference
path space
mean-field dynamics
interacting particle methods
Langevin algorithm
Wiener measure
entropic regularization
optimal transport
Schrödinger bridge

Abstract: Trajectory inference aims at recovering the dynamics of a population from snapshots of its temporal marginals. To solve this task, a min-entropy estimator relative to the Wiener measure in path space was introduced in [Lavenant et al., 2021], and shown to consistently recover the dynamics of a large class of drift-diffusion processes from the solution of an infinite dimensional convex optimization problem. In this paper, we introduce a grid-free algorithm to compute this estimator. Our method consists in a family of point clouds (one per snapshot) coupled via Schrödinger bridges which evolve with noisy gradient descent. We study the mean-field limit of the dynamics and prove its global convergence to the desired estimator. Overall, this leads to an inference method with end-to-end theoretical guarantees that solves an interpretable model for trajectory inference. We also present how to adapt the method to deal with mass variations, a useful extension when dealing with single cell RNA-sequencing data where cells can branch and die.

Trajectory Inference via Mean-field Langevin in Path Space
Lénaïc Chizat*
Stephen Zhang"
+.
Matthieu Heitz
Geoffrey Schiebinger
Abstract
Trajectory inference aims recovering the dynamics of population from snap- shots of its temporal marginals. To solve this task, a min-entropy estimator relative the Wiener measure in path space was introduced Lavenant et al. []] and shown to consistently recover the dynamics of a large class of drift-diffusion processes from the solution of an infinite dimensional convex optimization prob- lem. In this paper, we introduce grid-free algorithm compute this estimator. Our method consists a family of point clouds (one per snapshot) coupled via Schrödinger bridges which evolve with noisy gradient descent. We study the mean- field limit the dynamics and prove its global convergence to the desired estimator. inference method with end-to-end theoretical guarantees
Overall, this leads that solves an interpretable model for trajectory inference. We also present how to adapt the method deal with mass variations, useful extension when dealing with single cell RNA-sequencing data where cells can branch and die.
Introduction
Trajectory inference aims at recovering the dynamics of population of particles given samples from its temporal marginals at various time-points. This problem arises notably the analysis of single-cell seguuninunn data 4 where one has access via destructive measurement process to the cell state samples from a large population of cells that evolve in time. this application, one would like to recover the overall dynamics of the population well as the trajectories of individual cells so as to improve our understanding of certain biological processes, such as embryonic development or tumor progression.
Trajectory inference can be cast as a regression problem where the unknown the law a continuous
stochastic process. Let X be the ambient space containing the particles and Ω = C( X)) the path-space, i.c the set all possible continuous trajectories over the time interval [0, 1]. We consider the problem of recovering probability distribution over paths / E P(2), i.e. the law of a continuous
stochastic process (X, X))xejo.11. For inference to be well-behaved, one should look for distribution
P such that its time marginals consistent with the observed snapshots and such that the overall dynamics represents satisfies some notion of regularity.
The problem of trajectory inference has received significant attention over the past several years. However, while pioneering work has focused on creating new methods, the theoretical treatment has remained limited. One class of methods focuses on recovering potential energy landscape that best fits the observed marginals. For example, Hashimoto et al. [5] encode the potential with a neural network, TrajectoryNet (3) encodes the potential neural ODE, and JKOnet (6) encodes the potential with a neural network architecture based on the JKO scheme ? and input convex neural
"Joint first authors.
Institute of mathematies, EPFT,, Switzerland University Melbourne, Victoria, Australia 5 University of British Columbia, Canada
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

networks. However, these current approaches are all nonconvex, and therefore it can be difficult to
establish rigorous guarantees.
Some guarantees have been established for the equilibrium case by Weinreh et al. g and recently consistency was established for the non-equilibrium case by Lavenant et al. [. who, regularity prior, penalize the entropy of P relative to the Wiener measure on S2 (i.e. the law of the Brownian motion). Moreover, Bunne al. have recently shown how produce an improved reference process, beyond ordinary Brownian motion D
In this work focus the estimator introduced by Lavevant al. [.. which reconstructs the process P by minimizing the entropy relative to the Wiener measure over Q. By trading off data fitting with regularization, this approach generalizes the approach of Waddington-OOT [2] to datasets with many time-points yet possibly few samples per time-point, which is required for consistency [T. yields better performance in practice [10], and is achievable through single-embryo profiling [01] While the approach of Lavevant et al. leads to a consistent estimator, it also leads to a challenging infinite dimensional convex optimization problem. It is tackled in the original paper by discretizing the space and then applying convex optimization methods. The goal of the present paper is to show that the specific structure of this estimator makes it amenable to newly introduced class of grid-free stochastic methods called Hean-FFell Langevin (MFL) dynamics 12 [3 non-linear generalization of Langevin dynamics which enjoys quantitative global convergence guarantees [4 [S] Instantiated in our context, this method consists in family of point clouds (one per snapshot) coupled via entropy regularized optimal transport, a.k.a. Schrödinger bridges [16], which evolve with noisy gradient descent. Intuitively, these dynamics can be interpreted as a (non-linear) Langevin diffusion over the path space Ω. Our approach leads to an inference method with end-to-end theoretical guarantees that solves an interpretable model for trajectory inference. We illustrate on Figure | the recovered estimator and on Figure 2 the MFL optimization dynamics.
Organization of the paper In Section 2 of our theoretical contributions is in Section 3 where introduce the MFL dynamics and show its quantitative convergence towards the min- min-entropy estimator at an exponential rate. Numerical experiments are shown in Section 4
we present the estimator of Lavenant et al. [. The heart
Notation and blanket assumptions For two probability measures Ju, V, their relative entropy is H(llz) = J log(du/ (dddd//ddy if μ < and +00 otherwise. For n N, let [n|] == {1.... n). Throughout, the ambient space X is compact convex subset of Rdd or the d-dimensional torus T and (B:), is a Brownian motion. The path space is C((0, X): and laws on path space are noted by capital letters, P or R. We denote by Pt.....t (et, Gty) P the marginal of such laws under the (joint) evaluation maps : + w(t). We use boldface letters for families of probability measures μ P(X)
2 Min-Entropy Estimator in Path Space
2.1 Trajectory Inference as Stochastic Process Inference
Model of population dynamics The first step of any inference task is to determine a prior on the ground truth. For population dynamics with mass conservation, a natural prior is given by drift-diffusion processes. We thus model the population dynamics
dX2- - (t, + VidB,, Law(Xo) = poo EE 10, 1],
(I)
where r > 0 the temperature/ddiffsiviiii assumed known, and V/ E e² 1] x X) is unknown (in case X has boundaries, one should also introduce a reflection term, that we ignore in this section). As discussed in prior works [[] is hopeless to recover the divergence-free component of the drift: is thus natural to assume that the drift is given by the gradient of function V, called the Waddington potential or epigenetic landscape in the context of cell development. The noise level r > models the inherent randomness of the ground truth process. It turns out that this noise also has favorable effect on the algorithm we develop here. This model can be extended to allow particles to branch and die [17), this extension is discussed in Section 4.2
Model of measurements Let S 44 <... < LT S 1 be a family of measurement times. contrast to the classical field of inference for stochastic processes [18] where each realization (Xe)rejo.1)

---

Samples
MFL
gWOT
MEL
0.0
12
oWOT
0.30
1.0
-10
0.8
0.5
RMS Erengy
M
N #m:64
64
0.6
0.20
0.4
-05
0.2
0.15
-15
0.0
2
2
2
82
Figure 1: (left) Sampled time-series data and reconstructions by the proposed Mean-Field Langevin dynamics (MFL) and Global Waddington-Ol M (gWOT). Observed and reconstructed marginals are colored by the measurement time L, and illustrative sample paths are overlaid in grey. We compare scenario with limited data at intermediate tnee-point (N 1) to a uniformly sampled scenario (N 64). (right) RMS Energy Distance to the ground truth over marginals for N 20 ."......2 20. illustrating the robustness of MFL to the low-sample regime. See f4.1 details.
observed times (with large), in trajectory inference each realization is only observed once, iee observe (e(T)))) with independence for each couple (i,j). The measurements summarized by the family of snapshots fú, Σド= 6x,,, E P(CC) for é E [7]. where n, 2 is the number of particles observed time 4.-
Inference problem The general goal is to recover the law P Law(X) E P(2) of the stochastic
process given the snapshots (fe...... Pez.). ). This object P contains all the information about X: in particular it contains the marginals Pt of the process and the transition probability between each family of marginals Pt, Pt.....fy. Note that, diseussed in the introduction, some works 3 6 focus on directly recovering the Waddington potential within a parameterized class of functions. This is different problem, with encouraging empirical results but weaker known theoretical guarantees.
2.2 Min-Entropy Estimator
Let W P(22) be the law the (reflected) Brownian motion X with temperature 7 and uniform initialization. The min-entropy estimator introduced in [T. is defined as the unique minimizer R of the functional TT ((2) R defined
F(R) = Σ AiFit (R, (r,) + +
where À is the regularization strength, At; = (t++1 12-1)/2 (with the convention t-1 and 1 1) and Fit,, is divergence functional that quantifies how much RG, and fa, differ, parameterized by bandwidth a >> For a particular choice of fitting functional (see below), the authors prove the following result:
Theorem 2.1 (Consistency, (T. If (1:))e111 becomes dense [0, 1] as T grows, then
lim lim R"= P weakly, almost surely.
A 40
In this paper, we consider the following data fitting term
Fit, (R,, fi.) = llo [ exp -I- 2112 dR,()))) dprlv)
(3)
20
H H R2, R, * g.) + H(a)) +C
(4)
'ee propose this name because H² is the probability measure in path space with the smallest entropy relative the Wiener measure among all those that fit least equally well the observations.

---

(2ma²) ---1==)== = is the Gaussian density and C > 0 is constant. Eq. (3) defines Fit, and is always well-defined and finite, but Eq. 4 only makes sense when H ((t) < oc. This is the negative log- -gglikelihood under the noisy observation model X = Xt., + +ZZi, where
where 9a (z) (x)
Xe,5 is the observation and Zi,j sid N(O, 1) the noise. Intuitively, when minimizing this soft-min objective over measures R2 which are mixtures of Dirac masses, each of the observed particles from fy tends to attract the particles from RG that are the closest- even if they are far away in an absolute barely influences those that are the farthest.
This data-fitting functional Fit,, slightly differs from Hgg. Rr)R,), which is the one for which Thm. 2.1 has been established in [I] but it preserves its key properties, namely joint convexity in (Re: (e)) and linearity in jr. Our choice for Fit, is motivated by more favorable behavior when R₂ is discrete measure and by its natural statistical interpretation. In this paper, we focus on the optimization aspects and leave the statistical analysis of the estimator Rº with this specific choice of Fit,, for future work. Let us also mention that Thm. 2.1 has been established for X compact manifold without boundary, while under our assumptions Xt may have boundaries.
We now turn to the presentation of the algorithm. The data model of Eq. ① was introduced only to motivate the min-entropy estimator via Thm. 2.1 and plays no role in the rest of the paper.
Optimization Method: Mean Field Langevin Dynamics
In order to compute the min-entropy estimator R E P(2), one needs to minimize an entropy regularized functional over the space of paths I P(2) - R which is of the form
F(R) = Fit(Re, ..... R,,) + TH(R)WT)
(5)
where we have posed Fit .e...... Re,) : : At, (Re, ). This is the sum of a convex functional Fit and the relative entropy with respect W If, instead of P(2). the optimization space was PP(X), problem with this structure could be solved by Mean-Field Langevin (MFL) dynamics. More explicitly, MFL dynamics are designed to minimize problems of the form F(y)= = G(μ) + (H) where G P(2) + IR is "smooth" and H is minus the differential entropy. They enjoy global convergence guarantees when G is convex. Considering for some m E No the noisy gradient descent on the function (II,.... (..,...... Tmm) ) + Gi 1 >m
i= 8z, ). MFL dynamics are obtained in the mean-field m + oo and vanishing step-size limit, see Section 3.2 for details.
Inspired by this parallel, we now design a drift-diffusion dynamics in path space that converges to the minimizer of F. The main idea is reformulation of the problem as "reduced" problem over P P(X)" with structure amenable to MFL dynamics.
3.1 Reduced Formulation
We first introduce the reduced functional F, and then state its connection to F in Thm. 3.1
For J4 P(X), let I(H,)) be the set of transport plans between μ and v4 that is, probability measures on X x X with respective marginals and v. The entropy regularized optimal transport cost between μ and v is defined, for some T; 0, as
/ co (o, wdd v) + TTH(Tlu@))= min TH @ v). TEM((C))
Tr. (u, v) =
min >eTl(a)))
(6)
where Pz y) is the transition probability density of the (reflected) Brownian motion on X over the time interval [0, t] or cquivalently the heat kernel with no-flux boundary conditions and cr (ar 3) = -7; log(p=, (z, v)). In numerical experiments we use the approximation suggested by Varadhan's 's formula when Ti is small: C,, (x,
blly x 2 (up an additive constant which
19 This optimization problem is also known as the Schrödinger bridge problem: its solution gives the most likely evolution of cloud of particles following Brownian motion, conditioned on being distributed as μ at L = 0 and at t= 16.. Some background on this problem is given in Appendix A including the definition of the Schrödinger potentials (i, e)) used hereafter. Consider the function G P(X)) + R defined for
is irrelevant when one is interested in minimizers only)

---

= (pl )) that represents family of I" reconstructed temporal marginals, by
T-1
-_ (p) 1((11))
G(p) Fit(p) +
Σ 1411
where 7, = (t,-1 tz)T. We now introduce the reduced objective F P(X)) - R. defined F(H)= G(u) (H)
where H(u) = ET= fog( lo (x)) (xx) is minus the differential entropy of the family of (
measures μr.
The next result makes the link between minimizing 3, the objective in path space and F, the reduced objective (8) It can interpreted as Representer Theorem for the min-entropy estimator This theorem is straightforward to deduce from [20) Cor. 3.5], but the specific form of F that we use here is central for our subsequent algorithmic developments
Theorem 3.1 (Representer Theorem). Let Fit PP(X) + R. be any function.
(i) If3 admits a minimizer R then -....... Rix) is a minimizer for F.
(ii) Conversely, if F admits a minimizer μ³ E P(()) then a minimizer Rt for F is built as (²(+)= [ WT (T) (f1....,
where WT (-J21,...., .T) is the law of WT conditioned on passing through 21..... 2T at times .1,...,1y br respectively and Rei.... is the composition of the transport plans hi which are optimal in the definition of (1) 1 (((+))), for == -.....,
The composition of the transport plans is obtained as Re (dæ1,.... dær) = (1.2(dz),
YT - (dde((a7
where we have introduced the conditional probability (a.k.a. "disintegrations)) characterized by 7r,++ (dæ,, 1x;+1) 7 1 (dr;+1 (z4) (r4) (dz;). In probabilistic terms, the equality in (it) can be understood saying that conditional on passing through
.r) ) at times (f1..... .fr). the paths of R are Brownian bridges with diffusivity т. The proof of Theorem i. given in Appendix
The importance of the reduced problem comes from the following facts:
The optimization space has been "reduced" from P(2)) to P(X) This reduction is enabled by the Markovian property of W Moreover, Theorem 3.1 gives an explicit method to construct minimizer for F from minimizer for F and the associated optimal transport plans ()e(-1)) This fact was already exploited in Lavenant et [
The objective function F is the sum of two terms: "smooth" function G and differential entropy term TH. This is precisely the structure tackled by MFL dynamics. Observe how the entropy in path space H (PIW*) is split into parts: the Schrödinger bridges T,, included in the "smooth" term G, and minus the differential entropy H.
Let us now describe some useful properties of G and F. Hereafter, the first-variation of G: P(X)T at μ is the unique (up to an additive constant) function [[HI E e(2)) such that for all v E P(X)T.
lim 1 (C((1-e)) e)p av) <40
G(H))
Σ vl ( - (z).
(9)
Proposition 3.2. The function G is convex separately in each of its input (but not jointly), weakly continuous and ity first-variation is given for E P(CC)" and E [T] by
6Fit. ()) [1]]
Wi.5-1
5Fit
AG;
go(x y) (g.- (v)
+
8p(()) [14]
dft, (v)
t t2
t; 1
X
"We stretch bit the term "representer theorem" which is usually reserved to miite-iimessionll reductions. "The key difference with @ Prop. B.1] is that here T, involves the entropy of y relative to the product measure (instead of the volume measure), which makes H (44 appear with a positive sign in F' (instead of negative).
5

---

where (piss b4,)) E (X) are the Schrödinger potentials for TT, (000,,,0)) with the convention that the corresponding term vanishes when involves $1 or GT,T++1 The function F is jointly convex and admits a unique minimizer fe', which has an absolutely continuous density (again denoted by p²) characterized by
(n (p=)(0) xe- -voll for i E [T].
From now on, we thus focus minimizing the reduced functional F in Eq. 8 keeping in mind that this is sufficient to minimize F, which is our main goal.
3.2 Mean-Field Langevin Dynamies
As previously mentioned, the MFL dynamics is natural when it comes to minimize functionals of the form F. G (r + H.. Here, we are increasing the entropy factor by 6 >> (recalled as an index of F') which will be useful to obtain convergence guarantees because G is not convex but Fo TH is. Using the first-variation V [po] of G given in Prop. 3.2 the MFL dynamics is defined as the solution of the following non-linear SDE of McKean Vlasov type, for 8 2 0:
[dx = VV (=,](X )s + 22(r + e)dB. + d Law(X, = 1400 1 - = Lawaaa )x!),, E (7)
(10)
where ! is the boundary reflection in the sense of Skorokhod problem, see [21 22 (this term not needed when
the d-torus). Beware that the pseudo-time the optimization dynamics in Eq. (10) should not be mistaken with the pscudo-time t of the process in Eq. ( which is now represented by the discrete exponents i [7].
The family of laws (p.)s20 of this stochastic process are characterized by the following system of PDEs (understood in the sense of distributions, with no-flux boundary conditions):
0.p ( = V (VV() (m (T + !))
E [T)
(11)
which are coupled via the quantity VV(0) [.... The link between 10 and (11) follows from Ito- Tanaka's is formula, see e.g. (23 Lem. C3]. This a multi-species PDE where each of the species y(()
attempts minimize A/e Fit (-1,s,) + (7 + H) via a drift-diffusion dynamics, and is connected to 4e (i-1)) and 4 (6+1) via Schrödinger bridges.
3.3 Quantitative Convergence
Let us now state the main convergence result proved in Appendix
Theorem 3.3 (Convergence). Let H4 :P(XTT be such that F(PO) < Then for 2 there exists a unique solution (Hz)=20 to the MFL dynamics (I]). For > 0, X the d-torus and moreover assuming that po has bounded absolute log-density, it holds
F((..) min F. S <==CC(F((())) -mim
where C = =Be-ale for some a, B > O independent of Ho and Moreover, taking smooth time- dependent Eg that decays asymptotically as â log(s) for some ā > a, it holds Fo(u.)) - Fo(u²) X log ([[g(())) log(s) + 0 and μ, converges weakly the min-entropy estimator μt"
We can make the following comments:
Under these assumptions, and with e fixed, the convergence of μ. to the minimizer of F. in relative entropy and in Wasserstein distance, also holds, with the same rate [14] Note that our convergence argument slightly differ from these works, in that the functional G is not convex itself, and the noise/diffusion term is used both to convexify the objective (since G + HH is convex) and to make the dynamics converge (via the additional eH term).
By Thm. 3.1 one can map (Pe))20 refer to as "Mean-Fiie Field Langevin in Path Space" in the title.
dynamics in P(I): it is this dynamics which
This result gives a convergence rate for an optimization dynamics in a non-convex landscape, so it is not even obvious to have global convergence. The convergence rate depends exponentially on -1/e, a standard drawback of Langevin-like dynamics in absence of log-concavity.

---

Reduced objective
Iter
Iter 50
Ieer 500
Iter 2500
3.0
0.5
125
2.8
0.0
1.00
0.75
0.0
2.4
025
1000 2000
-2
Ieeraion
Figure 2: Optimization dynamics in the setting of Figure [ with 64. (left) Evolution of the objective function F(P,)) (Eq. (8) (right) Evolution of the reconstructed marginals [k] (Eq. (13) with the iteration number k, starting from isotropic Gaussians at k 1 and colored by measurement time t,
The only point in the proof where the assumption that X is the d-torus is needed is the technical Lem. C.3 and we believe that the same convergence statement holds for a convex compact domain of RR Further extending Thm. 3.3 to the non-comeen non-compact case would however raise more profound theoretical challenges (even the well well-posedness of the MFL dynamics is not clear in this case).
Summary of theoretical guarantees: Overall, the MFL dynamics of Eq. [1] with simulated
annealing (e, > 0) converges the unique minimizer μ of the reduced problem in Eq. (8), from which we can build by Thm 3.1 the min-entropy estimator R" E P(22) which is optimal for (5) In this sense, our method leads an inference method with end-to-end theoretical guarantees for trajectory inference. However, this statement is about an idealized dynamics which in practice has be discretized in time and with particles. the next section ($3.4), we will see that the resulting discretization error can be controlled, leveraging the long history of work on mean-field dynamics.
3.4 Discretization
In practice, an approximation of the MFL dynamics is obtained by running noisy gradient descent on the function Gm : (X + R defined as
Gm (m))== G(
where
(12)
m
where E N is be number of particles used to discretized each of the time marginals 4 Since it can be shown that m V x Gm (X) [V x(({()))))))))))))(())))))))){ (see e.g. [15] Prop. 2.4]), this leads the following update equations, for 0. E [T] and j E [m]:
X
V
[2]]]( [A]) 22 + c)2
x
Rid
(4)
14
(13)
E (T)
[
m
where 1) >> is step-size, 7 are independent standard Gaussian variables and one should moreover project all particles on each step case X has boundaries. The convergence of such numerical scheme towards the MFL dynamics is standard in the mean-field literature 24 23 and holds here thanks to the stability of VV which follows from 25
In the related context of wide neural networks, quantitative bounds on the discretization error were studied in [26] Thm. 2]. shown that for kn, with high probability, |G(p²) G(p.) 1/2 Ö 112 + 122) for some > independent of k., n) and m. Inspecting their proof, can be
seen that this rate depends on the rate at which G(u)) and VV [μ] can be estimated from independent samples of μ., which studied our case 27 28 The exponential dependence in the pseudo-time of these bounds make them however not precise enough to obtain interesting long-time guarantees and we leave for future work direct convergence analysis the fully discretized MFL dynamics (note that 14 Cor. 2] gives a guarantee for the time-discretized MFL dynamics).
7

---

Each iteration of Eq. (3) requires to solve an entropic optimal transport problem in order to obtain VV. As discussed in Appendix this can be solved precision é in time O(m² (T,ē)) using Sinkhorn's is algorithm.
4 Numerical Experiments
4.1 Simulated data
In our first experiment, we compare the behaviour of the MFL dynamics to that of Global Waddington- OT (gWOT) [ in setting with few samples per time-point. Although both methods minimize a functional of the form (2). gWOT works on a fixed, discretized support: the union of all observed sample points. Thus, gWOT should perform poorly when the support set has missing regions or "gaps". To simulate this, we simulated dataset with N; samples at time 4, with N₁ Noo = 64 and N; N, S i I for N E (2", .2..... 2º). The samples are drawn from the marginals of bifurcating stochastic differential equation (see Appendix G for details).
In Figure Iefft) we illustrate two extreme cases: N = 1 (very few samples at intermediate timepoints)
and N = 64 (uniform sampling over time) respectively, with = 0.05, λRWOT 0.0025. Visually, it seems that the output of MFL is robust the few-sample regime, with relatively little qualitative difference between the reconstructed trajectories for N = 1, 64. On the other hand, the performance of gWOT degrades visibly once the set of observed points is poor reflection of the support of the underlying process.
To examine this quantitatively, we applied both MFL dynamics and g gWOT for various values of N and computed the mot-meaa--suare (RMS) Energy Distance [29] over time to an approximate ground truth (see Appendix GG. As shown in Figure for large values of N we observe that both MFL and gWOT perform similarly, but with MFL outperforming gWOT for small N as expected.
Finally, Figure 2 illustrates the evolution of the reduced objective F (8) and the reconstructed marginals 1 [b] over the course of MFI. dynamics for N = 64, À 0.05. To estimate the entropy term (z),, we used the Kozachenko- Leonenko nearest neighbour estimator [30, Section 6] (note that estimating the entropy is not needed in the algorithm; it is only used here to illustrate the convergence).
4.2 Dealing with Branching
In the context of single cell RNA-sequencing data it is crucial to be able to deal with the birth and death of cells, a.k.a. branching, which does not occur homogeneously in the domain X. Assume that we dispose of a function g E (* 1] x X) such that g(t, x) is the prior knowledge about the instantaneous growth rate of the distribution of particles in x time in the model of Eq. (I). That is, in time dt the probability of a particle X, dividing is g(t, X, dt (17). We would like to incorporate this knowledge, and also allow for additional mass variations to account for the inaccuracy of our prior g. To this end, proceed heuristically by simply replacing Tr, (Hee)Met)) in Eq. O6 with the following "uubbalanced" Schrödinger bridge problem
min TEM, (xx) /
v)d((x v) + (H(())))) + pH + (7) @ (aana)
(14)
where ir. x exp( (op(-g(()((:::))
p)) are probability measures, p >> 0 is a parameter, M_ (2(2 is the set of nonnegative measures X and 72 are the marginals of This problem can be solved with a variant of Sinkhorn's algorithm (1]
t,)/2) and Pree exp(g(x)
The rationale behind this formula is as follows: (i) the modifications Pt. of Pt. intended to
approximate the growth process ŌzAr = g(t, ))t over the time interval [t4, 1++1] and (ii) we relax the marginal constraints, with parameter p >> O, to account for the potential inaccuracy of our prior g. We illustrate on Figure the practical advantage of taking into account branching in the algorithm.
Simulated data experiment Here we consider modified version of the bistable process from 1 Figure 18] (see Appendix for details). Figure shows paths sampled from the ground truth, sample points, and the inferred trajectories. Since the lower potential well is closer to the initial condition,

---

Ground truth paths
Sample
MFL Branching
12.
10.0
0
8
5.0
0.2
-1
-1
2.5
0.1
-2
0.0
-2
0.0
0.0
0.2
0.0
0.2
0.0
Figure 3: Accounting for branching rates allows for the separation of spatial dynamics from growth. (left) Ground truth paths and sample points. (right) Reconstruction produced by MFL dynamics with and without accounting for branching. Fraction of paths terminating in the upper (resp. lower) branch annotated in the plot.
Subsample
MFL
Full dataset
day
Error
2.5
gWoT
5
Subsomate
2.5
1.0
0.0
3
UMAP1
UMAP1
UMAP1
day
Figure 4: (left) Subsampled timepoints (Subsample) from the iPSC reprogramming dataset of [2]. with 100 cells at the first and last timepoints and 10 cells at each intermediate timepoint, shown beside inferred marginals obtained using MFL dynamics and the full dataset for those timepoints. (right) Best performance measured by Energy Distance of inferred marginals to the full dataset for MFL dynamics (Langevin) and Global Waddington-O1 []] (gWOT), with the subsampled timepoints (Sample) for reference.
we observe in the ground truth that = 72% of paths terminate the lower branch. This effect is opposed by the high rate of branching in the upper branch, which ultimately results in particles being sampled in the upper branch. By accounting for branching the MFL dynamics we are able to reconstruct dynamics that isolate the effect of the potential, evidenced by the comparable proportion (77%) of inferred paths ending near the lower branch. On the other hand, neglecting the presence of branching results in growth effects being confounded with the potential, and the proportion of paths ending near each branch are roughly equal.
4.3 Reprogramming dataset
We now consider the stem cell reprogramming dataset of (2). in which a population of differentiating cells were profiled over time course using single-cell RNA sequencing. For the purpose this example, restrict our attention to days 2.5-6 inclusive making total of 8 timepoints. As previously, we consider regime where few samples are observed as snapshots MFL dynamics to infer reconstructed marginals. We reason that if the underlying process is well described by a Schrödinger bridge (as was empirically validated in (2), then MFL dynamics should be able to "improve" on the sample marginals. As proxy for ground truth, we used the full dataset (consisting of 59 154 cells over the 8 timepoints). From this, we subsampled timepoints consisting of 100 cells days 2.5 and 6, and 10 cells at the remaining intermediate timepoints. After carrying out a series preprocessing steps (see Appendix H for details), we applied MFL dynamics to the subsampled data produce reconstructed marginals. Figure 4 left) shows an example of the sampled points, MFL reconstruction, and the full dataset.
the time-series and we apply
A reconstruction "error" was then computed the Energy Distance 20 of each reconstructed marginal to the corresponding snapshot in the full dataset. Tigurerrighh shows the error for each timepoint for each of MFL and gWOT, for the best performing parameters found in parameter

---

sweep. We observe that both MFL and gWOT improve upon the subsampled snapshots, but MFL by larger margin.
Conclusion
We introduced a grid-free numerical method for trajectory inference that computes the min-entropy estimator introduced in [T] with global and quantitative convergence guarantees in the mean-field limit, This method arises naturally when decomposing the optimization problem in suitable way and, in practice, outperforms the fixed-grid method of [T.
Concerning limitations, our method shares those of the min-entropy estimator since does not incorporate fine structural prior on the structure of the Waddington potential (i.e. it is fully nonpara- metric), it may suffer from a limited statistical efficiency in high dimension. An interesting research direction is to quantify this statistical performance and to adapt our algorithm to learn the Waddington potential V in a structured parameterized class of functions, jointly with the law on paths
On a more abstract level, our main insight is that min-entropy problems in Wiener space can be solved via multi-species diffusion dynamics coupled via Schrödinger bridges. This point of view raises interesting theoretical questions, e.g. can we rigorously interpret our dynamics as diffusion in path space?
Acknowledgements
GS was supported by a Career Award at the Scientific Interface from the Burroughs Wellcome Fund, a NERF Exploration Grant, a NSERC Discovery Grant, and a CIHR Project Grant.
References
[1] Hugo Lavenant, Stephen Zhang, Young- Hcon Kim, and Geoffrey Schiebinger. Towards a mathematical theory of trajectory inference. arXiv preprint arXiv:2102.09225 2021.
[2] Geoffrey Schiebinger, Jian Shu. Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. ppuimal-transpor analysis of single-cell gene expression identifies developmental trajectories reprogramming. Cell, 176(4):928-9) 2019.
[3] Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy, Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International Conference on Machine Learning, pages 9526-9536. PMLR, 2020.
Jeffrey A. Farrell, Yiqun Wang, Samantha J. Riesenfeld, Karthik Shekhar, Aviv Regev, and Alexander F. Schier. Single-cell reconstruction developmental trajectories during zebrafish embryogenesis. Science, 36006292)) 2018.
[5] Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative RNNs. In International Conference on Machine Learning, pages 2417-2426. PMLR, 2016.
[6] Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pages 65 11 1-6528. PMLR, 2022.
[7] Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker Planckk equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
[8] Caleb Weinreb, Samuel Wolock, Betsabeh K. Tusi, Merav Socolovsky, and Allon M. Klein. Fundamental limits dynamic inference from single-cell snapshots. Proceedings of the National Academy of Sciences, 115(10)E2467-E2477 2018.
[9] Charlotte Bunne, Ya-Ping Hsich, Marco Cuturi, and Andreas Krause. Recovering stochastic dynamies via Gaussian Schrödinger bridges. arXiv preprint arXiv:2202.05722 2022.
[10] Geoffrey Schiebinger. Reconstructing developmental landscapes and trajectories from single-cell data. Current Opinion in Systems Biology,200 2021.
10

---

[11] Markus Mitteenwwigg, Yoav Mayshar, Saifeng Cheng, Raz Ben- Ron Hadas. Yoach Rais, Elad Chomsky, Netta Reines, Anna Uzonyi, Lior Lumerman. Aviezer Lifshitz. Zohar Mukamel. Hashaar Orenbuch, Amos Tanay, and Yonatan Stelzer. single-embryo, single-cell timelvesolved model for mouse gastrulation, Cell, 184(11):2825 04.004.
(ttp:/28222222222222 2842.c22, 2021. ISSN 0092-8674, doi: https:///oi.orr//10.. 10 cell.
[12] Kaitong Hu, Zhenjie Ren, David Šiška, and Lukasz Szpruch. Mean-field Langevin dynamics and energy landscape of neural networks. In Annales de I'Institut Henri Poincaré, Probabilités et Statistiques. volume 57, pages 2043-2065. Institut Henri Poincaré, 2021.
[13] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. mean field view of the landscape two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665-E77)) 2018.
[14] Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field Langevin dynamics. In International Conference on Artificial Intelligence and Statistics, pages
PMLR, 2022.
[15] Lénaïc Chizat. Mcan-field Langevin dynamics: Exponential convergence and annealing. Transactions on Machine Learning Research, 2022. URL https: wwwp:/rumtememeeeemmmm idderrummmiiddiie htt.:/orrnmnnidDD0022211222a
[16] Christian Léonard. From the Schrödinger problem to the Monge-Kantorovich problem. Journal of Functional Analysis, 262(():1879- 1920. 2012.
[17] Aymeric Baradat and Hugo Lavenant. Regularized unbalanced optimal transport entropy minimization with respect to branching Brownian motion. arXiv preprint wxXiv:2111001666 2021.
[18] Malempati Madhusudana Rao. Stochastic processes: inference theory. Springer, 2000.
[19] James R. Norris. Heat kernel asymptotics and the distance function in Lipschitz Riemannian manifolds. Acta Mathematica, 179(1):79-103, 1997.
[20] Jean- Sann-avid Benamou, Guillaume Carlier. Simone Marino, Luca Nenna. An entropy minimization approach to second-order variational mean-field games. Mathematical Models and Methods in Applied Sciences, 29(08):1553-1555 2019.
[21] Hiroshi Tanaka. Stochastic differential equations with reflecting boundary condition in convex regions. Stochastic Processes: Selected Papers of Hiroshi Tanaka, 9:157, 1979.
[22] Louire-Louis Lions and Alain-Sol Sznitman. Stochastic differential equations with reflecting boundary conditions. Communications pure and applied Mathematics, 97(4):111-33) 1984.
[23] Adel Javanmard, Marco Mondelli, and Andrea Montanari. Analysis of two- layer neural network via displacement convexity. The Annals of Statistics, 48(6):3619 3642, 2020.
[24] Alain-Sol Sznitman. Topies in propagation of chaos. éenle é'ééé de probabilités de saint-flour XIX-1989, 165-251, Lecture Notes Math., 1464, 1991.
[25] Guillaume Carlier, Lénaic Chizat, and Maxime Laborde. Lipschitz continuity of the Schrüdinger map in entropic optimal transport. arXiv preprint XrXi:2220000225, 2022.
[26] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pages 2388-2464, PMLR, 2019.
[27] Eustasio del Barrio, Alberto Gonzalez-Sanz, Jean-Michel Loubes, and Jonathan Niles-Weed, An im- proved central limit theorem and fast convergence rates for entropic transportation costs. arXiv preprint arXiv:2204,099 10900, 2022.
[28] Philippe Rigollct and Austin Stromme. On the sample complexity of entropic optimal transport. arXiv preprint or Xiv:2206. 13472, 2022.
[29] Gábor Székely and Maria L. Rizzo. Energy statistics: A class statistics hased on distances. Journal of statistical planning and inference, 1249-11722 2013.
[30] Shashank Singh and Barnabás Póczos. Analysis of k-nearest neighbor distances with application to entropy estimation. arXiv preprint 1rXiv: 08578, 2016.
[31] Lénaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xave Xavier Vialard. Scaling algorithms for unbalanced optimal transport problems. Mathematics of Computation, 87(314):2563 2609, 2018

---

[32] Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with Wasserstein loss. Advances in Neural Information Processing Systems, 28, 2015.
[33] Marcel Nutz. Introduction to entropic optimal transport, 2021.
[34] Simone Di Marino and Augusto Gerolin. An optimal transport approach for the Schrödinger bridge problem and convergence of Sinkhorn algorithm. Journal of Scientific Computing, 85(2):1-28, 2020.
[35] Michel Ledoux. Concentration of measure and logarithmic Sobolev inequalities. In Seminaire de proba- bilites XXXIII, pages 120-216. Springer, 1999.
[36] Richard Holley and Daniel Stroock, Logarithmic Sobolev inequalities and stochastic Ising models. Journal of Statistical Physics, 46(5): 159--194,, 1987.
[37] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
[38] Jean Feydy, Thibault Séjoumé, Framçois-Xavel Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and MMD using Sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2681-2690. PMLR, 2019.
[39] Gabriel Peyré and Marco Cuturi. Computational optimal transport with applications data science. Foundations and Trends@ in Machine Learning. 1((-6))355-607, 2019.
[40] Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent better than by Sinkhorn's algorithm. In International conference on machine learning, pages 1367-1376. PMLR, 2018.
[41] Jason Altschuler, Jonathan Niles-Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. Advances in neural information processing systems, 30, 2017.
[42] Alexis Thibault, Lénaic Chizat, Charles Dossal, and Nicolas Papadakis. Overrelaxed Sinkhorn Knopp algorithm regularized optimal transport. Algorithms, 14(5):143, 2021.
[43] Pierre Ablin, Gabriel Peyré, and Thomas Moreau. Super-efficiency of automatic differentiation for functions defined as a minimum. International Conference on Machine Learning, pages 32-41, PMLR, 2020.
[44] Dimitris Bertsimas and John Tsitsiklis. Simulated annealing. Statistical science, 8(1):10-15, 1993.
12


Reviews:
[{"review_id": "zbqwxAqwhw", "reviewer": "Reviewer_qDLz", "summary": "In this work authors study the problem of fitting a law on the path space induced by gradient SDE with additive and reflective Brownian motion to the set of prescribed time marginals. This is similar to multi-marginal OT with a difference the transport plan is restricted the be the law induced by an SDE.  The problem has been recently studied in Lavenant et al where, essentially, an estimator used in the current paper has been developed and analysed. Authors observed (which is not new) that one can translate entropic optimisation problem on the Wiener space (as in Lavenant et al.) to the optimisation problem on the product space of probability measures supported on the compact set \\mathcal X (the compactness being a technical condition). This reformulation then open the avenue for mean-filed langevin dynamics (MFLD) (or gradient flow on W_2 space)  to be used as an optimisation ‘algorithm’ for which convergence along with the convergence speed can be deduced from recent works in the literature. Indeed mean-filed langevin dynamics has been recently studied by Mei et al (2018) and Hu et al (2019) where convergence has been proved rigorously. These result have been recently extended by Chizat 2022 and Nitanda et al 2022 that showed that under appropriate conditions convergence of MFLD is exponential. In the current paper authors showed, by assuming compactness of the space \\mathcal X, that results of Chizat and Nitanda et al apply.  \n\nIt’s worth remarking that conceptually this work is related to ‘Mean-Field Neural ODEs via Relaxed Optimal Control’ by Jabir at al 2019 and \n‘Gradient Flows for Regularized Stochastic Control Problems’ by Siska et all 2020, where detailed analysis of coupled MFLD has been studied. In these works the cost function was slightly simpler and significant amount of regularisation was required to establish convergence.  \n", "questions": "- It wasn’t clear where in the proof you used the fact that only SDE with drift of grade it type are needed? I would appreciate more discussion on that in the appendix \n- When stating theorems please make sure assumptions are clearly stated \n- In eqn 10 reflected BM is used. How do I know X lives in \\mathcal X ? Don’t I need to consider Skorohod problem?\n- Similarly to derive 11 from 10 Ito formula is used. Please Don’t you need Ito-Tanaka formula here? \n- Proof of Lemma B.2 : for completes please provide more details rather then just citing [21] and [23]\n- In deriving LSI compactness of \\mathcal X is critical. Can you comment more about the constants how the diameter of D affects convergence ?\n- please provide more details in the proof of Proposition D.1 rather than just citing other’s people work. That will make the paper easier to read. \n", "limitations": "Mathematically main limitation is the compactness of \\mathcal X, but the result is still very interesting.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Very interesting contribution tackling difficult problem of proving convergence rates for gradient flow algorithms for the inference on the path space. \n-  The paper is (mainly) well written making highly technical material accessible to broad audience \n- cell-RNA sequencing is a very promising area of applications for the ML models and I anticipate significant impact of this work on it\n\nWeaknesses:\n- The compactness of \\mathcal X is not a natural assumption but seem necessary for the analysis to go through \n- Some parts of the proofs are written in haste and I would appreciate more details (see comments bellow)\n- When citing other works, especially in the proofs, please give exact detail of the results used rather than just the name of the paper… ", "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": ["I don’t know"]}, "extra_text": {}}, {"review_id": "S9lgAihtW_", "reviewer": "Reviewer_mVx5", "summary": "This paper tackles the challenge of recovering the law of a stochastic process from sample observations of its marginals. The authors propose to approximate an existing, continuou-time, consistent estimator minimising a given regularised functional by another reduced, discrete-time estimator defined as the sum of a functional G and an entropy term H. The functional G is built from solutions to Schrodinger bridge problems between consecutive marginals. To obtain the reduced estimator, the authors consider a Langevin SDE of Mc-Kean-Vlasov type and show that its dynamics converge to the unique minimiser of the reduced optimisation problem. Numerical experiments are performed on simulated data from a bifurcating SDE as well as on real-world RNA-sequencing data. ", "questions": "Is it possible to provide a more concrete interpretation of the integral in L. 156?\n\nCan you provide additional details or comment on the time- and space-complexities of the proposed algorithm? In particular, how does it scale in the length and dimensionality of the paths?\n\nHow are the regularising constants chosen in the experiments?\n\nIn Figure 1, for large values of N, gWOT seems to start outperforming MFL. Does this trend continue if N is chosen in the order of 10^3, 10^4...?\n\nWould the proposed methodology apply if one considered a different reference measure than Brownian motion (BM)? In particular, what would happen if one considered instead a non-Markovian reference process, for example fractional BM, or other processes with memory?\n\nRegarding benchmarking as well as assessment of the success for the proposed method, have you considered divergences for probability measures supported on pathspace, such as the families of MMD distances studied in [1]?\n\n[1] Salvi, C., Lemercier, M., Liu, C., Horvath, B., Damoulas, T., & Lyons, T. (2021). Higher order kernel mean embeddings to capture filtrations of stochastic processes. Advances in Neural Information Processing Systems, 34, 16635-16647.", "limitations": "As the authors highlight in their conclusion, the proposed algorithm suffers from the typical drawbacks of min-entropy estimators which makes it not scalable to high-dimensional problems.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Recovering the law of a stochastic process from sporadic observations of its marginals is an important and difficult task, both at a theoretical and practical level, and this work makes good contributions in this direction. I enjoyed reading the paper and found its material original, well-written and the mathematical claims thoroughly justified in the appendix.\n\nOne possible weakness is the scalability of the proposed algorithm to high-dimensional and long trajectories, as noted by the authors themselves in their conclusion.\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": {}}, {"review_id": "EdiYctKNv-", "reviewer": "Reviewer_2ZNK", "summary": "The authors consider the problem of finding the dynamics of a population from snapshots of its temporal marginals. They introducing an estimator which which is the sum of the goodness of fit to data term and entropy-based term (a.k.a. prior/regularizer). Instead of minimization of the entropy-regularized functional over the space of all paths, they introduce the “reduced” functional for optimization that includes additional entropy-regularized Optimal Transport problems between intermediate positions. The reformulated problem turns out to be equivalent to the initial problem. This follows from Theorem 3.1 which the authors prove (they call it the Representer theorem as its idea is analogous the idea of the original representer theorem). The resulting “reduced” problem is practically more tractable as it considers only a finite amount of time moments. The authors use Mean-Field Langevin (MFL) dynamics to approach the functional and recover the trajectory of the stochastic process, finding the optimal corresponding marginal distributions simultaneously. In the experimental session, they compare proposed approach with Global Waddington-OT (gWOT) in the simulated and real data experiments in small dimensions.", "questions": "(1) Consistency Theorem 2.1 seems to be formulated in [1] for a different than (3) functional (lines 110-111). Does the consistency hold for (3) which the authors use?\n\n(2) How do you solve numerically PDE (11)?  \n\n(3) What is the computational complexity of the proposed method?\n\n(4) How well does the method scale to higher dimensions? It would be great to see the error dynamics on some simulated dataset with the increase of the dimension.\n\n(5) The unbalanced OT appears rapidly and out of nowhere. Is it true that all the theoretical results of the previous sections apply to formulation (14) instead of the original entropic balanced OT?\n", "limitations": "The authors sufficiently fully describe the limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n(1) While most other methods are based on recovering the Waddington potential directly, the current paper uses point of clouds (particles) to recover the desired dynamics.\n\n(2) The theoretical part of the paper is interesting, specifically I would like to highlight the Theorem 3.1 in the spirit of Representer which reduces optimization over paths to optimization over finite products of measure spaces at observation times. Additionally, the paper proves the result on the exponential convergence of the functional which they optimize (Theorem 3.3).\n \n**Weaknesses**\n \n(1) The method is point-cloud-based (particle-based) and it is presumably not very scalable for high dimensions as it would require exponential amount of particles.\n\n(2) The authors claim that existing approaches can be difficult to establish rigorous guarantees and establish guarantees for their method. However, when it comes to practice, it seems like in their method a large gap anyway appears between the theory and practice. For example, Thm. 3.3 is proved for measures, but in practice they are finitly approximated by particles (eq. (12)); entropy term is estimated by the nearest neighbor estimator (line 259); sinkhorn terms are also computed between discrete distributions – all these approximations introduce minor biases here-and-there which raises questions to which extent the proposed theory (Theorem 2.1, 3.3.) is still applicable and advantageous w.r.t. The other methods.\n\n(3) The comparisons are a little bit limited – only the gWOT method is included. What about (some) of the other related methods? Specifically, seems like [6] or [9] might be relevant.\n\n(4) The method has some hyperparameters which are not very clear how to choose.\n\n(5) There is not clear explanation of the final algorithm with description of all the inputs, outputs, parameters, etc. Seems like these details can be collected from the main text and appendix, but this is a negative factor for reproducibility and readability.\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