Trajectory Inference with Smooth Schrödinger Bridges
Authors: Wanli Hong
Yuliang Shi
Jonathan Niles-Weed
Venue: ICML 2025 poster (2025)
Keywords: Trajectory Inference
Schrodinger Bridge
Message Passing
Probabilisitic Graphical Models
Primary area: probabilistic_methods->gaussian_processes
TL;DR: We extend the multi-marginal Schrödinger Bridge problem by generalizing its reference process to a smooth Gaussian process, retaining computational tractability with polynomial-time complexity.

Abstract: Motivated by applications in trajectory inference and particle tracking, we introduce **Smooth Schrödinger Bridges**. Our proposal generalizes prior work by allowing the reference process in the multi-marginal Schrödinger Bridge problem to be a smooth Gaussian process, leading to more regular and interpretable trajectories in applications. Though naïvely smoothing the reference process leads to a computationally intractable problem, we identify a class of processes (including the Matérn processes) for which the resulting Smooth Schrödinger Bridge problem can be *lifted* to a simpler problem on phase space, which can be solved in polynomial time. We develop a practical approximation of this algorithm that outperforms existing methods on numerous simulated and real single-cell RNAseq datasets.

Trajectory Inference with Smooth Schrödinger Bridges
Wanli Hong 12 Yuliang Shi
Jonathan Weed 14
Abstract
Lavenant et al., 2024) builds on Schrödinger's remarkable thought experiment (Schrödinger, 1931; Schrödinger, 1932), formulating trajectory inference as a Kullback- mini- mization problem. Suppose that a set of particles is observed at times lo: lo....,4 1x E [0, 1] and that the particles' arrange- ment at time t& is represented by a probability measure pax for all & E [K]. The multi multi-marginal Schrödinger bridge corresponding to these observations is the solution of
Motivated by applications in trajectory infer- ence and particle tracking, we introduce Smooth Schrödinger Bridges. Our proposal generalizes prior work by allowing the reference process in the multi- mallinamaaiiall Schrödinger Bridge problem to be smooth Gaussian process, leading to more regular and interpretable trajectories in applica- tions. Though naïvely smoothing the reference process leads to a computationally intractable problem, we identify a class of processes (includ- ing the Matérn processes) for which the resulting Smooth Schrödinger Bridge problem can be lifted to simpler problem on phase space, which can be solved in polynomial time. We develop practical approximation of this algorithm that outperforms existing methods on numerous simulated and real single-cell RNAseq seq datasets.
min D(P(R) s.t. Pa = pkx Vk E [K], PEP(R)
(I)
where D is the Kullback-Leid the time marginal of P at time 4, and R is the law of refer- ence process, typically Brownian motion. (See Appendix A for a complete list of notation.) This formulation has an ele- gant
divergence, Ps denotes
interpretation: the Schrödinger bridge matches the observed marginal distributions while keeping partiele trajectories as faithful possible the prior R.
While extensively studied both theoretically and method- ologically, this classical Schrödinger Bridge (SB) approach has a critical limitation: its trajectories inherit the roughness of Brownian paths, leading to noisier estimators and less in- terpretable posterior paths; see Figure 1, left side. Moreover, when the experimenter aims to track the positions of indi- vidual particles in system evolving over time, inference with the Brownian motion prior fails to "borrow strength" from adjacent time points, resulting in less accurate results; see Figure 2, left side.
1. Introduction
The trajectory inference problem-reconstruutinn the paths of particles from snapshots of their evolution-is funda- mental to modern data science, with applications in single- cell biology (Huguet et al., 2022a; Saelens et al., 2019;
Schiebinger et al.. 2019; Sha et al.. 2023; Tong et al., 2020),
fluid dynamics (Brunton et al., 2020; Ouellette et al., 2006),
and astronomical object tracking (Kubica et al., 2007; Li-
The literature on trajectory inference contains various pro- posals for encouraging smooth paths, inspired by spline algorithms on Rd or the dynamics of physical systems. (A full comparison with prior work appears in Appendix B.) However, from the perspective of Schrödinger's original formulation, these modifications sacrifice the clear statis- tical interpretation of (1). We develop new approach: a smooth version of the Schrödinger Bridge problem in which the Brownian motion is replaced by a smooth Gaus- sian process. This proposal has an appealing statistical grounding: like Gaussian process regression (Rasmussen Williams, 2006), it provides a flexible and principled way to perform non- parametric estimation while incorporating prior smoothness assumptions. As Figures and 2, right side, show, the resulting estimates significantly outperform the vanilla Schrödinger Bridge, producing better estimates with smoother paths.
ounis et al., 2020). Given observations of a collection of indistinguishable particles at discrete times, the statistician aims infer the continuous trajectories that generated this data.
A leading approach (Chen et al., 2019; Chizat et al., 2022;
"Fqual contribution 'Center for Data Science, New York Uni- versity, New York, United States "Shanghai Frontiers Science Center of Artificial Intelligence and Deep Leaming New York Uni- versity Shanghai Shanghai, China 'Department of Mathematics, The University of British Columbia, Vancouver, Canada *Courant Institute of Mathematical Science, New York University, New York, United States. Correspondence to: Jonathan Niles-Weed <jnw www.eeiiw..su..........
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Smooth Schrödinger Bridges
Molio
poir 2m12
Brownian Hobion pror
Matern for wit v
38
-1.
Figure 1. Comparison between the classical SB and smooth SB with Matérn prior on a trajectory inference task for the petal dataset (Huguet et al., 2022a). Colored points are observed data and paths are inferred trajectories. Our smooth SB approach gen- crates much smoother and more concentrated trajectories.
0.
02
6.6 00 1.
Time
62
04
66
68
21
Figure 2. Comparison between the classical SB and smooth SB. The observations come from 20 independent trajectories a Matémn Gaussian process with 3.5. The pictures de- pict the posterior identification of particles obtained by solving the SB problem with two different priors. Our smooth SB approach recovers trajectories much more accurately.
Our contributions:
tween solutions to (1) and solutions
We propose a class of Gaussian processes suitable as priors for smooth SB problems, which we show can be lifted to simpler SB problems on phase space.
Poy Pa))P(RAAK-)) D(PK))/KK)) Pa MK, WK E [K]. (2)
min
Moreover, each measure P4 absolutely continuous with finite entropy and Riy) has density problem is equivalent to
then this
We show that message passing algorithm for the resulting lifted SB problem converges in time quadratic in K, an exponential improvement over the baseline approach.
K
min PI)EEPRR4K+3)) / ((w)Pk)((W) de +D
P(K]
X
Mk
k=0
We develop efficient approximation of our algo- rithm that outperforms existing methods in practice, in several cases by 2-5x.
Ps pk. Vk [K].
(3)
The reformulation Lemma 2.1 is essential because it elim- inates the need to optimize over probability measures on the deffnntediimensionall space Ω. Though (3) was derived for absolutely continuous measures, the expression is sensi- ble for arbitrary marginal measures. Following common convention the literature (see, e.g., Pooladian & Niles- Weed, 2024), we therefore take (3) as the basic definition of the multi-marginal Schrödinger bridge in what follows. Importantly, when the measures μ4 are finitely supported, the resulting optimization problem is fiite-dimensional.
2. Background
2.1. The Multi-Marginal Schrödinger Bridge
Let R be a measure the space Ω2 C((0, 1]; R)) of continuous Re-valued paths, and let jo.----, Pk be proha- bility measures on Rd Fix sequence of times to 0 < t1 ... < tx = 1. Given E Ω we write Wi = (t+) for E [K]. Given a probability measure P on Ω and E [K]. let Ps be the marginal distribution of at time 4. that is, let Ps be the element of ((Rd) obtained by pushing P forward under the map w + Wk. The multi-marging Schrödinger Bridge (1) exists under suitable moment ditions on μo...., Pk (Léonard, 2014), and the strict con- vexity of the -Libler divergence guarantees that when a solution exists, it is unique.
2.2. Sinkhorn's algorithm for multi-marginal transport
Specializing the case where each of the marginals Pi is supported finite set X's of size most n, the joint measure P(K can be represented by finite-albes expo- nentially large- order (K + 1) tensor, of total size n K+1
Though phrased as a minimization problem over the space P(2), the Schrödinger bridge problem admits a "static" reformulation as multi-marginal entropic optimal transport problem over the space (K+).. Write R[K] for the joint law of w = (wp..... wK) for R.
A direct method to solve (3) in this case uses Sinkhorn's 's cel- ebrated scaling Algorithm Sinkhorn's algorithm is based on the observation that the optimal solution Por to (3) is of the form Pirj (x) = [T) (*x) for all x = (Zo..... ....... xK) E I Xs. for some "scaling func- tions" ox X - R E [K]. Algorithm produces
Lemma 2.1. There is an one-to-one correspondence be-
2

---

Smooth Schrödinger Bridges
sequence of iterates ( k E [K]. 2 0, which converge to these optimal scalings.
smooth sample paths (see Lemma C.I). This observation suggests the pessimistic conclusion that no practical algo- rithm exists for solving Schrödinger bridges with smooth priors.
Algorithm 1 Vanilla Sinkhorn Algorithm
Input: Mo,,..,, [
However, the key contribution of this work is to show that this conclusion is false. The next section identifies a class of smooth priors for which (3) can be solved efficiently.
Initialize 8g (u) UK E (I+))* ( O.
(0)
while not converged do e 6+1 for k= = O.... K do
2.3. Autorcgressive Gaussian processes
Compute
(
using (4) with {u/²) ick
and
To develop our proposal for efficient smooth Schrödinger bridges, we focus on the class of continuous- time Gaus- sian anloregressive processes, classic model in statisties and signal processing (Phillips, 1959). They are also widespread choice in applications: as we show below (The- orem 2.3 and Proposition 2.4), the famous Matérn kernel gives rise to such processes. While it is known that such pro- cesses offer crucial efficiency benefits in Gaussian process regression (Gilboa et al., 2015), their algorithmic implica- tions for the SB problem have not been explored prior to this work.
(1-1)
(9) o
MR 2 $
end for end while Return (uk [E([K]
The key subroutine in Algorithm 1 is an operation called "marginalizzaionn:: given current iterates (#)) the marginalization along coordinate k is the function Sy : Xx + R defined by
Definition 2.2 (Rasmussen & Williams, 2006, Appendix B). Let be positive integer. A Gaussian process + w(t) defined on Cmn (R; Rd) is a Gaussian autoregressive process (GAP) of order m if is a stationary solution to the stochastic differential equation
Sp(xx) =
Σ exp(-C(x)) Z
II (,()))) (4)
The complexity of computing S, directly is exponential in K, since the sum in (4) has an exponential number of terms. Moreover, no general polynomial
dm dim ((1 w(() + am dim w()) + ..... + aaw(() = oE(l),
d
algorithms for even approximating (4) exist under standard computa- tional complexity assumptions (Altschuler & Boix-Adserå, 2023, Lemma 3.7). Together, these facts suggest that (3) will only be computationally feasible under special assumptions on R.
for some constants ao...., am-1 R and o E Rdxd, where {(t) denotes a white-noise process on Rd with in- dependent coordinates.
The key fact about GAPs, which we leverage to develop efficient algorithms, is that, even though a GAP w is typically not Markov, the process 1) := (w, dw/dt,.... "m/: taking values in phase space Rdxm is Markov process. We call ry the "lifted"" Gaussian process corresponding to w. The main observa- tion driving our practical algorithm for smooth Schrödinger bridges is that when R is a GAP, problem (3) admits an effi- cient algorithm obtained by lifting the optimization problem to phase space.
crucial observation (Altschuler & Boix-Adsers 2023; Chizat et al., 2022), which has driven the near-universal choice of Brownian motion as prior, is that when R is Markov process, (4) can be computed in Poly (K, (x)) time. The tractability of marginalization when R is a Markov process follows from the decomposition
exp(-C(x)) (x) Ro( II R (FR(ER-1), k=1
GAPs form rich class of smooth Gaussian processes. Moreover, they admit simple characterization in terms of their covariance functions. Recall that a Gaussian pro- cess w taking values in CIR; IRd) is characterized by two functions m : R + IRd and k R + daxd which satisfy
which shows that the exponential-size sum in (4) factors into K sums, each of which can be computed efficiently. The assumption that R is Markov extends even to works that consider priors other than Brownian motion (Bunne al., 2023; Vargas et al., 2021).
m(t) Bw((), k(s, t) = cov(w (s), (t),))
This raises what appears to be an inherent tension: efficient algorithms for (3) rely on the assumption that the process R is Markov; however, limiting to the case of Markov pro- cesses necessarily precludes consideration of priors with
The process is stationary if <(s,t) = k(t - s) for some k : R - ddd For notational convenience, we focus on the zero-mean case, where m = 0, though our results apply
3

---

Smooth Schrödinger Bridges
to the general case as well. Zero-mean, stationary Gaus- sian processes are characterized entirely by the covariance function k.
Section 4, we show that this reformulated problem can be solved by belief propagation algorithm in number of iter- ations that scales linearly with K. Finally, in Section 5, we develop practical approximation of the belief propagation algorithm, with overall runtime quadratie in K.
We summarize the important properties of GAPs in the following theorem.
In what follows, for notational convenience, we focus on the case d 1. (We discuss the runtime considerations assciated with larger d in Section 6.)
Theorem 2.3 (Saatçi, 2012, Section 2.2). Let w be zero- mean continuous, stationary Gaussian process on R& with covariance function Then w is a GAP of order m if and only if its spectral density (T) == 2 I_ k(t)e -ibs dt e ad is of the form
Let R be the law of mean zero GAP of order m for an integer m > 0. Theorem 2.3 guarantees that == (w, dw/dt,.... .itl.,,.,,"www//// dd""""/ddt=")) is a stationary Gauss- Markov process on R.,, whose law we denote by R. We write R for the joint law of the finite-dimensional vector r) == (To.... ................................ ........),))))))....))).)....... ) which is Gaussian with mean 0 and covariance matrix E E +1)m()))((+))
ST(T) oo - 15(??)
E
where f is polynomial of degree m.
Moreover, if w is a GAP of order m, then == (w, dt. .aa/dd....dddd///dd w/tt--) is a zero-mean, stationary Garko- process.
We first show that the smooth Schrödinger bridge prob- lem corresponding to R can be rewritten in terms of R. Suppose for concreteness that μu is supported on a finite set X3 S R. and write Xix] = oo Xx. We write Dik] :- R(m-1)(K+1), Given probability measure /' on phase space Z[K] == 2Kk] x 3/xx whose marginal dis- tribution on J/xx1 is absolutely continuous, we write p(x. y) for its density with respect to the product of the counting measure on XiK] and the Lebesgue measure on V/ We will use the variable z (x,y) to denote an element of 2|K]-
With Theorem 2.3 in hand, we easily obtain many examples of GAPs. Crucially, this includes Gaussian processes cor- responding to the Matérn kernel, the most popular smooth Gaussian process in applications, defined by the covariance function
k(t)& x o² K.,
(5)
where K2 is the modified Bessel function of the second kind and y and l are positive parameters.
We obtain the following:
Proposition 2.4 (Hartikainen & Särkkä, 2010, Section 4.1). Suppose is Gaussian process C(R; R) whose covari- ance function is Matérn kernel for smoothness parameter m 1/2, for m E N. Then is a GAP of order m.
Lemma 3.1. Assume that each of the marginals pa is sup-
ported on a finite set Nx. There is a one-to-one correspon- dence between solutions to (3) and solutions to
min D 2
1
Σ
z zp(z) H((),
By Theorem 2.3, the same holds true in the multidimen- sional case if each coordinate is taken to be independent of Matérn covariance.
MIM xEX1X
(6)
Per Mk: Vk [K],
Though the above characterization is formulated for station- ary processes, our algorithm also applies to the nonstation- ary case, for instance, to the case of integrated Brownian recovers the standard Schrödinger bridge, and = 2 gives rise to the so- called "momentum Schrödinger bridge." previously studied in (Chen, Conforti, Georgiou, and Ripani, 2019; Chen, Liu, Tao, and Theodorou, 2023). The integrated Brownian mo- tion prior possesses close connections to spline regression; see (Saatçi, 2012, Section 2.2.3) for more details.
where H(p) == = ExEXIRI p(z) log p(z) dy and the minimization is taken over all densities on [ under which the marginal law of 2y is equal to pk.
motion: ath (t) = (t).. Taking
We refer to (6) as the lifted Schrödinger Bridge problem, since it is obtained by lifting the optimization problem from densities on XiK] to densities on 2
At first sight, the lifted problem (6) is no improvement over (3)-indeed, the situation seems to have become worse due to the introduction of the continuous variables y. How- ever, we now show that unlike (3), problem (6) is directly amenable to efficient algorithms.
3. Lifting Schrödinger Bridges
The remainder of this paper is devoted to giving an efficient algorithm for smooth SBs whose reference process is GAP. In this section, we leverage the structure of GAPs to lift (3) to higher-dimensional problem, with better structure. In
The first step is leverage the Gauss-Markov property of Ri to simplify (6). Recall the fundamental fact that, under R, the law of n has the Markov property. Given z = y) 2 and E [K], we write z (I.,

---

Smooth Schrödinger Bridges
62
probabilistic strueture of w and the lifted process y under R means that can represent their joint distribution by simple hidden Markov model (see Figure 3), with ob- served variables fwx LK k== and hidden variables {ne } k== These correspond to variable nodes in Section 3, denoted by circles.
92-12
70-1
%
The dependencies among these variables nodes are repre- sented by factor nodes (squares in Figure 3): the factor nodes 0t.-1.1 connecting 7]k-1 and 77 enforce the joint law of the pair (1%=11 me., and the factor nodes or connecting Tis and w/ enforce the deterministic requirement that the first coordinate of 17,, equals wy. We associate to 12-1, the factor node potential 4x : Z x 2k given by
Figure 3. A portion of the graphical model corresponding to the joint law of and v).
for the kth coordinate of z, which takes values in Zx. The Markov property then implies that 2 is independent of 20,-. 11 conditioned on Zд. At the level of densities, this implies the decomposition
TE Zo + (zo. Z1)
Zn
L1 1z(00.20 21) == exp
=exx
2
r(z) r(zo)
II r(z 2)-1)
Qx (Z (Z 15
D3 1, = exp
h k> 1,
2
K
X r(z))
II exp
- 2
Qx
so that
K
k=1
r(z) x II
t (zk (zk-1,
where r(2o) x
20 20)) and each Qk is
quadratic function:
To find the optimal solution pd to (6), we use a belief propa- gation algorithm (Algorithm 2). The algorithm iteratively updates "messages" traveling between factor and variable nodes, which are depicted in Figure 3. The "vertical" mes- sages 8g. 7 are real-valued functions on Xr which encode information about the marginals of p", while the "hori- zontal" messages 6t 6% are real-valued functions on Z₁ which encode information about the joint distribution of Tis and its neighbors 7パ-1 and 7/+11 These messages are iteratively updated back and forth across the graph via the application of the following operators:
(z( = (zk- A A7 (zx- AaZx-1),
for matrices Ak = Ek - E Rmxm and As := Σx ΣS -1 - E mmx. These representations follow directly from standard for- mulas for Gaussian conditioning, and admit explicit expres- sions in terms of the kernel of the GAP (Saatçi, 2012, Section 2.3.1).
This decomposition represents the first term in (6) as a "tree- structured cost" (Haasler et al., 2021a), therefore rendering the marginalization in (4) amenable to belief propagation methods. We develop such a method in the next section.
Ie(ō% (055,8))(a (Zx) : : = / 0% (F2. (ak. yr)dy 32
4. Belief Propagation
10- (B) (2x 1)=
The goal of this section is show that a belief propagation algorithm can be used to efficiently solve the lifted prob- lem (6). Belief propagation (Kschischang et al., 2001; Pearl, 1982) is a canonical approach for performing inference in high-dimensional models whose densities possess simple factorizations, such as the one given in (7). The use of be- lief propagation algorithms for Sinkhorn- type problems is not new (Haasler et al., 2021a; Singh al., 2022; Teh & Welling, 2001); however, we stress that their application to Schrödinger bridges with smooth priors is novel.
Σ
Da Zs zx)o (2x))
Rr Rx (8 Bo- (Zxx)
Σ
Pz(ZE-1, 12,00 (2x 1) (2R i)dy
Jia-1
12-1 ENNe-
Since these operators involve manipulating functions on the space Zg. they are not directly implementable. In this section, we regard these operators as single basic operators for the purpose of complexity analysis. We develop effi- cient techniques to bypass their direct evaluation in the next section.
To develop our algorithm, we first reformulate (4) in the language of graphical models. (See Koller & Friedman, 2009; Wainwright& Jordan, 2007, for background.) The
5

---

Smooth Schrödinger Bridges
The main result of this section is that Algorithm 2 implicitly implements the Sinkhorn algorithm (Algorithm 1) for the multi-margind entropic optimal transport problem (3), with cost given by
(2) The solution of (6) given by
K
p(z) x II
24 (2- Zx)
(xx)
(12)
k=1
k=0
K II
where (EK) is the fixed point of Algorithm (2).
exp(-C(xx) =
24) dyK 1=1
The representation in (12) implies that the output of Al- gorithm 2 can be used to efficiently manipulate and sam- ple from the solution to the smooth SB problem; see Ap- pendix D for details.
9 This connection allows us to develop a rigorous conver- gence guarantee for Algorithm 2. Indeed, (9) implies that exp(-C(xx) is, up normalizing constant, the joint den- sity of R R). so that Algorithm 1, and hence Algorithm 2, solves the smooth Schrödinger Bridge problem.
5. Approximate Belief Propagation
The final ingredient of our algorithm is an approximation scheme to efficiently implement the operators in (8). We use the technique proposed by Noorshams & Wainwright (2013): decompose the continuous messages in suit- able orthonormal basis. The orthonormal decomposition method provides two key benefits. First, by truncating the series at sufficiently high order, we can create an accurate representation of the messages using only a finite number of orthogonal coefficients. Second, the key update rules in Algorithm described in (8) involve taking the T,2 inner product of two continuous messages. Expressing the mes- sages in an orthonormal basis simplifies these operations considerably.
Algorithm 2 Belief Propagation with Continuous Massages K- 1-1 distributions
1: Input: Factor node potentials {
{u&) on {Kk K k=
2: Initialize & = 0, x 1, oj = and ( Nx - IR+,
& E [K] arbitrary
3: while not converged do
4: ( l
5:
for k = K ..., do // left pass */}
5 t 1 )
6:
7:
end for
8:
70
(8)
< Iolōu 5t
9:
90 + na O Yu
Let { represent an orthonormal basis in the space L'R The subscript & indicates that one free to choose different bases for different e: [K].
10:
for k = l,... K do (/+ right pass $/}
11:
5 Ra(o5-1 pp
(0
(6) [6 5,
We express the horizontal messages of and 8f in terms of this basis as follows:
12:
7k
13:
8 + pa Ø %
14: end for
0 (2x) (mx)))(y)),
(13)
end while Return )EEKK)
16:
5% (21)) = (((m))
(14)
Theorem 4.1. Algorithm is equivalent to Algorithm with cost given by (9), in the sense that, if the initialization of Algorithm 2 and Algorithm I satisfy
where and denote the coefficients for rightward and leftward messages and the coefficients are given by
(((x))) = / 5t (x)
Uk (0)
vk E [K]
(10)
(x()) / o% (IX, o
then. for all 8> 0 and all l = ..., K, we have
Utilizing the orthonormal expansions from (13)-(14), we can re-write Algorithm so that operates directly on the coefficients 4 := 2 and Pュ := (r() 22. which are functions from X' Eg.
(8) ()) 7 S(() (
(2) Uk
(11)
In particular, we have the following two consequences:
(I) Algorithm achieves an s-approximate solution for (3) in O(KCmax/e-1) iterations, where Cmax = maxxeX (xx) and Ô(-) hides polylogarithmic factors the problem parameters.
First, by I2((2)) orthogonality, the update rule for 2 can be written
Ix 8, (2x) ((xx)r(xx))
6

---

Smooth Schrödinger Bridges
Note that in this notation, we can write
Algorithm 3 Approximate Belief Propagation Algorithm
Precompute To,..., Tx-1 E |RxxxxMXX XXM by (16). Initialize {rx(2x) and (4u())) as M dimen- sional vectors filled with I's. while not converged do
Be(mx) = (=(=x)/7r(xx) ((st)))(
We now consider the update rule for the left and right mes- sages. The following lemma shows how to express these updates in terms of coefficients.
for k: = K,..., do
Update 4 by (19) end for
Lemma 5.1. Fix functions o% = -i=1 lg 1 of and Sx.
Calculate Bu by (18)
Then the coefficients of o L B6) in the basis
for k: = 1..., K do
(of ( 1 are given by
Update ru by (20)
Calculate 3x by (18)
end for end while
(0k-1) = Σ
(k(
Σ j=1
(
Return V& and ly
(15)
where pif 1 (1 denotes
6. Time complexity and practical considerations
If Similarly, 5 Ei=1 oh-1. then the coeffi- cients of dg = Re(dx-1 Bx-1) in the basis (e are given by
(x de)x (16)
Implementing Algorithm 3 requires selecting bases (ooJ)1))e along with number of coefficients M. The computational complexity of the resulting algorithm scales directly with M, which we summarize in the following result.
(1x)
Theorem 6.1. Executing T iterations of Algorithm 3 takes O(TKn²M²) time. In particular, executing T = Ö(K 1)) iterations takes n m time.
Σ Th IEM
(x-1(Ex-1)
5=1
(17)
Theorem 4.1 suggests that OKC iterations suffice to obtain a e-approximate solution to (6); unfortu- nately, however, we lack a rigorous approximation guaran- tee quantifying the difference between the output of Algo- rithm 3 and that of Algorithm 2. Nevertheless, the success of our empirical results (Section 7) indicates that Algo- rithm 3 does offer a good approximation for the solution to the smooth SB problem. We leave the open question of demonstrating this fact theoretically to future work.
To obtain a practical procedure, we replace the infinite sums in (13)-(14) with finite approximations. Choosing a suffi- ciently large M, and using the first M orthonormal func- tions the basis approximate ot and 6% we repeat the previous steps to derive the update rules, expressed as matrix-vector multiplications.
(xx) + (1xx) (x (xx))
(18)
4- (x=k( + Σ
14), (mE)
(19)
The main tuning parameter of our algorithm is the choice of M. In principle, this choice should depend on the smooth- ness of the messages or and o the order m of the GAP, and the dimension d. Since 6 and dj. correspond to Gaus- sian densities, their expansion in many reasonable bases (for example, Fourier or Wavelet bases) will exhibit strong decay; however, since they are defined on 3%- R(m-1)d, standard smoothness arguments would predict that the nec- essary number of coefficients scales exponentially with the product md.
rk(*k) t
(20)
Σ
Bk-1 (Ix DrZ (#x- Tx- (32-1,
where 4- (xx-1)) and rx (ark) are vectors in R and (x--( xx) ) is a matrix in XMxM whose element in
row i and column j is given by (T
Upon convergence of this algorithm, we obtain the coeffi- cients тж and 4. and thereby obtain estimates of 37. As discussed after Theorem 4.1, these messages can be used directly for downstream tasks involving the Schrödinger bridge.
However, the dependence on d can be somewhat ame- liorated under additional assumptions. Suppose that the GAP has independent coordinates and the orthogonal bases have tensor product structure, so that each basis element
7

---

Smooth Schrödinger Bridges
E (2( satisfies oly) IT (y()) for (pl)) E (R= and where y(0) corresponds to the deriva- tives of the jth coordinate of w. In this case, the tensor r factors across each dimension into the tensor product of smaller tensors [(0
I,....d), d). Suppose the num- ber of coefficients we use is equal for each dimension. i.e.
TO) E R à xarà then the complexity of steps (19) and (20) drops to P(dn2M?²5 by taking advantage of this lower rank structure. In this important special case, therefore, the time complexity of our algorithm scales more benignly with d, which allows us to take larger M when rises. A record of the running time of one iteration of message passing against M and dimension d is presented in Table
Figure 4. Visualization of orbits in 3D space. The colors of trajec- tories depend on the starting point. The second row is the visu- alization of the XY-space projection of the corresponding above plot.
Table f. A comparison between smooth SB and two haseline parti- cle tracking methods. Information on evaluation metrics appears in Appendix F.
7. Experiments
We test our algorithm on two kinds of cow-dimensional smooth trajectory inference tasks. The first kind aims at tracking the exact trajectory of each individual particle (e.g. Figure 2), which we refer to as the One-By-One (OBO) tracking task in the following text. For this kind of task, we evaluate the performance by calculating the distance of each inferred trajectory and the ground truth and measure the per- centage of time that the algorithm tracks a partiele correctly. For the second kind, the task is to infer the group trajectories of point clouds whose evolution has geometric structure. We provide two ways to evaluate the performance for this kind of task, similar to the evaluation in Banerjee et al.. We first keep all the observations and visualize the trajectories and see if they form a pattern that is close to the ground-truth pattern. Secondly, we will leave out observations at a certain timestep and instead infer the position of particles at this timestep and evaluate the distance between the inferred and real observations, which we call Seave-One-Tie One-Timestep-O0 (LOT) tasks. Our code for reproducing these experiments is available Github.! !
Dataset
Method
JumpP
5p acc
Mean lg
Tri-stable Diffusion
SSB (ours)
1.12c-1
0.798
1.42c-2
BMSB
5.20e-1
0.171
6.79e-1
W2M
2.25e-2
0.956
1.37e-8
N Body
SSB (ours)
5.00e-4
0.999
5.48e-4
BMSB
1.14e-1
0.641
5.43e-1
W2M
1.08e-1
0.649
5.50e-1
2D Gaussian Process
SSB (ours)
5.60e-3
0.991
3.57e-4
BMSB
1.37e-1
0.612
2.82c-1
W2M
6.60e-2
0.760
2.32e-1
appear in Appendix F.
7.2. Point clouds trajectory inference
We also test our algorithm on five challenging baselines in the trajectory inference literature. We compare our al- gorithm quantitatively on tasks of LOT with three other state-of-the-arr algorithms: MIOFlow (Huguet et al., 2022b), DMSB (Chen et al., 2023) and F&S (Chewi al., 2021). visualization of the output of our algorithm appears in Fig- ure 5, 6 and the quantitative results for the LOT tasks are provided in Table 2. Our algorithm can recover the geomet- ric pattern in each low-dimensional dataset by generating smooth trajectories. We can also apply our algorithm to somewhat higher dimensional settings by appropriately tun- ing the number of coefficients the orthonormal expansion. In particular, in our experiments on the 10-dimensional Dyngen cycle dataset shown in Figure 6, we use 4 approx- imation coefficients for each of the first 5 dimensions and approximation coefficient for the each of the last 5 di- mensions. Despite the small number of coefficients, the resulting trajectories are still meaningful. Full experimental details appear in Appendix
7.1. One-By-One tracking
For OBO tasks, we consider three synthetic datasets where trajectories of particles intersect frequently. We compare our algorithm with the standard Schrödinger Bridge (SB) and a modification based on computing the W2 optimal matching at each step (W2M). We test on three data sets, two 2-dimensional data sets (Fig 7 in Appendix F) and a 3-dimensional data set consisting of the simulated orbits of an N-bodd physical system (Fig. 4). Smooth SB performs second-best on the Tri-stable diffusion dataset and substan- tially outperforms other approaches the more challenging N Booo and Gaussian Process data. Quantitative evalua- tions are summarized in Table 1. Full experimental details
'tttps: / www.//w/wwooooooooooooooooooo.oo http:taaww...oaaaaaaooooo/n//eee hataaaaniiiioonnn/ooooosss
8

---

Smooth Schrödinger Bridges
8. Discussions and Future Directions
We have presented a new method for trajectory inference and particle tracking based on smooth Schrödinger bridges, which achieves very good performance on number of chal- lenging benchmarks. The main limitations of our proposal are related to the approximate implementation developed in Section 5. As we have discussed, our proposal suffers from the curse of dimensionalityy because the number of coefficients typically scales exponentially with respect to both the order of the GAP (m) and the dimension of the observations (d). In numerical experiments, relatively small values of M (of order 1000) seem to perform well for problems up dimension 10. An important question for future work is to either develop an approach with better dimensional scaling or, alternatively, show that the exponen- tial scaling in dimension is unavoidable, as is the case for Wasserstein barycenters (Altschuler & Boix-Adserå, 2022).
Figure 5. Visualization of trajectory inference on various datasets by lifted SB, for 5D Dyngen Tree data, we visualize the 2D projection of the first three dimensions in the second row.
Implementing our approach also requires selecting suitable GAP to use as reference process. As our experiments make clear, it is not necessary that the reference process match the data generating process precisely (see, e.g., Figure 2). However, picking an appropriate variance for the Gaussian process is important for good performance (see Figure 16). In our experiments, choosing
Figure 6. Visualization trajectory inference on the 10D Dyngen cycle dataset projected by Phate, we visualize the 2D projection of the first dimension against the second to fifth dimension.
data where data is diagonal matrix containing the empirical standard deviation along each dimension, typically works well.
Finally, we have considered a definition of the Schrödinger bridge which enforces the strict marginal constraint Ps- = Pk- In applications, it is natural to assume that observa- tions of the particles are corrupted with noise, which mo- tivates a version of the SB with an approximate constraint Ps RJ Pa (Chizat et al., 2022: Lavenant et al., 2024). is possible to incorporate noisy observations into the graphical model framework we describe above by introducing suit- able potential at the factor nodes ag. A similar approach generalizes to other missing data problems, for instance, when some dimensions are not observed. We leave this extension to future work.
Table Performance comparison on the Leave-One One-Timestep Timestep-O0 Out task at step j between algorithm and the state-of-art algorithms. Our algorithm is typically best or second-best. DMSB failed to converge the Dyg dataset despite extensive tuning. Information on evaluation metrics appears in Appendix F.
Dataset
Method
W1(()
Mc(1)
My(4)
Petal
SSB (ours)
2.70e-2
2.21e-5
3.75e-3
j = 2
MIOFlow
2.22e-1
9.06e-3
1.1le-1
DMSB
2.10c-1
3.52e-3
3.68e-2
F&S
2.05e-2
2.85e-5
4.48e-3
EB j =
SSB (ours)
8.45e-2
3.46e-3
5.04e-2
MIOFlow
1.34e-1
1.36e-3
2.81e-2
DMSB
1.46e-1
9.43e-3
9.72e-2
Acknowledgements
F&S
8.72e-2
1.47e-3
3.88e-2
JNW acknowledges the support of National Science Foun- dation grant DMS-2339829.
Dyg Tree j =
SSB (ours)
9.81e-2
1.64e-3
3.51e-2
MIOFlow
2.33e-1
2.82e-2
1.73e-1
DMSB
F&S
9.78e-2
2.00e-3
4.64e-2
Impact Statement
Dyg Cycle j 7
SSB (ours)
1.94e-1
2.85e-2
1.7le-1
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
MIOFlow
4.25e-1
1.57e-1
4.11e-1
DMSB
F&S
2.84e-1
2.91e-2
1.70e-1
9

---

Smooth Schrödinger Bridges
References
Chen, T., Liu, G., Tao, M., and Theodorou, E. A. Deep momentum multi-marginal Schrödinger bridge. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Con- ference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, Decem- ber 10 16, 2023, 2023. URL http: pappers. nips. copppcccppppr..emc httpp/aapprrrrssaeeeeeeeeeeeeeee www.s/eeeeaaeeeeeee /ttps//aap/eeaa//2/33333 /hash/ 7260ffeece388810ffff030000000002 Abstract-Confereene html.
Altschuler, J. M. and Boix-Adserà, E. Wasserstein barycen- ters are NP-hard to compute. SIAM Journal on Mathe-
matics of Data Science, 4(1):179-203, February 2022.
ISSN 2577-0187. doi: 10.1137/21m13300062 URL http: / /dx. doi. /ttp:/c............1../.........
Altschuler, J. M. and Boix-Adserà, E. Polynomial-timg algorithms for multimarginal optimal transport problems with structure. Math. Program., 199(1-2, Ser. 1):1107- 1178, 2023. ISSN 0025-5610. 022-01868-7, URL https:/ /ttp:/ii.oo......///////. 1007/ 000-007022-088656
doi: 10.1007/s10107-
Chen, Y., Conforti, G., and Georgiou, T. T. Measure- valued spline curves: An optimal transport viewpoint. SIAM Journal on Mathematical Analysis, 50(6):5947- 5968, 2018. doi: 10. 00.117718M11166444 URL https: // //dpi. 1137
Banerjee, A., Lee, H., Sharon, N., and Moosmüller, C. Ef- ficient trajectory inference in Wasserstein space using consecutive averaging. In The 28th International Confer- ence on Artificial Intelligence and Statistics.
66249.
Chen, Y., Conforti, G., Georgiou, T. T., and Ripani, L. Multi- marginal Schrödinger bridges, pp. 725-732. Springer In- ternational Publishing, 2019. ISBN 9783030269807. doi: URL http: /dxx doi doi.org org/orrg100 10011 3- 030 116980-7_777
Benamou, J.-D., Gallouët, T. O., and Vialard, F.-X. Second-order models for optimal transport and cubic splines on the Wasserstein space. Foundations of Com- putational Mathematics, 19(5):1133-114 Oct 2019. ISSN 1615-3383. doi: 10.10077A112008-0199.09222 URL httphhtttss/.k..spttttee http://er....o//oooontoeeeeeeeee pdf.
Chewi, S., Clancy, J., Le Gouic, T., Rigollet, P., Stepa- niants, G., and Stromme, A. Fast and smooth in- terpolation Wasserstein space. In Banerjee, A. and Fukumizu, K. (eds.), Proceedings of The 24th In- ternational Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 3061-3069, PMLR, 13-15 Apr 2021. URL https: http:////occeeeeeeeeeeeeeeeessee mlr mww. v130 chewi 21a. html.
Borisov, I. S. On criterion for Gaussian random processes to be Markovian. Theory of Probability & Its Applica- tions, 2764),863-865, 1983. doi: 10.1137/1127097 URL https:/ / /oi. www.da......eeeeeo.ee...eoooee.e 10. 1137 01333112170777
Brunton, S. I.,, Noack, B. R., and Koumoutsakos, P. Machine learning for fluid mechanics. An- mual Review of Fluid Mechanics, 52(Volume 52, 2020):477-508, 2020. ISSN 1545-4479. doi: htwp:////......llaauueeeeeeeeee
Chizat, L., Zhang, S., Heitz, M.. and Schiebinger, G. Trajectory inference via mean-field Langevin in path space. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Ad- vances Neural Information Processing Systems, volume 35, pp. 16731-16742. Curran Associates, Inc., 2022. URL httpss/ htt.ttss://prccceeeieee ngs. neurips. cc/paper /c//speet_iiles/ppprr//22//// GaI1BBlcee71f6677777raa111111233 Paper-Conferrece..dde
URL
org/ content//ournals/10.. 114 annurev- fluid-010719-06620
htttps://www.annualll
Bunne, C., Hsich, Y., Cuturi, M., and Krause, A. The Schrödinger bridge between Gaussian measures has a closed form. In Ruiz, F. J. R., Dy, J. G., and van de Meent, J. (eds.), International Conference on Artificial Intelligence and Statistics, 25-27 April 2023, Palau de Congressos, Valencia. Spain, volume 206 of Proceedings of Machine Learning Research, pp. 5802-5833. PMLR, 2023. URL https: /pppcceedingss mlr a.lr.press/ v v206 bbbnnnunn2..a...mm html.
Clancy, and Suarez, F. Wasserstein- Fisherr-Faaaaaa splines.
arXiv preprint wxXiv:223.115225, 2022.
De Bortoli, V., Thortoo,, J., Heng, J., and Doucet, A. Diffu- sion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:1 1661151777779 2021.
Cannoodt, R., Saelens, W., Deconinck, L., and Sacys, Y. Spearheading future omics analyses using dyngen, a multi-modal simulator of single cells. Nature Communi- cations, 12(1):3942, 2021.
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Bois-
bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fa- tras, K., Fournier, N., et Pot: Python optimal trans-
10

---

Smooth Schrödinger Bridges
port. Journal of Machine Learning Research, 22(78):1-8, 2021.
Justiniano, J., Rumpf, M., and Erbar, M. Approximation of splines in Wasserstein spaces. ESAIM: Control, Opti- misation and Calculus of Variations, 30:64, 2024. ISSN 1262-3377. doi: 10. 10.105 1/oee///0022000000000 URL http: / //dx dww..di..... org/ / 0 051 //55//ccc/2/22008.
Gilboa, E.,, Saatci, Y., and Cunningham, J. P. Scaling multi- dimensional inference for structured Gaussian processes. IEEE Trans. Puttern Anal. Mach. Intell., 37(2):424-453 2015. doi: 10.1109/TPAMI.0011 2013.192 192. URL ht tps: //doi. wtt.l//rrg....119////1.. 1wt.:1109//AAAAA 2013.192.
Koller, D. and Friedman, N. Probabilistic g graphical mod- els. Adaptive Computation and Machine Learning, MIT Press, Cambridge, MA, 2009. ISBN 978-0-262-01319-2. Principles and techniques.
10.
Gushchin, N., Kolesov, A., Korotin, A., Vetrov, D. P., and Burnaev, E. Entropic neural optimal transport via diffusion processes. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 75517-75544, Curran Associates, Inc., 2023. URL https: /ttppc/////ccceiiggs.... neurips. cc http/ppe__fllleeppppr///////////
Kschischang, F. R., Frey, B. J., and Loeliger, H.-A. Fac- tor graphs and the sum-product algorithm. IEEE Trans. Inform. Theory, 17(2):498-519 2001. ISSN 0018-9448. doi: 10.1109/18.910522 URL https://dii..o.. org/ 10. 11w./18111... 910572.
fi
eeac 41 4a 4aac5141441114404444333333 4dbbbbbbb32244bbbbbb-bbbbbb
Kubica, J., Denneau, L., Grav, T., Heasley, J., Jedicke, R., Masiero, J., Milani, A., Moore, A., Tholen, D., and Wainscoat,
Paper-Conferrnce.. pdf.
Efficient intra- and inter-night linking of asteroid detections using kd- trees. Icarus, 189())1151-168. 2007. ISSN 0019- 1035. doi: http://oooong//lo.00jjjjaaa.....
Haasler, I., Ringh, A., Chen, Y., and Karlsson, J. Multi- marginal optimal transport with a tree-structured cost and the Schrödinger bridge problem. SIAM J. Control Op- tim., 59(4):2428 2453, 2021a. ISSN 0363-0129. doi: 10.1137/20M13201995 URL httpt:ttt///o http:///oo..o////////
008. URL WRLnttpps//www.ccieecedieeeeet cccienee/attidee//////0111000000
Lavenant, H., Zhang, S., Kim. Y.-H., and Schiebinger, G.
Haasler, I., Singh, R., Zhang, Q. Karlsson, J., and Chen, Y. Multi-marginal optimal transport and probabilistic graphical models. IEEE Transactions on Information Theory,67) 67(7):4647-4668, 2021b. doi: 10.1109/TIT:2021. 3077465.
Toward a mathematical theory of trajectory inference. The
Annals of Applied Probability, 34(1A), February 2024. ISSN 1050-5164. doi: 10. 2up12/23 URL http: //dx. doi org. 10. 1214 AAP11999..PP
Léonard, C. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete Contin. Dyn. Syst. 3444)11533-1574 2014. ISSN 1078- 0947. doi: 10.39344/ccds.2144.44.4 / ///oi. orgi.org/1 10. 3934 2014. 34.4..44.1333
Hartikainen, J. and Särkkä, S. Kalman filtering and smooth- ing solutions to temporal Gaussian process regression models. In 2010 IEFE International Workshop Ma- chine Learning for Signal Processing, pp. 379-384, 2010. doi: 10. 10.110//LSP.0000555599
URL https:
Lin, T., Ho, N., Cuturi, M., and Jordan, M. On the com- plexity of approximating multimarginal optimal transport. Journal of Machine Learning Research, 23(65):1-43, 2022. URL http: /m/r..mrr.../ org warg/papeene papers / v23/19- 1tml.
Huguet, G., Magruder, D. S., Tong, A., Fasina, O., Kuchroo, M., Wolf, G., and Krishnaswamy, Manifold interpolating optimal-transport flows for trajectory inference. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Ad- vances in Neural Information Processing Systems, volume 35, pp. 09705-29778... Curran Associates, Inc., 2022a. URL https: http::::s//pccceeeinng.oooo neurips.
Liounis, A. J., Small, J. L.., Swenson, J. C., Lyzhoft, J. R., Ashman, B. W., Getzandanner, K. M., Moreau, M. C., Adam, C. D., Leonard, M., Nelson, D. S., et al. Au- tonomous detection of particles and tracks in optical im- ages. Earth and Space Science, 7())e2019EA000845, 2020.
cc bfco3f076884888830888828877 Paper-Conference.p2
Huguet, G., Magruder, D. S., Tong, A., Fasina, O.,, Kuchroo, M., Wolf, G., and Krishnaswamy, S. Manifold inter- polating optimal-transport flows for trajectory inference. Advances in neural information processing systems, 35: 29705-2971 29718, 2022b.
Moon, K. R., Van Dijk, D., Wang, Z., Gigante, S., Burkhardt, D. B., Chen, W. S., Yim, K., Elzen, A. v. d., Hirn, M. J., Coifman, R. R., et al. Visualizing structure and transitions in high- nigh-dimensional biological data. Nature biotechnol- ogy, 37(1)): 17(112):482-1144 148--1492 2019.

---

Smooth Schrödinger Bridges
Noorshams, N. and Wainwright, M. J. Belief propagation for continuous state spaces: Stochastic message-passing with quantitative guarantees. J. Mach. Learn. Res., 14: 2799-2835, 2013. ISSN 1532-4435.
Schrödinger, E. Über die Umkehrung der Naturgesetze. Angewandte Chemie, 44(30):636-636, 1931.
Schrödinger, E. Sur la théorie relativiste de l'électron et I'interprétationn de la mécanique quantique. Ann. Inst. H. Poincaré, 2(4))269-310, 1932. ISSN 0365- 320X. URL http:/ /tt:://wwww www.:wumdam.oomm org/ id-AIHP 1EAAAHP_13222 2_4_269_0
Ouellette, N. T., Xu, H., and Bodenschatz, E. A quantitative study of three-dimensional Lagrangian particle tracking algorithms. Fluids, 40(2):301-313, Feb 2006. ISSN 1432-
Experiments in
Sha, Y., Qiu, Y., Zhou, P., and Nic, Q. Reconstructing growth and dynamic trajectories from single-cell tran- scriptomics data. Nature Machine Intelligence, 6(1):25- 39, November 2023. ISSN 2522-5839. doi: 10.1038/ 542256-023--07663-w URL http: http:/dxxxxxdooo/////o
1114. doi: 10. 00.10077/003488-55-00666
URL
attps://links Slttps///L....spiiiger.cco wttp://eer.occccooeeeeeee pdf/10.
Pavon, M.. Trigila, G., and Tabak, E. G. The data-driven Schrödinger bridge. Communications on Pure and Ap- plied Mathematics, 74(7):1545-1573, 2021. 2021.
Shen, Y., Berlinghieri, R., and Broderick, T. Multi Multi-marginal schrödinger bridges with iterative reference refinement. The 28th International Conference on Artificial Intelli- gence and Statistics.
Pearl, Reverend Bayes on inference engines: A distributed hierarchical approach. In Waltz, D. L. (ed.), Proceed- ings of the National Conference Artificial Intelligence, Pittsburgh, PA, USA, August 18-20, 1982, pp. 133-136.
Singh, R., Zhang, Q., and Chen, Y. Learning hidden Markov models from aggregate observations. Auto- matica, 137:110100, 2022. ISSN 0005 1098. doi:
AAAI Press, 1982. URL http: / / www. http://ww..aw.a..a.../.......... org/ Library/ /Atp://AAAI////1 1aw///a22222222222 032112222000000000 php.
Phillips, A. W. The estimation of parameters in systems of stochastic differential equations. Biometrika, 46(1-2):
URL https::/www.sccieneddiccc..
67-76, 1959. ISSN 0006-3444. doi: 10.2307/2332809. URL Btthhtsss//6o..ooo 10. 1t..11..220.11.
2307 /23228292
Sinkhorn, R. Diagonal equivalence to matrices with pre-
Pooladian, A.-A and Niles- Weed, Plug-in estimation of Schrödinger bridges. 08 2024. URL https::/aaaat
scribed row and column sums. Amer. Math. Monthly, 74: 402-405, 1967. ISSN 0002-9890. doi: 10.2307/2314570. URL https: / //w.i/doi/////////////.///./ org/10, 2307 00307/2214555 2314570.
org/pdf/244 /24dd .1666.p......
Teh, Y. W. and Welling, M. The unified propaga- tion and scaling algorithm. In Dietterich, T. G., Becker, S., and Ghahramani, Z. (eds.), Advances in Neural Information Processing Systems 14 [Neu- ral Information Processing Systems: Natural and Synthetic, NIPS 2001, December 3-8, 2001, Vancou- ver, British Columbia, Canada], pp. 953-906 MIT Press, 2001. URL https::/proceedii
Rasmussen, C. E. and Williams, C. Gaussian pro- cesses for machine learning. Adaptive Computation and Machine Learning. MIT Press, Cambridge, MA, 2006. ISBN 978-0-262-18253-9)
Saatçi, Y. Scalable inference for structured Gaussian pro- cess models. PhD thesis, Citeseer, 2012.
Saelens, W., Cannoodt, R., Todorov, H., and Saeys, Y. A comparison of single-cell trajectory inference methods. Nature Biotechnology, 37(5):547-554, April 2019. ISSN 1546-1696.
neurips wwtpc//ppppppp////eeaaeeeeeeeeee
Abstract. html.
doi: 10.1038/s41587-019- 0071-9. URL http:: /xx wtw.do.og.....000000 s41587 01 0019000--99
Tong, A. Huang, J., Wolf, G., Van Dijk, D., and Krishnaswamy, S. TrajectoryNet: A dynamic opti- mal transport network for modeling cellular dynam- ies. In Daumé, III, H. and Singh, A. (eds.). Pro- ceedings of the 37th International Conference on Ma- chine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9526-9536. PMLR, 13-18 Jul 2020. URL www.thttt::/ppoeeeeeete v1 9 /ttg//ong22a... html
Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subra- manian, V. Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P. Lee, L., Chen, J., Brumbaugh, J., Rigol- let, P., Hochedlinger, K., Jaenisch, R., Regev, A., and Lander, E. S. Optimal-transpont analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928-99433322) 2019, ISSN 0092-8674. doi: http:///loiiggg/0.0..0...c......
mlr press/
Vargas, F., Thodoroff, P., Lamacraft, A., and Lawrence, N. Solving Schrödinger bridges via maximum likeli- hood. Entropy, 23(P) No. 1134, 30, 2021. doi:
006. URL /RO////ttpsee//ww..eeeeeeeeeeeee siciaance/rrticeee//////22222220
12

---

Smooth Schrödinger Bridges
10.3399//633911134 URL http::///dci.rr//10. 3390/e23091134.
Wainwright, M. J. and Jordan, M. Graphical models, exponential families, and variational inference. Founda- tions and Trends@ in Machine Learning, 1(1-2):1-305, 2007. ISSN 1935-8245. doi: 10.1551122000000001 URL http:///c...o..rrr/0/...66///000
13

---

Smooth Schrödinger Bridges
A. List of notation
P(Z): the set of Borel probability measures on X.
D-I--): Kullback-Leibler divergence. For probability measures P and R,
D(P[R): = [log dP dP P&R an +
otherwise.
CRR; R C (R; (R; R"): continuous (respectively, m-times continuously differentiable) functions from IR to Rd.
[N], [ for a positive integer N, [N] is the set O....N while [N] is the set I....N.. N.
Rt: the set of positive real numbers
Xy: the support of Pk. assumed to be finite, Vk: the space R(m-1)d, identified with the possible values of ...,,,.,,,.. d w-imm- at 1.2- Zx: the phase space Xx Vk.
Xm -IkkE Nr. M/1 [TbErr] D'. 2pr
IIEET Zx.
if σ& = ok otherwise.
1o, (xt)): the indicator function 1o, (zx) =
① and point-wise multiplication and division
Poly(Ki..... Ke): (multivariate) polynomial function in parameters Ki,..., K₁.
B. Prior work
There has been significant prior work both the Schrödinger bridge problem and its applications in trajectory inference. Here, we briefly survey some important related contributions.
The Schrödinger bridge has been the subject of significant theoretical interest since its introduction; see (Léonard, 2014) for historical details and additional context. Currently, there are many methodological approaches to the Schrödinger bridge problem (see, e.g. De Bortoli et al., 2021; Gushchin et al., 2023; Pavon et al., 2021; Pooladian & Niles-Weedt 2024), almost all of which focus on the two-marginal case (in our notation K 1). As was forcefully pointed out by Altschuler Boix-Adserå, the multi-marginal case presents significantly greater computational challenges (Altschuler Boix-Adserà, 2022; 2023); indeed, no general-purpose algorithms with running time polynomial in K exist. However, it is possible obtain polynomial-time algorithms in certain special cases (Altschuler & Roix-Adserà, 2023). Our work shows that smooth Schrödinger bridges with GAP priors are such case.
The connection between multi-marginal entropic OT problems and belief propagation has been highlighted in a number of prior works dating back more than two decades (Singh et al., 2022; Haasler et al., 2021b; Teh & Welling, 2001), though apparently without connection to smooth Schrödinger Bridge problems. The exception is the "momentum Schrödinger bridge" of (Chen, Conforti, Georgiou, and Ripani, 2019; Chen, Liu, Tao, and Theodorou, 2023), which is, implicitly, example of smooth Schrödinger with prior given by integrated Brownian motion; the Sinkhorn algorithm proposed in (Chen al., 2019) can be viewed as ad hoc version of belief propagation for this special case. However, those works do not make the connection with smooth Gaussian processes, and their algorithm does not apply to more general GAPs.
The importance of the Schrödinger bridge problem for trajectory inference was implicitly recognized in the pioneering work (Schiebinger et al., 2019), and developed mathematically by (Chizat et al., 2022; Lavenant al., 2024). The fact that the vanilla Shrödinger bridge problem uses Markov prior is crucial for algorithms (Chizat et al., 2022), but has also been recognized as an undesirable feature that has led the development of different trajectory inference methods which give rise to smoother paths. Apart from (Shen al.), which proposes "robust" version of the Schrödinger Bridge problem in which a single reference process is replaced by family of (Markov) reference processes, these alternate methods largely abandon Schrödinger's 's formulation and develop very different techniques. Many of the most successful methods are based on measure-valuee generalizations of splines (Banerjee et al.; Benamou et al., 2019; Chen et al., 2018; Chewi et al., 2021; Clancy & Suarez, 2022; Justiniano et al., 2024). Other methods based on neural networks have also been proposed (Huguet al., 2022a; Tong et al., 2020). We compare these methods in our experimental section.
14

---

Smooth Schrödinger Bridges
C. Additional results and omitted proofs
The following lemma shows that any Markov Gaussian process with differentiable sample paths is essentially trivial. Lemma C.1. Let w R+ + be real-valued Gaussian process that is Markovian and almost surely differentiable. Assume Var[w(t)] for all t 2 Then w(t) ([)( for all t 2 0 almost surely.
Proof of Lemma C.I. Since w an almost surely differentiable Gaussian process, we know that the covariance kernel 8 (tt, 8) == ([)) = ((), is differentiable both coordinates and 0zy [(t,, s) exists. It is well known (see (Borisov, 1983)) that a real-valued Gaussian process is Markovian and only if, for any t S l2 S t3, we have
((1,, (a) (a 62).
(21)
Let t v and E > 0. From (21), we have ((s, K(, s = ((t, 8) (s, Dividing each side by E and taking the limit as + 0, we obtain
0 K(L, 8) = (k),,)
K(s,8) 8)
K (s,
(22)
Ox2
d02
Similarly, we also have tt,,t) 0 K(t, s) [t,, 8) 0z
K(t,t). Combining these two identities, follows that
021
2²
2 K(s,s).
K(t,t)K(ss,, 8)
K1,, 8) = K(t, 8)
K(t, 1)
(23)
dz Or2
0z
Note that a (t, s) = Cov(w/ (t), (orlw/(t),w((11) and 0z1 expectation.
8² K(t,s) = Cov(w'(t),u'l by interchanging differentiation and Da Dr dar
For any 8, by taking / 8 in (23), we obtain
(s)(w((s)) Cov(w'(s),
(24)
which further implies that
(s)(w(s),, 0. Var(w(s))
Var(w (s)l(() = Var(w'(s)
Let 9s: [0, + R+ be given by
gx(t): = (ar(w((s)l(())
(25)
We have shown that (b) = O. Moreover, for any t E s). we have
d
VVr(w'(>) (8)) Cor(u'(s),
16(1)
(26)
di
Var(w(t))
2
0
K(ttt)
② (t,,s)²
-
(t,,s)
g2
K(t,s) k((,,t)
(27)
K(t,
Oay
Dz2
dxz
dr,drz
2
θ Kt,tt)
0 K(t,s) s)²
δ ((t,s)
K(t,s) à K(t,t)
((s
(28)
K(t, 8)2
dx1
duz
Ox2
Ox1
K(s,8)
2
a
i K(t, s)2
2
k((,tt) =((t:))=0) =0,
K(,tt)
K (t,
(29)
k(t, 8)2
Ox
Dzz
drz
dx
where the third line uses (23) and the last line uses (22). Consequently, we have shown that Var(w' (s))(w(t)) = 0 for any I < s. We complete the proof by applying the Cauchy-Schward inequality: for any s >> O, we have
Var (W) Var (w(s) (0)) (0)) < s 1." Var (w (O)) de 0.
(30)
Therefore, w(s) = E[w((s (0) almost surely for all > 0.
15

---

Smooth Schrödinger Bridges
C.1. Proof of Lemma 2.1
This follows directly from known arguments in the two-marginal case (see Léonard, 2014), which we recapitulate below.
We first show the equivalence between (I) and (2), which requires assumptions the marginal measures. We can decompose the measure R as R(w) R(W R) (w), where R(w w) denotes the conditional law of w given the values and similarly for P. By the chain rule for Kullback-Leibler divergence, we obtain
w =
((P] R) [P[K] R(K) Px D(P(- w) R( w))
For any choice of Pkx. choosing P(w w) = R(w w) makes the second term vanish. Therefore, any solution P of (1) must be of the form Pix] [w) w) for solution Pix] of (2), and conversely.
For the second claim, assume that each of the measures Ps is absolutely continuous with finite entropy, and that R has density. In this case, any feasible solution (2), i... any PKK such that ((Pk) [k]]))A()))))))))) oo, must be absolutely continuous with respect to RiK! and hence have density well, which we also denote by P(K]. We obtain
D(P[]] R(K))= / log (w)
flog
exp( exp(-C(w)) (w]](w) du
f ((w))) dw + / Pk](W) log
Pwn](W) [kRE[K] Ma (wiz)
dw + / (w)(WW) log II kE[K]
Na (wk) dw.
The marginal constraints imply that
/ (w](W) log
II kE(A)
ph (w(ww) dw Σ
/ kE[K]
pa (w)) log un (wk) duik
Since each μa finite entropy by assumption, this sum finite and constant over the constraint set, so it can dropped from the objective without changing the solutions.
C.2. Proof of Lemma 3.1
Lemma 3.1 can be reformulated as follows: there one-to-one correspondence between solutions to
Σ
(og(x))og(r(())-
Σ Px log PIx((x)/ II (a(TT)) kE[K]
max P(x) xEM(
(31)
XEA(K]
Ps = pks Ak E [K]
and
Σ p(z) log(r(z)) dy
Σ p(z) log(p(z))dy
max
Jym XEX
ViKi XEXINI
(32)
Pz. = μr Vk [T),
where the maximization in (31) is taken over distributions in P(K(r)) and the maximization in (32) is taken over probability densities on k-
First, note for any P(rj satisfying the marginal constraints Ps = ur for all E [K],
Σ (x) log
II (x(x)) = Σ Σ
log = Σ Σ na (x(x)) (ox))
KE[K]
kE[K] x
(E(K)=MER
Therefore, the term ExEr Pig log HEE[KK] (M(z)) in (31) is constant on the feasible set and can be dropped from the objective.
Next, by the law of total probability,
/ Σ p(z) log(r(z))dy /
Σ p(z) log(r(x))
Σ p(z) log(r(ylx)) jym xEXm
/
Jim xEXxI
Jy xEXim
(33)
16

---

Smooth Schrödinger Bridges
and
Σ p(z) log(p(z))dy Yio xEX(T)
/ Σ log((((x)) + J XEX(ZI DIT XEX())
p(z) log(p(y(x)) (()))p(())))d dy.
(34)
where r(x) and r(ylx) denotes the marginal density of the x variables and conditional density of the y variables under R(rr, and analogously for p. It follows from exchanging the order of integration and summation that the first terms on the right-hand side of (33) and (34) recover the two terms in (31) (assuming we have dropped ΣxEx1e P(x] log [IRE MM (Tx) from (31)). We may combine the second terms of (33) and (34) to obtaín that the remaining term in the objective of (32) reads
Σ p(x) / p(y|x) log( r(y(x) x
p(y|x) dy.
(35)
y/m
follows from the strict concavity of logæ: that (35) most O and equals if and only p(y[x) r(y|x) for every and p(-[x))-almsst every y. Since is probability density function, the equality condition is equivalent to saying that for every E Krr we have p(-xx) r(-lx) Lebesgue almost everywhere. Therefore, if Pir] is maximizer for (31) then (x)(()) x) is maximizer for (32). On the other hand, p" (z) maximizer for (32), then Jyy, p" (z)dy is maximizer for (31).
C.3. Proof of Theorem 4.1
We first prove that if the initialization of Algorithm 1 and Algorithm 2 satisfy
- (0)
Vk
[K]
(36)
then, at the end of each while le loop in both algorithms,
( =s(
%
(1)
(0)
for all & E [K].
We proceed by induction. We first investigate the first for loop in Algorithm 2 (the "left pass"). By backwards induction on k = K,..., 1, we obtain that at the conclusion of the left pass
5t (Z f Σ (2--1
(ax)dyk.... dyk Vk E [K]+
(37)
where the integration is taken over (yk..... yK) E I1 V and the sum is taken over (IR. (IR...., IKK) E II X, Therefore, the updates to "Po and Bo satisfy
(2) (xo) Io,,, (2o)) %
K
Σ
II
$ (2,-1,
(x.))
ViAi t=1
Σ exp(-C(x))
II 8 (xx)
21
s000 (2o),
where the last equality uses the induction hypothesis a({ ((--1) 1) =0g (f-1) and consequently
882 (ooooo)//// 1/5 (ma) ( o6 (xa)
17

---

Smooth Schrödinger Bridges
We now show by induction on & that 7kk S and 3h v at the conclusion of the "right pass". We have already established the base case l O. For & 1, in the left pass, the updates the right pass satisfy
(8)
k
6. (zx)
/
Σ
II
-- (z, 2:) 2 (2,) dya.
(38)
Combining (37) and (38) implies that
(x+) 1
K-1
1-1
Σ
II (2 2
II
8 (x.)
II i=k+1
(x:)
ViKl
K
Σ
exp(-C(x))
8 (x))
II 2mk+11
1=0
% (mx)
where the final equality holds by induction. As in the case, this implies 3 v/ This proves the claim.
(2)
The first implication, the time complexity of Algorithm 2, follows directly from existing analysis of multi-margine entropic optimal transport problems (Altschuler & Boix-Adserå, 2023; Lin et al., 2022). Indeed, those works show that Algorithm 1 yields an e-approximate solution to a discrete multi-marginal entropic optimal transport problem with arbitrary cost tensor C' in Poly(K, K, /s) iterations. As each iteration of Algorithm corresponds to an iteration of Algorithm 2, a similar guarantee holds for Algorithm 2. In particular, (Lin et al., 2022, Theorem 4.3 and proof of Theorem 4.5) shows that KR/s iterations suffice, where R can be taken to be of order Cmax + log K KnCmmax/E).
The second implication, on the form of the optimum, follows from the same considcrations as Lemma 3.1. The proof of Lemma 3.1 establishes direct link between solutions of (3) and (6). Specifically, if Piry is the optimal solution for 3, then
p" defined by p" (z) r(y) x) Piry (x)(() { being the fixed point of Algorithm 1. Thus.
(x) is the corresponding solution for (6). Notice that P* (x) x r(x) II (T(),, with
p² x (2) II (xx)
II
t 24 1, 2%) II vi (mx),
k=0
k=1
k=0
proving (12). Since we have already shown that the iterations of Algorithms and 2 agree, they have the same fixed points, and oz Br. desired.
C.4. Proof of Lemma 5.1
We compute:
/
6 (xx-1) =
(24
12-1
Σ
EEEXs
Σ
# E j=1 Maxii
Σ (K(#K) Σ
N E
1=1
The computation for rf analogous.
18

---

Smooth Schrödinger Bridges
C.5. Proof of Theorem 6.1
For each iteration of Algorithm 3, we run 2K sub-iterations for updating the left and right messages 43 and Pk. The intermediate step (18) for the calculation of Sg has complexity O(nM). The major update steps (19) and (20) involve matrix- multiplication with matrix of size nM x nM, which has complexity O(n ((n Therefore, the total complexity of the algorithm is given by O(T O(TKn
D. Practical implementation with the Haar basis
The Haar Wavelet basis stand out as natural choice in the context of this work because of its positivity properties: since the messages need to be positive, using orthonormal bases that consist of positive functions guarantees that the approximation is positive as well. Recall that 2g is finite collection with members and 3%- = mm-i. Set sequence of positive integers M,,-.., M with each number tied to a dimension of D/s. Our goal is to utilize M = M1 x ... x Mn-1 orthonormal polynomials to execute the algorithm. Let i E [M1]* x ... x [Mm-1]* with žn representing the n-th component of We define the following:
6 == Z
(39)
Kn 2
Kn 2
B
03 + (7n - 1)
ox + 5k
n=1
Here, Z& serves as the normalization constant ensuring that of has unit L ~o-norm; effectivvly, Z' is the reciprocal of the volume of the hypercube B The parameter òş is chosen via:
Kn 8z = = max (ar((ne)) (2)). ..... Var(m (m))
(40)
2
where C a hyper-parameter that one can tune and a typical choice is 3. With the set { determined, we show how to compute ro.... T the start of Algorithm By setting o₃ to be very small, we can leverage this approximation:
Ta(a 2541, 1,3) = // (xww) (v (d
(41)
= Vol(B db (T 2
(42)
where vi is the midpoint of the hypercube Bj. Remember, &x (Zz, 1 indicates the conditional density of the reference process r(zx+1 2 for & 2 1, and shifts to the unconditional joint density r(zo, z1) at k = 0. Evaluating the relevant Gaussian densities allows us to compute [ efficiently.
When Algorithm reaches convergence, the representation (12) can be used to efficiently manipulate p (for instance, to generate samples). particular, (12) shows that p inherits the Markov property of r, so to sample from p suffices compute the pairwise marginals 22, 1141), which are proportional
K
Σ
Σ
II Am
( (Z--1, zx) II
(xx))dyo ... (or YK
SEIK]
7,EX%
igk,k+1
This may be computed efficiently by repeatedly contracting one coordinate at time, in the left pass of Algorithm 2.
When Algorithm 3 is implemented using wavelets, p(zę, Z++1) can be approximated by ( 2, 3), which approxi- mates p( v #*+1, -=++1,111+1)) v With p(*K, ", 5) in hand, we propose two trajectory inference methods. The first method, which is stochastic, generates trajectory samples following the distribution p(zz, Z++ 1), as outlined in Algorithm 4.
The other way to rebuild the trajectories involves utilizing argmax operations on the belief tensor, and is summarized as Algorithm
19

---

Smooth Schrödinger Bridges
Algorithm 4 Trajectory Inference Through Sampling
Input: -o, Px-1 [ and E Xn
Marginalize Po to probability distribution on 2o {yol. denote it as Pint Sample Yo proportional Pint (
for k= 0,.... K do
Let ii be the index of yı in {y}]
Sample (1+11. proportional to (T (gy
end for
Return (.o. kK} and yoo. .... yK
Algorithm 5 Trajectory Inference Through ArgMax
Input: Po: ... PK-1 [2xxxxXxx M and 20 E An
Marginalize Po probability distribution on Ag x (yi), denote Pint Set Yo as the maximizer of Pint (20:
for A O,... K do
Let l be the index of yk in {y1)
Set (1K+1, as the maximizer of (T (.)
end for
Return fz...... Fx and {yo.... .... yK
E. Log-domain Implementation of Algorithm
In the execution of the algorithm, some values could fall below machine precision, triggering numerical problems. Take, for instance, the wavelet basis: we aim to estimate Гę by computing 4 at a specific point. However, yı or Y&+1 drifts significantly away from its mean, 4 ends up zero. For the sake numerical stability, we derive the log-domain implementation of Algorithm in the section. Recall that M is the total number of orthonormal polynomials that we use perform the approximate message passing algorithm. Define the log version of quantities
(xx) (xx)) and (x(x)) log(rx(xx))
(43)
(xx) == log(cs (I+)) and f (IR. Fx+1) == log (I' (k41,
(44)
For h: 0, K, the log-domain updates are summarized as follows:
(xx) log
Σ
Σ exp (GC+1(06+1) + P e (xx+1) f (x 22+2))
(45)
#ki1 E 111
(x())) = log
Σ
Σ
exp ( ( + Po (#1-1.
(46)
M Σ
exp (M((T)) + t((an))
Ca (xx) = log(cx(<x)) = log
Note that one can use the logsumexp function in scientific computing packages to implement the log-domain updates given above. And, for the wavelet decomposition, f(ie) (x+1) can be efficiently approximated by calling a SciPy scipy. stats. norm. 1ogpdf function.
F. Experiment Details
We first summarize the different data sets, metrics, and experimental findings.
For all experiments, we consider datasets with the number of points being constant at each step and each particle has uniform weight. For the kernel, we use either Matern kernel with 1.5 or 2.5.
20

---

Smooth Schrödinger Bridges
Figure 7. Visualization of matching results on Tri-stable diffusion and 2D Matern datasets.
F.1. Individual particle tracking
Smooth Gaussian process: (2D version of Figure 2): Each particle follows 2D Gaussian process with dimension-
independent covariance matrix generated by Matern kernel with v = 3.5, f = 1. Our algorithm can recover most of the trajectories correctly by selecting smoother trajectories than the other two algorithms.
1.
Tri-stable diffusion process: Particles are initialized randomly around zero by normal distribution and follow the evolution of deterministic dynamic system. This dynamic is considered in (Lavenant et al., 2024), each particle will be absorbed into one of the three attraction points with increasing velocity as the distance to the attraction point decreases. We find when particles are close to each other and have similar velocities, algorithm has the chance to lose track of the particles during sampling. The 22 matching performs almost perfectly in this scenario by choosing the closest particle.
N-body physical system: This dataset simulates the orbits of eight planets that circle a star and the task is to keep track of each planet's orbit in the system. The 2D version is considered in (Chewi et al., 2021), but they only consider smooth interpolations between objects and provide no guarantee of the correctness of matching. The visualization result is presented in Figure 4. The W2 matching and standard SB make many mistakes when trajectories are close to each other while our algorithm gives near-perfect tracking.
Our quantitative evaluation appears in Table 1. We use several metrics to evaluate our results: JumpP is the observed likelihood that a particle transitions to a different path at each time increment. 5p acc quantifies the proportion of sequence of five consecutive steps that remain on the same trajectory. The Mcan f2 metric is the average of all & distances between each sampled trajectory and the ground truth trajectory that maintains the longest continuous alignment with the sample.
F.2. Point cloud inference
Petal: The Petal dataset was introduced in (Huguct et al., 2022b) which mimics natural dynamics of cellular differentiation and bifurcation.
Converging Gaussian: This dataset was constructed by (Clancy & Suarez, 2022). It models cellular dynamics using Gaussian point clouds that split and converge over time. Points are sampled evenly around the center of the circles.
Embryoid Body: This se- sc-RNAseq dataset records statistical data of embryoid body (EB) differenniation over a period of 27 days with 5 snapshots taken between days of 0-3, 6-9, 12-15, 18-21, and 24-27. The dimension of the original data was reduced to 2 using a nonlinear dimensionality reduction technique called PHATE (Moon et al., 2019) in (Tong et al., 2020). We inherit this 2D projection dataset from them.
Dyngen Tree (Dyg Tree): This is another sc- sc-RNAseq dataset crafted by (Huguet et al., 2022b) using Dyngen (Cannoodt et al., 2021). The data is embedded into dimension 5 also using PHATE. This dataset contains one bifurcation and is considered to be more challenging than the Petal dataset.
Dyngen Cycle (Dyg Cycle): Another 10-D sc-Re RNAseq dataset inherited from (Banerjee et al.).
21

---

Smooth Schrödinger Bridges
We subsample each dataset moderate size and make ny constant over all timesteps.
Our quantitative evaluation uses the following metrics: W is the Wasserstein-li distance, Mc is the Maximum Mean Discrepancy with Gaussian kernel K (z, y) = (x - (1))////////222 and M, is the Maximum Mean Discrepancy with identity kernel K (zz y) (z, y)-
Table 3 presents additional results on other LOT tasks.
Table Performance comparison on LOT tasks at step j between our algorithm and state-of-art algorithms.
Dataset
Method
Wi(4)
Mc(4)
M
Petal
SSB (ours)
2.98e-2
3.26e-5
5,53e-3
j = 4
MIOFlow
1.76e-1
9.97e-3
1.14e-1
DMSB
2.63e-1
9.18e-3
1.15e-2
F&S
2.44e-2
4.41e-5
7.33e-3
EB
SSB (ours)
6.00e-2
3.88e-4
2.04e-2
j 3
MIOFlow
1.29c-1
3.31e-3
5,26e-2
DMSB
2.46e-1
4.69e-2
2.37c-1
F&S
7.42e-2
1.46e-3
3,88e-2
Dyg Tree
SSB (ours)
1.09e-1
4.67e-3
6.73e-2
j = 2
MIOFlow
2.23e-1
1.57e-2 *
1.34e-1
DMSB
F&S
9.23e-2
1.43e-3
3.93e-2
F33. Implementation details
We ran our experiments MIOFlow experiments, we utilized an NVIDIA A A100-SXM4-80GB. In tracking individual particles, standard SB with Brownian motion prior and W2 matching implemented by Python's 's OT library (Flamary et al., 2021). Check out Table for detailed dataset info.
x86-64 setup. Smooth SB (ours) and F&S do not require GPU support; for DMSB and
Table The table provides detailed statistics of the datasets we experimented on. Since the number of particles is constant over time, to denote the number of observations each time step and K stands for the total number of time steps excluding the initial observation. The Figure column stands for the label of figures where this dataset is presented in the main body text.
Name
K
n
Dimension
Figure
Author
Smooth Gaussian Process
20
20
1
2
Smooth Gaussian Process (2D)
20
25
2
7
Tri-stable Diffusion
20
20
2
7
(Lavenant et al., 2024)
N Body
50
8
3
Petal
5
40
2
1,5
(Huguet et al., 2022b)
Converging Gaussian
48
2
5
(Clancy & Suarez, 2022)
Embryoid Body
4
100
2
5
(Tong et al., 2020)
Dyngen Tree
40
5
5
(Huguet al., 2022b)
Dyngen Cycle
15
25
10
6
(Banerjee et al.)
Throughout our experiments, we employ either a Matern prior with parameters v specified in (5), utilizing the wavelet basis. We conduct 200 iterations of message- paassigg algorithms (T = 200). The initial covariance, Eo,o, is derived from the stationary distribution's covariance of the lifted Gaussian process corresponding to
1.5, or v = 2.5, 2, as
22

---

Smooth Schrödinger Bridges
the kernel, ensuring Ege remains steady over time (refer Equation (2.39) in (Saatçi, 2012) for stationary distribution calculation). We assume all observations occur equally the time period i 2, hence the observation duration is dt 2/K, which is utilized to compute covariance between timesteps. For datasets with dimensions greater than one, the kernel is assumed to be identical and independent for each dimension, simplifying the tensor I' by factoring it across dimensions, thereby reducing the computational load in the matrix-vector multiplication during the message-passing phase. We assume an equal allocation of coefficients per dimension, thus M, = (M à for each i = 1,..., d. For the kernel with 2.5, we express M; = (mi; m22) as tuple where M,- = m m represents the count of approximation coefficients the velocity and acceleration dimension. Each task has fixed M, choice except for the Dyngen Cycle dataset, with the kernel's variance calculated
dotaa where data is the dataset's ss standard deviation and is adjustable hyper-parametto
Detailed parameter settings for each task can be found in 5 and 6.
Table 5. Detailed parameter setting of our algorithm for all the experiments for Figures 2, 7, and 4
Dataset Name
v
M
M,
Dimension
Smooth Gaussian Process
2.5
200
(40,5)
1
Smooth Gaussian Process (2D)
2.5
1024
(8,4)
Tri-stable Diffusion
1.5
1024
32
2
1
N Body
2.5
1728
(4.3)
4
Petal
1.5
400
20
2
0.5
Converging Gaussian
1.5
900
30
0.5
Embryoid Body
1.5
144
12
2
Dyngen Tree
1.5
1024
0.5
Dyngen Cycle
1.5
1024
(4) x5 + (1) x 5
10
2.5
Table Detailed parameter setting of algorithm for the LOT tasks in Table 3 and 2 where we only used Matern kernel with and all the parameters are the same as the corresponding experiment in Table 5 except for c.
1.5
Petal j = Petal j = 4 EB j = 2 EB
Dyg Tree j = Dyg Tree j = 2
Dyg Cycle = 7
0.5
0.45
1
1
0.3
0.25
2.5
Implementation of Sampling: In the point cloud matching, where creating trajectories inferring positions of the left-out timestep requires generating new points, we use conditional Gaussian sampling. Initially, Yı for all data points are sampled from calculated probability tensors P&- Under wavelet basis, the sample we get pins down the range of the Y& we then calculate the conditional Gaussian density for y& given X₃ and simply use sampling and rejecting to get a valid velocity sample. Given that the lifted Gaussian process Markovian, we condition on two consecutive observations generate extra points for constructing trajectories or estimating positions at unobserved timesteps.
F.4. Details of other algorithms
Standard Schrödinger Bridges: The single hyper-parameter in regular SB is the scaling factor for the correlation between successive time steps, defined cov(i, min(i,.). Table shows the configuration.
Table 7. Detailed parameter setting of standard Schrödinger Bridges
Smooth Gaussian process (ID & 2D) N-Body
Tri-stable Diffusion
S
0.5
0.3
0.03
23

---

Smooth Schrödinger Bridges
Natural cubic splines were interpolation method of choice, detailed in Chewi et al., 2021, Appendix D, and wrote
F&S:
code for implementation. The visual results of the experiments in Figure 5 are shown in Figure & 9. In the case of the Dyngen dataset, F&S's matching is significantly poorer than what our algorithm achieves. This underscores a crucial flaw of the standard W2 matching it lacks the capacity to split mass effectively when the data distribution is uniform. The primary culprit is Dyngen's asymmetric bifurcation structure.
Petal
Corverging Cauestor
Embryal body
Observation Ae to
Phate dmn
Mlate dmn33
Phata d
Phate dm
Phace d
Figure 8. Visualization of trajectory inference various dataset by F&S, for the 5D Dyngen Tree data, visualize the 2D projection of the first three dimensions in the second row.
Dyngen
mexcmm
Figure 9, Visualization of trajectory inference on various datasct by FSI, for the 10D Dyngen Cycle data, we visualize the 2D projection of the first dimension against the second to fifth dimensions.
MIO Flow 2: We build upon the baseline settings from (Huguet et al., 2022b), with tuning for optimal performance. Here's the prime parameter configuration we achieved:,
1. Petal (complete): >; = 1e 2, Aa = 25, other variables adhere to the default Petal (hold one one_out = F False) setup.
2. Petal (leave one out): > = 1e - 3, n. local.clloooo local..poocss = 40, n. n.coochs = 0, remaining parameters align with the default Petal (hold_one__r) True) configuration.
3. Embryoid Body (complete & leave one out): gae embeded dim 2, other settings match the default Embryoid Body configuration.
24

---

Smooth Schrödinger Bridges
4. Converging Gaussian (complete): Ae le Aa 15, nelocceessoohss 20, all others follow the default Petal (hold one out False) guidelines.
5. Dyg Tree & Cycle (complete): parameters match the default Dyngen (hold_one_out False) specifications.
6. Dyg Tree & Cycle (leave one out): locolll Lpochs (hold one out True) settings.
0, n_epochs 50, rest conforms the default Dyngen
Visual comparisons for these experiments, as seen in Figure 5, are presented in Figure 10 & 11. DMSB By primarily
Petal
Corvering Cauaxier
Embrynid body
Oovervation Me t>
Eyngen
Phote din 3
priate dtm
Tilettiie
Phate dif
Maae dm
Phace d
Figure 10. Visualization of trajectory inference various dataset by MIO, for the 5D Dyngen Tree data, visualize the 2D projection of the first three dimensions the second row.
mhite di11
Tokcork
Thine
Prata
Figure I.. Visualization of trajectory inference various dataset by MIO, for the 10D Dyngen Cycle data, visualize the 2D projection of the first dimension against the second to fifth dimensions.
using the default parameters from (Chen et al., 2023), alongside some refinements, here's the top-notch parameter set:,
1. Petal (leave one out): epoch 2, num stage 13, num marg other values are consistent with the default Petal setup.
Thaps//gitthh.com//aaannnChhhhne
25

---

Smooth Schrödinger Bridges
2. Embryoid Body (leave one out): n-epoch 2, numm_sagge configuration.
13, other settings align with the default Petal
tuning this algorithm as each network requires approximately 5 hours for training. We anticipate enhanced DMSB performance with more extensive data and training duration. Moreover, the trajectory generation process of DMSB does not capture the geometry of the dataset well (see Figure 12 as illustration), hence we do not provide additional visualization for this algorithm here.
1-
-1
Figure 12. Trajectories drawn by DMSB reflect the underlying geometric structure of the Petal dataset well.
F.5. Additional Results and Analysis
Extended quantitative results one-by-one particle tracking: We provide more metrics in Table 8 for evaluating the performance of one-by-one particle tracking, which is an extend version of Table 1.
JumpP: the observed likelihood that particle transitions to different path at every increment in time.
3p acc: the proportion of a sequence of three consecutive steps that remain on the same trajectory.
öp acc: the proportion of a sequence of five consecutive steps that remain on the same trajectory.
Traj acc: the fraction of trajectories sampled that are correctly matched at each step.
For given sampled trajectory, we match it with ground-truth trajectory that maintains the longest continuous alignment with the sample. We say this ground truth trajectory is the matched trajectory of the sample.
Max & the maximum & distance found between a sampled trajectory and its matched ground truth across every sample.
Mean lz: the average & distance found between a sampled trajectory and its matched ground truth across every sample.
TrajKL: the KL-divergence between the matched trajectory histogram and the uniform distribution over the ground- truth trajectories.
Visualization of LOT tasks: We provide the visualization of the sampled point in the Leave-One-Out tasks in Figures 13 and 15 corresponding to results in Table 2. For all the cases, our algorithm is able to draw similar patterns to the ground truth.
Impact of scale hyper-parameter ct Recall that the variance of the Matérn kernel in our experiments is chosen to be Co for a positive hyperparameter c. We need to choose an appropriate such that the algorithm is able to search for
26

---

Smooth Schrödinger Bridges
Dataset
Model
JumpP
3p acc
5p acc
Traj acc
Max 6
Mean E
TrajKL,
ID Matérn
SSB (ours)
4.77e-2
0.930
0.893
0.530
1.912e-01
2.552e-03
0.000360
Process
BME
6.08e-1
0.191
0.079
0.001
1.463e+00
2.708e-01
0.036132
W2
2.18e-1
0.621
0.412
0.050
4.641e-01
1.526e-01
0.138629
Tri-stable
SSB (ours)
1.12e-1
0.852
0.798
0.489
1.020e-01
1.419e-02
0.011770
Diffusion
BME
5.20e-1
0.300
0.170
0.018
6.787e-01
6.595e-02
0.047383
W2
2.25e-2
0.971
0.956
0.800
1.377e-08
2.112e-09
0.000000
N body
SSB (ours)
5.00e-4
0.999
0.999
0.983
2.083e-01
5.475e-04
0.000100
ВME
1.14e-1
0.796
0.641
0.000
9.719e-01
5.426e-01
0.026030
W2
1.08c-1
0.804
0.649
0.000
8.679c-01
5,495e-01
0.173287
2D Matérn
SSB (ours)
5.60e-3
0.993
0.991
0.904
4.181e-03
3.571e-04
0.000000
Process
BME
1.38e-1
0.764
0.612
0.179
1.418e+00
2.819e-01
0.038109
W2
6.60e-2
0.867
0.760
0.360
1.015e+00
2.321557e-01
0. 10994
Table 8. Performance metrics for different models across various datasets for partiele matching.
the next matching point that constructs a smooth interpolation. If is too small, then the velocity space of the algorithm searching will not be enough incorporate the correct particles, this leads to the consequence that some of the particles will not be matched at all. If e is too large, then it will result in randomness in the algorithm, which will hurt the performance of OBO tracking and blur the underlying geometric structure in pattern matching, See Figure 16 on the Petal dataset for illustrations. In practical implementation, we increase if we find all the sampled velocities are concentrated at the center of the velocity grids and decrease if they are concentrated at the border instead.
Time and Space Complexity Analysis: our algorithm spends most of the time running the message passing algorithm. The running time complexity of one iteration of message passing is given by O(K 2²² Since we use dimension independent kernel, the tensor can be factorized across dimensions, we take M, = M and single iteration in this case has complexity OdTMM*+ + this allows us to take more coefficient when dimension increases. We give the run time record for one iteration message passing when 10 and N 20 in Table
Table 9, Running time record (in seconds) of one iteration of message passing, M, is taken to be round(M Vd), One observes that the running time decreases fixed M when increases due to the dimension independence.
d=1
d=2
d=3
M=1024
80.91
6.32
5.21
M=512
20.87
2.59
1.29
M=256
5,36
0.82
0.45
M=128
1.42
0.26
0.26
27

---

Smooth Schrödinger Bridges
Petbly 2
P
Ground Ruth
Ground Truhh Interned
Inferned
Embryddd Body
Embryoid Bocy
Ground Buth
Ground Tuth Inforeed Tinendep
Indereed
Figure 13. Visualization of sampled timestep in the LOT tasks for Dyngen datasets. The blue 'x' points are the left-out ground truth, the star points are inferred timestep by our algorithm.
Dyngen
Ground Tuth inferred Timesteo
Ground Tuuh Inferred Timestep
Ground Trush Timestep
Phate dm
Prote din
phote dim
Oyngon--2
Ground Tuth Infered Timmtes
Ground Tuth Infered Timemte
Ground Truth Inferned Timentap
Phote
Phote
Plote
Figure 14, Visualization of sampled timestep in the LOT tasks for Dyngen Tree dataset. The blue 'x' points are the left-out ground truth, star points are inferred timestep by our algorithm.
28

---

Smooth Schrödinger Bridges
Dyngon Cycle
Ground Tuth Timestop
Plane 2
Phate dmn
Phate dim
Plate der
Figure Visualization of sampled timestep in the LOT tasks Dyngen Cycle datasets, The blue points are left-out ground truth, the star points are inferred timestep by our algorithm.
c=1.1
c=025
0
c=1
c=2
Figure 16. Visualization of matching on the Petal dataset. We observe when = 0.1 some observations have no matching at all, while when 1,2, the matching contains randomness and mistakes.
29


Reviews:
[{"review_id": "gi64gi488U", "reviewer": "Reviewer_XtBd", "summary": "The paper proposes solving multi-marginal Schrödinger bridges w.r.t. a reference process based on autoregressive Gaussian process. Interpretation to phase space is constructed, which leads to a tractable algorithm based on probabilistic graphical models and belief propagation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# General remarks\n\nWe thank the reviewers for their helpful comments and questions. We are gratified that the reviewers found our method \"interesting\" (XtBd), \"novel,\" \"very useful\" (N7xY), and theoretically \"sound\" (okGt).\n\nWe would like to address an important issue that arose in several reviews: **dimension dependence**.\n\nFirst, we agree that we should highlight more explicitly that our experiments were conducted on low-dimensional data. We plan to include this fact prominently in our revision.\n\nSecond, the reviewers are correct that the dependence on dimension is poor. The running time scales quadratically with $M$, the number of coefficients used in the approximation algorithm (see Theorem 6.1). As mentioned in Section 6, $M$ will typically be exponential in the dimension, limiting our approach to relatively low dimensional problems.\n\nHowever, we offer several arguments supporting the interest of our method:\n\n1. A low-dimensional algorithm is still valuable in applications. For example, in the fluid dynamics and astronomical object tracking examples mentioned in the introduction, the data lies in dimension 2 or 3. Our method makes an important contribution to these inherently low-dimensional problems. Based on our experiments, smooth SBs are now the leading method for N-body tracking problems in 3 dimensions.\n\n2. Existing approaches often fail, even in low dimension. In the 5-dimensional Dyngen dataset, our algorithm (Figure 5) provides significantly more reasonable results than current leading approaches (F&S, Figure 7; MIO, Figure 8). Another algorithm, DMSB, fails to even converge (Table 2). This indicates the low-dimensional setting is far from solved.\n\n3. A small number of coefficients go a long way. In our experiments on the 5-dimensional Dyngen dataset, we use 1024 coefficients (only 4 per dimension), yet our approach still brings substantial benefits. We have also conducted experiments on 10-dimensional data using 1024 coefficients. Our method outperforms the other methods on this data in a leave-one-out task. We will include this new, higher-dimensional experiment in our revision.\n\n4. It is common to assume that high dimensional data has low intrinsic dimension. Many approaches to trajectory inference begin with non-linear dimension reduction. For example, the pipline developed by Schiebinger et al. (2019) begins by reducing dimension to 30. While this is still larger than we can handle, \"working up\" from low-dimensional examples can yield biologically relevant insights as computational power improves.\n\n# Response to XtBd\n\n**Comment:** \nIn the end of Sec 2.2 and Footnote 1, it's a bit unclear why \"smooth\" paths corresponds to stationary processes.\n\n**Response**\nThe footnote about stationary processes has been removed. To clarify, we've added a new lemma in the appendix proving that any Markov Gaussian process with differentiable paths is essentially trivial:\n\nLet $\\omega: t \\mapsto \\omega(t)$ be a real-valued Gaussian process that is Markovian and a.s. differentiable. For arbitrary $t \\ge 0$ such that $\\mathrm{Var}(\\omega(t)) > 0$, we have $\\omega(s) = \\mathbb{E}[\\omega(s)|\\omega(t)]$ for all $s \\ge 0$ a.s.\n\nIn other words, any such Gaussian process is essentially deterministic, indicating that obtaining smooth paths requires moving beyond Markov processes.\n\n**Comment:**\nRelation to [2006.14113]?\n\n**Response:**\nWe appreciate the reviewer mentioning this important work, which we failed to cite though we cited its companion works (''Learning hidden markov models from aggregate observations'' and ''Multimarginal optimal transport with a tree-structured cost and the Schrödinger bridge problem\"), which have the same authors.\n\nIn [2006.14113], the authors develop a belief propagation algorithm for multi-marginal optimal transport. The use of belief propagation for multi-marginal OT has a long history (Teh & Welling, 2001), and like [2006.14113], we use this connection for an efficient algorithm. However, we believe we're the first to exploit this connection to develop an algorithm for the Schrodinger Bridge problem with a smooth prior.\n\nOur Algorithm 2 is indeed analogous to Algorithm 2 in [2006.14113], tailored to our setting:\n1. The potential functions are:\n   - Between $\\eta_{k-1}$ and $\\eta_k$: Given by $\\Phi_k$ (defined in Section 4)\n   - Between $\\eta_k$ and $\\omega_k$: Given by the indicator function $\\mathbf{1}_{\\{z_k^{(0)} = x_k\\}}$\n2. Our message passing variables ($\\delta^-$, $\\delta^+$, $\\beta$, $\\gamma$) correspond to their factor-to-variable messages $m_{\\alpha \\rightarrow j}\n3. Our operators $\\mathcal{L}$, $\\mathcal{I}$, $\\mathcal{R}$ implement the operations in equations (38a,b)\n\nAn important technical difference is that since our $\\eta_k$ variables are continuous, the operators $\\mathcal{L}$, $\\mathcal{I}$, $\\mathcal{R}$ cannot be directly implemented. This fact requires us to develop an approximation scheme, which gives rise to our practical implementation detailed in Section 5.", "claims_and_evidence": "- In the end of Sec 2.2 and Footnote 1, it's a bit unclear why \"smooth\" paths corresponds to stationary processes. As mentioned by the authors before Sec 3, there are a few prior works that encourage smoothness through non-stationary processes. \n\nOverall, I think this paper presents an interesting theoretical finding yet falls short on the experiment sides.", "methods_and_evaluation_criteria": "Y", "theoretical_claims": "Y", "experimental_designs_or_analyses": "- Experiment results are in rather low dimension (<= 5). Dimension reduction is performed on some sc-RNAseq datasets, e.g., Embryoid Body was initially reported in MIOFlow with dimension 200. How does the complexity of the proposed method scales in terms of dimension?", "relation_to_broader_scientific_literature": "Numerical methods for solving trajectory inference problems could be beneficial, as these problems are widespread across various scientific domains.", "essential_references_not_discussed": "N", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "Typo in L268: Sec 3 should be Fig 3", "questions_for_authors": "Can the authors comment on the relation to https://arxiv.org/pdf/2006.14113? This prior work also introduced a belief-propagation based method for solving multi-marginal OT.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sXHA0QA12y", "reviewer": "Reviewer_N7xY", "summary": "The authors proposed a novel method to learn smooth trajectories in an SB problem, extending the usual SB method to allow non-Markovian reference by lifting it to phase space, that also effectively extending momentum SB. The method is accompanied with an approximated belief propagation algorithm.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to N7xY\n\nWe thank the reviewer for these insightful comments.\n\n**Comment:**\nWhat is the assumptions on the marginals in Lemma 2.1 (and 3.1)? In fact I would like the authors to double check if there is any missing assumptions for the marginals across the paper.\n\n**Response:**\nThank you for allowing us to clarify this important point, which we'll make clearer in the revision.\n\nA complete proof of Lemma 2.1 is now included in the appendix. Lemma 2.1 consists of two parts:\n\n1. The first part demonstrates that the multi-marginal Schrödinger bridge problem with finite marginal constraints (1) can be reduced to a finite-dimensional KL divergence minimization problem (2). This reduction requires no assumptions about the marginals.\n\n2. The second part addresses a technical challenge: we want to define a multi-marginal Schrödinger bridge problem with the prior $R$ being a GAP and the marginals being discrete. However, the KL divergence in the optimization problem is ill-defined when $R$ has absolutely continuous marginals and the marginals are discrete. We address this by proving that problem (2) is equivalent to formulation (3) *if* the marginals have densities. Since (3) is well defined even for discrete marginals, we take (3) as the **definition** of the multi-marginal Schrodinger Bridge for general marginal distributions. This is in fact a standard approach in statistical analysis of the Schrodinger Bridge problem, see, e.g., Pooladian & Niles-Weed (2024).\n\nThe equivalence between (2) and (3) uses two assumptions:\n* The marginals $\\mu_k$ are absolutely continuous with respect to the Lebesgue measure, with finite entropy $\\int \\mu_k(x) \\log \\mu_k(x) dx$\n* The joint marginals of the prior $R$ are absolutely continuous\n\nIn the remainder of the paper, we adopt the setting described in the beginning of Section 2.2 where the marginals are sampled data (i.e., $\\mu_k$ is discrete). In particular, this is true of Lemma 3.1. Note that, before the lemma, we defined $p(x, y)$ to refer to a \"mixed\" discrete-continuous density, where $x$ coordinates are discrete but $y$ coordinates are continuous. We will emphasize this point in the revision.\n\n**Comment:**\nIs the Markovian assumption in the GAP prior essential? How about the Gaussianity? \n\n**Response:**\nThe GAP itself is not Markov; however, the Markovian structure of the **lifted** process is crucial for our algorithm because:\n    \n1. It enables factorization of the joint distribution into pairwise terms, allowing efficient message passing.\n\n2. It ensures the solution $p(z)$ of (6) is Markovian (by Theorem 4.1), enabling efficient trajectory sampling.\n\nGaussianity in the GAP prior provides computational efficiency through closed-form message passing updates but isn't essential. The method can be generalized to work with any process that can be lifted to a Markov process on the phase space by replacing the potential functions between $\\eta_k$ and $\\eta_{k+1}$ with the conditional density $r(\\eta_{k+1} | \\eta_k)$ of the desired prior process.\n\n**Comment:**\nDoes the algorithm translate to general missing data problem?\n\n**Response:**\n**Yes**. As the reviewer mentioned, not observing higher-order derivatives is a form of missing data handled by our algorithm. Other forms of missing data can be incorporated by modifying the factor nodes $\\alpha_k$ in the graphical model (Figure 3). Currently, these nodes are indicators enforcing that only positions $\\omega_k$ are observed. By modifying these factor nodes, one can incorporate other types of missingness in the observations. We plan to pursue this in future work.\n\n**Comment:**\nIs it true that an inferred trajectory necessarily pass one particle at each time point? And does is makes sense in the context of single cell example? Is there an assumption that each time snapshots has same number of particles? \n\n**Response:**\nYes, in the current formulation, inferred trajectories must pass through exactly one particle at each observed time point, due to the marginal constraints.\n\nFor point cloud matching applications, while tracking individual \"particles\" may not be physically meaningful, our algorithm remains valuable for inferring collective motion patterns:\n\n1. We can infer velocity and acceleration fields that characterize point cloud dynamics\n2. We can interpolate the point cloud distribution at unobserved times by sampling from the solution\n\nOur results in Table 2 show that this approach yields good results.\n\nRelaxing the requirement that trajectories pass through observed points can be achieved by modifying the factor nodes $\\alpha_k$ to model observational noise. This allows inferred trajectories to deviate from exact observed positions. We've added this extension as a remark in the revision, and plan to pursue it in follow-up work.\n\nOur method does not require the same number of particles across different time snapshots. Both the theoretical framework and algorithm support arbitrary discrete marginal distributions.", "claims_and_evidence": "The idea is well supported with theoretical results and empirical claim of algorithm is support by experimental evidences.\n\nThe only part being a bit hard to follow is that whether marginals are the continuous, population version or the discrete, sample version in theories. It might be helpful to spell that out more explicitly when using them.", "methods_and_evaluation_criteria": "The evaluation criteria makes sense for empirical tests. The only potential complain I have is that the dimension tested are relatively low.", "theoretical_claims": "I tried to check the proofs and did not find obvious errors. However Lemma 2.1 only has proof sketch and I want to see a fully spelled out proof for completeness. Especially spell out the assumptions on $\\mu_k$, which appears to be continuous in the sketch but said to be general in lemma statement. In fact I would like the authors to double check if there is any missing assumptions for the marginals across the paper. \n\nThis is the main reason I put the current score and I am more than happy to change if assumptions are better spelled out.", "experimental_designs_or_analyses": "I checked all the experiment and did not find obvious issues.", "relation_to_broader_scientific_literature": "The method extend momentum SB that are SoTA for smooth trajectory inference and can use quite general non-Markovian references. I think the latter is more important if it can be served as a foundation for more generalization to missing dimension problems.", "essential_references_not_discussed": "Not I am aware of immediately.", "other_strengths_and_weaknesses": "Strengths\n- The HMM view is very useful and I think it has potential to be extended to other setups beyond smoothness. \nWeaknesses:\n- I am unsure how scalable the algorithm to dimension is. Single cell data can have much higher dimensions. Even if not, having an algorithm work in several dimensions can be useful for a range of applications still.", "other_comments_or_suggestions": "- There are missing cross references at then end of page 19 and caption of Table 16.\n- I would like to have a full proof instead of proof sketch of Lemma 2.1, mostly to be sure about assumptions went in.", "questions_for_authors": "These questions might be out of the scope of the current paper but I am curious\n\n1) it appears to me that the lifted SB problem having Markovian reference is key for the algorithm. It might be viewed as a missing data problem that we did not observe any higher order derivatives (which is momentum for momentum SB problem). Does the algorithm translate to general missing data problem that 1) some dimension are missing and 2) if they are observed the prior should be Markovian? Or is it also essential to have Gaussianity in GAP? \n\n2) Is it true that an inferred trajectory necessarily pass one particle at each time point? This seems inherent from lemma 2.1. I understand this corresponding to interpolating empirical distributions. I wonder however, is this a desired property for all applications, e.g., in single cell examples/point cloud task one loses particle identity by killing the cell, so each (latent) trajectory is evaluated once and the chance a latent trajectory pass two observational data is zero, but in application e.g., individual particle tracking and N body problems one loses identity by permutation and we kept the same set of particles and we indeed want trajectories to pass exact one particle each snapshot. Is there an assumption that each time snapshots has same number of particles?\n\n3) Related to 2), Lemma 2.1 (maybe also 3.1?) seems to assumed $\\mu_k$ to be absolutely continuous, I wonder from when on the theories start to use empirical measures instead of population version?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IZEj5rdn1r", "reviewer": "Reviewer_okGt", "summary": "**Disclaimer: Despite the forthcoming criticisms, I find this paper intriguing and recommend its acceptance.**\n\n\n------------------------------\n\nThis paper introduces a class of smooth Gaussian processes as priors for Schrödinger bridges and designs efficient algorithms for their computation. A key insight is that, while these processes are inherently non-Markovian—potentially leading to exponential-time complexity in terms of the number of marginals—the proposed class can be transformed into a **Markovian** form. This transformation enables the use of efficient message-passing algorithms. The authors validated their methods through simulations and experiments on low-dimensional real-world problems.\n\n## update after rebuttal\n\nI would like to thank the authors for their detailed rebuttal. As stated previously, I find this paper interesting and continue to advocate for its acceptance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for these insightful comments.\n\n**Comment:** Which processes are absolutely continuous with respect to the GAP? Is the minimizer of the SB problem with GAP prior smooth?\n\n**Response:** When $R$ is a smooth prior, the solution to (1) is also smooth, in the same sense. To be more specific, a sample path of an order $m$ GAP is $m-1$-times differentiable with probabiliy 1. In other words, there is an event $A \\subseteq \\Omega$ such that each $\\omega \\in A$ lies in $C^{m-1}$ and $R(\\Omega \\setminus A) = 0$. Since the solution $P$ to (1) is necessarily absolutely continuous with respect to $R$, we must have that $P(\\Omega \\setminus A) = 0$ as well, so that $P$ is concentrated on $C^{m-1}$ paths. We will clarify this.\n\nAlternatively, the explicit form of the solution to the smooth SB problem given in Theorem 4.1 shows that the optimal $P$ is a mixture of Gaussian processes, which inherit the smoothness of $R$.\n\n**Comment**: Recovering the dynamic SB solution from the static one requires an additional flow-matching-style training step, which can be computationally expensive.\n\n**Response:** We would like to draw a distinction between two senses of \"recovering\" the dynamic SB, focusing for a moment on the standard SB. The *solutions* of the static and dyanmic SB problems actually coincide (by convolving the solution to the static problem with a Brownian bridge, see Leonard 2014). This is the argument we use in the proof of Lemma 2.1. In this sense, the dynamic SB is \"recovered\" directly from the static problem, and in particular, this allows easy sampling from the dynamic SB. However, the reviewer is correct that recovering the dynamic SB *as a stochastic process* (e.g., recovering the drifts in the corresponding SDE) is not so direct, though it can still be accomplished from the static problem (as in Pooladian & Niles-Weed, 2024).\n\nIn our work, we recover a smooth dynamic SB in the first sense (with very reasonable computational complexity): given the solution to the static problem, we \"recover\" the corresponding dynamic solution in the sense that we can easily sample from it (using standard techniques in GP regression). In particular, we can sample from intermediate, unobserved times. (See our \"leave-one-out\" experiments in Tables 2 and 3, which shows that this leads to accurate reconstruction on synthetic data.)\n\nIt is less clear how to recover it in the second sense, since the smooth dynamic SB is not given by the solution to an SDE. We agree that it would be interesting to pursue this direction in future work.\n\n\n**Comment:** What is the practical motivation for using smooth priors?\n\n**Response:** Our motivation is twofold:\n\n1. In many systems, the only physically relevant models involve particles whose velocities vary continuously, i.e., whose paths are at least $C^1$. This is never the case with a BM prior. As our experiments show (e.g., Figure 4), this is particularly relevant when trajectories of two particles cross and separate: the BM prior easily loses the identity of particles between timesteps when two trajectories intersect, but assuming that the velocities are continuous allows the particles to be identified.\n\n2. Existing work on trajectory inference largely focuses on smooth trajectories (see [F&S](https://proceedings.mlr.press/v130/chewi21a.html), [WLR](https://arxiv.org/abs/2405.19679) as examples), not for physical but for essentially *statistical* reasons. Like classic cubic splines, at a heuristic level, fitting data with smooth paths helps to ameliorate overfitting and ensure interpretability of the resulting estimates. Making this intuition rigorous in a particular statistical setting is a very interesting open problem.\n\n**Comment:**\nThe most significant limitation of the proposed framework is that it essentially reduces to a one-dimensional Schrödinger bridge. The theory does not extend to high-dimensional Gaussian processes with dependent coordinates.\n\n**Response:**\nThe fact that the theory does not extend to processes with dependent coordinates is not quite true.\nFirst, in Theorem 2.3, the covariance matrix $\\sigma$ need not be diagonal; this allows the coordinates of the white-noise process defining the GAP to be correlated. More generally, if we consider a large covariance matrix $\\mathbb{R}^{md \\times md}$ that describes the covariance among all entries of hidden variables across dimensions, it is not necessary that this matrix be block diagonal (corresponding to independent coordinates). From an algorithmic perspective, this only changes the value of $\\Gamma$ (see equations 19 and 20); the algorithm need not be changed.\n\nHowever, we agree that from a computational perspective, the use of independent coordinates is beneficial, since the tensor $\\Gamma$ factors. This greatly reduces the space complexity of storing $\\Gamma$, and simplifies some message passing updates.\n\nWe also point the reviewer to our general comments above on dimension dependence.", "claims_and_evidence": "### Unsupported Claims  \n\nSeveral claims in the paper lack sufficient justification:  \n\n1. **Motivation for Using Smooth Priors**  \n   The core motivation of the paper is to replace the standard Brownian motion-driven prior with a smooth one. While this idea seems reasonable, it is not thoroughly justified. Many trajectory inference problems indeed require smooth trajectories, but the authors should provide concrete examples illustrating when and why this is beneficial.  \n\n   Furthermore, the authors should explicitly state that their RNA-seq experiments are only toy versions, as they are reduced to dimensions 2 and 5 (a fact buried in the appendix). In reality, biological processes often exhibit strong microscopic randomness, making it unclear why smooth trajectories are preferable in this context.  \n\n2. **Continuity of Processes Relative to the Prior**  \n   A key missing question is: *What processes are absolutely continuous relative to the prior introduced in this paper?* This is crucial because, even if the prior is smooth, the minimizers of the Schrödinger bridge (SB) problem might still concentrate on non-smooth trajectories. If this is the case, it challenges the authors’ claim that they are performing *smooth* trajectory inference.  \n\n3. **Static vs. Dynamic Schrödinger Bridges**  \n   The authors appear to focus only on the **static** Schrödinger bridge, which corresponds to multi-marginal optimal transport with entropy regularization. While it is well known that the **values** of the static and dynamic SB problems coincide, recovering the latter from the former requires an additional flow-matching-style training step, which can be computationally expensive. This raises concerns about whether the comparisons in the experimental section are fair, especially regarding diffusion-based methods that train dynamic SBs directly.  \n\n4. **A Fundamental Limitation: Essentially a One-Dimensional Schrödinger Bridge**  \n   The most significant limitation of the proposed framework is that it **essentially reduces to a one-dimensional Schrödinger bridge**. The theory does not extend to high-dimensional Gaussian processes with dependent coordinates, which is a major drawback for real-world applications. This limitation likely explains why the experiments are restricted to problems of at most dimension 5.", "methods_and_evaluation_criteria": "The experiments are primarily conducted on synthetic data and low-dimensional settings. The efficiency of the proposed algorithm in high-dimensional scenarios remains uncertain.", "theoretical_claims": "The theoretical claims are sound and interesting.", "experimental_designs_or_analyses": "The experimental design seems flawed; see my criticisms in **Claims And Evidence** and **Methods And Evaluation Criteria**.", "relation_to_broader_scientific_literature": "The key contribution of this paper is the introduction of a novel class of smooth priors for the Schrödinger bridge problem, which, to the best of my assessment, is both original and valuable.", "essential_references_not_discussed": "The references are appropriately cited.", "other_strengths_and_weaknesses": "I have outlined several criticisms in **Claims and Evidence**. However, I believe this paper makes an interesting contribution due to the following strengths:  \n\n- **Novelty of the setting** – As the authors point out, introducing a smooth prior for the Schrödinger bridge problem is a novel and underexplored direction.  \n- **Simplicity of the algorithm** – The proposed algorithm is a straightforward modification of existing methods for Sinkhorn-type problems, making it both easy to understand and implement.  \n- **Clarity of presentation** – The paper is well-written and effectively communicates its ideas.", "other_comments_or_suggestions": "There is a \"Table ??\" at the bottom of p.19, line 1043.", "questions_for_authors": "Please see **Claims And Evidence**.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "w1n4NMPPSS", "reviewer": "Reviewer_1euq", "summary": "The authors proposed the Schrodinger Bridge (SB) with smooth priors guided by the Gaussian process to reduce the exponential cost in K for solving multi-marginal SB problems. The problem is first discretized (suffers from exponential cost) and then lifted to high dimensions and solved by belief propagation methods, resulting in polynomial costs.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for these  comments.\n\n**Comment:**\n\nCan we use the probability flow ODE to avoid Brownian motion?\n\n**Response:**\n\nWhile probability flow ODEs are indeed powerful tools, they are not ideal for our specific goals for several reasons:\n\n1. Inferring Individual Trajectories:\n    - Probability flow ODEs focus on evolving the entire distribution and don't preserve individual particle identities or trajectories\n    - Our method explicitly tracks individual particles, which is crucial for understanding the motion of individual particles (e.g., in N-body systems)\n\n2. Statistical Framework:\n    - The smooth Schrödinger bridge has an elegant quasi-Bayesian interpretation: the Schrödinger bridge matches the observed marginal distributions while keeping particle trajectories as faithful as possible to the prior $R$\n    - Probability flow ODEs lack such statistical guarantees when used for smoothing\n\n3. Other Advantages of our Method:\n    - Our method naturally provides velocity and acceleration estimates. These emerge directly from the HMM model we use\n    - Post-hoc smoothing via ODEs would require additional assumptions and computation\n\nThis makes our approach more suitable for applications requiring detailed trajectory analysis and closer adherence to the prior distribution. \n\n**Comment:**\n\nCan we tackle the continuous multi-marginal SB directly without tackling the discrete optimal transport? \n\n**Response:** \n\nThe goal of this paper is to infer particle trajectories that are as close to the prior path distribution as possible while matching the observed empirical distributions. Empirical distributions are discrete by nature, so a discrete multi-marginal Schrödinger bridge problem is a natural starting point. However, our algorithm can be extended to continuous marginals (population distribution) by using the following approach:\n1. Replace the messages $\\beta$ and $\\gamma$ with their continuous versions\n2. Replace sums with integrals in (8) on the right-hand side of the equation\n\nThis continuous extension of our algorithm yields scaling functions $\\beta^*_0, ..., \\beta^*_K$ that correspond to the multi-marginal Sinkhorn algorithm for continuous distributions. Crucially, this modification preserves the computational efficiency of the discrete version, reducing the exponential complexity of multi-marginal Sinkhorn to polynomial time. Furthermore, similar to how we represent $\\delta$'s using orthonormal bases in Section 5, the continuous $\\beta$ functions can also be expressed through basis expansion, maintaining both mathematical rigor and computational tractability.\n\n**Comment:**\n\nI suggest the authors make a diagram of the flow of methods to facilitate the reading.\n\n**Response:**\nWe appreciate this suggestion and will add a diagram to the revised version of the paper. The diagram will illustrate the flow of methods, including the relationship between the multi-marginal Schrödinger bridge problem, the static problem, and the lifted space. This will help clarify the overall structure and flow of our approach for readers.", "claims_and_evidence": "The paper is a bit hard to read.\n\nThe multi-marginal SB problem is \n1) first discretized into discrete space in sec 2; \n2) augmented with a new continuous variable and was related to a sequential conditional structure via the Gauss-Markov property; \n3) was solved using a linear Belief Propagation (BP); \n4) a practical (BP).\n\nOverall, I think the method is quite hard to understand, especially when I am not an expert in the Gaussian process and have no knowledge in  belief propagation.", "methods_and_evaluation_criteria": "NA", "theoretical_claims": "Didn't check.", "experimental_designs_or_analyses": "I don't know much about the usage of multi-marginal SB or the applications in RNA-Seq.", "relation_to_broader_scientific_literature": "This paper aims to alleviate the exponential cost when it comes to many marginals.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "I am a bit confused by the usage of first considering the discrete state space and then uplifting the problem to a higher dimension and even augmenting with some continuous variables.\n\nIt seems like we don't need the discretization step. As such, the exponential cost does not exist actually when we don't consider the discrete-type problems.", "other_comments_or_suggestions": "The first paragraph in section 3 is quite important. I suggest the authors make a diagram of the flow of methods to facilitate the reading.", "questions_for_authors": "1) limitation: its trajectories inherit the roughness of Brownian paths, leading to noisier estimators and less interpretable posterior paths. This problem can be easily solved by the probability flow ODE, why cannot we use the prob flow ODE to avoid the Brownian motion?\n\n2) Can we tackle the continuous multi-marginal SB directly without tackling the discrete optimal transport? So many tools have been used, which may affect the adoption of this method.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper considers trajectory inference in the space of probability measures. The traditional approach of multi-marginal Schrodinger bridge is improved by the introduction of a class of smooth Gaussian processes as priors. Efficient algorithms are developed based on lifting and, consequently, (approximate) belief propagation. Reviewers and AC largely found the idea interesting, the method effective, and the work rigorous, although multiplier reviewers also had concerns about scalability with dimension. Overall, the merits overweigh imperfection, and I recommend acceptance as poster (not higher because the concerns around scalability have not been completely resolved).
Decision: Accept (poster)