Light Schrödinger Bridge
Authors: Alexander Korotin
Nikita Gushchin
Evgeny Burnaev
Venue: ICLR 2024 poster (2024)
Keywords: Schrödinger Bridge
Optimal transport
Entropy regularized OT
Unpaired Learning
Primary area: generative models

Abstract: Despite the recent advances in the field of computational Schrödinger Bridges (SB), most existing SB solvers are still heavy-weighted and require complex optimization of several neural networks. It turns out that there is no principal solver which plays the role of simple-yet-effective baseline for SB just like, e.g., $k$-means method in clustering, logistic regression in classification or Sinkhorn algorithm in discrete optimal transport. We address this issue and propose a novel fast and simple SB solver. Our development is a smart combination of two ideas which recently appeared in the field: (a) parameterization of the Schrödinger potentials with sum-exp quadratic functions and (b) viewing the log-Schrödinger potentials as the energy functions. We show that combined together these ideas yield a lightweight, simulation-free and theoretically justified SB solver with a simple straightforward optimization objective. As a result, it allows solving SB in moderate dimensions in a matter of minutes on CPU without a painful hyperparameter selection. Our light solver resembles the Gaussian mixture model which is widely used for density estimation. Inspired by this similarity, we also prove an important theoretical result showing that our light solver is a universal approximator of SBs. Furthemore, we conduct the analysis of the generalization error of our light solver. The code for our solver can be found at https://github.com/ngushchin/LightSB.

Published as a conference paper at ICLR 2024
LIGHT SCHRÖDINGER BRIDGE
Alexander Korotin' Nikita Gushchin", " Evgeny Burnaev
Skolkovo Institute Science and Technology, "Artificial Intelligence Research Institute k.koorott serroti@sskoltene ru, n.gushchiinskolteech.
ABSTRACT
Despite the recent advances in the field of computational Schrödinger Bridges (SB), most existing SB solvers are still heavy-weighted and require complex op- timization of several neural networks. It turns out that there is principal solver which plays the role of simple-yet-effective baseline for SB just like, e.g.. k- method in clustering, logistic regression in classification gorithm discrete optimal transport. We address this issue and propose a novel fast and simple SB solver. Our development which recently appeared in the field: (a) parameterization of the Schrödinger po- tentials with sum-exp quadratic functions and (b) viewing the SegrSdirde potentials as the energy functions. We show that combined together these ideas yield lightweight, simulation-free and theoretically justified SB solver with a simple straightforward optimization objective. As result, it allows solving SB in moderate dimensions in matter of minutes on CPU without a painful hy- perparameter selection. Our light solver resembles the Gaussian mixture model which widely used for density estimation. Inspired by this similarity, we also prove important theoretical result showing that our light solver is uni- versal approximator of SBs. Furthemore, we conduct the analysis of the gen- eralization error of light solver. The code for our solver can be found https: // //g//gtttttt github. com wwwp///wwcccnnnnnnnn..o.. hhtghhhhhiiiiggittttt
Sinkhorn al-
smart combination of two ideas
Figure Unpaired male >> female translation by our LightSB solver applied in the latent space of ALAE for 1024x 1024 FFHQ images. Our LightSB converges 4 cpu cores in less than I minute.
INTRODUCTION
Over the last several years, there has been a considerable progress in developing the computational
approaches for solving the Schrödinger Bridge problem (Schrödinger 1931 1932 SB), which is
also known as the dynamic version of the Entropic Optimal Transport (Cuturi 2013 EOT) problem. The SB problem requires finding the diffusion process between two given distributions that is max- imally similar to some given prior process. In turn, EOT problem is a simplification in which user interested only in the joint marginal distribution on the first and the last steps the diffusion.
"qqual contribution.

---

Published as a conference paper at ICLR 2024
Historically, the majority of studies in EOT/SB sub-field of machine learning concentrated around studying EOT between discrete probability distributions, see (Peyré et al.) 2019) for a sur- vey. Inspired by the recent developments in the field of generative modeling via diffusions (Ho) jet al. 2020; |Rombach et al. 2022) and the diffusion-related nature of SB, various researches started developing solvers for continuous setups of EOT/SB,
(Gushchin et al. (023b) for survey. This implies that a learner has only a sample access to the (continuous) distributions and based on has recover the entire SB process (e.g., its drift) between the entire distributions. Such solvers are straightforwardly applicable to image generation (Wang et al. 2021 De Bortoli et al. [222]] and image-to image translation tasks (Gushchin et al. 2023a) as well as to important smaller scale data transfer tasks, e.g., with the biological data (Vargas et al. 2021 Koshizuka & Sato 2022).
Almost all existing EOT/SB solvers (see $4) have complex neural networks parameterization and
many hyper-parametert they expectedly require time-consuming training/inference procedures. Al- though this unavoidable in large-scale generative modeling tasks, these techniques look too heavy- weighted when user just wants to learn EOT/SB between some moderate-dimmensooaal data distri- butions, e.g., those appearing in perspective biological applications OT Tong et al.) 2020; Vargas et al. 202 Koshizuka & Sato 2022, Bunne et al., 2022 2023 Tong et al. 2023). In this paper, address this issue and present the following main contributions:
1. We propose a novel light solver for continuous SB with the Wiener prior, i.e., EOT with the quadratic transport cost. Our solver has a straightforward non non-minimax learning objective and uses the Gaussian mixture parameterization for the EOT/SB ($3.1 3.2). This development allows us to solve EOT/SB between distributions in moderate dimensions in a matter of minutes thanks to avoiding time-consuming max-min optimization, simulation of the full process trajectories, iterative learning and MCMC techniques which are in use in existing continuous solvers ($4)
2. We show that our novel light solver provably satisfies the universal approximation property for EOT/SB between the distributions supported on compact sets (33.3).
3. We derive the finite sample learning guarantees for our solver. We show that the estimation error vanishes at the standard parametric rate with the increase of the sample size (Appendix (A).
4. We demonstrate the performance of our light solver in a series of synthetic and real-data experi- ments 5
including the ones with the real biological data (35.) considered in related works.
Our light solver exploits the recent advances in the field of HOT/SB, namely, using the log-sum- exp quadratic functions parameterize Schrödinger potentials for constructing EOT/SB benchmark distributions (Gushchin et al.) (023b) and viewing EOT as the energy-based model (Mokrov et al., 2024) which optimizes certain Hullback-Leible divergence. We discuss this in $3.1
BACKGROUND: SCHRÖDINGER BRIDGES
In this section, we recall the main properties of the Schrödinger Bridge (SB) problem with the Wiener prior. We begin with discussing the Entropic Optimal Transport (Cuturi 2013: Genevay] 2019), which is known as the static SB formulation. Next, we recall the dynamic Schrödinger bridge formulation and its relation to EOT (Léonard 2013; Chen et al. 2016).. Finally, we summarize the aspects of the practical setup for learning EOT/SB which we consider throughout the paper.
We work in the D-dimensional Euclidean space (IR Il- [I). We use PR to denote the set of absolutely continuous Borel probability distributions on RD which have finite second moment. For any p E (R. we write p(z) to denote its density at a point x E RR."" In what follows, KL is a short notation for the well-celebrated Kullback-Leible divergence, and (2 SS) is the density at a point x E IR." of the normal distribution with mean r E R2D and covariance 0 x
IR DXD.
Entropic Optimal Transport (EOT) with the quadratic cost. Assume that Po E P2)) P1 E (2) have finite entropy. For e > 0, the EOT problem between them is to find the minimizer of
min xell(po.p1)
-00 21 21 (2o, 21 ddzddz + eKL (T x p1)
(1)
2
RD R
where I[ppo P1) is the set of probability distributions (transport plans) on R D. x RR whose marginals po and p1. respectively, and po P1 is the product distribution. KL in (T) is assumed to be equal to +oo if π is not absolutely continuous. Therefore, one can consider only absolutely continuous म.
The minimizer x² of (T) exists; it is unique and called the EOT plan.
2

---

Published as a conference paper at ICLR 2024
(Dynamic) Schrödinger Bridge with the Wiener prior. We employ Ω as the space of R -aluued functions of time t E [0, 1] describing trajectories in R² which start at time t = 0 and end at t 1. In turn, we use P(S2) to denote the set of probability distributions on L2, i.., stochastic processes. We use dW, to denote the differential of the standard Wiener process. For process T & P(2), we denote its joint distribution at / 0, by I E P(RDD IR.)). In turn, we use T T2a,2, to denote the distribution of T for t (0, 1) conditioned on T's values To. 21 t 0.
Let W E P(2) denote the Wiener process with volatility e > which starts at Pu at = 0. Its differential satisfies the stochastic differential equation (SDE): dW; VedW,. The SB problem with the Wiener prior We between po, Pr is to minimize the following objective:
min KL (T])W))
(Po(P)))))
where F(po, P1) C. P(2)) is the subset of stochastic processes which start at distribution po (at the time t = 0) and end at P1 (at f = 1). This problem has the unigue solution which is a diffusion process T described by the SDE: dZ, 9" (Z,, t) t) + dW; (ééonadd 2013 Prop. 2.3). The process 7 called the Schrödinger Bridge and g" RD x 1] R is called the optimal drift.
Equivalence between EOT and SB. It is known that the solutions of EOT and SB are related to
each other: for the EOT plan n" and the SB process 7 it holds that x¹ = -=" Hence, solution π to EOT I can be viewed as a part of the solution T to SB (2). What remains uncovered by EOT in SB is the conditional process |=0.- Fortunately, it is known that Ti = W sq.2, i.e.,
simply matches the "inner part" of the Wiener prior W which is the well- Brownian Bridge (Pinsky& & Karlin 2011 Sec. 8.3.3). Hence, one may treat EOT and SB as equivalent problems.
Characterization of solutions (Léonard 2013). The EOT plan ' = =x"
nT has a specific form
π T=' 21) x0 (zx) exp ( I| 12/22) 00 (x1),
R.D. + R4 are measurable functions called the Schrödinger potentials. They are defined up to multiplicative constants. The optimal drift g of SB can be derived from o*:
where 15*
gx, cv log /
N(x' (1 ))* (z')dz'.
(4)
To use 4 it sufficies to know 6² up to the multiplicative constant.
Computational Schrödinger Bridge setup. Even though SB/EOT have many useful theoretical properties, solving the SB/EOT problems remains challenging in practice. Analytical solution is
available only for the Gaussian case (Chen al. 2015 Janati al. 2020 Mallasto et al.) 2022
Bunne et al., (023) plus for some manually constructed benchmark pairs of distributions Po.
(Gushchin et al. 2023b). Moreover, in real world setting where SBs are applied (Vargas al. 2021
Bunne et al.) 2023 Tong et al.) (2023), distributions po and p1 are almost never available explicitly but only through their empirical samples. For the rigor of the exposition, below we formalize the typical EOT/SB learning setup which we consider in our paper.
We assume that the learner has access to empirical samples X = (arg- "....,,.. ~ Po and XM
of of ... Af P1 from the (unknown) data distributions Po- P1 P2. (RPP.. These samples (train data) are assumed to be independent. The task of the learner is to recover the solution (process T or plan " to SB problem (2) between the entire underlying distributions Po, Pl.
The setup above is the learning from empirical samples and is usually called the continuous OT. In the continuous setup, it is essential to be able to do the out-of sample estimation, e.g., simulate the SB process trajectories starting from new (test) points snew ~ Po and ending at P1. In some practical use cases of SB, e.g.. generative modeling (De Bortoli et al.) (221I) and data-to-data translation (Gushchin et al. 2023a). user is primarily interested only in the ends of these trajectories, i.e., new synthetic data samples from P1. In the biological applications. it may be useful to study the entire trajectory as well (Bunne et al., 2023 Pariset et al. (023). Hence, finding the solution to SB usually implies recovering the drift g" of SB or the conditional distributions x ((11 of the EOT plan.
3 LIGHT SCHRÖDINGER BRIDGE SOLVER
In 3.1 we derive the learning objective for our light SB solver. In 3.2 we present its training and inference procedures.
$3.3 we prove that our solver is a universal approximator of SBs.
3

---

Published as a conference paper at ICLR 2024
3.1 DERIVING THE LEARNING OBJECTIVE
The main idea of our solver is to recover parametric approximation = nt" of the EOT plan. Then this learned plan Te will be used to construct an approximation Ta T the entire SB process 3.2 To learn Ta, we want to directly minimize the KL divergence between xt and TØ:
KI. (N (| - min. DEG
(5)
This objective is straightforward but there is an obvious obstacle: we do not know the EOT plan
r². The magic is that optimization (5) can still be performed despite not knowing n². To be-
gin with, recall that we already know that the EOT plan x" has a specific form (3). We define
u" (xx) def exp(- Ilzoli ) (za) and ex (x1) def exp(- | (x1). These two are measurable 2e
2e
functions RD + R4, and we call them the adjusted Schrödinger potentials. Now reads as
" <"(o x1) = u² (2o) exp (moo (toe//////u (x1) =
π (21 x exp ((20, (()))=)))))"( (z1).
(6)
Our idea to exploit this knowledge to parameterize Te. We define
exp (2)))
( x1) = (2(20)==((11<00 = po
Co (1(
i.e. we parameterize v" as Ug. In turn, cz(z)) def fgo exp (200 (10. vo (2) dz is the normalization.
Our parameterization guarantees that SRD To(To, x1 )dr; = Po (2)). This is reasonable as in practice a learner is interested in the conditional distributions x (x1| rather than the density π (aro 21) of the plan. To be precise, we parameterize all the conditional distributions no (1 (200) via common potential tg. Below we will see that it is sufficient for both training and inference. In Appendix E we show that within our framework is easy also parameterize the density of po in (7). Surprisingly, this approach naturally coincides with just fitting a separate density model for Po-
Now we show that optimization can be performed without the knowledge of x². Proposition 3.1 (Feasible reformulation of the KL minimization), For parameterization 7 it holds that the main KL objective (5) admits the representation KL (T" [|(7T) L()) - L", where
def
log co(xo)po (oo)))o(()))) - fee (p (di)dar,
(())
JRD
and L* E IR is constant depending distributions Po,Pl and value 0 but not on 0.
We that minimizing KL equals to the minimization of the difference in expectations of log ca (200) and log vy (z(x1) w.r.t. Po. Pl, respectively. This means that the objective value (8) admits Monte- Carlo estimation from random samples and (8) can be optimized by using the stochastic gradient descent w.r.t. O. Yet there is still an obstacle that co may be hard to compute analytically. We fix this below.
We recall 6 with To = and see that x" (21 = 0) x 8" (21), i.e., v" can be viewed as an unnormalized density of a distribution. Inspired by this observation, we employ the (unnormalized) Gaussian mixture parameterization for the adjusted potential Ng:
Uo (11) def
where / def = {as, Tk. S are the parameters: ag 2 0, Tk E RD and symmetric 0 < S₁ E DxD. Note that multiply St in ( by just to simplify the further derivations. For parameterization conditional distributions ( (m) and normalization constants co(x)) are tractable.
Proposition 3.2 (Tractable form of plan's components). For the Gaussian Mixture parameterization of the adjusted Schrödinger potential vg in 7 it holds that
K
x (2) (mn) =
as (2
where
Te(#o)
def Fh + SEZ0.
co (zo)
k=
T Sg.o + 2r 20),
(x0) def
def āi (z0)-
Or exp
co(20)
2c
4

---

Published as a conference paper at ICLR 2024
The proposition provides the closed-form expression for ca())) which is needed to optimize (8
Relation to prior work. Optimizing objective (8) and using the Gaussian mixture parameterization
(9) for the Schrödinger potentials are two ideas that appeared separately in (Mokrov et all, 2024)
and (Gushchin et al. 2023b), respectively, and in different contexts. Our contribution is to combine these ideas together to get an efficient and light solver to SB, see the details below.
(a) Energy-based view on EOT. In (Mokrov et al. 2024), the authors aim to solve the continuous
EOT problem. Using the duality for weak OT, the authors derive the objective which matches our 8 up to some change of variables, see their eq. (17) and Theorem 2. The lack of closed form for the normalization constant forces the authors to use the complex energy-based modeling (LeCun et al. |006||Song & Kingma| [021] EBM) techniques to perform the optimization. This is computationally heavy due to using the MCMC both during training and inference of the solver. Compared to their work, we employ a special parameterization of the Schrödinger potential, which allows to obtain the closed form expression for the normalizing constant, In turn, this allows us to optimize ( directly with the minibatch gradient descent and removes the burden of using MCMC at any point. Besides, our solver yields the closed form expression for SB (see 33.2 while their approach does not.
(b) Parameterization of Schrödinger potentials with Gaussian mixtures. In (Gushchin et al.) 2023a), the authors are interested in manually constructing pairs of continuous probability distribu- tions for which the ground truth EOT/SB solution is available analytically. They propose generic method to do this (Theorem 3.2) and construct several pairs to be used as a benchmark for EOT/SB solvers. The key building block in their methodology is to use kind of the Gaussian mixture param- eterization for Schrödinger potentials, which allows to obtain the closed form EOT and SB solutions for the constructed pairs (Proposition 3.3, Corollary 3.5). Our parameterization (9) coincides with theirs up to some change of variables. The important difference is that we use it to Tearn the EOT/SB solution (via learning parameters 0 by optimizing (8) for a given pair of distributions po- P1- At the same time, the authors pick the parameters 0 at random to simply set up some potential and use it to build some pairs of distributions with the EOT plan between them available by their construction.
To summarize, our contribution is to unite these two separate ideas (a) and (b) from the field of EOT/SB. We obtain a straightforward minimization objective (8) which can be easily approached by standard gradient descent
This makes the process of solving EOT/SB easier and faster.
3.2 TRAINING AND INFERENCE PROCEDURES
TRAINING. As the distributions po, Pi are accessible only via samples X = (=b. and X = (z)......
" ~pp (recall the setup in $ we optimize the empirical counterpart of x My
N
M
((0) def
log co co(2x)
log Uo (o(x))) ~ L()).
(10)
N 1
M
1
We use the (minibatch) gradient descent w.r.t parameters 0. To further simplify the optimization, we consider diagonal matrices S₁ in our parameterization ( of ug. Not only does it help to drastically reduce the number of learnable parameters in but it also allows to quickly compute S in (D) time. This simplification strategy works reasonably well in practice, see beloo. Importantly, it is theoretically justified: in fact, it suffices to even use scalar covariance matrices S = A > 0 in vg, see our Theorem 3.4 The other details are in Appendix [
INFERENCE.
The conditional distributions (0(1) are mixtures of Gaussians whose parame-
ters are given explicitly in Propostion 3.2 Hence, sampling IΙ given zo straightforward and lightspeed. So far we have discussed EOT-related training and inference aspects and skipped the question how use me to set-up some process To T approximating SB. We fix it below.
With each distribution Ta defined by 7 via formula (7). we associate the specific process T = To whose joint distribution at 0, 1 matches Ty and conditionals satisfy To To W ja.1 Infor- mally, this means that we "insert" the Brownian Bridge "inside" the joint distribution my at t = 0, 1. Below we show that this process admits the closed-form drift go and the quality of approximation of T by To is the same as that of approximation
π by πe.
Proposition 3.3 (Properties of To). Let vg be an unnormalized Gaussian mixture given by 9 and To given by 77) Then To introduced above is a diffusion process governed by the following SDE:
To
dx,= go (X,, t)dt VcdW,, Xo Po:
(11)
We discuss the generalization properties of our light SB solver in Appendix A
5

---

Published as a conference paper at ICLR 2024
go(a,t) deff VV ((1))0,((1-t)1)) Σ ( (auNTraD.css))MMMttttt
with A def aل-ي + and (,(x,)) (x,t) def
cl-1)² x + S Tk. Moreover, it holds that
KL (T - KL (- (T=)
(12)
The proposition provides closed form for the drift go of To for all (x, RD x [0, 1]. Now that we know what the process looks like, is straightforward to sample its random trajectories starting given input points To. We describe
ways for based on the well-known schemes.
Euler-Maryama simulation. This is the well-celebrated time-discretized scheme to solve SDEs. Let At I be the time discretization step for an integer S > Consider the following iteratively constructed (for E 1. S - 11) sequence starting at .o:
((s+1)A Tsde + go (I. st + VEDiS. with ~ 10,,1,,
(13)
where &. are i.i.d random Gaussian variables. Then the sequence [ is time-discretized approximation of some true trajectory of T starting from Tg- Since our solver provides closed form ge (Proposition 3.3) for all t [0, 1]. one may employ any arbitrary small discretization step At.
Brownian Bridge simulation. Given start point zo, one can sample an endpoint 3 ~ mo (2 (20)
from the respective Gaussian mixture (Proposition 3.2) What remains is to sample the trajectory from the conditional process T which matches the Brownian Bridge W W
we already have some trajectory To, To.-.....--3 , I1 with O < t < ... < l < |z Suppose that we have only 20. x1), and we want to refine the trajectory by inserting point at Following the properties of the Brownian bridge, it suffices to sample
(initially = O,
4
1'
(xt, - xe,,, (t' (t)+4)
Tt. (1e
(14)
1++1 4
f
Using this approach, one may sample arbitrarily precise trajectories of without any discretization errors. Unlike (13), this sampling technique does not use the drift ge and to get random sample at any time t one does not need to sequentially unroll the entire prior trajectory at [0, t).
3.3 UNIVERSAL APPROXIMATION PROPERTY
Considering plans vo with the Gaussian mixture parameterization (.. it is natural to wonder whether this parameterization universal. Namely, we aim understand whether Te can approximate any Schrödinger bridge arbitrarily well if given a sufficient amount of components in the mixture Ug. We provide a positive answer assuming that Po: P1 are supported on compact sets. This assumption is not restrictive as in practice many real-word sol-world distributions are compactly supported anyway.
Theorem 3.4 (Gaussian mixture parameterization for the adjusted potential provides the universal approximation of Schrödinger bridges). Assume that pa and p1 compactly supported. Then for all 8 > there exists a Gaussian mixture Ug (9) with sealar covariances Sk = A > 0 of its components that satisfies the inequality KL (T (T)) KL (π* [[TT) < 8.
Although this result looks concise and simple, its proof is quite challenging. The main cornerstone proving the result is that the key object be approximated, i.e., the adjusted potential n", is just a measurable function without any nice properties such as the continuity. To overcome this issue, we employ non-trivial facts from the duality theory for weak OT (Gozlan et al. 2017).
We also highlight the fact that our result provides approximation of T on the non-compact set. Indeed, while Po and P1 are compactly supported, Té's marginals at all the time steps t E (0, 1) always supported on the entire RD which is not compact, recall, c.g., (14). This aspect adds additional value to our result as usually the approximation is studied in the compact sets. To our knowledge, our Theorem 3.4 is the first result about the universal approximation of SBs.
4 RELATED WORK
Over several recent years, there has been a notable progress in developing neural SB/EOT solvers. For review and benchmark of them, we refer to (Gushehin al.) 2023b). The dominant majority of them have rather non-trivial training or inference procedures, which complicates the usage of them in practical problems. Below summarize their main principles.

---

Published as a conference paper at ICLR 2024
Allows to sample from
Recovers
Strishes the universal apgroximation
Works for ressonably small
Solver
Property
minimax objective
iterativve objective
simulation based training
the drift (x,
x (-(x)
use simulation inderence
Seguy et
2013
Donulls al
X022
x
Maksov 0.
Cuslichua et
20/24
x
7
Ramooi
Chen ct 20218
/
/
Kuta al 20:4
LightB (ours)
Table I: Comparison of features of existing BOT/SB solvers and our proposed light solver.
Dual form solvers for FOT. The works (Genevay et al., 2016 Seguy et al. 2018 Daniels et al.,
2021 aim solve EOT T i.... the static version of SB. The authors approach the classic dual EOT problem (Genevay] 2019 $3.1) with neural networks, They recover the conditional distributions (21 11' [ro)dr; without learning the ac-
r" [2 [2n) only the barycentric projections 2g + fgo 21 "1"
tual SB process T Unfortunately, both solvers do not work for small c due to numerical instabilities (Gushchin al. 2023b Table 2). At the same time, large is of limited practical interest the EOT plan is nearly independent, i.e., n (1) Rs (oo(r))p (4r) and x" (x) (ma) Rs p1 (21). In this case, it becomes reasonable just to learn the unconditional generative model for P1. This issue is addressed in the work (Mokrov et al.) 2024) which we discussed in $3.1 There the authors consider the weak OT dual form (Backhoff- Veraguas et al., 2019. Theorem 1.3) for EOT and demonstrate that it can
be approached with energy-based modeling techniques (LeCun et al. 2006 EBM). Their solver as well as (Daniels et al. 2021) still heavily rely on using time-consuming MCMC techniques.
solvers can be also adapted to sample trajectories from SB by using the Brow- nian Bridge just as we do $3.2 However, unlike our light solver, these solvers do not provide an access to the optimal drift g" Recently, (Gushchin et al.) 2023a) demonstrated that one may reformulate the weak EOT dual problem so that one can get the SB's drift g" from its solution as well. However, their solver requires dealing with challenging max-min optimization problem and requires simulating the full trajectories of the learned process, which complicates training.
The
Iterative proportional fitting (IPF) solvers for SB. Most SB solvers (Vargas et al., 2021 De Bor- toli et al. 2021 Chen et al., 2021a 2023) directly aim to recover the optimal drift g as it can be later discussed in 63.2 Such solvers mostly based on the
used simulate the SB trajectories
iterative proportional fitting procedure (Fortet 1940, Kullback 1968 |Ruschendor]] 1995), which is also known the Sinkhorn algorithm (Cutur) (2013) and, in fact, coincides with the well-known algorithm (Dempster et al., 1977, EM), see (Vargas & Nüsken, [203]] Proposition 4. for discussion. That is. the above-mentioned solvers learn two SDEs (forward and inverse processes) and iteratively update them one after the other (IPF steps). The first two solvers do this via the mean- matchhin regression while the others optimize a divergence- based objective, see (Chen et al.) a0 $1). (Chen et al.) [023] 55). All these solvers require performing multiple IPF steps. At each of the steps, they simulate the full trajectories of the learned process which in- troduces considerable computational overhead since these processes are represented via large neural networks. In particular, due to the error accumulation as IPF proceeds, it is known that such solvers may forget the Wiener prior, see (Vargas & Nüsken 2023 1..1.2) and references therein.
Other solvers for EOT/SB. In (Shi et al. 2023), new approach to SB based on alternative Marko- vian and Reciprocal projections is introduced. In (Tong et al. (023). the authors exploit the property They propose a
that SB solution T consists of entropic OT plan π and Brownian bridge W
Oree Trgestory statis Trpctter and (91ta) t0guttey
Traectory cojectery of
"uadtey and (itee)) "agectry of
(a) x ~ Po. y ~ P1.
cb)== 2-
(c) c = 0.01.
(d) ( 0.1.
Figure 2: The process To learned with LightSB (ours) in Gaussian + Swiss roll example.
7

---

Published as a conference paper at ICLR 2024
c=0.1
D=2
D=16
D=
64
128
D=2
D=16
=
D= 128
D=
D=16
D=
D
128
Best solver
1.94
13.67
11.74
11.4
1.04
9.08
18.05
15.23
1.40
1.27
2.36
1,31
LightSB
0.03
0.08
0.28
0.60
0.05
0.09
0.24
0.62
0.07
0.11
0.21
0.37
Estd
00.01
#0.04
40.02
#0.02
#0.003
+0.006
#0.007
+0
+0.007
#0.02
+0.01
+0.01
=0.01
Table 2: Comparisons of cBW3 UVVPU 4 (%) between the optimal plan n² and the learned plan Ne.
new objective to learn T in the form of SDE the EOT plan गt' is known. Since x² is not actu- ally known, the authors use the minibatch OT to approximate it. In (Kim et al., 2024), the authors propose exploiting the self-similarity of the SB problem and consider the family of SB problems on intervals {[4,, 1) to sequentially learn T However, they add an empirical regularization which may bias the solution. Besides, there exist SB
series of conditional distributions =t;; 1
solvers for specific setups with the paired train data available (Somnath et al. 2023 Liu al. 2023
Summary. We provide a Table D with the summary of the features of the discussed EOT/SB solvers. Additionally, in Appendix F we mention other OT solvers which are related but not closely relevant to our paper because of considering non EOT/SB formulations or non-continuous settings.
5 EXPERIMENTAL ILLUSTRATIONS
Below we evaluate our light solver on setups with both synthetic ($5.1
5.2
and real data dis-
tributions (55. $5.4. The code for our solver is written in PyTerch available at https: The experiments are issued in the form of convenient *. .ipynb notebooks. Reproducing each experiment requires few minutes on CPU with cores. The implementation and experimental details are given in Appendix D
//github. com/ hgushchin/Lightt
5.1 MIV-DDIMENSONNA EXAMPLES
To show the effect of on the learned process Te. we give toy example of mapping 2D Gaus-
sian-Swwiss Roll with our light solver for = 2 - 10-3 100 10-2,,0 10011 see Fig. 2 As expected, for small c the trajectories are almost straight, and the process Tg nearly deterministic. The volatility of trajectories increases with c, and the conditional distributions Ty Ty(I1 (Z1 become more disperse.
5.2 EVALUATION ON THE EOT/SB BENCHMARK
To empirically verify that our light solver correctly recovers the EOT/SB, we evaluate on a recent EOT/SB benchmark by (Gushchin et al. 20236, 85). The authors provide high- nigh-dimensional contin- uous distributions (Po. P1) for which the ground truth conditional EOT plan x² (-(zo) and SB process T* are known by the construction. Moreover, they use these pairs to evaluate many solvers from $
We use their mixtures benchmark pairs (see their 54) with various dimensions and e, and use the same conditional B BVV- metric (see their 55) to compare our recovered plan no with the ground truth plan x² In Table 2 we report the results of our solver vs. the best solver in each setup according to their evaluation. As clearly seen, solver outperforms the best solver by considerable margin. This is reasonable as the benchmark distributions are constructed using the similar principles which solver exploits, namely, the sum-exp (Gaussian mixture) parameterization of the Schrödinger potential. Therefore, our light solver has a considerable inductive bias for solving the benchmark.
5.3 SINGLE CELL DATA One of the important applications of SB is the analysis of biological single cell data (Koshizuka & Sato, 2022 Bunne et al., 2021 2022). In Appendix C we evaluate our algorithm on the popu- lar embryonic stem cell differentiatio
dataset which has been used in many previous works (Tong
ct al. 2020; Vargas et al., [202]] Bunne et al. 2023 Tong et al. 2023); here we consider the more high-dimennsonal dataset from the Kaggle completion "Open Problems Multimodal Single-Cell Integration" (MSCI) which was first used in (Tong et al. 2023). The MSCI dataset consists of
single-cell data from four human donors at time points (days 2, 3. and 7). We solve the SB/EOT problem between distribution pairs at days 2 and 4, 3 and 7, and evaluate how well the solvers re- cover the intermediate distributions at days 3 and correspondingly. We work with PCA projections
Setup
Solver type
DIM
50
Discrete ECT
Sinkhom
Cuturi 20
GPC Vio0)
2.34 (G0 st
2.24 (2.5
(6 m)
Comiinuus BUT
based
Makzow al
GPU VIRI
2.39 0. 10 (19
2.32
0.15
1.46 1 0.20 (15
Coniiuous E01
Minimax
Grobchin et all Arssa [I GPU VIDO
2.44 0. 13:43 m)
2,24
13
m)
1.32 1 0.08
Comiiuous BUT
IPs
Vargas ot al,
2021 GPU ViBI
3.14 4: 0.27 (8 m)
2.56
4
U.20
2.05 + u.19
Cominuuus
minimization
LightSB (ours)
cures)
0.27
(65
0.26
19
Table 3: Energy distance (averaged for two setups and 5 random seeds) on the MSCI dataset along with 95%-confidence interval (+ intervals) and average training times (s - seconds, m minutes).
8

---

Published as a conference paper at ICLR 2024
(a) Male >>
Female.
(b) Female - Male.
(c) Adult - Child.
(d) Child - Adult.
Figure 3: Unpaired translation by our LightSB solver applied in the latent space of ALAE for 1024x 1024 FFHQ images. Our LighuSB converges 4 cpu cores in less than I minule.
with DIM = 50, 100, 1000 components. We use the energy distance (Rizzo & Székely) 2016 ED) as metric and present the results for different classes of SB solvers in Table Jalong with the training time. We see that LightSB achieves similar quality to other EOT/SB solvers but faster and without GPU. The details of used preprocessing. hyperparameter and baselines are in Appendix D.4
5.4 UNPAIRED IMAGE-TO-IMAGE TRANSLATION
One application which is frequently considered in EOT/SB papers (Daniels et al., 2021; Chen et al.) 2021b) is the unpaired image-to-image translation (Zhu et al. 2017). Our solver may hard to apply to learning SB directly in the image space. To be precise, it is not designed to be used in image spaces just like the conventional Gaussian mixture model is not used for image synthesis.
Still we show that our solver might be handy for working in the latent spaces of generative mod- els. We consider the task of male male-female translation. We pick pre-trained ALAE autoencoder (Pidhorskyi al. 2020) for entire 1024 x 1024 FFHQ dataset (Karras et al. 2019) of 70K human faces. We split first 60K faces (train) into male and female subsets and use the encoder to extract 512-dimensional latent codes { ( N 1==1 and {2|h 1==1 from the images in each subset.
Training. We learn the latent EOT plan Fe(2112)) by using the above- mentioned unpaired samples from the latent distributions. The training process takes less than I minute on CPU cores.
Inference. To perform male-ffemale translation for new male face xgew (from 10K test faces), we (1) encode it via zB znew Enc(x)ee), (2) sample 21 ~ [2 ([21laaeu) and then (3) decode Iι Dec(21) and return it. Note that here (unlike $5.33 the process To is not needed, only Tg is used.
Results. The qualitative results are given in Fig. T and Fig. 3a Furthermore, in Fig. 3] we provide additional examples for other setups: female- -aaee and child'-aduul.. For brevity, we show only translated images per an input image. In Appendix H.] we give extra examples and study the effect of Our experiments qualitatively confirm that our LightSB can solve distribution translation tasks in high dimensions 512)., and it can be used to easily convert auto auto-encoders to translation models.
6 DISCUSSION
Potential impact. Compared to the existing EOT/SB solvers, our light solver provides many ad- vantages (Table Dp. It one-step (no IPF steps), does not require max-min optimization, does not require the simulation of the process during the training, provides the closed form of the learned drift go of the process To T and of the conditional distributions (1(2) R x (x112z) (zo) of the plan. Moreover, our solver is provably universal approximator of SBs. Still the key benefit of our light solver is its simplicity and ease to use: it has straightforward optimization objective and does not use heavy-weighted neural parameterization. These facts help our light solver to converge in matter of minutes without spending a lot of user/researcher's 's time on setting up dozens of hyperpa- rameters, e.g., neural network architectures, and waiting hours for the solver to converge GPUs. We believe that these properties could help our solver to become the standard easy-to-use baseline EOT/SB solver with potential applications to data analysis tasks.
Limitations and broader impact are discussed in Appendix

---

Published as a conference paper at ICLR 2024
7 REPRODUCIBILITY
The code for our solver is available at
1. To reproduce experiments from :5.1 it is enough to train LightSB model by running notebook notebooks then run notebook sstsbookss/swsss..rooo..o
ppyb with hyperparameters described in $D22 and ipynn ipynb to plot Fig. P
2. To reproduce experiments from 55. needed to install Entropic OT benchmark from github and then run notebook Light EOT www.hh/w.ccaa...oo ipy.mnrrkkiiy.b........ ipynb with hyperparameters described in DD3 to reproduce ported metrics in T
3. To reproduce experiments from Appendix C
it is needed to install library from and then to run note- book notebooks/LLigheeee wtt.ssssttttttt..stsis..sssss wtw.siiggn.elll...... e_cell. ..ppel...ppp......... All required data is already prepro- cessed and located data folder.
To reproduce experiments from 5..3 it needed to download data from https: htt.://a......cccccoooooooo////a
and ipynb to preprocess data.
then to run notebook aattaa/aata..reproceet The experiments with LightSB solver can be reproduced by running the notebook The experiments with Sinkhorn solver can be
notebooks / Light L.B.MSCI. reproduced by running the notebook wwweeookks/Shthhoon...ss.. ipynb
The code for ALAE is already included in our code and reproduce experi- ments from $5.4 it is first necessary to load the ALAE model rumning the script ALAE/training wtttafats//towllo...... FFHQ dataset from ALAE and these data can be downloaded directly using notebook notebooks t ipynb. To train the LightSB model is necessary to run the notebook notebooks Light aea.aal..iee.. ipynb. The same notebook also contains code for plotting the results of trained models.
We have already coded the
8 ACKNOWLEDGEMENTS
The work supported by the Analytical center under the RF Government (subsidy agreement 000000D733332115500002 Grant No. 70-2021-00145 02.11 02.11.20
REFERENCES
Brandon Amos. On amortizing convex conjugates for optimal transport. The Eleventh Interna- tional Conference on Learning Representations, 2022.
Arip Asadulaev, Alexander Korotin, Vage Egiazarian, Petr Mokrov, and Evgeny Burnaev, Neural op- timal transport with general cost functionals. In The Twelfth International Conference Learning Representations, 2024. URL
net. firee/.footee forumm www.ruumiiiniieeeeeeeee
Julio Backhoff- Veraguas, Mathias Beiglböck, and Gudmun Pammer. Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):1-28, 2019.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong- Kjong-Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single-cell perturbation re- sponses using neural optimal transport. bioRxiv. pp. 2021-12 2021.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal trans- port modeling of population dynamics. In International Conference on Artificial Intelligence and Statisties, pp. 6511-666 1-5528.. PMLR, 2022.
10

---

Published as a conference paper at ICLR 2024
Charlotte Bunne, Ya-Ping Hsieh, 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, pp. 5802-5833. PMLR. 2023.
Tianrong Chen, GurnnHorg Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using tations, 2021a
sdes theory. In International Conference Learning Represen-
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Optimal steering of linear stochastic system to final probability distribution, part IEEE Transactions on Automatic Control, 61(5): 1158-1169, 2015.
Yongxin Chen, Tryphon T Georgiou. and Michele Pavon. On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169:671 099:671-691,220 2016.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrodinger bridge. SIAM Review, 63(2):249-313 2021
Yu Chen, Wei Deng, Shikai Fang, Fengpei Li, Tianjiao Nicole Yang, Yikai Zhang, Kashif Rasul, Shandian Zhe, Anderson Schneider, and Yuriy Nevmyvaka. Provably convergent schrödinger bridge with applications to probabilistic time series imputation. 2023.
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport. arXiv preprint arXiv:2305.14777, 2023.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large-scale optimal transport. Advances in neural information processing systems, 34:12955-12965 2021.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695- 2021.
Nabarun Deb, Promit Ghosal, and Bodhisattva Sen. Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections. Advances Neural Information Processing Systems, 34:29736-29753, 2021.
Arthur P Dempster, Nan M Laird, and Donald Rubin. Maximum likelihood from incomplete data via the algorithm. Journal of the royal statistical society: series B (methodological), 39(1): 1-22, 1977.
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's 's algorithm. Interna- tional conference machine learning, pp. 1367-1376. PMLR, 2018.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estima- tion and its applications. Transactions on Machine Learning Research, 2023, ISSN 2835-8856. URL / net forum? iediddd2maa add=mmmss@sc33 Featured Certification.
http:///peereeiiieee
Kilian Fatras, Younes Zine, Rémi Flamary, Remi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein: asymptotic and gradient properties. In International Conference on Arti- ficial Intelligence and Statistics, pp. 2131-2141. PMLR, 2020.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar's Z. Alaya, Aurélie Boisbunon, Stanis- las Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Foumier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021. URL http::ttt///. jmml wwwp/pppprs///222222 451 html

---

Published as a conference paper at ICLR 2024
Robert Fortet. Résolution d'un système d' é'quutttons de m. schrödinger. Journal de Mathématiques Pures et Appliquées, 191-4))83-105, 1940.
Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Extremal domain translation with neural optimal transport. In Thirry-seventh Conference on Neural Information Processing Systems, 2023. URL https: https://openievieett wttp:/ewwiwwwweeeeeeeeeeeeeeeeee mwt//fouut 1it.mmmii-vv222ioo338...
Aude Genevay. Entropy-regularized optimal transport for machine learning. PhD thesis, Paris Sciences et Lettres (ComUF), 2019.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. In Advances in neural information processing systems, pp. 3440-3448, 2016.
Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, 273(1 3333773332733, 2017.
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnaev. En- tropic neural optimal transport via diffusion processes. Advances Neural Information Pro- cessing Systems, 2023a.
Nikita Gushchin, Alexander Kolesov, Petr Mokrov, Polina Karpikova, Andrey Spiridonov, Evgeny Burnaev, and Alexander Korotin. Building the bridge of schr odinger: A continuous entropic optimal transport benchmark. In seinth- Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023b.
Pierre Henry-Labordere. (martingale) optimal transport and anomaly detection with neural net- works: A primal-dua algorithm. Available at SSRN 3370910, 2019.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in 2020.
Neural Information Processing Systems,
Can-Christian Hütter and Philippe Rigollet. Minimax estimation of smooth optimal transport maps. 2021.
Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic optimal transport be- tween unbalanced gaussian measures has a closed form. Advances in neural information process- ing systems, 33:10468- 2020.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Beomsu Kim, Gihyun Kwon, Kwanyoung Kim, and Jong Chul Ye. Unpaired image-to-image trans- lation via neural schrödinger bridge. In The Twelfih International Conference on Learning Rep- resentations, 2024. URLhttps:
wwwp///eee.....eeeemme//. net forum?iG=uQBui 7ELXfO
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1:Xiv:1412.6988 2014.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. International Conference on Learning Representations, 2021a. URL WRRL/httts:///peeeeeeee
net wett//ooommddddbooaaaaaa.. EXsa
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34: 114499--4605, 2021b.
Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative net- works for barycenter estimation. In Thirty-Sixi Sixth Conference on Neural Information Processing Systems, 2022a. U http:::///peereviieee net. nerm/o d= EnzxTnaN
12

---

Published as a conference paper at ICLR 2024
Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasser- stein GANs are not optimal transport? In Thirty-sixth Conference on Neural Information Pro- cessing Systems Datasets and Benchmarks Track, 2022b. BeLLattpss///peneeviie net/
forum? Expt tGll
Alexander Korotin, Daniil Selikhanovych, and Exgeny Burnaev. Kernel neural optimal transport. In International Conference on Learning Representations, 2023a. URL /
dim Zuc Uma
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In In-
ternational Conference on Learning Representations, 2023b. URL net Dott/Eooumniic-cCBBBLLLNM
Takeshi Koshizuka and Issei Sato. Neural lagrangian schr)""(o} dinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations, 2022.
Solomon Kullback. Probability densities with given marginals. The Annals of Mathematical Statis- lics, (9)::; 1233-1243, 1968.
Fabian Latorre, Leello Tadesse Dadi, Paul Rolland, and Volkan Cevher. The effect of the intrinsic dimension the generalization of quadratic classifiers. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https://openreview net worumoorriiooo hKvt sqIte
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. tutorial on energy-based learning. Predicting structured data, 100), 2006.
Christian Léonard. survey of the schr odinger problem and some its connections with optimal transport. arXiv preprint arXiv:1308.0000
2013.
Guan-Homg Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. Pbb: Image-to-image schr\." odinger bridge. arXiv preprint arXiv:230.00877, 2023.
Xingchao Liu, Chengyue Gong, et al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2022.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672- 6681. PMLR, 2020.
Anton Mallasto, Augusto Gerolin, and Hà Quang Minh. rntropy- 2- wasserssnnin distance between gaussian measures. Information Geometry, 5(1):289-323, 2022.
Tudor Manole, Sivaraman Balakrishnan, Jonathan Niles-Weed, and Larry Wasserman. Plugin esti- mation of smooth optimal transport maps. arXiv preprint 107.i:22107.23660 2021.
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018
Petr Mokrov, Alexander Korotin, Alexander Kolesov, Nikita Gushchin, and Evgeny Burnaev. Energy-guided entropic neural optimal transport. In The Twelfth International Conference on Learning Representations, 2024. URL
eww. ew net forum?rdd
dotUs2eVs7
Quang Minh Nguyen, Hoang H Nguyen, Yi Zhou, and Lam M Nguyen. On unbalanced optimal transport: Gradient methods, sparsity and approximation error. arXiv preprint arXiv:2202.03611, 2022.
T Tin Nguyen, Hien D Nguyen, Faicel Chamroukhi, and Geoffrey J McLachlan. Approximation by finite mixtures of continuous density functions that vanish at infinity. Cogent Mathematics & Statisties, 175)615, 2020.
13

---

Published as a conference paper at ICLR 2024
Matteo Pariset, Ya-Ping Hsich, Charlotte Bunne, Andreas Krause, and Valentin De Bortoli. Unbal- anced diffusion schr)," odinger bridge. arXiv preprint arXiv:2306.09099 2023.
Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008.
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport. Foundations and Trends in Machine Learning, 11(5-6):355-60)) 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 14113, 2020.
Mark A.
Pinsky and Samuel Karlin. brownian motion and related processes. In Mark A. Pinsky and Samuel Karlin (eds.), An Introduction to Stochastic Modeling (Fourth Edition), pp. 391-446. Academic Press, Boston, fourth edition edition, 2011. ISBN 178-0-12-38- 1416-6. doi: https://doi. 0tg///0..066/398880-133334446660 ccieenee/arricle///1////82888282 66000008
URL https: www www.eceeiiii....c... rect. com/
Aram-Alexandre Pooladian and Jonathan Wiles-Weed Entropic estimation of optimal transport
maps. arXiv preprint arXiv:2 1020099.00044 /2004, 2021.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional conference on machine learning, pp. 1530-1538. PMLR, 2015.
Philippe Rigollet and Austin J Stromme. On the sample complexity of entropic optimal transport.
arXiv preprint urXiv:2206. 13006.13440 13472, 2022.
Maria L Rizzo and Gábor J Székely. Energy distance. wiley intendisciplinary reviews: Computa- tional statistics, 8(1):27-38, 2016.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2021.
Ludger Ruschendorf. Convergence of the iterative proportional fitting procedure. The Annals of Statistics, pp. 1 160-1177, 1995.
Erwin Schrödinger. Über die unkehrung der naturgeselze. Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u..., 1931.
Erwin Schrödinger. Sur la théorie relativiste de l'électron et l'interprétation de la mécanique quan- tique. In Annales de l'institut Henri Poincaré, volume 2, pp. 269-310, 1932.
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Con- ference on Learning Representations, 2018.
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algo- rithms. Cambridge university press, 2014.
Yuyang Shi, Valentin De Bortoli, Andrew Camphell, and Araud Doucet. Diffusion schrödinger bridge matching, In Thirth-seventh Conference on Neural Information Processing Systems, 2023. URL http:htts///ppeeeiiieee htt.:iiwwwweeeeeeeeeeeeeeeeeeeee forum id -dgrruni-- JTS
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion schr\," odinger bridges. arXiv preprint arXiv:2302. 1419, 2023.
Yang Song and Diederik P Kingma. How to train your energy-based models. arXiv preprint /0t::2/011002888 2021.
14

---

Published as a conference paper at ICLR 2024
Austin Stromme.
Sampling from schrödinger bridge. In International Conference Artificial Intelligence and Statistics, pp. 4058 4067. PMLR, 2023.
Amirhossein Taghvaei and Amin Jalali. 2- 2-Wasserstein approximation via restricted convex poten- tials with application improved training for GANs. arXiv preprint arXiv: (eXX:::902..0197 2019.
Alexander Tong, Jessie Huang. Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet. A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9526-9536. PMLR, 2020.
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free schr'." odinger bridges via score and flow matching. arXiv preprint arXiv:2307.33672 2023.
Francisco Vargas and Nikolas Nüsken. Transport, variational inference and diffusions: with appli- cations to annealed flows and schr\," odinger bridges. arXiv preprint arXiv:2307.0n 0000000 2023.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9):113 134, 2021.
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International Conference on Machine Learning, pp. 10794-10804. PMLR, 2021.
Yiling Xie, Yiling Luo, and Xiaoming Huo. An accelerated stochastic algorithm for solving the optimal transport problem. arXiv preprint arXiv:2203.000 00813, 2022.
Karren D Yang and Caroline Uhler. Scalable unbalanced optimal transport using generative adver- sarial networks. International Conference on Learning Representations, 2018.
Jun- Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-imes tonggge translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223-2232. 2017.
15

---

Published as a conference paper at ICLR 2024
A GENERALIZATION PROPERTIES OF OUR LIGHT SOLVER
In theory, to recover the optimal plan x one can solve (()) + ming which is equivalent to the
direct minimization of KL (m=|| w.r.t. (Proposition 3.1 According to ( [()) consists of the difference of integrals of log co (To) and log uy (21) over the distributions Po and pr, respectively. In practice, there are several sources of errors which do not allow to perfectly optimize the objective.
1. Statistical (estimation) error. Since distributions Pu, P1 are accessible only via empirical sam- ples X (=b)-. ................................ To Po and X {z).... {2i,...,=j]] x") Pt, one is forced optimize the empirical counterpart ((0) of L(0). In this objective, the integrals over po: Pi are replaced with their estimates using samples X x² recall (10). For given samples X X we denote
Ū - O(x', x!) arg min ((0).
(15)
Usually, Z(6) is called the empirical risk and @ is the empirical risk minimizer.
Approximation error. The parametric class for Uo over which one optimizes the objective is restricted. For example, we consider (unnormalized) Gaussian mixtures Uy with K components (parametrized with f = {as, rk, S {k) This may lead to irreducible error in approximation of the OT plan x² with Ty due parametric restrictions. In our setup, the quantity
C( f min L(0)- L*
(16)
is the approximation error. Here o" arg ming C()) is the best approximator (in a given class).
Optimization error. practice, we solve Σ(0) ming with the gradient descent. The opti- mization w.r.t. is non-convex and there are general guarantees of convergence to the global empirical risk minimizer 0. This may introduce an additional optimization error. Analysing this quantity is a too general question of the non-convex optimization and goes far beyond the scope of our paper. Therefore, for further analysis assume this error to be zero
Given the gap between the theoretical objective C(0) and its empirical counterpart Z(0). . it natural to wonder how close is the recovered Tg to π² We aim to obtain the bound for the expected KL between n² and Th (or, equivalently, 7 and Te), i.e., EKL (x (|=a) IKLL (T* [1T)), where the expectation is taken w.r.t. the random realization of the train data X .x.. This quantity is the
natural definition of the generalization error in our setting. Note that
EKL (T" (T')ITT) 11TT)) Prop. 3.3
EKL (TT" (|TT])
3.
EE [C(()) L*]
E ())( 2(00)) + E ((0)) ) -- = C(() =
L(
[(0)) - L-]
(17)
Statistical error
Approximation enror
Thanks Theorem 3.4 we already known that the second term (the approximation error) can be made arbitrarily small if we pick a Gaussian mixture with sufficiently large amount of compo- nents. Hence, our analysis below focuses on bounding the statistical error and understanding the rate of its convergence to zero as a function of available sample sizes N, M. Our following theorem demonstrates that the statistical error decreases at the usual parametric rate.
Theorem A.1 (Bound for the statistical error). Assume that po: Pt are compactly supported. Assume that the considered parametric class (3 0) consists of (unnormalized) Gaussian mixtures with K components with bounded means ||r||| R (for some R. > 0), covariances sl Sy 3 S1 (for some < s S S) and weights a I ag I A (for some 0 < a S A). Then the following holds:
Co C1
E I 0) [(O'
<
VN VM
where constants Co, C1 depend only on K, R, 8, S, a, A, Po- PL, e but not on sample sizes M, N.
The proof is given in the next Appendix section. In the future, it would be interesting to study the trade-off between the statistical error and approximation enror rather than study these instances separately we do in paper. However, providing such analysis will probably require making stronger assumptions (e.g., smoothness) on the distributions Po- P1 and the true optimal adjusted Schrödinger potential v*. We leave this interesting question for the future work.
16

---

Published as a conference paper at ICLR 2024
B PROOFS
B.1 PROOFS FOR THE RESULTS IN THE MAIN TEXT
Proof of Proposition 3.7 In the derivations below, we use H-) denote the entropy, i.e., the minus KL divergence with the Legesgue measure. We obtain
/ RRXXRO T" "'(00, v1) log nt (x0, 21) dradei
KL (=((x[[TT)
To (1)
-H(TT)
drgder x T' (1), log (p0(2o)T((z1(T0) ) =
-H(T*)
x (0 x1) log (oxo))ddd JRDXRD T' x1) (z)
RR XRD
=
exp (())))))))v((z1) drgde
-H(T*)
T (zo) log Po (o(or)do -
x² x'(o, log
co(
H(p))
+ H(=')+((po) -6 S (10, (100) x )drada + (18)
del
J
XRD T'(00, x1) log co (ooo)ddd - /
"-' >"(00, x1) log vo( )drgdx
RD RD
/ po( log (oo)
-L* +
P1(z1) log (x)1)dd = L(0) L" (19)
20
RD
which is exactly what we need.
Proof of Proposition 3.2
We use equation
for To æ1) and equation 9 for 2o (x1) to derive:
To <e(2)(z0) = exp (1r0, m1) (1)o
1
K
exp (100,
ax (xalax cSx) =
co(x0)
co(x))
Ose exp (20, ((or)11/))))
Co (ro) k=1
SE' T (a) -kk)
(27) D/222/255 -D/2))1-11( exp ((00 exp( (01 rxx)T
co (zo)
e
k=
K
(
04 (2T)) (/2 1/2 -1/2-111pep exp
(2.g) 12 - (x) rra) Si (0) - nw))) =
co(ar)
2c
1=1
1
1
a6 (2TT)= D/22 -//2 -1/2
(200 x1 If S + 2r S xι - r/ S r1))
exp
ce (z0)
2
k=1
K
1
ax((2) 1 D/2 D/2 SSa
-1/2
exp
(
1
(-x
T
S
+ 2 (Sa20 + rx) S- 20 - 1/ S (w.)) =
co(
2r
1=1
=x<(xo)
K
(2T)
-D2 eS
1/2 -1/2
exp
-- (x) - - (zo)) Sp (x1 - rk()))
(x
20
1
(-r T Sx nx + r (xo) (xo))) =
exp
2c
17

---

Published as a conference paper at ICLR 2024
K
SSS 1x + r (xo)) 1//
Og exp
1.
W (1)),
co( (20)
2c
k=1
1
T 1
nh + (Saao + r4) r (xo) (Sk20 + rx)
ag exp
Co (r))
2e
k=1
To SAI0 + 2 20
or exp
(x))
co (20)
2c
(m6)
au (an))))))
Cp (Fo)
Since Sap (x(21 mo)ddr 1, we see that = k=1 Oeg (zo) and conclude the proof.
Proof of Proposition 3.3 Define pe = J2D Te. Te Z11 ddr, as the density of the second marginal of
TT. From the OT benchmark constructor (Gushchin et al., 2023b] Theorem 3.2). follows that
constructed is the unique EOT plan between Po and pa: just set I" (xx1 del' ||x,1] + log va (111)
the mentioned theorem. Thus To the Schrödinger bridge between po and ps by its construction.
Then the fact that Ta is given by SDE (1) follows from the direct integration of (4 using (x(2)) def
exp( [| (1)( as the Schrödinger potential: 2c
(0 EV log / NN('11z, (1 t)eID)oo(Z) (z')d'''' )dr'
RRD
|
V a log
N/x'l lx, (1 t)(TD) expl
)uo "a''d'''
RD
2c
|
c log
N1x'la, t)(n)) exp(
Σ
(2 (x)) (5))dr'
JRD
2c
eV log Σ [ak / N(x'lz, (1 t)elD) (2')ri,, eSa) exp(
k
||'0|1 )dx')
RD
2e
K
$ V log (2m) 2 [11 1)en
k=1
V. of
(x' - 2T(T(a' - x)
(x' S (a' Tk)
xT >da' })
exp(
RD
2c 2e -()
2c
cV, log (exp(
TT
(aul<Syl - exp(
r ST 1rk
rky
2c t)
2e
1=1
/ RD
1
S ) x'|
1
1 S rr)]"
exp(-
{11 In +
(To]".o "=x' "'')d'''))
(20)
c11 t)
A
(he(u.))
x x
K
-p 1r rkg
ev log
exp(
(axleS>1 exp(-
2e(1 1
2c
km
[ - 1 (2T)
exp( 1if (z. ()(( (1, (x,{)))))))
(21)
2
18

---

Published as a conference paper at ICLR 2024
K
T S 'T
Vv log (2T) exp(
20 I
C4 (2T) 3 |e5xl =exp(
2c(1
2c
Ma])),(11-)1o
N(ru(0,eS.)
(2x) 3/2 A 2 expl ( (x, (x,
())
(22)
N(h(2.0)0.AL)
<V log (A(x)0, t)Ib)
Σ km
In the transition from 20 to 21
we use the integral formula from (Petersen et al. 2008 Sec 8.1.1).
In the transition from to we simply multiply the expression under V log by (2x) 22 as this does not change the expression.
Finally, with the measure disintegration theorem Vargaa
2021
Appendix C), we obtain
KL (T" (T)) KL (T (T)) +
KL (Tr 10 Te (æo, 21 diddzodr1
KL (I [[TT))
RPXRO
where we cancel out the KL term it coincides with KL (we [z0.20 | (120221) 0.
Proof of Theorem 3.4 It is intuitively clear that if we are able to approximate of arbitrarily well (in some sense) via v9, then we also achieve small KL (T* [[=)) as ma explicitly depends on up. The challenge here is that v" is just measurable function without any prior known properties, e.g., continuity. Hence, approximating it with continuous mixture some reasonable norm, c.g., the uniform norm ||- oloc. may be even impossible. This emphasizes the challenge of deriving the desired universal approximation result and points to necessity to use more tricky strategies.
Recall that for all 5 O need to find unnormalized Gaussian mixture to va(a) such that KL (xx ([=)) < 8. To begin with, pick any such 6 >>> 0 and fix it until the end of the proof.
Stage 1. This stage is about employing certain known facts from the EOT duality. Let us use
Cost(TT") def
1/2||20 - 2 (xo x1) + eKL (T* IlPo x P1)
(23)
JRRRR
to denote the optimal value of (T) We start from considering the equivalent reformulation of [
Cost( =
{
{
min =Ell(po-P1)
2 II.0 21112 01 dx cKL po p1)
R.D
(H(1)) + min rEIl(po-PI) Jgp { aro a1 = (2n eH (cH(((((1o)) (1o( (24)
where (-|22)) denotes conditional distribution of x given Zg. Term in (24) known as the weak representation of EOT, see (Gushchin et al.] 2023b Eq. (3) and (5)) for extra discussion, and admits a dual form (Backhoff- Veraguus et al. 2019 Eq. (1.3)):
J sup J(f) def (RC)
(EC,,(D)) ( (fe jc (co))o((z))d [ (())))(()))e)), (25)
sup
where C2 del is the so-called weak (entropic) C-transform of which is defined by
Rf:RD R continuous s.t. 3a, 3.7 :R: all + 3. $.)) 5 and po
(RP)(RD) der inf ( 1/0 2 1|:00 - x1 2112g(21)dm, - cH()) - -J001S51 [ [(o1)g(2))d (x))
(26)
19

---

Published as a conference paper at ICLR 2024
Here P2(RR) are all the probability distributions whose second moment is finite. We slightly abuse the notation as we write ((21) although here is not necessarily absolutely continuous. How- ever, if q does not have density, then -cH (4) +o0, which is a bad option for the minimization problem (26). Therefore, one may consider E P2. (TR C (2) in (26). The advantage of
EOT compared to many other OT formulations is that the minimizer of (26) is available explicitly:
((())) 1/2|/r x1 ||
d (11) def
(1(z))
exp
(27)
(Mokrov al. 2024 Proof of Theorem 1). The mentioned paper considers the compact subsets of RDd but their derivation is generic and works for our non-compact case as well. Here
[(x 10 - x11² dx
7 (30)
def
exp
(28)
JRD
c
is the normalizing constant. It is finite thanks to the upper boundness of f due to belonging to (TR...) Due to the same reason, it is not hard check that q has finite second moment. If we further follow the mentioned work and plug q in (26). we get (2o) = -e log Z (zo).
From 25 and the definition of the supremum, it follows that for all 8' > O there exists some function
e C2 (T. P) for which the following inequality holds:
J(Î) -E f log z (zo) (zo)dzo + / (()1)p((z)dde) Cost(TT)) - -
JRD
For our needs, we pick 8' & and suitable f for it and move on to the next stage.
Stage 2. Let 7 E be an upper bound for f. i.c., f(z1) $ 9 for all 21 R.. It exists thanks to I (RD)... Recall that pı is compactly supported by the assumption of the theorem. Let R be some radius such that the zero-centered ball of this radius contains the support of pr. We define
f(x)) def f(x1)) - max(0, |>||| 1²]} s f(ai) S7.
We see that
vz fo J.P jc (z00)oo(20)dd 2 f 80 (2o))o(zo)daa) (29)
/
By the construction of f. holds that f(æ1) f(21) when 2n is in the support of p1. Thus,
[ ((11)p((2))))) JRP
{ f(z1)1((2)ddr;
(30)
We combine 29 with (30 and see that J()) 2 J(F > J 8 Cost(=")) - (pi) 8'.
We note that Po is compactly supported, and z is continuous (w.r.t. x6) and non-negative. There- fore, there exists a constant 2min > 0 such that 2J (fo) 2 Zmin when zo belongs to the support of pn- Analogously, since p1 is compactly supported, we may find a positive constant emin 0 such that exp(/(xx)/ 3/4) 2 emin for all 2 in the support of p1. We fix constants 2min, for next steps.
Right now we derive
( ma 11 -R-)- R2) S exp ? R exp(-Ilax11/A))
exp (f(m))/) S exp
This means that x1 + exp (F(z))/)) is normalizable density (go exp (fer))/4)zz < oc) because its density is bounded by unnormalized Gaussian density. Additionally, see that exp (F(es)/e) vanishes infinity. Thus, for every 8" 0 there exists unnormalized Gaussian mixture ug "0(8") Nguuyen al. 2020 Theorem 5a) which is 8--close exp(f/<) in the uniform norm D;
| exp(f//)]bs>
sup (vg(2)) exp(f(en)/A) <
(31)
RIERD
"The result of (guyen et al. 2020) considers the approximation of normalized mixtures. This detail is not important and the result straightforwardly extends unoommaizedd mixtures. One can first normalize the mixture, approximate it and then re-scale both the target and the approximator back to the original seale.
20

---

Published as a conference paper at ICLR 2024
From the statement of the mentioned theorem it also follows that may pick all the covariances be scalar, i.e., 1z((51) Σk=1 (2N(11 (x1l%s cof )) for some K and 2gg E R.D, in the mixture "p (L,.... ..t....,K H). Indeed, just recall the definition of M in (Nguyen et al,, 2020 and oF E R2 (A: E
put g to be standard D-dimensional normal distribution. For further derivations, we pick
(2iTe) 2/2
8" del min {5/2(
-1
Emin
(32)
Zmin
and its respective mixture V0 with scalar components' covariances.
We define 22(21)
def
25(21) exp( ly ). It is again an unnormalized Gaussian mixture because it is product of two 2e
unnomalized Gaussian mixtures. Besides, it also has scalar covariances of its components because
multiplier exp(-
|| II" Y's covariance is scalar itself. More precisely, we have
2c
|1 12
K
|mmil? 2E
(x() det (z(21) exp(-
N.N(aalpu exp(
2e
(V2Te)*
D
BeN(zilu (
k=1
(Tt)) B0 c(1 oo)
Pkk
Ok
Σ
ax (N (33)
1+
o
d
Here in transition to 33
we use the formulas from Petersen et al. 2008 $8.188.. We derive that
2t 2=(00)= = / exp (f(es)/e) exp (
1/2||10 - x1
dx >
JRS
J (ug(zi) - 8d) exp ( 1/21110 - x211²
ddr
1 0z(21) exp (
1/211200 1r - F1 12 do 8" / exp ( 1/211.10 - xz1
dx
R20
=(2
| 1|
f vo(x1) exp (12o. -8" (2TE)) 0/2
exp(
2c
or, equivalently,
|.001
|
2²(z)) 8" (2(xT) 2/2 > exp(
vg (21) exp v1)o = exp(-
)co(xo)
2e
RD
2e
=co(xo)
(34)
Since + log z is a Euin support of Po. we may write
Lipschitz function on [Amin: +oc) and zJ 2xxo) 2 Zmin for To in the
8d (2TTT)) 2 log (Z((z)) + (2(c6) - log (z()))
2min
We use this inequality with 34 to derive
8" (2((TT))// log (z + (2TT))
||
log (zJ (zr))
+ log co
2min
2c
for vo in the support of Po- We integrate this expression for Fag ~ Po and obtain
log (200)oo(T0)d + /
8" (2(2TT) D/2
||
Po(xo)dro + / log co (2)))
2min
_BD
2c
21

---

Published as a conference paper at ICLR 2024
After regrouping the terms, get
|
(zo)doo
/
log (ogcolr))oozz))dda + (2TEE) 2/22
/ JD
log 23 (2o)
R2
2e
JRD
Zmin
(35)
Now we study another expression. Recalling that + log(z) is _ -Lisschittzz for 2 6min we get
Cam
8"
log (xp((()))))) log [exp(F(aa)/o) - 8"] = f(xa)/e log [ [exp(F(mn\/e) 8"1 (36)
Cmin
for all 21 in the support of P1. Here we also use the fact that
exp(f(m))<<) - 8" 2 exp(F(ms)/e) Cmin 2 20min
Crnin 2 Cmin
by the choice of 8", recall the definition of Cmin and see (32). We recall
31 to get
36
8"
log (g( log [exp(T(aa)/c) 5"1
M(m))/
Cmin
We exploit this observation to derive
/
|
/ log vo(2)) (on )dr =
/ec log (p >
2c pı (x1)de +
f(a1) 8"
f
f(x1)
8"
Pp (11)dr1
p1(z))dz
(37)
emin
Emin
We sum
(37 with 335 and get
=-2(0)
log (o p (x1)dxr
log co(xx) dro
- log zJ (ro)Po(zo)das) + / F(x1) Di (2) )d.
8"
8" (2TEe) 2/2
Cmin
2min
||
po(
pr (x) (x1)de1
2c
J80
2c
8"
8" 8"((TT)) P/2
|| po(
| | 112
+J
Pı (x) 1idddde
emin
Zmin
R
2c
JRD
2r
8e
8"
8" (2(TTE) 0/2
||
||
|
11
[Cost(T*) cH(pi)
Po Po(
pt (x (d)ddr
2
Cmin
Zmin
JRD
JRD
2c
/
|
1
||
1|
8
8"
8" (2TTe) D/2
-ost(z")-+FFt) -HH((p))
po(
Pi
(2) )kr
JRD
2
220
2
2
Gmin
Zmin
8"
8" (2TTc) P/2
(38)
Omin
Emin
where L matches the constant defined in (88 Indeed.
-L- -H( -E-1 S (2o. x1) =
=KI( (ppo H(())
(P1) + KL( llpo p1) --1 {
(20, (Fo. xn )drgdr
JRDRRD
'(Coa(sr) (CCo(())++1 + fgp (eool:Po((oo)dde 1/2 Jgo || (wr )dzs)
Cost(T") - <(H(p)) {
[
1
Do ( daoo - /
|
P1 (x )dx
RDD
2
JRD
2
22

---

Published as a conference paper at ICLR 2024
At the same time, from 19 and 38
we derive that
(2TTe) 2/2
6
8
KL ( mo) = ( L
5
+
Cmin
2min
2
2
Thus, Gaussian mixture is one that we seek for. This finishes the proof.
B.2 PROOF OF THEOREM A. APPENDIX A
The proof follows from combining the two auxiliary facts below.
Proposition B.1 (Rademacher bound on the statistical error). It holds that
E [[((0) [(+*)] v 4R (Vo Po) (RMM(V1p))
where Vo flog colo θ). V1 (log 08 0 Q and R (V.V))) p) denotes the well-celebrated Rademacher complexity (Shalev-hwwrrzz & Ben-David 2014 {26) of the functional class V w.r.l. to the sample size N of distribution p.
Proof of Proposition B.7 This result can be derived exactly the same way (Mokrov et al. 2024 Theorem 4) or (Taghvaei Jalali 2019 Theorem 3.4).
Proposition B.2 (Rademacher complexity bound for constrained log-sum-exp quadratic functions). Let < a A, let 0 < u S U, let 0 < w S W and V > O. Consider the class of functions
K
V = fo + log Σ
URæ z + of x + wy) with ag exp (I-Upz
(39)
k=1
uI UK- We say that such class is the class of constrained log-sum-ews quadratic functions. Assume that is compactly supported and the support lies zero-centered ball of radius O. Then
UT; [|v]|| V;W S wk W;a. on S A}.
Rx (v. S
VN' where the constant C depends only on K, u, U, o, A,V, w, W, P but not on the sample size N.
Proof of Proposition B.2 The Rademacher complexity of linear constrained functions + 1 + wg well known and is bounded by 0 (/))) (Shalev-Shwartz & Ben-David 2014 826.2).
The complexity of the constrained quadratic functions
x² URæ is also O ( which follows
from their representation using the Reproducing Kernel Hilbert spaces (RKHS), see (Latorre al.) 2021 Lemma 5 & Eq. 24) and additionally (Mohri et al. 2018 Theorem 6.12). Hence, by the well- known additivity of the Rademacher complexity also follows that the complexity of constrained forms + x U + oi x + wy is also bounded by O( ( Since x, Uk, We are bounded, the function x + cxp (a UN + of x + wx) is Lipschitz in x with the shared Lipschitz constant for all admissible Uk, Uk. wk. Therefore, the Rademacher complexity of such functions is also O ( recall the Talagrand's contraction principle (Mohri al. 2018 Lemma 5.7). The same applies to ag exp (~"U,æ + off wx) cxp (a Un vo + [wy log owl)
as these are also constrained exp-quadratic forms but with slightly adjusted constraints on the bias parameter Wi. Using the additivity the Rademacher complexity again, we see that K-sums
K
ag exp Uxx x + o x + wx)
of such functions also have complexity bounded by O (.) The remaining step is to note that each such function is both lower and upper bounded (by some positive numbers depending constraints), hence the logarithm of such functions is also a Lipschitz operation with some finite
the
Lipschitz constant. Thus, the complexity of + log Σ 44 exp (xT J + ofr + wx) is also O the constant hidden in O(-) incapsulates the dependedce on K, u, U, A, V. W, P. O
23

---

Published as a conference paper at ICLR 2024
Finally, can prove the bound on the estimation error in Theorem A.I
Proof of Theorem A.7 Just note that both Vo and 11 are the constrained classes of log-sum-exp quadratic functions in the sense of Proposition B.2 and apply Proposition B.1 For V1 this directly follows from the assumptions of the current Theorem. For Vo it follows from the the fact that co is also lug-sum-exp quadratic function with constrained parameters (our Proposition 3.2)
C EMBRYONIC STEM CELL DIFFERENTIATION SINGLE CELL DATA
For the embryonic stem cell differentiation single cell setup we use code and data from
to work with the embryonic stem cell differentiation dataset and to evaluate our light solver.
The provided data shows the cell differentiation collected at five different intervals (tg: day to 3, f1: day 6 to 9, t2: day 12 to 15, t3: day 18 to 21, tq: day 24 to 27). These collected cells were anal- ysed by scRNAseq subjected to quality control filtering and then represented feature vectors using Principal Component Analysis (PCA).
Solver
Wi metric
UFCFFM
0.79
[SF]-M-EXnG
0.795 =0.00 060
LightSB (ours)
0.823 +0.017
Reg. CNF
0.825 1
N.Nel
1.
Following the above-mentioned works, consider solving the problem of transporting the cell distribution at time 4-1 to time f++1 for E [1,2,3] Then we predict the cell distributions the
DSB
0.862 +0.023
FCFM
0.872 20.087
MSF] Geo
0.879 1.0.148
MSFT"M-Siak
1.198 +0.342
intermediate time & and compute the Wasserstein-1 (W, distance between the predicted distribution and the ground truth distribution. We average over all 3 setups i E 2, and present our results in Table 4 To estimate the standard deviation, we run experiments with different seeds for each i E [1, 2, 2]. For LightSB we use K = 100, Ir 10-2_ batch size 128 and do 2 10 gradient steps. We use results for other methods from Tong et al, 2023), whose authors were the first to consider this setup in (Tong et al. (2020). Our solver (c=0.1) performs at the level of the best other methods, while converging only in minute on CPU cores and learning only ~100 parameters.
SB-CFM
1.221 10.380
DSBM
1,755 #0.429
Table The quality of intermedinte distribution restoration of single-cell data by different methods.
0.1.
D DETAILS OF THE EXPERIMENTS D.1 GENERAL IMPLEMENTATION DETAILS
To minimize 10
we parameterize Th and S₃ of Ug 9 and use the Adam optimiser (Kingmal Ba 2014). We parameterize ag as using the logarithm log og; we parameterize r directly as vector; we parameterise the matrix S₁ in the diagonal form with the values log( (S&) ) on its diagonal.
Initialization. We initialize log ag by log K F& by using random samples from pi and log( (Sx,,,, by log 0.1 (it can be tuned hyperparameter but even without any tuning it works well with this initialisation on every considered experimental setup).
D.2 DETAILS TOY EXPERIMENTS
We use K 500 in all the cases. For 1 and e 10-2, we use lr 10-3 and for e 2-10-3
we use 10 and batchsize 128. We do 10s gradient steps.
D.3 DETAILS OF EVALUATION ON THE BENCHMARK
We use K = 50 in all the cases. For c = 10-1, we use Ir = 10-3, 3, For = 2- we use Adam with Ir 10 and batch size 128. We do 10 gradient steps.
In Table 5 we additionally present results of the non-conditional BW3-UVP metric. LightSB beats other solvers or performs at the same level for all setups.
24

---

Published as a conference paper at ICLR 2024
6=1.1
10
D
1= 16
D 128
17 16
17 64
D 128
D 16
64
D 128
solver
0.016
0.25
0.005
0.56
0.01
0.15
0.23
LightSB
0.005
0.017
0.037
0.069
0.004
0.01
0.03
0.07
0.03
0.04
0.17
0.30
+ std
=0.002
00.007
+0.007
+0.008
+0.002
=0.004
+0.006
=0.007
#0.01
20.01
=0.01
40.02
Table 5: Comparisons of BW3-UVP 4 (%) between the target PI and learned right marginal of my.
D.4 DETAILS OF MULTIMODAL SINGLE-CELL INTEGRATION EXPERIMENTS
We use data from the Kaggle competition "Open Problems Multimodal Single-Cell Integration".
The data describes gene expression of cells days 2, 3. 4 and which containt 6071, 7643, 8485, 7195 data points, respectively. Analogously to Tong et al. (2023), we use only CITEseq expression data; remove the donor-dependent bias we select only one donor with ID 13176. To preprocess the data, we use PCA projections with 50, 100, 1000 components. Then for each case we consider setups: data from day and day 4 as a distribution pair for the SB problem with data from day 3 for evaluation, and data from day 3 and day 7 as a distribution pair for the SB problem with data from day for evaluation. In each setup, we normalize the data by scaling it to the sum of each feature variance over the concatenated data from the start, end, and evaluation days (after PCA projection). For the evaluation, we take the prediction for the evaluation day for all cells from the start day and calculate the energy distance with the ground truth distribution.
For all described setups we use c = 0.1 in our and baseline solvers. For our LightSB solver we use K 10, Ir 10- and batchsize 128. We do 10 gradient steps.
Baselines. We compare LightSB with SB/EOT algorithms from three different classes: maximin (Gushchin et al. 2023a), Langevin-based (Mokrov al. 2024 and IPF-based (Vargas al. 2021 For completeness, we also add the popular discrete HOT Sinkhorn solver (Cuturi) 2013).
1. Maximin solver. For (Gushchin et al. 2023a) solver we use the official code from
https: //tpt/gttttt////// www.///./.c/.eeeee wwaoo/nguscccinng
coieurraapp ima wtw.:ppptmmmaaaaaaaaooomoo Transport
We use the same hyperparameters for this setup as the authors (Gushchin et al. 2023a Appendix E) use in their nigh-dimensiioal Gaussian setup. The only exception is the number of discretiza- tion steps N, which we set to 100 as well as for SB solver (Vargas et al., (202T) below.
based solver. For (Vargas et 2021) solver we use the official code from
Instead of Gaussian processes which the authors use, we use the same neural network as in (Gushchin et al. 2023b) to get better scalability. We use N 100 discretization steps, 50 IPF iterations, 10 epochs on the each IPF-iteration and 128 samples from distributions pu and p1 in each of them. We use the Adam optimizer with I = 10 for optimization.
3. Langevin-based solver. For (Mokrov al. 2024) solver, we use the official code from
guideduuidde Ent
We take the advantage of the author's setup from their 2D Gaussian-+Swinn lowing experimental framework, we adapt the original code by increasing the dimensionality of the learned fully-connected NN potentials. The chosen hidden directionalities for the poten- tials are [256, 256, 256] for D 50, 100 and [2048, 1024, 512] for D 1000. We choose all hyperparameters of the method in concordance with their code for ε = 0.1 EOT regularization coefficient, except Ir. We pick Ir 5 10-4 1 for training stability reasons. The numbers of training iterations are N = 8K for D = 50, 100 and N 4K for D = 1000. We get predictions for intermediate distributions by using Brownian Bridge (analogous to (1)).
experiment. Fol-
4. Discrete solver [ To run the Sinkhorn algorithm (Cuturi, 2013), we use the ot... sinkhorn with parameters met hod= "sinkhorn. log" and stop wwweshhhdeeeeeeeeee procedure from
Discrete OT neither can be straightforwardly used to infer trajectories for new oout-oftrain-samplee input cells, nor provides the trajectories for existing cells. In setup, the latter issue can be overcome by inserting
25

---

Published as a conference paper at ICLR 2024
Python OT Package (Flamary et al. 2021 Note the default threshold parameter is 10 but found that the algorithm stucks at tolerance z 10 10ma; hence, we increased the tolerance.
Remark. We use the full-dataset (a.k.a. full-batch) discrete OT. We found that for e 0.1 (which use for all the solvers) converges very slowly, even requiring more time to converge than our light solver in dimension 1000. This is explainable the convergence of Sinkhorn algorithm notably degrades when + O; it empirically seems like this degradation worse that in our solver. We demonstrate the convergence plots of the Sinkhorn vs. our light solver in Figure 4
Co Corvergance spoed compur DIK--50
Cervergancce spood cempo Ean DrM-100
Gorveegg spocd comserfion No1000
Snbhom (ghU) dime5g
Sekhaam IS1UL dine1000
Seehom ISPU) m0000000 UgriS (C"U) dmm-1000
ug152 (OhU) dim-00
Ugnis (OUU, dim-1000
24
23
7.1
2.2
23
2.2
Figure 4: Convergence speed comparison on MSCI dataset, starting day 3, ending day and evaluation day 4.
D.5 DETAILS OF IMAGE DATA EXPERIMENTS
We use official ALAE code and model from
and neural network extracted attributes for the FFHQ dataset from
wtatffff--fettueeeeeeeeeeeee dataset
We use K 10, Ir 10-3 and batchsize 128. We do 10 gradient steps.
E COMPLETE PARAMETERIZATION OF THE EOT PLAN
As we pointed in $3.1 our solver obtains an approximate density xg (2 (mg) of conditional dis- tributions π (xx (20) but does not recover the density of the entire plan ( xi) Rs T²" r0(2 æ1).
However, in some practical applications it may be needed to recover this density. Fortunately, this requires only minor modification of our light solver. Indeed, consider the parameterization
exp (Jna, (ll.)))))ee(()) cę(xo)
The To P. (mo)
(ao)
(40)
which is analogous to but we also introduce a density model Pa for the left marginal of the plan. By repeating the derivations of the proof of Proposition 3.1 it is not hard
that
KL (I [[ = KL (palle>) + [C(0) - --
(41)
which means that learning of parameters w turns to be separate density estimation problem. It can be solved, e.g., with the well-celebrated expectation-maximizztion algorithm (Dempster et al, 1977) for Gaussian mixture models with normalizing flow (Rezende & Mohamed 2015).
RELATED WORK: OTHER OT SOLVERS
For completeness of the exposition, we mention OT solvers which are not directly relevant to our EOT/SB setup because of considering non-entropic OT discrete OT settings.
the Brownian Bridge (analogously to [1 on of pairs of samples from the recovered discrete EOT plan (Stromme) 2023). Sampling from this bridge, one may construct approximation of the distribution at the intermediate time of interest.
26

---

Published as a conference paper at ICLR 2024
Discrete OT solvers. There exist many OT solvers (Cuturi 2013 (Nguyen et al. 2022 Xie et al. 2022) for the discrete OT setup (Peyré et al. 2019). This setup requires finding discrete matching between given train samples but does not require generalization on the test data. Hence, discrete solvers are not relevant to our continuous EOT/SB setup ($2) It is worth mentioning that there are several works studying the statistical properties of OT and developing out-of-sample estimators based on the discrete/batched OT solutions (Hütter & Rigollet,
Dvurechensky al. 2018)
2021 Pooladian & Niles-Weed, 2021 Manole et al. 2021 Deb al.) (021), Rigollet Stromme [2022) Fatras et al.] (2020). Despite having good theoretical properties, they mostly estimate the barycentric projection but not the entire plan .².
Other continuous OT solvers. Above in the work, we discuss the EOT/SB solvers but there ex- ist many papers proposing neural solvers for other OT formulations: unregularized OT (Henry-
Labordere 2019 Makkuva et al. 2020 Korotin et al. 202 b. 2022b Fan et al., 2023 Liu et al.)
2022 Gazdieva et al., 2023 Rout al. 2021
Amos 2022), weak OT (Korotin et al. 2023b a)
unbalanced OT (Choi et al. 2023 Yang & Uhler 2018 general OT (Asadulaev et al. 2024). These
works are of limited relevance as they do not solve EOT/SB.
G LIMITATIONS AND BROADER IMPACT
Limitations. We summarize and list some limitations of our light solver.
1. Analogously to the well-celebrated Sinkhorn algorithm (Cuturi, 2013) for the discrete OT, our solver may experience computational instabilities when applied very small regularization coef- ficients c > 0. This is due to the necessity to compute the exponent of values which proportional to co when computing Co or Ug in (8). However, actually works well for reasonably small e.
experiments show ($5), our light solver
2. Our main optimization objective (8) is not necessarily convex w.r.t. the parameters 0, i.e., the gradient-based optimization methods may experience local minima. While we mention this issue, we do not consider it serious enough: anyway, many existing machine learning algorithms have non-convex objectives (k-means, Gaussian mixture models, deep neural networks, etc.) and do not necessarily converge to the global optimum but still work well in downstream tasks. Note that the existing alternative continuous EOT/SB solvers ($4) also have non-convex objectives.
3. Our solver uses a kind of a Gaussian mixture approximation of EOT/SB which may be too restrictive to apply our solver to large-scale generative modeling problems unlike the complex neural EOT/SB solvers which we mention in $14 But this is very natural: default Gaussian mix- ture model for density estimation is also not used for modeling complex real-data distributions, e.g., images. At the same time, such models still play irreplaceable role in many smaller scale problems due to its extremal simplicity and ease of use. Therefore, we hope that our solver will play an analogous role in the field of computational continuous EOT/SB.
4. Our work considers SB with the Wiener prior W', which is equivalent to EOT with the quadratic cost c(z, 3) parameterization (9) is useful as it provides the closed-fon foreed-fmmmma expression for conditional distribu- tions no (at1 ( and drift go (Proposition 3.2). This is due the fact that the product of the unnor- malized Gaussian mixture density eo with the unnormalized normal density exp( 1 || æı ll?) is again a Gaussian mixture. We do not know if our light solver can be easily generalized to more general costs priors. We leave this question open for future studies.
[20 21112 and entropic regularization value In this case, the Gaussian mixture
Broader Impact. Deep learning models become more complex, and may negatively affect the Earth's ss ecology as the required computational clusters need increasing amounts of energy to work and water to cool them. fact, sometimes deep neural networks (DNNs) are applied when they may be unnecessary, so they pointlessly burn resources. Our work investigates SB and its applications to data, where DNNs are widely used. Our proposed light solver demonstrates that SB can be well-learned without DNNs and in few minutes even CPU. This helps to avoid tome-consuming GPU-training of previous solvers and decrease the negative impact on nature.
Potential negative broader impact of our research is the same that of most of the other ML searches. Namely, the constant advances in the field of the Al and the implementation of ML models in production pipelines may lead to transformation or replacement of some jobs in industry.
27

---

Published as a conference paper at ICLR 2024
H ADDITIONAL EXPERIMENTAL RESULTS
H.1 DEPENDENCE ON THE PARAMETER
In Figure 5 we show how the solution learned by LightSB depends on the parameter e in the Adult->child experiment. As expected, the diversity increases with the increase of e.
H.2
ADDITIONAL IMAGE GENERATION RESULTS.
In Figure 0 we show more samples from LightSB trained on every considered image setup.
28

---

Published as a conference paper at ICLR 2024
(a) LightSB Adult Child, = 0.1. Almost no diversity.
LightSB Adult Child, 0.5. Resonable diversity.
(c) LightSB Adult - Child, = 1.0. Moderate diversity.
(d) LightSB Adult Child, 10.0. High diversity.
Figure 5: LightSB Adult + Child for different e.
29

---

Published as a conference paper at ICLR 2024
(a) LightSB Male Female, 0.1 more samples.
LightSB Female Male, 0.1 more samples.
(c) LightSB Adult
Child, e more samples.
(d) LighTB Child + Adult more samples.
Figure 6: LightSB more samples for every considered setup.
30


Reviews:
[{"review_id": "iEOwitaSaN", "reviewer": "Reviewer_rvYk", "summary": "Summary\n1) The authors propose a novel light solver for continuous SB with the Wiener prior, i.e., EOT with the quadratic transport cost. The solver has a straightforward non-minimax learning objective and uses the Gaussian mixture parameterization for the EOT/SB, which avoids the time-consuming max-min optimization, simulation of the full process trajectories, iterative learning, and MCMC techniques that are in use in existing continuous solvers.\n2) The authors show that their novel light solver provably satisfies the universal approximation property for EOT/SB between the distributions supported on compact sets.\n3) The authors demonstrate the performance of the light solver in a series of synthetic and real-data experiments, including the ones with the real biological data considered in related works.", "questions": "Questions:\nHow is this method compared to Langevin dynamics/Langevin Monte Carlo methods in terms of runtime and empirical results?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths:\n- The authors provide a light solver for SB that does not rely on neural network parametrization.\n- This work provides a universal approximation for the solver which seems to be rather non-trivial. But to be honest, I'm not familiar with the proof so I'm not sure of the technical depth.\n- A minor thing but I really appreciate the authors giving a very clear description of the limitations of the solver.\n\nFollowing the rebuttal, the authors also provided statistical guarantees for the method.", "weaknesses": "Weaknesses:\nThis work does not have a finite-time nor finite-sample convergence guarantee. I believe that with additional assumptions one can obtain convergence guarantees as in [1]. Furthermore, there have been multiple works on the sampling complexity of quadratic cost EOT, and given the equivalence between EOT and SB, I believe that obtaining guarantees should be feasible.\n\nI also think that the experiments are not comprehensive enough. While the results are nice, I'm wondering what is the runtime performance of the light solver against other solvers. I cannot vouch for the method if I don't know how the method would improve in terms of the quality of the result and the computational cost.\n\nThe authors did mention that the work relies on Gaussian mixture parameterization and the entropic cost. It is indeed a limitation but personally, I think it is fine. Nevertheless, I will still raise this as a potential weakness.\n\nAll and all, I like the paper but it does not convince me enough to vouch for its acceptance yet.\n\nMinor comments:\nI think the literature on EOT is not sufficient. The authors should include gradient-based EOT solvers such as [2] (both of these methods achieved the optimal O(n^2/eps) complexity, which is stronger than that of APDAGD and seems to have good performance) and a gradient-based entropic UOT solver [3].\n\nReferences:\n[1] Chen, Y., Deng, W., Fang, S., Li, F., Yang, N. T., Zhang, Y., … Nevmyvaka, Y. (2023). Provably Convergent Schrödinger Bridge with Applications to Probabilistic Time Series Imputation. Retrieved from http://arxiv.org/abs/2305.07247\n[2] Xie, Y., Luo, Y., & Huo, X. (2023). An Accelerated Stochastic Algorithm for Solving the Optimal Transport Problem. Retrieved from http://arxiv.org/abs/2203.00813\n[3] \"On Unbalanced Optimal Transport: Gradient Methods, Sparsity and Approximation Error\".\nQuang Minh Nguyen, Hoang Huy Nguyen, Lam Minh Nguyen, Yi Zhou", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nie7wyjszo", "reviewer": "Reviewer_o9rU", "summary": "Current SB solvers often have complex neural parameterization and time-consuming training procedures due to minimax. This work proposed a novel parameterization technique and a non-minimax training method to bypass the issues above, with the aim at offering a lightweight and easy-to-use SB baseline.", "questions": "See the weaknesses above.", "limitations": "", "rating": 8, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strong motivation: the complex neural parameterization and time-consuming training procedures do hinder the application of current SB methods.\n\nClear presentation: the background knowledge and preliminaries are clearly presented. The connections with important references, such as (Tong et al., 2023), are highlighted and summarized in Table 1. The narrative is highly structural, with words and paragraphs to offer overviews of important sections.\n\nSoundness: the work seems highly self-contained: learning objectives, training/inference strategies, and theoretical properties. Limitations are also well discussed in appendix, some of which are examined through experiments.", "weaknesses": "1. Equation (5) seems incomplete or not well-defined. The equation is a crucial part of the methodology, and its clarity is essential for both understanding the method and replicating the results.\n\n2. It is pointed regarding other SB solvers that `they expectedly require time-consuming training/inference procedures.`. The authors state that their method avoids `time-consuming max-min optimization, simulation of the full process trajectories, iterative learning, and MCMC techniques,` which are commonly employed in existing solvers. While this claim holds conceptual interest, empirical evidence to substantiate this would significantly strengthen the paper.\n\nIn summary, although there exist some issues with this work, most of them seem correctable. I look forward to seeing the author feedback and the revised manuscript.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PFm3cOi7FH", "reviewer": "Reviewer_cSx4", "summary": "This paper proposes a fast Schrodinger Bridge (SB) solver based on the parameterization of Schrodinger potentials. SB is a dynamic version of the Entropic Optimal Transport (EOT), and there exists plenty of solvers for these problems, they requires complex parameterization by the form of neural networks / or is sensitive with the entropy regularization parameters, and thus are costly in terms of evaluation on larger scale datasets. The authors of this work propose a solution to this problem by consider a settings where the continuous SB is associated with a Wiener prior (as a reference measure). The key idea is to consider the parameterization of the Schrodinger potential as a mixture of Gaussian, and rewrite the SB optimization objective following this, with easy to compute mean and (scalar diagonal) covariance matrix. The authors provide some theoretical analysis of their algorithm, along with empirical demonstrations on synthetic dataset and realistic dataset (single cell data population dynamic and image-to-image translation).", "questions": "See weaknesses section.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 1, "strengths": "- Algorithms based on well-studied theory of SB and EOT.\n- The paper is well-written with clear structure.\n- Well-motivated problem: lightweight solver is much needed for problem that usually requires heavy computational power.", "weaknesses": "* **Novelty of the paper:** it seems like the paper is just a combination of the two previous works [1, 2]. Admittedly, the authors have clearly elaborated in their paper of such cases, but I do not see the clear novelty in the methodology part. Even on the proof of the universal theorem of the Gaussian mixture model, one can see that half of it is straightforward calculation from the two aforementioned paper.  \n* **Unclear benefit the of LightSB solver in realistic setting:** for the single cell dataset using W1 distance as a metric, LightSB's performance leaves some gaps with the two best methods, however I do not see author's comment about this. The other benchmark on unpaired image-to-image translation is very hard to judge, as for this task there is no quantitative metric to compare with other solvers. I do not know why the authors omit unconditional image generation tasks, as this is an important and popular benchmark that has FID as a standard metric. Morevoer, there have already exist results on some of the neural EOT solvers or the diffusion SB solver (using iterative proportional fitting)/diffusion SB matching (using iterative Markovian fitting), or flow matching/rectified flow in this task, so it is easy to compared with the baseline. \n*  **Questionable theoretical result:** in the proof the universality of the Gaussian mixture parameterization for SB (theorem 3.4), in the paragraph below equation (29), I fail to understand why the authors wrote \n\n> \"Besides, it also has scalar covariances of its components because multiplier $exp(− |x_1|^2 /2\\epsilon )$’s covariance is scalar itself\" \n\n, but what I understood in the paper's settings is that $x_1 \\sim \\pi_1$ is an unknown distribution, without assumptions on its parametric form. This is a key argument, as without it the factorization failed to be what the authors claimed in the statement of the theorem, I hope the authors could clarify this to me, otherwise it would be a hole in the proof of an important analysis of the paper.    \n\n[1] Petr Mokrov, Alexander Korotin, and Evgeny Burnaev. Energy-guided entropic neural optimal transport. arXiv preprint arXiv:2304.06094, 2023.  \n[2] Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems, 2023a.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5jdbuMGSff", "reviewer": "Reviewer_KWE5", "summary": "This paper proposes a novel method to approximate the Schrödinger Bridge problem. Following the connections of the Schrödinger Bridge problem to entropic optimal transport, the authors want to learn the true entropic OT plan. They parametrize the optimal entropic OT plan as the product of two probabilities, the source measure $p_0$ and a conditional plan $\\pi(x_1|x_0)$ (up to a normalization constant). Following this novel parametrization, they show that they can optimize the parametrized plan without the knowledge of the true entropic OT plan. To deal with the normalization constant, they use a mixture of Gaussian representations to get a tractable form of each of the components. They then explain the inference and training procedure. Finally, they perform several experiments to show the practicability of the proposed method (2D synthetic data, an EOT benchmark, single-cell dynamics and unpaired image-to-image experiment).", "questions": "1. How does your method depend on the dimensionality of the data? I think that this is a legitimate question, especially with the normalization constant being approximated by a mixture of Gaussians. I recommend adding more dimensionality in some experiments (like single-cell by considering a larger number of principal components. Maybe 5, 50, 200, 500, 1000?) to study how the proposed method performs as the dimension grows. I also recommend adding this discussion to the limitation paragraph in Appendix E. \n\n2. On the single-cell experiment of the unpaired image-to-image experiment, could you show the training speed to reach convergence and the number of iterations it took? I checked the appendix and I did not find such plots. Maybe it would be interesting to compare to other OT solvers (sinkhorn or stochastic variants) on a simple problem to see the different behaviour. (I acknowledge that the Sinkhorn solvers are only usable in a discrete setting.) \n\n3. I found the limitation discussion in Appendix E interesting and it could have been in the main paper. I recommend moving it to the main paper. Maybe some discussions about related work could go in the appendix instead. \n \nI am ready to reconsider my score if the authors can reproduce some of the competitor results and consider the highest dimensional setting of single-cell data experiments (dimension of +1000). It will depend on how well the proposed algorithm will behave on medium dimensionality.\n\n\n----- EDIT POST REBUTTAL -----\n\nThank you for your answer. I have read the rebuttal.\n\n[Single-cell experiments] Thank you for the novel experiments on the single-cell trajectory problem. I find them interesting and encouraging, especially with the different dimensions. \n\nI understand the motivation for using a different metric than W_1 to compare the different approaches. However as it is not the standard metric, it is hard for me to compare your method with other standard single-cell trajectory methods. As the authors-reviewer discussion has ended (due to a late rebuttal submission), I would have appreciated seeing the Wasserstein 1 metric for the different approaches. Indeed, the novel dataset was first considered by [Tong et al., 2023] where they used the W_1 distance. As the authors did not reproduce their method with the considered metric, it is hard to compare with more standard approaches on single-cell datasets like ODE/SDE-based approaches [Tong et al., 2023]. Unfortunately, two of the considered competitive methods are not standard in the single-cell trajectory literature. Therefore, I still think that the experimental section lacks reproduced competitors especially as the metrics are different.\n\n[Novel theoretical section] Thank you for the novel appendix H. It brings some light to the proposed methodology. I suggest the authors include Theorem H.1 in the main paper.\n\nIn my opinion, this is a borderline paper. I acknowledge that the proposed approach is new and interesting but I still think that it lacks a more rigorous experimental section to understand its practical performance with the current literature. Therefore, I will keep my score.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "i) The parametrization of the entropic OT plan is novel (to the best of my knowledge) and interesting.\n\nii) The derived loss is very interesting as learning the parametrized model does not require knowing the true entropic OT map.\n\niii) The discussion on the normalization constant and the proposed parametrization (with a mixture of Gaussians) to overcome this issue is appealing (even if the computation of the normalization constant is based on existing works).\n\niv) The method has been tested in different experiments.\n\nv) The paper is clear and easy to read.", "weaknesses": "1. In the single-cell experiments, the competitor results are taken from another paper [Tong et al., 2023]. As the evaluation is performed on a leave-one-domain out and training on the others, it is questionable that the training procedure was the same. I believe that some of these competitors should be reproduced by the authors to ensure that the training setting is similar, especially as the unpaired image-to-image experiment is a quantitative experiment without competitors.\n\n2. The experiments were performed in a relatively small dimensional setting: the unpaired image-to-image experiments used a pre-trained feature extractor of dimension 512, the single-cell data experiment used the representation of the 5 whitened principal components (ie dimension of 5), and the EOT benchmark data were (at most) of dimension 128. Therefore, the question of the performance of the proposed method in high-dimensional data is legitimate, especially as the authors use a mixture of Gaussian representation.\n\n3. Little is said about the statistical estimation (with respect to the data dimension and sample size) of the proposed parametrization to the true entropic OT plan $\\pi^\\star$. It is known that it also suffers from the curse of dimension [2,3].\n\n4. Some related work is missing and should be discussed and mentioned [1,2,3,4]\n\n[1] Stochastic optimization for large-scale optimal transport, Genevay et al, Neurips 2016\n[2] On the sample complexity of entropic optimal transport, Rigollet et al.\n[3] Minimax estimation of smooth optimal transport maps, Hutter et al.\n[4] Learning with minibatch Wasserstein, Fatras et al., Aistats 2020", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Jp8ToWoM0p", "reviewer": "Reviewer_o21Y", "summary": "The paper proposes a fast, light-weighted, solver for Schrödinger bridge by using a Gaussian mixture parametrization on the energy potentials. This simplifies computation of the static EOT map that would otherwise be intractable, resulting in efficient learning process. Extensive experiments are conducted on low/mid dimensional benchmark, single-cell dataset, and unpaired image translation in latent space.", "questions": "- Can the author provides image experiments without the latent space? I suggest smaller dataset such as AFHQ 32 or 64 for faster evaluation given the limited rebuttal period. Otherwise, can the author provide and include discussions on the scalability of the proposed method? Given Thm 3.4, it seems like the method could be applied to these scenarios.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed method is notably simple and elegant. It effectively combines key insights from previous works and elevates them to address a significant problem within the SB community.\n\n- The paper is well-written, and the thorough comparison to related works, especially in Table 1, is particularly valuable. Additionally, the comprehensive discussions in the appendix about limitations and broader impacts are appreciated.", "weaknesses": "- The proposed method on image dataset requires pretrained latent space (512 dimension) that is already structurally informative.\n\n- On the discussion of tractable / real-world SB given pairing, a few important references such as \"Aligned SB\" and \"image-to-image SB\" are missing. \n\n- Given that the proposed method is computationally light weighted, it'll be beneficial to have some quantitative comparison (actual runtime, memory etc) to prior works.\n\n- All \"Schrodinger\" should be changed to Schrödinger.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)