Go With the Flow: Fast Diffusion for Gaussian Mixture Models
Authors: George Rapakoulias
Ali Reza Pedram
Fengjiao Liu
Lingjiong Zhu
Panagiotis Tsiotras
Venue: NeurIPS 2025 spotlight ()
Keywords: Schrodinger Bridge problem
Optimal transport
Primary area: theory
TL;DR: We propose a method to approximate the solution to Schrodinger Bridge and multi marginal momentum Schrodinger Bridge problems with discrete or continuous Gaussian Mixture Model marginal distributions.

Abstract: Schrodinger Bridges (SBs) are diffusion processes that steer,  in finite time, a given initial distribution to another final one while minimizing a suitable cost functional. Although various methods for computing SBs have recently been proposed in the literature, most of these approaches require computationally expensive training schemes, even for solving low-dimensional problems. In this work, we propose an analytic parametrization of a set of feasible policies for steering the distribution of a dynamical system from one Gaussian Mixture Model (GMM) to another. Instead of relying on standard non-convex optimization techniques, the optimal policy within the set can be approximated as the solution of a low-dimensional linear program whose dimension scales linearly with the number of components in each mixture. The proposed method generalizes naturally to more general classes of dynamical systems, such as controllable linear time-varying systems, enabling efficient solutions to multi-marginal momentum SBs between GMMs, a challenging distribution interpolation problem. We showcase the potential of this approach in low-to-moderate dimensional problems such as image-to-image translation in the latent space of an autoencoder, learning of cellular dynamics using multi-marginal momentum SBs, and various other examples. The implementation is publicly available at https://github.com/georgeRapa/GMMflow.

Go With the Flow: Fast Diffusion for Gaussian Mixture Models
George Rapakoulias
Ali Reza Pedram 1,2
Fengjiao Liu"
Lingjiong Zhu
Panagiotis Tsiotras
Department of Aerospace Engineering, Georgia Institute of Technology, Atlanta, GA
School of Computer Science, University of Oklahoma, Norman, OK 3 Department of ECE, FAMU-FSU College of Engineering, Tallahassee, FL.
4 Department of Mathematics, Florida State University, Tallahassee, FL.
Abstract
Schrödinger Bridges (SBs) are diffusion processes that steer, in finite time, given initial distribution to another final one while minimizing a suitable cost functional. Although various methods for computing SBs have recently been proposed in the literature, most of these approaches require computationally expensive training schemes, even for solving low-dimensional problems. In this work, we propose analytic parametrization of set of feasible policies for steering the distribution of dynamical system from one Gaussian Mixture Model (GMM) another. In- stead of relying on standard non-convex optimization techniques, the optimal pol- icy within the set can approximated the solution of -ow-diensional linear program whose dimension scales linearly with the number of components in each mixture. The proposed method generalizes naturally more general classes of dynamical systems, such as controllable linear time-varying systems, enabling ef- ficient solutions to multi-marginal momentum SBs between GMMs, challenging distribution interpolation problem. We showcase the potential of this approach in translation in the
www-to-moderate dimensional problems such as latent space of autoencoder, learning of cellular dynamics using multi-margine momentum SBs, and various other examples. The implementation is publicly available https:/ //g:::///tti http///gtttbbbtt thub. com/ georgeRapa./GMf/
Introduction and Background
The problem of finding mappings between distributions of data, originally known the Optimal Transport (OT) problem mathematics, has received significant attention in recent years in multiple research fields, due to its application in problems such generative AI (Ruthotto & Haber, 2021; Arjovsky et al., 2017), biology (Bunne et al., 2023b; Bunne & Rätsch, 2023; Tong et al., 2020), mean field problems (Liu et al., 2022) and control theory (Chen et al., 2015a,b; Rapakoulias & Tsiotras, 2024) among many others. Despite appearing static in nature, reformulating OT in the context of dynamical systems imbues it with further structure and unlocks tools from the literature on dynamical systems that can be employed for its efficient solution (Benamou & Brenier, 2000).
To the stage, consider two distributions Po: Pl, supported the d-dimensional Fuclidean space, denoted by R.d, and consider the regularized version of the static OT optimization problem, known the Entropic Optimal Transport (EOT) problem (Peyré Cuturi, 2019):
min (p. R 2 |00 1l x1) cH(T),
(1)
"Corresponding author: grap@gatech.c edu 39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

where (xo, 21) ) is the transport plan (also referred to as coupling) between Po. Pl, [p P1) is the set of all joint distributions with marginals Pa. Pt, and H is the differential entropy, defined by H() - - ʃ p(x) log p(x) dr. The corresponding dynamic formulation of the EOT problem is known as the Schrödinger Bridge Problem (SBP) (Léonard, 2014; Chen et al., 2021). When formulated as a stochastic optimal control problem, the SBP is given by
[x) (1) [%= 1 2c
[xz = (x(34) dt + vedw, In po,
min JSB A Exerp Nell
s.t.
(2)
Pt,
where the objective is to find an optimal drift function u, (22z), also referred as the control policy in the context of control applications, belonging to a set of adapted finite-energy policies U, such that, when applied to the stochastic dynamical system defined by the first constraint in (2), the marginal distribution specified in the second constraint is guaranteed, i.c., for initial conditions sampled at time 0 from Po. the state at time will be distributed according to pl. and the cost JsB in (2) will be minimized.
The increased practical applications of EOT and SBs in multiple machine learning problems, espe- cially in high-dimensional generative applications where the boundary distributions po: Pi are only available through a finite number of samples, have led to the development of a multitude of algo- rithms over recent years. The state-of-the-art methods for solving SBs leverage the properties of problem (2), such as the decomposition of the optimal probability flow into conditional problems that are easier to solve, sometimes even analytically (Chen et al., 2016; Lipman et al., 2023; Liu al., 2023). In this category of methods, recent technique known as Diffusion Schrödinger Bridge Matching (DSBM) (Shi et al., 2023; Peluchetti, 2023), or its deterministic counterpart, known as Flow Matching (FM) (Lipman et al., 2023) or Rectified Flow (RF) (Liu et al., 2023), leverages the decomposition of the optimal probability flow to a mixture of flows conditioned on their respective endpoints and retrieves an approximation of the optimal solution to (2) as a mixture of conditional policies that are easy to calculate. Theoretically, one needs to combine an infinite number of condi- tional flows to retrieve the true flow, due to the continuous support of the boundary distributions. To overcome this issue, a neural network is usually trained to approximate this infinite mixture.
While the DSBM and the various Flow Matching algorithms have proven very effective in high- dimensional problems, the efficient solution of SBs in simpler problems is hindered by the lack of dlosed-form expressions in all but very few special cases with Gaussian marginal distributions (Bunne et al., 2023a). To tackle this problem and avoid costly neural network training in smaller problems, recent methods such as Light-SB (LSB) (Korotin et al., 2024) and Light-SB Matching (LSBM) (Gushchin et al., 2024) have been proposed to obtain quick and efficient solutions to SBs within seconds, for problems with low-complexity boundary distributions, such as mixture models. These methods work by an efficient parametrization of the Schrödinger potentials, key component of the SB. Because this parametrization does not lead to closed-form expressions for the boundary distributions of the SB, the calculation of its parameters is carried out through optimization.
Inspired by the flow decomposition idea behind DSBM and FM methods, and motivated by the need to obtain light- and fast SB solvers for a wide class of SB problems, in this paper, we solve the problem of finding a policy that can efficiently steer the distribution of a dynamical system from Gaussian Mixture Model (GMM) to another one, using mixture of conditional policies that can each steer the individual components of the initial mixture to the components of the terminal mixture. This approach, which is tailored to GMMs, separates the problem of fitting the boundary distributions to the data and solving the SB, resulting in improved accuracy with regard to the marginal distribution fitting. More specifically, we claim the following main contributions:
1. We present a computationally efficient, training-free method to solve the Schrödinger Bridge and the multi-margin Momentum Schrödinger Bridge problems in the case where the boundary distributions are Gaussian Mixture Models.
2. In contrast to existing approaches, our method can handle both stochastic and determin- istic versions of the problem (2). Based on a control-theoretic formulation, our approach also naturally generalizes to dynamical systems with a general Linear Time- Varying (LTV) structure, with the control input and stochastic component having different dimensions than the state, which could be of interest in Mean Field Games (MFG), multi-agent control ap- plications (Ruthotto et al., 2020; Liu et al., 2022; Chen, 2024), and higher order distribution interpolation such as Wasserstein splines (Chen et al., 2018).
2

---

3. We demonstrate the substantial potential of our algorithm in low-dimensional prob- lems, moderate-dimennionall image-to-image translation tasks, and multi- murti-marginal diffu- sion learning problems. Specifically, we show that our approach outperforms state-of-the- art lightweight methods for solving the SB problem both in terms of training speed and accuracy of the learned boundary distributions, when these are available through samples (40% better FID scores in the image translation task and one order of magnitude better MMD scores in the multi-marginal diffusion-learning problems).
4. Finally, we extend our method to problems with continuous GMM marginal distributions, a wide class of distributions that can capture multiple useful distributions with heavy tails, and we use our approach construct upper bounds on the 2- Wasserstein distance and approximate the displacement interpolation between Student-t distributions.
2 Preliminaries
2.1 Diffusion Schrödinger Bridge Matching and Flow Matching
The composition of diffusion processes as mixtures processes conditioned on their endpoints was originally proposed by Peluchetti (2021) as a simulation-free algorithm for generative modeling applications. The concept was later tailored to solve the SBP the DSBM algorithm (Shi et al., 2023), proposed concurrently by Peluchetti (2023). Similar simulation-free methods have also been proposed to solve variants of the same problem Albergo & Ganden-Fijnden (2023) and Liu et al. (2024); Theodoropoulos et al. (2025) for the stochastic bridge setting, as well as in Liu et al. (2023) and Lipman et al. (2023) for the deterministic setting. For a more comprehensive overview along with comparisons with other available methods, we refer the reader to (Shi et al., 2023, Section 5) and (Peluchetti, 2023, Section 5).
Given problem (2), the main idea is to decompose the problem into sequence of elementary condi- tional subproblems that are easier to solve, and then express the solution as a mixture of the solutions of the conditional subproblems. This idea has intuitive motivation: Informally, finding policy that transports the state distribution from an initial density to a target density can be separated into two problems. First, one needs to figure out a transport plan solving the "who goes where" problem and then one needs to compute a point-to-point optimal policy, that solves the "how to get there" problem (Terpin et al., 2024a). In many cases, the two subproblems 2021, 2016; Terpin et al., 2024a); most importantly, however, computing the point-to-poin optimal policy can be solved analytically for simple dynamical systems, such the one in (2).
decoupled (Chen et al.,
More precisely, the optimal probability flow Pž of Problem (2) is known (Föllmer, 1988; Chen al., 2021) to admit the decomposition
/ Rdd
(x
W Welmo (x) dx (No,
(3)
where W Wal (x) is the probability density of the unforced dynamics dx, = Vedw, dww, namely the Brownian motion kernel, pinned 2o for t and at x1 for 1, and म. (xo, 21) is the entropic optimal transport plan between Pa. P1 solving (1). Dai Pra (1991) showed that W Welno (x) solves the following optimal control problem
min ( æ1) A E [ [ (2))) 1.jo,EE
{ [d: = 1 dt + Vedw, { zo 8za1 21
s.t.
(4)
8215
where 8z0: 6z, are Dirac delta functions centered on zo and 21, respectively. Assuming Pejo.1 (1)) and 1 (x) solve (4), one can construct feasible solution for the original problem (2) using any transport plan q 21) I[po, p1), i.e., any joint distribution between the desired boundaries Po, Pl. using the mixtures
(x(x)) = J Pro. (2) xn) dæo dri, Showing that the flow (5a) is feasible solution to (2) for any valid coupling g( amounts to verifying that the flow (x()) satisfies the boundary distributions po. Pi at times t = 0 and 1,
(5a)
u. ux(x) = f (x) Paoo, (2)
dag dz
(5b)
(x))
3

---

respectively, To prove that the policy (5b) produces (5a). it suffices to show that the pair (5a), (5b) satisfies the FPK PDE (Lipman al., 2023: Liu al., 2024). When q( 21) = म. (ag. 21 (5a) reduces to (3), and (5b) recovers the optimal solution to (2) (Shi et al., 2023; Peluchetti, 2023).
2.2 Schrödinger Bridges with Gaussian Marginals
The SBP with Gaussian Marginals, henceforth referred extensively studied in the literature and can be solved either analytically for simple choices of prior dynamics (Bunne et al., 2023a) or convex semidefintte optimization problem for general linear dynamical systems both for continuous and discrete time cases (Chen et al., 2015b; Liu et al., 2025). Because we the GSB a building block to construct policy that works with general GMM boundary distributions, we briefly review the available methods for its solution here. To this end, consider the optimization problem with Gaussian marginals
the Gaussian SB (GSB), has been
E % [x)(())) at
dæ (xrr dt + Vedw, 2i0 ~~ Eo), In ~ ((n, Ei),
min JGSB A uell
s.t.
(6)
where ya, Co, Mi, E1 are the means and covariances of the initial and final Gaussian boundary dis- tributions, respectively.
Proposition (Bunne et al., 2023a, Theorem 3) The optimal solution Problem (6) is given by 4z)z)= = K2 - M+) + with
Po + tpr. = -o. and K, = - --1 where
Σ, = - --))² E + (1 t){(C, C. +el),
S. t(E, CC.) (1 tEo C.) etl,
with C. = 1 D.E. +l) and D (42 E, Eij e² 1 Furthermore, the optimal value of the cost JGSB in (6) is given by the following proposition.
Proposition 2. Consider Problem (6) with 0. Then, the optimal value for the cost JGSB is
JGSB || 2 + tr (ro) + tr(E,) (tr - log det M2. + log det Σ1) + +.,
(8)
where M. + [[+ + 16 ) and c is constant independent of the boundary distributions.
In the limit o 0, (8) reduces to the well known Bures Wassereassrr distance (Bhatia al., 2019), defined by
BW(N Eo)l Ei)) A || poll tr(En) + tr(Ei) 2tr (2) Σ
2.3
Momentum Schrödinger Bridges with multiple Gaussian marginals
Other than Problem (6), we will also make use of the solution to the corresponding Gaussian multi- marginal Momentum SB (GMSB) (Chen et al., 2019), which is a variation of Problem (6) where the dynamics include a momentum term, while the goal is to match multiple marginal distributions regular time intervals. More specifically, the GMSB problem reads
E [%% [lue(aro we)llad )
min wEll
(10a)
s.t. dx, vs dt, du, uo(ee, 02) dt + Vedw, It, ~ N(u,, Er,), i= I,..,..,
(10b)
(10c)
where N is the number of marginal constraints and the joint space of position and velocity, namely Te, U. Rdd is referred to as the phase space. Compared the standard SB problem, and apart from having multiple marginal distributions, the GMSB problem only constrains the position component of the phase space, namely 2t. Even when the marginals are Gaussian, a closed-form solution to (10) is unknown; however, the problem can be solved efficiently using semidefinite programming, In the special case where the noise parameter e zero, problem is known as the Gaussian Wasserstein spline Problem (Chen et al., 201 9018), and an efficient semidefinite formulation for solving it is given Chen al. (2018). Since the semidefinite formulation for solving (10) is well-studied problem, due to space considerations, we defer it to Appendix D. Finally, we note that generalization of the

---

Problems (6) and (10) is achieved by a Linear Time Varying (LTV) structure in the prior dynamics of the bridge, ie., replacing the first constraint in (6) with
dx, - Acx, dt + Brut dt + Dr dw
(1I)
where 24 E R", A, Raxd U ER" B, Ráxm, D; E Rd We E RR. Bridges with prior dynamics of the form (11) have been extensively studied in the context of control theory, with the corresponding literature known as Covariance Steering (CS) (Chen et al., b015a,b; Bakolas, 2018; Liu et al., 2025). CS problems can be formulated convex programs for both continuous and discrete-time cases (Chen et al., 2015b; Liu et al., 2025), and therefore attain an efficient and exact calculation, which will exploit in the sequel.
3 Fast Diffusion for Mixture Models
3.1 Gaussian Mixture Schrödinger Bridge
Equation (5b) expresses the policy of Problem (2) as an infinite mixture of conditional, point-to-point policies. In this section, we extend this idea construct mixture policy, consisting of conditional policies each solving a Gaussian bridge sub-problem of the form (6). To this end, consider the problem
JGMM & E J (1))())))
min well
(12a)
s.t. dæ; (xx))) dt + Vedw,
(12b)
20 Em No aiN (mb 28),
Ex af (n),
(12c)
The main result summarized in the following theorem.
Theorem 1. Consider problem (12), with No components in the initial mixture and N₃ components in the terminal mixture. Assume that Utys the conditional policy that solves the (i.j)-GSt prob- lem, that is, the bridge from the i-th component of the initial mixture, to the j-th component of the terminal mixture and let the resulting probability flow be Pelgj- Furthermore, let )aj 2 0 such that, for all j E (1,2,...,N1}, Ei Aij = of and such that, for all i E ... No). E3 Aij ob.
Then, the policy
(a)cy()))ss ue(x)) Eis Melis (z) Ei, Pris)
(13)
a feasible policy for Problem (12), and the corresponding probability flow pr(x) Eis (())()))
(14)
The mixture policy (13) is weighted average of conditional policies, weighted according to (1)))))( while the denominator Σi,; Palas (x) (ij is just a normalizing constant. Since Pejay (x) is a Gaussian distribution centered at the mean the (i. Gaaussin bridge at time t, this weighting scheme prioritizes the conditional policies whose mean closer to the value of x at the time t.
Equation (13) provides set feasible of solutions to Problem (12), for all values of Asy satisfying the conditions of Theorem Obtaining the optimal policy within this set challenging, in general. Alternatively, we can formulate tractable problem by minimizing an upper bound to the original minimum effort cost function (12a), which is linear with respect to the transport plan Aiy. More formally, we have the following theorem.
Theorem 2. Let Jij be the optimal cost of solving the (i, (i.J)-Gaussian Gaussian bridge subproblem of the form (6) with marginal distributions the i-th component of the initial and the j-th component of the terminal mixture. Then, the cost function of the linear optimization problem
min Ag 20
Jor EEs Jes
(15a)
s.t.
Σ,>== (1,2.... No), and =f (1,,,, ...,,....AAN. ,N), (15b)
provides upper bound for (12a), ie., JaMM S JOT, for all positive values of >ig satisfying (15b).
5

---

For clarity of exposition, we defer the proofs of Theorems 1, 2 to Appendix A, along with an optimality analysis of the upper bound of Theorem 2.
In practice, to use policy (13) in problems where the boundary distributions are available only through samples, GMMs are first fitted in the samples of po. P1 using the Expectation Maximization (EM) algorithm (Bishop & Nasrabadi, 2006), and then (13) is calculated using Theorems 1, 2. For clarity of exposition, we present an overview this approach in Algorithms 1, 2 for both training and inference, accompanied by the corresponding theoretical complexity analysis in Appendix A.6.
Algorithm GMMflow training
Input: Samples from boundary distributions Pu. P1; number of GMM components No and N1,
noise level
{ais Mg: 20 No
i=1 < (po,, No) fits a GMM to initial dataset
(4 21
115=1 (M, N1) // fits GMM to final dataset
Ni
for (4,j) E {I,.... No} x {l..... N1} do /compptee in parallel
{ueky Pelis, Jij) CS Eg mú, Mi) //soves the (i,j)-th conditional GSB
end for
A1 SOLVE (15) USING (Jej. ag.
return Ugji Pejeys Aij
Algorithm GMMflow inference
Input: Component-level solutions Ptlej. transport plan Xij, Initial condition xo Po. SDE integrator sde_int().
(z) + Compute (13)
Te + sde. sde_inte((((_eeeeeeeeeeeeeeedd_s int((innnnnnnnn..... To, t [0,
return Ø1
3.2 Multi-Marginal Problems
In this section, we generalize the results of Section 3.1 to solve the martinlargiinal momentum SBs problem (Chen et al., 2023, 2019) with GMM marginal distributions. That is, we solve
V. [lu(( v))l dt
min
JCMM A E
(16a)
well
s.t.
dx; dt, du melr,, v2) dt + Vedue, F1 ~ - NW E 1=
(16b)
(16c)
where the i-th marginal mixture assumed to have N; components. Similarly to Theorem 1, we will combine conditional GMSBs of the form (10) to build feasible set of policies solving (16).
To facilitate notation, we denote by i (ii,..., ias) the index of the conditional multi-marginal GMSB and use the notation (ili, = k3 denote the set of all values of i such that iz = k. With this notation in mind, we provide the following generalizations of Theorems and 2:
Theorem 3. Consider problem (16), with M marginal mixture distributions, each having N; Gaus- sian components, where i = 1.... M. Let i = (f1...id be M-dimensional index, Weil be the conditional policy that solves the i-GMSB problem, that is, the Gaussian multi-marginal momentum Schrödinger Bridge going through the (i....., iM) components of the marginal mixture models, and let the resulting probability flow be Pejz. Furthermore, let >j 2 be such that, for all j = 1,.... and for all k -,...., Ny.
Emi,= ò === a5.
(17)
Then, the policy
us ux(x, v) E, Wes(s, v) En Pai(
(18)
a feasible policy for Problem (12), and the corresponding probability flow is
pc(x,v) = PiP(,I,, v\
(19)
To approximate the optimal multi-n multi-marginal transport plan A₁, we use the following upper bound.

---

Theorem 4. Let J1 be the optimal cost of solving the i-GMSB subproblem of the form (10) with marginal distributions (i1....,MM) components of the marginal mixtures. Then, the cost finction of the linear optimization problem
min 2 0
Jor - Ei AJd1
(20a)
s.L.
Eiti,= \ = o 1....... k k,k=l,...,...
(20b)
provides an upper bound for (16a), that is, JCMM I JOT- for all values of \ satisfying (20b).
In practice, in order to use policy (18) for inference, the initial conditions for the SDE (16b) must be defined. Specifically, to fully define the initial conditions in the state space [zo, unl. given an initial position sample æo, the corresponding velocity
must be estimated. To this end, note that given X the state distribution Pe(are v2) is a fully defined GMM given by Equation (19). that the conditional distribution pe (vel 22x) is also GMM, whose parameters can be easily computed. Since this calculation is trivial, we defer it to Appendix C.3. For completeness, we provide the multi- marginal inference algorithm in Algorithm 3. The multi-marginal training algorithm is omitted due to its similarity to (2).
easy to show
Algorithm 3 Multi-marginal GMMflow inference
Input: Component- solutions Ugjs Ps multi-marginal transport plan A;, sample from To
(o())) SDE integrator sde_int().
ue (x) <- Compute (18)
Sample Po(Po) aro using Equation (B.19d) [mgg v/] <- sde_int (Equuation (16b), (To; vo]. e [0, 1))
return Z1z
3.3
Continuous Gaussian Mixtures
The results of Section 3.1 can be extended to problems with continuous GMM boundary distribu- tions. Specifically, we consider a bridge of the form
E =[/ [1u
min well
(21a)
s.t.
dx. ue( dt,
(21b)
((t))) mm
i= 0,
(21c)
where the boundary distributions (21c) continuous GMMs with mixing measures P(wo), P(w)) respectively. We keep the dynamics (21b) deterministic to simplify the analysis, although all the results carry over to the general case of stochastic dynamics well. This more general formulation specializes to problem (12) when P(uo), P(w1) have discrete support; however, it includes many other scenarios where the mixing distributions are continuous. Specifically, the generalization of Theorem is as follows.
Theorem 5. Consider Problem (21) and assume that Ue is the conditional policy that solves the (wo, wi w1)-GSB problem, that is the bridge from the initial Gaussian distribution with parameter wo the terminal Gaussian distribution with parameter and let the resulting probability flow be Pt wo.wr Furthermore, let A (wo, wi) be any coupling such that its marginal distributions are Po and Pi respectively. Then, the policy
Petero ue (x) dA(wo, wr)
(x)
Jgm xcR
We/woow. (xr)
(22)
I XBe Pt] worws (x) dA(wo: W1)
a feasible policy for Problem (21) and the corresponding probability flow is
J
Pi (x)
Pelwo,ws wi (x) dAA dAAwo, w1).
(23)
xRm
Furthermore, Theorem 2 generalizes the following.
7

---

Theorem 6. Let (wo, w1) be the optimal cost of solving the (wo, (ooo,ii)Gaaaale Gaussian bridge subproblem. Then, the optimal transport problem:
JOT A min AETI(Po,P)) Pi) R XR"
Jww. wl) dA(wo. w1),
(24)
provides an upper bound for Problem (21), that is, JGMM S JOn, where 11(Po, P1) represents the set of all couplings with marginals Po and Pt.
Problem (24) is challenging to solve in general. However, in many practical cases, such as for Student-t boundary distributions, the parameter spaces for wo. wi are one dinemmmmniionll and (24) can be solved in closed form. We explore this interesting direction in Appendix B to approximate the Wasserstein-2 distance and the displacement interpolation between heavy-tail distributions.
4 Related Work
Although the idea of creating a mixture policy from elementary point-to-point policies at the heart of flow-matching, to the best of our knowledge, its benefits for solving problems with mixture models with a finite number of components have not been explored. In an early work, Chen et al. (2016) developed an upper bound on the 2- Wasserstein distance between Gaussian mixture models, i.e,, the static, deterministic version of Problem (12), that matches the upper bound of Theorem 2. However, the problem of finding a policy that solves the dynamic problem was not explored. Fo- cusing on works concerning dynamic problems, the concept of constructing stochastic differential equations (SDEs) mixtures of Gaussian probability flows can be traced back to mathematical fi- nance applications (Brigo, 2002; Brigo et al., 2002). More recently, and in the context of generative applications, Albergo & Vanden-Eijnden (2023) developed similar expressions for finding a policy for steering between mixture models using an alternative conditional solution for the Gaussian-te -to- Gaussian bridge, focused on the case of deterministic, fully observable dynamical systems with no prior dynamics, and the component-level transport plan was not optimized.
Our work reveals some similarities in scope and structure with LSB (Korotin et al., 2024) and LSBM (Gushchin et al., 2024). Similarly to LSB and LSBM, we aim to provide numerically inexpensive tools to solve the SBP in low-to-moderate dimensional scenarios. Moreover, our feedback policy in Equation (13) is a mixture of affine feedback terms weighted with exponential kernels; this is also the case in LSB and LSBM (Korotin al., 2024, Proposition 3.3). The formulations, however, are otherwise quite distinct. Specifically, LSB/LSBM works by modeling the so-called Schrödinger potentials using GMM. This results in flows with boundary distributions that have mixture- like structure and are optimal by construction, but whose marginals are neither amenable to exact calculation, nor are GMMs, in general. Finally, to approximate the optimal flows for boundary distributions available only through samples, both works solve non-linear optimization problems, which are prone to converging to locally optimal solutions. In contrast, our approach works by first pre-fitting GMMs to the data using the Expectation Maximization (EM) (Pedregosa et al., 201 1), and then computing an optimal policy by solving exactly a linear program. Finally, being based on a control-theoretic framework, our method generalizes effectively to partially observable and multi- marginal distribution matching problems as shown in Section 3.2, while extending LSB and LSBM to handle such problems is non-trivial.
5 Experiments
2D Problems and Benchmarks. We first test the algorithm in various 2D "toy" problems as shown, for example, in Figure for a Gaussian-to-Gaune levels. To assess optimality, we evaluate the resulting transport cost for policy (13) for each noise level and compare with the upper bound from (15a). We also run a series of EOT benchmarks and compare them with state-of-the-art neural approaches such the DSB (De Bortoli et al., 2021) and DSBM (Shi et al., 2023) algorithms. Due to space considerations, we defer their discussion to Appendix C, along with further experiments studying training and inference time scaling with spect to the problem dimension and the number of GMM components. To evaluate the performance on problems with many GMM components, we tested the algorithm on the distributions depicted in Figure 2, where we first pre-fit 500-component GMMs in the initial and terminal samples.
Mixture problem for various noise
Image-to-Image Translation. Following Korotin et (2024), we use our algorithm in the latent space of an autoencoder to perform a nam-tn and adult-to-child image translation task. We
8

---

-3
-10
-10
-3
10
-10
-10
10
-2
-3
Figure 1: Gaussian to 8- 8-Gaussians with no prior dynamics and for various levels of noise.
Gr
DCSL
Figure 2: GT to DCSL Schrödinger Bridge with zero prior dynamics and 0.1.
Original
LightSB
GMMflow
Original
LightSB
GMMflow
Figure 3: Adult Child (left) and Man Woman (right) image translation task.
the pre-trained ALAE autoencoder (Pidhorskyi al., 2020), trained the FFHQ dataset (Kar- ras et al., 2021). The latent space of the autoencoder 512-dimensional. We start by fitting a 10-component mixture model to the embeddings of each image class, with diagonal covariance ma- trices to facilitate matrix inversions in the Gaussian-tur Ceussian-to-Gaausiin policy calculations summarized in Proposition 1, and then apply Algorithms 1, for 0.01. The results are illustrated in Figure 3.
To test how well the generated images match the features of the given target distribution, we calculate the Fréchet inception distance (FID) scores (Heusel al., 2017) between the actual and the generated images of a given class, using 10,000 samples from each distribution. The FID scores correspond Wasserseenn distance between the images of the classes, evaluated the latent space of the Inception network. To further test how close the transformed images are to the target class, we also calculate the empirical Bures- Wasserstein distance between the transformed images and the real images of the target class, directly in the latent space of the ALAE autoencoder and report it as ALAE-B B.W.. We compare against two sfate-of sfate-of-the-ae ---thhe-tt lightweight SB solvers, namely, LSB and LSBM, with results shown in Tables and 2.
the empirical
Table 1: Man-to- Man-to-Woman Woman FID comparison
Table 2: Adult-to-Child FID comparison
M-W
FID
ALAE-BW
T. Cost
A-C
FID
ALAE-BW
T. Cost
LSB
4.94
28.9
8.23
LSB
6.62
31.00
8.18
LSBM
4.98
28.3
8.18
LSBM
6.61
30.99
8.19
GMMflow
3.04
9.3
9.05
GMMflow
3.50
8.54
9.33
Qualitatively, our algorithm performs more aggressive feature changes compared to the baseline method, illustrated in Figure 3. Quantitatively, the features of the transformed images better

---

capture the true distribution of the features of given target class, given the almost 40% better FID scores and 65% better ALAE-BW scores provided in Tables and We note that the improvement in the FID and ALAE-BW scores comes with slight increase in the average transport cost, which we measure by EE, Io - ll. as reported in the last columns of Tables and 2. We attribute this significant performance gain to our use of the EM algorithm for the GMM pre-fitting, which is less prone to converge locally optimal values, compared to LightSB's s maximum likelihood objective, or the LSBM's bridge matching objective. We further remark that our approach takes 63% less time to train, noted in Table 5 in Appendix C.2.
Multi- Problems. A key challenge in SBs is learning a system's 's underlying diffusion process, given samples from partial observations of the distribution of its state, measured at regular time intervals (Chen et al., 2023). This challenge arises, for example, in learning the dynamics of large cell populations throughout their different developmental stages (Bunne et al., 2022; Tong et al., 2020; Terpin et al., 2024b). To showcase the effectiveness of our approach in such problems, we consider the scRNA-seq dataset from (Moon et al., 2019), with the pre-proe pre-processing detailed in Tong et al. (2020). The dataset contains samples from the first 100 Principal Components (PC) of individual cell proteins, grouped at 5 regular time intervals, denoted by 11....,45
For our setup, we keep the first 5 PCs from the marginal distributions, and prefit 5-component GMMs in each marginal. We use the second-order model (16b) capture the prior dynamics and the structure of the system; however, we note that any LTV model with structure of the form (11) would be applicable. To solve the resulting multi-marginal momentum SBs, we first compute the cost of each GMSB and then solve (20). Computing the GMSB cost for all the combinations of components the most computationally expensive part of our approach. By parallelizing this computation, the total training time is approximately 8 minutes on an 17-1 CPU.
We visualize the data generated by our method in Figure 4. and provide standard performance met- rics in Table 3. Specifically, following Chen et al. (2023), we use the Sliced Wasserstein Dis- tance (SWD) and Maximum Mean Discrepancy (MMD) metrics averaged over the 4 predicted time marginals of the dataset. Although there are no light-weight solvers for second-order multi-marginal problems, we compare our method against DMSB (Chen et al., 2023), NSBL (Koshizuka & Sato, 2023), and MIOFlow (Tong et al., 2020). Even though our method requires minutes to compute on CPU while neural methods take an hour to train on high-end GPU, the proposed approach outperforms the baselines by one order of magnitude in the MMD metric due to the accurate fitting of the GMMs. The SWD metric, although acceptable, is bounded by the expressivity of the GMMs to capture the higher-order distribution structure. The metrics for MIOFlow, NSBL, and DMSB in Table 3 are taken from (Chen et al., 2023).
Real Data Generated Data
Real Data
Generated Data
Table 3: SWD and MMD indices averaged over the 4
time steps
Method
SWD
MMD
PC
MIOFlow
0.38
0.28
NSBL
0.24
0.10
DMSB
0.22
0.06
1
PC
PC3
3
GMMflow
0.37 14E-3
0,0038 47E-4
Figure 4: GMM momentum multimagrinal SB: Real vs Generated data.
6 Conclusion and Limitations
This paper introduces novel, efficient method for solving SB problems with GMM boundary distri- butions by utilizing mixture of conditional policies, each solving a Gaussian bridge subproblem. In the same way that low-dimensional methods such as Gaussian distributions and mixture models are highly used in statistics and machine learning, we believe that our approach will be a valuable tool in many useful practical problems related to optimal transport, distribution interpolation, distributional control, and related applications, given its very low computational complexity. and excellent empir- ical performance in complicated statistical problems. The main limitation of our work stems from its dependence on GMM marginal distributions. Since this class of distributions not designed for use in very high-dimensional problems, we do not expect our method to be applicable in such areas, but rather to work as an efficient tool for obtaining rapid SB solutions to smaller problems.
10

---

Acknowledgments
The authors would like to thank Peter Garud for his help with the coding of the multimarginal SB implementation. Support for this work has been provided by ONR award N00014-18 1144-18---2228 and NASA ULI award #80NSSC20M0163. This article solely reflects the opinions and conclusions of its authors and not of any NASA entity. George Rapakoulias acknowledges financial support from the A. Onassis Foundation Scholarship. Lingjiong Zhu is partially supported by NSF grants DMS-2053454 and DMS-2208303.
References
Michael Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In International Conference on Learning Representations, Kigali Rwanda, May 2023.
David F. Andrews and Colin L. Mallows. Scale mixtures of normal distributions. Journal of the Royal Statistical Society, Series B (Methodological), 36(1):99-102, 1974.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, volume 70, pp. 214-223, Sydney. Australia, 2017. PMLR.
Efstathios Bakolas.
Finite-horizon covariance control for discrete-time stochastic linear systems subject to input constraints. Automatica, 91:61-68, 2018.
Jean-David Benamou and Yann Brenier. computational fluid mechanics solution to the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-38 2000.
Alain Bensoussan, KCJ Sung, Sheung Chi Phillip Yam, and Siu-Pang Yung. Linear-quadratic mean field games. Journal of Optimization Theory and Applications, 169:496-529, 2016.
Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the Burcs- Wasserstein distance between positive definite matrices. Expositiones Mathematicae, 37(2):165-191 2019.
Christopher M Bishop and Nasser M Nasrabadi. Pattern Recognition and Machine Learning, vol- ume 4. Springer, 2006.
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, Cam- bridge, UK, 1st edition, 2004.
Damiano Brigo. The general mixture diffusion SDE and its relationship with uncertain-volatiliy option model with volatility decorrelation. Available at SSRN 455060, 2002.
Damiano Brigo, Fabio Mercurio, and Giulio Sartorelli. Lognormal-mixtue dynamics under differ- ent means, 2002.
Charlotte Bunne and Gunnar Rätsch. Neural optimal transport predicts perturbation responses at the single-cell level. Nature Methods, 20:1639-1640, 2023.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal trans- port modeling of population dynamics. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. 6511-6528. PMLR, 28-30 Mar 2022.
Charlotte Bunne, Ya-Ping Hsich, Marco Cuturi, and Andreas Krause. The Schrödinger bridge be- tween Gaussian measures has a closed form. In International Conference on Artificial Intelligence and Statistics, volume 206, pp. 5802-5833, PMLR, 2023a.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia Del Castillo, Mitch Levesque, Kjong- Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single- cell perturbation responses using neural optimal transport. Nature Methods, 1799 2023b.
Tianrong Chen, Guan-horng Liu, Molei Tao, and Evangelos A Theodorou. Deep multi-marginal momentum Schrödinger bridge. In Proceedings of the 37th International Conference on Neural Information Processing Systems, volume 36, pp. 57058-57086 Curran Associates, Inc., Decem- ber 2023.
11

---

Yongxin Chen. Density control of interacting agent systems. Transactions on Automatic Control, 69(1):246-260, 2024.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Optimal steering of linear stochastic system final probability distribution, part Transactions on Automatic Control, 61(5):1158- 1169, 2015a.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Optimal steering of linear stochastic system to a final probability distribution, part II. Transactions on Automatic Control, 61(5):1170- 1180, 2015b.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Optimal transport over linear dynamical system. Transactions on Automatic Control, 62(5):2177-2152, 2016.
Yongxin Chen, Giovanni Conforti, and Tryphon T. Georgiou. Measure-valued spline curves: An optimal transport viewpoint. SIAM Journal on Mathematical Analysis, 5046)55947,5968, January 2018.
Yongxin Chen, Giovanni Conforti, Tryphon Georgiou, and Luigia Ripani. Multi-marginal Schrödinger bridges. In International Conference on Geometric Science of Information, pp. 725- 732, Toulouse, France, August 2019. Springer.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard Sinkhorn meets Gaspard Monge on a Schrödinger bridge. SIAM Review, 63(():249-313 2021.
Paula Cordero-Encinar, O. Deniz Akyildiz, and Andrew B. Duncan. Non-asymptotic analy- sis of diffusion annealed Langevin Monte Carlo for generative modelling. arXiv preprint arXiv:2502.09336, 2025.
Paolo Dai Pra. A stochastic control approach reciprocal diffusion processes. Applied Mathematics 1991.
and Optimization,
Valentin De Bortoli, James Thortton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems, volume 34, pp. 17695-17709, held virtually, 2021. Curran Associates, Inc.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanis- Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy. levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021.
Hans Föllmer. Random fields and diffusion processes. Lect. Notes Math, 1362:101-200. 1988.
Nikita Gushchin, Alexander Kolesov, Petr Mokrov, Polina Karpikova, Andrei Spiridonov, Evgeny Burnaev, and Alexander Korotin. Building the bridge of Schrödinger: A continuous entropic optimal transport benchmark. Advances in Neural Information Processing Systems, volume 36, pp. 18932-18963, New Orleans, USA, 2023. Curran Associates, Inc.
Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal Schrødinger bridge matching. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 17100-17122, PMLR, 21-27 Jul 2024.
Dan Hendrycks and Kevin Gimpel. arXiv: 16Gv:1606.000 08006.08415,,0 15, 2016.
Gaussian error linear units (GELUs). arXiv preprint
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local Nash equilibrium. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
12

---

Tero Karras, Samuli Laine, and Timo Aila. A style-batd base- generator architecture for generative adversarial networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 43(12); 4217-4228, 2021.
Juno Kim, Jachyuk Kwon, Mincheol Cho, Hyunjong Lee, and Joong-Ho Won. -vvariational au- toencoder: Learning heavy-tailed data with student's i and power divergence. In International Conference on Learning Representations, Vienna, Austria, May 2024.
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light Schrödinger bridge. In The Twelfth International Conference on Learning Representations, Vienna, Austria, May 2024.
Takeshi Koshizuka and Issei Sato. Neural Lagrangian Schrödinger bridge: Diffusion modeling for population dynamies. In The Eleventh International Conference on Learning Representations, Kigali, Rwanda, May 2023.
Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems-A, 34(4):1533-1574, 2014.
Yaron Lipman, Ricky Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In International Conference on Learning Representations, Kigali Rwanda, May 2023.
Fengjiao Liu and Panagiotis Tsiotras. Reachability and controllability analysis of the state covari- ance for linear stochastic systems. arXiv preprint arXiv:2406. 14740, 2024.
Fengjiao Liu, George Rapakoulias, and Panagiotis Tsiotras. Optimal covariance steering for
discrete-time linear stochastic systems. Transactions on Automatic Control, 70(4):2289-2300 2025.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized Schrödinger bridge. In Advances in Neural Information Processing Systems, volume 35, pp. 9374-9388, Louisiana, LA, 2022. Curran Associates, Inc.
Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A Theodorou, and Ricky TQ Chen. Generalized Schrödinger bridge matching. In International Conference on Learning Representations, Vienna, Austria, 2024.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations, Kigali Rwanda, May 2023.
Anton Mallasto, Augusto Gerolin, and Hà Quang Minh. Entropy--regularized 2- Wassersteii distance between Gaussian measures. Information Geometry, 5(1):289-323, 2022.
Kevin R Moon, David Van Dijk, Zheng Wang, Scott Gigante. Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, Natalia B. Ivanova, Guy Wolf, and Smita Krishnaswamy. Visualizing structure and transitions in high-dimensional biological data. Nature Biotechnology, 37(12))1482-1144 2019.
ApS Mosek. Mosek modeling cookbook, 2020.
Kushagra Pandey, Jaideep Pathak, Yilun Xu, Stephan Mandt, Michael Pritchard, Arash Vahdat, and Morteza Mardani. Heavy-tailed diffusion models. In International Conference on Learning Representations, Singapore. 2025.
Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, Jake Vanderplas, Alexandre Passos, David Cournapeau, Matthieu Brucher, Matthieu Perrot, and Edouard Duch- esnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12(85): 2825-2834 1.
Stefano Peluchetti. Non-denoising forward-time diffusions. arXiv preprint arXiv: 2312. / /4589, 2021.
13

---

Stefano Peluchetti. Diffusion bridge mixture transports. Schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023.
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport: With Applications Data Science. Foundations and Trends in Machine Learning, 2019.
Stanislav Pidhorskyi, Donald A Adjeroh, and Gianfranco Doretto. Adversarial latent autocncoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 14104-14113, Seattle, Washington, June 2020.
George Rapakoulias and Panagiotis Tsiotras. Discrete-time optimal covariance steering via semidef- inite programming. In 62nd Conference on Decision and Control, pp. 1802-1 187, Singapore, 2023.
George Rapakoulias and Panagiotis Tsiotras. Discrete- maximum likelihood neural distribution steering. In IEEE 63rd Conference on Decision and Control (CDC), pp. 2195-22090000 MiCo, Milan, Italy, 2024.
Lars Ruthotto and Eldad Haber. An introduction to deep generative modeling. GAMM-Mittellnngen 44(2):2021100088, 2021.
Lars Ruthotto, Stanley Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung. A machine learning framework for solving sighedimensional mean field game and mean field control prob- lems. Proceedings of the National Academy of Sciences, 117(17)9183-9199, 2020.
Filippo Santambrogio. Optimal Transport for Applied Mathematicians, volume 87. Springer, 2015.
Augustinos D Saravanos, Yihui Li, and Evangelos A Theodorou. Distributed hierarchical distribu- tion control for very-large-scale clustered multi-agent systems. In Robotics: Science and Systems XIX, Daegu, Republic of Korea, July 2023.
Simo Särkkä and Amno Solin. Applied Stochastic Differential Equations, volume 10. Cambridge University Press, 2019.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. In Advances in Neural Information Processing Systems, volume 36, pp. 62183- 62223. Curran Associates, Inc., 2023.
Antonio Terpin, Nicolas Lanzetti, and Florian Dörfler. Dynamic programming in probability spaces via optimal transport. SIAM Journal on Control and Optimization, 1183-1168,, 2024a.
Antonio Terpin, Nicolas Lanzetti, Martín Gadea, and Florian Dorfler. Learning diffusion at light- speed. In Advances in Neural Information Processing Systems, volume 37, pp. 6797-6832. Curran Associates, Inc., 2024b.
Panagiotis Theodoropoulos, Nikolaos Komianos, Vincent Pacelli, Guan-Horng Liu, and Evangelos Theodorou. Feedback Schrödinger bridge matching. In The Thirteenth International Conference on Learning Representations, Singapore, 2025.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajecto- ryNet: A dynamic optimal transport network for modeling cellular dynamics. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9526-9536. PMLR, 13-18 Jul 2020.
Pauli Virtanen, Ralf Gommers, Travis E Oliphant, Matt Haberland, Tyler Reddy, David Coumapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan van der Walt, Matthew Brett, Joshua Wilson, K. Janod Millman, Nikolay Mayorov, Andrew R.J. Nelson, Eric Jones, Robert Kern, Eric Larson, CJ Carey, Ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E.A. Quintero, Charles R Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: fundamental algorithms for scientific computing in Python. Nature Methods, 17(3):261-2) 2722 2020.
Richard Lee Wheeden and Antoni Zygmund. Measure and Integral, volume 26. Dekker New York, 1977.
14

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract introduction accurately reflect paper's contributions and scope?
Question:
Answer: [Yes]
Justification: An accurate overview of the method, along with all the main claims about of the proposed approach are clearly stated in the abstract.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The main limitations of the paper stemming from the applicability Gaus- sian mixtures to high-dimensional problems have been well discussed.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specificatiin, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might be used reliably provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
15

---

judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification:
All proofs are provided in the Appendices.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide extensive details about all the experiments we conducted in the Appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on contribution, reproducibility can accomplished various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data also be provided via
often one good way to accomplish this, but reproducibility detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
16

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or a way re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide link public implementation of our method in the abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips.cm cc/
for more details.
While we encourage the release code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https: hat.ps.cc/ubbll. http:/uullcccgggl wtticc//ggu///eeeeeeeeeeeeeeeeee /oaeee/ee/ouebbdmmccccc ssi onPolicy) for details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
We provide enough details in the experiments section of our paper, as well as in the Appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
17

---

The full details can be provided either with the code, in the appendix, or supple- mental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide confidence intervals for the SWD and MMD indices in Table 3. We argue that the FID scores in Tables and do not need a confidence interval because they are computed with a very large number of samples (10,000).
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK report I-sigma error bars, but one should state The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative етror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
We provide details about the computational resources used conduct the experiments.
Justification:
Guidelines:
The answer means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
18

---

Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http:////......c/eiiccceeaaeeaea
Question:
Answer: [Yes]
Justification:
We have read and comply with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there is a special consid- eration due to laws or regulations in their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impact of the method the end of our paper.
Guidelines:
The answer NA means that there societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not include pretrained Large Language Models or large datasets.
Guidelines:
The answer NA means that the paper poses such risks.
19

---

Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We properly cite any third-party tool used our paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the pack- age should be provided. For popular datasets, paperswi whrowwitthhheeeee com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification:
We provide public implementation of our algorithm in the abstract.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
20

---

Question: For crowdsoureing experiments and research with human subjects, does the pa- per include the full text of instructions given participants and screenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification:
We do not have any crowdsourcing experiments or experiments including
humans.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not have any experiments including humans.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
We do not use LLMs any part of our research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non non non-standard components.
Please refer to our LLM policy (https: //tp://neeuiippee ccc/onnferences/22255 for what should should not be described.
21

---

Go With the Flow: Fast Diffusion for Gaussian Mixture Models
Supplementary Material
The supplementary document organized as follows.
In Appendix A, we provide the technical proofs of the results in the main paper, an opti- mality amalysis of Theorem 2, and the complexity analysis of Algorithm 1.
In Appendix B, we discuss continuous mixtures and applications heavy-tailed distribu- tions.
In Appendix C, we provide additional numerical experiments and some details for the periments of Section 5 in the main paper.
In Appendix D, we provide technical background for the Gaussian Schrödinger Bridges with LTV prior dynamics.
A Proofs
All proofs are carried out for a stochastic Linear Time- Varying (LTV) system
dæ; Acx dt + B,u(xz) dt De dw.
(A.1)
where 2, E IRd A₁ E Rddd U R", B, Rdxm Dj gdxg and dw is the q-dimensional
Brownian increment having the properties E[dw] = E(dw dt] = 0 and Ejdw dw') Idt. The
dynamical system (12b) just special case of (A.1) with At 0, B, I, Dr vel, while the second order model (16b) can be captured by
A = 0 1
0 I
Dr = ve 0 1
B
=
(A.2)
A.1 The Beakker-Planc--ollmoee Equation
The equation describing the propagation of the distribution of the state of the dynamical system (A.1), known as the Fokker-Planck Kolmmckooo (FPK) equation (Särkkä Solin, 2019) is:
Ppt 0t
+ Σ 0x,
a
(pe(Aro) + Bru((T)))
82
((D,D))p O.
(A.3)
dz,0x,
where, for simplicity, we write P: p(t, This equation can be written more compactly using standard vector notation as follows
Op. vV (pr (o((+BB((())))) 1tr (D.D.V*p.) = 0, Øt
(A.4)
where p2pt denotes the Hessian of the density with respect to the state x at time t. the specific case where D; ve 1, equation (A.4) reduces to the well-known equation
Op.
+ V. (pe(Ace + Baua(1))) App = 0,
(A.5)
dt
where A denotes the Laplacian operator.
A.2 Proof of Proposition? 2
Let Po ~ (Ng, Ea), PI ((,,, Σn) be two multivariate Gaussian measures in Rd, and consider the entropy regularized 2- Wasserstein distance problem between Pu and Pl. defined by
/ || dz v) + <DKL(T @ p1),
(s(Poll)) & min TEll(po.P).
(A.6)
where DKT. denotes the KL-divergence operator.
In (Mallasto et al., 2022, Theorem 2), it is shown that (A.6) admits the following closed-form solu- tion
W5 (5(Polle1)
|l - = 2 + tr(E)) tr(E,) (tr -MMM--log log det M, dllog 2 -2d), (A.7)
22

---

where M. = + (I (1/e))
Returning to Problem (6), let P, W be the path measures corresponding to the controlled SDE of (6) (expressed the first constraint) and the Brownian motion with covariance c/ respectively, with initial conditions sampled from Po. Let D(po. P1) denote the set of all path measures with marginals Po: P1 times 0 and Then, the SB problem (6) admits the following equivalent representations:
{
| vll
d dr (z( H(x) + (()) log (2Te)
inf DKL. DKL. (PIIW)) inf (EP(pa.)t)
2e
(A.8a)
inf PED(Po.PL)
JGSB
(A.8b)
2x
where (A.8a) is due to the disintegration of the path measures, and (A.8b) is due to Girsanov's Theorem. We refer the reader Chen et al. (2021) for detailed derivation. Solving (A.8b) for JGSB, we obtain
JOSB inf
{ { -
( (TH 20
(A.9)
up constant independent of the parameters of pa, Using the closed form solution for (A.7),
and noting that
DKL (TllPo & p1) = (m)) + (Po) + H(p)), and that the differential entropy of the multivariate Gaussian distribution
(A.10)
1
H(p)) log det Σ₁ + log 2ne.
(A.11)
we conclude that the optimal cost of Problem (6) is equal to JCSB 2 (Po (I))) (p1)
|| μ 2| tr(Eo) tr(Ei) (trM2. log det M2. log det En),
(A.12a)
up numerical constant independent of po, P1. which concludes the proof.
A.3 Proof of Theorem
First, notice that the probability flow (14) satisfies the constraints (12c) for all feasible values of Xij since
Po
Σ PolssAs =
Σ Nib Eb)d.s E Nuá: Ej)ag
(A.13a)
4
Σ 6d
PilsAAi = EN( Ei = = Nud. N(d
(A.13b)
Therefore, it suffices to show that the policy (13) produces the probability flow (14). Following the approach of Lipman et al. (2023) and Liu et al. (2024), show that the pair (13), (14) satisfies the FPK equation. We start from the FPK equation describing a conditional flow and sum over all conditional variables to retrieve the unconditional flow. Specifically, given that the individual
policies Ueji solve the Gaussian Bridge subproblems (6), the pair (Prjajs Mejis) satisfies the FPK equation for the dynamical system (A.1), that is,
Ops ij
(peiss (Air B,un))) (D.D O. (A.14)
at
Multiplying equation (A.14) by Xij and summing over i.j, we obtain
Opnps
Σ Xrj ij
V
(Pets (Apr + Brusis)) (D.D.V"(A1s))) =0,
(A.15)
23

---

which implies that
(Axz
Σ
Pansdiss
+V.
Acx
Σ
Palisdes + Be Σ
))
D.D.
Σ
Puissius)
=0.
(A.16)
2
This can be further simplified as
Pt Aex + Be E Velis Eis Pelesdes) ))
0p: +
Pelestes
tr (p)) (A.17)
Ōt
which yields that
Op:
+ (pi(Arot Brus)) 2
tr (D.D)V²(a)) = O.
(A.18)
0
This completes the proof.
A.4 Proof of Theorem
It suffices to show that the (12a) is upper bounded by the cost of (15a). Substituting policy (13)
to the cost (12a) we obtain
JGMM =Eeerrpe
Σ
(x))
(xz)
dt
(A.19a)
Eis Peles (1re)(
= =fo' S|moo mm())
Σ
dx dl
(A.19b)
Eis Peiks (2:) (a)x)(())))
S/¹ / pelax)
Eis Jueke (x)]] PetK)(E))di
dx dt
(A.19c)
Ens Pelay (z)))
L'' / E1l [|)b)s(>>)|l (2))(()) dr dt
(A.19d)
Aij Exrrpiss Imnnt(l))) (2))) dt
(A.19e)
Σ
AgJJ Jor,
(A.19f)
where (A.19b) is due to Fubini's theorem (Wheeden & Zygmund, 1977, Theorem 6.1) and (A.19c) makes use of the discrete version of Jensen's 'ss inequality (Wheeden & Zygmund, 1977, Theo- rem 7.35).
A.5 Optimality of the Upper Bound of Theorem 2.
To assess the optimality the upper bound introduced Theorem 2, we study the gap between Jor and JGMM in the following theorem.
Theorem 7, In the setting of Theorem 2, let (1)e((r),, Uejes (x) be the solution of the (i.J)-GSB and ue(2), pe(x) defined in (13), (14). Then, the following bound holds for Jor,
S Jor JGMM S -f' JRP
ΣΣ P'
Jap Peig-4lje ugp23 112 mminf (PepPelag P dx dr (A.20)
JJ 3'
where the dependence on omitted for notational convenience.
24

---

Proof. We first express uz(x) (x) as expectation, i.c.,
Σ (x) ( Pc(z)
u (z)
[x]]
where wr(x)) follows discrete distribution defined by (we(x) (z))(() w.p. (()(())\s//(()))} Note that for random variable æ IR,, the variance decomposition yields
||]]|| [[||=|]]] E||l E[=]]]]]
Using the last equation and the expression (x) above, written as an expectation, we obtain
JGMM / pc(x)) (x(x)) da dt
f./ f Σ
AisPans XesPons (mnR)))) dr di -fo' Σ (p) we(x)) do di Rd i.
dod
-
Jor Σ
AupP uell dr dt.
The fact that the second in the last equation is non-negative justifies the upper bound in The- orem 2. Next, we show that when the conditional densities are well separated, the second term in the last equation becomes arbitrarily small, and the bound becomes tight. Expanding the term side the norm in the integral of the last equation, and dropping the dependence on æ for notational convenience, we obtain
| u, 11
Ueli.s -
E DARy 15
Pt
U1 sjPt Evy
Meloy 3o Ptle3''
Pt
2
Aue Pe sts'
|aaaaa - 1g
(due to Jensen's 's inequality)
Pt
A Ptlee
Mejerr 11?
Pr.
Pt
Substituting this upper bound in the expression for JGMM and rearranging, we get
0 5 Jor
JGMM
<
A )ypmne)lluuss())) ux(x))1 dx dt
RR
1 4 Pepee dx
8'
Pe dr,
where the last inequality comes from the inequality N This completes the proof.
4:05
< min(a;, az) for all positive numbers
E. 0.
25

---

Letting
I min (ispess Pep di
IJ min( MistAys Popyy dr
that is, the normalized distribution of the minimum of the densities Pelej Pojars and assuming that Eo Belluu))(()) Mele (x)))]] < co, since the policies Utjay are affine with respect to z, we conclude that
Jor JGMM + 0 as
Paps)) +
((..)), (d',5")
(A.21)
where (V,,,)) v) denotes the total variation between two probability measures.
A.6 Training and Inference complexity of GMMffow
In this section, provide computational complexity analysis of Algorithm I with respect to the number components in each mixture and the problem dimension. The computational complexity of fitting GMM using the EM algorithm scales O D2)) (Pedregosa et al., 2011) 1), where / is the number of EM iterations, N is the number of data points, K is the number of Gaussian components (modes), and D is the dimensionality of the data. Once the GMMs are fitted, solving a linear program with No X N1 variables, where No and N₁ denote the number of modes in the input and output distributions, respectively. Modern solvers such MOSEK (Mosek, 2020) effi ciently solve LP problems using interior-point methods, which have a computational complexity of ((ViN3 (Boyd& & Vandenberghe, 2004), where I represents the number of constraints.
Regarding the computational complexity of inference, each evaluation of the GMMffow policy, i.e., equation (13), scales linearly with the number of components in each mixture and the SDE integra- tion also scales linearly with the number of discretization time steps. In practice, when implementing the GMMflow policy, only a small number of GSB policies are computed since the component level transport plan Aı sparse. Moreover, this computation is done in parallel for all conditional policies together. This results in very fast, practically constant-time inference regardless of the component number or problem dimension.
A.7 Proof of Theorem
We start by noting that can write the dynamical system (16b) in the form of (A.1) with
01 0
0
4
B, =
D.
(A.22)
0
I
Vel
Due to (19) the joint density of the phase space (are, vi)) is given by
(
Dav Yee
Pe(xx v) = Σ A N
(A.23)
i
We now note that for all j -..... M, the position marginal Øt, at time ty is distributed as
N, Σ
It. Σ AN 26
(A.24)
k=1 (kiy=k)
k=
and therefore the flow (19) satisfies the constraint (16c). Using a similar approach to the proof of Theorem we will show that (18) produces the probability flow (19) by summing over all con- ditional GMSB flows. To facilitate notation, we will denote the phase space by 2 E R2d i.c., 2 o). Given that the individual policies deji solve the GMSB subproblems (10), the pair (Paji Mejs) satisfies the FPK equation for the dynamical system (A.1), that is,
Opys Öt
1
+ V. (Pejs (Az + Buqj)) - tr (DDTV = O. 2
(A.25)
Multiplying equation (A.25) by A and summing over i, we obtain
a Pen
Σ
A
V - (pus (Az + Buex))
(DD(((((()11))) 0,
(A.26)
a
26

---

which implies that
① +V.
Az E Pilids + B Σ uniieeldi
Ōt
[ Eouid)) (
2
DDTF2
(A.27)
This can be further simplified as
+ (pe
))
Ppe +
Pelidi
Az Σ 4zli
(DD'V =0,
(A.28)
E1 Pelsd1
which yields that
Op:
V. (pr(A=+Buu)
tr (DD) = 0.
(A.29)
Ot
This completes the proof.
A.8 Proof of Theorem
The proof is similar to that of Theorem 2. Substituting policy (18) to the cost (I6a) we obtain
2
Pep(ZE)
JCMM =Ezeupe
Σ (z
dt Eeprli
(A.30a)
S/' pc(2)
Prja(2)Ab
Σ mel((=)
d=
(A.30b)
Σi EiP()(2)21
/ po(2) En [| (epp())
dz dl
(A.30c)
EiPrn(=)d
Σ ||uqa(=)]| (ep()))
(A.30d)
Σ A [(lmantte)) at
(A.30e)
-DS= JoT,
(A.30f)
where (A.30b) is due to Fubini's theorem (Wheeden & Zygmund, 1977, Theorem 6.1) and (A.30c) makes use of the discrete version of Jensen's inequality (Wheeden & Zygmund, 1977, Theo- rem 7.35).
A.9 Proof of Theorem
The proof is similar to the (discrete) GMM case. First, notice that
Po dA(wp, w1) = f (dt
gm mm P0
(A.31)
P1 mm Pilworer dA(wo, w1) Rm
(A.32)
Next, notice that Pelwaaws and Welwaww w. satisfy the FPK equation:
i Pt
V (Ptheo,,n, (Arz, + B wo an )) tr (D.DIV2 )) = 0.
(A.33)
Ot By taking the expectation with respect to the distribution A A(wa, wi) in (A.33), we get
J JE X
apejugg
V
(ot (ADIz + Beue
Ot
(D.D)."(P
dA (A(wo,wi)= (A.34)
2
27

---

which implies that
①
/
Pe/wo,en dA dA (Wo, w1)
Öt
m. xkm
V.
Arte
Pelaonee, dA(wo, w1) + Br
JRm xRm
us wo, Ptfwe, dA(wo,; wi)
JRm xRm
tr Di (1en Pe) wi dA dA (wo: 21))
=
(A.35)
which yields that
Ppr
[
Ptl wo wy dl(wo, w1)
01
+ V.
Acate + Br
JRm xpm
Welwoow
J xR Pt/ve wy dA (wo, wi)
tr (D:))V((p:)) =0.
(A.36)
Hence, we conclude that
Ppt 01
(pi(A, + B,u))) - 2 (D.) O.
(A.37)
A.10 Proof of Theorem 6
The proof is similar to the (discrete) GMM case. We can compute that
an (x) dd(ww, w1)
JGMM Eexcopp
ue/woowu (x)
dl
(A.38)
JRm xR Pt/awo.co (x)dA(wo, w1)
f/ [ Pt
Wel uro i (x) Petwo.ws (2)) ww)
drdt
(A.39)
Rm
R xBm
Rm Rmn am Pe mo (x)) wa)
! J Pt
/ellaa (x)ll Pt Pelwa, (z)dA(wo) (oo wi)
R se R.
dadl
(A.40)
JEmxam Pelwores (ar (WW)) w1)
( JRP RR XRm
Ielefroow (x) Pr woo (1) (r)dA(wo, wi drdd
(A.41)
V/ Ieeeoroo (x) dA w1).
E Ea Semuptlsooe Pa
(A.42)
Rm RRm
Hence, for any E (Too Pı).
JGMM < [ J [wo, wi )dA(ddAA.ooooo wi).
(A.43)
By taking the infimum over A E II(Po, P1) in (A.43), we conclude that JGMM S Jor.
B Continuous Gaussian Mixtures
Theorem 6 reduces the high-dimensional dynamic optimal transport problem (21) to the simpler, static OT problem (24) in the space of couplings between the parameter distributions, i.e., IIPPo Pi)- Although in general, problem (24) is still difficult to solve, in many practical applications the pa- rameter spaces Wo, Wi are one-dimensional and J(wo, w1) ) has tractable closed form. Under these conditions, and provided the distributions Po. P1 admit positive densities po. P1. (24) can solved in almost closed form (Santambrogio, 2015).
Our motivation for the extension continuous mixtures stems from the fact that many heavy-tail distributions, such as the multivariate Student-t distribution and the alpha-stable distribution, can be expressed in the form of continuous Gaussian mixtures. In this context, one can use a general- ized version of Theorems I and 2 to create tractable upper bounds on the 2- 2-Wasserstein distance between such distributions and approximate the corresponding optimal transport map, displacement interpolation, and flow fields, respectively.
28

---

and PDF given by B.1 Multivariate t-Distribution
The Student-t distribution has been used as a heavy-tailed alternative to the Gaussian distribution, as generative prior distribution in diffusion models and related generative models in the recent litera- ture; see e.g., Kim et al. (2024); Pandey et al. (2025); Cordero-Encing et al. (2025). In this section, will explore the cases of Theorems and to the case of Student-t boundary distributions.
To this end, let zo, x1 follow d-dimensional multivariate t-distributions with parameters Vo, po: 2o and P1, MM1, En respectively.
multivariate t-distribution can be viewed as a generalized Gaussian mixture model; see, for ex- ample, Andrews & Mallows (1974), More specifically, let uo, in follow a gamma distribution," i.e.
Ug
(/,
u1 Gamma(p) /2/
(B.2)
and conditional on uo. L1, zi ~ ~N(H, un Eo) and ø1 ((1, "1 E1) respectively. Then, wf = Wi follows the distribution InverseGamma( /b /2/2, vo/2) and w/ un follows the distribution InverseGamma(1//)) (1/)), that is, w has the probability density function
(2/2 T(vo/2) /2)
(1/x) +l -
(B.3)
and w/ has the probability density function
(2/2)) (( (1/z) + -
-
(B.4)
Let Po,Pi.PO Pr the CDFs and PDFs of wo, W1 respectively. Starting with the CDF of Wi, we have
Po(x) = P(w S x) P(w S <x²) L.
(2/vo) u0/2
T( T(
T(vo/2) (1/y) (1/4)
+1 e- F0B dy
2 2n
(B.5)
T Ua
which implies that wp Po has the probability density function
d P(w S x) = 2r (2/2))2022 20/2 2(vo/2) (1/2²)
= -i: 2(2/vo) 2o/2
Po(x) = dr
(1/2))
u0 mou
l(vo/2)
(B.6)
Similarly, wi ~ P1 has a CDF given by
I( 2 (122
P1(x)
(B.7)
I(4,0)
(1 2(2/v1) T(m/2) (1/x) ~7
(B.8)
With equations (B.5)-(B.8) in mind, consider Problem (21) with Student-t distributions with param- eters vo, Mo Eg and vı, M1, E,, that is, continuous mixtures of the form
Po(x) / (xx po. w you ppo(uu )
(B.9a)
(1(x)= = / N((x: M1: E E1)p (wi )dw1,
(B.9b)
"The density of a multivariate t-distribution with v degrees of freedom, and scale and location parameters n, Σ respectively, given by
(( 111/2)/d/23-//1151 1/2 + (x T (x m)
(u4d)/2
(B.1)
Here Gamma(a, b) denotes gamma distribution with probability density functional proportional to c- -bx where is the shape parameter and is the inverse scale parameter.
29

---

where Po(wo): Pi (wi) are given by (B.6) and (B.8) respectively.
Considering noise-free dynamics to simplify the respective formulas, in (21b), the (ooo-w))-SS for the boundary distribution parametrization (B.9), admits the following closed form
Σ t0
(1 u + = + (1 -t)twgry (C+CC).
(B.10a)
M4] w1
{1 t)1m0 + ttm
(B.10b)
Kg jung,en wn
5!
(B.10c)
0rjorow1 wF
M1 1o:
(B.10d)
S Salwo:
= = - C² - (1- C)2
(B.10e)
where C= 2 D
D = (20 Dj Eg ) Equations (B.10) yield
Pej
(x) =N (x: Mallog Eejwoows
(B.1la)
Ut 100.00
(z) K k (x Me/mro, + M1 140-
(B.11b)
Furthermore, the optimal cost J Jww, w1) ) in (24) be calculated using (9) and equals
Jww, wn) I1 n 2 w (Ft) + wftr(En) tr()).
(B.12)
Focusing on problem (24), known that when the transport cost has the form J(wo. wi) = h(wo wi) where h is a convex function, the corresponding optimal transport plan is given by (Santambrogio, 2015, Theorem 2.9)
A' w1) = (P)'
(z), (z)) (n0 1]),
(B.13)
where [n0, 1]) })) is the uniform measure over the set [0, 1), and Po P are the corresponding inverse CDFs of (B.5), (B.7). The cost function (B.12), does not satisfy this condition, since cannot be written as perfect square for general scaling matrices Eo. Ej. We resolve this issue through the following proposition.
Proposition 3. When solving the OT problem (24), the transport cost (B. 12) and the cost function Jun, w1) = =wwo- w2l are equivalent, i.e,, they result in the same optimal coupling A
Proof. Equation (B.12) can be written in the form:
Jwo, wi) = |l40 - f1ll + w tr(22 - D) + tf D) + wwo w1 12 tr()..
(B.14)
In (B.14), the terms 11 iu 1| witr(Eo- D), wftt(EI- do not contribute the optimization problem in (24), since they are constant for any feasible transport plan A E [I(po) P1) due to the fixed boundary distributions. By dropping these terms, as well the positive scaling constant tr (D), we obtain the desired result.
Equation (B.13) implies that the optimal value for Problem (24) is given by
Jôr - J (LP,'(w)) (T)
(B.15)
while the optimal control policy ut (I) and the respective density pf resulting from substituting to the optimal transport plan A to the formulas (22) and (5) of Theorem are given by
(T()) = S P (w))) (w) (x)dw,
(B.16a)
and
ut (x(x)) L' "M]pp"(w)) (w) (x) (mn)) (x)
dw.
(B.17)
Since equations involve only oieedimensi integrals, they be easily computed numerically using quadrature. Although P. PT not available in closed form, they can be
30

---

obtained in most scientific computing packages such as scipy (Virtanen et al., 2020) by properly scaling the quantile function of the inverse gamma distribution.
To illustrate this approach, we caleulate the upper bound (B.15) and the true Wasserstein-2 distance for a one-dimensional problem between two Student-t distributions with parameters μ4 = //2 = O, Σ1 = 1, Σ₂ {0.25, 1, 1,4), 1/2 for various values of 1/2 E [2.5, 10) and report the results in Figure 5. We note that our approach works for arbitrary Student-t distributions in any dimension, however, we study the ID case in this example be able calculate the exact Wasserstein distance and quantify the tightness of our upper bound. Although rigorously studying the tightness of the bound (B.15) remains an open problem, evident in Figure 5, it closely approximates the true Wasserstein distance, at-least in this simple ID scenario.
3: - 1, 5o 0.25
h=15-1
1, 12=4
1.50
1.5
0.5
Jor
0.4
0.75
Jor 0.0
Jar
0.5
4
10
4
6
10
Figure 5: Comparison between true Wasserstein distance and upper bound (B.15) for ID Student-t distributions.
We believe this result, along with the policy (B.17) and the interpolation (B.16), could be useful tools in developing simulation methods for training diffusion models with heavy-tail prior dis- tributions, or for creating tractable OT flows between mixtures of Student-t distributions. We leave these interesting directions as future work, since they are not immediately related to computationally inexpensive diffusion model training, which is the main theme of the rest of this paper.
C Additional Experiments and Implementation Details
C.1 Additional Details on 2D Problems
To compare our approach for the problem of Figure with state-of-the-art neural SB solvers we used the original implementations of the DSB* (De Bortoli et al., 2021) and DSBM (Shi et al., 2023). The network architecture used for both algorithms is the fully connected DNN of De Bortoli et (2021) with 128 de-demensional sinusoidal temporal encodings, 256 neurons in the encoder layer, (256, 256} neurons in the decoder layers, and SiLU activation functions (Hendrycks & Gimpel, 2016). We run all algorithms and report the results in Table 4. For the zero noise case, i.e., e = O, DSB and DSBM are not applicable, so we approximate the true OT cost using discrete optimal transport, calculated using the POT library (Flamary et al., 2021), 1), using 10,000 samples from each distribution. As evident from Table 4, although the DSB and DSBM algorithms can approximate the true SB cost, they fail to retrieve the true optimal solution for larger values of the noise parameter, due to their non-convex loss functions.
Table 4:
Transport cost comparison for the problem in Figure
c
Jor (12a)
JGMM (15a)
DSBM
DSB
OT
0
100.06
89.45
84.87
0.1
100.15
89.32
84.26
98.62
1
102.28
89.28
131.50
100.82
TO
162.13
T16.60
133.04
244.31
Furthermore, regarding the example problem in Figure 2, we provide additional details about the approximation of the boundary distributions as mixture models in Figure 6.
C.2 Image-to-Image Translation Details
To better evaluate the performance of the proposed approach in the Image-to- Image translation task, we provide further examples in Figures 9 and 10 as well as approximate training and inference times
bridge
'https///githul www.//oooonyooeeee sht/ pytorch
31

---

Samp tram A
mexture DRN,
Sumpes A.-
Teroonal micture GMN
Samples from Py
GT GT
DCSL
DCS
DCSL
Figure 6: GT to DCSL distribution steering details.
for our approach, and compare them with the training and inference times of LightSB in Table 5. For our approach, training time consists of the time required fit the GMMs in the latents of the FFHQ dataset for the two boundary distributions, and the solution of the linear program (15). As inference time, we consider the time taken for the integration the SDE (or ODE for e = 0) (12b) with the mixture policy (13). We observe that while inference time is small for solving the deterministic (optimal transport) problem, i.e., for O, integrating the stochastic dynamical system for positive values of requires more time due to the small time step required for SDE integration. The quality of the produced images was not found to be affected by this parameter, implying that be used for fast, deterministic inference, while positive value of will allow for some randomness the generated images. We also note that the faster training time for our approach is mainly due to the very fast convergence of the EM algorithm, which is also less likely converge to local minima, compared to the standard maximum likelihood method for fitting distribution to data. All tests were conducted on desktop computer with RTX 3070 GPU.
could
Table 5: Training and inference time comparison with state of the art. Inference time is measured for a batch of 10 images and uses GPU parallelization for calculating the elementary GSB policies of the mixture policy (13).
Training [s]
Inference (e = 0) [s]. Inference (e = 0.1) [s]
LightSB Ours
57
0.02
17
0.06
0.2
C.3 Multi-Marginal Problems Details
PC1
PC
predicted
true
GMMs
PC4
3
PC
PC
PC
PC
Figure 7: Additional visualization of results in 5-d scRNA problem: predicted vs true distributions for all time marginalss overlayed with the 3-sigma bound for each Gaussian component of the pre- fitted GMMs.
KKoootin et al. (2024)
32

---

Solution times. To solve each marginal GMSB we use the semidefinit formulation detailed in Section D and used Mosek (2020) to solve the resulting semidefinite program. Specifically, we assume the GMMs between the five temporal marginals are spaced 1 time unit apart, resulting in problem horizon of 4 time units. We use a coarse temporal discretization with time-step A/ = 0.1 (i.e., 10 time steps between [t,, ()) to evaluate the cost tensor for the optimization problem (20) and a fine resolution discretization of At = 0.01 for the final policy calculation, solving only for the GMSBs with non-zero transport parameter A1. There are total of 3, 125 combinations of GMSBs for this problem, and the computation of each one using the coarse time grid takes roughly 0.35 s on an Intel i7 12-th generation CPU with 32 GB of RAM memory, giving total of 18.5 minutes of calculations, if all GMSBs are solved serially. This computational overhead can be greatly decreased if the GMSBs solved in parallel. In our setup, parallelized the calculation using MOSEK's built-in capabilities and solved them in batches of 12, using 2 CPU threads per problem. This brought down the total calculation time under minutes. For the final policy calculation, there only 21 active GMSBs in the mixture policy (18), each taking 6 seconds to compute. The total run time for our algorithm for this problem, adds up 8 minutes for this problem, which is considerably lower than the corresponding neural methods (24 minutes on GPU for the DMSB algorithm (Chen al., 2023)).
Velocity inference. After calculating the conditional GMSBs and solving (4), the marginal mix- ture distribution for the entire phase-space is fully defined for the entire time horizon of the problem through Equation (19). Given position sample To at time t = 0, the correspondigg velocity com- ponent can be inferred using conditional GMM sampling. Specifically, considering that the joint distribution of the phase space time is
Pe(a, v) = Σ N
I
(B.18)
it is easy to show that the density of (oloz) is also a Gaussian Mixture Model, since
pc(u|z) = Pc(xx,)) v)
(B.19a)
pc(x)
ΣNN
[
Fee
(B.19b)
E NNN
(z: Pejp Exe
NN (w: Meji
ole
XX1e N
xx; Mijp: 2/
(B.19c)
Σ NWW
[z; 1iip Ing
AW x; Mij
N
0; Heli ofar
20
(B.19d)
ENN 2; 19j
E
where Equation (B.) is due to the Bayes rule, with (Bishop & Nasrabadi, 2006)
m/a Mej
14.
x m
and
r/i
For our problem, initial position, and then use the joint position-velocte initial conditions to calculate the sample's ss trajectory by integrating (12b).
Equation (9.) to sample the initial velocity of new sample given its
Visualization of results. To better visualize our results, we provide more information about the predicted distributions in Figure 7, overlaid with the fre-fitted GMMs at each time step. It easy to visually confirm that the 5-component mixtures capture the marginal distributions in all time-steps accurately, and since our method exact, there is minimal distribution mismatch in the predicted marginals, as confirmed quantitatively by the indices in Table 3.
33

---

C.4 Performance on EOT Benchmarks
To further evaluate the optimality of the proposed approach, we tested the algorithm the Entropic Optimal Transport benchmark detailed in Gushchin et al. (2023). The benchmark provides pair of boundary test distributions po- P1. where po is scaled Gaussian distribution and p1 a mixture- distribution that is easy to sample from, and an optimal conditional transport plan x (xr (zo),, which Gaussian Mixture Model and is known in closed form. For the pair po, Pt, the optimal policy solving (2) can be calculated explicitly, allowing direct comparisons with our approach. The metric we use measure the optimality of our approach is the Bures- Wasserstein Unexplained Variance Percentage (cBW-UVP) (Cani)GGGsieii (Gushchin et al., 2023), defined by
100%
/
cBW-UVP(1,, (
xT)
dro.
(B.20)
Var(pi)
which measures the distance between conditional transport plans, evaluated using the Bures- Wasserstein metric.
To the method of Gushchin al. (2023), we first obtain samples from the two boundary test distributions and then fit mixture models on them using EM. We then deploy policy (13), and report the values of the cBW-UVP index between the known optimal conditional transport plan n² (21 (2)), and the conditional transport plan resulting from the integration of the policy (13), denoted ( [x))) We use 1,000 initial condition samples 2o, and for each sample, we draw 1,000 21 samples from the distributions (((x (2)) and (<x) (xo) to compute the empirical Bures- Wasserstein distance in (B.20). The results are reported in Table 6 for problems of various dimensions and noise levels, along with many other available methods for solving the same problem (Gushchin et al., 2023, Table 5). We note that although our approach requires virtually no training compared computationally expensive neural OT and SB approaches, it outperforms many of these algorithms, outlining its excellent performance in problems where GMMs accurately capture the marginal distributions of the problem.
Table 6: Comparisons of cBW2-UVP (%) between the optimal plan x* and the learned plan n. Colors indicate the ratio the metric to the independent baseline metric: ratio 0.2, ratio (0.2, 0.5), ratio > 0.5.
c=011
c=1
10
D=2 D=16 D=61
D= 128
D=2
D=16
D=
D=
128
D=2
D=16
D=
D= 128
LSSTT]
SCONES]
34,88
71.34
59.12
136.44
32.9
50.84
60.44
52.11
[NOT]
1.94
13.67
11.74
11.4
4.77
23.27
41.75
26.56
2.86
4.57
3.41
6.56
EgNOT
129.8
75.2
60.4
43.2
80.4
74.4
63.8
53.2
4.14
2.64
2,36
1.31
ENOT
3.64
13.6
12.6
1.04
21.6
1.4
2.4
19.6
MLE-SB
4.57
4.13
D
73.54
59.7
1386.4
1683.6
33.76
70.86
53.42
156.46
FB-SDE-A
8644
1166.82
1566.44
30.62
131.72
FB-SDE-J
51.34
89.16
119.32
173.96
29.34
69.2
155.14
177.52
[SBM
5.2
16.8
37.3
35
0.3
9.7
3.7
3557
15000
SF M-Sink
[LightS]]
0.03
0.08
0.28
0.60
0.05
0.09
0.24
0.62
0.07
0.11
0.21
0.37
[ightSB-M (MB)
0.005
0.27
0.63
0.002
0.04
0.12
0.36
0.04
GMMMBow (ours)
10.35
14.68
11.15
11.2
5.78
7.20
0.16
0.28
1.43
2.77
Independent coupling
166.0
152.0
126.0
110.0
86.0
80.0
72.0
60.0
4.2
2.52
2.26
2.4
To further benchmark our algorithm with respect to run-times, we provide wall-clock times for both training and inference with respect to the number of components and the problem dimensionality for the boundary distributions provided in the EOT benchmark. We report these values Table
below.
Table 7: Training time for for EOT benchmark.
Dim
comp
5
10
20
50
100
2
0.209
0.0595
0.1083
0.2303
1.0256
16
0.0567
0.0948
0.1413
0.4218
3.3918
64
0.182
0.2228
0.7217
1.2875
2.1431
128
0,2802
0.4562
0.7713
1.6164
3,4101
34

---

Table
Inference time for for EOT benchmark.
Dim # comp
5
10
20
50
100
2
0.037
0.031
0.031
0.032
0.030
16
0.032
0.032
0.032
0.033
0.032
64
0.032
0.032
0.032
0.032
0.039
128
0.032
0.033
0.036
0.032
0.082
c-0.10
-5
-10-
-10
-10
Figure 8: Gaussian B-Gaussians with LTT prior dynamics.
We note that because the EOT benchmark uses initial Gaussian distribution and GMM terminal distribution, there is no point in reporting metrics such as marginal distribution accuracy or transport plan optimality. since these will perform best when the number of components used in GMMflow matches the setting of EOT benchmark. Furthermore, exploring how well a GMM approximates general distribution beyond the scope of our work; therefore, we do not provide experiments that explore this issue.
the number of components increases is well-studied problem and goes
C.5 Problems with LTI Prior Dynamics
To test the algorithm on more complicated dynamical systems, we use the 4-dimensional Linear Time-Invarian (LTI) system
dæ; Ar, dt + Bug dt + D Ddw,
(B.21)
with
2S S
L 02
-1
0 T2 7
A
$
B
= cT
and boundary distributions
Pu Σ k=0
V
(B.22a)
P1 (O,, 0.414)
(B.22b)
We note that solving problem (12) with the dynamical system (B.21) in place of (12b) is not currently solvable using any mainstream neural SB solvers because the stochastic disturbance du in (B.21) does not enter through the same channels as the control signal u, and the state x;. The only available method solve this problem is detailed in Chen al. (2016), which, however, assumes access to the solution of the static EOT problem (1) with boundary distributions (B.22), and a closed form of the probability density transition kernel included by the dynamical system (B.21) for u, 0. The results of our approach are illustrated in Figure 8.
To solve the Gaussian Bridge sub-problems with a dynamical system of the form (B.21) we use discrete-time convex formulation of Rapakoulias & Tsiotras (2023). We also include brief overview of the method in Appendix D. Each continuous-time Gaussian Bridge is discretized (in the temporal dimension) into 101 steps over uniform intervals of size Al 0.01. We used MOSEK (Mosek, 2020) to solve the resulting semidefinite programs.
35

---

D Gaussian Bridge for Linear Time- Varring Systems
In this section, we briefly review the available methods in the literature to solve the Gaussian Bridge problem with general LTV dynamics of the form (A.1). That is, we consider the problem
[. [lu
min E well
(C.la)
dx. = AcT dt + B(u( dt + D
(C.1b)
30 NNo, Ea), x1 ((),,))))
(C.1c)
The solution of problem (C.1) is used to solve the Gaussian Bridge problem for the example in Section C.5 and is relevant to applications with prior dynamies of more general structure such as mean field games (Bensoussan et al., 2016) and large multi-agent control applications (Saravanos al., 2023) higher-order distribution interpolation problems (Chen al., 2018, 2019). The existence and uniqueness of solutions for problem (C.1) are studied in Chen et al. (2015a); Liu et al. (2025); Liu & Tsiotras (2024). Since the state of (C.1b) remains Gaussian throughout the steering horizon, i.e., It ~ (Re Ex), the problem simplifies to that of the control of the first two statistical moments of the state, namely the mean Pr and the covariance E, Using a control policy parametrization of the form
(C.2) allows for the decoupling of the propagation equations for the mean and covariance of the state. More specifically, applying (C.2) to (C.1b), the equations describing the propagation of bt and Σ yield (Särkkä & Solin, 2019, Section 5.5)
Uy (x()) K2(x -4) + Ut:
ΣS (A: + B,K,)2, + Ec(A, B,K,) D.D!,
(C.3a)
A = Appe + B,g,
(C.3b)
Expanding the expression (C.3a) and performing the change of variables U, K,E,, we obtain
Σ A,E, E,A + B,U, U.B.B -
(C.4)
which is linear in U,, E. Furthermore, substituting (C.2) into the cost function (C. and using the eyclic property of the trace operator along with the standard properties of the expectation yields
E L
[|ux(=)]1 dt J. u/ tr (KEE,K) dt dt= L. i U2 tr (U,E, UU] dt. (C.5)
Equations (C.3b), (C.4), (C.5) can used to reformulate problem (C.1) to a simpler optimization problem in the space of affine feedback policies, parameterized by U2 and us To be more precise, problem (C.1) reduces to
[ u71e + tr (U,E, U?! dt,
min N..
(C.6a)
Σ A, + ETA + B,U, B. - D.D.
(C.6b)
jr Awd Biu
(C.6c)
which can be further relaxed to a convex semi-definite program using the lossless convex relax- ation (Chen al., 2015b)
min / uTue+ tr(Y E.
(C.7a)
U.E. UT < Yr,
(C.7b)
$. A,E, E,A] + B,U, + B.B. D.D.
(C.7c)
i4 Arj. + B
(C.7d)
after noting that the constraint (C.7b) can be cast as a Linear Matrix Inequality (LMI) using Schur's complement as
Σ U U, Y =
Problem (C.7) is still infinite dimensional since the decision variables are functions of time t E [0, 1]; however, it can be discretized, approximately using a first-order approximation of the derivatives in (C.7c), (C.7d) (Chen et al., 201.5b) or exactly using a zero-order hold (Liu et al., 2025; Rapakoulias & Tsiotras, 2023), and solved to global optimality using a semidefinite programming solver such as MOSEK (Mosek, 2020).
36

---

Original
LightBB
Go With the Flow
Figure 9: Further examples for the Image-to- tomee translation task.
37

---

Original
LightSB
Go With the Flow
Figure 10: Further examples for the adult-to-child Image-to- Io-mm translation task.
38


Reviews:
[{"review_id": "2Lgu9WYCJt", "reviewer": "Reviewer_HrsH", "summary": "The authors propose a lightweight method for approximating solutions to SB problems. Their approach has them approximate the boundary distributions by Gaussian Mixtures (GM) and then approximately solving the SB problems for the GMs. Their theoretical results motivate the algorithm for doing the latter, showing that any solution to the GM SB problems has a policy and flow that can be described as a combination of the policy and flow for the SB solution between each of the individual Gaussians. They provide experiments comparing their results to existing SB algorithms, demonstrating improvements.", "questions": "- Does the final algorithm optimise J_OT to obtain \\lambda_{ij}?\n- If so, is it clear how close the objective is to J_GMM or how well the approximation is actually solving the GM SB problem (not just the OT problem). Could you please point to evidence for this?\n- Is \\rho_{ij} and u_{ij} obtainable in closed-form? If not, how difficult is it to approximate?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n- The algorithm appears to be highly efficient and most SB algoriths are highly inefficient\n- The approach of using GMs as the endpoints is simple and yet appears to be quite novel to the literature. The proposal of this could motivate future work\n- The paper is largely well-written with the preliminaries section especially serving as a high quality introduction to the field\n- The theoretical results regarding SB for GMs is interesting in its own right and could motivate future work\n- The extension of the work to momentum SB and continuous GMs demonstrates a deep exploration of the ideas considered in the work.\n\nWeaknesses\n- The algorithm they ultimately use is not very clearly presented or discussed much. It would be good to include a clear presentation of this following the theoretical results. For example, it would be good to know how J_OT should be optimised\n- The Comparison seems somewhat limited, focussing primarily on comparison with other light methods. A more thorough comparison with standard SB algorithms would help.\n- It is not clear to me how close J_GMM is to the proxy J_OT. There are some experiments in table 4 that provide some results in this direction, but a more thorough discussion and set of experiments about how much the proposed method actually solves the GM problem would greatly improve to the validity of the paper.", "was_revised": true, "final_justification": "The paper proposes a fairly naive methodology, but it does seem to work well. I think its simplicity is a strength and not a weakness and the fact it works well will be interesting to the ML community. Also the quality of presentation, the depth to which the idea was explored and the fact that rigorous theoretical results were provided make this worth accepting.", "rebuttal": "We would like to thank the Reviewer for his/her constructive comments and for the positive evaluation of our work.\nIn the following, we address the comments and questions of the Reviewer one by one.\n\n1. **Optimization of $J_{\\\\mathrm{OT}}$ to obtain $\\lambda_{ij}$**: In order to obtain the weights $\\lambda_{ij}$ of the $(i, j)$-GSB subproblem, the algorithm indeed optimizes the cost function $J_{\\mathrm{OT}}$ of the linear optimization problem given by Equation (15) in Theorem 2. To solve Problem (15), each term $J_{ij}$ is first computed through Proposition 2 and then, using the known constants $\\alpha_0^i, \\alpha_1^j$ the linear program (15) is solved using off-the-shelf LP solvers such as MOSEK.\n\n2. **Bounds on the optimality gap**:  We thank the Reviewer for this question. As the Reviewer points out, some empirical evidence about the optimality gap is presented in Table 4 and also in the extensive benchmark testing in Section C.4. \nTo gain a deeper theoretical understanding for the cases when Theorem 2 introduces some sub-optimality, we provide the following alternative proof of Theorem 2, along with a result proving that when the mixture components are well separated, Theorem 2 yields the optimal solution. \n\nSpecifically, in the setting of Problem (12), let $\\rho_{t|ij}, u_{t|ij}$ be the solution of the $(i,j)$-GSB, with boundary distributions the $i$-th component of the initial and the $j$-th component of the terminal mixture, let $\\rho_t = \\sum_{ij} \\lambda_{ij} \\rho_{t|ij}$ and let\n\\\\[\nu_t(x) = \\sum_{i,j} {u_{t|ij}(x) \\frac{ \\rho_{t|ij}(x)\\lambda_{ij}} {\\rho_{t}(x)}} = \\mathbb{E} [ \\omega_t(x)],\n\\\\]\nwhere $\\omega_{t}(x)$ follows a discrete distribution defined by $\\\\{ \\omega_t(x) = u_{t|ij}(x) \\\\, \\\\, \\mathrm{w.p.} \\\\, \\\\, \\frac{ \\rho_{t|ij}(x)\\lambda_{ij}} {\\rho_{t}(x)} \\\\}$.\n\nNote that for any random variable $x \\in \\mathbb{R}^n$, the variance decomposition yields \n\\\\[\n    \\\\| \\mathbb{E}[x] \\\\|^2 = \\mathbb{E}[ \\\\|x\\\\|^2] - \\mathbb{E}[ \\\\|x  - \\\\mathbb{E}[x]\\\\|^2 ].\n\\\\]\nUsing the last equation and the expression $u_t(x)$ above, written as an expectation, we obtain\n   \\\\[ \n   J_{\\mathrm{GMM}} \\triangleq  \\int_0^1 \\\\!\\\\! \\int_{\\mathbb{R}^n} \\\\!\\\\! \\rho_{t} (x) \\\\| u_t(x) \\\\|^2 \\\\, \\mathrm{d} x \\\\, \\mathrm{d} t \n   \\\\]\n   \n   \\\\[ \n   \\Rightarrow J_{\\mathrm{GMM}} =\\int_0^1 \\int_{\\mathbb{R}^n}  \\sum_{i,j} \\lambda_{ij} \\rho_{t|i,j} \\\\|u_{t|i,j}(x)\\\\|^2  \\\\, \\mathrm{d} x \\\\, \\mathrm{d} t - \\int_0^1 \\int_{\\mathbb{R}^n} \\sum_{i,j} \\lambda_{ij} \\rho_{t|i,j} \\\\|u_{t|i,j}(x)- u_t(x)\\\\|^2  \\\\, \\mathrm{d} x \\\\, \\mathrm{d} t \n   \\\\]\n\\\\[  \\Rightarrow  J_{\\mathrm{GMM}} = J_{\\mathrm{OT}} - \\int_0^1   \\int_{\\mathbb{R}^n} \\sum_{i,j} \\lambda_{ij} \\rho_{t|ij} \\\\|u_{t|ij} - u_t \\\\|^2  \\\\, \\mathrm{d} x \\\\, \\mathrm{d} t  \\\\]\nThe fact that the second term in the last equation is non-negative justifies the upper bound in Theorem 2. \nNext, we show that when the conditional densities are well separated, the second term in the last equation becomes arbitrarily small, and the bound becomes tight.\nExpanding the term inside the norm in the integral of the last equation, and dropping the dependence on $x$ for notational convenience, we obtain\n \\\\[ \\\\|u_{t|i,j} - u_t\\\\|^2 \\\\]\n \\\\[ = \\bigg\\\\|u_{t|i,j}-  \\sum_{i'j'} u_{t|i'j'} \\frac{\\lambda_{i'j'} \\rho_{t|i'j'}}{\\rho_t} \\bigg\\\\|^2   \\\\]\n \\\\[ = \\bigg\\\\| \\frac{ u_{t|ij} \\rho_{t}- \\sum_{i'j'} u_{t|i'j'} \\lambda_{i'j'} \\rho_{t|i'j'} }{ \\rho_t} \\bigg \\\\|^2 \\\\]\n\\\\[ \\leq\\left(\\sum_{i'j'}\\Vert u_{t|i,j}-u_{t|i'j'}\\Vert\\frac{\\lambda_{i'j'} \\rho_{t|i'j'}}{\\rho_t}\\right)^{2} \\\\]\n\\\\[ \\leq\\sum_{i'j'}\\Vert u_{t|i,j}-u_{t|i'j'}\\Vert^{2}\\frac{\\lambda_{i'j'} \\rho_{t|i'j'}}{\\rho_t} \\\\]\n\\\\[ \\sum_{\\\\substack{i' \\neq i \\\\ j' \\neq j}}\\Vert u_{t|i,j}-u_{t|i'j'}\\Vert^{2}\\frac{\\lambda_{i'j'} \\rho_{t|i'j'}}{\\rho_t} \\\\]\nwhere in the 4th equation we applied the discrete version of Jensen's inequality. \nSubstituting this upper bound in the expression for $J_{\\mathrm{GMM}}$ and rearranging, we get \n\n \\\\[ 0 \\leq J_{\\mathrm{OT}} - J_{\\mathrm{GMM}} \\\\]\n \\\\[ \\leq \\int_0^1 \\int_{\\mathbb{R}^n} \\sum_{i,j} \\lambda_{ij} \\rho_{t|i,j} \\\\|u_{t|i,j}(x)- u_t(x) \\\\|^2  \\\\, \\mathrm{d} x \\\\, \\mathrm{d} t \\\\] \n \\\\[ \\leq  \\int_0^1  \\int_{\\mathbb{R}^n} \\sum_{i,j}  \\sum_{\\substack{  i' \\neq i \\\\ j' \\neq j} } \\big \\\\|u_{t|ij}  -  u_{t|i'j'} \\big \\\\|^2 \\frac{ \\lambda_{i'j'} \\lambda_{ij} \\rho_{t|i'j'} \\rho_{t|i,j}  }{  \\rho_{t}  }  \\\\, \\mathrm{d} x \\\\, \\mathrm{d} t \\\\]\n  \\\\[ \\leq  \\int_0^1 \\int_{\\mathbb{R}^n} \\sum_{i,j}  \\sum_{\\substack{  i' \\neq i \\\\ j' \\neq j} } \\big \\\\|u_{t|ij} - u_{t|i'j'} \\big \\\\|^2 \\min \\\\{ \\lambda_{ij} \\rho_{t|ij},  \\lambda_{i'j'} \\rho_{t|i'j'} \\\\}  \\\\, \\mathrm{d} x \\\\, \\mathrm{d} t  \\\\]\nwhere the last inequailty comes from the inequality $\\frac{a_i a_j}{\\sum_i a_i} \\leq \\min\\\\{a_i, a_j\\\\}$ for all positive numbers $\\\\{a_i\\\\}_{i=1}^N$. \n\nIntuitively, when the components are well separated, the term $\\min \\\\{ \\lambda_{ij} \\rho_{t|ij},  \\lambda_{i'j'} \\rho_{t|i'j'} \\\\}$ becomes arbitrarily small, making the bound tight. \nMore rigorously, noting that, for any vectors $a,b$, $\\min\\\\{a,b\\\\} = \\frac{1}{2}(a+b) - \\frac{1}{2}|a-b|$, \nwe get\n\n\\\\[ \\int_{\\mathbb{R}^n} \\min \\\\{ \\lambda_{ij} \\rho_{t|ij},  \\lambda_{i'j'} \\rho_{t|i'j'} \\\\} \\\\, \\mathrm{d} x \\\\]\n\\\\[ =\\frac{1}{2} (\\lambda_{ij} + \\lambda_{i'j'}) - \\frac{1}{2} \\int_{\\mathbb{R}^n} | \\lambda_{ij} \\rho_{t|ij} -  \\lambda_{i'j'} \\rho_{t|i'j'} | \\\\, \\mathrm{d} x \\rightarrow 0 \\\\]\nor that \n \\\\[ \\frac{1}{2} (\\lambda_{ij} + \\lambda_{i'j'}) - \\mathrm{TV}(\\lambda_{ij} \\rho_{t|ij}, \\lambda_{i'j'} \\rho_{t|i'j'}) \\rightarrow 0 \\quad \\text{as} \\quad  \\mathrm{TV}(\\rho_{t|ij}, \\rho_{t|i'j'}) \\rightarrow 1 \\\\]\n\nSince $\\int_{\\mathbb{R}^n} \\min \\\\{ \\lambda_{ij} \\rho_{t|ij},  \\lambda_{i'j'} \\rho_{t|i'j'} \\\\} \\\\, \\mathrm{d} x\\rightarrow 0$ for every $t\\in[0,1]$ and\n$\\int_{\\mathbb{R}^n} \\min \\\\{ \\lambda_{ij} \\rho_{t|ij},  \\lambda_{i'j'} \\rho_{t|i'j'} \\\\} \\\\, \\mathrm{d} x\n\\leq\\int_{\\mathbb{R}^n} \\{ \\lambda_{ij} \\rho_{t|ij}+  \\lambda_{i'j'} \\rho_{t|i'j'} \\} \\\\, \\mathrm{d} x=\\lambda_{ij}+\\lambda_{i'j'}$\nuniformly in $t$, by bounded convergence theorem, $\\int_{0}^{1}\\int_{\\mathbb{R}^n} \\min \\\\{ \\lambda_{ij} \\rho_{t|ij},  \\lambda_{i'j'} \\rho_{t|i'j'} \\\\} \\\\, \\mathrm{d} x\\mathrm{d}t\\rightarrow 0$.\n\nGiven that the the policies $u_t$ are affine functions of $x$ (see Proposition 1), $\\mathbb{E}\\_Q [ \\\\| u_{t|ij}(x) - u_{t|i',j'}(x) \\\\|^2] < \\infty$  where $Q = \\\\frac{\\\\min \\\\{ \\lambda_{ij} \\rho_{t|ij}, \\lambda_{i'j'} \\rho_{t|i'j'} \\\\} }{ \\int \\int \\min \\\\{ \\lambda_{ij} \\rho_{t|ij}, \\lambda_{i'j'} \\rho_{t|i'j'} \\\\} \\\\, \\mathrm{d} x \n\\\\, \\mathrm{d} t} $, i.e., the normalized distribution of the minimum of the densities $\\rho_{t|ij}, \\rho_{t|i'j'}$.\nHence, we conclude that  \n\\\\[ J_{\\mathrm{OT}} - J_{\\mathrm{GMM}} \\rightarrow 0  \\quad \\text{as} \\quad  \\mathrm{TV}(\\rho_{t|ij}, \\rho_{t|i'j'}) \\rightarrow 1. \\\\]\n\n3. **Closed-form expressions of $\\rho_{ij}$ and $u_{ij}$**: Both the policy $u_{t|ij}$, the probability density $\\rho_{t|ij}$, and the cost $J_{ij}$ for the $(i, j)$-GSB subproblem can be obtained in closed form through Propositions 1 and 2 for the case of the Gaussian SB. For the more general cases of the Gaussian multi-marginal momentum SB or the more general SB with linear prior dynamics, the policy $u_{t|ij}$, $\\rho_{t|ij}$, and $J_{ij}$ are obtained by solving a semidefinite program. We refer to Appendix D for a detailed discussion.  \n\n\n4. **Clear presentation of the algorithm**: We thank the reviewer for this suggestion.\nWe agree that an algorithmic block would improve the readability of our work. To this end, we will include a pseudocode block of our method in Section 3 of our manuscript. \nTo avoid a duplicate response, we refer the Reviewer to our response to Reviewer VvBB for a preview of our algorithm. \n\n\n5. **More thorough comparison with standard SB algorithms**: In our original submission, we compared our method against neural approaches in the 2D experiments (Section C.1, Table 4, compared against DSB, DSBM), the EOT benchmark (Section C.4, Table 6, compared against all state-of-the-art methods), and the multi-marginal RNA example (Section 5, Table 3, compared against NLSB, MIOFlow, and DMSB). \nWe stress that for the multimarginal momentum SB problem, there are no available light-weight approaches in the literature.  \nIn the image translation example, we only compared against LSB and LSBM since these match the scope of our algorithm. Due to space and time limitations, we have not compared against neural approaches in this experiment; however, we point the Reviewer to the excellent work [1] for a comparison of LSBM with neural approaches for the image translation task. \n\n[1] Gushchin, N., Kholkin, S., Burnaev, E., \\& Korotin, A. (2024, July). Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning."}, {"review_id": "aO0BJ7WJU8", "reviewer": "Reviewer_PeSU", "summary": "The paper develops a computational method to compute suboptimal solutions of the schrödinger bridge problem between gaussian mixture models, and showcases the potential of the problem in low-dimensional settings. The key idea is to exploit the decoupling of particles dynamics with the assignment problem, building on top of recent work on the topic.", "questions": "1. Can the authors discuss further their suboptimality gap? Related: How does the approximation of GMM further exacerbate this gap?\n2. Can the authors provide comparisons (at least in the deterministic case) with the approach in https://epubs.siam.org/doi/abs/10.1137/23M1560902?\n3. Can the authors clarify the setting for the RNA dataset?", "limitations": "The authors adequately discuss the limitations.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "The paper is very focused on a specific problem with clearly stated limitations. The theory is correct (beside a few minor clarifications needed) and gradually introduces more complexity, rendering the paper very pleasant to read. I identified the following weaknesses, that the authors may consider to address to strengthen their contribution:\n\n1. Although reading the entire paper makes the scope of the work very clear, I believe it could be stated better in the introduction: low-dimensional problem in which the marginals can be well approximated by GMM, and in which there is a need for a good solution, fast, but not necessarily optimal. To build on this, the authors could also exemplify a few of these applications.\n2. In the contributions, it is stated that the paper proposes a solution to the SB problem with GMM, but technically, it only provides a suboptimal solution. Moreover, the suboptimality gap is not quantified and may be very large. Carefully addressing this gap or clarifying this potential limitation will strengthen the paper for other people to build on top.\n3. The authors cite https://epubs.siam.org/doi/abs/10.1137/23M1560902, but perhaps they could also contextualize their result with respect to the recipe provided in https://epubs.siam.org/doi/abs/10.1137/23M1560902. Namely, the mixing discussed in Go with the Flow is an approximation of the mixing derived in https://epubs.siam.org/doi/abs/10.1137/23M1560902, hence the suboptimality. Note that the work cited is for deterministic case, another suboptimality comes from the stochastic dynamics in this type of mixing. When comparing with the literature, in the deterministic case, they could consider comparing to https://epubs.siam.org/doi/abs/10.1137/23M1560902, which provides an exact solution. They could compare it also in the stochastic case, to have an idea of a different cost approximation.\n4. The suboptimality of the derived bound is unclear. Can the authors comment on it?\n5. The efficiency benefits are somewhat minor (the algorithms considered are all very fast), so it is hard to tell if the benefits come from algorithmic speedups or from implementation details without a complexity analysis.\n6. The setting for the RNA dataset is not very clear (it is almost never clear in the literature either….). Can you please clarify the objective, how do you evaluate, etc? Also, for the experiments, while I appreciate the numerous experiments I believe that it would be better to defer more to the appendix and discuss more thoroughly less experiments in the main text.\n\nSome other comments:\n\n1. The authors present 6 theorems. While I appreciate the increasing complexity, reversing the order of the theorem would allow to discuss some of the results as corollary and gain some space for more thorough discussions in the main text. It would also allow to avoid essentially duplicated proofs. I would also make a Lemma instead of a Theorem for the first theorem of each “pair” of theorems.\n2. Can you please double check Proposition 1? The expression for the velocity field seems wrong (independent of t). Or could the authors perhaps clarify?\n3. Before equation 11, the authors refer to Problem 6 and 10, but then mention replacing the first constraint only in 6 (typo? perhaps you want to add 10 as well). I would also not use the word sequel as it made me think of a follow up paper.\n4. Perhaps stop the sentence “Problem (24) is challenging to solve in general.” instead of the “,”.\n5. By the way, I find this decoupling approach very interesting. It also relates to a recent work on a similar, yet slightly different setting https://arxiv.org/abs/2406.12616. There the decoupling is somewhat dual (for estimation).", "was_revised": true, "final_justification": "The authors addressed all my concerns and accordingly I think the paper deserves acceptance.", "rebuttal": "We would like to sincerely thank the Reviewer for his/hers positive and in-depth evaluation of our work, for finding our work pleasant to read, and for the multiple sincere suggestions to improve the quality and presentation of our work. We will address all issues raised by the reviewer in detail.\n\n   1. **Suboptimality gap of our method**. As mentioned in Theorems 1, 3, and 5, the mixture policy in equation (13) introduces a family of feasible policies for the SB, which is smaller than the class of all feasible solutions for the SB. In general, the globally optimal solution of the SB might not belong to this family, which introduces some suboptimality with respect to the transport cost, in exchange for numerical tractability. To empirically quantify this optimality gap in our original submission, we provide extensive benchmark testing in Appendix C.4 in Table 6. In practice, we observe that although the mixture policy calculated by Theorems 1, 2 may, in general, be suboptimal, in practice results in better coupling and transport costs, especially for larger values of the noise parameter $\\epsilon$, than most neural SB or EOT methods, due to the non-convex loss functions associated with the corresponding methods. To improve the clarity of our work and strengthen our contribution, **we have added a new result to our paper**, which explores the suboptimality introduced by the upper bound of Theorem 2. This new result shows that when the components of the GMMs are well-separated, the upper bound introduced in Theorem 2 is tight, and the mixture policy is optimal within the feasible set introduced in Theorem 1. To save space and avoid a duplicate response, we refer to point 2) in our response to Reviewer HrsH for the rigorous proof of this statement.\n\n   2. **Comparison with the approach in [1]** We thank the Reviewer for this suggestion. As the reviewer mentions, the cited work is for deterministic, discrete-time problems while our scope is continuous-time deterministic/stochastic problems. Furthermore, in order to apply the result of [1], one needs to be able to explicitly calculate the Wasserstein-2 optimal transport plan between Gaussian Mixture Models in order to implement the policy in [1, Example 5.2], which is not available in closed form for GMMs. \n   Therefore, the only alternative would be to use the empirical W-2 distance using samples. We note that for the 2D examples shown in Figure 1, we have already performed this calculation in Appendix C.1, Table 4 (see first line for the deterministic case).    \n\n\n  3. **RNA experiment setup** We thank the reviewer for this comment. We will present a brief overview of the experiment here and add more details in the main text of the paper. From an optimization perspective, in the RNA example, we solve an instance of Problem (16). The dataset, i.e., the marginal distributions, consists of 5-dimensional distributions at 5 time steps, $t_1, ..., t_5$. In the distribution of each time step, we first fit a 5-component GMM to obtain the GMM parameters of the constraint (16c). To compute $\\lambda_{\\mathbf{i}}$ in Theorem 4, one first needs to calculate $J_{\\mathbf{i}}$, which is the cost of each Gaussian Multi-marginal momentum Schrodinger Bridge (GMSB) of the form (10). While there is no closed solution for the GMSB, it can be efficiently calculated as a Semidefinite Program (see Appendix D for the details). Once the cost tensor $J_{\\mathbf{i}}$ is computed, $\\lambda_{\\mathbf{i}}$ is computed by solving the corresponding Problem (20), which then yields the flow through Equation (18). To evaluate our results, besides visual means through Figures 4 and 7 (see supplementary material), we use the Sliced Wasserstein Distance (SWD) and the Maximum Mean Discrepancy (MMD) between the true marginal distributions and the inferred marginal distributions. To calculate the inferred distributions, we use the first marginal at time $t_1$ as an initial condition for the position variable, i.e., $x_1$, and infer the velocity by sampling from Equation (B.19d) (this is easy since (B.19d) is a GMM). We then integrate the flow in order to sample from the marginals at the rest of the time steps $t_2,...,t_5$. The MMD and SWD both measure distributional discrepancy between the inferred and target distributions in a given time step. We average the calculated SWD and MMD values for the four predicted time steps (we skip the first since it was used for initial conditions) and report the values at Table 3, along with other baselines. \n\n**Other comments**: \n    \n   4. **Computational complexity of our method**. We thank the Reviewer for suggesting adding a theoretical analysis of the computational complexity to our work. The computational complexity of fitting a GMM using the Expectation-Maximization (EM) algorithm scales as $O (I N  K  (D + D^2))$ [1], where $I$ is the number of EM iterations, $N$ is the number of data points, $K$ is the number of Gaussian components (modes), and $D$ is the dimensionality of the data. This complexity highlights that GMM fitting becomes expensive with higher dimensions or a large number of modes. Once the GMMs are fitted, our method solves (LP) problem with $N_0 \\times N_1$ variables, where  $N_0$ and $N_1$ denote the number of modes in the initial and final distribution, respectively. Modern solvers like MOSEK or Gurobi efficiently solve LP problems using interior-point methods, which have a computational complexity of $ O(\\sqrt{l} N^3)$ [2], where  $l$ represents the number of constraints.   This ensures scalability for moderate numbers of modes, but it can become computationally demanding as the number of modes in the GMMs exceeds $1,000$ for each boundary distribution. In practice, in the 2D example of our paper (see Figure 2 and Figure 6), we tested the algorithm in problems with 500-component GMMs and validated that training takes under a minute. \n\nRegarding the computational complexity of inference, each evaluation of the GMMflow policy (Equation 13) scales linearly with the number of components in each mixture and the SDE integration also scales linearly with the number of discretization time steps. In practice, when implementing the GMMflow policy, we only compute a small fraction of conditional GSB policies since the component-level transport plan is very sparse. Moreover, this computation is done in parallel for all conditional policies together. This results in very fast, practically constant-time inference regardless of the number of components in the boundary GMMs or problem dimension. \n\nTo further benchmark our algorithm in practice, we provide wall clock times for both training and inference as a function of the number of components in the boundary distributions (same number for initial and final) and the problem dimensionality. We use the boundary distributions provided in the EOT benchmark for this test. We report these values here: \n\n**Training times in seconds** \n\n | Dim\\ \\# comp | 5      | 10     | 20     | 50     | 100     |\n |------------|--------|--------|--------|--------|---------|\n | 2          | 0.209  | 0.0595 | 0.1083 | 0.2303 | 1.0256  |\n | 16         | 0.0567 | 0.0948 | 0.1413 | 0.4218 | 3.3918  |\n | 64         | 0.182  | 0.2228 | 0.7217 | 1.2875 | 2.1431  |\n | 128        | 0.2802 | 0.4562 | 0.7713 | 1.6164 | 3.4101  |\n\n\n**Inference times in seconds**\n\n | Dim\\ \\# comp | 5     | 10    | 20    | 50    | 100    |\n |------------|-------|-------|-------|-------|--------|\n | 2          | 0.037 | 0.031 | 0.031 | 0.032 | 0.030  |\n | 16         | 0.032 | 0.032 | 0.032 | 0.033 | 0.032  |\n | 64         | 0.032 | 0.032 | 0.032 | 0.032 | 0.039  |\n | 128        | 0.032 | 0.033 | 0.036 | 0.032 | 0.082  |\n\n  5. **Merging theorem proofs** We thank the reviewer for this suggestion. Although we believe that keeping the theorems separate makes the flow of the paper easier to follow, we agree that the proofs of Theorems 1,5, and Theorems 2, 6 are similar. We will consider merging them in the final version of our work for conciseness. The multi-marginal versions of the same theorems are considerably different, which prompts us to keep the proofs separate.\n\n  6. **Proposition 1 is correct** We thank the reviewer for this comment. We confirm that Proposition 1 is correct. The vector field $u_t(x) = K_t(x - \\mu_t) + v_t$ is correct and will, in general, be time varying since $K_t$ and $\\mu_t$ are time-varying. We note that for the Gaussian Schr\\\"odinger Bridge, the feed-forward term $v_t= \\mu_1 - \\mu_0$ is not time varying; however, we are keeping the dependence in time since for more general bridges, such as the Gaussian Momentum SB or the SB with general linear prior dynamics, this term will be time varying. \n\n  7. **Readability and scope** To increase readability and better convey the scope of our work, we have altered the abstract to better highlight our motivation, namely, the efficient solution of low-dimensional SB problems. \n  In the revised version of the manuscript, we have changed the sentence before Equation (6) to read \"the first constraint in (6) or (10b).'' Furthermore, we will refrain from using \"sequel'' to avoid confusion.\n\n  \n[1] Pedregosa, Fabian, et al. ``Scikit-learn: Machine learning in Python.\" Journal of Machine Learning Research 12 (2011): 2825-2830.\n\n[2] Boyd, Stephen P., and Lieven Vandenberghe. Convex optimization. Cambridge University Press, 2004."}, {"review_id": "5P83fsD9Dg", "reviewer": "Reviewer_nHAK", "summary": "The authors provide a specific formulation of optimal transport that is useful in low-dimensional regimes and avoids the computationally complexity suffered by more established models in high-dimensional regimes. In particular, they propose a policy which is theoretically feasible and an associated probability flow for the setting of transportation between mixed Gaussian conditionals, by solving Gaussian bridge sub-problems. They support their theoretical findings with experiments comparing their methodology to others' on toy data, image matching, and single-cell data with their methodology outperforming contemporary methods in every instance.", "questions": "N/A", "limitations": "Yes, the authors have addressed limitations. The other limitations are included in the weaknesses mentioned in the previous paragraph.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "The authors concisely summarize contemporary methods for optimal transport between probability distributions and describe the gap their paper intends to close. They provide theoretical evidence supporting their proposed transport plan and make a clear delineation between existing methods and what their method offers computationally.  The empirical results are strong; their proposed transport plan outperforms contemporary methods across a diverse set of experiments across different metrics. A potential improvement, as mentioned by the authors, is to not constrain the proposed methodology to the setting of Gaussian Mixture Models. It could also be worth empirically exploring around what dimension the proposed model loses its efficacy through ablation studies.", "was_revised": true, "final_justification": "The authors have sufficiently defended their work and the paper is deserving of acceptance.", "rebuttal": "We sincerely thank the Reviewer for the positive evaluation of our work and for the suggestion to further test our algorithm through ablation studies to study how it scales in practice on problems with many components. To address this suggestion, we performed a run-time analysis for training and inference with respect to the problem dimension and nubmer of GMM components. To avoid duplicating the response, we refer the Reviewer to point 4. in our response to Reviewer PeSU."}, {"review_id": "RmlW6t7QpH", "reviewer": "Reviewer_kjHY", "summary": "The authors address the Schrödinger Bridge (SB) problem for Gaussian mixture endpoints (SB-GM). They observe that the SB between two individual Gaussians (SB-G) admits a closed-form solution, and hence propose to parametrize the feasible SB-GM solutions as a weighted mixture of SB-G solutions. By doing so, they mitigate the original optimization problem and optimize only over the mixture weights to receive the feasible solution. They the authors define a linear optimization problem for optimization over the mixture weights and rigorously show that the linear optimization problem objective upper-bounds the true optimal cost for the SB-GM. Building on this insight, the authors extend their framework to both the multi-marginal momentum SB and to SBs between infinite Gaussian mixtures.\n\nIn practice, they approximate arbitrary distribution endpoints $p_0$ and $p_1$ by finite Gaussian mixtures, construct SB-G solutions for every pair of mixture components, and recover the approximation of overall SB-GM by solving a linear optimization problem over mixture weights.  An analogous scheme is developed for the multi-marginal momentum SB. The method is validated on 2D toy problem, Image-to-Image translation in the latent space of ALAE and learning dynamics of cell populations.", "questions": "- Is the optimal solution for the SB problem with gaussian mixture endpoints lies in the set of solutions defined in Theorems 1, 3, 5?\n- Is the $J_{\\text{OT}}$ in Theorems 2, 4, 6 transport cost for the corresponding $\\lambda$ defined solutions?\n- Are there any results on tightness of upper bounds in Theorems 2,4,6?\n- What is the computational complexity of the linear optimization algorithms used, or that could be potentially used? Including worst cases. Is there any results on training time and method quality w.r.t. number of gaussian mixtures used for approximation (preferably on EOT bench)?\n- Is there results on transport cost for the ALAE experiment?\n- Why do you have relatively high SWD and very low MMD on cell population dynamics learning experiment?\n- How do you apply NLSB [5] to multi-marginal SB problem in the cell population dynamics learning experiment? As far as I understand NLSB is general SB method. If you consider the general SB method, then I think you also should consider methods [1, 2, 3, 4].\n- What potential algorithms can be considered for solving the stated in Theorem 6 optimization problem?", "limitations": "Mostly yes, but I would like authors to clarify the complexity of linear optimization methods they use w.r.t. number of components in Gaussian Mixtures. Depending on the answer, there is a possibility one could consider it a limitation.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n\n- The proposed algorithm is both novel and computationally efficient, demonstrating strong performance in low-dimensional settings.\n- The authors’ approximation strategies for diverse SB formulations—including the multi-marginal momentum SB and the infinite-mixture SB—represent valuable conceptual contributions.\n- Empirical results on image-to-image translation, EOT benchmark (Appendix C)  and cell-population dynamics modeling show that GMMflow mostly matches state-of-the-art methods.\n- The proposed algorithm tries to advance the area of SB methods specifically crafted for low dimensional cases.\n\n**Weaknesses**:\n\n- Authors do not provide guarantees for a mixture of SB-G solutions to reach the optimal solution for the SB-GM problem. In general, only the upper bound on the optimal transport cost can be received that way. For the general SB the approximation of its distribution endpoints by the gaussian mixtures is already an approximation and the mixture of SB-G solutions adds a second approximation on top.\n    - Would be nice to know more about the solutions gotten as the result of solving linear optimization problems in Theorems 2, 4, 6 and their relation to the optimal ones. Maybe there is some cases where you can find exactly the optimal solutions. For other related questions see **Questions** section.\n- The number of variables in the linear optimization problem (mixture weights) grows linearly with the number of components in Gaussian Mixtures used for approximation of distributions. In the two distribution case number of variables is: $(M_1 \\times M_2)$, or in the multi-marginal case  $(M_1 \\times M_2 … \\times M_k)$. This seems like a lot of mixture weights, especially when one considers complex distributions that require a lot of Gaussian Mixture components to be approximated accurately. Additionally one has to solve the linear optimization problem, with its own computational complexity, to optimize over these mixture weights.\n    - Would like to see a study on the growth of Gaussian Mixtures components along with  training time and method quality.\n- As far as I understand the inference of the SB-G mixture solution requires the simulation of the SDE. Which is rather unconvenient, especially since LSB/LSBM and Gaussian to Gaussian SB solution transport plans can be inferred without the SDE simulation.\n- The evaluation:\n    - In the Image-to-Image translation experiment only the target distribution matching  metrics (FID, ALAE-BW) are shown. Which is weird since the target matching is guaranteed, up to GMM approximation error, by the construction of the method. On the other hand, the optimality of transport cost is not guaranteed. Notably by the Figure 3 one can see that LightSB transport plan has lower transport cost.\n    - The only experiment that quantitatively evaluates the transport cost is EOT benchmark shown in Appendix C. The performance of method becomes better with growth of $\\epsilon$, while the lower $\\epsilon$ examples do mostly evaluate the transport cost of the method and higher $\\epsilon$ examples do mostly evaluate the ability of method to fit the ending ($p_1$) distribution. This experiment results can be seen as proof that the proposed method is bad as transport cost but good at distribution fitting.\n    - In the EOT benchmark experiment some baselines are missing: DSBM [1], SF2M-sink [2], LightSB [3], LightSBM [4]. All the metrics are already computed, see Table 1 [4].\n\n**Typos/Minor text concerns**:\n\n- Line 128: $p_i, p_f$ should it be $p_0, p_1$?\n- Line 298, 300: Bures-Wasserstein distance\n- Line 332, 337, Table 3: NSBL → NLSB\n\n[1] Shi, Y., De Bortoli, V., Campbell, A., & Doucet, A. (2023). Diffusion schrödinger bridge matching. *Advances in Neural Information Processing Systems*, *36*, 62183-62223.\n\n[2] Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., ... & Bengio, Y. (2024, April). Simulation-Free Schrödinger Bridges via Score and Flow Matching. In *International Conference on Artificial Intelligence and Statistics*(pp. 1279-1287). PMLR.\n\n[3] Korotin, A., Gushchin, N., & Burnaev, E. Light Schrödinger Bridge. In *The Twelfth International Conference on Learning Representations*.\n\n[4] Gushchin, N., Kholkin, S., Burnaev, E., & Korotin, A. (2024, July). Light and optimal schrödinger bridge matching. In *Forty-first International Conference on Machine Learning*.\n\n[5] Koshizuka, T., & Sato, I. Neural Lagrangian Schr\\\"{o} dinger Bridge: Diffusion Modeling for Population Dynamics. In *The Eleventh International Conference on Learning Representations*.", "was_revised": true, "final_justification": "The authors have adressed my concerns.", "rebuttal": "We thank the reviewer for the careful consideration of our work, for finding our results novel and valuable conceptual contributions, and for recognising the computational efficiency of our method. We will address all the comments raised by the reviewer in detail. If all the Reviewer's concerns are well addressed, we would appreciate a positive reflection in his/her's score. \n\n1. **Optimality of our approach and tightness of the upper bound in Theorem 2**: As mentioned in Theorems 1, 3, and 5, the mixture policy in equation (13) introduces a family of **feasible** policies for the SB, which is smaller than the class of all feasible solutions for the SB. In general, the globally optimal solution of the SB might not belong to this family, which introduces some suboptimality with respect to the transport cost, in exchange for numerical tractability. To empirically quantify this optimality gap in our original submission, we provide extensive benchmark testing in Appendix C.4 in Table 6. In practice, we observe that although the mixture policy calculated by Theorems 1, 2 may, in general, be suboptimal, in practice results in better coupling and transport costs, especially for larger values of the noise parameter $\\epsilon$, than most neural SB or EOT methods, due to the non-convex loss functions associated with the corresponding methods. To improve the clarity of our work and strengthen our contribution, **we have added a new result to our paper**, which explores the suboptimality introduced by the upper bound of Theorem 2. This new result shows that when the components of the GMMs are well-separated, the upper bound introduced in Theorem 2 is tight, and the mixture policy is optimal within the feasible set introduced in Theorem 1. Due to space limitations, and to avoid a duplicate response, we refer to point 2. in our response to Reviewer HrsH for the rigorous proof of this statement.\n\n2. **$J_{OT}$ and transport cost of $\\lambda$ policy** As Theorem 2 suggests, the cost $J_{OT}$ is an upper bound of the true transport cost of the mixture policy (13). This bound becomes tight when the components of the solution are well-separated. In general, using the mixture policy is guaranteed to yield a cost that is at least smaller than the upper bound $J_{OT}$. This can also be validated experimentally, for example, through Table 4 (see columns 1st and 2nd columns for the upper bound $J_{OT}$ and the true cost $J_{GMM}$ respectively).  \n\n3. **Computational complexity of solving Linear Programs (LPs)** We thank the Reviewer for this comment. Modern solvers like MOSEK or Gurobi efficiently solve LP problems using interior-point methods, which have a computational complexity of $ O(\\sqrt{l} N^3)$ [6], where  $N$ is the number of decision variables and $l$ represents the number of constraints. In practice, this allows solving LPs with millions of variables easily. In practice, in the 2D example of our paper (see Figure 2 and Figure 6), we tested the algorithm in problems with 500-component GMMs (250,000 decision variables in the corresponding LP),  and validated that training takes under a minute for both GMM fitting and solving the LP. For a full complexity analysis of our method, we refer the Reviewer item 4. in our response to Reviewer PeSU. \n\nTo further benchmark our algorithm in practice, we provide wall clock times for both training and inference with respect to the number of components and the problem dimensionality for the boundary distributions provided in the EOT benchmark. We report these values here: \n\n **Training times in seconds** \n\n | Dim\\ \\# comp | 5      | 10     | 20     | 50     | 100     |\n |------------|--------|--------|--------|--------|---------|\n | 2          | 0.209  | 0.0595 | 0.1083 | 0.2303 | 1.0256  |\n | 16         | 0.0567 | 0.0948 | 0.1413 | 0.4218 | 3.3918  |\n | 64         | 0.182  | 0.2228 | 0.7217 | 1.2875 | 2.1431  |\n | 128        | 0.2802 | 0.4562 | 0.7713 | 1.6164 | 3.4101  |\n\n **Inference times in seconds**\n\n | Dim\\ \\# comp | 5     | 10    | 20    | 50    | 100    |\n |------------|-------|-------|-------|-------|--------|\n | 2          | 0.037 | 0.031 | 0.031 | 0.032 | 0.030  |\n | 16         | 0.032 | 0.032 | 0.032 | 0.033 | 0.032  |\n | 64         | 0.032 | 0.032 | 0.032 | 0.032 | 0.039  |\n | 128        | 0.032 | 0.033 | 0.036 | 0.032 | 0.082  |\n\nWe note that when implementing the GMMflow policy in practice, we only compute a small fraction of conditional GSB policies since the component-level transport plan is very sparse. Moreover, this computation is done in parallel for all conditional policies together. This results in very fast, practically constant-time inference regardless of the component number or problem dimension.\n\nFinally, we note that because the EOT benchmark uses an initial Gaussian distribution and a GMM terminal distribution, there is no point in reporting metrics such as marginal distribution accuracy or transport plan optimality as a function of the number of GMM components, since these will perform best when the number of components used in GMMflow matches the setting of EOT benchmark. Furthermore, exploring how well a GMM approximates a general distribution as the number of components increases is a well-studied problem and goes beyond the scope of our work, so due to time constraints, we will not provide experiments that explore this issue. \n\n4. **Evaluation of the transport cost for the ALAE experiment** To further evaluate the transport cost of our method, we will calculate the empirical Wasserstein transport cost of our method and compare it with LSB and LSBM. More specifically, we will report the values of $\\big( \\int\\|x_0-x_1\\|^2 \\mathrm{d} \\pi(x_0, x_1) \\big)^{\\frac{1}{2}}$ for different levels of the regularization constant $\\epsilon$, using the transport plan from GMMflow, LSB and LSBM. We used 10,000 samples to calculate the empirical expectation. \n\n |M2W             | GMMflow | LSB  | LSBM   |\n |----------------|---------|------|--------|\n |$\\epsilon=0.1$  | 11.14   |10.47 |\t10.44|\n |$\\epsilon=0.01$ | 9.05    |8.23  | 8.18   |\n |$\\epsilon=0$    | 8.75    |-     | -      |\n\n |A2C             | GMMflow | LSB  | LSBM   |\n |----------------|---------|------|--------|\n |$\\\\epsilon=0.1$  |11.41    |10.45\t| 10.43  |\n |$\\epsilon=0.01$ |9.33\t |8.19\t| 8.18   |  \n |$\\epsilon=0$    | 9.04    |  -   |  -     |\n\nWe note that although this experiment agrees with the intuition of the Reviewer, namely that GMMflow results in higher transport costs than LSB/LSBM, since the difference is small, it is hard to attribute this to any potential suboptimality introduced by the upper bound of Theorem 2 or the restriction of the SB policy in the feasible set introduce by Theorem 1. Specifically, the 40\\% lower FID and 70\\% lower ALAE-BW metrics in Table 1 of the paper indicate that GMMflow learns considerably different marginal distributions than LSB/LSBM. Therefore, it is not clear whether the increased transport costs are due to the different marginal distributions or the suboptimal transport plan.\n\n5. **High SWD and very low MMD**: We thank the reviewer for this observation. We attribute this to the nature of using GMMs to capture complicated distributions. The very low MMD indicates the accurate fitting of GMMs, since the MMD metric is moment-based and captures the low-order information in the distributions. On the other hand, since the true distributions are not GMMs, the SWD metric shows that some higher-order information is lost. \n\n6. **NLSB algorithm in the RNA example**. The Neural Lagrangian Schr\\\"odinger Bridge (NLSB) solver is more complicated than a standard SB solver since the cost function it optimizes contains density-dependent and velocity regularization terms. We refer the Reviewer to [5] for the NLSB implementation details in the cellular dynamics example. We argue that existing implementations of methods [1-4] are not directly applicable to the multimarginal setting unless the problem is solved pairwise for any two consecutive distributions, reducing the problem to a series of standard, two-marginal SBs. \n\n7. **Infinte component GMM computation**. While solving the optimization problem (24) in Theorem 6 is challenging, in general, the scope of our work focuses on the special case where the parameter spaces $w_0, w_1$ are 1-dimensional. As detailed in Appendix B of our supplementary material, this contains multiple important classes of problems, such as the computation of the Wasserstein distance between multivariate Student-t distributions. In this setting, Problem (24) can be solved in closed form using Proposition 3 from the supplementary material, yielding a comprehensive upper bound in the Wasserstein distance between Student-t distributions. \n\n8. **Minor Typos** We thank the reviewer for pointing these out. We will correct these typos in the final version of our work. \n\n\n[1] Shi, Y., De Bortoli, V., Campbell, A., \\& Doucet, A. (2023). Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 62183-62223.\n\n[2] Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., ... \\& Bengio, Y. (2024, April). Simulation-Free Schrödinger Bridges via Score and Flow Matching. In International Conference on Artificial Intelligence and Statistics(pp. 1279-1287). PMLR.\n\n[3] Korotin, A., Gushchin, N., \\& Burnaev, E. Light Schrödinger Bridge. In The Twelfth International Conference on Learning Representations.\n\n[4] Gushchin, N., Kholkin, S., Burnaev, E., \\& Korotin, A. (2024, July). Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning.\n\n[5] Koshizuka, T., \\& Sato, I. Neural Lagrangian Schr\"{o} dinger Bridge: Diffusion Modeling for Population Dynamics. In The Eleventh International Conference on Learning Representations.\n\n[6] Pedregosa, Fabian, et al. “Scikit-learn: Machine learning in Python.” Journal of Machine Learning\nResearch 12 (2011): 2825-2830"}, {"review_id": "Eb3ITwweqb", "reviewer": "Reviewer_VvBB", "summary": "The authors proposes a training‑free method to solve Schrödinger Bridge problems when both endpoints are Gaussian Mixture Models (GMMs). The method expresses a global policy as a weighted mixture of closed‑form gaussian bridges, and chooses weights via a low‑dimensional linear program whose size grows with the number of mixture components. The paper gives convincing numerical results on a set of toy and real world problems including image-to-image translation.", "questions": "- The results for image-to-image translation tasks are compelling. Moreover it is surprising that you can get away with only a 10-component diagonal GMM. Do the authors imagine you do full image generation on some constrained datasets in a latent space? At what point does the method break down? It would be interesting to see the GM assumptions pushed to its limits.\n\n*Possible action items (see weaknesses above)*\n- Runtime analysis in terms of number of components, number samples, etc. (what ever is most relevant) for both training and inference.\n- Algorithm block \n- Error vs number of mixture components.\n\n(I am not insisting that all of these must be done, I understand time constraints, but I do believe they would improve the paper.)", "limitations": "- The author discuss, but obvious limit is the Gaussian assumption.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "*Strengths*\n- The paper is clear and easy to read.\n- Training-free methods and Gaussian Mixture Models are probably under-rated in a world dominated by expensive generative methods. I like to see this kind of paper. \n- The paper server as a good introduction to the area\n- The derivation of the method is clear and the theory is substantial.\n- The results for image-to-image translation tasks are compelling.\n\n*Weaknesses*\n- The clarity of the paper would be enhanced with an Algorithm block describing the full implementation. \n- It would be valuable to know more about the sensitivity when boundary data are poorly approximated by a GMM. It would be great to see an experiment with error on the y-axis and number of GMM components on the x-axis for some of the examples. Is most of the error induced by fitting the marginals, or even if the marginal are well approximated by the GMM can other sources of error dominate?\n- The paper is sold as a \"fast\" method yet there is no runtime analysis, empirical or theoretical. This type of analysis would shed light on key questions such as how does the runtime scale with dimension and number of samples etc. Similarly an experiment with runtime on the y-axis and components on the x-axis would be useful.", "was_revised": true, "final_justification": "The work is strong and well presented. The approach is reasonable. A good edition to the conference.", "rebuttal": "We would like to sincerely thank the Reviewer for their thoughtful and constructive feedback. We are especially grateful for the recognition of the clarity of our presentation and the potential of training-free approaches.\nWe will address all points raised by the Reviewer in detail: \n\n\n1) **Regarding the use of 10-component diagonal GMMs in image-to-image translation tasks**: We appreciate the reviewer’s observation and insightful question. In our image-to-image translation experiments, both the input and output images are first mapped from their high-dimensional pixel space to a low-dimensional latent space of dimension 512 using a pretrained encoder. The GMM fitting and flow matching are then performed entirely in this latent space, where the distributions tend to be significantly smoother and more structured than in the original pixel domain. This dimensionality reduction, along with the regularity of the learned latent representations, helps explain why a relatively small number of diagonal components (e.g., 10) is sufficient to capture the essential modes of the data.\n\n\n2) **Runtime analysis**  We thank the Reviewer for suggesting we add a theoretical analysis of the computational complexity of our algorithm to our work. The computational complexity of fitting a GMM using the Expectation-Maximization (EM) algorithm scales as $O (I N  K  (D + D^2))$ [1], where $I$ is the number of EM iterations, $N$ is the number of data points, $K$ is the number of Gaussian components (modes), and $D$ is the dimensionality of the data. This complexity highlights that GMM fitting becomes expensive with higher dimensions or a large number of modes. Once the GMMs are fitted, our method solves (LP) problem with $N_0 \\times N_1$ variables, where  $N_0$ and $N_1$ denote the number of modes in the initial and final distribution, respectively. Modern solvers like MOSEK or Gurobi efficiently solve LP problems using interior-point methods, which have a computational complexity of $ O(\\sqrt{l} N^3)$ [2], where  $l$ represents the number of constraints. This ensures scalability for moderate numbers of modes, but it can become computationally demanding as the number of modes in the GMMs exceeds $1,000$ for each boundary distribution. In practice, in the 2D example of our paper (see Figure 2 and Figure 6), we tested the algorithm in problems with 500-component GMMs and validated that training takes under a minute. \n\nRegarding the computational complexity of inference, each evaluation of the GMMflow policy (Equation 13) scales linearly with the number of components in each mixture and the SDE integration also scales linearly with the number of discretization time steps. In practice, when implementing the GMMflow policy, we only compute a small fraction of conditional GSB policies since the component-level transport plan is very sparse. Moreover, this computation is done in parallel for all conditional policies together. This results in very fast, practically constant-time inference regardless of the number of components in the boundary GMMs or problem dimension. \n\nTo further benchmark our algorithm in practice, we provide wall clock times for both training and inference as a function of the number of components in the boundary distributions (same number for initial and final) and the problem dimensionality. We use the boundary distributions provided in the EOT benchmark for this test. We report these values here: \n\n**Training times in seconds** \n\n | Dim\\ \\# comp | 5      | 10     | 20     | 50     | 100     |\n |------------|--------|--------|--------|--------|---------|\n | 2          | 0.209  | 0.0595 | 0.1083 | 0.2303 | 1.0256  |\n | 16         | 0.0567 | 0.0948 | 0.1413 | 0.4218 | 3.3918  |\n | 64         | 0.182  | 0.2228 | 0.7217 | 1.2875 | 2.1431  |\n | 128        | 0.2802 | 0.4562 | 0.7713 | 1.6164 | 3.4101  |\n\n\n**Inference times in seconds**\n\n | Dim\\ \\# comp | 5     | 10    | 20    | 50    | 100    |\n |------------|-------|-------|-------|-------|--------|\n | 2          | 0.037 | 0.031 | 0.031 | 0.032 | 0.030  |\n | 16         | 0.032 | 0.032 | 0.032 | 0.033 | 0.032  |\n | 64         | 0.032 | 0.032 | 0.032 | 0.032 | 0.039  |\n | 128        | 0.032 | 0.033 | 0.036 | 0.032 | 0.082  |\n\n3) **Algorithmic block** In response to the suggestion to improve clarity with an algorithm block, we have now added the following algorithm, describing the full implementation,  to the revised version of the paper.\n\n\n**Go With the Flow Algorithm**\n\n**Input** Initial dataset $\\\\rho_0$, final dataset $\\rho_1$; number of GMM components $N_0$ and $N_1$.\n\n**Step 1:** $ \\\\{\\alpha_0^{i}, \\mu_0^i, \\Sigma_0^i \\\\}_{i=1}^{N_0}\\leftarrow$EM$(\\rho_0, N_0)$ (Fits a GMM to initial dataset)\n\n**Step 2:** $ \\\\{\\alpha_1^{j}, \\mu_1^j, \\Sigma_1^j \\\\}_{j=1}^{N_1}\\leftarrow$ EM$(\\rho_1, N_1)$ (Fits a GMM to terminal dataset)\n    \n**Step 3:** $\\forall (i,j)\\in N_0 \\times N_1$, \n    $\\\\{u_{t|ij}, \\rho_{t|ij}, J_{ij}\\\\}  \\leftarrow$ CS $((\\mu_0^i, \\Sigma_0^i),( \\mu_1^j, \\Sigma_1^j))$ (solves the $(i,j)$-th conditional GSB)\n    \n**Step 4:** $\\lambda_{ij} \\leftarrow $  solve (15) using $ \\\\{J_{ij}, \\alpha_0^{i}, \\alpha_1^{j}\\\\}$ \n\n\n**Step 5:**  Return $ \\\\{u_{t|ij}, \\rho_{t|ij}, \\lambda_{ij} \\\\}$\n    \n\n4) **Sensitivity of GMMs**: We agree that understanding the sensitivity of our method with respect to the quality of GMM approximation is an important direction. Although we have not conducted a dedicated sensitivity analysis experiment, our empirical results (e.g., in image translation tasks) suggest that relatively low-component GMMs can still produce compelling results, indicating a good degree of robustness. That said, we acknowledge that performance may degrade under severe underfitting, and a more systematic study of this sensitivity is a valuable avenue for future work.\n\n\n[1] Pedregosa, Fabian, et al. “Scikit-learn: Machine learning in Python.” Journal of Machine Learning\nResearch 12 (2011): 2825-2830."}]

Author rebuttal:

Meta-review: The paper considers the Schrödinger Bridge (SB) problem in two setups: bi-marginal and momentum multi-marginal. For each setting, the authors propose a lightweight algorithm (in the spirit of LightSB work [1]) to obtain an approximate SB solution. The main idea is to approximate marginals with Gaussian Mixture Models (GMMs) and exploit the closed-form SB solutions available between individual Gaussian components.

In the bi-marginal case, the authors first construct a process that is a weighted mixture of pairwise SBs between the components of the two marginals (the weights are optimized). They then construct an "average" process represented by an SDE with an analytically computable drift, made possible by the Gaussian and Gaussian Mixture structure. This SDE is proposed as an approximate SB solution between the mixtures. The multi-marginal case is completely analogous but is more technically complex due to the increased number of marginals.

Overall, the reviewers highly rated the results from the paper: the core idea is interesting, the extension to the multi-marginal case is inspiring (as such extensions are rare), and the resulting algorithm is efficient for low-dimensional problems where GMMs are effective (as demonstrated in experiments). A clear limitation, however, is the method's lack of scalability to high dimensions due to its reliance on GMMs. Furthermore, as noted by several reviewers, there was an initial lack of a clear theoretical understanding of how close the proposed solution is to the true SB. Despite these limitations, the reviewers and meta-reviewer find the paper to be technically solid, original, and well-executed. Its strengths are deemed to outweigh its limitations, and acceptance is recommended.

**P.S.** From the perspective of related work [2], it seems that the proposed algorithm can be viewed as a Markovian projection of a special reciprocal process constructed as a mixture of SBs. 

[1] Korotin, A., Gushchin, N., & Burnaev, E. Light Schrödinger Bridge. In The Twelfth International Conference on Learning Representations.

[2] Shi, Y., De Bortoli, V., Campbell, A., & Doucet, A. (2023). Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 62183-62223.
Decision: Accept (spotlight)