Generalized Schrödinger Bridge Matching
Authors: Guan-Horng Liu
Yaron Lipman
Maximilian Nickel
Brian Karrer
Evangelos Theodorou
Ricky T. Q. Chen
Venue: ICLR 2024 poster (2024)
Keywords: Distribution matching
diffusion models
generalized Schrödinger bridge
stochastic optimal control
Primary area: generative models
TL;DR: We propose Generalized Schrödinger Bridge Matching, a new distribution matching algorithm for training diffusion models enhanced with task-specific optimality structures.

Abstract: Modern distribution matching algorithms for training diffusion or flow models directly prescribe the time evolution of the marginal distributions between two boundary distributions. In this work, we consider a generalized distribution matching setup, where these marginals are only implicitly described as a solution to some task-specific objective function. The problem setup, known as the Generalized Schrödinger Bridge (GSB), appears prevalently in many scientific areas both within and without machine learning. We propose Generalized Schödinger Bridge Matching (GSBM), a new matching algorithm inspired by recent advances, generalizing them beyond kinetic energy minimization and to account for nonlinear state costs. We show that such a generalization can be cast as solving conditional stochastic optimal control, for which efficient variational approximations can be used, and further debiased with the aid of path integral theory. Compared to prior methods for solving GSB problems, our GSBM algorithm always preserves a feasible transport map between the boundary distributions throughout training, thereby enabling stable convergence and significantly improved scalability. We empirically validate our claims on an extensive suite of experimental setups, including crowd navigation, opinion depolarization, LiDAR manifolds, and image domain transfer. Our work brings new algorithmic opportunities for training diffusion models enhanced with task-specific optimality structures.

Published as a conference paper at ICLR 2024
GENERALIZED SCHRÖDINGER BRIDGE MATCHING
Guan-Horng Liu'; Yaron Lipman"3, 2-3 Maximilian Nickel', Brian Karrer', Evangelos A. Theodorou', Ricky T. Q. Chen'
'Georgia Institute of Technology Weizmann Institute of Science FFAIR, Meta {ghliu, evangelos. theodorou @eattec edu (ylipman, maxn, bri briankarrer, rtqichen) @meta.com com
ABSTRACT
Modern distribution matching algorithms for training diffusion or flow models directly prescribe the time evolution of the marginal distributions between two boundary distributions. In this work, we consider generalized distribution matching setup, where these marginals are only implicitly described as a solution to some task-specific objective function. The problem setup, known as the Gen- eralized Schrödinger Bridge (GSB), appears prevalently in many scientific areas both within and without machine learning. We propose Generalized Schrödinger Bridge Matching (GSBM), new matching algorithm inspired by recent ad- vances, generalizing them beyond kinetic energy minimization and to account for task-specific state costs. We show that such generalization conditional stochastic optimal control, for which efficient variational approxima- tions can be used, and further debiased with the aid of path integral theory. Com- pared to prior methods for solving GSB problems, our GSBM algorithm better preserves feasible transport map between the boundary distributions throughout training, thereby enabling stable convergence and significantly improved scala- bility. We empirically validate our claims on an extensive suite of experimental setups, including crowd navigation, opinion depolarization, LiDAR manifolds, and image domain transfer. Our work brings new algorithmic opportunities for training diffusion models enhanced with task-specific optimality structures.
be cast solving
INTRODUCTION
The distribution matching problem- learning transport maps that match specific distributions is a ubiquitous problem setup that appears in many areas machine learning, with extensive appli- cations in optimal transport (Peyré & Cuturi, 2017), domain adaptation (Wang & Deng, 2018), and generative modeling (Sohl-Dickstein et al., 2015; Chen et al., 2018). The tasks are often cast as learning mappings, Xo +> X1, such that Xo ~ μ X1 butions μt, v. For instance, diffusion models construct the mapping as the solution to a stochastic differential equation (SDE) whose drift uf (X,) IRd x [0, 1] + R2 is parameterized by 6:
follow the (unknown) laws of two distri-
dx, ui (X.)dt + adW,, Xo μ, X1
(I)
The marginal density pt induced by (1) evolves the Fokker Plank equation (FPE; Risken (1996)),
ge (x, = - (u, (x 2 (X,, Po = H.Pi =
(2)
and prescribing Pt with fixed uniquely determines a family of parametric SDEs in (1). Indeed, modern successes of diffusion models in synthesizing high-fidelity data (Song et al., 2021; Dhari- wal & Nichol, 2021) attributed, partly, to constructing P. mixture of tractable conditional probability paths (Liu et al., 2023b; Albergo et al., 2023; Tong et al., 2023). These tractabilities enable scalable algorithms that "match of given p."". hereafter referred to matching algorithms.
Alternatively, one may also specify Pt implicitly as the optimal solution to some objective func- tion, with examples such optimal transport (OT; Villani et al. (2009)), or the Schrödinger Bridge
Work done in part as a research intern at FAIR, Meta.
We adopt constant R throughout the paper but note that all analysis generalize time-dependent dt.

---

Published as a conference paper at ICLR 2024
problem (SB; Schrödinger (1931); Fortet (1940); De Bortoli al. (2021)). SB generalizes stan- dard diffusion models to arbitrary μ and v with fully nonlinear stochastic processes and, among all possible SDEs that match between μ and v, seeks the unique un that minimizes the kinetic energy.
While finding the transport with minimal kinetic energy can be motivated from statistical physics or entropic optimal transport (Peyré & Cuturi, 2019; Vargas et al., 2021), with kinetic energy often being correlated with sampling efficiency in generative modeling (Chen et al., 2022: Shaul et al., 2023), it nevertheless limits the flexibility in the design of "optimality". Indeed, kinetic energy corresponds to the squared-Euchi -uurddnn cost in OT, which, despite its popularity, is merely one among numerous alternatives available for deployment (Di Marino et al., 2017). On one hand, it remains debatable whether the & distance defined in the original data space (e.g. pixel space for images), as opposed to other metrics, are best suited for quantifying the optimality of transport maps. On the other hand, distribution matching in general scientific domains, such as population modeling (Ruthotto et al., 2020), robot navigation (L.iu al., 2018), or molecule simulation (Noé et al., 2020), often involves more complex optimality conditions that require more general algorithms to handle.
To this end, we advocate a generalized setup for distribution matching, previously introduced as the Generalized Schrödinger Bridge problem (GSB; Chen et al. (2015); Chen (2023); Liu et al. (2022)):
1 (x))) + vi(x,) dt subject to (1) or, equivalently, (2).
min 6
E
(3)
GSB is a distribution matching problem- as it still seeks a diffusion model (1) that transports μ to v. Yet, in contrast standard SB, the objective of GSB involves additional state cost V2 which affects the solution by quantifying penalty equivalently a
to general decision- making problems. This state cost can also include distributional properties of the random variable X2. Examples of Ve include, e.g., the mean-field interaction in opinion propagation (Gaitonde et al., 2021), quantum potential (Philippidis et al., 1979), or a geometric prior.
optimality (3) and feasibility (2). Within the set of feasible solutions that satisfy (2), GSB considers the one with the lowest objective in (3) to be optimal. Therefore, it is essential to develop algorithms that search within the feasible set for the optimal solution. Unfortunately, existing methods that approximate solutions to (3) either require relaxing feasibility (Koshizuka & Sato, 2023), or, following the design of Sinkhorn algo- rithms (Cuturi, 2013), prioritize optimality over feasibility. While an exciting line of new matching algorithms (Peluchetti, 2023; Shi et al., 2023) has been developed for SB, i.e., when V22 == 0, to ensure that the solutions are proximity to the feasible set throughout training, it remains unclear whether, how, these "SB Matching" (SBM) algorithms can be extended to handle nontrivial V/-
Solving GSB problems involves addressing two distinct aspects:
We propose Generalized Schrödinger Bridge Matching (GSBM), a new matching algorithm that generalizes SBM to nontrivial V,, We discuss how such generalization can be tied to conditional stochastic optimal control (CondSOC) problem, from which existing SBM algorithms can be derived special cases. We develop scalable solvers to the CondSOC using Gaussian path approximation, further debiased with path integral theory (Kappen, 2005). GSBM inherits a similar algorithmic characterization to its ancestors (Liu et al., 2023b; Shi et al., 2023), in that, during the optimization process, the learned p and pl induced by the subsequent solutions of uf remain close to the boundary marginals (H,, v) and preserve them exactly under stricter theoretical conditions; see Sec. 3.1 for details. This distinguishes GSBM from prior methods (e.g.. Liu et al. (2022)) that learn approximate solutions to the same problem (3) but whose subsequent solutions only approach (1,, v) after final convergence. This further results in a framework that relies solely on samples from μ. 1/ wittout knowing their densitics-and enjoys stable convergence, making it suitable for high dimensional applications. A note on the connection to stochastic optimal control and related works (Liu et al., 2022) can be found in Appendix A. Summarizing, we present the following contributions:
We propose GSBM, new matching algorithm for learning diffusion models between two distri- butions that also respect some task lusk-specific optimality structures in (3) via specifying V,-
GSBM casts recent matching methods as conditional stochastic optimal control problems, from which nontrivial Vi can be incorporated and solved at scale via Gaussian path approximation.
Compared to prior methods, e.g., Liu et al. (2022), that also provide approximate solutions to (3), GSBM enjoys stable convergence, improved scalability, and, crucially, maintains a transport map that remains much closer to the distribution boundaries throughout the entire training.
2

---

Published as a conference paper at ICLR 2024
Algorithm match (implicit)
Algorithm match (explicit)
Require: s.t. Po = p4 =
Require: := Ep. [Pt]a.1] s.t. Po =V, 4loo11 repeat
repeat
Sample Xu ~ po. X2 ~ Pt. X1 P
Sample Xo, X1 ~ pa,1, X₂ ~ Pr
Compute 1 given (Xo, Xt, X1) Take gradient step w.r.t. Lesplicn (f) until converges return x??"
Take gradient step w.r.t. Limplici (0) until converges return s
Through extensive experiments, we showcase GSBM's remarkable capabilities across variety of distribution matching problems, ranging from standard crowd navigation and 3D navigation over LiDAR manifolds, to high- dinhnsninnsi opinion modeling and unpaired image translation.
2 PRELIMINARIES: MATCHING DIFFUSION MODELS GIVEN PROB. PATHS
As mentioned in Sec. 1, the goal of matching algorithms is to learn SDE parametrized with uf such that its FPE (2) matches some prescribed marginal Pt for all [0, 1]. In this section, we review two classes of matching algorithms that will play crucial roles in the development of our GSBM.
Entropic Action Matching (implicit). This is recently proposed matching method (Neklyudov et al., 2023) that learns the unique gradient field governing an FPE prescribed by pt. Specifically, let
(X,): : Fd [0, 1] + be parametrized function, then the unique gradient field Vs (X,) by
which the probability density of the pushforward from μ at time 0 to / matches pt m--miniiizes
Limplicin (0) Ey, (Xo)] E. (x1)] $ Ep.
ass
o²
+
As dt.
(4)
2
Neklyudov al. (2023) showed that Cimplici (0) implicitly matches unique gradient field Vs; with least kinetic energy, Le., it is equivalent to minimizing fo Ep. V V dt, where
s = arg min
Ep
2 Iu (x,))
dt subject to (2).
Bridge & Flow Matching (explicit). If the Pt can be factorized into := Epo (P: (xeaas / }]] where the conditional density Pe (1), denoted an SDE, d.X, (X, (ro, Fi )d + adW,, then it can be shown (Lipman et al., 2023; Albergo & Vanden-Eiinden, 2023; Liu et al., 2023a) that the minimizer of (see Appendix C.1 for the derivation)
denoted with the shorthand P:jo.1 associated with
Lexplici (0) /'"
} Ilu (X. xn )1 dt,
B. Epgyp.t
similar to Vs also satisfies the FPE prescribed by Pt- In other words, uf preserves pr for all t.
Implicit explicit matching losses. While Entropic Action Matching presents general match- ing method with the least assumptions, its implicit matching loss Limplicit (4) scales unfavorably to high-dimensional applications, due the need to approximate the Laplacian (using egg. Hutchinson (1989)), and also introduces unquantifiable bias when optimizing over a restricted family of func- tions such deep neural networks. The explicit matching loss Lexplisir (5) offers a computationally efficient alternative but requires additional information, namely ux(X, (x .1) = 14 Remark- ably, in both cases, the minimizers preserve the preseribed marginal pt. Hence, if Pu μ and pi these matching algorithms shall always return feasible solution-a diffusion model that matches between y and v. We summarize the aforementioned two methods Alg. and 2.
GENERALIZED SCHRÖDINGER BRIDGE MATCHING (GSBM)
We propose Generalized Schrödinger Bridge Matching (GSBM), novel matching algorithm
that, in contrast to those in Sec. 2 assuming prescribed pt. concurrently optimizes p. to minimize (3) subject the feasibility constraint (2). All proofs be found in Appendix B.
3

---

Published as a conference paper at ICLR 2024
3.1 ALTERNATING OPTIMIZATION SCHEME
Let us revisit the GSB problem (3), particularly its FPE constraint in (2). Recent advances in dy- namic optimal transport (Liu al., 2023b) and SB (Peluchetti, 2023) propose a decomposition of this dynamical constraint into two components: the marginal Pt. 1 E (0. 1), and the joint coupling between boundaries Po,1. and employ alternating optimization between uf and ps. Specifically, these optimization methods, which largely inspired recent variants of SB Matching (Shi et al., 2023) and our GSBM, generally obey the following recipe, alternating between two stages:
Stage I: Optimize the objective, in case, (3) w.r.t. the drift uf given fixed Pre[o.1]-
Stage 2: Optimize (3) w.r.t. the marginals PrE(0.1) given the coupling p8, defined by uf.
Notice particularly that the optimization posed in Stage resembles the matching algorithms in Sec. 2. We make the connection concrete in the following proposition:
Proposition 1 (Stage 1). The unique minimizer to Stage coincides with V (X,).
This may seem counter-intuitive at first glance, both the kinetic energy and V show up in (3). This is due to the fact that uf no longer affects the value of Ep, (Xx)] once Pt is fixed, and out of all uf whose probability density matches pt, the gradient field is the unique minimizer of kinetic energy. We emphasize that the role of Stage 1 to learn a uf given the prescribed p. and provides better coupling P which is then used to refine Pt in Stage 2. Therefore, the explicit matching loss (5) provides just the same algorithmic purpose. Though it only upper-bounds the objective of Stage 1, its solution often converges stably from measure- theorrete perspective (see Appendix C. for more explanations). In practice, we find that Lexplicit performs as effectively as Limplicit. while exhibiting significantly improved scalability, hence better suited for high-dimensional applications.
We now present our main result, which shows that Stage 2 of GSBM can be cast as a variational problem, where V, appears through optimizing conditional controlled process.
Proposition 2 (Stage 2: Conditional stochastic optimal control; CondSOC). Let the minimizer to Stage factorized by (x,) = pe(
pr((eear Xe|=0, æ1 )p6, (Ir, xn )dzgdzı where p8,1 is the boundary
coupling induced by solving (1) with uf. Then, Pe (X, (ro .n) = P:jo,1 solves:
min J = L/' (x((mm))))) Majo,1
1e (Xe 21 21)1 VX(),) dt
(6a)
2
S.T. dx, = (X, 21 )tt + adW,
Xo Zo. X1 σ
(6b)
Note that (6) differs from (3) only in the boundary conditions, where the original distributions p, v are replaced by the two end-points (Ip, x1 drawn from the coupling induced by u/- Generally, the solution to (6) is not known in closed form, except special cases. In Lemma 3, we show such a case when V(x) is quadratic and a>0 Note that as (xx) vanishes, 1 in Lemma collapses to the Brownian bridge and GSBM recovers the match- ing algorithm appearing in prior works (Liu, 2022; Shi et al., 2023) that approximate OT/SB, sum- marized in Table 1. This suggests that our Prop. 2 directly generalizes them to nontrivial V,,
Table 1: Solutions to (6) w.r.t. different V/, and how they link to different methods, including Rectified flow (Liu et al., 2023b), DSBM (Shi et al., 2023), and our GSBM.
Method
Vi(x)
Pc #1)
RecFlow (a=0)
straight line (Lemma a, a++0)
0
DSBM
Brownian bridge (Lemma with a-+0)
0
(e>0)
GSBM
quadratic
Lemma 3
(a20)
arbitrary
Sec. 3.2
Lemma 3 (Analytic solution to (6) for quadratic V and > 0). Let Vxx) a a0l0x//1 a, 0,
then the optimal solution to follows a Gaussian path X; ~ N(caa + Ec 78 Id). where
sinh(n(1-t)) sinh
(sinh(n(111) et == 02.a. e: cosh(n(1-tt) - cc cosh 17, 7t- =
C
"1
These coefficients recover Brownian bridges as a + 0 and, if further a + O. straight lines.
Notablyy, the absence of V, in optimization was previously viewed issue for Sinkhorn methods aimed solving (3) (Liu al., 2022). Yet, in GSBM, appears naturally from how the optimization is decomposed.
4

---

Published as a conference paper at ICLR 2024
Spline Moon Ak
Spline Std v
CondSCC Pi(
Algorithm 3 Sp1 neOpt
Instialzation
Require: Iо, 21, {Xa,) where Oct1< <tp=1
Initialize f4s, Tt. with (9)
for m = 0 to M do
Sample X; ~ (N 72 Ia) with (7)
Optimized
Compute (x()) and uy (xx I1) with (8)
Estimate the objective I in (6a)
Take gradient step w.r.t control pts. (Xi,, 7o
Figure 1: Example of spline optimization (Alg. 3) for 14 E 12.7% E R, and the re-
end for
return Prjo,1 parametrized by optimized p4, 7
sulting CondSOC (6) solution.
3.2 APPROXIMATE SOLUTIONS TO CONDSOC FOR GENERAL NONLINEAR STATE COST
We develop methods for searching the approximate solutions to CondSOC in (6), when V, neither quadratic degenerate. As need to search for every pair (zo, 21) drawn from P 1. seek efficient methods that are parallizable and simulation-Fii -utdiee of
outside
sampling from pÓ.1-
Gaussian path approximation. Drawing inspirations from Lemma 3, and since the boundary con- ditions of (6) are simply two fixed points, we propose to approximate the solution to (6) Gaussian probability path pinned at To, 211 (X
= [[ 72 Ia).
s.t.
71 = 0,
where 14 E Rdd and "tt E R are respectively the varying mean and standard deviation. immediate result following from (7) is a closed-form conditional drift (Särkkä & Solin, 2019) (see Appendix C.2 for the derivation):
U4 (M, Xo, = Sepe + (X, pt); where ar == 0c7c
(8)
2
Notice that E IR is a time-varying scalar. Hence, however complex 14, may be, the underlying SDE (with the conditional drift u, (Xy| x1) in (8)) remains linear. Also note that this drift, since it is gradient field, is the kinetic optimal choice out of all drifts that generate pe Pr(X./zo. x1).
Spline optimization. To facilitate efficient optimization, we parametrize μt, os respectively as d- and 1-D splines with some control points Xt, Rd and nt E IR sampled sparsely and uniformly along the time steps 0 t < < tK 1:
Pe Spline(t; Zo. 21)
Ve := Splinet; 70=0,
Notice that the parameterization in (9) satisfy the boundary in (7), hence remains as feasible solution to (6b) however (Xt,, Yta } change. The number of control points K is much smaller than discretization steps (K <30 for all experiments). This significantly reduces the memory complexity compared to prior works (Liu et al., 2022), which require caching entire discretized SDEs.
Alg. summarizes the spline optimization, which, erucially, involves no simulation of an SDE (6b). This is because optimize using just independent samples from Pıju,1, which are known in closed form. Since the CondSOC problem relaxes distributional boundary constraints to just two end- points, our computationally efficient variational approximation holds out very well in practice. Fur- thermore, since we only need to optimize very few spline parameters, we did not find the need to consider amortized variational inference (Kingma & Welling, 2014). Finally, note that since solve CondSOC for each pair (1), ~ P8_1 and later marginalize to construct Pt, we need not explicitly consider mixture solutions for Prjo,1 1- long we sample sufficiently many pairs of (Io, 2. In practice, we initialize { } from uf and Yt, = a 4 -1-111 (x) from the standard deviation of the Brownian bridge. Figure/ demonstrates 2D example.
Resampling using path integral theory. In cases where the family Gaussian probability paths not sufficient for modeling solutions of (6), a more rigorous approach involves the path integral theory (Kappen, 2005), which provides analytic expression to the optimal density of (6) given any sampling distribution with sufficient support. We discuss this in the following proposition.
This induces computational overhead, as {Xe,) intermediate steps when simulating To, ~ P0.1-
5

---

Published as a conference paper at ICLR 2024
Algorithm 5 Generalized Schrödinger Bridge Matching (GSBM)
Initialize uy and set Pt as the solution (6) with independent coupling p0,1 := p 3
2: repeat
3: u match (pt) or match (Pt. uzjo.1)
Alg. 1 1
4: Sample X1, X1, from up on timesteps Octι <... < tkel
5: Optimize P Me, Yt SplineOpt (Xo X1, (X(,})
Alg.
6: Determine 1 from μt, ?te using (8)
7:
if apply path integral resampling then
optional step b Alg.
8: Prjo,1 + Impt Sample (Ptjo,1)
9: end if
10: until converges
Proposition 4 (Path integral solution to (6)). Let r(X [o, 11) be a distribution absolutely contin-
uous w.r.t. the Brownian motion, denoted q(X [o0), where shorthand X = =Xxe(u,1). Suppose is associated with an SDE dX,, = dtt + adW, Xo=xo, X1=I1, a>0. Then, the optimal, solution (6) can be obtained via
y" (X (100,11) z (aro in (x|o, 11);
(10)
where Z is the normalization constant and w is the importance weight:
(([ x1) =xexp exp (-Jo (-f-=t vi((.)) (1ll((X)))))) di L." (x.) Taw:).
(II)
02
In practice, r( (1|rr, .11) can be any distribution, but the closer it is to the optimal p², the lower the variance of the importance weights w (Kappen Ruiz, 2016). It is therefore natural to consider us- ing the aforementioned Gaussian probability paths as rx[[2g, properly optimized with Alg. 3, then resampled following proportional to (10)-this is equivalent in expectation performing self-normalized importance sampling but algorithmically simpler and helps reduce variance when many samples have low weight values. While path in- tegral resampling may make Lexplicit less suitable due the change in the conditional drift gov- erning p" (X (ro, I1) from (8), still observe empirical, sometimes significant, improvement (see Sec. 4.4). Overall, we propose path integral resampling as an optional step in our GSBM algo- rithm, it requires sequential simulation." In practice, we also find empirically that the Gaussian probability paths alone perform sufficiently well and is easy to use at scale.
Algorithm 4 Impt Sample
Require: Pijo, Mer 7hs a 0
Sample Xię from (6b) given (8)
Compute w with (11) and Z Σ, w return Resample Pejo,1 with (10).
3.3 ALGORITIM OUTLINE & CONVERGENCE ANALYSIS
We summarize our GSBM in Alg. 5, which, as previously sketched in Sec. 3.1, alternates between Stage (line 3) and Stage 2 (lines 4-9). In contrast to DSBM (Shi et al., 2023), which constructs analytic 1 and i, when V/2 = O, our GSBM solves the underlying CondSOC (6) with nontriv- ial V. We stress that these computations are easily parallizable and admit fast converge due to our efficient parameterization, hence inducing little computational overhead (see Sec. 4.4).
We note that GSBM (Alg. 5) remains functional even when a=0, although the GSB problem (3) was originally stated with a>00 (Chen et al., 2015). In such cases, the implicit and explicit matching still preserve pt and correspond, respectively, to action (Neklyudov et al., 2023) and flow matching (Lipman et al., 2023), and CondSOC corresponds to a generalized geodesic path accounting for Vi
Finally, we provide convergence analysis in the following theorems.
Theorem 5 (Local convergence). Let the intermediate result of Stage and after repeating n times, be 0r, Then, its objective value in (3), L(O"), is monotonically non-increasing grows:
L( 2 L(0n+1).
Theorem 6. The optimal solution GSB problem (3) is a fixed point of GSBM, Alg.
"We note that simulation of trajectories X ==Xıeto.11 from (6a,8) can he done efficiently by computing the covariance function, which requires merely solving an ID ODE; see Appendix C.2 for detailed explanation.

---

Published as a conference paper at ICLR 2024
6-0.00-0,25
t-0.25-0.50
5-0.55-0.75
t-0.55-100
Optimaltty leasiblityy
GSBM Dreepesss
TANN Figure Simulation of SDEs with the of after long
Mi
7k-
1x
Training Time (hours)
Figure 2: Feasibility vs. optimality on three
training. Notice how DeepGSB diverges drastically from our GSBM, which satisfies feasibility at all time.
crowd navigation tasks with mean-field cost.
LiDAR surface & tangent plane
GSBM; objective value (3) 6209 (left: 3D view, right: 2D view)
DeepGSB; obj.=7747 (2D view)
Po and (2D view)
Figure 4: Crowd navigation over LiDAR surface. Height is denoted by the grayscale color.
EXPERIMENT
We test GSBM variety of distribution matching tasks, each entailing its state cost V1():.. By default, we use the explicit matching loss (5) without path integral resampling, mainly due to its scalability, but ablate their relative performances in Sec. 4.4. Our GSBM compared primarily to DeepGSB (Liu et al., 2022), a Sinkhorr-insppree machine learning method that outperforms existing deep methods (Ruthotto al., 2020; Lin al., 2021). Other details are in Appendix D.
4.1 CROWD NAVIGATION WITH MEAN-FIELD AND GEOMETRIC STATE COSTS
We first validate our GSBM in solving crowd navigation, canonical example for the GSB prob- lem (3). Specifically, we consider the following two classes of tasks (see Appendix D.1 for details):
Mean-field interactions. These are synthetic dataset in IR2 introduced in DeepGSB, where the state cost V3 consists of two components: an obstacle cost that assesses the physical constraints, and an "mean-field" interaction cost between individual agents and the population density pe:
Vx(x) = (x) + Limernex (2; pc), Limerract (x;; (x: log Pc())
(entropy)
(congestion) (12)
Eygee || p 2+1
2
Both entropy and congestion the costs incurred for individuals to stay in densely crowded regions with high population density.
fundamental elements of mean-field games. They measure
Geometric surfaces defined by LIDAR. A more realistic scenario involves navigation through complex geometric surface. In particular, we consider surfaces observed through LiDAR scans of Mt. Rainier (Legg & Anderson, 2013), thinned 34,183 points (see Fig. 4). We adopt the state cost
(x) Lmaniiola (z) (x), Lmanitola [m( = Lncigh((x) exp
( (x)). (13)
where x(x) projects to approximate tangent plane fitted by k-nearest neighbors (see Fig, and Appendix D.3) and (x)) refers to the z-coordinate of T (x)), i.e. its height.
Figure tracks the feasibility and optimality, measured by W2(rf, v) and (3), of three mean-field tasks (Stunnel, Vneck, GMM). On all tasks, our GSBM maintains feasible solutions throughout training while gradually improving optimality. In contrast, due to the lack of convergence analysis, training DeepGSB exhibits relative instability and occasional divergence (see Fig. 3). As for the ge- ometric state costs, Fig. demonstrates how our GSBM faithfully recovers the desired multi-mode distribution: It successfully identify viable pathways with low state cost, one of which bypasses the saddle point. constrast, DeepGSB generates only uni-modal distributions with samples scat- tered over tall mountain regions with high state cost, yielding higher objective value (3).
7

---

Published as a conference paper at ICLR 2024
Diffusion Schrödinger Bridge Matching (V, 0)
Generalized Schrödinger Bridge Matching (V: 0)
Epoch
Egoch Epoch
Epoch
Epoch Epoch
t=1.00
Figure 5: Comparison between DSBM (Shi et al., 2023) and our GSBM on: (lestmost 5 columns) the generation processes and (rightmost columns) their couplings (X1 during training. By constructing Vi via latent space, GSBM exhibits faster convergence and yields better couplings.
1-0,00
c-o 25
t-0.50
t-0.75
t-1.00
t-0.00
T-025
t-0.50
t-0.75
t-1.00
Table 2:
FID values
DSBM
of dog->ooe >aat for DSBM and our GSBM.
( 21). Instead of using linear interpolation as in DSBM, GSBM optimizes P:jo,1 w.r.t. (6) where V, is defined via a latent space, thereby exhibiting more semantically meaningful interpolations.
Figure 6: Mean of Px(
DSBM
GSBM
14.16
12.39
4.2 IMAGE DOMAIN INTERPOLATION AND UNPAIRED TRANSLATION
Next, we consider unpaired translation between dogs and cats from AFHQ (Choi et al., 2020). We aim explore how appropriate choices of state cost Vú help encourage more natural interpo- lations and more semantically meaningful couplings. While the design of V, itself is an interest- ing open question, here we exploit the geometry of a learned latent space. To this end, we use a pretrained variational autoencoder (Kingma & Welling, 2014), then define V conditioned on an interpolation of two end points (notice that x and respectively belong to image and latent spaces):
Vi(1 (o, 21) | 1ecoder(=;)) 21 = I(t, Encoder Encoode((11).
(14)
Though [tt, 20, 21) can be any appropriate interpolation, (Shoemake, 1985) to be particularly effective, due to the Gaussian geometry in latent space. As the high dimensionality greatly impedes DeepGSB, we mainly compare with DSBM (Shi et al., 2023), the special case of our GSBM when V4 is degenerate. As in DSBM, we resize images to 64x64.
find the spherical linear interpolation
Figure 5 reports the qualitative comparison between the generation processes of DSBM and GSBM, along with their coupling during training. meaningful Ve, GSBM typically converges to near-optimal coupling early in training, and, as ex- pected, yields more interpretable generation processes. Interestingly, despite being subject to the same noise level (o=0.5 5 in this case), the GSBM's generation processes are generally less noisy than DSBM. This is due to the optimization of the CondSOC problem (6) with our specific choice of Vt. As shown in Fig. 6, the conditional density Pe (x, 21) used in GSBM appears eliminate unnatural artifacts observed in Brownian bridges, which rely on simple linear interpolation in pixel space. Quantitatively, our GSBM also achieves lower FID value as measure of feasibility, as shown in Table 2. Finally, note that the inclusion of V2 and the solving of CondSOC increase wallclock time by a mere 0.5% compared to DSBM (see Table 5).
clear that, with the aid of semantically
Deept
GSBM
4.3 HIGH-DIMENSIONAL OPINION DEPOLARIZATTIN
dynamics (15)
Target
[Obj - 914.1
(Ob) - 505.5
D5 PCA vs.
Finally, consider high-dimennn -ith-dmmensional opinion depolar- ization, initially introduced in DeepGSB, where an opin- ion X, 1000 is influenced by polarizing effect (Schweighofer et al., 2020) when evolving through in- teractions with the population (see Appendix D.2):
Figure 7: Distribution of terminal opin- ion X, E IR 1000 and their directional similarities (DS; see Appendix D.2).
dX, Spularize (X,;Pe Pe at odW,
(15)
8

---

Published as a conference paper at ICLR 2024
0.0
0.1 1.0
Table 3:
Relative runtime be-
tween combination of matching losses and PI resampling on solv- ing Stunnel, relative to Lexplicit
Cexpliit
Cipplict
Our GSBM reveals how the optimal solution drastically changes w.r.t the levels of noise (σ) when navigating through a narrow passway surrounded by the obstacles (shown in gray).
Figure 8:
without PI 100%
276%
with PI 108% 284%
Table 4: How the objective value (3) changes when enabling Pl resampling on each matching loss and a in Stunnel. Per- formance is improved by deceasing (-) the objective values.
Table 5: Percentage of time spent in different stages of Alg. 5, measured on the AFHQ task.
0.5
1.0
2.0
ma t ch
Simulate of
Solve (6)
(line 3)
(line 4) (lines 5-6)
Lesplicit
5.9
-5.2 3.2=2.m 2.6
1.5 + 3.5
Limplicit -112.4 +: 53.7
-20 = 8.0
-0.7 +: 6.0
64.3%
35.2%
0.5%
Without any intervention, the opinion dynamics in (15) tend to segregate into groups with diamet- rically opposed views (first column of Fig. 7), as opposed to the desired unimodal distribution pr (second column of Fig. 7). To adapt our GSBM for this task, we treat (15) as base drift, specifi- cally defining ug (X+) == fpotanie (t) + rf (Xe), and then solving the CondSOC (6) by replacing the kinetic energy with I |lo defined in (12). As shown in Fig. 7. both DeepGSB and our GSBM demonstrate the capability to mitigate opinion segregation. However, our GSBM achieves closer proximity to the target Pl, indicating stronger feasibility, and achieve almost half the objective value (3) relative to DeepGSB.
Similar to DeepGSB, we consider the same congestion cost V,
4.4 DISCUSSIONS
Effeet of noise (o). In the stochastic control setting (Theodorou et al., 2010), the task spec-speifiic value of plays crucial role in representing the uncertainty from environment the error in exe- cuting control. The optimal control thus changes drastically depending on a. Figure 8 demonstrates how our GSBM correctly resolves this phenomenon, problem discussed in Kappen (2005). In the absence of noise (a=0), it is very easy to steer through the narrow passage. When large amounts of noise is present (o=1.0), there is a high chance of colliding with the obstacles, so the optimal solution is to completely steer around the obstacles.
example of the famous "drunken spider"
Ablation study on path integral (PI) resampling. In Table 4, we ablate how the objective value changes when enabling PI resampling on different matching losses and noise a and report their performance (averaged over 5 independent trails) on the Stunnel task. We observe that PI resam- pling tends to enhance overall performance, particularly in low noise conditions, at the expense of slightly increased runtime of 8%. Meanwhile, as shown in Table implicit matching (4) typically requires longer time overall (x2.7 even in two dimensions), compared to its explicit counterpart.
Profiling GSBM. The primary algorithmic distinction between our GSBM and previous SB matching methods (Shi et al., 2023; Peluchetti, 2023) lies in how Prjo,1 and urjo, are computed (see Lemma 3), where GSBM involves solving an additional CondSOC problem, iee., lines 5-6 in Alg. 5. Table 5 suggests that these computations, uniquely attached GSBM, induce little compu- tational overhead compared to other components in Alg. 5. This computational efficiency is due to our efficient variational approximation admitting simulation-free, parallelizable optimization.
CONCLUSION AND LIMITATION
We developed GSBM, new matching algorithm that provides approximate solutions to the Gen- eralized Schrödinger Bridge (GSB) problems. We demonstrated strong capabilities of GSBM over prior methods in solving crowd navigation, opinion modeling, and interpretable domain transfer. It should be noticed that GSBM requires differentiability of V/ and relies on the quadratic control cost to establish its convergence analysis, which, despite notably improving over prior methods, remains as necessary conditions. We acknowledge these limitations and leave them for future works.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We acknowledge the Python community (Van Rossum & Drake Jr, 1995; Oliphant, 2007) and the core set of tools that enabled this work, including PyTorch (Paszke et al., 2019), functorch (Ho- race He, 2021), torchdiffeq (Chen, 2018), JAX (Bradbury et al., 2018), Flax (Heek al., 2020), Hydra (Yadan, 2019), Jupyter (Kluyver et al., 2016), Matplotlib (Hunter, 2007), numpy (Oliphant, 2006; Van Der Walt et al., 201 I), and SciPy (Jones et al., 2014).
REFERENCES
Michael S Albergo and Eric Vanden- Eijnen. Building normalizing flows with stochastic inter- polants. In International Conference on Learning Representations (ICLR), 2023.
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Fijndenn. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint arXiv:2303. 08797, 2023.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: In metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Ap- plications, 12(3):31-3326, 1982.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com/ google/jas
Ricky Q. Chen. torchdiffeq, 2018.
URL httpp:://gttaub..cmm/ttttt wttp://w..coottttt...o..o.ooo
torchdi ffeq.
Ricky T. Q. Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv
preprint wrXhv:2302.03660 2023.
Ricky Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of Schrödinger bridge using forward-backwart SDEs theory. In International Conference on Learning Represen- tations (ICLR), 2022.
Yongxin Chen. Density control of interacting agent systems. IEFE Transactions Automatic Control, 2023.
Yongxin Chen and Tryphon Georgiou. Stochastic bridges of linear systems. IEEE Transactions on Automatic Control, (2): -310)1,11111,,-, 2015.
Yongxin Chen, Tryphon Georgiou, and Michele Pavon. Optimal steering of inertial particles diffus- ing anisotropically with losses. In American Control Conference (ACC). IEEE, 2015.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems (NeurIPS), 2013.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems (NeurIPS). 2021.
Prafulla Dhariwal and Alex Nichol. Diffusson models beat GANs on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
10

---

Published as a conference paper at ICLR 2024
Simone Di Marino, Augusto Gerolin, and Luca Nenna. Optimal transportation theory with repulsive costs. Topological optimization and optimal transport, 17:204-255 256, 2017.
Robert Fortet. Résolution d'un système d' équations de M. Schrödinger. Journal de Mathématiques Pures et Appliquées, 19(1-4):88) ((1-4):83-100 105, 1940.
Jason Gaitonde, Jon Kleinberg, and Êva Tardos. Polarization in geometric opinion dynamics. In ACM Conference on Economics and Computation, pp. 499-519, 2021.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL http: /ttp:///ittiie github. com/ /oog/eoolle///oooooooooooooooooe flax.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. 3-vae: Learning basic visual concepts with a con- strained variational framework. In International Conference on Learning Representations (ICLR), 2016.
Richard Zou Horace He. functorch: Jax-like composable function transforms for pytorch. https:
// / wtt.hub.ccm/pytooeett... http./cccc/yyocccc/uttoooccchc functorch, 2021.
John Hunter. Matplotli:: 2d graphies environment. Computing in science & engineering, 9(3): 90, 2007.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Cratistics-imullati and Computation, 1853) 1989.
Kiyosi Itô. On stochastic differential equations, volume 4. American Mathematical Soc., 1951.
Eric Jones, Travis Oliphant, and Pearu Peterson. {SciPy): Open source scientific tools for (Python). 2014.
Hilbert Kappen. Path integrals and symmetry breaking for optimal control theory. Journal of Statistical Mechanics: Theory and Experiment, 2005(11):P11011: 2005.
Hilbert Johan Kappen and Hans Christian Ruiz. Adaptive importance sampling for control and inference. Journal of Statistical Physics, 162(5):1244-126) 2016.
Diederik P Kingma and Max Welling. Auto-Encodinn variational bayes. In International Conference Learning Representations (ICLR), 2014.
Thomas Kluyver, Benjamin Ragan-Kelley, Fernando Pérez, Brian E Granger, Matthias Bussonnier, Jonathan Frederic, Kyle Kelley, Jessica B Hamrick. Jason Grout, Sylvain Corlay, et al. Jupyter notebooks-a publishing format for reproducible computational workflows. In ELPUB, pp. 87-90, 2016.
Takeshi Koshizuka and Issei Sato. Neural Lagrangian Schrödinger bridge: Diffusion modeling for population dynamics. In International Conference on Learning Representations (ICLR), 2023.
Nicholas Legg and Scott Anderson. Southwest Flank of Mt.Rainier WA, 2013. URL https: / //opentooporaaphy. org/ wwtpoogtg/////oeeeeee/ooe/ee att.//etat/ooooeo 052013 26910. 1. Accessed on 2023-09-12.
Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems, 2013.
Christian Léonard, Sylvie Rælly, and Jean-Claude Zambrini. Reciprocal processes. A measure- theoretical point of view. Probability Surveys, 2014.
David Levin. The approximation power of moving least-squares. Mathematics of Computation, 67 (224):1517-1511, 1998.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: arXiv:/805.0w 2018.

---

Published as a conference paper at ICLR 2024
Xuechen Li, Ting-Kat Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2020.
Alex Tong Lin, Samy Wu Fung, Wuchen Li, Levon Nurbekyan, and Stanley Osher. Alternating the population and control neural networks to solve high-dimensional stochastic mean-field games. Proceedings of the National Academy of Sciences, 118(31), 2021.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. International Conference on Learning Representations (ICLR), 2023.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos A Theodorou. Deep generalized Schrödinger bridge. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Guan-Homg Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar. 12SB: Image-to-Image Schrödinger bridge. In International Conference on Ma- chine Learning (ICML), 2023a.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint DrXiv:2209.14577 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. International Conference on Learning Representations (ICLR), 2023b.
Zhiyu Liu, Bo Wu, and Hai Lin. A mean field game approach to swarming robots control. In American Control Conference (ACC). IEEE, 2018.
lyya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Representations (ICLR), 2019.
Alain Mazzolo and Cécile Monthus. Conditioning diffusion processes with killing rates. Journal of Statistical Mechanics: Theory and Experiment, 2882(,) 2022.
Kirill Neklyudov, Daniel Severo, and Alireza Makhzani. Action matching: A variational method for learning stochastic dynamics from samples. In International Conference on Machine Learning (ICML), 2023.
Frank Noé, Alexandre Tkatchenko, Klaus-Robert Müller, and Cecilia Clementi. Machine learning for molecular simulation. Annual Review of Physical Chemistry, 71:361-390, 2020.
Masashi Okada and Tadahiro Taniguchi. Variational inference mpc for bayesian model-based rein- forcement learning. In Conference on Robot Learning (CoRL), 2020.
Travis E Oliphant. A guide NumPy, volume 1. Trelgol Publishing USA, 2006.
Travis Oliphant. Python for scientific computing. Computing Science & Engineering, 9(3): 10-20, 2007.
Etienne Pardoux and Shige Peng. Backward stochastic differential equations and quasilinear parabolic partial differential equations. In Stochastic Partial Differential Equations and Their Applications. Springer, 2005.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, al. Pytorch: An imperative style, high- performance deep learning library. In Advances in neural information processing systems, pp- 8026-8037, 2019.
Stefano Peluchetti.
Non-Denoising forward-time diffusions, 2022.
URL https://
openreview. net/ http://orrnnnni/iieeeeee "dru
Stefano Peluchetti. Diffusion bridge mixture transports. Schrödinger bridge problems and generative modeling. arXiv preprint arXiv:2304.009/7, 2023.
12

---

Published as a conference paper at ICLR 2024
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Center for Research in Eco- nomics and Statisties Working Papers, 2017.
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With upplications to data sci- ence. Foundations and Trends@ in Machine Learning, 11(5-6):355-66 6073666767 2019.
Chris Philippidis. Chris Dewdney, and Basil Hiley. Quantum interference and the quantum poten- tial. Nuovo Cimento B, 52(1):15-28) 1979.
Konrad Rawlik, Mare Toussaint, and Sethu Vijayakumar. On stochastic optimal control and re- inforcement learning by approximate inference. In International Joint Conference on Artificial Intelligence (IJCAI), 2013.
Hannes Risken. Fokker-Plande equation. In The Fokker-Planck Equation, pp. 63-95. Springer, 1996.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional networks for biomed- ical image segmentation. In International Conference on Medical Image Computing and Computer-assisted Intervention. Springer, 2015.
Lars Ruthotto, Stanley Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung. A machine learning framework for solving high-dimensional mean field game and mean field control prob- lems. Proceedings of the National Academy of Sciences, 117(17):1183-1193, 2020.
Simo Särkkä and Arno Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Erwin Schrödinger. Über die Umkehrung der Naturgesetze, volume IX. Sitzungsberichte der Preuss Akad. Wissen. Phys. Math. Klasse, Sonderausgabe, 1931.
Simon Schweighofer, David Garcia, and Frank Schweitzer. An agent-based model of multi- dimensional opinion dynamics and opinion alignment. Chaos: An Interdisciplinary Journal of Nonlinear Science, 30(9):093139, 2020.
Neta Shaul, Ricky T. Q. Chen, Maximilian Nickel, Matthew Le, and Yaron Lipman. On kinetic op- timal probability paths for generative models. In International Conference on Machine Learning (ICML). 2023.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Ken Shoemake. Animating rotation with quaternion curves. In Conference on Computer Graphics and Interactive Techniques, 1985.
Jascha Sohl-Dickstein. Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing (ICML), 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations (ICLR), 2021.
Evangelos Theodorou. Nonlinear stochastic control and information theoretic dualities: Connec- tions, interdependencies and thermodynamic interpretations. Entropy, 17(5):3352-33 -37-332-2-375.00 2015.
Evangelos Theodorou, Jonas Buchli, and Stefan Schaal. A generalized path integral control ap- proach to reinforeement learning. Journal of Machine Learning Research (JMLR), 11(Nov): 3137-3181, 2010.
Emanuel Todorov. Linearly-solvale Markov decision problems. In Advances Neural Information Processing Systems (NeurIPS), 2007.
13

---

Published as a conference paper at ICLR 2024
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. In International Conference Machine Learning (ICML), Workshop Track, 2023.
Stefan Van Der Walt, S Chris Colbert, and Gael Varoquaux. The numpy array: a structure for efficient numerical computation. Computing in Science & Engineering, 13(2):22, 2011.
Guido Van Rossum and Fred Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Francisco Vargas, Pierre Thodoroff, Neil D Lawrence, and Austen Lamacraft. Solving Schrödinger bridges via maximum likelihood. Entropy, 2021.
Cédric Villani et al. Optimal transport:
Old and new, volume 338. Springer, 2009.
Mei Wang and Weihong Deng. Deep visual domain adaptation: A survey. Neurocomputing, 312: 135-153, 2018.
Holger Wendland. Scattered data approximation, volume 17. Cambridge university press, 2004.
Omry Yadan. Hydra a framework for elegantly configuring complex applications. Github, 2019.
URL httpttttt////thhhie wwtp:////.ccccoocererrrcceece/ee
Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and HJB equations, volume 43. Springer Science & Business Media, 1999.
Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. Proceedings of the IEEE international conference computer vision, pp. 2223-2232, 2017.
14

---

Published as a conference paper at ICLR 2024
A REVIEWS ON STOCHASTIC OPTIMAL CONTROL AND RELATED WORKS
Solving Generalized Schrödinger bridge by reframing as stochastic optimal control The so-
lution to Generalized Schrödinger Bridge (GSB) problem (3) can also be expressed as the solution a stochastic optimal control (SOC) problem, typically structured as
min (-g
Ep.
1l (X)I1 + vi(X,)) di + Ep, [x(x)),
(16a)
s.t. dX;,, = (x))) + adW,, Xo ~M
(16b)
where we see that the terminal distribution "hard constraint" in GSB (3) is instead relaxed into a soft
"terminal cost" Φ(.) : TRd R. Problems with the forms of either (16) or (3) are known be tied
to phearll-solvvbll Markov decision processes (Todorov, 2007; Rawlik et al., 2013), corresponding to a tractable class of SOC problems whose optimality conditions-thh Hamittn--lccoiiBelll equations-adine admit efficient approximation. This is attributed to the presence of the ---norm control cost, which can be interpreted as the KL divergence between controlled and uncontrolled processes. The interpretation bridges the SOC problems to probabilistic inference (Levine, 2018; Okada & Taniguchi, 2020), from which machine learning algorithms, such as our GSBM, can be developed.
However, naively transforming GSB problems (3) into SOC problems tial issues. The design of the terminal cost is extremely important, and in many cases, it is an intractable cost function as we do not have access the densities μu and Prior works have mainly stuck to simple terminal costs (Ruthotto et al., 2020), using biased approximations based on batch estimates (Koshizuka & Sato, 2023), or using an adversarial approach to learn the cost function (Zhu et al., 2017; Lin et al., 2021). Furthermore, this approach will necessitate differentiating through an SDE simulation, requiring high memory usage. Though memory--fficicee adjoint methods have been developed (Chen et al., 201 8; Li et al., 2020), they remain computationally expensive to use at scale.
introduce many poten-
In contrast, our GBSM approach only requires samples from μ and v. By enforcing boundary distri- butions as a hard constraint instead of a soft one, our algorithm finds solutions that satisfy feasibility much better in practice. This further allows us to consider higher dimensional problems without the need to introduce additional hyperparameters for tuning terminal cost. Finally, our algorithm dras- tically reduces the number of SDE simulations required, as both the matching algorithm (Stage 1) and the CondSOC variational formulation (Stage 2) of GSBM can be done simulation-free.
Comparison to related works Table compares to prior learning based methods that also pro- vide approximate solutions to (3) to our GSBM. Meanwhile, Fig. summarizes different classes of matching algorithms.
Table 6: Our GSBM features hetter feasibility, requires only samples from μt, v, has local conver- gence analysis, and exhibits much better scalability.
Requirement from Convergence analysis
Feasibility (2)
Dimension d
distributions μr,
NLSB (Koshizuka Sato, 2023)
(reeaxed)
samples & densities
/local)
5 1000
DeepGSB (Liu et al., 2022)
F only in limit
samples & densities
X
GSBM (this work)
F
only samples
V
(locall
> 12K
KE lask task- specific loss
Generalized OT & SB
optimality structure
(Ganaralized Schrödin ger bridge matching; ouns)
kinetic energy
optimal transport (OT)
Schrödinger bridge (SB) IDBM, DSBM
Rectified Flow
fixed
action matching AM, Entropic AM
flow matching FM, Riamannian FM
bridge matching 1SB, Aligned 5B
JMM Extreee AM toetddr clud. 20031 PM Expren vtal. 20231 Sammen t0hv à Ligmenn 20031 1155 to al. 2325:
Figure 9: Summary of different matching algorithms. Notice that OT and SB are w.r.t. E costs.
15

---

Published as a conference paper at ICLR 2024
B PROOFS
Proposition 1 (Stage 1). The unique minimizer to Stage coincides with (X,.
Proof. By keeping Ps fixed in (3), the value of Ep, Vx.) no longer relies ut rendering (3) the same optimization problem as in the implicit matching, see (4) and Sec. 2. As the implicit matching (4) admits a unique minimizer Vst. conclude the proof.
Proposition 2 (Stage 2: Conditional stochastic optimal control; CondSOC). Let the minimizer to Stage 2 be factorized by px(X)) = pe((Xlloo 21 DP (Io, 21 )dzddz where p6,1 is the boundary coupling induced by solving (1) with ut. Then, pe X2 [o. x1) = P:jo,1 solves:
min J = J (x((xa))))
(X. a1)1 Vi(X,) dt
(6a)
Majo.1
s.l. dX, = (1) + odW,, Xo No. X 21
(6b)
Proof. Let us recall the GSB problem (3) in the form of FPE constraint (2):
1 (x,)) + vi(Xi) dt
min the
Epp.
(17a)
0
a s.t.
Ot Pc (
= v. (u,( pe (x))
Apr(Xe)
Po 11P1
(17b)
Under mild regularity assumptions (Anderson, 1982: Yong & Zhou, 1999) such that Leibniz rule and Fubini's ' Theorem apply, we can separate Po, as Specifically, the objective value (17a) can be decomposed into
as it is assumed to be fixed out of the marginal
di S Epo.1
Ep.
1
12
+ V1
Bpgaaa
|luall2 + v.. de, 2
by law of total expectation
2
(21luell++ + Vi dt by Fubini's Theorem
Epo.
Eppsio.
which recovers (6a). Similarly, each term in the FPE (17b) can be decomposed into
a (x = 0
01. pc (X 21 21) (2o, xn )dro, = Exo,
ot
Po
(ur(X)) (c()))) = v. (uixx.) / Pe (o0, x pp. (xo, x1 )
= =v. (f
)o. )dzo.
= Epo.. [V- (ue P().))]
(x,(X,) = (Vpe(X.))
(X(el2o. x1)dzo,1)
=
Epo.1
V Vpro..1]
Epo..
Aprja.1
and, finally, (x(o)) = Ep. 18zo (x)] and (1()) Epo.. S0, (z)). Collecting all related terms yields:
(we P:jo.) Apq0.1 Po = 65. = 6zy;
P:jo.1 Øt
which is equivalent to the conditional SDE in (6b). Note that we denote u, as u. (x ( to emphasize the fact that when Prjo. (xx is factorized of p. (X,). the GSB problem (3) can be factorized into mixture of SOC problems, each with end-point constraint #1).
16

---

Published as a conference paper at ICLR 2024
Remark (PDE interpretation of Proposition 2 and (6)). The optimal control to (6) is given by (X, 21) = o²V log V2 (X,(z))) where the tare-ng potential (x(X, solves a par- tial differential equation (PDE) known as the Hamilton-Jacobi-Boor Bellman (HJB) PDE:
Ø1 (.(((=)) - - AV (x)a))
+ (((x)) (I(()) - ōm (x).
(18)
In general, (18) lacks closed-form solutions, except in specific instances. For example, when V is quadratic, the solution is provided in (19). Additionally, when V is degenerate, (18) simplifies to heat kernel, and its solution corresponds to the drift of the Brownian bridge utilized in DSBM (Shi
et al., 2023). Otherwise, one can approximate its solution with the aid of path-integral theory, as shown in Prop. 4).
Lemma 3 (Analytic solution to (6) for quadratic V and a > 0). Les ((x) == a a|]0 a, > 0.
then the optimal solution (6) follows Gaussian path X; ~ ~N(coo + G. 78 Ia), where
sinhh(((1tt))
e cosh(n(1111) - Ce cosh 7t - sinh(n (11
6 11 = ov ov
sinh
These coefficients recover Brownian bridges
O and, if further a + 0, straight lines.
This is direct consequence of conditioned diffusion processes with quadratic killing rates (Mazzolo & Monthus, 2022). Specifically, Mazzolo & Monthus (2022, Eq. (84)) give the analytic expression to the optimal control of (6), when V2 == al|oz||?
Proof.
ue (x(Xx/ x1)
sinh +)) Fi tanh(n(11t)) X 7= 22.
17
(19)
As (19) suggests linear SDE, its mean be Fd solves an ODE (Särkkä & Solin, 2019):
dyt
dt sinh sinh(n(1-()) 21 (anh((((1(t)
whose analytic solution is given by
31 S
90-
dr +
(inb(((1-T))
14
(20)
9t
where K depends on the initial condition and
tanh()(1-7)) dr) )
7
( exp [In - ( sinh(r((10))
sinh
9t ==exp
Note that (*) is due to change of variable V, sinh(m((17)). Substituting Ie back to (20), and noticing that zo po 2 K, yields the desired coefficients:
1
sinh(n(1-t))
#r c + etT1,
Ct sinh
e cosh(n(1-tt) a cosh
(21)
Similarly, the variance Eg e daxd of the SDE in (19) solves an ODE
dE,
20 E, 2
dt tanh(n(11-))
Repeating similar derivation, as in solving fr, leads to
g² Σ, sinh 2 (n(1 (coth)(())))))))))))))))))))) - coth 17) la (in)()(1-
o²
(22)
17
which gives the desired 7. since Σ, = -
We now show how these coefficients (21) and (22) recover Brownian bridge as a or, equivalently,
7 = /222 approaches the zero limit.
(1-1)) - ((-rt) (+) 1 11(1 - t) - (1 (01-7(1- ((1-))) t)) -- e --
lim C = lim
7-0
n-+0
1 17 - (1 (1-1)
lim e, lim
(cosh(n(1-t)) cs cosh +0
(1 - t) -1=t,
4-00
lim 7t lim σ 40
sinh(n(1 -t))
(a#)
0v(11 t)t,
7-0
17

---

Published as a conference paper at ICLR 2024
where (+) and (**) are respectively due to exp(2) ~ 1 + and
(imh(((1-6))
e 1) v(2-) 1 + 7/1 -1 (1-1(111-1 (-11))) =
lim 7-00
lim
7
2n
2rj
Hence, we recover the analytic solution to Brownian bridge. It can be readily seen that, if further + O, the solution collapses to linear interpolation between æo and 2T
Remark (Lemma 3 satisfies the boundary conditions in (6b)). One can verify that co = eg = 0, and 70 = O. Furthermore, we have satisfies the boundary condition in (6b).
0, = 1, and 7 0. Hence, as expected, Lemma
Proposition 4 (Path integral solution to (6)). Let r(X [r0, x1) be a distribution absolutely contin- uous w.r.t the Brownian motion, denoted q([[zo) where shorthand X Xxe(0,.1- Suppose is associated with an SDE dX, = vx(X,)) + odW, Xo=ro. XI=I1, 0>0.. Then, the optimal, path-ini integral solution (6) can be obtained
y" (X(o0,z1) = w wo In (X|o
(10)
where Z is the normalization constant and w is the importance weight:
(( == exp (-J' 0z Vi(X.) (11x(()))1) dt - (xx)
1
(11)
Proof. As the terminal boundary condition of (6b) is pinned at 21, we can transform (6) into:
min /' 0.1 J
1 vi(),) dt - log 1m, (X1),
Ex
11 (X, æo: )| + σ²
(23a)
20
s.t. dX,- = un (X, Io, (o0 21 )dl odW, Xo = 00
(23b)
where 1-(-) is the indicator function of the set A. Hence, the terminal cost "- log L, (x)" vanishes at T2 but otherwise explodes. Equation (23) a valid stochastic optimal control (SOC) problem, and its optimal solution can be obtained as the form of path integral (Kappen, 2005):
1 p" (X 21) Z'
exp (-o,
[ vi(x)) dt 12, (X1))(XXzz))
1
(24)
where (((11)) the density of the Brownian motion conditioned on Xo = zu, the normalization constant Z' is such that (24) remains
proper distribution. and shorthand X X 1) We highlight that Equations (23) and (24) are essential transformation that allow us to recover the conditional distribution used in prior works (Shi al., 2023) when V == 0 (see the remark below)
Directly re-weighting samples from q using (24) may have poor complexity as most samples are assigned with zero weights due 1z, (-). A more efficient alternative is to rebase the sampling distribution in (24) from g to an importance sampling r( [2o, x1) such that Xo = Ty. Xi = I and absolutely continuous with respect to q. Then, the remarkable results from iiformation-thooeete SOC (Theodorou et al., 2010; Theodorou, 2015) suggest
p' (' 21) Z exp (- vix.)))
(((10)) r(X ( 21), r( Joo x1)
(25)
where the Radon-Nikodym derivative & can be computed via Girsanov's theorem (Särkkä & Solin, 2019):
q(((z)) exp (-/," 20 10(x)) [u
(26)
r(X I .1)
Substituting (26) back to (25) concludes the proof.
Remark (Equation (24) recovers Brownian bridge when V 0). One verify that, when Vi 0, the optimal solution px Io, æ1) x (((22)) (X,) = g(X I xi ) is simply the Brownian motion conditioned on the end point being 21, which is precisely the Brownian bridge.
18

---

Published as a conference paper at ICLR 2024
Theorem 5 (Local convergence). Let the intermediate result of Stage and 2, after repeating times. be 8". Then, its objective value in (3), ) is monotonically non-increasing as n grows:
[(0") 2 L(On+!).
Proof. Let pf. and u. be the marginal distribution and vector field, respectively, after n steps of alternating optimization according to Alg. 5, i.e., p. and uf" satisfy the FPE. Furthermore, let P6,1 and Pf be the coupling and conditional distribution, respectively, defined by uf" and let Pajo, be the solution to CondSOC (6). This implies that the marginal distributions at step + 1 are given by
Pt gu+l
f 1811r o
(27)
Now, under mild assumptions (Anderson, 1982; Yong & Zhou, 1999) such that all distributions approach zero at a sufficient speed as ||>||| - oo, and that all integrands are bounded, we have (note that inputs to all functions are dropped for notational simplicity):
/
L(0
Pt 8m
12+ V1 drædft
f
gr 1
gr
(2 511 112 +++1 +V: d dz dt d by Fubini's Theorem
P1jo.1
P
g"
Pipo.
( 21mm²" 112 11 V3 do
1
1
1 by optimizing (6)
f /
gm-1
112 11 Vi dr;dt
by Fubini's Theorem and (27)
Pt
ga-1
1
go+t
112 V/ dr
(28)
C(O
where we factorize the solution of Stage 1 by rf" = 1 d The inequality in (28) follows from the fact that u as the solution to the proceeding Stage 1, finds the unique minimizer that yields the same marginal as pg=+1 while minimizing kinetic energy, i.e.
gh+t f
Pt
1 8" 12 dx. f pt"
gn+
||2 dæy.
Theorem 6. The optimal solution to GSB problem (3) is a fixed point of GSBM, Alg. 5.
Proof. Let pf and of be the optimal solution to the GSB problem (3), and suppose pt dr Prrdde It suffices to show that
1. Given Zo, 31 p the conditional distribution Płjo,2 is the optimal solution to (6).
2. Given pt, both matching algorithms (Alg. and 2) return u7.
The first statement follows directly from the Forward forward-Backward Backward SDE (FBSDE) representation of (3), initially derived in Liu et al. (2022, Theorem 2). The FBSDE theory suggests that the (conditional) optimal control bridging any Zp, I1 ~ Pō,1 satisfies a BSDE (Pardoux & Peng, 2005) that uniquely associated with the HJB PDE of (6), i.e., (18). This readily implies that p is the solution to (6).
Since " is known to be a gradient field (Liu et al., 2022, Eq. (3)), it must be with the solution returned by the implicit matching algorithm (Alg. 1), as Limplicu returns the unique gradient field that matches pf. On the other hand, since the explicit matching loss (5) can be interpreted as a Markovian projection (Shi et al., 2023), i.e., it returns the closest (in the KL sense) Markovian process to the reciprocal path measure (Léonnrd, 2013; Léonard al., 2014) defined by the solution (6). Since
19

---

Published as a conference paper at ICLR 2024
the solution to (6), as proven in the first statement, is simply P the closest Markovian process by construction u,- Hence, the second statement also holds, and conclude the proof. It is important to note that, while Lexplicit and Limplicit do not share the same minimizer in general, they do the equilibrium Pi.
C ADDITIONAL DERIVATIONS AND DISCUSSIONS
C.I EXPLICIT MATCHING LOSS
How (5) preserves the prescribed pt. A rigorous derivation of (5) can be found in, egg., Shi et al. (2023, Proposition 2), called Markovian projection. Here, provide an alternative derivation that follows closer to the one from flow matching (Lipman et al., 2023).
Lemma 7. Let the marginal Pt be constructed from mixture of conditional probability paths, ie., Pe(x)) := po [Pe( T1)], where pe (a T1 = Peju,1 is the time marginal of the SDE, dX, ue( Xx K2| x1 )dl + odW, Xo = zD, X1 = 21, then the SDE drift that satisfies the FPE
prescribed by Pt. given by
ug (I) Ep [me (2r/20 2n Drooo (x/k0, 21)]
(29)
((())
Proof. It suffices to check that ut (1) satisfies the FPE prescribed by Pr:
/ ( 0t Pr Po.1d
or P((x)
/ - (urjo PHoo.) do. + / (6o0 Apioa) po.1d2o.1
(9)
(u. pc) + 2
0² Apr.
immediate consequence of Lemma is that
Lesplici (0) = Ep. 2 (xi) - (x,) + 0(1,
(30)
where O(1) is independent 0. Hence, the uf preserves the prescribed pr.
Relation to implicit matching (4). Both explicit and implicit matching losses are associated to some regression objectives, except w.r.t different targets, i.e., u² in (30) vs. Vs in Sec. 2. As pointed out in Neklyudov et al. (2023), the two targets relate to each other via the Helmholtz de- composition (Ambrosio et al., 2008, Lemma 8.4.2), which suggests that of = Vst + we where wi is the divergence- vector field, i.e., V (w.Pe) O. Though this implies that Leaplica only upper-bounds the kinetic energy, its solution sequential, by alternating between solving (6) in Stage remains well-defined from the measure perspective (Peluchetti, 2022: 2023; Shi al., 2023). Specifically, the sequential performs alternate projection between reciprocal path measure defined by the solution to (6), in the form p; == E [P.j..],, and the Markovian path measure, and admits convergence to the optimal solution for standard SB problems.
C.2 GAUSSIAN PROBABILITY PATH
Derivation of analytic conditional drift (8). Recall the Gaussian path approximation in (7):
X, = pe + 7Z NO, Ia).
which immediately implies the velocity vector field and the score function (Särkkä & Solin, 2019; Albergo et al., 2023):
0.7 (X, 14), Vlogp((X.)
1 (X, p4)
0,X, O.44
71
77
20

---

Published as a conference paper at ICLR 2024
Ground Truth
Alg.
Algorithm CovSample (line in Alg. 4)
Require: SDE (6b) with Pe:
Discretize [t,.... ,tk] 0 1 <
< tk <1
Solve Di according the 1D ODE in (32)
Compute covariance matrix C₁ R by (33)
Perform Cholesky decomposition C Likg
Figure 10: How Alg. 6 faithfully re- covers the covariance matrix of, in this case, Brownian bridge.
Sample X₃ Mr L,Z,, Ze E R [Zi]ss ())
return X;
We can then construct the conditional drift:
a²
o² (X, mc).
" (x(xx|oro==31 x1) 0,X, logP:(X,) 0,14
(31)
7t
2%/
Notice that (31) is of the form of gradient field due to the linearity in X₁. One can verify that substituting the Brownian bridge, H4 == (1 () + tæı and == / t), to (31) indeed yields the desired drift -1=X:
Efficient simulation with analytic covariance function (Footnote 4). Proposition 4 requires samples from the "joint" distribution in path space. This requires sequential simulation, which can scale poorly to higher-dimensional applications. Fortunately, there exists efficient computation to (6b) with the (linear) conditional drift uejo,1 given by (8), as the analytic solution to (6b) reads
(Io +
[ e b fee
-% u,
X, = al
9t :
(32)
IR is defined in (8), b. 0.p atte, and f dW, is the Ito stochastic integral (Itō, 1951). The covariance function between two time steps 8, I, such that < S 1, can then be computed by
where at
74 17 27,
o²
t f
- 9- odW, (" 3. odW>)
Cov(s,()
odW,, f e odW,)
g t
(i)
c +0e
2u- dddt
0c 7
where (i) due to the independence of the Ito integral between [0, s) and [s, f). (fi) is due to the Ito isometry, and, finally, (iit) is due to substituting 7 Var(s) == 2ag S e 2g. dddt.
Repeating the same derivation for t, such that 0 < between any given two timesteps 8, E [0, 1] can be written cleanly as
1, the covariance function of (6b)
Cov(s, 6, L) = 2
(st)(s.t) e.Bmax -miin(a)))
(33)
which, crucially, requires only solving an "1D" ODE, dg. = a,tt, go We summarize the algo- rithm in Alg. 6 with example of Brownian bridge in Fig. 10. We note again that all computation parallelizable over the batches.
C.3 ANALYTIC SOLUTION TO (6) FOR QUADRATIC V AND o =0
Recall the finite-horizer and continuous-timg linear quadratic regulator, defined as:
min x Fx1 L [x] Qx u, Rue] dt s.t. de Ax, + Bug, We
(34)
21

---

Published as a conference paper at ICLR 2024
whose optimal control ui -R-1BPar can be obtained by solving Riccati differential equation:
P. P. P.A P.BR Q:
P = F,
or, equivalently, Lyapunov differential equation (one can verify that H; Pr for all f):
H1 AH, H,A BR HQQH, H1 =
(35)
The end-point constraint can be accounted by setting F - x., suggested in Chen & Georgiou (2015), yielding H, = With that, the solution to (6) for quadratie V(x) = a||=|| and a == can be obtained by solving the matrix ODE in (35) with A 0, B Ia. R L. and Q = ald, which admits an analytic solution in the form of hyperbolic tangent functions, similar to (19).
D EXPERIMENT DETAILS
D.1 EXPERIMENT SETUP
Table 7: Hyperparameters of the SplineOpt (Alg. 3) for each task.
Stunnel
Vneck
GMM
Lidar
AFHQ
Opinion
Number of control pts. K
30
30
30
30
8
30
Number of gradient steps M
1000
3000
2000
200
100
700
Number of samples Optimizer
4
4
4
4
4
4
SGD
SGD
SGD
mSGD
Adam
SGD
Baselines. All experiments DeepGSB are run with their official implementation" and default hyperparameters. We adopt the "actor-critic" parameterization as generally yields better perfor- mance, despite requiring additional value networks. On the other hand, implement DSBM by ourselves, as our GSBM can be made equivalent to DSBM by disabling the optimization of the CondSOC problem in (6) and returning the analytic solution of Brownian bridges instead. This al- lows us to more effectively ablate the algorithmic differences, ensuring that any performance gaps attributed the presence of V,, All methods, including our GSBM, are implemented PyTorch (Paszke et al., 2019).
Network architectures. For crowd navigation (Sec. 4. and opinion depolarization (Sec. 4.3), we adopt the same architectures from DeepGSB, which consists of 4 to residual blocks with sinusoidal time embedding. For the AFHQ task, we consider the U-Net (Ronneberger et al., 2015) architecture implemented by Dhariwal Nichol (2021). 6 All networks are trained from scratch, without utilizing any pretrained checkpoint, and optimized with AdamW (Loshchilov & Hutter, 2019).
Task-specific noise level (o). For crowd navigation tasks with mean-field cost, we adopt {1.0, whereas the opinion depolarization task uses a 0.5. These values are inherited from DeepGSB. On the other hand,
and 0.5 respectively for LiDAR and AFHQ tasks.
GSBM hyperparameters. Table 7 summarizes the hyperparameters used the spline optimiza- tion. By default, the generation processes are discretized into 1000 steps, except for the opinion depolarization task, where we follow DeepGSB setup and discretize into 300 steps.
GSBM implementation (forward & backward scheme). In practice, we employ the same "for- ward and backward" scheme proposed in DSBM (Shi et al., 2023), parameterizing two drifts, one for the forward SDE and another for the backward. During odd epochs, we simulate the coupling (line 4 in Alg. 5) from the forward drift, solve the corresponding CondSOC problem (6), then match the resulting Pt with the backward drift. Conversely, during epochs, we follow the reverse pro- cess, matching the forward drift with the Pt obtained from backward drift. The forward-backward alternating scheme generally improves the performance, as the forward drift always matches the ground -oound-truth terminal distribution P1 = y/ (and vise versa for the backward drift).
'https: :/ hww.thhh.b...hhh......coeoo u / DeepGSH, under Apache License.
'https:/ //gi thub. com openai did-di f fuusi on, under MIT License.
22

---

Published as a conference paper at ICLR 2024
Crowd navigation setup. All three mean-field tasks Stunnel, Vneck, and GMM are adopted from DeepGSB, as shown in Fig. 11. We slightly modify the initial distribution of GMM to testify fully multi-model distributions. As for the mean-field interaction cost (12), we consider entropy cost for the Vneck task, and congestion cost for the Stunnel and GMM tasks. We adjust the multiplicative factors between Lobstacle and Linteract to ensure that noticeable changes occur when enabling mean-field interaction; see Table 8. These factors are much larger than the ones considered in DeepGSB (A2 < 3). In practice, we soften the obstacle cost for differentiability, simi- lar to prior works (Ruthotto et al., 2020; Lin et al., 2021), ) and approximate (a)) se E(s..1)) Pe (I (ro, T1) as the mixture of Gaussians.
Stunnel
Vneck
GMM
Figure 11: Initial and terminal dis-
tributions for each crowd naviga- tion task with mean-field state cost. Obstacles are marked gray.
Table 8:
Multiplicative fac- tors of the state cost for each crowd navigation task: V((2;)
Aobs oosstacc (4) Linterect (2; pt).
AFHQ setup. As state cost Ve is defined via latent space, we pretrain a 3-VAE (Higgins et al., 2016) with = 0.1. On the other hand, the spherical linear interpolation (Slerp; Shoemake (1985)) refers to constant-speed rotational motion along a great circle are:
Stunnel
Vneck
GMM
Aabs Aint
1500
3000
1500
50
sin 6) 12)
sin(t ?
Islep(( 20; 21)
sin 52 20
sin 21;
arccos((zo: (10)))))
(36)
D.2 OPINION DEPOLARIZATION
t- 0.00 (w)
0.25
t 0.50
t- 0.75
t 1.100
PCA vs.
DS
Figure 12: Simulation of the polarize drift (37) in 1000
Polarize drift in (15). We use the same polarize drift from DeepGSB, based on the party model (Gaitonde al., 2021). At each time step f, all agents receive the same random information & Rd sampled independently of pt. then react to this information according to
Fpolarizo (2; Ex) = Eyno. [a(x, y, E,) (() = {
if sign((æ, sign (w, (v)
otherwise
(37)
where ÿ := :0 a/Ilul and a(z, y, &t) the agreement fundtion-iinicates whether the two opinions æ and y agree on the information & Hence, (37) suggests that the agents are inclined to be receptive to opinions they agree with while displaying antagonism towards opinions they disagree with. This known to yield polarization, shown in Figure 12.
Directional similarity in Figures and 12. Directional similarity is standard visualization for opinion modeling (Schweighofer et al., 2020) that counts the histogram of cosine angle between pairwise opinions. Hence, flatter directional similarity suggests less polarize opinion distribution.
D.3 APPROXIMATING GEOMETRIC MANIFOLDS WITH LIDAR DATA
Since LiDAR data collection of point clouds in [-5, 5]3 C R33 we use standard methods for treating it more Riemannian manifold. For every point in the ambient space, define the projection operator by first taking a k-nearest neighbors and fitting 2D tangent plane. Net.N( (z)
(.1...., ] be the set of k-nearest neighbors for query point R3 in the ambient space. We
23

---

Published as a conference paper at ICLR 2024
then fit a 2D plane through moving least-s approach (Levin, 1998; Wendland, 2004),
arg min a.b,c
Σ w (a), bor x (.(,)/2
(38)
where the superscripts denotes the I, y: coordinates and use the weighting w(2, x.) exp(-Ilæ- - x with T = 0.001. We solve this through a pseudoinverse and obtain the ap- + by + 2. When k - oc, this tangent plane smooth: however, we find that using k = 20 works sufficiently well in our experiments, and our GSBM algorithm is robust the value of k. The projection operator I (2) then defined using the plane,
proximate tangent plane
+ <(x) = x - |12122
n,
where n [a b -11
(39)
This projection operator is all we need to treat the LiDAR dataset as manifold. Differentiating through will automatically project the gradient onto the tangent plane. This will ensure that optimization through the state cost V, will be appropriated projected onto the tangent plane, allowing us to optimize quantities such as the height of the trajectory over the manifold.
The exact state cost we use includes an additional boundary constraint:
V(x)= > [[(()) =l| exp(T(x)(e))
(gm(n-5/o.1) + (1 sigm(r+-//.1)
DE (a())
Loaadiial
Draigtt
Leuudayy
(40)
where "sigm" is the sigmoid function, used for relaxing the boundary constraint. The loss Looundry simply ensures don't leave the where the LiDAR data exists. We set À 5000. The gradient of (x) moves x in a direction that is orthogonal to the tangent plane while the gradient of (z) moves along directions that lie the tangent plane. If only have mannfflld- then CondSOC essentially solves for geodesic paths and we recover an approximation of Riemannian The state cost Lneight depends on (xz) as this ensures we travel down the mountain slope when optimizing for height. Other state costs can of course also be considered instead of Lneight-
Flow Matching (Chen Lipman, 2023) (parameterized in the ambient space) when
D.4 ADDITIONAL EXPERIMENT RESULTS
Xi)
X1|Ko,
Stunnel
X1 ==0
XiKKo.
Spiral - Moon
Figure 13: Illustration of how different couplings, X1|zo, emerge with nontrivial state costs V V(x).
Nontrivial V(x) induces different couplings. Figure 13 illustrates how different couplings, X, |xo, emerge with nontrivial state costs Vxx). Specifically, we color different regions of the initial distribution (lefumost column) with different colors and track their pushforward maps. We consider the same Vxx), i.e., obstacle and congestion costs, for Stunnel (top row) and adopt quadratic cost for Spiral - Moon (bottom row). In these two particular cases, having nontrivial V(x) encour- ages stronger mixing, thereby yielding a different coupling (rightmost column) compared to the one induced without ((x) (middle column).
24

---

Published as a conference paper at ICLR 2024
DSBM (V, = 0)
GSBM (V₃ = 0)
t-0.00
t-
t-0.75
t-0.00
t-0.75
Figure 14: Additional comparison between DSBM (Shi et al., 2023) and our GSBM through the mean of (x(,, (0,21) with randomly sampled (1). While DSBM simply performs linear interpolation between To and æ1, our GSBM optimizes w.r.t. (6) where V is defined via latent space, thereby exhibiting more semantically meaningful interpolations.
Table 9: Quantitative comparison between NLSB (Koshizuka & Sato, 2023) and our GSBM. Notice that our GSBM consistently ensures feasibility. In contrast, as NLSB approaches the GSB problem by transforming it into stochastic optimal control (SOC) problem with soft terminal cost (see Appendix A), it may overly emphasize on optimality at the expense of feasibility.
Feasibility W(r v)
Optimality (3)
Stunnel
Vneck
GMM
Stunnel
Vneck
GMM
NLSB (Koshizuka Sato, 2023)
30.54
0.02
67.76
207.06
147.85
4202.71
GSBM (ours)
0.03
0.01
4.13
460.88
155.53
229.12
Stunnel 2-0.33-0.66
Vneck t-0.33-0.66
--0.00-0.33
r-0.66-1.00
t-0.00-0.33
t-0 r-0.66-1.00
1-0.00-0.33
r-0.33-0.66
t-0.66-1100
Figure 15: Performance of NLSB (Koshizuka & Sato, 2023) on the same crowd navigation tasks. Notice how NLSB trapped some local minima Stunnel and completely failed on GMM.
Additional interpolation results on AFHQ. Figure 14 provides comparison results between DSBM (Shi et al., 2023) and our GSBM on Pe(X, 21 with randomly sampled (Ir,
Additional comparisons on crowd navigation with mean-field cost. Table 9 and Fig. 15 report performance of NLSB" (Koshizuka & Sato, 2023)-an adjoint-based method for approximat-
hhtpps://gg thub. wtw.coooeeee./eeeeo koshi zuka 7 sbs under MIT License.
25

---

Published as a conference paper at ICLR 2024
ing solutions to the same GSB problem (3). Figure 16 provides additional comparison between
DeepGSB (Liu et al., 2022) and our GSBM. should be obvious that our GSBM outperforms both methods.
DeepGSs
GSBM (ours)
1=0:00-0:25
1=0.5-00.50
t=0.50-0.75
t=0.75-1.00
0:0.5000:25
t=0.25
t=0: 50-055
t= 1=0.75-1.00
DeepGSB
GSBM (ours)
1-0.00-0.25
1-0.55-0.50
1-0.75-1.00
2-0.00-0.25
7-0.55-0:50
1-0.75-1.00
113 Treining
2/3 Treining
SS Preming
DeepGSB
0000000-.25
t=.525-0.50
t=0.50-0.75
1=0.75-1.00
1=0.00--.25
0.0.25-0.50
t=0.5000.75
t=0./5-1.00
222 Sairinn
3/3 Training
Figure 16:
Additional comparison betwee DeepGSB (Liu et al., 2022) our GSBM on (top to bottom)
Stunnel with a = 2.0 and Vneck with a = 11.0,.00, Notice again how the training of DeepGSB exhibits relative instability and occasional divergence.
26


Reviews:
[{"review_id": "0dYaHgk6PF", "reviewer": "Reviewer_7Jx7", "summary": "The authors consider so-called generalised Schrödinger Bridge Matching problem: unlike standard SBM, which can be formulated as an optimisation problem of a drift of some diffusion, the corresponding optimisation problem in case of GSBM involved an additional state cost. The authors proposed an iterative algorithm to solve the problem. One of interesting tricks, proposed to increase computational efficiency of the solution, is a spline-based  approximation. The authors verified the algorithm for GSBM problem on a number of toy and practical tasks.", "questions": "- how do computational efficiency and accuracy scale w.r.t. dimensionality? sample size?\n\n- how do specific strategies for placement of spline control points influence accuracy (see sec. 3.2)?\n\n- Gaussian path approximation in (7) looks like a strong assumption. It is not clear from theorem 5 whether it is always possible to decrease L in case we base our computational algorithm on such assumption. Are there any other assumptions such that we can get closed form expressions like (8)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the authors consider an import problem statement, which is a very natural generalisation of the SBM problem\n\n- the proposed approach is theoretically sound\n\n- the authors test their approach on a diverse set of problems\n\n- the text is easy to follow", "weaknesses": "- the authors did not investigate computational limits of the proposed approach \n\n- it seems that the proposed approach does not work in case of image data, only if we consider some latent space like in sec. 4.2. Any ideas how we can improve here?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cORsfXyuKy", "reviewer": "Reviewer_cBUM", "summary": "This paper focuses on the problem of controlled measure transport. Given a time dependent density $p_t(x)$, how do you adapt $p_t(x)$ with respect to some potential function $V_t(x)$ so that it arrives at a target density $p_{T}(x)$ and does so in some measure of least cost (or as they refer to it, energy). \n\nThey rely on recent distribution matching algorithms to do this, writing the discovery of some vector field $s_t(x)$ and the least cost set of $p_t(x)$ that correspond in Fokker-Planck sense to this $s_t(x)$ as the solution to two constrained minimization problems. \n\nThe reviewer interprets the \"generalized\" in the title to refer to the fact that the paper moves away from the notion of just connecting two densities, but also opens the avenue for doing this e.g. with obstacles (a la optimal control problems). \n\nThey demonstrate their method e.g. on toy blockade examples and with the effects of stochasticity on the energy minimization, as well on image-to-image translation examples and a cool LiDAR surface example.", "questions": "- A comment rather than a question: For an exciting read on the problem, the authors should see [6] , it's quite fun.\n- Can the authors come up with a problem for which conditional density is actually far from Gaussian, or show that any of their examples currently are far from Gaussian? \n- Can the authors try to quantify the \"semantically meaningful\" coupling?\n\n\n\n\n\n[6] Stochastic control liaisons: Richard Sinkhorn meets Gaspard Monge on a Schroedinger bridge, *Yongxin Chen, Tryphon T. Georgiou, Michele Pavon*, 2020.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The authors exposit their work quite well (with adjustments needed mentioned in the weaknesses and questions section). The experimental section is thorough and highlights the myriad applications of the perspectives presented in the paper. With the proper contextualization, the method presented is interesting and falls in line well with a series of other related works, highlighting the benefits of controllable paths (for the limited set of potentials the consider).", "weaknesses": "While the experiments are cool and thorough, and while the framing of the optimization problem is coherent, the paper diminishes itself by its framing. *This is an addressable issue* and the reviewer thinks it can be straightforwardly, especially during a rebuttal period, as it's not an experimental thing. The framing issue is that it oversells itself in a way that misleads the reader. \n\nI think a good principle of academia (that we can all reasonably agree with) is to only assert what we can reasonably back-up and to not mislead. This paper does not provide a numerical solution to optimal transport or the schrodinger bridge. In the introduction, it casts OT and Schrodinger Bridge as edge-cases of their variational problem. *However, they do not provide an actual solution to that variational problem*. The statement of this problem has been well known. The **application** of it in the context of generative modeling and dynamical systems is compelling (even if approximate! you don't need to solve OT/SB to do cool control!) but it does not need to be couched as if it is solving those problems. It is still compelling to say \"we came up with an approximate way to do stochastic optimal control in the context of deep generative modeling\" without saying \"we have provided a variational method for solving optimal transport and SB.\"\n\nThe issue is that the constraint (6b) is actually quite hard to enforce. The Gaussian approximation is fine. But the authors should really be careful in stating what it is they are doing and how it relates to the problem they are seeking to solve, because gaussian paths are not going to do it. The paper that this one primarily cites (Flow Matching for Generative Modeling) seems to do this too -- in which it specifies an \"optimal transport path\" that isn't actually providing a means to solve OT.\n\nThe reason the reviewer stresses this point is two-fold: the first is that mathematicians have spent many years carefully characterizing the hardness of these problems and approximate solutions without overstating them. The second is that doing this induces confusion to new readers who come into the field. This is a damaging effect, it muddles points and slows down new members, and one the reviewer has already had to deal with from students misinterpreting the claims of OT in these types of papers.\n\nThis can be addressed by simply toning down some of claims, and including these caveats, e.g. in the bullet points on page two summarizing the papers contributions.\n\n**References**\n\nThere are some necessary adjustments to the citations needed, but that can be cleared up. Here are some things that should be addressed:\n\n- At the beginning of section 3.1: Neither [Liu 2023b] nor [Peluchetti 2022] propose any sort of entropic optimal transport. Perhaps there was a typo in the citep. While Liu 2023b proposes a procedure for OT (when the iterating field is a gradient field), nothing about it is entropic.\n- There seems to be some missing citations with regards to the matching algorithms. On page one following the line \"a mixture of tractable conditional probability paths (Ho, Song, Lipman)\" should include [1,2,3] listed below. These should likewise be included Bridge and Flow matching (explicit) section. There is also related follow-up work by Tong et al listed as [4]\n- The Schrodinger Bridge problem in the introduction should not be attributed to the machine learners. It's been studied for decades. The original citation is provided in [5], and a follow-up is provided in [6]. \n\n\n**Other adjustments**\n- The discussion about feasibility seems not so useful. The statement of the problem the paper wants to solve is that the solution is feasible (it's chicken-egg). That's the point of a constraint in an optimization problem. If the one paper the authors want to compare this work to (the Generalized Schrodinger Bridge one) fails to meet this, it doesn't seem worth taking the reader down that line when the problem statement assumes it. Table 1 could be in an appendix :) \n- There are some minor things throughout the text to adjust, which are semantically inexact. For example: in the line in the \"Implicit action matching\" paragraph:  \"...then the unique gradient field $\\nabla s^\\theta_t(X_t)$ that matches $p_t$ can be obtained by minimizing...\". The noun comparison is wrong here. $\\nabla s^\\theta_t(Xt)$ does not match a probability density. The learning problem is to find the $\\nabla s^\\theta_t(X_t)$ for which the probability density *of the pushforward* from the initial density $p_0$ to $\\tilde p_t$ by this $s_t^\\theta$ matches $p_t$. Densities match densities.\n- The authors should adjust the first paragraph of the second page to highlight that optimal transport as a field does consider a wide variety of other costs besides the $\\mathcal W-2$ cost. There are e.g. repulsive costs, e.g. in the study of quantum dynamics, etc. The notion of optimality still holds.\n- Lastly, are the authors sure that a reasonably different coupling $p(x_0, x_1)$ emerges at the minimizer? Neither set of experiments seems to suggest this. While it is true in the image-to-image translation, that the intermediate density $p_t(x)$ has a more semantically meaningful translation, the marginal endpoints (the $x_0$ associated with $x_1$) for both algorithms look pretty much identical. Moreover, Figure 1 highlights an equivalent picture. The initial and final empirical densities (light green and dark green) are identical even after learning. \n\n\n**The reviewer is more than happy to improve their score if the framing, citations, and \"other adjustments\"**. As of now I would ideally mark this as a 4/10 and not a 5/10, but that isn't an option any more it seems.\n\n[1] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. Xingchao Liu, Chengyue Gong, Qiang Liu, Sept 2022.\n\n[2] Building Normalizing Flows with Stochastic Interpolants. Michael S. Albergo and Eric Vanden-Eijnden, Sept 2022.\n\n[3] Stochastic Interpolants: A Unifying Framework for Flows and Diffusions. Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, March 2023.\n\n[4] Improving and generalizing flow-based generative models with minibatch optimal transport. *Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, Yoshua Bengio*, July 2023.\n\n[5] Schr¨odinger, E., “Uber die Umkehrung der Naturgesetze,” ¨ Sitzungsberichte der Preuss Akad. Wissen. Phys. Math. Klasse, Sonderausgabe, Vol. IX, 1931, pp. 144–153.\n\n[6] R. Fortet, R´esolution d’un syst`eme d’equations de M. Schr¨odinger, J. Math. Pure Appl. IX (1940), 83-105.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MgZW3Y0LOw", "reviewer": "Reviewer_AA2i", "summary": "The paper considers a novel method for the generalized Schrodinger bridge method, which allows for incorporation of a control cost into the interpolation. Their method works by alternating between optimizing the drift and the marginals, with the second done via a tie to conditional stochastic optimal control. In this second step, the optimal probability path between two fixed endpoint marginal samples is approximated with a Gaussian spline (splines separately for the time-varying means and standard deviations). The method is tested and compared primarily against DeepGSB in crowd navigation, image interpolation, and opinion depolarization.", "questions": "In relation to the first item above: \n1. Can you prove convergence to a global optimum for a particular class of state cost V?\n2. If not, do you have any commentary on other strategies for initialization of the spline optimization?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a method which improves upon previous methods by ensuring guaranteed feasibility and stronger performance in high dimension. The method also seems to clearly outperform comparable methods in the various application settings.", "weaknesses": "* The convergence guarantees obtained are merely local, and do not guarantee convergence to a global optimum. \n* As noted in the conclusion, the method requires a differentiable state cost.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gayKO7lDJ5", "reviewer": "Reviewer_fcZ9", "summary": "This paper propose Generalized Schrodinger Bridge Matching (GSBM) that can be view as solving conditional stochastic optimal control, for which efficient variational approximations can be used, and further debiased with the aid of path integral theory.  Compared to other methods for solving GSB problems, GSBM algorithm preserves a feasible transport map between the boundary distributions throughout training, thereby enabling stable convergence and significantly improved scalability.", "questions": "Does (6) have an analytic solution only if V_{t} is quadratic and \\sigma>0?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper has rigorous logic, clear expression and abundant experimental details. This work broadens the scope of application of Schrodinger bridge matching (SBM) algorithm and converts the existing optimal transport mapping and SBM methods into the special case of variational formulas. In addition, the feasibility conditions are met at any time during the entire training process, which is a big difference from the previous methods. Overall, this work opens up new algorithmic opportunities for training diffusion models with task-specific optimal structures.", "weaknesses": "The Generalized Schodinger Bridge Matching (GSBM) proposed in this paper is not applicable to any form of V_{t}, but requires the differentiability of V_{t} and relies on quadratic control costs to establish its convergence analysis. Although it is a significant improvement over previous GSB solvers, the Generalized Schodinger Bridge Matching (GSBM) method is not applicable to any form of V_{t}. But differentiability is still a necessary condition.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)