Deep Momentum Multi-Marginal Schrödinger Bridge
Authors: Tianrong Chen
Guan-Horng Liu
Molei Tao
Evangelos Theodorou
Venue: NeurIPS 2023 poster (2023)
Keywords: Schrödinger Bridge
Trajectory Inference
Optimal Transport

Abstract: It is a crucial challenge to reconstruct population dynamics using unlabeled samples from distributions at coarse time intervals. Recent approaches such as flow-based models or Schrödinger Bridge (SB) models have demonstrated appealing performance, yet the inferred sample trajectories either fail to account for the underlying stochasticity or are unnecessarily rigid. In this article, we extend SB into phase space and propose $\underline{D}$eep $\underline{M}$omentum Multi-Marginal $\underline{S}$chrödinger $\underline{B}$ridge (DMSB), a novel computational framework that learns the smooth measure-valued spline for stochastic systems that satisfy position marginal constraints across time. By tailoring the celebrated Bregman Iteration and extending the Iteration Proportional Fitting to phase space, we manage to handle high-dimensional multi-marginal trajectory inference tasks efficiently. Our algorithm outperforms baselines significantly, as evidenced by experiments for synthetic datasets and a real-world single-cell RNA sequence dataset. Additionally, the proposed approach can reasonably reconstruct the evolution of velocity distribution, from position snapshots only, when there is a ground truth velocity that is nevertheless inaccessible.

Deep Multi-Marginal Momentum Schrödinger Bridge
Tianrong Chen, Guan-horng Liu, Molei Tao, Evangelos A. Theodorou Georgia Institute of Technology, USA (tianrong. chen, mtao, hteeeeoesseeeeoooooaaaaaeeeh..
Abstract
It crucial challenge to reconstruct population dynamics using unlabeled sam- ples from distributions at coarse time intervals. Recent approaches such as flow- based models or Schrödinger Bridge (SB) models have demonstrated appealing performance, yet the inferred sample trajectories either fail to account for the underlying stochasticity are unnecessarily rigid. In this article, We extend the approach in [1] to operate in continuous space and propose Deep Momentum Multi-Marginal Schrödinger Bridge (DMSB), novel computational framework that learns the smooth measure-valus spline for stochastic systems that satisfy position marginal constraints across time. By tailoring the celebrated Bregman Iteration and extending the Iteration Proportional Fitting to phase space, we manage handle nigh-dimensiooal multi-marginal trajectory inference tasks efficiently. Our algorithm outperforms baselines significantly, as evidenced by experiments for synthetic datasets and real-world single-cell RNA sequence dataset. Addition- ally. the proposed approach can reasonably reconstruct the evolution of velocity distribution, from position snapshots only, when there is ground truth velocity that is nevertheless inaccessible.
Introduction
We consider the multi-marginal trajectory inference problem, which pertains to elucidating the dynamics and reactions of indiscernible individuals, given static snapshots of them taken at sporadic time points. Due to the inability of tracking each individual, one considers the evolution of the statistical distribution of the population instead. This problem received considerable attention, and associated applications appear various scientific areas such as estimating cell dynamics [2, 3], predicting meteorological evolution [4], and medical healthcare statistics tracking [5]. [6, 7] constructed an energy landscape that best aligned with empirical observations using neural network. [8, 9] learn regularized Neural ODE [10] to encode such potential landscape. Notably, in the aforementioned work, the trajectory of samples is represented a deterministic way. In contrast, [11, 12] employ Schrödinger Bridge (SB) to determine the most likely evolution of samples between marginal distributions when individual sample trajectories are also affected by environmental stochasticity. Yet, these approaches scale poorly w.r.t. the state dimension due to specialized neural network architectures and computational frameworks.
SB can be viewed nonlinear SDE that yields a straight path measure between two arbitary distributions. The straightness implied by achieving optimality of minimizing transportation costs (i.e. 2- 2-Wasserstein distance (W2)). We note SB is often related to Score-based Generated Model (SGM), both of which can be used for generative modeling by constructing certain Stochastic Differential Equation (SDE) that links data distribution and a tractable prior distribution (i.e. 2 marginals). SGM accomplishes the generative task by first diffusing data prior through pre-specified linear SDE, during which neural network is also learned to approximate the score function. Then this score approximator is used reverse this diffusion process, and consequently establish the generation. tritically-damper
solution to the entropy regularereg optimal transport problem. SB seeks
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Table Comparison between different models in terms of optimality and boundary distributions po and Pt. Our DMSB extends standard SB, which generalizes SGM beyond Gaussian priors, to phase space, similar to CLD. However, unlike CLD, DMSB jointly learns the phase space distributions, i.e., po( u) go (1/2)) and Po( v) ps(2)) (v(z). In other words, DMSB infers the underlying phase state dynamics given only state distributions.
Models
Optimality
Po(-)
P.(.) (O, Σ) W( Σ) @ (O ps( PB p=(2)go(o(z)
SGM [20]
(A(
CLD [13]
Pa(x) @ (0, Σ)
SB [14]
4 2 + kinks
(
DMSB (ours)
W2 smooth
PA(x)g0(U|z)
Langevin Diffusion (CLD) [13] extends the SGM SDE to the phase space by introducing an auxiliary velocity variable with a tractable Gaussian distribution at both the initial and terminal time. The resulting trajectory in the position space becomes smoother, stochasticity is only injected into the velocity space, and the empirical performance and sample efficiency are enhanced due to the structure of the critical damped SDE. The connection between SGM and SB has been elaborated in [14, 15] and scalable mean matching Iterative Proportional Fitting algorithm (IPF) is proposed to estimate SB efficiently in high dimensional cases. Applications of SB, such managetto-image transformation [16, 17], RNA trajectory inference [11], solving Mean Field Game[18), Riemannian interpolation [19], demonstrate the effectiveness of SB in various domains.
In this work, we start with SB phase space (termed momentum SB, mSB in short), and then further investigate mSB with multiple empirical marginal constraints present in the position space, which was formulated as multi-marginal mSB (mmmSB) [1]. This circumvents the need for expensive space discretization which does not scale well to high dimensions. We also address the challenge of intricate geometric averaging in continuous space setup by strategically partitioning and reorganizing the constraint sets. Furthermore, by incorporating the method of half-bridge IPF. The optimality of transportation cost in SB leads to straight trajectories, and if one solves N 2- 2-marginal SB problems and connect the resulting trajectories to match N+1 marginals, the connected trajectories will have kinks at all connection points. On the contrary, in mmmSB, the optimality of transportation cost leads to smooth measure-spline over the state space that also interpolates the empirical marginals. Therefore, this approach is highly suitable for problems originated from physical systems and/or those that should have smooth trajectories, such as trajectory inference in single-cell RNA sequencing. Our research will emphasize on solving mmmSB efficiently in high-dimensions (thus the approach will differ from that in the seminal work [1]; see Sec.4). The differences between our algorithm and prior work are demonstrated in 1able. and the main contributions of our work are fourfold:
enhance the algorithm's 's computational efficiency
We extend the mean matching IPF to phase space allowing for scalable mSB computing.
We introduce and tailor the Bregman Iteration [21] for mmmSB which makes it compatible with the phase space mean matching objective, thus the efficient computation is activated for high dimensional mmmSB.
We show how overcome the challenge of sampling the velocity variable when it not available in training data, which enhances the applicability of our model.
We show the performance of proposed algorithm DMSB on toy datasets which contains intricate bifurcations and merge. On realistic high-dimension (100-D) single-cell RNA-seq (scRNA-seqq) datasets, DMSB outperforms baselines by a significant margin in terms of the quality of the generated trajectory both visually and quantitatively. We show that DMSB able to capture reasonable velocity distribution compared with ground truth while other baselines fail.

---

2 Preliminary
2.1 Dynamical Schrödinger Bridge problem
Dynamical Schrödinger Bridge problem has been extensively studied in the past few decades. The objective of the SB problem solve the following optimization problem:
min DEL (T](E),
(1)
TEll(po.P))
where π E [l pr) belongs set of path measures with its marginal densities t and T being Po and Pr. & is the reference path measure (i.e., [14] sets & as Wiener process from Po). The optimality of the problem (1) is characterized by set of PDEs (3).
Theorem 2.1 ([22]). The optimal path measure πT in the problem (1) is represented by forward and backward stochastic processes
dx;- [2 V log V. dt /2 dww, Xy po-
(2a)
dx; = 2 v log T + 22 dw,, XT ~ pr
(2b)
in which V V C are the solutions to the following coupled PDEs,
ay, A. s.t. (O, (0, po(-), WI, (T, -
av,
Öt A
Ot
(3)
The stochastic processes of SB in (2a) and (2b) equivalent in the sense of 't' [0, x(2
(26) Pi = p. Here p stands for the marginal distribution of SB at time 4, which also represents the marginal density of stochastic process induced by either of Eq.2. The potentials V and VG explicitly represent the solution Tokker-Plank Equation (FPE) and Hamiton--Jacoii-ellln after exponential transform [14] where FPE describes the evolution of samples density and HJB represents for the optimality of Eq. I. Furthermore, the marginal density also obeys factorization of y = L T Such rich structures of SB will later on be used to construct the log-likelihodd objective (Thm.B.1) and Langevin sampler for velocity ($4.4).
equation (HJB)
To solve SB, prior work have primarily used the half-bridge optimization technique, also known as Iterative Proportional Fitting (IPF), in which one iteratively solves the optimization problem with one of the two boundary conditions [14, 15, 23],
(d+1) = arg min (KL(T||T rEll(-.P))
(d+2) = arg min Dx D (T|( =Ell(po.))
with initial path measure πt (0) : By repeatedly iterating over aforementioned optimizations until the algorithm converges, the SB solution will be attained SBB lim r(d) [24]. In addition. [25] shows that the drift term in SB problem can also be interpreted as the solution Stochastic Optimal Control (SOC) problem by having optimal control policy z = 2 Vx log (T, x:):
dx; 2,dl + /2dw, Xn Pu.
z" a'(x)= E arg min E ZE
2 2d
s.t
Pr.
This formulation will be used later for constructing phase space likelihood objective function in $3. Regarding solving the half-bridge problem, abundant results exist the literature for the vanilla SB deseribed above [14, 15, 23), but will be solving different SB problem; Prop.4.1 for formulation and $.4 for a solution.
2.2 Bregman Iterations for Multiple Constraints
Bregman iteration [21] can be viewed multiple marginal generalization of IPF, and it is widely used to solve entropy regularized optimal transport problem [I] with multiple constraints. The algorithm can efficiently solve problems in the form of,
inf KL (T|E), xEK.
3

---

SB Problem
Momentum SB Problem
Marginal Momentum SB Problem
min 11=1haxett Jo
min | 1111tt
min |alll 30 dx, v,et dv, 2,6t + gdw, & - P. i [0.1....
Problem Formulations
dx, z,dt gddr, Xn Pa- X4 Pr
dx, velt zdtt gdw;
[g Vul
@
1x vyl
Solutions Likelihood Objective IPF (Chen al.)
Baseline:
Prop
Phase Space Likelihood Objective IPF (ours)
Decouple constraints OptimisetSSe
DMSB (our algorithm)
(Chen al.)
(ours)
Snxphhot
Toy Example
Snapehot
Snapshot
Snapshot
Generated Samples
A summary of various SB problems and corresponding algorithms. The toy example in the 3rd row illustrates that vanilla SB determines 'straight' paths (modulo fluctuations due noise) between pairwise empirical marginals, while our multi-marginal momentum SB approach establishes smooth measure-spline between marginals in the position space (albeit still stochastic, the path is smooth between any pair of adjacent 2 marginals, because noise is added to velocity, and the path is also smooth across different pairs of adjacent 2 marginals per design.
Figure I:
where K. is the intersection of multiple closed convex constraint sets KG: K. = nt KG Bregman Projection (BP) defined as optimization w.r.t one of the constraint NCt,
L (E) == arg min K L(T(e),
=EK,
and d-th Bregman Iteration (BI) is recursively computing BP over all the constraints in I::
v0 L, n
PKL (m (m((.n-))),
The initial condition for (d 1)-th BI is x (d+1.0) (d.L), Under certain conditions (see e.g., [24)), one has that (d...) converges to the unique solution:
(d.L)
(E) as d
Remark 2.2. One BI traverses all constraints via multiple BPs, and each BP solves optimization problem with one constraint. can notice that the BI will become the aforementioned IPF procedure solving SB problem (1) by defining T. 2, 1C IT(po )), K2 (1))
Table 2: Mathematical notation.
Notation
Definition position variable
Notation
Definition
P
position distribution p(x)
velocity variable
velocity Distribution (v)
concatenation of [x, v
distribution of μ(x, v)
m
3 Momentum Schrödinger Bridge
We first describe how to conduct half-bridge IPF training in the phase space, which can be used to solve momentum SB (mSB) problem with two marginals constraints. This scalable phase space half-bridge technique will then be applied to multi-marginal cases (Sec.4). Fig. demonstrates how we develop an algorithm based on [14]. Notations used following sections are listed in Table.2. mSB extends SB problem to phase space, which consists of both position and velocity.

---

We will first consider boundary distributions that depend on both x and v, although eventually we will use this as module find transport maps between two distributions that only depend on position x, as velocity v is auxiliary variable artificially introduced for obtaining smooth transport. Conceptually, as an entropy regularized optimal transport problem, SB tries obtain the straightest path between empirical marginals of positions x with additive noise, but mSB aims at finding the smooth interpolation between empirical marginals of x [26] conditioned on boundary velocity distributions (see Fig.1). Such smooth measure-valued splines in the position space are obtained by the optimization problem in the phase space [I]:
dx:
0
0
0
0 0
(( 8.l π = Law(x, v): (EI((H.)TT)
min
di
+
di
+
dw
dv:
gb
Z1
0 gt
dmm.
fiv.c)
s(r)
2(4)
B())
Similar to Theorem 2.1, one can derive set of PDEs using the potential functions (t,, x, v) and (t,,,, x, v), and subsequently apply IPF procedure to solve the problem. The formulation of the phase space PDE can be found in Appendix.B.. B.2. Such PDE representation of mSB results in straightforward yet innovative log- oiglliiel training that enables efficient optimization the IPF.
Proposition 3.1 (likelihood bound). The half-bridge IPF in phase space
π. (d+1) arg min (KI(TT|T<())
(d+2) arg min DKL(T] (KI(T||(((tt) =Ell(-PT)
#Ell(por)
represents the bound of the likelihood and gives approximate likelihood training:
Z, == arg min logp(mno. 0)
Ż, == arg min - log pl mn, T))
Ze
%
T
where log 0) x Em,
2 2r + Zt gv log 2r11? de.
and m samples from:
dm
f 8Zi dt + gg()dd,, my PT
(5)
2, - 0 and P. is the density of path measure induced by eq.5 at time A similar result for log( mmr T) can be obtained a similar derivation.
Proof. See Appendix B.1.
Remark 3.2. After optimizing Z,, the reference path measure becomes eq.5, which implies и E
[[,, (r)), ), i.e., the constraint in half-bridge IPF is satisfied. A path measure æ is induced by either Ze or
2, As being mentioned in Remark.2.2. One half-bridge IPF is basically one BP and one IPF is one BI. Prop.3.1 provides convenient way perform one BP in the form of π arg min,ek, DKL(m]H by maximizing log-likelihood given constraint K. and reference path measure म.
Prop.3. provides an alternative way to conduct the BI which will be heavily used in mmmSB $3, and it is computationally efficient after parameterizing and discretization (54.4).
Deep Momentum Multi-Marginal Schrödinger Bridge
We first state the problem formulation of momentum multi-marginal Schrödinger Bridge (mmmSB). Different from previous two marginals case, we consider the scenario where N + 1 probability measures 14, are lying time 1, In addition, velocity distributions are not necessarily known.
Proposition 4.1 ((I). The dynamical mmmnB with multiple marginal constraints reads:
N-1
min ((T) = Σ
KL (etsstineeeetaee S.1 π E KK := N N
(6)
0
5

---

Fori [O...N] min J(x)
Forj [N...]] min J(a)
(d)-th BI
bboundary
x Mpridge
x² x
st E Kg
Sohved by Sec4-3
Abourday
(d + 1) Bl
Porj [0 min J(n)
Por E [0...V] min J(x)
min J(se)
Kbridge
:Solved Sec.44
x² Karidge
2*
x² K
Figure 2: The procedure details the Bregman Iteration (BI) employed in DMSB. The gray and blue blocks represent the BP step performed under Kioundary constraint for forward and backward policies, respectively. The red block signifies the BP step executed under the Koridge constraint. Algorithms for training and sampling can be found in Appendix.D.
{
Tegsty dm = Mto: / Moodvoo = Pto
where:
Kit =
{ Tey sity dmnx = Mens
Kix
Mey dvey = pew)
{
Tesste dm = Mtt / re, = Mez [u,dv: =Prs (7)
=m]
Kr.
and K. is the intersection close convex set of Key.
The problem described in Prop.4.1 can be solved by classical BI algorithm integrated with Sinkhorn method [1]. However, due to the curse of dimensionality and unfavorable geometric explicit solution, the BP cannot be applied in high-dimensional and continuous state space directly. To tackles these difficulties, we parameterize the forward and backward policies z, and 2, by pair of neural networks. We further decouple and resemble the constraints by which it enables the scalable likelihood IPF and avoids the geometric averaging issue under mmmSB context.
4.1 Decoupling and Reassembling Constraints
We decompose the constraint set (7) by
XCO IC. Te ps,, dm = M4: Mr, dv.. = pt:
Teates dm = Pr. Pr, dv,, = Pt.
KG, k
where
(8)
K?, { Teates dm = Te; dm,
One can notice that the KP, and KC), share similar structure as simpler boundary marginal conditions
Krto and Key- hence we can get rid of the notorious geometric averaging (see $4 in [1]). Notably, this
type of constraint provides an opportunity to utilize Proposition 3.1 1. for optimization, but the joint
distribution of and is still absent. We classify the constraints into two categories:
houunaayy =
n KC,, K, /Vr E {0, Khridge = (n
By following BI ($2.2), we execute optimization w.r.t. (6) while projecting the solution to subset of Kooundary or Karidge iteratively. The sketch can be found in Fig.2. The next sections will provide more details on obtaining the joint distribution μ and optimizing within each constraint set.
Hereafter, we only demonstrate the optimization for forward policy 2t given reference path measure driven by fixed backward policy 2,. The procedure be applied for the 2, and vice versa.
4.2 Optimization in set boundary
We first show how optimize forward policy z, w.r.t. objective function (6) given the reference path measure मे driven by fixed backward policy 2, under one subset of Kooundary-
Proposition 4.2 (Optimality w.r.t. Kbountar). Given the reference path measure मे driven by the backward policy 2, from boundary Pt.. in the reverse time direction, the optimal path measure in the forward time direction of the following problem
N-1
s.t E { Mester dm Meo. /
min [(T): = Σ KL (
[u.dv:: Pts

---

Pt. Teaten
is
Techeo Tudtrs dm dVe;
When Wedtaa = 16 the
the following equations need to hold Vt E 4114:
| + 2 - VV log =0 Pt. (V, Xt,) qr, (Ve, (xt,);
(9a)
(9b)
where Pt and di denote the marginal density and conditional velocity distribution of the reference path measure at time t, respectively.
Proof. See appendix.B.5
Remark 4.3. When the ground truth distributions of velocity Yr, are available, one can simply sample from "/t since the joint distribution je available in this case. In order to matching the reference path measure in KI. divergence sense, one needs to match both the intermediate path measure eq.9a and the boundary condition eq.9b. In the traditional two-boundary SB case, matching the boundary condition is often disregarded due to either having predefined data distribution or tractable prior. However, our specific case, as the velocity not predefined, it becomes imperative to address this issue and optimize it through the application of Langevin dynamics.
4.3 Optimization in set Rriddee
The formulation of optimization under Kiriigge is similar to the previous section but differs by the boundary condition 10b):
Proposition 4.4 (Optimality w.r.t. Kondge), Given the reference path measure ß driven by the backward policy 2, from boundary Hıy in the reverse time direction, the optimal path measure the forward time direction of the following problem
min [(T)
E KL (Teeste (Festees), s.1. Kbridge = (KT)))
qtg Ttgtxx TigaN F d dV
is:
when Tigity = boaty the following equations need to hold vI E [lo, IN]:
|| 2, vV log 1e Pto (Vea, Xt,) = 9r, (Veg Xtn)
(10a)
(10b)
Proof. See appendix.B.6
Conceptually, the above optimization objective with Koridge constraint aims at finding continuous path measure close to reference path measure म while any intermediate marginals constraints will not be considered. The boundary condition of reference path measure in the next iteration Pto (Vto Xt, is determined by eq. 10b. Fortunately, the empirical samples from this distribution are available, though the analytic representation of the distribution ito (Vto X1o) is unknown. Hence we can utilize these samples as empirical sources from boundary distribution gr, (V,,, for the next BP. For further explanation and intuition, one can find Appendix.
4.4 Parameterization and Training Objective Function
Inspired by the success of prior work [14], we parameterize path measure a by forward policy zi or backward policy 2, combined with one of constraints in Kroundary or Khnidge (see Fig.8 in Appendix for visualization). We adopt -aler-Maruyana discretization and denote the timestep 8, Notably, eq.9b and eq. 10b can be implied by minimizing phase space NLL in Prop.3.1. This leads to the following objective function, termed as phase space mean matching objective, which will be used to train neural networks that represent z, and 2, after time discretization:
LMM = E [|6,22 (m(+6,) 5,2 (me+6,) (m, + 5z22 me+6.) [12]
The velocity boundary condition for the reference path measure in the succeeding BP is encoded in eq.9b or eq. 10b, but the representation of conditional distribution eq.9b not clear. We leverage the favorable property of SB to parameterize and sample from such distribution.
7

---

Eareetrttitt (h)
2 [cconttrant p))
40fonta p)
Figure 4: Validation of our DMSB model on complex GMM synthetic dataset. The velocity and position of the same sample correspond to the same shade level. Upper: Samples' evolution in the position space. Bottom: Learnt samples' evolution in the velocity space.
Proposition 4.5 (127, 28]). T and N shares same path measure, then
9.0 Pt. (V,, x q. (V x o Ri (Ve,)x);)
where:
logop .0.0 2 22 /g.
(11)
Prop.4.5 suggests that one can use Pt. (Vr, xx,) == Pt. to imply condition (9b) and obtain samples from such distribution by simulating Langevin dynamics. Namely, first sample position from
ground truth X₁, ~ Pt., and then sample Vt, ~ p using eq. 1. One can further adopt the same regularization [29] to enforce the condition of Prop.4.5.
4.5 Training Scheme
Here introduce the scheme to traverse BI (sce Fig.2). In BI, all constraints must be iterated once. For the sake of LMM, the reference path measure should be induced by opposite direction. A single BI cannot be recursively repeated due the conflict of reference path measure direction. For example (see 2ig.2), at the end of d-th BI, स is yielded by forward policy while the first BP of d-th BI also optimizing forward policy which violates LMM. Instead, we reschedule the optimization order. Specifically, in (d + 1)-th BI, we optimize backward policy at the first BP and the last BP.
Experiments
Setups: We test DMSB on 2D synthetic datasets and real-world seRNA-seq dataset [30]. We choose state of the art algorithms MIOFlow [9] and NLSB [11] as our baselines. We tune both models to the best of our hardware capacity. We choose Sliced- Wasserstein Distance (SWD)[31) and Maximum Mean Discrepancy (M2D)[32) together with visualization as our criterion. The detailed setup of training and evaluation can be found Appendix.C.
Synthetic Datasets:
The Petal [9] and Gaussian Mixture Model (GMM) dataset are simple yet challenging, as they mimic natural dynamics arising in cellular differentiation, including bifurca- tions and merges. We compare our algorithm with MIOFlow in Fig.3. DMSB can infer trajecto- ries aligned with ground truth distribution more faithfully at timesteps when snapshots are taken. In GMM experiments (see Fig.4), we choose standard Gaussian at initial and terminal time steps while four- modal GMM and eight-modal GMM are placed at inter- mediate time steps. Besides good position trajectory, it
DMSBicus)
MiOFiow!
almost serendipity that DMSB can also learn the rea- sonable velocity trajectory without any access ground truth velocity information. This paves the way for our later velocity estimation for the RNAsc dataset. SeRNA- seq Dataset: The emergence of single-cell profiling tech- nologies has facilitated the acquisition of high-resolution single-cell data, enabling the characterization of individ- ual cells at distinct developmental states [7]. However, because the cell population is eliminated after the mea- surement, one may only gather statistical data for single samples at particular timesteps, which neither preserves any correlations over time nor provides access the ground truth trajectory. The diversity of embry- onic stem cells after development from embryoid bodies, which comprises mesoderm, endoderm.
Figure 3: Comparsion with MIOFlow and ground truth challenging petal dataset. DMSB is able to generate tra- jectories whose time marginal matches ground truth faithfully and outperforms prior work. Time is indicated by colors.

---

Table Numerical result of MMD and SWD on 100 dimensions single-cell RNA-seq dataset and results for leaving out (LO) marginals at different observation. DMSB outperforms prior work by large margin for both metrics and all lea case. See Appendix. 4 for Results over 3 seeds.
MMD 4
SWD 4
Algorithm
w/o LO
LO-t1
LO-t2
LO-t3
w/o LO
LO-t1
LO-tz
LO-t3
NLSB[[1]
0.66
0.38
0.37
0.37
0.54
0.55
0.54
0.55
MIOFlow[8]
0.23
0.23
0.90
0.23
0.35
0.49
0.72
0.50
DMSB(ours)
0.03
0.04
0.04
0.04
0.20
0.20
0.19
0.18
neuroectoderm, and neural crest in 27 days, is demonstrated by the scRNA-seq dataset. The snapshot of cells are collected between (tg: day 0 to 3, ty: day 6 to 9, t2: day 12 to 15, ty: day 18 21,ta: day 24 to 27). Snapshot data are prepossessed by the quality control [30] and then projected to feature space by principal component analysis (PCA). We inherit processed data from [8]. We validate DMSB on 5-dim and 100-dim PCA space to show superior performance on high-dimension problems compared with baselines. We further show that DMSB can estimate better velocity distribution compared with baselines when the ground truth is absent during training and testing.
We testify the performance of our model by computing MMD and SWD with full smapshots and when one of snapshots is left out (LO). We postpone the comparison of all the models on 5-d RNA space to the appendix (see Fig. 9 and Table.6) because the problem is relatively simple and all models accurate trajectory. Table.3 summarizes the average MMD and SWD between estimated marginal and ground truth over different snapshot timesteps. DMSB outperforms prior work by large margin in high (100) dimensional scenarios. The visualization (Fig.5) in PCA space further justifies the numerical result and highlights the variety and quality of the samples produced by DMSB.
infer
NLSB[10]
MIOFlow/B]
DMSB(ours)
Ground Truth
PP24
PC3
PC3
PC3
PC3
PC5
PCS
PCS
PCS
PCS
RC
Velsaly
Truss Mebcly
Figure 5: Comparison of population-level dynamics on 100-dimensional PCA space at the moment of observation for scRNA-seq data using MIOFlow, NLSB, and DMSB. We display the plot of the first principle components (PC), Baselines can only learn the trajectory's 's fundamental trend, whereas DMSB can match the target marginal along the trajectory across different dimensions. The right figure shows Kernel Density Estimation [33] of samples generated by DMSB and ground truth at lg and t4. The generated samples for all timesteps and comparison with baseline are in Appendix.F.
Interestingly, Fig.4 demonstrates that DMSB can reconstruct reasonable evolution of the velocity distribution which was not accessible to the algorithm. We further validate such property in 100-D RNAse dataset. During the training and testing, all the models do not have access to the ground truth velocity. We run the experiments of 100-D and 5-D RNAsc datasets and average the discrepancy between ground truth velocity and estimated velocity over snapshot time. The numerical values are listed in the 7 and Table.6. The plot of velocity and position can be found in Fig.9 and Fig. 10. The plot illustrates that while all models are capable learning reasonable trajectories, only DMSB has the ability to estimate plausible velocity distribution. This property holds even for

---

100-D RNA dataset (see Fig.5,11,12). This is notable, despite the velocity estimated by DMSB does not perfectly match the ground truth. because should be noted that the proposed phase space SDE and the optimality of OT are artificial and may not necessarily represent the actual RNA evolution. Moreover, as individual evolutions cannot be tracked, possibilities such as (A-+A, B B-B] versus (A (A-BB, B B->A can not be discerned, which renders exact velocity recovering almost impossible.
Conclusion and Limitations
In this paper, we propose DMSB, scalable algorithm that learns the trajectory which fits the different marginal distributions over time. We extend the mean matching objective to phase space which enables efficient mSB computing. We propose a novel training scheme fit the mean matching objective without violating BI which is the root of solving mmmSB problem. We demonstrate the superior result of DMSB compared with the existing algorithms.
main limitation this work is, the rate of convergence to the actual mmmSB has not been quantified after neural network approximations convergence of mean matching iteration, supporting its outstanding performance [14], the iteration still fails converge the actual SB [34] precisely due practical neural network estimation errors accumulating over BI. However, recent work [35] shows the convergence of SB when training error exists. In addition, DMSB cannot simulate the process with death and birth of cells which can be potentially described as unbalanced optimal transport [36].
introduced. Even though [15] theoretically analyzed the
7
Acknowledgement
This research was supported by the ARO Award W911NF2010151, and the DoD Basic Research Office SwardHQ00332110000
10

---

References
[1] Yongxin Chen, Giovanni Conforti, Tryphon T Georgiou, and Luigia Ripani. Multi-marginal schrödinger bridges. In International Conference Geometric Science of Information, pages 725-732 Springer, 2019.
[2] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176 (4):928-99)) 943, 2019.
[3] Karren D Yang and Caroline Uhler. Scalable unbalanced optimal transport using generative adversarial networks. arXiv preprint arXiv: /447, 2018.
[4] Mike Fisher, Jorge Nocedal, Yannick Trémolet, and Stephen J Wright. Data assimilation in weather forecasting: a case study in pde-constrained optimization. Optimization and Engineer- ing, 10(3):409-426, 2009.
[5] Kenneth G Manton, XiLiang Gu, and Gene R Lowrimore. Cohort changes in active life expectancy in the us elderly population: Experience from the 1982-2004 national long-term care survey. The Journals of Gerontology Series B: Psychological Sciences and Social Sciences, 63(5): S281, 2008.
[6] 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.
[7] Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal opti- mal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pages 6511 1-6528. PMLR, 2022.
[8] Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajecto- rynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pages 9526-9536. PMLR, 2020.
[9] Guillaume Huguet, Daniel Sumner Magruder, Oluwadamilola Fasina, Alexander Tong, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference. arXiv preprint 2.Kiv:2200.1492228 2022.
[10] Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, pages 6572-6588 2018.
[11] Takeshi Koshizuka and Issei Sato. Neural lagrangian schr')" odinger bridge. arXiv preprint arXiv:2204.00 04853, 2022.
[12] Lénaïc Chizat, Stephen Zhang, Matthicu Heitz, and Geoffrey Schicbinger. Trajectory inference via mean-field langevin in path space. arXiv preprint arXivv2205.07146. 2022.
[13] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically-damped langevin diffusion. arXiv preprint arXiv:21 2021.
[14] Tianrong Chen Guan-Horg Liu and Evangelos A Theodorou. Likelihood training of schrödinger bridge using hewaaddbackwaad sdes theory. arXiv preprint 1xXi::2110.11291 2021.
[15] Valentin De Bortoli, James Thomton, Jeremy Heng, and Amaud Doucet. Diffussin schrödinger bridge with applications to score-based generative modeling. arXiv preprint a7X1v72106.01357. 2021.
[16] Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Amaud Doucet. Conditional simulation using diffusion schrödinger bridges. Uncertainty in Artificial Intelligence, pages 1792-1802. PMLR, 2022.

---

[17] Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. 2 sb: Image-to-imad schr)." odinger bridge. arXiv preprint arXiv:2302.05872 2023.
[18] Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos A Theodorou. Deep generalized schr\," odinger bridge. arXiv preprint vrXiv:2209.09893, 2022.
[19] James Thornton, Michael Hutchinson, Emile Mathieu, Valentin De Bortoli, Yee Whye Teh, and Arnaud Doucet. Riemannian diffusion schr\." odinger bridge. arXiv preprint arXiv:2207.03024 2022.
[20] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 1. 2020.
[21] Lev M Bregman. The relaxation method of finding the common point of convex sets and its application to the solution problems in convex programming. USSR computational mathematics and mathematical physics, 7(3):200-21,, 1967.
[22] Michele Pavon and Anton Wakolbinger. On free energy, stochastic control, and schrödinger processes. In Modeling, Estimation and Control of Systems with Uncertainty, pages 334-348. Springer, 1991.
[23] Francisco Vargas. Machine-learning approaches for the empirical schrödinger bridge problem. Technical report, University of Cambridge, Computer Laboratory, 2021.
[24] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Itera- tive bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111AA 133, 2015.
[25] Paolo Dai Pra. A stochastic control approach reciprocal diffusion processes. Applied mathematics and Optimization, 23(1):313-322 1991.
[26] Jean-David Benamou, Thomas O Gallouët, and François-Xavier Vialard. Second-order models for optimal transport and cubic splines on the wasserstein space. Foundations of Computational Mathematics, 19(5): 1133:111311143, 13-1143, 2019.
[27] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-322 1982.
[28] Edward Nelson. Dynamical theories of Brownian motion, volume 106. Princeton university press, 2020.
[29] Hung-Yu Tseng, Lu Jiang, Ce Liu, Ming-Hsuan Yang, and Weilong Yang. Regularizing generative adversarial networks under limited data. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7921-7931, 2021.
[30] Kevin R Moon, David van Dijk, Zheng Wang. Scott Gigante, Daniel Burkhardt, William Chen, Kristina Yim, Antonia van den Elzen, Matthew Him, Ronald R Coifman, et al. Visual- izing structure and transitions in high nigh-dimensional biological data. Nature biotechnology, 37 112:114224492,,01 2019.
[31] Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasser- stein barycenters of measures. Journal of Mathematical Imaging and Vision, 51:22 2015.
[32] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
[33] Murray Rosenblatt. Remarks on some nonparametric estimates of density function. The annals of mathematical statistics, pages 832 837, 1956.
[34] David Lopes Fernandes, Francisco Vargas, Carl Henrik Ek, and Neill DF Campbell. Shooting schrödinger's 's cat. In Fourth Symposium Advances in Approximate Bayesian Inference, 2021.
12

---

[35] Yu Chen, Wei Deng, Shikai Fang, Fengpei Li, Nicole Tianjiao Yang. Yikai Zhang, Kashif Rasul, Shandian Zhe, Anderson Schneider, and Yuriy Nevmyvaka. Provably convergent schr)," odinger bridge with applications to probabilistic time series imputation. arXiv preprint arXiv:2305.0724f: 2023.
[36] Yongxin Chen, Tryphon Georgiou, and Michele Pavon. The most likely evolution of diffusing and vanishing particles: Schrodinger bridges with unbalanced marginals. SIAM Journal on Control and Optimization, 60(4):2016-200 20029,,,,,,,,,,,. 2022.
[37] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. arXiv e-prinis, pages arXiv-2101, 2021.
[38] Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and HJB equations, volume 43. Springer Science & Business Media, 1999.
[39] Kenneth Caluya and Abhishek Halder. Wasserstein proximal algorithms for the schrüdinger bridge problem: Density control with nonlinear drift. IEEE Transactions on Automatic Control, 2021.
[40] Ioannis Exarchos and Evangelos A Theodorou. Stochastic optimal control via forward and backward stochastic differential equations and importance sampling. Automatica, 87:159-165, 2018.
[41] Tianrong Chen, Ziyi Wang, Ioannis Exarchos, and Evangelos Theodorou. Large-scale multi- agent deep fbsdes. In International Conference on Machine Learning, pages 1740-1748. PMIR, 2021.
[42] Tianrong Chen, Ziyi Wang, and Evangelos Theodorou. Deep graphic fbsdes for opinion dynamics stochastic control. In 2022 IEEE 61st Conference on Decision and Control (CDC), pages 4652-4659. IEFE, 2022.
[43] Hopf Eberhard. The partial differential equation ut+ uux= μxx. Communications on Pure and Applied Mathematics, 3(3):201-25) -230,, 1950.
[44] Julian D Cole. On quasi-linear parabolic equation occurring in aerodynamics. Quarterly of applied mathematics, 9(3):225-25 236, 1951.
[45] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
[46] Francisco Vargas, Pierre Thodoroff. Neil D Lawrence, and Austen Lamacraft. Solving schrödinger bridges via maximum likelihood. arXiv preprint 100.0001,,0000
13

---

A Appendix
B Proof in $3 and $4
Before stating our proofs, we provide the assumptions used throughout the paper. These assumptions are adopted from stochastic analysis for SGM [27, 37, 38], SB [39]. and FBSDE [40-42].
(i) Pt, with finite second-order moment for all EG.
(ii) f and g are continuous functions, and lg()) > 0 is uniformly lower-bounded w.r.t. t.
(iii) A1 [0, have V Pe 1), V. log ((sss)); V. log ()) 1 and ((. ())))) Lipschitz and at most linear growth w.r.t. x and v.
(iv) v V C...2
(v) 3 > 0: pr. (m) ((exp-1lmll as m - co. 00.
Assumptions (i) (ii) (iii) are standard conditions in stochastic analysis to ensure the existence- uniqueness of the SDEs; hence also appear in SGM analysis [37]. Assumption (iv) allows applications of Itō formula and properly defines the backward SDE in FBSDE theory. Finally, assumption (v) assures the exponential limiting behavior when performing integration by parts. w.o.l.g, we denote = [v, Ov
B.1 Proof of Proposition.3.1
The results of the Prop.3.1 is part of results of Prop.B.4 which gives the results for both forward and backward likelihood objective.
Theorem B.1. The optimization problem
min (') / | lalgudxdvdd
(12)
Ou(me) - [I(f gu]] p} + 2 Avm.
8.l
(13)
No = -(0xx)) #1 = p(T,x,v),
will induce the coupled PDEs,
(m,,)
- [[ff +2 11 + g²AA/H
(14)
Øt
2
@Φ(m.)
1 Vo v T
1 2 Avo.
(15)
Øt
2
and the optimal control of the problem is
a' = vvv
14

---

Proof. One can write the Lagrange by introducing lagrangian multiplier Φ:
C(u, a, ②) = L/¹ /
J. / J
0p
|a
Budxdvdt
dvdxdt
2
0l
/
1
+
9 Aml Vm |(f gu) mi
2
dvdxdt
J "R"XR"
06
||a
luxddddd
M
dvdxdt
xR=
2
Ot
[
OVm [(f gu) nl|
Amp
dvdxdt
uudxdddd 1.¹ x xB"
|la
dvdxdt
R xR= 2
Ot
[1
- "moT((f + gu)l / [ A
dvdxdt
JR
{
dø
|all
Vxx و A2 - gv a e dvdxdf
Ōt
By taking the minimization within the bracket, The optimal control is.
a' V By Plugging it back, the optimality of the aforementioned problem is presented as:
0p(m,)
V. [(f + g2 Vv pl gg Avp.
1 V V
1
x 2 9 L2Ar
Ot
Theorem B.2. The optimal forward and backward processes are represented as:
dm; = [s - gu!- dt + g())dw; (forward)
(16)
dm, = [f + yu dl + g(()dw, (Backward)
(17)
in which f = =1v, T Optimal control is expressed as,
u Zi
(18)
V log V
u b+ Z.
21
0
gv log V0
0
(19)
where V and V are the solution following PDEs,
0v,
1
-VV.V
0V. 1
A. - VV
(20)
0t 2
s.t v, v, v, = p(x.v.0)
V, p(x, v.
Proof. By Lemma.B.I notice that the optimal control is: g V cb.
By leveraging Hopf-Cole [43, 44] transformation, here we define
V exp (0),
T μ exp (-Ø).
15

---

Then we can have the following expressions:
V = Vx
AV = (VV)
Σ a
(exp (d) Vol
3m:
06:
(Vo);
exp (6)
0m, exp(o)
(0)
3m:
= exp (6) [| + AΦ]
V = exx (-6) (- (v) + +exp( VM
= exp(-o)(-1) ( + Vp)
A = V (vū)
Σ am:
lexp (-
(-HVV Vp)]
Σ (I [ exp(-0) Vol.
8m;
0
VΦ); VOT θ
+ +cxp(-0)
[v/l]
im,
3m:
Om,
= (-p) [M|| V T Vn Ve Aμ - HÁN
-00 Vμ4
= (-6) [M|] V 24 V + 4/ ~1
Thus, we can have.
OV = exp (0) as
Ot
Øt
exp (6) (-1 19 -
1 VxO xの v V 2 Avō)
2 A. V T
2V
0p
do
exp
(-0)
Ø1 p nexp (-x)
16
06
01l
μ dΦ
= exp (-Ø)
0t
Ot
v [l(f [[(f+gu) gu) p} + 30"Amm + 11 (110 V 302 AND
exp
(-Ø)
exp
(-0)
V. (. V + 2 Am/ + 11
~ V
+ V x + 11 9 A.
1
exp
(-0)
9 ~ V 9² 22A V +
2 -2 マ/ V v V の + 1 9² A ④
= exp (-0)
1 A
5 2 T Vx + / Av/l + こリ 1 + vv vV Vx2
exp (-Ø)
1 11V g.V 4 VV - 11 / + 1
1
- - V + an VV
2
1 2 A.
V
VT V
16

---

Then we can represent the optimal control as:
0
= Ze =
(21)
Zi
Cope-Cole
O
(22)
Vv
vV log V.
Then the solution of such mSB is characterized by the forward SDE:
dm, = f + +=f"" dt + g(t)dwe,
(23)
Due to the structure of Hopf-Cole transform, one can have
Pi SB (2.(23) L
Pt.
(24)
According to [27, [27,28], the reverse drift of such SDE (eq.23) u should admits,
n
u b.
gV, log SB
(25)
0 9 vV. log
u b=
(26)
v log V + Vv log V. 0
(27)
vV log 1
which yields The backward optimal control
b-
Ze
0 2,
0 gV log V
(28)
Thus, the optimal forward and backward process is
dm;
f + guf fo
dt + g(t)dwr
(29)
dm,
[f gu dt + g()
(30)
And vi and T satisfy following PDEs.
0v; 1
01 7 A. V.V.V
0V. 9" A. V2K.V
Ot
Lemma B.3. By specifying f [v, TT, The PDE shown 20 can be represented by following SDEs
dx dv
(0
0
- v log w/
dt +
9
dw
(31)
dy 2
|ltt + z Tdw;
(32)
dÿ
|2|2 + z 2 + V.
g2 dt + z dw;
(33)
s.t :e:eep (yo + Fo) = p(x v, 0), exp (yr + (T) = v(
(34)
Where:
y = y(x,, v,t) = log (xx,, Ve, t),
z(x,, Ve, = VV, log (xxx, vo,
y = 5(x,, vr,t) = log (X+, V2, t). Ve, vv log V( V.,t)
17

---

Proof. One can write
Ölog v
V 20
Öt
1
A.
-Vxlog Vi vv
2
1 V. v/
lo V
" 2 Tr
Olog v
V TT v
29 D. v0
Ōt
- = VT
29 TT
By applying Itô's lemma,
log
llo vp dt
+
Vx Vx v g2 log 1
90 Ar log v dt + [ log v "] gdw:
1
Öt
1 v. dt
Vx log T
g Tr
1 9d Tr
1 V
1 V V. w
Vx v
+ ا V log 11 +
d + 9 [V log w dw,
V. log 1
2 g Tr V v/ T
di + g [V log v7 dw:
V log v dt + g [V. log | dw.
Similarly, one can have,
log Th Öt
1 Vx log i T - - log v VV. log T + 9 A. log e dt + V log T-| gdw,
d log
dt +
Vx log L T
v+
29 Tr
dt
+ vx log T v + g² V,
log v V, log vه + 29² A. log v dt + g Vr log TT| dw,
Noticing:
= 22 T +
1
V
log V0 = Tr V v 1 v log 612
2
1
V. log 1/12 + Ay log v0
Following the above derivation, one can have,
1 0
d dog T-
Vx log T v
Tr
v i dt
1 7 A. log e dt + g V. log T1 dw,
V. x log T v 9" V log v V, log きे
V. log V. V. log v2 +
1 V
V. log $12 + 2 220 A, log T di + g [v. log TT| dw:
1
18

---

By defining
y = y(e,, v, - log (xx,, Ve, ()
z(xe, Ve, - gV. log (xx: Vos 1) 2(xe, ve, vV log V (x1, v2,()
y = (xx, ve,() = log V (Xe, (),
One can conclude the results.
dx
-9" V. log w/ dt 0 0 g
dw
dv
dy
||||| + z dw:
2
1
dy =
|211 + z + V. 92 dt + dw.
s.t :.teexp (yo + Vo) =
exp exp + Pr) = v.
Proposition B.4. The fog-likelihood at data point my can be expressed as
log p(mp: 0) Em, (17)) [log p(pg. T)] - 1-" Em,-(17)
1
+ 1121112 + +z,20 + Vv gāg dt
2
E (17) [logr(m pogg((mm mT: T)]
E Em
2 1z
- gV log (17) + 2, 112 Vl log (17) 2112 dt
mean matching objective
1/¹ (m)
5112: vV (lg + z, 1 dl
mean matching objective
log (m,, T) = Em, ((6)) dogp(mo. 0)]
Em
2 |2, 12t tt + _ 12/11 +2, + V. gTz dt
Em,,(10) [log papp(mm mn, 0)] -
E Em
14: 112
2 25
V log (1)) + 2, 1
2 V. log 2/11 dt
2
matching objective
x Em,~(16)
Tol [2. vv log (10) (16) + zz
matching objective
By maximizing the log-likelihood at time I. then iteratively, (2),2e) will converge to the solution of phase space SB.
19

---

Proof. from Lemma.BB3, one can have:
p(go 0)
E yo + Yol
E ly + Yrl
E
dl
+
[2
+ 2 + Vv goz dl
2
2
T
1
1
E [logg(mm) T)]
IE
2
+
1
2 2 + V glz dt
= E [log;(mm:: T))
IE
24
11
+
2
12/1
27 (gV log (1)) + 2 2e di.
E [log pl mT, T)]
E
2 2
vv log +
SB
1 +
2 11 Vv log =2:1² di
similar result can be obtained for log pm,; T).
One can notice that the likelihood objective is continuous time analog of the mean matching objective proposed [15], and iterative optimization between logp(mo. 0) and log p(mp, 7)) are the continuous analog of IPF, Hence, the convergence proof will keep valid (see Proposition 4 in [15]).
The equivalence of KL divergence optimization in IPF and likelihood optimization is widely analyzed in [14, 15, 18]. The objective function will eventually boil down to the mean matching objective shown in the above proposition.BB B.4.
Proposition B.5 (Optimality w.r.t Kobondary). Given the reference path measure driven by the policy &. from boundary Ptas in the reverse time direction, the optimal path measure the forward time direction of the following problem
min ((T) = Σ KL (Teata (1 s.l eff Teatey dmean Me, / me, Pr: }
N-1
E
1=0
Pt, Te,t44 11 Teote dm dvt."
When Weydini Fiptom the following equations need hold Vt E [4,4++)
1|z + 2. - V. log Pr. 0, Pt. (V,, xx,) qr, (Ve, (x+,);
(35a)
(35b)
where Pt and qe denote the marginal density and conditional velocity distribution of the reference path measure at time t and t., respectively.
Proof. Due to the similarity of optimization for Krounday, the close form solution of the next path measure is (see $4 in [1] for detail):
Pt: Regttaa
Teste
J Weetter dm dv.,
By denoting the transition kernel of parameterized SDE driven by backward policy 2, as g(-).), and the time range between t; and f-11 is discretized into S interval by EM discretization. Then one can
20

---

get
Pr. Rentes
I Regites dm dvs,
Pt, (Xe, )9, (ma, (me)+5)) ... M.41 6, (mi.4: -6, Me Meue (mz+41)
qt, (xt.)
Pt. (xx, )as, (Net Ve, Kti Ve Mesta, +de Ve.tb.) 6.) U141 8e (M 11111116. -8. M 1 (m...))
qt, (xt, 19,44, (Xt,-62) V1,-6.)
Pt. (Xi, )q,, (Xi, Vu, Ko, Ve,+6,) ... 96.11 (m4411
me M4 (m,4.)
qu, (X1, 1914,, (Xt;+6,) Ve.+6,)
Pt. (xx,)) (Vess Xo Vest (x.))
11,111111 10.11166111.. (M
m. Me (man1)
4444, (m)+
44446, (me.+6, (2,6))))
(m
M Mete (ma
Pe,, (x,)), (Ve,, Necter Ve,+6, xx,
(m,++6,)
Miss, (me,+d, m2).+2)))
(me
M Meter (mans)
(Xr,)vv) Xt, ( m
(me++6,)
(x.))v)(( Xt, (g.), (m.)) Mere, (Me.+25, (me,,
9/+11 -5s (M
&
Maa Mete (ma..)
ah-n, (1,,+26))
(36)
Doing eq.36 revursively
S-1 -x,,(xx))yy, (V(e, Xx, +6 mm))) II
do (( mm)
8-1
-Pr,(Xx)g(V+, Xx,) II
(m ((+12))) (
a=0
According to [15], given the policy 2, the transition kernel qs ( ( can be esti- mated by 2, (see Proposition 3 in [15]]and it can be treated as the label for the forward policy 2 for construct following objective function for
all Thus, if Nestee is aligned with Thgt 1 then policy z₁:
C Σ m + 8.Z, (m, m.+5. + 8, (me+5.) - (mm, + 8e 8.Ze++. (m)) ))
(37)
3
> 66 (m,) + 8 Z.44, (m:) (m,+6, m. 82 Z.Ze4 (m{+6,))
(38)
~ Σ Z, (m.) Z,46, (m:)= - Vv log (17) 112
(39)
due to the special structure of Z, and Z,
(40)
Σ I. (m:)) + 2c+5, (m.) V. log Pt
(17)
(41)
Where 2. corresponds FE. BR, and B&+1 in [15] respectively. Furthermore, we need find a density function Pt, (Ve, xx,) which satisfies
Pt. (x,, )r, (vi, (x,,) Pt. Pt, (x,, (x..))..) Pt, (v,, ge, (vr,
to be the new boundary condition.
21

---

Proposition B.6 (Optimality w.r.. Krndgg). Given the reference path measure π driven by the policy 2, from boundary Pex in the reverse time direction, the optimal path measure the forward time direction of the following problem
N-1
min [=
Σ
KL (Te,sis1 (estee)) s.t πT E Korridee = 1n
2=0
is:
Tigts
Ito Tio:NN
J T d dVto
when SigtN toty the following equations need to hold Vt E [to, tyl:
I|z + Z, vV log 1e 1 =0 Pto (V Xta) = dto (Vea> Xta)
(42a)
(42b)
Proof. Same proof as B.5.
Remark B.7. The optimizer of such problem can be represented
π Nex R.y
(43)
which can also be represented as,
x² x== adpex (Fe/)e))
(44)
Where the notation R represents for the time reversal. The Proposition.4.4 is basically using neural network Z to approximate eq.44.
C Experiment Details
We test DMSB on 2D synthetic datasets and realworld ScRNA-seq dataset [30]. We parameterize z(t,, m; 0) and ((, m; Φ) with residual-based networks for all datasets (see.fig.6) 6)). The network adopts position encoding and is trained with Adam [da] on one Nvidia 3090 Ti GPU. We use constant g(t) for simplicity though the framework can adopt time varying function g(t). We the time horizon T = LN = 1 N and interval 6g = 0.01. We use EM discretization throughout the whole paper. For seRNA-seq dataset, we split data into train and test subsets(85% and 15%).All the experiment results are simulated by all-step push forward from initial data points at time t = to-
MIOFlow and NLSB setup: We use the official implementation of NLSB and MIOFlow.For MIOFlow, we report the best performance for all experiments w/GAE(or AE) and w/o GAE(or AE) embedding. For NLSB, we enlarge the size of the neural network to the best of our GPU capacity for a n00-dimensional scRNA- seq dataset and report the best performance during the training.
We evaluate the velocity of NLSB, as SDE model, by its estimated drift term at time steps == (1, 3, 4, 5). Because MIOflow w/ GAE simulates trajectories the latent space, we estimate the velocity by using the forward finite difference technique with discretization 1E 3 see after mapping from the latent code to the original space. We run the experiments of 100-D and 5-D RNAsc datasets and average the discrepancy between ground truth velocity and estimated velocity over snapshot time. The numerical values are listed in the 7 and 6able.6. The plot of velocity and position can be found in Fig.9 and Fig.10. 10. We do not want to underestinate any prior work and tried out best to tune the prior work. Feel free to communicate with the first author if one can reproduce better results in the experiment section, and we are willing to update it.
Metrics and Evaluations The 1- W-Ssererttein Distance suffers from the curse of dimensionality seriously. In the main paper, we are using Sliced- Wasserstein Distance (SWD) and Maximum Mean Distance as our criterion for 100-dim RNA dataset. An example is listed in the following toy code. One can notice that W1 suffers from the curse of dimensionality seriously, the distance between two gaussian samples is even larger than the distance between gaussian and zeros (See following code snapshot). Hence such metric is not suitable for high dimension (2 100) dataset evaluation even though some papers report W1. In order to better evaluate our model compared with baselines, are using W1, Energy Distance, Max-sliced Wasserstein distance, Sliced- Wasserstein Distance and MMD. Our metric is adapted from Geoloss (W, and Energy). POT (Sliced Wassersetein and Maximun-Sliced Wasserstein) and this repo (MMD).
Trajectories Cache Similar to prior work [14, 15), we also need to cache the trajectories for training purposes. We cache 4096 trajectories for each Bregman Projection.
22

---

Special Clarification for NLSB
We evaluate the velocity of NLSB, as an SDE model, by its estimated drift term time steps t= (1, 3. 4). It may not be reasonable to consider the drift tern as the real velocity, but the drift term can certainly depict a trend of SDE, so we still provide the result here.
from ot ot. sliced import sliced_ aasserstein_distanne
a=torch. randn (1000, 100)
#1000 gaussian samples with dimension 100
b-torch zeros (1000, 100)
#1000 zeros samples with dimension 100
c-torch. (1000, 100) #1000 gaussian samples with dimension
('SWD distance between and is: O' format (Loss (a, b)))
print ('SWD distance between and is: O formaormattaae (Loss (a, c)))
#SWD distance between and is: 1.0433608293333926
#SWD distance between and is: 0.11096614599227900
Listing 1: Distance compute by SWD distance with 1000 samples and 100 dimensions.
from geomloss import Samplesloss
a=torch. randn (1000, 100) # #1000 gaussian samples with dimension 100
b-torch. (1000 100) #1000 zeros samples with dimension 100
(1000, 100) #1000 gaussian samples with dimension 100 .p=1)
Doss-Saapleslossssssttese
print distance between and is: ff'fformat (Loss (a,b)))
print('Wr distance between and с is: 0 format (Loss ())
#W1 distance between and b is: 9.781816339892578
#N1 distance between and is: 11. 71.7304002145999
Listing 2: Distance compute by W1 distance with 1000 samples and 100 dimensions.
Training:W: We use Exponential Moving Average (EMA) with a decay rate of 0.999. Table.7 details the hyperparameters used for each dataset. The learning rate for all the datasets is set to be 2e-4 and the training batching size is 256. For computation efficiency, we cache large batch size of empirical samples from reference trajectory and sample training batch size from cache data. The hyperparameters can be found in Table.7.
m/
ResBlock
ResBlock
layers
Output
Position
layers
encoding
Figure Neural network architecture for all experiments. The network size parameters) are varying between different tasks.
Figure 7: Training Hyper-parameters
Dataset
time steps
# BI
g(t)
Parameters
SNR
v, Langevin
Semicircle
15
2000
0.2
1.21M
3
0.15
Petal
30
2000
0.2
1.21M
2
0.15
GMM
15
2000
0.2
1.21M
0.15
scRNA (100 dim)
15
4000
0.4
1.34M
0.15
Langevin sampling:Thee Langevin sampling procedure for the velocity is summarized in Given some pre-defined signal-to-noise ratio (we set snr =0.15 15 for all experiments), the Langevin noise scale at each time step and each corrector step is computed by
0, |z([, m + E(t m,
(45)
23

---

D Algorithms
Algorithm Sampling Procedure DMSB
Input: Policies z()-;0) and ()) Total sampling step S
Data distributions
Initializing
4,
velocity distrihutions 7/+, = (0, I) if they are not avaliable. for = 0 to - do
if s==0 then Sample position data Xtg from Pto-
if ground truth velocity distribution Tho avaliable then
Sample velocity data Vt, from 7%o
else
Sample velocity data Vt. by Langevin simulation conditioning Xt, (Algorrtthhm)) end if
m.g = [xto viol"
end if
Simulating dynamics: dm: [f,,) g(t) di + s(t)dwr(eq.16)
end for
return mce(to.ta)
Algorithm 2 Langevin Sampler at f, marginal constraint
policies u(,, ( 0) and ()). Previous timestep predicted velocity va.- Sample position from ground truth Xs; Pt.
Input:
for step = 0 to Langevin steps do Sample NO, I).
Construct new m.; = [xt, vc.]
Compute V. log P. zs [z(1s, m.;) + 1/9..
Compute with (45).
Langevin Sampling Vt, t- Vt, ot. Vv log V20ge. end for
return m.; [x1,, v
24

---

Algorithm 3 DMSB Training
Input: + Marginal position distribution Pt, i E [0, Parametrized policies z(... ;0) and z(.... : The number of Bregman Iteration B. Initialize postion and velocity at time step t; m. == Nome for the first iteration.
if Use ground truth velocity then
set prior velocity: Ty: = 74 else
set initial velocity = (0
end if
for 0
do
N to z² Opt.SubSee((e tx- Zref = n.; Zopt = =2", m None) [Optimize Khooanday
end for
for k
0 to N do OptSubSet(Le, 11+1, Zref = 2 Zopp 22, 11 0, None) [Optimize Keoundury
z
end
z", m = Opt.SubSet(tx, tor Zref = z², Zopt z2, Φ m m) [Optimize Kansaee
for k 0 to N 1 do
xz" 0. z OptSubSet((e. 2ref 2 Khopt
m [Optimize Keoundary]
end
for N to 1 do z² OptSubSet(ta, tu Bref = z², Zopt z², c m None) [Optimize Khoundday)
end for
z",
OptSubSet(to. tx Zref Zupt z"
[Optimize Kandge]
Algorithm 4 Function OptSubSet (Optimization for subsets)
imput: Initial time 4; and terminal time 15. Reference path measure boundary condition Pt., Reference path measure driver z..,. Policy being optimized zopi and corresponding parameter 7). Empirical sample form last iteration
output: Zopt sampless m., from reference path measure. if in is None then
Sample position data xt, from Pt.-
if velocity distribution "e, avaliable then
Sample conditional velocity data Vt, from Tt.
clse
Sample velocity data Vi, by Langevin simulation conditioning (see Algorithm.2.) end if
m: [Xe, Ve.
m, =m
end if
Sample trajectory meelt;.ty) from m. using Zr-f
Compute L = aLCMM + (1 arCrg
(Regularization of SB Lrvy [29] is optional)
update
25

---

E
Additional Diagram
Reference path A,s
H
122
Optimized forward path x Optimized backward path x
1
Hia
M
Current boundary condition
12
io
Hi.
Boundary condition of Next BP
Mi
H
Hr.
Optimization w.r.t boundary
11
P
1/
Optimization w.r. bridge
Figure 8: The detailed example diagram of Fig.2. We demonstrate an example of marginals case. The training scheme can be extended general marginals easily. The figure consists of BIs that differs by the training order. Given the reference path measure, we first run the Bregman Projection (BP) within the subset of Kboundary sequentially and end up with the constraint Kordgge
F
Additional Experiment
Table 4: Our algorithm results over seeds. Numerical result of MMD and SWD on 100 dimensions single-cell RNA-seq dataset and results for leaving out marginals at different observation. DMSB outperforms prior work by a large margin for both metrics and all leave-out case.
LO
Metries
fg
t3
14
Avg
w/o MMDI
0.021+1E-3
0.02915E-3
0.038=2E-3
0.034=2E-3
0.032=3E-3
SWD1
0.114:15B-2
0.155:22-22
0.19::3E-2
0.155 11-2
0.16:::2B-2
w/ LO-t
MMD
0.09+1E-3
0.019=1E-2
0.032-2E-2
0.029=2E-2
0.042=2E-2
SWDL
0.140:2E-2 0.135#1E-2
0.1912E-2
0.155=1E-2
0.13313E-2
w/
LO-t2
MMD
0.021+11-3 0.06515E-3
0.03212E-3
0.02122-3
0.033=3E-3
SWDL
0.100455-2 0.20242E-2
0.13#3E-2
0.191=1E-2
0.155=2E-2
W/ LO-13
MMD[
0.025::2H-3
0.026+2H-2
0.075#1H-2
0.029=2E-2
0.040=2H-2
SWDL
0.12412B-2
0.14
0.27+2D-2
0.18+1B-2
0.179-3B-2
Table 5: Numerical result of Wasserstein-1 (Wi), MMD, energy distance and Max-sliced Wasserstein distance (MWD) on position of 5 dimensions single-cell RNA-seq dataset using 500 generative samples and 500 ground truth data.
Dim=5
Energy 1
MMD
W1 4
SWDJ
MWD
NLSB MIOFLOW
0.04
0.10
0.74
0.24
0.48
0.09
0.28
0.79
0.388
0.66
(MSs(
0.03
0.06
0.67
0.22
0.41
Table 6: Numerical result of Wasserstein-11 (W1), MMD, energy distance and Max-sliced Wasserstein distance (MWD) on the velocity of dimensions single-cell RNA-seq dataset using 500 generative samples and 500 ground truth data.
Dim=5
Energy 1
MMD
W, 4
SWDJ
MWD 1
N
0.44
1.37
1.75
0.83
1.40
MIOFLOW
0.68
2.11
1.88
0.94
1.54
(Mrss)urs)
0.40
0.85
1.67
0.74
1.43
'See special clarification (Appendix.C) for the velocity generated NLSB
26

---

Table 7: Numerical result of Wasserstein-1 (Wi), MMD, energy distance and Max-sliced Wasserstein distance (MWD) on the velocity of 100 dimensions single-cell RNA-seq dataset using 500 generative samples and 500 ground truth data.
Dim=100
Energy
MMD J
SWD4
MWD І
NLSB²
2.12
1.6
0.94
1.27
MIOFLOW
2,41
1.89
5.66
DMSB(ours)
0.36
0.18
0.39
0.78
NL.SB
MIOFlow
DMSB
Ground Truth
PC1
PC1
PC1
PC1
PC3
PC3
PC3
Figure 9: Comparison of population-level dynamics on 5-dimensional PCA space the moment of observation for scRNA-seq data using MIOFlow, NLSB, and DMSB. We display the plot of the first 4 principle components (PC). All method performs well under this experiment setup.
NLSB
MIOFlow
DMSB
Ground Truth
PCI
PC1
PCI
PC1
PC3
PC3
PC3
PC3
Figure 10: Comparison of estimated velocity on 5-dimensional PCA space at the moment of obser- vation for scRNA-seq data using MIOFlow, NLSB, and DMSB. We display the plot of the first principle components (PC). For the results of NLSB, see special clarification of NLSB in Appendix.C C
"See special clarification (Appendix.C) C) for the velocity generated NLSB
27

---

NLSB
MIOFlow Flow
DMSB
Ground Truth
PC2
PC1
PCT
P.I
PCI
PC4
PC3
PC3
PC6
PC5
PC5
PCS
PCS
Figure 11: Comparison of estimated velocity 100-dimensional PCA space the moment of observation for scRNA-seq data using MIOFlow, NLSB, and DMSB. We display the plot of the first principle components (PC). For the results of NLSB, see special clarification of NLSB in Appendix. C
groand trath
preskct
PCJ
PCI
PCI
PCI
PCI
PCO
PCO
PCO
PCO
PCO
PC3
PC3
PC3
PC3
PC3
PC2
PC2
PC2
PC2
PC2
PC5
PC5
PCS
PCS
PCS
PC4
PC4
PC4
PC4
PC4
Figure 12: Comparison of estimated velocity 100-dimensional PCA space the moment of observation for scRNA-seq data using DMSB with ground truth. We display the plot of the first 6 principle components (PC).
G Intuitions of Propositions and Theorems
1. Remark for Proposition 3.1: Within each half-bridge IPF, the variable Z,(or 2,) essentially learning the reverse-time stochastic process induced by 2, This process can also be viewed minimizing the approximated parameterized negative log-likelihood.
28

---

2. Remark for Proposition 4.2: In order to matching the reference path measure in KL diver- gence sense. one need match both the intermediate path measure eq.(9a) and the boundary condition eq. eq.(9b). In the traditional two boundary SB case, matching the boundary condition is often disregarded due to either having a predefined data distribution or tractable prior. However, in our specific case, as the velocity is not predefined, it becomes imperative to address this issue and optimize through the application of Langevin dynamics.
3. Remark for Proposition 4.3: Following the same argument the remark Proposition 4.2, it becomes evident that, in this particular scenario, there is no need to account for the data distribution since there are no position constraints when optimizing with K. Consequently, the optimal solution will inherently align faithfully with the reverse diffusion and adapt to the boundary conditions imposed by the reference path measure.
4. Remark for Proposition 4.5: We indeed underexplained important nontrivial fact (thank you so much for catching it): the unique structure of SB leads to a beautiful fact that the score is proportional to the sum of the forward and backward drift terms. This facilitates the sampling of velocity. Specifically. the score function can be obtained using eq(11), I), as supported by the findings in eq(24). can also be understood the one realization of Nelson duality (see Lemma 1 in [46] and [28]).
H Complexity
Here we provide the complexity of our algorithm.
Table 8: Time complexity w.r. Dimensionality ( (Marginals=5)
# dimensions
5
100
Train Sampling
24min
25min
33min
44min
1sec
1.6sec
2.0 sec
2.02sec
Table 9: Time complexity w.r.t number of marginals (Dim=100)
Marginal
2
3
5
Train
32min
25min
33min
44min
Sampling
2.02sec
1.6sec
2.0
2.02sec
29


Reviews:
[{"review_id": "nlfv2nmkcR", "reviewer": "Reviewer_17Zf", "summary": "The authors extend the diffusion Schrodinger bridge methodology to mult-marginal setting whereby each marginal is ordered sequentially, in addition, introducing  momentum into the diffusion Schrodinger bridge framework. The method shows excellent performance on trajectory inference tasks.", "questions": "- Why use mean matching of [1] rather than divergence based training advocated in [2]? I agree they are theoretically equivalent up to an integration by parts but practically different. Similarly, the authors state the method is based on [2] (e.g. table 1 and line 116) but then if using the training objective given by [1] is it not more accurate to say the method is based on [1]?\n- What is the complexity and training time for this method? It appears that one must perform multiple reverse time diffusions per IPF step.\n\n- [1] Bortoli et al https://arxiv.org/abs/2106.01357\n- [2] Chen et al https://arxiv.org/abs/2110.11291", "limitations": "- The authors discuss lack of theoretical convergence results\n- Computational complexity and time to train is not discussed", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The proposed work fills a gap in the literature, and the authors demonstrate how diffusion model based Schrodinger bridges [1] and implementation of [2] can feasibly be plugged into the multi-marginal setting [3]. \n- The method shows excellent performance on trajectory inference tasks, outperforming other methods. The approach feels quite natural.\n\n[1] Bortoli et al https://arxiv.org/abs/2106.01357 \\\n[2] Chen et al https://arxiv.org/abs/2110.11291 \\\n[3] Chen, Multi-marginal Schrodinger bridges, 2019", "weaknesses": "- The computational complexity of the algorithm is not discussed. It appears multiple (2N reverse diffusions?) must be computed per IPF step. This seems quite costly. I would appreciate discussion and more transparency on this.\n\n- It is very difficult to clearly see what is the main algorithm and training procedure from the main text. I see algorithms are provided in the appendices but these are also not very well detailed in my opinion. The paper would improve in clarity and reproducibility if this was given more attention. This is especially true as no code is provided. Indeed, some training details like caching trajectories are mentioned in passing (as was introduced in [1], which should be cited) but it is not clear from this paper what the authors mean and without a reference it is difficult to know.\n\n- Possible error in likelihood proof / explanations\n\nIn line 454, $\\mathbb{E}[y_0]$ (and for $\\tilde{y}_0$) is decomposed into $\\mathbb{E}[y_T]$ minus the path integral of the SDE for $y$. This seems fine. However, my understanding is that unless the diffusion has fully converged to the prior (i.e. after convergence of the IPF, or with infinite regularization in SB), the term  $\\mathbb{E}[y_T]=\\mathbb{E}[\\log p(m_T)]$ is from the terminal point in the diffusion (which is not fixed until after convergence) hence must depend on the network parameters used to define the diffusion. This is commented on in Theorem 2 of [3]. This term is dropped in [2] and in this paper, which makes the likelihood interpretation incorrect for training each step of the IPF.\n\nInstead further work is required, as in [3], to show that time-reversal is a bound of the likelihood and gives **approximate** likelihood training.\n\nThe training procedure will still work as it coincides with the time-reversal of a diffusion and the same training procedure as given in [1] but the likelihood explanation does not seem justified as it is currently written. \n\nIt appears other reviewers had similar concerns with [2], and also that the overall procedure really is not likelihood training but IPF with time-reversal where each time-reversal can be viewed as approximate likelihood. https://openreview.net/forum?id=nioAdKCEdXB\n\n\n-------- \n\n\nI am conflicted by this paper. Whilst overall I think it is a good contribution methodologically and believe it can be justified with extra work. I have concerns about clarity and more importantly that this likelihood interpretation is misleading and possibly incorrect.\n\nI would be happy to increase my score if the authors address this. I am happy to be corrected if I misunderstand the likelihood interpretation.\n\n- [1] Bortoli et al https://arxiv.org/abs/2106.01357\n- [2] Chen et al https://arxiv.org/abs/2110.11291\n- [3] Song et al, Maximum likelihood training of score-based diffusion models, 2021", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# To Reviewer 17Zf\nWe would like to express our sincere gratitude for your valuable feedback and comments. We truly appreciate the time and effort you invested in assessing our submission.\n\nPlease kindly see our itemized replies below in order to address the reviewer's concerns.\n#### **1. The computational complexity of the algorithm is not discussed. It appears multiple (2N reverse diffusions?) must be computed per IPF step. This seems quite costly. I would appreciate discussion and more transparency on this.**\nIt is a very vital question which we should include in the main paper. The short answer is, yes, it will have 2N reverse diffusions. However, it will only increase marginal computation cost compare with prior simulation-based work [1],[2]. The reason is ---\n\nWe discretize time into $S$ steps (where $S$ is set to 200 for the Petal experiment and 400 for other instances). The reference path measure $\\\\bar{\\\\pi}\\_{t\\_i:t\\_{i+1}}$ is simulated using $S/(N-1)$ timesteps, given that we have $N$ marginals. To facilitate the training of $z_t^{\\theta}$, we must simulate the reference path measure $\\\\bar{\\\\pi}\\_{t\\_i:t\\_{i+1}}$ for all $i \\\\in [0,N-1]$. Consequently, the total number of simulation steps remains at $S$, which is not prohibitively large. However, it is important to note that additional simulations are necessary for the velocity component. As illustrated in Figure 7 of the appendix, this incurs an overhead of (2(two networks inference)x$N$ NFE) when compared to prior work [1],[2].\n\n[1] Valentin et al. \"Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling\"\n\n[2] Chen et al. \"Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory\"\n\n#### **2. Unclear Algorithm and missing detail of caching trajectorties and source code**\nWe genuinely apologize for any confusion caused by the complex presentation of the algorithm. In the revised version, we will provide a clearer and more straightforward explanation of the algorithm to improve its comprehensibility for all readers. Additionally, we will elaborate further on specific techniques, such as the trajectory cache, which are commonly employed in prior works but may not be familiar to a broader audience.\n\nWe have provide the codebased to the ACs according to the NeurIPS2023 Rebuttal instruction. Furthermore, when the paper could be accepted, we will make the code publicly available to facilitate reproducibility and promote transparency in our research. Thank you for your understanding, and we are committed to making the necessary improvements to enhance the quality of our work.\n\n#### **3. clarity and more importantly that this likelihood interpretation is misleading and possibly incorrect.**\nYes, it is misleading and we will add a remark in a revised version to improve the clarity. Our optimization is to find a path measure that best approximates (in KL) the time reversal of the reference path measure obtained in the previous iteration. This optimization approximates likelihood training.\n    \n#### **4.Why use mean matching of [1] rather than divergence based training advocated in [2]?**\nThis work is inspired by both [1,2] and we think [1] will lead to more scalable training objective since there is no expensive auto-grad computation. We will update the Line 116 to better state the importance of [1].\n\n[1] Valentin et al. \"Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling\"\n\n[2] Chen et al. \"Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory\"\n\n#### **4.Complexity?**\nThank you for bringing this matter to our attention! We concur that it is important to provide a clear disclosure of the algorithm's complexity within the paper, so here we provide some real numbers. These reports will be integrated into the paper in the future revision. \n\nIt is noteworthy that our model demands more training time in comparison to previous approaches. However, it is worth highlighting that the superior results we achieve significantly surpass the baseline performance. Here are detailed numerical values. We record the wall time when all evaluation criteria (MWD,SWD and MMD) no longer drop (i.e. empirically convergent). In terms of sampling, we record the wall time for sampling one forward trajectory with 256 batch size.\n\nAll of the results are based on single-cell RNA dataset.\n\n##### **Complexity w.r.t Dimensionality (marginals = 5)**\nDimension | train | Sampling\n--- | --- | --- \n|5     |24 mins|1 sec\n|10     |25mins|1.6 sec\n|50     |33 mins|2.0 sec\n|100     |44 mins|2.02 sec\n##### **Complexity w.r.t Marginals (dim=100), Remark: we keep the same time steps, that leads to almost same sampling time.**\nNumber of Marginal | train | Sampling\n--- | --- | --- \n|2     |32mins| 2.02  sec\n|3     |33 mins|2.02  sec\n|4     |38 mins|2.02  sec\n|5     |44 mins|2.02  sec\n\n##### **Training Time Comparison(dim=100, marginals=5)**\nalgorithm | DMSB | NLSB  | MIOFlow |\n--- | --- | --- | --- \n|Training     |44 mins|30 mins|20mins\n|Sampling     |2.02 sec|1.6 sec|2.12 sec|\n    \n\n#### **4. The authors discuss lack of theoretical convergence results**\n\nAs stated in the conclusion section, while our performance results are promising, we acknowledge that the theoretical convergence analysis is currently lacking. However, a recent work [3] addressed the convergence proof in scenarios where the exact solution at each Bregman Iteration is not readily available. As that was almost concurrent, we hope we could consider it as a direction for future research.\n\n[3] Provably Convergent Schrödinger Bridge with Applications to Probabilistic Time Series Imputation. arXiv:2305.07247"}, {"review_id": "wXoFcUps0u", "reviewer": "Reviewer_5VcR", "summary": "The paper addresses the topic of multi-marginal trajectory inference in high dimensions using Schrödinger Bridge (SB). In particularly, the authors focus on the so-called momentum SB in phase space where the resulting trajectories in position space are smooth interpolations between the intermediate marginals. The motivation for this is that in real physical systems smooth trajectories are often more likely than abrupted changes in drift direction at intermediate marginals. After an introduction of the preliminaries of Schrödinger Bridge and Bregman iterations, the phase space formulation of SB is introduced, which this work then builds upon. By restructuring the problem formulations and introducing an efficient training scheme for the involved function approximators, the authors show how to solve the multi-marginal momentum SB problem in a computationally efficient way. This is then utilized in the experiments which comprise both a synthetic as well as a 100-dimensional real-world use-case.", "questions": "- My main suggestion to the authors would be to make the ideas behind their original propositions accessible to a wider audience by providing some intuition and explanation beyond the mathematical proofs.\n- In addition, open-sourcing the code would reduce the entry barrier for practitioners for applying the method.\n- In the abstract it reads \"In this article, we extend SB into phase space\". Is that claim really justified? I had the impression that SB in phase space was introduced in reference 20 and that the novel contribution is in the derivation of a computationally efficient way for half-bridge iterative proportional fitting via negative log-likelihood minimization.", "limitations": "Limitations of the presented approach are discussed in the \"conclusion and limitations\" section.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Originality: The authors make several important technical contributions that allow for a computationally efficient solution of the (multi-marginal) momentum SB problem in high-dimensional phase spaces. Related work and the foundations the authors build upon are discussed and cited in the manuscript. \n- Quality: All original propositions are supported with detailed proofs in the appendix (which I did not check). The experimental results show a convincing improvement over the state of the art. In the conclusion, limitations of the work are addressed.\n- Clarity: The paper has a good introduction of the work it builds on. In addition, the novel contributions are clearly stated. The submission is well written and quite accessible on pages 1-4 where the relevant prior work is explained. \n- Significance: The presented contributions are likely to enable machine learning practitioners to apply the momentum (multi-marginal) SB method to many interesting high-dimensional trajectory-inference problems. The advances over the state of the art are demonstrated convincingly on both an artificial and a real-world use-case. Further, the approach is also applicable if the marginal velocity distributions are not available, which is probably relevant in many real world use-cases where velocity information can be much harder to obtain the positional information. ", "weaknesses": "- Clarity: The novel contributions are likely difficult to understand for a non-expert reader. While proofs are presented in great detail in the appendix, a more intuitive explanation or interpretation of propositions 3.1, 4.2, 4.4, 4.5 is lacking. In addition, reproduction and/or extension of the results is probably a significant effort because the source-code is not provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# To Reviewer 5VcR\nWe deeply thank the reviewer for all the comments. The summary is accurate and the questions are interesting and helpful. \n\nPlease kindly see our itemized replies below for addressing the reviewer's concerns.\n\n#### **1. a more intuitive explanation or interpretation of propositions 3.1, 4.2, 4.4, 4.5 is lacking.**\nWe express our gratitude for the thoughtful suggestion, and we intend to incorporate intuitive explanations as well as additional remarks in the future revision. These elucidations will serve to enhance the clarity and comprehensibility of our work, contributing to its overall accessibility. Namely, we are going to add clarification as following:\n\n1. Remark for Proposition 3.1: Within each half-bridge IPF, the variable $Z_t$ (or $\\hat{Z}_t$) is essentially learning the reverse-time stochastic process induced by $\\hat{Z}_t$. This process can also be viewed as minimizing the approximated parameterized negative log-likelihood, as correctly pointed out by reviewer 17Zf.\n3. Remark for Proposition 4.2: In order to matching the reference path measure in KL divergence sense, one need to match both the intermediate path measure eq.(9a) and the boundary condition eq.(9b). In the traditional two boundary SB case, matching the boundary condition is often disregarded due to either having a predefined data distribution or a tractable prior. However, in our specific case, as the velocity is not predefined, it becomes imperative to address this issue and optimize it through the application of Langevin dynamics.\n3. Remark for Proposition 4.3: Following the same argument as the remark of Proposition 4.2, it becomes evident that, in this particular scenario, there is no need to account for the data distribution since there are no position constraints when optimizing with $K_{bridge}$. Consequently, the optimal solution will inherently align faithfully with the reverse diffusion and adapt to the boundary conditions imposed by the reference path measure.\n4. Remark for Proposition 4.5: We indeed underexplained an important nontrivial fact (thank you so much for catching it): the unique structure of SB leads to a beautiful fact that the score is propotional to the sum of the forward and backward drift terms (kindly see Lemma B.3 and Line 448 for details). This facilitates the sampling of velocity. Specifically, the score function can be obtained using eq(11), as supported by the findings in eq(24). It can also be understood as the one realization of Nelson duality (see Lemma 1 in [1] and [2]).\n\n[1] Vargas et al. \"Solving Schrodinger Bridges via Maximum Likelihood\"\n\n[2] Nelson, Edward. \"Dynamical theories of Brownian motion.\"\n\n\n#### **2. In addition, reproduction and/or extension of the results is probably a significant effort because the source-code is not provided.**\nThanks for the suggestion. We have provided the code-base to the ACs according to the NeurIPS 2023 Rebuttal Instruction. Furthermore, when the paper could be accepted, we will make the code publicly available to facilitate reproducibility and promote transparency in our research. Thank you for your understanding, and we are committed to making the necessary improvements to enhance the quality of our work.\n\n#### **3. In the abstract it reads \"In this article, we extend SB into phase space\". ,Is that claim really justified? I had the impression that SB in phase space was introduced in reference 20 and that the novel contribution is in the derivation of a computationally efficient way for half-bridge iterative proportional fitting via negative log-likelihood minimization.**\n\nThe reviewer is correct and our use of word was imprecise. We will replace it in the abstract by\n\"We extend the approach in reference [20] to operate in continuous space\". In the main text we will continue this clarification with: \"This circumvents the need for expensive space discretization which does not scale well to high dimensions. We also address the challenge of intricate geometric averaging in continuous space setup by strategically partitioning and reorganizing the constraint sets. Furthermore, we enhance the algorithm's computational efficiency by incorporating the method of half-bridge IPF.\" \nin order to state the contributions of our work more precisely.\n\nThe helpful comment is very much appreciated!"}, {"review_id": "SoJP36dE71", "reviewer": "Reviewer_A2vV", "summary": "In this paper, the authors present an algorithm, DMSB (Deep Multi-Marginal Momentum Schrödinger Bridge), to approximate solutions to an extension of the Schrödinger Bridge (SB) problem into phase space (mmmSB), where (i) marginal constraints on the position are given across time and (ii) stochasticity is only introduced on the velocity variable (which makes trajectories smoother). This framework aims at solving multi-marginal trajectory inference problems, i.e., inferring likely stochastic dynamics of particles on a time interval, given snapshots of them at certain time steps. This work proposes several contributions: (i) it extends the approach proposed by [1] (which solves a single-variable formulation of SB) to phase space (Proposition 3.1), (ii) it presents an efficient numerical scheme for the Iterative Proportional Fitting (IPF)-type algorithm presented by [2] and restated in Proposition 4.1, to solve mmmSB in practice, (iii) it shows great performance on realistic high-dimensional single-cell RNA sequencing ($d=100$) compared to previous baselines [3,4].\n\n[1] Likelihood training of schrödinger bridge using forward-backward sdes theory, Chen et al., 2021.\n\n[2] Multi-marginal Schrödinger bridges, Chen et al., 2019\n\n[3] Manifold interpolating optimal-transport flows for trajectory inference, Huguet et al., 2022.\n\n[4] Neural Lagrangian Schr$\\backslash$\" odinger bridge, Koshizuka et al., 2022. \n", "questions": "- What is the reference measure $\\xi$ used in the momentum-like setting (see Line 125 and Eq (6)) ? As far as I understand, it is the path given by the Brownian motion propagated along the velocity variable with the corresponding ODE $d x_t= v_t dt$. The authors should state it explicitly.\n- In my opinion, $\\bar{\\pi}$ is the solution of the problem given in Proposition 4.4 without any other constraint. Does the solution given in the main paper actually include the constraint on the boundary ? What is $q_{t_0}$ in Line 187 ?\n- I recommend the authors to specify the set of measures on which the optimization is performed in Propositions 4.1, 4.2 and 4.3: is it a collection of couplings or joint distributions over all the N states ?\n- As far as I understand, the constraints given by $K_{bridge}$ enable to smooth out the trajectories of the particles along the whole time interval. To prove the efficiency of their method, could the authors compare their method with the procedure where the only constraints are given by $K_{boundary}$ (which should work too if I am not wrong) ?\n- What is the definition of $\\mathcal{L}_{reg}$ (Algorithm 4) ?\n- How many Bregman iterations did you perform in your experiments ?\n\nMajor comments:\n- I recommend the authors to give the proof of Proposition B.6 in the appendix, since it is not straightforward.\n- For sake of clarity, the authors should introduce in the main paper a consistent notation to refer to couplings (eg, $\\pi_{t_1, t_2}$) and joint distributions (eg, $\\pi_{t_1:t_2}$) in order to avoid any confusion.\n\nMinor comments:\n- The authors should insist more on the (stochastic optimal control) formulation of the SB problem that they are solving in this paper. I don't think it is well highlighted in Figure 1.\n- There is inconsistency of the control variable ($a$ or $u$) in Appendix B.1.\n- I think there is a typo between $\\hat{m}$ and $\\bar{m}$ in Algorithms 3 and 4.\n- For sake of clarity, I think that the system of PDEs in Eq (20) should be displayed in the main paper.", "limitations": "Unlike [1] or [5] where experiments on Celeba dataset ($d\\geq3072$) are performed, the dimension does not exceed 100 in the experiments of this paper. In my opinion, it is crucial to present experiments with such order of dimensionality to study the scalability of DMSB.\n\n[1] Likelihood training of schrödinger bridge using forward-backward sdes theory, Chen et al., 2021.\n\n[5] Score-based generative modeling with critically-damped langevin diffusion, Dockhorn et al., 2021.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 1, "contribution": 3, "strengths": "This work provides a well-motivated and theoretically well grounded adaptation of the SB problem to the multi-marginal setting with smooth trajectories. The authors derive a solid framework and present convincing numerical experiments on real-world data.", "weaknesses": "- In my opinion, this paper is not easy to follow and some statements are provided without context or comment, which consists in the main limitation for me at this stage. Since I find the content pretty good, I will increase my score if the authors accept to be clearer in the main paper, by answering my questions and following my general comments given below.\n\n- This paper lacks a theoretical result of convergence (at least non-quantitative, in the ideal setting where the neural networks would perfectly fit the drift terms of the forward and backward SDEs). It is not clear to me how the convergence result stated in [5] extends to the current setting.\n\n[5] Iterative Bregman projections for regularized transportation problems, Benamou et al., 2015.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# To Reviewer A2vV\nWe extend our sincere gratitude to the reviewer for the valuable comments. Please kindly find below our itemized responses, presented in an effort to address each of the reviewer's concerns.\n\n#### **1. This paper lacks a theoretical result**\nAs discussed in the conclusion section, it is important to note that, despite achieving satisfactory performance, we currently lack theoretical convergence results for our approach. However, the recent work of [1] has provided valuable insights into the convergence proof even in scenarios where the exact solution at each Bregman Iteration is unavailable. As that was almost concurrent, we intend to explore and investigate this direction as part of our future research efforts.\n\n[1]  Provably Convergent Schrödinger Bridge with Applications to Probabilistic Time Series Imputation. arXiv:2305.07247\n \n\n#### **2. What is the reference measure ?**\nThank you for bringing this to our attention. Indeed, your observation is accurate. The path corresponds to the trajectory obtained by propagating the Brownian motion along the velocity variable:\n\n$dx_t=v_tdt$\n\n$dv_t=g_tdW_t$\n\nand it yields analytically available statistics which is aligned with the one proposed in [2] (see eq(12)). The results can be obtained by using the calculation in Section 6.1 of [3] after setting $g_t=\\sqrt{\\epsilon}$. We will add this and the detailed derivation in the future revision.\n\n[2] Chen et al. \"Multi-marginal Schrödinger bridges\"\n\n[3] Särkkä et al. \"Applied Stochastic Differential Equations\"\n#### **3. Is the solution of the problem given in Proposition 4.4 without any other constraint. Does the solution given in the main paper actually include the constraint on the boundary ?**\nYes, the solution of the problem in Prop 4.4 is without any boundary constraint. Specifically,\n- The solution of Proposition 4.2 in the main paper **does** include the constraint. The constraint is implied in the $\\rho_{t_i}$ in the numerator.\n- The solution of Proposition 4.4 in the main paper **does not** include the constraint. Hence, the boundry condition is implied in the $q_{t_0}$ in the numerator which is induced by the reference path measure eq(10b).\n\n\n#### **4. I recommend the authors to specify the set of measures on which the optimization is performed in Propositions 4.1, 4.2 and 4.3?**\nApologies for any confusion. As presented in the paper, we have decomposed the optimization of Proposition 4.1 into two distinct propositions, namely Proposition 4.2 and Proposition 4.3. In the case of Proposition 4.2, the optimization process involves joint distributions $\\pi_{t_i:t_{t+1}}$, considering two marginals at $t_i$ and $t_{i+1}$. Meanwhile, Proposition 4.3 entails the optimization of the joint distribution $\\pi_{t_0:t_N}$.\n\n#### **5.Could the authors compare their method with the procedure where the only constraints are given by $K_{boundary}$?** \nWe have conducted the experiments per suggestions provided in the review, and we can verify that the suggested algorithm to compare is functional, as evidenced by Figure 2 in the rebuttal PDF. However, importantly its performance is not on par with our proposed approach, particularly in terms of convergence speed and final performance.\n\n#### **6. What is the definition of  $L_{reg}$ (Algorithm 4) ?**\nSorry for the missing reference. $L_{reg}$ means the regularization mentioned in Ln 210.Specifically, it is the regularization term proposed in [4], which can potentially force the the condition of Prop.4.5 (see line 210).\n\n[4] Ki-Ung Song, \"Applying Regularized Schrödinger-Bridge-Based Stochastic Process in Generative Modeling\"\n\n#### **7. How many Bregman iterations did you perform in your experiments ?**\nWe are using 15 Bregman Iterations for all the experiments and 30 Bregman Iterations for Petal experiment.\n\n#### **8. Provide the proof of Proposition B.6?**\nAs we are unable to update the revision at this stage, we provide here a sketch of the proof and will provide a full version in a future revision:\n\nThe results can be derived similarly to the proof of Proposition B.5, but with variations beginning from line 480. The constraint in $K_{boundary}$ necessitates decomposing the joint distribution into $p_{t_i}(x_{t_i})$ and $\\hat{q}(v_{t_i}|x_{t_i})$. However, in the optimization of $K_{bridge}$, no constraints remain, implying that the global minimization of the optimization will precisely match the reverse diffusion while satisfying the boundary condition. Interestingly, despite lacking density information from the boundary condition of the reference path measure, we fortunately possess samples from it, which can be utilized to construct the reference path measure for the next iteration. To achieve this, it is imperative to ensure that the joint distribution at the boundary remains equivalent (i.e., as given by eq.(42b)).\n\n#### **9 consistent notation to refer to couplings (eg, ) and joint distributions?**\nwe apologize for any confusion caused. As per your suggestion, we have rectified the notation in all sections of the rebuttal reply. The updated version will be included in the forthcoming revision of the document.\n\n#### **10. There is inconsistency presented variable and PDE should be put in the main paper**\nThanks for the careful reading! It is indeed a typo and it is easy to fix. We will update the confusing typo and variable and display the PDE in the main paper once we can update the revision and if we do not exceed the page limit.\n\n#### **11. Image experiments?**\nRegrettably, due to time constraints and limited hardware resources (one Nvidia 3090Ti GPU), we are unable to conduct image experiments during the rebuttal phase. However, we will thoroughly consider this aspect for future investigations should the opportunity arise. Thank you for your understanding and consideration of our constraints."}, {"review_id": "c722FWdmGs", "reviewer": "Reviewer_3Whf", "summary": "The paper aims at solving efficiently in high dimensions the multi marginal momentum Schrödinger Bridge, that is Schrödinger Bridge in phase space with multiple marginal constraints. They also tackle the issue of marginals constraints where only the positions are enforced. They reach this objective by proposing a new algorithm: DMSB.\n\nThe paper first focuses on adapting the iterative proportional fitting (IPF) algorithm to the momentum Schrödinger Bridge with fully specified marginals in phase space. They show that the IPF can be seen as an alternating maximization of log-likelihood.\n\nNext they focus on the multi-marginal momentum Schrödinger Bridge. They propose a new formulation of the marginals constraints in order to apply their version of IPF as log-likelihood maximization via policy specification. Moreover they prove that the optimality conditions for the optimal bridge (under their new set of constraints) allows to sample the velocities when they aren't initially specified.\n\nIn order to train efficiently in high dimension they adapt a neural network parametrization of the policies known in the classical SB setting to the momentum SB setting.\n\nFinally they test the the performance of DMSB on synthetic data. The performance on high dimensional data is higher than the baselines algorithm. Moreover they are able to recover velocity distributions which is a novelty.\n\nThe contributions can be summarized as follows:\n- Extend the IPF to the momentum SB\n- Introduce a new set of constraints for the Bregman iterations which allows for a neural net parametrization of the policies and sampling of velocities when they are not specified.\n- Proposal and benchmark of a new algorithm DMSB made to solve momentum multi-marginal SB\n\n\n\n\n", "questions": "How long does training take for scRNA-seq using DMSB compared to MIOFlow and NLSB ?", "limitations": "The author address the limitations of the works which are mainly the proofs of the theoretical convergence of the algorithm in the neural nets parametrization case of the momentum multi marginal SB. Proofs which are key to back up the interesting practical results.\n\nThey also sharply point out the limitation which is tied to their use case which is the impossibility for mmmSB to model for the death and birth of cells.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The presented use case of their algorithm is cell profiling over time. In that regard the proposed algorithm performs well above baselines and adresses the following technical challenges: sampling of the velocities, curse of dimensionality, partial informations on the marginals.\n\nThe sampling of velocities and curse of dimensionality are both tackled by a combination of proposition 3.1 and proposition 4.2/4.3. Proposition 3.1 proposes a novel and interesting formulation of the half bridge using log-likelihood. This is combined with a novel way to deal with the marginal constraints which is a key advance towards efficient high dimensional computing with respect to prior works.\n\nGlobally the article puts swiftly together multiple well known ideas and adds key ingredient (decoupling of marginal constraints, log-likelihood IPF) in order to produce an algorithm which removes limitations from prior algorithms. Those limitations being :\n- lack of scalability\n- need for fully specified velocities\n- robustness to missing marginals\n\nFinally the benchmark contains multiple metrics which is helpful in the high dimensional context since the curse of dimensionality renders some metrics less meaningful as they point out.\n\nThe explicit description of the algorithm in the appendix is appreciated because it makes the article whole: it contains the theoretical and practical aspects.", "weaknesses": "Though inspired by a prior work the section on the neural nets parametrization lacks clarity. How does the log-likelihood minimization is tied to the mean matching objective and thus the loss Lmm? This link is succintly pointed out in the appendix B.5 equation (38).\n\nThe training scheme section goes rapidly over how the neural networks are trained. The figure 8 in appendix clarifies a lot the training. The explanation on the discretization method is unclear in section 4.5.\n\nFinally the main advantage of the algorithm DSMB is the scalability however there are no complexity analysis of the algorithm with respect to dimensions and number of marginals. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# To Reviewer 3Whf\nWe deeply thank the reviewer for all the comments. The summary is accurate and the questions are interesting and insightful. \n\nPlease kindly see our itemized replies below in order to address the reviewer's concern.\n#### **1. How does the log-likelihood minimization is tied to the mean matching objective and thus the loss Lmm?**\n\nOur understanding is that the minimization of the mean matching objective ($L_{MM}$) corresponds to optimizing the lower bound of the log-likelihood, as noted by reviewer 17Zf. This was proved in Proposition B.4.\n\n#### **2. The training scheme section goes rapidly over how the neural networks are trained. The figure 8 in appendix clarifies a lot the training. The explanation on the discretization method is unclear in section 4.5.**\nOwing to constraints on page count, Figure 8 has been placed in the appendix; however, it will be reintegrated into the main body of the paper should additional pages become available.\n\nRegarding discretization, the simplest Euler-Maruyama method has been consistently employed throughout the entirety of this manuscript, although our approach can use better discretization as well. \n\nSpecific details pertaining to total time and timesteps can be referenced in Appendix Figure 7. We will further explain it in the main paper.\n\n#### **3 Complexity (Dimension, marginal, time)**\nThank you for bringing up this matter! We concur that it is important to provide a clear disclosure of the algorithm's complexity within the paper so here we provide the empirical complexity results. These analyses will be integrated into the paper in the future revision. \n\nIt is noteworthy that our model demands more time and memory in comparison to previous approaches. However, it is worth highlighting that the superior results we achieve significantly surpass the baseline performance. Here are detailed numerical values. We record the wall time when all of evaluation criterion (MWD,SWD and MMD) does not drop empirically (empirically converges). In terms of sampling, we record the wall time for sampling one forward trajectory with 256 batch size.\n\nAll of the results are based on single-cell RNA dataset.\n##### **Complexity w.r.t Dimensionality (marginals = 5)**\nDimension | train | Sampling\n--- | --- | --- \n|5     |24 mins|1.62 sec\n|10     |25mins|1.63 sec\n|50     |33 mins|1.84 sec\n|100     |44 mins|2.01 sec\n##### **Complexity w.r.t Marginals (dim=100), Remark: we keep the same time steps, that leads to almost same sampling time.**\nNumber of Marginal | train | Sampling\n--- | --- | --- \n|2     |32mins| 2.01 sec\n|3     |33 mins|2.01 sec\n|4     |38 mins|2.01 sec\n|5     |44 mins|2.01 sec\n\n##### **Time Complexity Comparison(dim=100, marginals=5)**\nalgorithm | DMSB | NLSB  | MIOFlow |\n--- | --- | --- | --- \n|Training     |44 mins|30 mins|20mins\n|Sampling     |2.01 sec|1.6 sec|2.12 sec|"}, {"review_id": "0EALpXhbCI", "reviewer": "Reviewer_EhNn", "summary": "The paper tackles the multi-marginal Schrodinger bridge space in phase space, by proposing a computationally tractable solver (DMSB). It leverages alternating Bregman projections to adapt the iterative proportional fitting algorithm in order to deal with multiple convex constraint sets.\n\nThe DMSB framework constrains the stochasticity of particle dynamics to their velocity component and therefore ensures smooth trajectories. The authors show how the proposed algorithm better describes scRNA-seq data by reconstructing lower-variance trends in the dynamics and (approximately) recovers the velocity field of particles.\n", "questions": "**Experiments**\n\nIt is unclear why the colors used in Figure 3 vary across the rows. Could it be possible to identify the position and velocity of each group of points by using the same colors in both graphs?\n\nContrary to what was stated in the text (“it is almost serendipity that DMSB can also learn the reasonable velocity trajectory without any access to ground truth”, line 230), I don’t find the approximate recovery of velocities surprising. Given the formulation of momentum Schrödinger bridges (lines 125-126), the algorithm must find a reasonable velocity field, in order to produce meaningful trajectories.\n", "limitations": "As stated by the authors, the benefits brought by DMSB are not detectable when learning simple trajectories in relatively low-dimensional spaces. It is therefore unclear whether many practical problems would benefit from it since this comes at the expense of doing without convergence guarantees.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper builds on well-known results involving Schrödinger bridges (SBs) and convex optimization to practically solve the multi-marginal SB problem while ensuring smooth trajectories. It extends SBs [13] by considering random velocities, rather than random infinitesimal displacements. Furthermore, it generalizes the popular critically-damped Langevin diffusion (CLD) [12] framework to non-gaussian velocity marginals at the extremes.\n\nA key contribution consists in performing 2 alternate optimization stages over the two sets $k_{\\text{boundary}}$ and $k_{\\text{bridge}}$, containing appropriately-partitioned constraints. The extension of the mean-matching objective to the phase space is also novel but relatively straightforward.\n\nThe content is well-motivated and appropriately described. Both mathematical results and algorithmic solutions are accompanied by remarks and informal insight, which help to understand their purpose and significance.\nSeveral tables and diagrams further clarify the presentation by (i) precisely stating the notation used (Table 2), (ii) the problem description (Figure 1), (iii) the comparison against previous work (table 1), and the structure of the algorithm (Figure 2).\n", "weaknesses": "The rationale behind the decoupling of constraints, i.e. the requirement that marginalizing joint densities over the previous/next timesteps yield the same distribution, could be made more explicit to the reader. In particular, it would be advisable to explain why the solution proposed by the authors avoids the “geometric averaging issue” (line 156) which affects instead the solution in [20].\n\nI think that Proposition 4.5 should be followed by a more extensive discussion on the differences between constraints in Eqs. 9b and 10b: Why a conditional distribution appears in the former and a joint one in the latter?\n\n**Miscellaneous**\n\nIn addition, I point out minor inaccuracies found while reviewing this draft:\n- The quantities optimized in the formulas in the first row of Figure 1 are incorrect ($u$ and $a$ do not appear anywhere in the optimization objective).\n- Proposition 4.1 looks more like a definition (and it is not even stated as a proposition in the work cited as a source).\n- Wrong sections are cited on the right of Figure 2.\n- There are spelling mistakes in lines 34, 66, 89, and 246.\n- I would suggest refining the references to the Appendix, by clearly specifying if the link points to a proposition or a section (e.g., in line 128).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# To Reviewer EhNn\nWe express our gratitude to the reviewer for their valuable feedback. The summary provided is accurate, and the questions raised are both intriguing and perceptive.\n\nKindly find below our itemized responses, organized in order to address each of the reviewer's concerns.\n#### **1. In particular, it would be advisable to explain why the solution proposed by the authors avoids the “geometric averaging issue” (line 156) which affects instead the solution in [20].**\nIn Section 4 of reference [20] (page 7), the authors introduce an analytic solution that incorporates the geometric averaging of two densities in both the numerator and denominator. The result is based on the constrain set $K_i$ (using notation from [20]), which accounts for the boundary from two directions, leading to the inclusion of geometric averaging in the analytic solution. While it is feasible to estimate this value in the discretized space, addressing this matter in the continuous space poses a challenge, and currently, no optimal solution exists for it.\n\nOne of the key highlights of our paper is the successful avoidance of this issue. We achieve this by decomposing the $K_i$ constraints into separate $K_{boundary}$ and $K_{bridge}$ components, allowing us to address and resolve the problem effectively through our proposed approach.\n\n#### **2.  Why a conditional distribution appears in the former and a joint one in the latter?**\nOne can obtain similar results as in the proof of Proposition B.5, but the outcome differs from that presented in line 480. Due to the constraints imposed on $K_{boundary}$, it becomes necessary to decompose the joint distribution as $p_{t_i}(x_{t_i})$ and $\\hat{q}(v_{t_i}|x_{t_i})$. However, in the optimization of $K_{bridge}$, there are no remaining constraints, indicating that the global minimization of the optimization will precisely match the reverse diffusion while simultaneously meeting the boundary condition.\n\nInterestingly, even though we lack density information from the boundary condition in the reference path measure, we fortuitously possess samples from it, which can be leveraged for constructing the reference path measure for the subsequent iteration. To achieve this, it is essential to ensure that the joint distribution at the boundary is equal (i.e., eq. (42b)).\n\n#### **3 Miscellaneous**\nThanks for the careful reading! We will update in the revision.\n#### **4 It is unclear why the colors used in Figure 3 vary across the rows. Could it be possible to identify the position and velocity of each group of points by using the same colors in both graphs?** \nWe intend to utilize distinct colors to represent position and velocity in the revised version. The draft plot with the desired color scheme can be found in Figure 1 of the attached Rebuttal PDF file.\n\n#### **5 convergence guarantees?** \nRegarding theory, as discussed in the conclusion section, despite the satisfactory performance of our approach, we acknowledge that we currently lack theoretical convergence results. However, we find encouragement in the work of [1], which sheds light on the convergence proof even when the exact solution at each Bregman Iteration is not available. As [1] was almost concurrent, we hope we could consider analysis based on it as a direction for future research.\n\nRegarding empirical performance, we mostly agree with reviewer; however, even in very low dimension, our algorithm still poses marginal advantage as  evidenced by Table 5.\n\n[1]  Provably Convergent Schrödinger Bridge with Applications to Probabilistic Time Series Imputation. arXiv:2305.07247"}]

Author rebuttal:
# To All Reviewers
We thank all reviewers for their valuable comments. We are excited that the reviews identified the novelty of our contribution, appreciated our experimental validations and acknowledging the significancy of our work. 

The common criticisms rather came from insufficient complexity analysis and inacurate description for certrain concepts, about which we apologize. 

In a revised version, we will integrate the clarifications suggested by all the reviews to make the paper easier to be understood, together with additional experimental validations suggested by the reviewers. 

An itemized summary for each reviewer is listed below, which in our humble opinion resolves all raised concerns. We sincerely thank you for your time and help in all cases.

## Summary of Revision for all Reviewers
- In alignment with NeurIPS 2023's code disclosure guidelines, we have provided the Area Chair with the anonymous GitHub link to our codebase. Additionally, we have included the command line instructions necessary to replicate the results in the paper, along with an estimate of the expected waiting time.
- We will add the complexity analysis as the reviewers suggest.
- We will release the code when the paper is accepted.
- We will correct the inaccurate statement of the Likelihood as being suggested by reviewer 17Zf.
- We will update the notation for the coupling and the joint distribution. Specifically, $\pi_{t_i,t_{j}}$ will represent for the coupling of two marginals, and $\pi_{t_i:t_{j}}$ will be the joint distribution (path measure) as being suggested by reviewer A2vV.
- We will correct all the typos mentioned by reviewers. Thanks again for your careful and helpful reviews.


Meta-review: The paper deals with the problem of multi-marginal Schrödinger Bridge (SB), employing a phase space formulation, bridging some gap in the literature. 

While all reviewers concur that the approach introduced in the paper is indeed novel, interesting, and well-supported, my assessment reveals that there is room for enhancing the paper's clarity. Furthermore, the recommended changes provided by the reviewers should be meticulously incorporated into the final version of the manuscript. I strongly stress to the authors the significance of giving due attention to these aspects and ensuring they are not disregarded.
Decision: Accept (poster)