Entropic Neural Optimal Transport via Diffusion Processes
Authors: Nikita Gushchin
Alexander Kolesov
Alexander Korotin
Dmitry P. Vetrov
Evgeny Burnaev
Venue: NeurIPS 2023 oral (2023)
Keywords: Optimal transport
Schrödinger Bridge
Entropy regularized OT
Neural Networks
Unpaired Learning

Abstract: We propose a novel neural algorithm for the fundamental problem of computing the entropic optimal transport (EOT) plan between probability distributions which are accessible by samples. Our algorithm is based on the saddle point reformulation of the dynamic version of EOT which is known as the Schrödinger Bridge problem. In contrast to the prior methods for large-scale EOT, our algorithm is end-to-end and consists of a single learning step, has fast inference procedure, and allows handling small values of the entropy regularization coefficient which is of particular importance in some applied problems. Empirically, we show the performance of the method on several large-scale EOT tasks. The code for the ENOT solver can be found at https://github.com/ngushchin/EntropicNeuralOptimalTransport

Entropic Neural Optimal Transport via Diffusion Processes
Nikita Gushchin Skoltech"
Alexander Kolesov Skoltech" Moscow, Russia a. Moolesov@skoltech.crg
Alexander Korotin
Skoltech, AIRI Moscow, Russia
Moscow, Russia gushchindskoltech. ru
korotin@skoltech..rr
Dmitry Vetrov
Evgeny Burnaev
HSE University! AIRII
Skoltech; AIRI' Moscow, Russia burnaev@@kolteehh.rrr
Moscow, Russia vetrovd@yandex.rg
Abstract
We propose a novel neural algorithm for the fundamental problem of com- puting the entropic optimal transport (EOT) plan between continuous proba- bility distributions which are accessible by samples. Our algorithm is based the saddle point reformulation of the dynamic version of EOT which is known as the Schrödinger Bridge problem. In contrast to the prior methods for large-scale EOT, our algorithm is end-to-end and consists of a single learn- ing step, has fast inference procedure, and allows handling small values of the entropy regularization coefficient which is of particular importance in some ap- plied problems. Empirically, we show the performance of the method on sev- eral scale EOT tasks. The code for the ENOT solver can be found https: / thub. com httpg//wwhhhhhhtttoiiioaaaaaaaaa
X ~ Po t=0.1 t=0.2 t=0.3 t=0.4 t=0.5 t=0.6 t=0.7 t=0.8
t=0.9
t=1.0
E= 0
E=1
==
Figure 1: Trajectories of samples learned by our Algorithm for Celeba deblurring with 10.
1
Introduction
Optimal transport (OT) plans fundamental family of alignments between probability distributions. The majority of scalable neural algorithms to compute OT are based on the dual formulations of OT, see (28 for survey. Despite the success of such formulations in generative modeling 42 BII, these dual form approaches can hardly be generalized to the popular entropic OT [12. This is due to the numerical instability of the dual entropic OT problem [14] which appears for small entropy
"Skolkovo Institute of Science and Technology i Artificiil Intelligence Research Institute THSE University
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

regularization values which are suitable for downstream generative modeling tasks, At the same time, entropic OT is useful as it allows to learn one-to-many stochastic mappings with tunable level of sample diversity. This is particularly important for ill-posed problems such as super-resolution [37]
Contributions. We propose a saddle-point reformulation of the Entropic OT problem via using its dynamic counterpart known as the Schrödinger Bridge problem ($4.1) Based on our new reformulation, we propose novel end-to-end neural algorithm to solve the related entropic OT problem for pair of continuous distributions accessible by samples ($4.2). Unlike many predecessors, our method allows handling small entropy coefficients. This enables practical applications to data + data mapping tasks that require slight variability in the learned maps. Furthermore, we provide an error analysis for solving the suggested saddle point optimization problem through duality gaps which are the errors of solving the inner and outer optimization problems ($4.3). Empirically, we illustrate the performance of the method on several toy and large- scaeesccal EOT tasks
2 Background
Optimal Transport (OT) and Schrödinger Bridge (SB) problems imply finding an efficient way to transform some initial distribution Fo to target distribution P1. While the solution of OT only gives the information about which part of Po transformed to which part of Pi, SB implies finding a stochastic process that describes the entire evolution from Po to P1. Below we give an introduction to OT and SB problems and show how they are related. For a detailed overview of OT, we refer to and of SB to [32
19
2.1 Optimal Transport (OT)
Kantorovich's 's OT formulation (with the quadratic cost). We consider D-dimensional Euclidean spaces X, y' and use P2) P2())) denote the respective sets of Borel probability distributions on them which have finite second moment. For two distributions Po E P2) PI E P22))), consider the following minimization problem:
1 - y11²
inf rEIT(Po. xxx
dm (2((,1)))
where II,o,, Fi) c P2(X' D) is the of probability distributions on X' y with marginals Po and P1. Such distributions π E [I(Po, P1) are called the transport plans between Po and P1. The set II F1) is non-empty as it always contains the trivial plan Fo P1. A minimizer x¹ of ( always exists and is called an OT plan. If is absolutely continuous, then x² is deterministic: its conditional distributions are degenerate, i.c., TT (+|x)= = ôT-(2) for some T X' + y (OT map).
Entropic OT formulation. We use (TT to denote the differential entropy of distribution and KL(TT|<)) to denote the Kullback-Leibler divergence between distributions т and n'. Two most popular entropic OT formulations regularize (I) with the entropy H or KL-divergence between plan π and the trivial plan P. x P2, respectively (e > 0):
I
x - vll
inf TEIL(Po.P)) /
di d v) - =H((T),
(2)
XxY
2
I v11
inf xEIl(Fo.P.) xxy
di 8) + +KL x P1).
(3)
2
Since и E IT(Po, IP)), it holds that are equal up to an additive constant when Po and IP, are absolutely continuous. The minimizer of (2) is unique since the functional is strictly convex in thanks to the strict convexity of H (T). This unique minimizer x² is called the entropic OT plan.
i.e. both formulations
2.2 Schrüdinger Bridge (SB)
SB with the Wiener prior. Let Q be the space of IR. valued functions of time [0, 1] describing some trajectories in RDD, which start at time = 0 and end at time ( t= 1. We use P(2) to denote the Ω. We use dW; to denote the differential of the standard Wiener process. Let W E P(2) be the Wiener process with the variance e which starts at Pg. This diffusion process can be represented via the following stochastic differential equation (SDE):
set of probability distributions
W : dxx, VedW,, No ~ Po.
(4)

---

We use Q)(T]IQ) to denote the Kullback-Leibe divergence between stochastic processes 7' and Q. The Schrödinger Bridge problem was initially proposed in 1931/1932 by Erwin Schrödinger 44
can be formulated as follows [1 Problem 4.1]:
inf WL(T (1)
(5)
where F(PPo P1) C P(22 is the set of probability distributions on Q having marginal distributions Po and P1 t and = respectively. Thus, the Schrödinger Bridge problem implies finding stochastic process 7' with marginal distributions Fo and IP, at times = 0 and = 1, respectively, which has minimal KL divergence with the prior process W
Link to OT problem. Here we recall how Scrödinger Bridge problem (5) relates to entropic OT problem 2 This relation is well known (see, e.g.. 132) [11] Problem 4.2]), and we discuss it in detail because our proposed approach ($4) is based on it, Let x.' denote the joint distribution of a stochastic process T at time moments & 0, and = πं denote its marginal distributions time moments f = 0, 1, respectively. Let 1 denote the stochastic processes conditioned on values x, at times 0, 1, respectively. One may decompose (L(TI[W>s) as [48 Appendix C):
S Jxxy (L
KL(TIIW9) KL KL(""||
[ | (r.y (r.y)dd ddTTd...)))) (z.y).
(6)
i.e. KL divergence between T and W is
of two terms: the first represents the similarity of the processes at start and finish times / = 0 and / = 1, while the second term represents the similarity
of the processes for intermediate times (0, 1) conditioned on the values t 0, 1. For the first term, it holds (see Appendix A [I]] Eqs 4.7-4.91):
1 -l1l 2E
KL((T" KL(T"I|>>>
dn d
((T)) ) +
(7)
Xxx
where C' is a constant which depends only on Po and e. In 32 Proposition 2.3], the authors show that if 7 the solution to 45), then Tiz. W Ja.y Hence, one may optimize (5) over processes T
for which T TI W( for every x, y and set the last term in (6) zero. In this case:
Ja.y
inf KL(T]|W">) = TEF(Fo.P1)
KL(T T TeF.(P.PP))
inf
inf KL(# 117""117")
(8)
x" P.)
ie., it suffices to optimize only over joint distributions xJ at time moments t = 0, 1. Hence, minimizing (8) is equivalent (up to an additive constant C') to solving EOT ( divided by e, and their respective solutions N- and x* coincide. Thus, SB problem (5) can be simplified the entropic OT problem 2 with entropy coefficient e.
Dual form of EOT. Entropic OT problem 8 has several dual formulations. Here recall the one which is particularly useful to derive our algorithm. The dual formulation follows from the weak OT theory 7 Theorem 1.3] and explain it in detail in the proof of Lemma B.3
ssup ( 1/ 8 (x)dPo(x) + / ((y)dP
inf KL(T|IT" W rEI(P..P.))
B
where (x) def inf infrep.o) [KL( W (++x) - Sy ((y)du(u)) The sup here taken over B
belonging the set of functions
(,, def y continuous s.t. Bu, U, w :: u/l 11 1 3() S w).
i.e., B should be continuous with mild boundness assumptions.
Dynamic SB problem (DSB). It known that the solution to the SB problem (5) belongs the class D(PP) of finite-energy diffusions T', [32] Proposition 4.1] which are given by:
Ty : dX, f(X,, + VedWe, Xo Po. ET, Er../
L [If((. < co.
(9)
where f: P x 10 + RDD is the drift function. The last inequality in means that Tj finite-energy diffusion. Hence, optimizing only over finite-energy diffusions rather than all possible
3

---

processes is enough to solve SB problem (5). For finite-energy diffusions, it is possible to rewrite the optimization objective. One can show that (T, W() between processes T and W is 40)
1 Er, Fr/l, ./ [I((x,,
KL(TyIIW>> = 2r
(10)
By substituting (10) in (5) SB problem reduces to the following problem [I] Problem 4.3]:
1T, L. t)l
inf P))
inf KL(T)I]W>)) P.) 2
(11)
where P1) P(2) is the set of finite-eneryy diffusion on 52 having marginal distributions Po and 11 at t = 0 and f t= respectively. Note that D(Fo, P1) c o(Poo. Since the problem (IT) is the equivalent reformulation of (5), its optimal value also equals (8) However, solving this problem,
as well as ( is challenging it is still hard to satisfy the boundary constraints.
3 Related Work
In this section, we overview the existing methods to compute the OT plan or map. To avoid any confusion, we emphasize that popular Wasserstein GANs [S] compute only the OT cost but not the OT plan and, consequently, are out of scope of the discussion.
Discrete OT. The majority algorithms in computational OT are designed for the discrete setting where the inputs Po, Pi have finite supports. In particular, the usage of entropic regularization (2) allows to establish efficient methods [13] to compute the entropic OT plan between discrete distributions with the support size up to 100-106 points, see [41] for survey. For larger support sizes, such methods are typically computationally intractable.
Continuous OT. Continuous methods imply computing the OT plan between distributions Po. Pi which are accessible by empirical samples. Discrete methods "as-is" are not applicable to the continuous setup in high dimensions because they only do a stochastic matching between the train samples and do not provide out-of-sample estimation. contrast, continuous methods employ neural networks to explicitly or implicitly learn the OT plan map. As result, these methods can map unseen samples from Po to P1 according to the learned OT plan or map.
There exists many methods to compute OT plans [51 36 38 26
28 27
16
19 18 but they
particular, they mostly focus on computing the deterministic OT plan (map) which may not exist. Recent works [31 B0 design algorithms to compute OT plans for weak OT 20, 77. Although weak OT technically subsumes entropic OT these works do not cover the entropic OT 2 because there is no simple way to estimate the entropy from samples. Below we discuss methods specifically for EOT and DSB (TT.
consider only unregularized OT (I) rather than the entropic one
3.1 Continuous Entropic OT
In LSOT [45), the authors solve the dual problem to entropic OT ( The dual potentials are then used to compute the barycentric projection æ + Jy y dr (v|zx) i.c., the first conditional moment of the entropic OT plan. This strategy may yield a deterministic approximation of n" for small e but does not recover the entire plan itself.
In 14 Figure 3), the authors show that the barycentric projection leads to the averaging artifacts which make it impractical in downstream tasks such as the unpaired image super- ruser-resolution. To solve these issues, the authors propose a method called SCONES. recovers the entire conditional distribution x (v(zz) of the OT plan n² from the dual potentials. Unfortunately, this is costly. During the training phase, the method requires learning a score-based model for the distribution Pp. More importantly, during the inference phase, one has to run the Langevin dynamic to sample from x' (v(2:)
The optimization of the above-mentioned entropic approaches requires evaluating the exponent of large values which are proportional to -1 [45] Eq. 7]. Due to this fact, those methods are unstable for small values e in (2) In contrast, our proposed method (#4) resolves this issue: technically, works even for c = O.
3.2 Approaches to Compute Schrödinger Bridges
Existing approaches to solve DSB mainly focus on generative modeling applications (noise -> data). For example, FB-SDE [10] utilizes data likelihood maximization to optimize the parameters of

---

forward and backward SDEs for learning the bridge. On the other hand, MLE-SB [88] and DiffSB [5] employ the iterative proportional fitting technique to learn the DSB. Another method proposed by (50 involves solving the Schrodinger Bridge only between the Dirac delta distribution and real data to solve the data generation problem.
Recent studies [34 33 have indicated that the Schrödinger Bridge can be considered as a specific instance of more general problem known as the Mean-Field Game []] These studies have also suggested novel algorithms for resolving the Mean-Field Game problem. However, the approach presented by [3 cannot be directly applied to address the hard distribution constraints imposed on the start and final probability distribution as in the Schrödinger Bridge problem (see Appendix H). The approach suggested by [34] coincides with that proposed in FB-SDE [10] for the SB problem.
The Algorithm
This section presents novel neural network- based algorithm recover the solution T,. of the DSB problem (TT) and the solution x of the related EOT problem (8. $4.1 we theoretically derive the proposed saddle point optimization objective. In $44.2 we provide and describe the practical learning procedure to optimize it. In $4.3), we perform the error analysis via duality gaps. In Appendix B we provide proof's of all theorems and lemmas.
4.1 Saddle Point Reformulation of EOT via DSB
For two distributions Fo and P1 accessible by finite empirical samples, solving entropic OT "as-is" non-trivial. Indeed, one has to (a) enforce the marginal constraints и E ([ and (b) estimate the entropy (T)) from empirical samples which challenging. Our idea below is employ the relation of EOT with DSB to derive an optimization objective which in practice can recover the entropic plan avoiding the above-mentioned issues. First, we introduce the functional
II(T)
ET, 2e /' (I(Xx,,)1l2 - -/, B(w)d=) (v) + J (w)ddr(()) (12)
(TB,Ty) del
This functional can be viewed as the Lagrangian for DSB 11 with the relaxed constraint d (y) d.P. (2), and function 3 : + R (potential) playing the role of the Langrange multiplier.
Theorem 4.1 (Relaxed DSB formulation). Consider the following saddle point optimization problem:
sup inf C(),Ty)
(13)
8 T₁
where sup is taken over potentials 3 E 2, and inf is taken over diffusion processes Ty E D(Po).
Then for every optimal pair (3", (B+TTry)) for (1). i.e., and Ty- E arginf ,(B
argsup inf C(B,Tg)
B T;
T
it holds that the process T3 is the solution SB 17
Corollary 4.2 (Entropic OT as relaxed DSB). (33 solves 13 then TTz- y is the EOT plan 2)
Our results above show that by solving
13 one immediately recovers the optimal process Ty- in
DSB 11 and the optimal EOT plan nt Tj. The notable benefits of considering 13) instead
of (2) and (1] is that (a) it is as an optimization problem over (3, Ty) without the constraint
dm (v) d.P (y). and (b) objective 12 admits Monte-Carlo estimates by using random samples from Po: Ty, IP. In beloow we describe the straightforward practical procedure to optimize (13) with stochastic gradient methods and neural nets.
Relation to prior works. In the field of neural OT, there exist so many maximin reformulations of OT (classic (28 42 [6 19 22), weak B1 30), general (6)) resembling our (13) that reader may naturally wonder (1) why not to apply them to solve entropic OT? (2) how does our reformulation differ from all of them? It is indeed true that, e.g., algorithm from 3I 6 mathematically covers the entropic case (2). Yet in practice requires estimation of entropy from samples for which there is no easy way and the authors do not consider this case. These methods be hardly applied to EOT.
In contrast to the prior works, we deal with entropic OT through connection with the Schrödinger Bridge. Our max-min min reformulation is for DSB and it avoids computing the entropy term H (*)
5

---

Il dt which can be in EOT. This term is replaced by the energy of the process E,, [L. straighfforwardly estimated from the samples of Ty, allowing to establish a computational algorithm.
4.2 Practical Optimization Procedure
To solve 12 we parametrize drift function f(z, t) of the process T and potential B( I by neural
nets fo R
x x10,1] + R and Bo IR_ - IR. We consider the following maximin problem:
sup inf { 1 Bry, B 2e
1// [I,o((X, (l
/, Bo(v))B((y) - J, Bolu)d Ti (v)
(14)
8 To
We use standard Fuler-Maruyama (Ful-Mar) simulation (Algorithm 2 in Appendix C for sampling from the stochastic process T by
Algorithm 1: Entropic Neural OT (ENOT)
Input:
samples from distributions Po. P1;
solving its SDE (9). To estimate the value of Ery (G (X,, 1) ()] in 140. we utilize the mean value of [1, t)ll over time of trajec- tory X, that is obtained during the simulation by Euler-Maruyama algorithm (Appendix (c) We train fo and & by optimizing (12) with the stochastic gradient ascent-descent by sampling random batches from Po. PI. The optimiza- tion procedure is detailed in [ We use fr.m denote the drift at time step for the m-th object of the input sample batch. We use the averaged of the drifts an estimate of fo (l, ))l dt dt in the training objective.
Wiener prior noise variance c 2 0;
drift network fa R x [0, 1] - RD:
beta network 3 : IR + R:
number of steps N for Eul-Mar (App C
number of inner iterations Ky.
Output: drift fă of TTa solving DSB repeat
Sample batches Xo ~ Po. Y P1;
[Xm: Jn <- Ful-Mar( (x, Tj.); Lg + [x ZEXN
Σ ()(z))
M Σ Ap(y):
yEY
Update by using aLg de
for k = to K do Sample batches Xo Po. Y ~ Pi;
Remark. For the image tasks ($5.3) $5.4.. we find out that using a slightly different parametrization of T3 considerably improves the quality of our Algorithm, see Appendix
{Xn, J. t Eul-Mar( (o, Tio):
N-1
fal
KL
+ 1 Σ
llfm,m|l?
N
==0
m=1
It should be noted that the term LJ, (X,, dtt] is not multiplied by 2. in the algorithm since this does not affect the optimal Ty. solving the inner optimization problem, see Appendix D
C f- KL [x
Σ Bo(x);
=EXN
Update 0 by using
OLo.
30
until converged;
Relation to GANs. At the first glance, our method might look like typical GAN as it solves maximin problem with the "discriminator" Agg. and SDE "generator" with the drift fo. Unlike GANs, in our saddle point objective (12), optimization of "generator" Ty and "discriminator" are swapped, i..., "generator" is adversarial to "discriminator", not vise versa, as in GANs. For further discussion of differences between saddle point objectives of neural OT/GANs, see 31, 64.3), 42 54.3), [16].
4.3 Error Bounds via Duality Gaps
Our algorithm solves maximin optimization problem and recovers some approximate solution (â, TT)) Given such pair, it is natural to wonder how close is the recovered T3 to the optimal Ty-. Our next result sheds light on this question via bounding the error with the duality gaps.
Theorem 4.3 (Error analysis via duality gaps). Consider pair (B, Ty). Define the duality gaps, i.e.
errors of solving inner and outer optimization problems by:
def C(3, (B) =
inf ((,TTT). T/
and
E2 sup inf ((B,T,) - inf C(3,TT) T)ED(Po)
def
(15)
Then holds that
Prv(T Ty.) (.)) S VEI and
Prv( " Tr) S VEI +
(16)
where we use -Tv(:-)) to denote the total variation norm (between the processes or plans).
In practice, Po ()) since can choose 0, v min min(float32), max(float32).

---

Fiudd distr Eutiion
Riled distr out
Fitee distnibttion
Vinput dictribttion
-2
-2
-1
Target distribution
Tajectoriss dogetidoto
Trajectories its Target data
Trajectories
-2
(a) Po. y
(b) ENOT (ours), (c) ENOT (ours), 0.01 (d) ENOT (ours),
Figure 2: Gaussian > Mix of 8 Gaussians. The process learned with ENOT (ours) for 0=0,0.01,
Relation to prior works. There exist deceptively similar results,
38 Theorem 3.6), 42 Theorem
4.3], 66 Theorem 4], 6 Theorem 3]. None of them are relevant to our EOT/DSB case.
In 38
[6], the authors consider maximin reformulations of unregularized OT (1. i... non- their notation) to be a convex function which in
entropic. Their result requires the potential |3 (f.
practice means that has to employ ICNNs [4] which have poor expressiveness 22 17 26] Our result free from such assumptions on B. In (6), the authors consider general OT problem 39] and
require the general cost functional to be strongly convex (in some norm). Their results also do not apply to our case as the (negative) entropy which we consider is not strongly convex.
5 Experimental Illustrations
In this section, we qualitatively and quantitatively illustrate the performance of our algorithm in several entropic OT tasks. Our proofs apply only to EOT > 0), but for completeness, present results = 0, i.e., unregularized case (T). Furthermore, we test our algorithm with c = on the Wasserstein-2 Benchmark (28). see Appendix [ We also demonstrate the extension of our algo- rithm to costs other than the squared Euclidean distance in Appendix [] The implementation details given Appendices and The code written PyTorch and is publicly available
also
https: http:///tttb.ccooogghhha//raratt
5.1 Toy 2D experiments
Here we give qualitative examples of our algorithm's performance on toy 2D pairs of distributions. We consider two pairs Po, P1: Gaussian - Swiss Roll, Gaussian + Mixture of Gaussians. We provide qualitative results in Figure 2 and Figure (Appendix E}. respectively. In both cases, we provide solutions of the problem for 0, 0.01, 0.1 and sample trajectories. For e 0, all the trajectories are straight lines as they represent solutions for non-regularized OT (T) see 443 $5.4]. For bigger e, trajectories, as expected, become more noisy and less straight.
5.2 High-dimensional Gaussians
For general continuous distributions Po. P1, the ground truth solution of entropic OT (2) and DSB (IT is unknown. This makes challenging to assess how well does our algorithm solve these problems. Fortunately, when P. and P1 are Gaussians, there exist closed form solutions of these related problems, see (55 and [9] Thus, to quantify
performance of our algorithm, we consider entropic OT problems in dimensions D E (2, 16, 64, 128) with 1 for Gaussian PP A0 Σo), P. P0, E. We pick E.g, Σn at random: their eigenvectors are uniformly distributed on the unit sphere and cigenvalues are sampled from the loguniform distribution log 2, log 2).
7

---

Dim
16
64
128
Dim
2
16
64
128
0.09
0.012
0.03
ENOT (ours)
(#-00006)
(#0.02))
(*0.00)
(#0.08)
(ours)
(=0.003)
(#0.01))
(+0.01.4)
(+0.004)
LSOT 43
1.82
6.42
32.18
64.32
LSCT [S
6.77
14.56
25.56
47.11
SCONES
1.87
6:27
6.88
SCONES
0.92
5.33
MLEESB 48
0.41
0.50
1.16
2.13
MI.F-SR
0.3
0.9
1.34
1.8
DirSB
1.11
2:20
DJRSB
1.7
2.43
FR-SDE-A
0.87
0.94
1.85
1.95
FR-SDE-A
0.75
1.36
2.45
2.64
FB-SDEJ
(0)
0.03
0.05
0.19
0.39
FB-SDEJ
10
0.07
0.22
0.34
0.58
Table I: Comparisons of B + (%)
Table 2: Comparisons of BW3 3VV-UVP + (%)
between the target P1 and learned marginal TT.
between the the EOT plan and learned plan π
time
0.2
0.4
0.6
0.8
0.01
0.023
0.042
0.067
0.096
ENOT (ours)
(1.0.001)
(1.0.005)
(1.0.007)
(10.015)
(10.019)
1SOT [45]
N/A
N/A
N/A
N/A
6.42
SCONES [4]
N/A
N/A
N/A
MLE-SB 48
0.10
0.23
0.30
0.36
0.50
DiASB [5
0.19
0.48
0.68
0.91
1.11
FB-SDE-A TO]
0.17
0.45
0.61
0.77
0.94
FB-SDEJ 10
0.18
0.32
0.31
0.17
0.05
Table 3: Comparisons of BW2- -UP J. (%) between the learned marginal distributions and the ground
truth marginal distributions at the intermediate time moments 0, 2 ....... in dimension D 16.
Metries. We evaluate (a) how precisely our algorithm fits the target distribution IP, on RRD; (b) how well it recovers the entropic OT plan T' which is Gaussian distribution on IR RP; (c) how accurate are the learned marginal distributions at intermediate times t = 0, TO ......
In each of the mentinnd cases, we compute the BW2-UVP (2 55] between the learned and the ground truth distributions. For two distributions X and X. it is the Wasserstein-2 distance between their Gaussian approximations which further normalized by the variance of the distribution X:
100% 2 Var(x)
BW DWZ-UVVP(XX2
WH(N(uR. Ex), Ex)).
(17)
We estimate the metric by using 10a samples.
Baselines. We compare our method ENOT with LSOT 45). SCONES 14,, MLE-SB 48 D55 FB-SDE [10] (two algorithms, A and J). Results given in Tables 2 and B LSOT and SCONES solve EOT without solving SB, hence there are no results in Table 3 for these methods. Our method achieves low BW? values indicating that it recovers the ground truth process and entropic plan fairly well. The competitive methods have good results in low dimensions, however they perform worse in high dimensions. Importantly, the OT plan (Table ( and the marginal distributions of the Schrodinger bridge (Table 33. To illustrate the stable convergence of ENOT, we provide the plot of BW- between the learned plan and the ground truth plan for ENOT during training in Figure 6 (Appendix
methods scores the best results in recovering
5.3 Colored MNIST
In this section, we test how the entropy parameter affects the stochasticity of the learned plan in higher dimensions. For this, consider the entropic OT problem between colorized MNIST digits of classes "2" (P)) and "3" (P2).
Effect of parameter For 0, 1, 10, we learn our Algorithm I on the train sets of digits "2" and "3". We show the translated test images in Figures 3a 3b and 3cl respectively. When e = there is no diversity in generated "3" samples (Figure 3a). the color remains since the map tried to minimally change the image in the RGB pixel space. When e = 1, some slight diversity in the shape of "3" appears but the color of the input "2" is still roughly preserved (Figure Bc). For higher the diversity of generated samples becomes clear (Figure 3c). In particular, the color of "3" starts to slightly deviate from the input "2". That is, increasing the value of the entropy term in () expectedly leads to bigger stochasticity generated samples for test datasets by ENOT to show how diversity changes for different (Table
the plan. We add the conditional LPIPS variance [24. Table 1] of
We provide examples of trajectories learned by ENOT in Figure Z (Appendix F))
Metrics and baselines. We compare our method ENOT with SCONES [4], and DiffSB 15] as these are the only methods which the respective authors applied for results, use the FID metric (23 which is the Bures-Waa Wasserstein (Freschet) distance between the
tasks. To evaluate the

---

X-
TAR
T2X1
7,00
X-
TAR
T:X1
7000
X-Pt
TAXI
TAKI
X-P
TO0
TX0
2
3
3
3
2
3
3
3
2
3
3
3
2
2
3
2
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2 3
3
3
7
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
2
2
2
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
2
(a) ENOT (ours) samples. (b) ENOT (ours) samples. (c) ENOT (ours) samples,
(d)
D T3 samples,
(c)
DiNSS T samples,
0, HID:6.0
= 1, F.F:66.28
<
10, FID:6.9
1. FID-93
10, FID:105
DOTIC COTUK
DOTTXI
X -i DOTIO COTUK DOTUX
X-Pt PL
D073X DOTIXI bomx.
S12
S00
SI2
S13
SKd
Sun
SXI
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2 3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
3
2
3
3
(1) DOT samples. FID:NIA
(g) DOT samples, FID:NIA
(h) DOT samples.
(i) SCONES 14 samples (i
SCONES [4] samples 100, FID:14.22
100 FID:N/A
25, FID:14.73
Figure 3: Samples of colored MNIST obtained by ENOT (ours) and DOT for different E.
distributions after extracting features using the Inception V3 model 47
We measure test FID for
every method and present the results and qualitative examples in Figure 3 There are no results for
SCONES with e = 0, 1, 10, since it is not applicable for such reasonably small due computational instabilities, see [14 65.11. DiftBB [15] can be applied for small regularization so we test c = 1, 10.
By the construction, this algorithm is not suitable for e == O.
Our ENOT method outperforms the baselines in FID. DiffSB 15] yield very high FID. This is
presumably due to instabilities of Dif which the authors report in their sequel paper (46) SCONES yields reasonable quality but due to high c = 25, 100 the shape and color of the generated images "3" starts deviate from those of their respective inputs "2".
For completeness, we provide the results of the stochastic matching of the test parts of the datasets 1, 10 (Figures 31
by the discrete OT for e = 0 and EOT [2] for out-of-sample estimation, obtained samples "3" are just test samples of ingh Discrete OT is not a comperitor here as it does not generate new samples and uses target test samples. Still it gives a rough estimate what to expect from the learned plans for increasing e.
3g
3h.. This is not the
(this setup is unfair).
5.4 Unpaired Super-resolution of Celeba Faces For the large-scale evaluation, we adopt the exper- imental setup of SCONES [14. We consider the problem of unpaired image super-resolution for the 64 x 64 aligned faces of CelebA dataset [35).
Colored MNNST Celeba
2.0 10"
3.4 2
5.1-
Table 4: LPIPS variability of ENOT samples.
We do the unpaired train-test split as follows: we split the dataset into 3 parts: 90k (train A1), 90k (train BI), 20k (test C1) samples. For each part we do 2x bilinear downsample and then 2x bilinear upsample to degrade images but keep the original size. As a result, we obtain degraded parts AO, BO, CO. For training in the unpaired setup, we use parts AO (degraded faces, Pp)) and B1 (clean faces, P1). For testing, we use the hold-out part CO (unseen samples) with C1 considered as the reference.
We train our model with e=0, 1, 10 to and test how restores C1 (Figure 4ap from CO images (Figure 4b) and present the qualitative results in Figures 4f learned by ENOT in Figure [
4g
4h
We provide examples of trajectories
Metrics and baselines. To quantify the results, as in [4]. we compute the FID score 23 between the sets of mapped CO images and CI images (Table The FID of ENOT is better than FID values of the other methods, but increases with probably due to the increasing variance of gradients during training As in $5.1 and $5.33 the diversity of samples increases with c. Our method works with small values of and provides reasonable amount of diversity in the mapped samples which grows with (Table 4). As the baseline among other methods for EOT we consider only SCONES, as is the only EOT/DSB algorithm which has been applied to data-xdata task at 64 x 64 resolution. At the same time, we emphasize that SCONES is not applicable for small due to instabilities, see [14
1..1]. This makes it impractical, as due to high its produces up-scaled images (Figures 4c

---

X-
Y- Po
S(X)
S(X)
S(X)
S(X)
T(X)
T(X)
T(X)
T[X]
T(X)
(a)
(b)
(c) SCONES,
100
(d) ICNN
(e) AugCycleGAN
T(X)
T(X)
TAX)
TAX)
TXX)
TAX)
TAXI
TXX)
T(X)
T(X)
T(X)
T(X)
(1) ENOT (ours),
(g) ENOT (ours). c =
(h) ENOT (ours), 10
Figure 4: Faces produced by ENOT (ours) and SCONEs for various c.
Figure 4a shows test degraded images (CO), 4b their original high-resolution counterparts (CI).
Method
ENOT. = 0
ENOT. E
ENOT, c 10
SCONES 14 c
100
AugCyoloGAN 2
ICNN [
FID
3.78
7.63
11.8
18.88
15.2
22.2
Table 5: Test FID values of various methods in unpaired super-resolution of faces experiment.
nearly random and not reflect the attributes of the input images (Figure 4ap. do not provide
results for DiftSB [15] since it already performs bad on Colored MNIST ($5.3) and the authors also not consider any image-to-image apart of grayscale 28x28 images.
For completeness, we present results on this setup for other methods, which do not solve EOT: ICNN- based OT (38) and AugCycleGAN (2). ICNN 44d learns a deterministic map. AugCycleGAN (4c) learns stochastic map, but the generated samples differ only by brightness.
6 Discussion
Potential impact. There is a lack of scalable algorithms for learning continuous entropic OT plans which may be used in data->data practical tasks requiring control of the diversity of generated samples. We hope that our results provide new direction for research towards establishing scalable and efficient methods for entropic OT by using its connection with SB.
Potential social impact. Like other popular methods for generating images, our method can be used simplify the work of designers with digital images and create new products based on At the same time, our method may be used for creating fake images just like the other generative models.
Limitations. To simulate the trajectories following SDE (9) use the Euler-Maruyama scheme. It straightforward but may be imprecise when the number of steps is small or the noise variance e is high. As result, for large c, Algorithm I may be computationally heavy due the necessity to backpropagate through a large computational graph obtained via the simulation. Employing time and memory efficient SDE integration schemes is promising avenue for the future work.
ACKNOWLEDGEMENTE This work was partially supported by Skoltech NGP program (Skoltech- MIT joint project).
10

---

References
[1] Yves Achdou, Pierre Cardaliaguet, François Delarue, Alessio Porretta, and Filippo Santambro- gio. Mean Field Games: Cetraro, Italy 2019, volume 2281. Springer Nature, 2021.
[2] Amjad Almahairi, Sai Rajeshwar, Alessandro Sordoni, Philip Bachman, and Aaron Courville. Augmented cyclegan: Learning many-to- mappings from unpaired data. In International Conference on Machine Learning, pages 195-204. PMLR, 2018.
[3] Brandon Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference Learning Representations, 2022.
[4] Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pages 146-155 JMLR. org. 2017.
[5] Martin Arjovsky and Leon Bottou. Towards principled methods for training generative adver- sarial networks. In International Conference on Learning Representations, 2017.
[6] Arip Asadulaev, Alexander Korotin, Vage Egiazarian, and Evgeny Burnaev, Neural optimal transport with general cost functionals. arXiv preprint arXiv:2205. 15403, 2022.
[7] Julio Backhoff- 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-22) 2019.
[8] Jean-David Benamou and Yann Brenier. computational fluid mechanics solution to the kongge-kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
[9] Charlotte Bunne, Ya-Ping Hsich, Marco Cuturi, and Andreas Krause. The schrödinger bridge between gaussian measures has a closed form. In International Conference on Artificial Intelligence and Statistics, pages 5802-5833, PMLR, 2023.
[10] Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-bace backwarrd sdes theory. In International Conference on Learning Represen- tations, 2021.
[11] Yongxin Chen, Tryphon Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge schrodinger bridge. SIAM Review, 63(2):249-311 2021.
[12] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
[13] Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International conference on machine learning, pages 685-693. PMLR, 2014.
[14] 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.
[15] Valentin De Bortoli, James Thomton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to sasee generative modeling. Advances in Neural Information Processing Systems, 34: 34:17695-17709 2021.
[16] Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions on Machine Learning Research, 2023. Featured
Certification.
[17] Jiaojiao Fan, Qinsheng Zhang, Amirhossein Taghvaei, and Yongxin Chen. Variational wasser- stein gradient flow. In International Conference Machine Learning, pages 6185-6215. PMLR, 2022.
[18] Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev, Extremal domain translation with neural optimal transport. In Advances in Neural Information Processing Syslems, 2023.
[19] Milena Gazdieva, Litu Rout, Alexander Korotin, Alexander Filippov, and Evgeny Burnaev. Unpaired image super-resolution with optimal transport maps. arXiv preprint arXiv:2202.01 2022.
11

---

[20] Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, 273((1):3327-34055 2017.
[21] A Hitchhiker's s Guide. Infinite dimensional analysis. Springer. 2006.
[22] Pierre Henry-Labordere. (martingale) optimal transport and anomaly detection with neural networks: A primal-dual algorithm. arXiv preprint arXiv:1904.0m 04546, 2019.
[23] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local nash equilibrium. In Advances
in neural information processing systems, pages 6626-6633 2017.
[24] Xun Huang, Ming- Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image- to-image translation. In Proceedings of the European conference on computer vision (ECCV), pages 172-189, 2018.
[25] Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems, 33:10468-10044 2020.
[26] Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations, 2021.
[27] Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back!
wasserstein GANs are not optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022.
[28] Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34: 34:14593-14605, 2021.
[29] Alexander Korotin, I Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-
2 barycenter estimation without minimax optimization. International Conference on Learning Representations, 2021.
[30] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport.
In The Eleventh International Conference on Learning Representations, 2023.
[31] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In The Eleventh International Conference on Learning Representations, 2023.
[32] Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. arXiv preprint arXiv: 1308. 0215, 2013.
[33] Alex Tong Lin, Samy Wu Fung, Wuchen Li, Levon Nurbekyan, and Stanley J Osher. Alternating the population and control neural networks to solve high-dimensional stochastic mean-field games. Proceedings of the National Academy of Sciences, 118(31):e222771311 18, 2021.
[34] Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized schrödinger bridge. Advances in Neural Information Processing Systems, 35:9374-9388, 2022.
[35] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
[36] Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Large- scale optimal transport via adversarial training with cycle-consistency. arXiv preprint 27X1v:2003.06635, 2020.
[37] Andreas Lugmayr, Martin Danelljan, and Radu Timofte. Ntire 2021 learning the super-resolution space challenge. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 596-612, 2021.
[38] Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pages 6672-6681. PMLR, 2020.
[39] François- Pierre Paty, Alexandre d' Aspremont, and Marco Cuturi. Regularity as regularization: Smooth and strongly convex brenier potentials in optimal transport. International Conference on Artificial Intelligence and Statistics, pages 1232. PMLR, 2020.
12

---

[40] Michele Pavon and Anton Wakolbinger. On free energy, stochastic control, and schrödinger processes. In Modeling, Estimation and Control of Systems with Uncertainty: Proceedings of Conference held in Sopron, Hungary, September 1990, pages 334-348. Springer, 1991.
[41] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Foundations and Trends in Machine Learning, 11(--6)335-6607) 2019.
[42] Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport
maps. In International Conference on Learning Representations, 2022.
[43] Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58- 63):94, 2015.
[44] Erwin Schrödinger. Über die umkehrung der naturgesetze. Verlag der Akademie der Wis- senschaften in Kommission bei Walter De Gruyter Company, 1931.
[45] Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Conference Learning Representations, 2018.
[46] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[47] Christian Szegedy, Wei Liu, Yangging Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1-9, 2015.
[48] Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9): 134, 2021.
[49] Cédric Villani. Optimal transport: old and new, volume 338. Springer Science Business Media, 2008.
[50] Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. International Conference on Machine Learning, pages 10794-10804 PMLR, 2021.
[51] Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computation of large scale optimal transport. volume 97 of Proceedings of Machine Learning Research, pages 6882-6892, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
13

---

A Extended background: KL divergence with the Wiener process plan
This section illustrates that 7 holds. Consider a process T E F(Po), i.., T is a probability distribution on Ω with the marginal Po at O.
Let W be the Wiener process with variance starting Po. i.e., it satisfies dX, VedW, with Xu Po- Hence, n (y/l=) is the normal distribution da W (v/ NVyla, cf). Then KL("")
dy between joint distributions at times t 0 and 1 of these processes is given by:
KL KL(T"I
log
dm (a, 1) da (x,
{
log d(z, yl
dx (x, v) du (x, (18)
xxxy
d yl
xxx
where dx(2.y) [z,0]
denotes the joint density of distribution We derive
-f log
dn W (x, dx (x.1) = - log
dn W (ylæ) du d
de dm (I,
Xxx
dlz, yl
dy
(x)
-J Xxx
log dm d (u|z) dm (x, 1 Jy
dn W'
(x) di dn (y)az) dPo(2)
log
dy
dx
du d d dn (anT((Z,)))) (2, / log du d (x)
log
1dx ldzT(ylzz) dPo(x)
xxy
dy
de
W (u)z) di ( (I,
1xx log di (x)
Jx log
[log
dPo(z)
dy
de
- -/x/1 log
dr W' (ylx) d
dm (x. v) / log dPolz)
dPo(#)
dy
de
-J log
dn W (y/z²) do dx.""))))) H(P))
xxxy
dy
- log
(2TT) -P exp
|| 1|1-~yll u113
do
(r, H(PP)) =
xxx
2c
D log((2T)
/
||0 - y11
dm (). + H(Po).
XxX
2c
After substituting this result into
18 one obtains:
|| v/12
D log((nr) + H(P)).
<L( KL(x"||zz"')
JxxV
dx (x, - -H(TT)++
(19)
20
B
Proofs
In this section, we provide the proof for our main theoretical results (Theorems 4. and 4.3 The proofs require several auxiliary results which we formulate and prove in $ B..1 and B.2
In I we show that entropic OT can be reformulated as a maximin problem. This is a technical intermediate result needed to derive our main maximin reformulation of SB (Theorem 4.1). More precisely, in $B.22 we show that these maximin problems for entropic OT and SB are actually equivalent. By using this observation and related facts, in B.3 we prove our Theorems and
B.1 Relaxation of entropic OT
To begin with, we recall some facts regarding EOT and SB. Recall the definition of EOT
| y112
inf xeIl(F..::)
dm (x, v) (T)).
(20)
Xxx
14

---

Henceforth, we assume that Po and P1 are absolutely continuous. The situation when Po or P1 is absolutely continuous is not of any practical interest: there no of I1) for which the differential entropy H is finite which means that (20) equals to + for every E IT(Po. IP), i.e., every plan is optimal. In turn, when Po and P1 are absolutely continuous, the OT plan is unique thanks to the strict convexity of entropy (on the set of absolutely continuous plans).
Recall equation 19 for KL(T|IT""))
KL(T|| W = /ews
||x - -yll
dr(z, v) - H(T) + C.
(21)
2e
We again note that
{/
> v11 dr(x, v) eld(()) =(((m)) +C,
inf KL(T||= =fI(
inf wEll(Po.P.)
xxxy
2
ie., problems 20 and can be viewed equivalent as their minimizers the same. For convenience, we proceed with inf infrel((P.P5)) KL(T||>">> and denote its optimal value by L", i.e.
del
inf =EM((a.P.))
KL(T||>"'))
For given 33 E ()), we define an auxiliary joint distribution (x. = dr (sls))))))) where (ylar) is given by
dm (y|z) = exp(3 (ylx))
C3
where (z del Ju exp(3(y))dr W (v/x). Note that C3 < oo since 3 E Co,2(D)) is upper bounded. Before going further, we need to introduce several technical auxiliary results.
Proposition B.1. For P2()) and x E X it holds that
(L(u|| (-1=)) S, Bly)du(u) = KL((II= - log Cū-
(22)
Proof of Proposition B.I.
We derive
du(y) KL[v||> W (+x)) - I B())du(v) = =/8 log ky(z)))z)
du(v) - I S(u)du(u) Cīdu(y)
S log exp(B(y)) W (y1x) du(y) / log Cz expe( W (y/2z) du(y)
dv(v)
I log d (
du(y) dv(y) - log Cj = KL(V|| (-|z)) log C3.
Lemma B.2, For π E IT(P)), i.e., probability distributions E P2(x x y) whose projection to X
equals Pu. we have
/ B(y)dr(() KL(T||=0) / log CzdPo(z).
KL KL(T|ITW'))
(23)
Proof of Lemma B.2 For each I e X. we substitute v <(11x) 22 and integrate over Po- For the left part, we obtain the following:
[ (KL(T(())))>** (-Ix)) - / 3(y)dT(y)x)) dFo(z) =
(I(T(11))1>"' (11))))((
Jxxy (())
15

---

JucS log dT(y)=) dir (z) S,/ B(y)dTi
d (y|+))
Le log da (ylz)
di (T(y)=)))P((z)
d (y)l
((y)dTi())
JanD log
du (r. v) di (I. v)
S,- ((y)dT (
da W (x, 1)
(I(T||>>'>) / B (
For the right part, we obtain:
Jx [KL(T())))] log C3 )
S KL KL(T(-|2))] x
log C3dPo((x)
x / log dr(yla)
dm (y
log CzdPP(())
dr (y(o))
dr (y)la) dddd J log (x) (opl::dPPP(a
dm (F)
log CzdPo(()
[ 1xxO
da (z. de (I,, y) / log (z))
log dn
(L(T||=0) - / log C)dPo(x))
Hence, the equality 23 holds.
Now introduce the following auxiliary functional L:
[(3, m) def (L(T||T - -/, (())di())) + / (( 1
Recall that и denotes the second marginal distribution æ. We use this functional derive the saddle point reformulation of EOT.
Lemma B.3 (Relaxation entropic optimal transport). holds that
L
inf KL(T||n"") suup inf xPI)(Fo.P))
x),
(24)
& aEIl(Po)
where sup is taken over potentials 3 E Co,2()) and inf over E I(Po).
Proof of Lemma
B.3
We obtain
// (ol
inf =EII(Po.P.)
(L(T|[=W")
inf rEl((F.F..)
inf I (1)=)) xEI((FoPPi)
(25)
where C(v, p) def (L[v|| W (wlx)). The last problem in (25) is known as weak OT 7 20 with weak OT cost C. For a given 3 C),)()), consider its weak C-transform given by:
g (x) def inf [C(x, S ((y)dv(y)
(26)
xEPa())
Since C: X'X'> P2())) R is lower bounded (by zero), convex in the second argument and jointly lower semi-continuous, the following equality holds [7] Theorem 1.3]:
L inf / xeIl(FoPPi)
(F sup f 8 (x))PP(()) J B(y)dPP(V); (27)
16

---

where sup is taken over 3 E (o,) We use our Proposition B.1 to note that
[uLT" " (ulz)) / B(y)du(v)} -
(x)x inf vEPx())
inf (KL(v||T(((z)) log 0ß} log Cg. MEPa(D)
This allows to derive
1 8 (x)dPo(z) + / 3(y)dPi(y)
(28)
Jx log CzdPo(o) B(u)dP
=0 inf TEll(Po)
{L(T||=0)} / log CzdPolæ) / B(U)dPi())
Do not depend on
[L[[]|=T) - / log CzdFo(z)+ + / ((y)dP, (1)
}
inf =ETT(Fo)
/ B(y)d((v) +
/ ((y)dPP())
inf (P
KL(T||a b W
inf sell(P.)
[(3
(29)
Here in transition to line (29 we use our Lemma B.2 It remains to take sup, in equality between
28 and (29) and then recall (27 to finish the proof and obtain desired (24.
Thus, we can obtain the value L by solving maximin problem (4 with only one constraint II(.))) Moreover, our following lemma shows that in all optimal pairs (3*, T²" which solve maximin problem (24), n² necessary the unique entropic OT plan between Po and Pt.
Lemma B.4 (Entropic OT plan solves the relaxed entropic OT problem), Let x be the entropic OT plan between Po and Pj. For every optimal 3 argsup, infFF L(B, R), we have
arginf E(3",n) m).
(30)
nell(Po)
Proof of Lemma B.4 Since 3 optimal, know from Lemma B.3 that inf =en(F)) C(3", ②) C.. Thanks to n² IT(Po, IP), we have nf P1. We substitute x² to [(8 T) and obtain
=d (Fi(
((W))PP(W)) - S B(y) dm KL (T"" L (31)
C(3", KL(T [|="""
The functional π + C(3",m T) is strictly convex (in the convex subset (IPo)) of distributions a for which [L[[||>"") < 00). Thus, it has unique minimizer, which is x*
From our Lemmas B.3 and B.4 it follows that to get the OT plan म" one may solve the maximin problem 24 to obtain optimal saddle point (5* T²). Unfortunately, it is challenging to estimate KL(m]] from samples, which limits the usage of this objective in practice.
B.2 Equivalence of EOT and DSB relaxed problems
Below we show how to relax SB problem [T] and link its solution to the relaxed entropic OT 24
For given 3 C3,2(?), we define auxiliary process 8 such that its conditional distributions [ = We a., and its joint distribution Th at t = 0, given by 3
To simplify many of upcoming formulas, we introduce C₃ fx log CādPo(x), Also, we introduce F(Po) to denote the set of processes starting at IP at time = 0.
17

---

Lemma B.5 (Inner objectives of relaxed EOT and SB are KL with T and π T² For π (1Pp))
and T F(Po), the following equations hold:
[(B,x) (L(=|T= - Cs / B(w)dPi())
(32)
(B8.))-K((TIEEE - -Cg Str ((u))dP(())
(33)
Note that the last two terms in each line depend only on 3 but not on π or T.
Proof of Lemma B.S The first equation (32 directly follows from Lemma B.2 Now we prove 33
((B J ((y)dP = KL(TIIW> / B(y)dT]()) (v)
KL (..ll// )dt (x, / B(y)dTT] (v)
KL(T"1|>"*)
(34)
x'x
KL
|1
Cs + / JXXY
KL |z,0 [ (a..y ddx (x,
KL
Cp
KL(T|zz8/
) (xr, 3) KL(T) Cp. 11 )
(35)
xxy
In the transition to line 34
we use the disintegration formula (6. In line 35 we use the definition
of 8 i.e., we exploit the fact that
[z.9 = W m.y and again use (6
As result of Lemmal B.5
we obtain the following important corollary.
Corollary B.6 (The solution to the inner problem of relaxed SB a diffusion). Consider the problem
[(B,T). TEF(Fo)
inf
(36)
Then T the unique optimizer of (36 and holds that T E D(Po), i.e, is a diffusion process:
T arginf L(B,T) = arginf L(3,Ty). TET(P))
(37)
TyED(P,)
Proof. Thanks (33). see that T the unique minimizer of (6).. Now let Q def TT" Then
/ (())d (w))
T arginf C(B,T) = arginf TEF(P)) TeF((Po)
KL(T]|W>>>
J B(y)d (3)
arginf (e.(P..))
KL(T]|WW))
arginf KL(TI]W">= (EF()..))
=Const, since of =f="
J// t)((
arginf KL(TyIIWS) = arginf TyED(Po.Q)
- ET,
(38)
Ty TyED(Fo. 2e
In transition to (8). use the fact that the process solving the Schrödinger Bridge (this time between Po and ( with the Wiener Prior a diffusion process (see Dynamic SB problem in $22
for details). As result, obtain 8 E D(Pu, Q) C D(Pp)) and finish the proof.
Below we show that for a given 8, minimization of the SB relaxed functional L(3, Ty) over T3 is equivalent to the minimization of relaxed EOT functional C(3, T) 24 with the same 3. Lemma B.7 (Equivalence of the inf values of the relaxed functionals). It holds that
inf C(B,Ty) inf [(B,T) -C + ((y)dPP (v). -EIT(P))
(39)
TT)D(Po)
Moreover, the unique minimizers given by 1 E D(Po) and n. Il(Fo), respectively.
18

---

Proof of Lemma B.7 Follows from Lemma B.3 and Corollary B.6
Finally, we see that both the maximin problems are equivalent. Corollary B.8 (Equivalence of EOT and DSB maximin problems). It holds that
L = sup inf C(B,Ty) = sup inf ((B xell(Eo)
(40)
a TyED(Fo)
Proof of Corollary B.8 We take sup, of both parts in equation (39
Also, it follows that the maximization of infr,eD(s, (3) over / allows to solve entropic OT.
B.3 Proofs of main results
Finally, after long preparations, prove main Theorem 4.1
Proof of Theorem 4.1 and Corollary 4.2 From our Lemma B.7 and Corollary B.8 it follows that
8 Eargsup inf C(3, ( + 3* E Eargsup inf [(3, x), E TFED(P))
B *EIT(Po)
i.e. both maximin problems share the same optimal 3*, Thanks that the process T E DFoo) and the plan ot TT" E I1(Fo) are the unique minimizers of problems
Lemma B.7
we already know
inf L(8", " inf [(B",T) TyED(Bo)
xell(Pa)
respectively. Therefore, 7--=70 and, in particular, TT = TT Moreover, since (8", ( TTy') is an optimal saddle point for C, from Lemma B.4 we conclude that -r = x², i.e., --r- is the EOT plan between Po and PT. In particular, -")- II(Po, P1) which also implies that Ty- E D(Po, The last step is to derive
dP (v)
/ (())dd (v) J ((()) (w) dm Ty (v) KL(Tr.IIW9)
L- = - KL(T>-IW99) -KL
O since Ty. CD((o.P))
which concludes that Tyy is the solution to SB
Proof of Theorem 4.3 Part From Lemma B.5 and Corollary
B.6 follows that that infr, C(3,
has the unique minimizer 8 whose conditional distributions are |=.9
1 W
ja.y"" Therefore,
6 C inf 2(3,
/ B(y) (v) - C3 + / B(o)dPi()) = KL(TT, (41)
KL(T,|IT C2+
Part Now consider 62. We know that
c = KL(Ty.I|IW)) =
(L. W W. Ty (x, KL( || W²)
Xxy
From Lemma B.5 and Corollary
B.6
we also know that
+ - inf ((3,TT) = -C3 Ty
Blu)dP.(.).
Therefore:
inf C(3,TT.)) = T W + C3 - B())dPP)())
E2 = C"
19

---

KL( ||T +
/ log CzdFoo(z) - (()d (v)
x
[ KL(TTM " (
/ log CjgPa(())) (ly) (v)
KL(T T5" (+|z)] (Hz))ddF((()) / log CzdPo(x) / B(y)d= (y)x)dPo(x))
/ |KL(T (-1))|| W' (-1x)) + log C3 - I Bly) (u/z) dd
KL(TT>> (+1x)) T (-1x))
log Cj + log Cj ddPo(x)
KL( (+<>)====>>>>>>>> (+|=)||= (-Ix))d(P(()) KL T )--
KL fzh | T dd T (x, = KL 11 (42)
XxY
=0, since Ty- |0, -W
Thus, we obtain 62 = yT (T
Part 3. By summing
41 and
42
and using the Pinsker inequality, we obtain
e1 + E2 = - (T [ 2 2pfv 2T z
[PIv(TT) T + (Tv )²² 2 (T Ty-). TT..
(43)
Here use the triangle inequality in line (43) Therefore, prv (T..Tyy)) S VEI +
Part 4. By summing
41 and (42 and using the Pinsker inequality, obtain
E + 62 = KL(T)|IT)) + KL(TT-1 (T3))
KL(TT>||TT
KL
KL (l Ta. dd T/ (I,
X'Y
"L T ||
KL IT
dn (x. 2
XxX
KL(x"/1| T
KL T ||
2pt (T 1, -. T. 20 - )²
[PTv(N "s TTh
+ pi n T" 11² PT (,
Thus, pTv( =10) S VEI - E2.
C Euler-Maruyama
In our Algorithm M at both the training and the inference stages, we use the Euler-Maruyama Algorithm 2 to solve SDE.
D Drift Norm Constant Multiplication Invariance
Our Algorithm aims to solve the following optimization problem:
Br, J. ClIf(X,, (lFdt) + / B(y)dP J. ((y)ap" (v)
sup inf 8 TyED(Po)
# (cC(8,T))
with C 1. the same time, we use C the actual value of C' >> 0 does not affect the optimal solution T to this problem. Specifically, if
2e our theoretical derivations 12 We emphasize that
20

---

Algorithm 2: Euler-Maruyam algorithm
Input :batch of initial states Xo at time moment = SDE drift network fo RRP x [0, 1] + RD; number of steps for the SDE solver N 1; noise variance 2 0. Output Output:batches :bathhes
j ]== of intermediate states at z simulating the proccess =X,== fxx, +
batches {fn In=o of drift values f(Xns tn) at = m=1 N simulating the process;
4 NV for t = 1.2...... N do
for p 1.2..... [Xoo do
Sample noise W from from.((., (0, ; S1-1, 5 1(X,,1,1-1) X1. + XL-1, + At + VeAtW;
(8,TTi is the optimal point for the problem with C = 1, then (Č3rTTr is the optimal point for
C = C. Indeed, for pair (3, Ty) it holds that
(B,Tg) = EE, 1// [if(X.,t)lPdtt + I// ((u)dPi(M) - -1, dl (v)
ā ET, 1. (1I(((x)))Pdd) + (S())dR(() (3SS))RF ())
1 δ Er, Ō11((x., () + B(y)dP - B Bl))ar"" - (3,T7). (44)
I/.
where we use del ČB. Hence problems sup, infr, L (3,Ty) and sup3 infr, L (3, viewed as equivalent in the sense that one can be derived one from the other via the change of variables and multiplication by Č > O. For completeness,
(B,T) Ty) can be
also note that the change of variables ++ 3 actually preserves the functional class of i.e., 3 (V)) - 3 C,,)
in the objective (2) and consider L for optimization. ie., use C 1 in Algorithm ! Still the dependence on remains in sup, infr, L (B) as Ty E D(Fo) is a diffusion process with volatility e. Interestingly, this point of view (optimizing _ instead of ) technically allows to consider even e 0. In this case, the optimization is performed over deterministic trajectories T' determined by the velocity field fX,,, The problem saddle point reformulation of the unregularized OT with the quadratic cost in the dynamic form, also known as the Cenamou-Brenteee formula (43] $6.1]. This particular case is out of scope of our paper (it is not BOT/SB) and we do not study the properties of L in this case. However, for completeness, we provide experimental results for e = O.
For convenience,
get rid of dependence
sup, infr, (3, (3,TT)))))))))r) may be viewed
E ENOT for Toy Experiments and High-dimensional Gaussians
In 2D toy experiments, we consider 2 tasks: Gaussian - gaussians and Gaussian - Swiss roll. Results for the last one (Figure 5) are qualitatively similar to results of the first one (Figure 2). which we discussed earlier II..I). D. For both tasks, we parametrize the SDE drift function in Algorithm [ by feedforward neural network fo with 3 inputs, 3 linear layers (100 hidden neurons and ReLU activations) and 2 outputs. As inputs, we 2 coordinates and time value (as is). Analogically, we parametrize the potential by feedforward neural network Sy. with 2 inputs, 3 linear layers (100 discretization
hidden neural and ReLU activations) and outputs. In all cases, we use N steps for solving SDE by Fuler-Maruyama Algorithm Adam with Ir = 10-1, 1, batch size 512. We train the model for 20000 total iterations of Bg, and on each of them, we Ky 10 updates for the SDE drift function fa.
21

---

Tinput diceribttion
Fitted distribtiion
Fitted distribution
Fitted distribution
-2
-2
Trjjectorizs Torget data
Trajectorics
Tajecterics
-1
-1
-]
-2
Target
-2
Input and target samples
ENOT (ours),
(c) ENOT (ours), 0.01
(d) ENOT (ours), e
Figure Gaussian + Swiss roll, learned stochastic process with ENOT (ours).
In the experiments with high-dimensional Gaussians, experiments but chose N = 200 discretization steps for SDE, all hidden sizes in neural networks 512, and we train our model for 10000 iterations. To illustrate the stability of the algorithm, provide the plot of BW3-UVP (%) between the ground truth EOT plan R and the learned plan π of ENOT during training for DIM 128 in Figure
exactly the same setup for toy 2D
ENOT (ours) BW-UVP plan during the training
3.0
2.5
BW-UVP plan
2.0
1.5
1.0
0.5
0.0
2000
4000
6000
8000
10000
Outer problem steps
Figure 6: Bw?-UvP J (%) between the the EOT plan π² and the learned plan π of ENOT and MLE-SB during the training (DIM = 128).
ENOT for Colored MNIST and Unpaired Super-resolution of Celeba Faces
5.4). we find out that using the following reparametrization of Euler- Maruyama Algorithm considerably improves the quality of our Algorithm In the Fuler-Maruyama Algorithm 2 instead of using a neural network to parametrize drift function s(X,, () and calculating the next state as X X, + f(X,, At + VeAt, we parametrize gX,,, f) X, + f(X,, A) by neural network fjo and calculate the next state as X go(X,, + VEAi. In turn, the drift function given by f(X,, 1) (X,,, - X,. Also, we not add noise the last step of the Euler-Mane Maruyama simulation because find out that it provides better empirical performance.
For the image tasks (55.3
22

---

X 0
t=
t=
0=0.
t=0.4
t=0.5
t=0.6
t=0.7
t=0.8
t=0.9
t=1.0
2
2
2
2
2
2
3
3
3
3
3
2
2
2
2
3
3
E-10
2
3
Figure 7: Trajectories from our learned ENOT (ours) for colored MNIST for different 6.
We use WGAN-QC discriminator's ResNet architecture for the potential 3. We use UNet f go (Xx,,) t) of SDE in model. To condition it t, we first obtain the embedding of by using the positional embedding | Then we add conditional instance normalization (CondIN) layers after each UNet's s upscaling block We use Adam with Ir 10-4 batch size 64 and 10:1 update ratio for fe/Sg. For < = 0 and
as
our model converges in z 20000 iterations, while for c 10 takes 70000 iteration to convergence. The last setup takes more iterations to converge because adding noise with higher variance during solving SDE by Marlyama Algorithm increases the variance of stochastic gradients.
In the unpaired super-resolution of Celeba faces, we use the same experimental setup as for the colored MNIST experiment. It takes x 40000 iterations for 0 and s 70000 iterations for 1 and = 10 to converge. In Figures [ we present trajectories provided by our algorithm for Colored MNIST and Celeba experiments.
Computational complexity. In the most challenging task ($5.4). ENOT converges in one week on 2x A100 GPUs.
G Details of the baseline methods
In this section, we discuss details of the baseline methods with which compare our method.
G.1
Gaussian case ($ 5.2
SCONES 14 We use the code from the authors' repository
for their evaluation in the Gaussian case. We employ their configuration blob/main/con coniiiinnooiiig./o -yy
LSOT 45 We use the part of the code of SCONES corresponding learning dual OT potentials blob/ hoob/mman/aae /pptt py and the barycentric projection blob/main//i boob/main/bproj py in the Gaussian case with configuration blob/main/ config..ppp py.
FR-SDE-J 10
We utilize the official code from
https://githib..cm//hll/////
with their configuration blob/mainn/configg//eeee htwphtttcccceeebbr.eccooom config.py py for the checkerboard-to-nniss toy experiment, changing the number of steps of dynamics from 100 to 200 steps. Since their super-parameters are developed for their 2-dimensional experiments, we increase the number of iterations for dimensions 16, 64 and 128 to 15 000.
We also take the code from the same repository above. We base configura- wwwn_moottttooomooo to__tt____e appiral_conff ig- for the moon-to-spiral experiment. As earlier, we increase the number of steps of dynamics up 200. Also, we change the number of training epochs for dimensions 16, 64 and 128 to 2.4 and 8 correspondingly.
FB-SDE-A [0]
tion on the authors' one (blob/man imain/conniiii configs////// ponfigsddeffull deffaultmmoon______
Selhub..om//aryiee/WWa..CC
23

---

DifIS [5] We utilize the official code from
yaml for toy problems. We increase the with their configuration plob/main/conf /at.oe///ddaat//dddd/dd Catigggrrriooo////aaaa//aaaaaaar amount of steps of dynamics 200 and the number of steps of IPF procedure for dimensions 16, 64 and 128 to 30, 40 and 60, respectively.
MLE-SB 488. We use the official code from
with syper-parameteet from blob/maain/notebooks/22 Toy Data/2d_examples. ipynb. We set the number steps to 200. As earlier, we increase the number steps IPF procedure for dimensions 16, 64 and 128 to 1000, 3500 and 5000, respectively.
G.2 Colored MNIST 5.3
SCONES 14
In order prepare a score-based model, we use the code from
httr:///gthhuboooo////rroooggogo
with their configuration blob/masterr/coffii///caaa SCONES from the official repository for their unpaired Celeba super-resolution experiment halob/main/ccones///oooo wwwnffigsssuuprrssssttoe. images instead of 64 x64 celebrity faces.
Next, utilize the code of
005, yml. We adapt for 32x32 ColorMNIST
DifISB 15
We use the official code with their configuration blob/main/conf /tt//miiitttee yaml adopting for three-channel ColorMNIST images instead of one-channel MNIST digits.
G.3 CelebA (95.4)
For the SCONES, we use their exact code and configuration from blob/main aeeiii/ssconsse crroone/cooofgg/suuuuooooooso Latree_______ 0__ yml. As for the score-based model for celebrity faces, we pick the pre-trained model from
SCONES 14
https: https//ggghhhaccccc///oorrrrrrgp comm//rroongrrru//ccnn
It is the one used by the authors of SCONES in their paper.
Augmented Cycle GAN (2) We the official code from
with their default hyper-parameters.
ICNN [38. We utilize the reworked implementation by
which is non-minimax version 26 of ICNN-based approach (38]. That is,
blob/main/ http///ootboooo//ttttttttee haggeate_iinggeebeennar.r... benchmark. .imhhaarrr....bbbb...... and only change the dataloaders.
H Mean-Field Games
This appendix discusses the relation between the Mean-Field Game problem and Schrödinger Brdiges.
H.1 Intro to the Mean-Field game.
Consider a game with infinitely many small players. At time moment 0, they are distributed according to Xo Pp. Every player controls its behavior through drift of the SDE:
dX,, t, + 22vdW,
24

---

Here is the distribution of all the players at the time moment t. When we consider specific player, we consider Pt as parameter. Each player aims minimize the quantity:
EI (T(Xe, + ((xx: ((Xr,PT)
Here L(x, a. P) p) is similar to the Lagrange function in physics and describes the cost of moving in some direction given the current position and the other players' distribution. The additional function f(Xt,Pt) Pe) is interpreted as the cost of the player's i interaction at coordinate x with all the others. Now we can introduce the value function ((x, t), which for position x and start time f returns the cost in case of the optimal control:
in J." (L(X,, f(Xe ((x.))))))txtxxxprrr)x g(Xr,Pr)).
((2, t)
Before considering the Mean-Field game, we need to define an additional function (æ, p. p). is similar the Hamilton function and is defined the Legendre transform of Lagrange function L:
H (o, P.
def supl-ap L(x, a)
Mean-Field game implies finding the Nash equilibrium for all players of such the game. It is known T that the Nash equilibrium is the solution of the system of Haaiiton-Jacohi-Bellna Fokker-Planee (FP) PDE equations. For two functions Hx, P) and [xx, p). Mean-Field game formulates as a system of two PDE with two constraints:
(HJB) and
0 vā + HI, V p) = f(z, (HJB)
- --p div Hz,, Vo)) = (FP)
s.t. p(x, 0) po Φ2x T) g(2, p(-,))
The solution of this system is two functions p(z, and (zx, t), which describe all players' dynamics. Also, in Nash equilibrium, the specific player's 's behavior is described by the following SDE:
dX. VX,, Vo(X,,(),) p)dt /2vdWe.
H.2 Relation our work.
In recent work 34], the authors show that the Schrodinger Bridger problem could be formulated as Mean-Field game with hard constraints on distribution p(- T) Ptarget (- T) via choosing proper function g(z, p(-, T)) such as:
oc,
if p(. + Ptorget (-.T) D(-T)) Ptorget (:, T))
g0x, p( T))
0,
Also, the authors proposed extension of DiffSB
34 algorithm for the Mean-Field game problem.
In 33 the authors in their experiments consider only soft constraints on the target density. More precisely, they consider only simple constraints such as g(x.p) = || Starget ll2, where Starget is given shared target point for every player, and every player is penalized for being far from this. Such soft constraint force players to have delta distribution at point Ztarget-
To solve the Mean-Field problem, the authors parameterize value function d(x, by a neural network and use different neural network Ng to sample from Pt- The authors penalize the violation of Mean-Field game PDEs for optimizing these networks. After the convergence, one can sample from the distribution P: by using neural network Ng. Approach [33] has the advantage that authors do not need to use SDE solvers, which require more sleps with growing parameter u of diffusion operator. However, computation of Laplacian and divergence for high-dimensional spaces (e.g.. space 12228- 2222-ddimesioonl space of 3x64x64 images) at each iteration of the training step may be computationally hard, restricting the applicability of their method to large-scale setups.
25

---

In our approach, we initially work with the SDE:
dx, a pi) + 22dW,,
which describes the player's behavior and use neural network to parametrize the drift a. We consider only hard constraints on the target distribution, f f(x,, Pt) 0 and L(X,, at Pt) = ,| since this variant of Mean-Field game is also the particular case of Schrodinger Bridge problem and is equivalent to the entropic optimal transport. Since we do not need to compute Laplacian or divergence, our approach scales better with the dimension. However, for high values of diffusion parameter (which is equal to the fe in our notation, where e the entropic regularization strength), our approach needs more steps for accurate solving of the SDE to provide samples, mentioned in limitations.
I Extending ENOT to other costs
In the main text, we focus only on EOT with the quadratic cost c(x, y) = 2 1 y/l² which coincides with SB with the Wiener prior W However, one could use different prior Qu instead of W in (5):
Q dX, v(X,, t)dt + VedW,,
and solve the problem
inf BT, / IIf(X,,) o(X,, dd..
inf KL(T))IQ.) (P
P..)(P..P.) 26
Here we just use the known expression o6 for (L.(T,IIQ.) between two diffusion processes through their drift functions. Using the same derivation as in the main text 22. it can be shown that this new problem is equivalent to solving the EOT with cost c(1, y) = log π 2. (y/æz) where Th (y/2z) is conditional distribution of the stochastic process Qu at time / = given the starting point æ at time
( t= 0. For example, for W (which we consider in the main text) we have
cxx, v) - log W (u)x)
2c (y- x) 2) + Const,
i.e., we get the quadratic cost. Thus, using different priors for the Schrodinger bridge problem makes it possible to solve Entropic OT for other costs. We conjecture that most of our proof's and derivations can be extended to arbitrary prior process Qu just by slightly changing the minimax functional (T))
2e ( En, 1. [1)(X..t) u(Xe, ((xx,))ll(t) d)] + S. Bo(y)dP 1
sup inf Ty
-J--5-- 3 (o())d (w)
We conduct toy experiment to support this claim and consider Qu with 00.1 and ozz f) log p(x), where log p(x) is 2D distribution with a wave shape, see Figure Intuitively, means that trajectories should be concentrated in the regions with high density of p. In Figure e there the gcey-scale color map represents the density of p. start points (Pu) are green, target points (I1) are red, obtained trajectories are pink and mapped points are blue.
J ENOT for the unregularized OT (E 0)
Our proposed algorithm is designed to solve entropic OT and the equivalent SB problem. This implies that > 0. Nevertheless, our algorithm technically allows using even in which case presumably computes the unegularized OT map for the quadratic cost. Here we present some empirical evidence supporting this claim as well some thcoretical insights.
EMPIRICAL EVIDENCE. We consider the experimental setup with images from the continuous Wasserstein-2 benchmark 228 $4.4]. The images benchmark provides 3 pairs of distributions (Early, Mid, Late) for which the ground truth unregularized OT map for the quadratic cost is known by the construction. Hence, we may compare the map learned with our method
0) with the true one.
We train our method with
on each of benchmark pairs and present the quantitative results in Table 6 We use the same L2-UVP metric (28] 54.2) as the authors of the benchmark. As the baselines, we include the results of MM:R method from 28 and the method from ( Both methods
minimax and have some similarities with our approach. As can see, ENOT with 0
26

---

start points
2.5
Finish points
Lorgel points
10-
0.5
0.0
3.00
-2.5
-2.0
-5
-05
0.5
1.0
Figure 8: Toy example with ENOT (ours) for the complex prior Qu dX, = vxxe t)dt SedWe.
Benchmark
Early
Mid
Late
28
1.4
0.4
0.22
0.61
0.20
0.09
ENOT (ours)
0.77
0.21
0.09
Table 6: Comparison on W2 benchmark. *Results are taken from
Table 2].
works better than the MM:R solver but slightly underperforms compared to 3 This evaluation demonstrates that our method recovers the unregularized OT map for the quadratic cost with the comparable quality to the existing saddle point OT methods.
THEORETICAL INSIGHTS. We see that empirically method with 0 recovers the unregularized OT map. At the same time, this is not supported by our theoretical results as they work exclusively for and rely the properties of the KL divergence.
Overall, it seems like for our method yields a saddle point reformulation of the Benamou- Brenier (BB) [8] problem which is also known as the dynamic version of the unregularized OT (e = 0) with the quadratic cost. This problem can be formulated as follows:
inf L (I (1Xx, s.t. :y:dX, 5(X, Xo ~ Po. X1 ~ P1, (45)
ie., the goal is to find the process Ty of the minimal energy which moves the probability mass of Po to F1. BB 45 is very similar to DSB 11 but there is no multiplier 1, and the stochastic process Ty restricted to be deterministic (c 0). It is governed by vector field f. Just like the DSB (IT) is equivalent to EOT (),. is known that BB (45) is equivalent to unregularized OT with the quadratic cost (c = 0). Namely, the distribution T" is the unregularized OT plan between Po and P..
In turn, our Algorithm for e optimizes the following saddle point objective:
Er..I. 1// (((X) B(y) (w)) (())(()))) (u) (0)) (46)
sup inf C(3,T) Tj
def sup inf B
where Ty dXx, = f(x,, dt with Xo Po (the constraint X1 IP, here is lifted) and C2 (J).
Just like in the Entropic case, functional L can be viewed as the Lagrangian for BB (45) with B
playing the role of the Lagrange multiplier for the constraint d= (v) dl (v)- Naturally, it is expected that the value (45) coincides with (46). and we provide a sketch of the proof this fact.
Overall, the proof logic is analogous to the Entropic case but the actual proof is much more technical as we can not use the K L-divergence machinery which helps to avoid non-uniqueness, etc.
Step 1 (Auxiliary functional, analog of Lemma B.3 We introduce an auxiliary functional
E(8, H) def [ 1
H()) ((H))))PPo(z) + +1 (1)ddPP
27

---

where 3 is potential and H R - R is a measurable map. This functional is nothing but the well-known max-min reformulation of static OT problem (in Monge's form) with the quadratic cost 3 Eq. 4]. 28 Eq.9]. Hence,
1 / || H(x)ll HPPP=
sup inf [(3, H) H
inf
=
Step 2 (Solution of the inner problem is always an OT map). An existence of some minimizer H = = in i L(3, H) can be deduced from the measurable argmin selection theorem, e.g., 22 Theorem 19]. For this H we consider IP def
"P Recall that
/
| H(x)ll 2
a arginf L(3, = arginf H
((H())))
H
x
Here may add the fictive constraint HHPO P' which is anyway satisfied by H and get
/
|| (x))1 B(H(
|| - -1) dFo(.).
H8 arginf HPP
arginf HHPo=P
The last equality holds since f B(H ((H()))dPo(z) = dF (y) does not depend on the choice of H due to the constraint HiP, = P.. The latter is the OT problem between I and 1e and we see that H is its solution.
Step (Equivalence for inner objective values). Since H is the OT map between Po, P (it is unique as Fo is absolutely continuous [43),, it can be represented as an ODE solution Tya to the Benamour Brenier problem between Po. P., i.e., Tya dX, (X,,, t)dt and HX
Xo + (X,, Furthermore, in this case, Xo (Xo)ll fo I1 ()l Hence, can be derived that
inf 2(8 H) = inf ((3,T)) H
T
Step (Equivalence the saddle point objective). Take sup over Co,2(D) and get the final equivalence:
sup inf 2(B, H) sup inf ((B,Ty) =
$
Step 5 (Dynamic OT solutions contained in optimal saddle points). Pick any optimal 8 argsup, inf и C(3, Ty-) and let Ty- be any solution to the Knnamou-Brenice problem. Checking
that T E inf L(3", Tr) can be done analogously to [31 Lemma 4J, 42 Lemma 4.1].
The derivation above shows the equivalence of objective values of dynamic unregularized OT (45) and our saddle point reformulation of BB (45). Additionally, it shows that solutions Ty- can be recovered from some optimal saddle points [3 Tr-) of our problem. At the same time, unlike the EOT case (e > 0), it is not guaranteed that for all the optimal saddle points (3*, Ty.) it holds that
Ty. is the solution to the BB problem. This aspect seems to be closely related to the fake solutions issue in the saddle point methods of OT 30 and may require further studies.
28


Reviews:
[{"review_id": "Q0FKRNpp9Z", "reviewer": "Reviewer_TSgQ", "summary": "The main idea of the paper is to estimate a stochastic map for the entropic optimal transport problem using its connection to the dynamic Schrödinger bridge (SDB) problem. The authors formulate the SDB as a saddle point problem of an associated Lagrangian. Then they recover the transport plan as the joint distribution of the solution to the dynamic Schrödinger bridge problem's initial and final values, while the transport is encoded in the drift term of the learned stochastic process that is the solution of the SDB. Compared to previous methods, the method at hand offers more stability for small entropic regularization coefficients. The errors on the drift term solution and on the transport map are quantified given the corresponding duality gaps of the inner and outer optimization problems of the saddle point objective. Finally, the approach is supported by experimental evaluation.", "questions": "* Computation of $\\mathbb{E}\\left[\\int_0^1\\Vert f(X_t,t)\\Vert^2{\\rm d}t\\right]$ : in line 197, it is indicated that the mean of the $f(x,t)$ is used. Is this justified by the Riemann integral discrete approximation? If so, does a trapezoidal rule for approximating the integral improve the result?\n* Is it straightforward to generalize the approach to costs other than the squared Euclidean distance?  Does it fundamentally change the nature of the associated stochastic process \n* Did the authors try to apply the method to the domain adaptation (DA) problem as several DA methods rely on optimal transport ?\n----------\nI have read the authors's rebuttal. They have addressed my concerns.", "limitations": "The impact of the paper and the limitations of the contribution are clearly discussed in Section 6.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is extremely well-written and reader-friendly.\n\n- Several remarks are made to facilitate reading and to provide intuition about technical notions.\n- A guarantee on the quality of the saddle point solution is provided.\n- Addressing the small $\\epsilon$ case, which is a source of instability of several other methods.", "weaknesses": "* In line 223, it is mentioned that the negative entropy is not strongly convex. This is false as the function $p\\mapsto x\\ln(x)$ has second derivate $x \\mapsto \\frac{1}{x}$ which is bounded from below by $1$ on the interval $[0,1]$. See for example Section 4.1 of reference [2]. As a result, the comparison to [1]  (reference [5] in the paper) needs to be reconsidered.\n\n* There is no concluding section.\n\n* Some experimental section metrics are not introduced.\n  \n  * The $\\rm{BW_2^2-UVP}$ is not introduced even in the appendix, although a reference for it is given. \n  \n  * FID: the previous remark applies.\n\n* The parametrization $g(X_t,t) = X_t + f(X_t,t)\\Delta_t $ should have been indicated in the main paper rather than in the appendix as although it is mathematically equivalent to the parametrization presented in the main paper, it allowed better results on the CelebA dataset according to the appendix\n  \n  ## Minor remarks\n\n* Problem with links: For some reason, the bibliographic references links along with links to equations and sections etc. are not working.\n\n* $\\pi^{W^\\epsilon}$ is introduced for the first time in Equation (8) without being defined.\n\n* $W_{|x,y}$ is not explicitly defined, although one can infer its meaning from the definition of $T_{|x,y}$.\n\n* Theorem 4.1: I think it should be \"every pair $(\\beta^*,T_{f^*})$ for (13)\" rather than \"for (12)\" as problem (13) is a saddle point problem.\n\n* Reference to Algorithm 2: in line 195, Algorithm 2 is referenced. However, it is not indicated that it is written in the appendix.\n\n* Suggestion: index $m$ can be removed in the \"$\\widehat{KL} \\leftarrow$\" line of Algorithm 1  since the sum terms are already indicated to be the values of $f_n$, or it is possible to indicate $\\sum_{m=1}^{|f_n|}$.\n\n* Line 158: I think it should be added \"that is bounded from above\" to \"a continuous function\".\n\n## References\n\n[1] Asadulaev, A., Korotin, A., Egiazarian, V., & Burnaev, E. (2022). Neural optimal transport with general cost functionals. *arXiv preprint arXiv:2205.15403*.\n\n[2] Peyré, G., & Cuturi, M. (2019). Computational optimal transport: With applications to data science. *Foundations and Trends® in Machine Learning*, *11*(5-6), 355-607.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer TSgQ, thank you for your comments. Here are the answers to your questions.\n\n**(1) In line 223, it is mentioned that the negative entropy is not strongly convex. This is false as the function $p\\mapsto x\\ln(x)$ has second derivate $x \\mapsto \\frac{1}{x}$ which is bounded from below by $1$ on the interval $[0,1]$. See for example Section 4.1 of reference [2]. As a result, the comparison to [1] (reference [5] in the paper) needs to be reconsidered.**\n\nRecall that the negative (differential) entropy of a distribution with density $p(x)$ is given by $-H(p) = \\int p(x)\\log p(x)dx$. The negative entropy is indeed $\\frac{1}{M}$-strongly convex if we only consider distribution $p$ whose density $p(x)$ is bounded by a constant $M>0$ (this follows from your argument). However, we work with general continuous distributions ($\\mathbb{P}_0,\\mathbb{P}_1,\\widehat{\\pi}$, etc.) which may not satisfy this assumption. For example, $\\mathcal{N}(x|0,\\sigma)$ for small $\\sigma$ may have density at $0$, which is greater than any given constant $M$. Therefore, the differential entropy is not $\\frac{1}{M}$-strongly convex for any $M>0$.\n\n**(2) Some experimental section metrics are not introduced.**\n\n*We will add the explanations and definitions of the metrics to Appendix.*\n\n**(3) The parametrization $g(X_t,t) = X_t + f(X_t,t)\\Delta_t $ should have been indicated in the main paper rather than in the appendix as although it is mathematically equivalent to the parametrization presented in the main paper, it allowed better results on the CelebA dataset according to the appendix.**\n\n*We will include a comment regarding the parametrization utilized in the primary text (to Section 4.2).*\n\n**(4) Problem with links: For some reason, the bibliographic references links along with links to equations and sections etc. are not working.**\n\nThank you for pointing this out. We will fix it.\n\n**(5) Minor remarks.**\n\nThank you for your comments. We will fix all the issues and think it will further improve the clarity.\n\n**(6) Computation of $\\mathbb{E}\\left[\\int_0^1\\Vert f(X_t,t)\\Vert^2{\\rm d}t\\right]$ in line 197, it is indicated that the mean of the \n is used. Is this justified by the Riemann integral discrete approximation? If so, does a trapezoidal rule for approximating the integral improve the result?**\n\nYes, we used the mean as a discrete approximation of $\\mathbb{E}\\left[\\int_0^1\\Vert f(X_t,t)\\Vert^2{\\rm d}t\\right]$. We have not tried to use other types of approximation of this integral that might improve the proposed algorithm. Thank you for your suggestion.\n\n**(7) Is it straightforward to generalize the approach to costs other than the squared Euclidean distance? Does it fundamentally change the nature of the associated stochastic process.**\n\nYes, it is quite straightforward. We focus only on EOT with the quadratic cost $c(x,y)=\\frac{1}{2}\\|x-y\\|^{2}$ which coincides with SB with the Wiener prior $W^{\\epsilon}$. However, one could use a different prior $Q_v$, given by the SDE: \n$$\nQ_v: dX_t = v(X_t, t)dt + \\sqrt{\\epsilon}dW_t,\n$$\nand solve the problem\n\n$$\n\\inf_{T_f \\in \\mathcal{D}(\\mathbb{P}_0, \\mathbb{P}_1)} \\text{KL}(T_f || Q_v) = \\inf\\_{T\\_f \\in \\mathcal{D}(\\mathbb{P}\\_0, \\mathbb{P}\\_1)} \\frac{1}{2\\epsilon} \\mathbb{E}\\_{T\\_f}[\\int\\_{0}^1 ||f(X_t, t) - v(X_t, t)||^2 dt].\n$$\n\nHere we just use the known expression to $\\text{KL}(T_f|| Q_v)$ between two diffusion processes through their drift functions as we did in our paper. Using the same derivation as in our paper, it can be shown that this new problem is equivalent to solving the EOT with cost $c(x,y) = -\\log \\pi^{Q_v}(y|x)$, where $\\pi^{Q_v}(y|x)$ is a conditional distribution of the stochastic process $Q_{v}$ at time $t=1$ given the starting point $x$ at time $t=0$. For example, for $W^{\\epsilon}$ (which we used) we have $c(x,y) = -\\log \\pi^{W^{\\epsilon}}(y|x) = \\frac{1}{2 \\epsilon}(y-x)^T(y-x) + \\text{Const}$, i.e., we get the quadratic cost. Thus, using different priors for the Schrodinger bridge problem makes it possible to solve Entropic OT for other costs. \n\nIt seems that all our proofs can be extended to any prior process $Q_v$ just by slightly changing the minimax functional:\n\n$$\n\\sup\\_{\\beta} \\inf\\_{T\\_{f}} (\\frac{1}{2\\epsilon} \\mathbb{E}\\_{T\\_f}[\\int\\_{0}^1 ||f(X\\_t, t) - v(X\\_t, t)||^2 dt] + \\int\\_{\\mathcal{Y}} \\beta\\_{\\phi}(y) d\\mathbb{P}\\_1(y) - \\int\\_{\\mathcal{Y}} \\beta\\_{\\phi}(y) d\\pi\\_1^{T\\_f}(y)). \n$$\n\nWe conducted a toy experiment to support this claim and consider $Q_v$ with $\\epsilon=0.01$ and $v(x, t) = \\nabla \\log p(x)$, where $\\log p(x)$ is a 2D distribution with a wave shape, **see figure 2 of the attached pdf.** Intuitively, it means that trajectories will be concentrated in the regions with a high density of $p$. In Figure 2, There the grey-scale color map represents the density of $p$, start points ($\\mathbb{P}\\_0$) are green, target points ($\\mathbb{P}\\_1$) are red, obtained trajectories are pink and mapped points are blue.\n\n**(8) Did the authors try to apply the method to the domain adaptation (DA) problem as several DA methods rely on optimal transport?**\n\nNo, we did not try.\n\n**Concluding remarks.**\nWe would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns about our work. If so, we kindly ask that you consider increasing your rating. We are also open to discussing any other questions you may have.\n\n**Additional references.**\n\n[1] Asadulaev, A., Korotin, A., Egiazarian, V., \\& Burnaev, E. (2022). Neural optimal transport with general cost functionals. arXiv preprint arXiv:2205.15403.\n\n[2] Michele Pavon and Anton Wakolbinger. On free energy, stochastic control, and schrödinger\nprocesses. In Modeling, Estimation and Control of Systems with Uncertainty: Proceedings of a\nConference held in Sopron, Hungary, September 1990, pages 334–348. Springer, 1991.\n"}, {"review_id": "Vk7geFl7uS", "reviewer": "Reviewer_LQN6", "summary": "This paper focuses on neural optimal transport  and more particularly through a \"dynamic schrodinger bridge\" approach. I am not an expert in this particular topic, but I must say that the authors manage to make it quite readable and a good introduction to the methodology.\n\nAs far as I can tell, the particularity of the approach is to start from a known key constrained optimization problem in eq (11) that has been presented before and that is progressively explained. Then, departing from previous methods, the authors propose a Lagrangian approach, for which we now have an unconstrained problem and the need to train not only the drift function that allows sampling, but also some \"beta network\" (coined in as the beta network), in a min-max optimization approach, that is similar --but different -- from the usual GAN methodology.\n\nAfter the rigorous theoretical treatment, the authors present some very nice experiment that seem to support their claims and the interesting features of their approach. I think the paper is quite challenging, but that it is also very stimulating and inspiring.\n\n-- \nafter rebuttal, I am still happy with the paper. Maintaining my score.", "questions": "- p3 \"Hence, one may optimize (5) over processes T for which T|x,y=W_xy for every x, y and set the last term in (6) to zero\". How do you actually do that ? By enforcing gaussian dynamics ?\n- the paragraph \"continuous OT\" reads a bit weak to me. All the references are just given in a row, without great care for actually discussing them and their connections with the proposed approach.\n\nIn algorithm 1, I have some questions.\n- do you actually need to store the gradients when computing {X_n, f_n} for the computation of L_\\beta ? It looks to me that Eul-Mar(X_0, T_{f_\\theta}}) can actually run in some `no_grad` environment, or am I mistaken ? This would mean there is some possibility in strongly parallelizing this or even use some external workers to compute that ?\n\n- On the contrary, in the inner loop (over k), you really need to store the gradients, right ? Is it necessary to use many inner iterations K_f ? Maybe I'm mistaken but I don't see that discussed, although this looks like a key computational burden when I have a look at algorithm 1, right ? Is it feasible to just record gradient for the last k steps ? for some of them ? Some thing that could help going faster ?\n\n- I have no clue what this BW2-UVP metric is. Could you at least give us a hint that would help us avoid checking reference [25] ?\n\n- You must check the references. Most of them are badly formated. You have \"schr\\\"odinger\" everywhere and you are missing many uppercases for proper nouns", "limitations": "I would say that the limitations are clearly mentioned", "rating": 9, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I feel slightly uncomfortable in assessing whether the proposed method is new or is not exactly or if the authors miss some particular recent and/or relevant method in their references, mostly because I cannot be considered an expert on the topic. Still, it looks to me that they are doing their best at providing a very objective account of the relevant literature and giving many pointers, so I am assuming they can be trusted when they present their contributions.\n\nThe proposed method seems quite ok to implement (cf Algorithm 1) and the results are very nice. I am not really aware of the appropriate metrics and usual evaluation criteria that should be used, but I felt compelled by the experiments, which made me want to try things out. I guess this is the most important aspect.\n\nTo summarize what I feel are the strengths of the paper:\n- good theoretical overview, motivations and derivations\n- the resulting method seems simple to implement\n- interesting performances", "weaknesses": "The paper is a bit difficult to follow, but I don't think it should really be modified to be simpler. I guess it's mostly a matter of myself not being an expert in the topic.\n\nStill, some changes here or clarifications here and there could help. I will mention them below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer LQN6, thank you for your comments. Here are the answers to your questions.\n\n**(1) p3 \"Hence, one may optimize (5) over processes $T$ for which $T_{|x,y}=W_{|xy}$ for every $x, y$ and set the last term in (6) to zero\". How do you actually do that ? By enforcing gaussian dynamics ?**\n\nWe did not intend to do that in practice. And we do not do it in our algorithm. We just noticed that it can be considered as such a constrained optimisation. In fact, there seems to be no straightforward way to parameterise such a family of processes. However, as we noted in the main text (lines 76-78), if $T^*$ is a solution to SB with prior $W^{\\epsilon}$ for any two distributions $\\mathbb{P}\\_0$ and $\\mathbb{P}\\_1$, then $T^*\\_{|x,y} = W^{\\epsilon}\\_{|x,y}$  [28, Proposition 2.3].\n\n**(2) the paragraph \"continuous OT\" reads a bit weak to me. All the references are just given in a row, without great care for actually discussing them and their connections with the proposed approach.**\n\nSince many recent papers appeared on continuous OT, we focused the discussion only on the most relevant EOT and SB papers and only briefly mentioned the rest. Following your suggestion, we will extend the discussion of methods for solving other types of continuous OT (unregularized OT).\n\n**(3) do you actually need to store the gradients when computing ${X_n, f_n}$ for the computation of $L_\\beta$ ? It looks to me that Eul-Mar($X_0$, $T_{f_\\theta}$) can actually run in some no grad environment, or am I mistaken ? This would mean there is some possibility in strongly parallelizing this or even use some external workers to compute that ?**\n\nWe do not have to store the gradients in the computation of ${X_n, f_n}$ for the computation of $L_\\beta$ in Algorithm 1, and we indeed do not store them in our implementation. However, we have not yet tried to use computational schemes with external workers in our work. We appreciate the idea and will try it.\n\n**(4) On the contrary, in the inner loop (over k), you really need to store the gradients, right ? Is it necessary to use many inner iterations $K_f$ ? Maybe I'm mistaken but I don't see that discussed, although this looks like a key computational burden when I have a look at algorithm 1, right ? Is it feasible to just record gradient for the last k steps ? for some of them ? Some thing that could help going faster ?**\n\nWe have to store them, but we can easily solve the memory problem using gradient checkpointing [1]. We have implemented it in our code. The hyperparameter $K_f$ affects the quality of the solution of the inner problem and can be tuned. We noticed that $K_f=1$ is too small, and the algorithm can easily diverge, but for $K_f=10$, the problem vanishes. We noticed in the limitations section that backpropagations through SDE may be computationally heavy. More efficient SDE solvers can be used to overcome this limitation.\n\nAnother possibility is to consider approaches from stochastic optimal control or reinforcement learning, since with the fixed potential $\\beta$, the inner problem can be formulated as a problem from these fields, as noted by **reviewer KyaP**. For example, one could consider $-||f(X_t, t)||^2\\Delta t$ as a reward in an intermediate step and $\\beta(X_{predicted})$ as an additional reward in the final step. In this case, one could use Q-learning or Actor-Critic approaches to learn from a segment of a trajectory without propagating through the whole trajectory.\n\n**(5) I have no clue what this BW2-UVP metric is. Could you at least give us a hint that would help us avoid checking reference [25] ?**\n\nIt is the Wasserstein-2 distance between distributions $\\mathbb{P}$ and $\\mathbb{Q}$ that are coarsened to Gaussians and normalized by the variance of the distribution $\\mathbb{Q}$:\n$$\n\\text{B}\\mathbb{W}\\_{2}^{2}\\text{-UVP}\\big(\\mathbb{P}, \\mathbb{Q} \\big) = \\frac{100 \\\\%}{\\frac{1}{2}\\text{Var}(\\mathbb{Q})} \\mathbb{W}\\_{2}^{2} (\\mathcal{N}(\\mu\\_{\\mathbb{P}}, \\Sigma\\_{\\mathbb{P}}), \\mathcal{N}(\\mu\\_{\\mathbb{Q}}, \\Sigma\\_{\\mathbb{Q}})).\n$$\n\nWe will add this definition to the main text.\n\n**(6) You must check the references. Most of them are badly formated. You have \"schr\"odinger\" everywhere and you are missing many uppercases for proper nouns.**\n\nThanks so much for pointing out this. We will check and correct the references.\n\n**Concluding remarks.** \nWe would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns about our work. We are also open to discussing any other questions you may have.\n\n**Additional references.**\n\n[1] Chen, Tianqi, et al. \"Training deep nets with sublinear memory cost.\" arXiv preprint arXiv:1604.06174 (2016).\n[2] Christian Léonard. A survey of the schr\\\" odinger problem and some of its connections with\noptimal transport."}, {"review_id": "E3UXp5elIo", "reviewer": "Reviewer_WYHi", "summary": "Inspired by how Sinkhorn duals are derived the authors adapt said derivation to the path measure and via the disintegration theorem they derive a novel unconstrained min-max objective for solving the Schrodinger bridge problem, the authors then proceed to showcase their method in eOT based tasks, introducing a new gaussian benchmark and displaying competitive results to previous approaches.  Additionally, the authors quantify errors in sampling and transport of approximate minimisers to their proposed schemes.\n", "questions": "1. Some of the methods you compare to are closed form / non-iterative in the way they solve the half bridges e.g. [1] using GPs.  A thing to note is iterative approaches such as SGD can keep resampling from P_0 and P_1 in the toy tasks nonstop, effectively making the dataset set size somewhat proportional to the iterations, this strongly benefits generalisation capabilities, especially in higher dimensions in contrast to a the GP approach in [1] which is likely fitted on a small dataset. If we assume that the dataset is fixed across epochs K_f * batchsize would give a proxy as to the size of the dataset used in the DL approaches , from inspecting the code provided in the supplementary zip in particular the high dim Gaussians and toy examples IPYNB this quantity seems to be above 5000, whilst I suspect for [1] it might be lower. Even then I suspect you resample from the toy distributions (Gaussians) every time as indicated in your pseudocode in which case the dataset sizes are simply not comparable at all and the non-iterative approaches like MLE-SB are prone to be affected significantly by generalisation error (which scales slowly in number of samples for kernel methods, which also don't overcome the curse of dim as well as neural networks)  in particular in high dimensions.\n2. Are the same number of timesteps used across the DL and the GP approaches when training ? computationally it feels unlikely the gram matrix would fit in memory for the given timesteps used with the DL approaches. \n3. A suggestion here would be to compare all approaches in an additional table (e.g. in the appendix) under budgets in the number of steps and number of samples (and fairly ensure these are the same across methods), this can be useful to quantify the data efficiency of each approach. Furthermore, the settings under which [1] was run (steps and dataset size) should be reported.\n", "limitations": "Finding saddlepoints / solving min-max problems is typical quite a challenging task and was often a huge challenge in the training stability of GANs (very high variance training results). In contrast to IPF which is much nicer from an objective viewpoint (sequence of regression losses), this approach poses the question of how stable training is.\n\nAs none of the experiments have error bars (on training runs) it is difficult to see if the proposed approach is robust and “easy” to train.  I would suggest to the authors provide such results and in addition maybe comparisons to training loss with Chen 2021 joint or De Bortoli 2021.\n", "rating": 9, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The extension of the Sinkhorn dual to the dynamic setting is rather elegant and certainly novel in the way it is carried out \n2. The paper is excellent on the presentation side in regards to technical ideas, whilst the contributions are novel/creative they are presented in such a way that understanding them was not difficult.\n3. The new formulation allows for a novel duality gap analysis, which is one of the few works analysing learned SBP methods in the approximate setting. \n4. From a purely conceptual viewpoint the work is great and rather complete, just some clarifications/additions on the experimental side and motivations could be enhanced.", "weaknesses": "Outside of a concern in how different methods are compared (detailed in the questions), this paper is overall well written and has mostly sound experiments. \n\nFrom the method standpoint, there are several potential weaknesses and lacking ablations which I will detail in the limitations.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer WYHi, Thank you for your comments. Here are the answers to your questions.\n\n**(1) Comparison with MLE-SB.**\n\nComparing entropic OT methods is difficult because they are based on different principles: IPF-based (MLE-SB, DiffSB, FB-SDE), dual form based (LSOT, SCONES), semi-dual form based (ENOT, ours), and each of them has different hyperparameters that need to be tuned. **It seems like our paper is unique in the field which performs such a comprehensive comparison.** We have taken the most similar hyperparameters to those from the authors' repositories. In our initial results, all the methods (except LSOT and MLE-SB) work well and give good results. LSOT performs poorly because it only learns the barycentric projection, not the EOT plan. At the same time, it was not very clear to us why the results for MLE-SB were bad.\n\nIt appears that [1] is the paper in which MLE-SB was first introduced, as it was the only one that used GP for training. Indeed, to ensure a fair comparison of the methods, rather than the way of solving regression problems, it would be better to use a neural network for all methods. This is because GP can be challenging to scale when used for setups with many time steps ($200$ for high-dimensional Gaussians) and a moderate amount of samples. \n\nWe reran the experiments for MLE-SB using neural network parametrization of the drift function (as well as for other methods) and using $200$ time steps as we used for all SB methods. We use $1000$ IPF iterations and sample $512$ samples from each distribution $\\mathbb{P}_0$ and $\\mathbb{P}_1$ at each IPF iteration. **The updated results are shown in Tables 1, 2, and 3 of the attached PDF file.**\n\nAs we can see, MLE-SB performs similarly or better than the other IPF-based method (DiffSB) and can also solve the problem in this setup. Now all the considered methods that learn the EOT plan work well on the Gaussian setup, and the goal of this experiment is achieved. The small residual error of all methods seems to be related more to hyperparameter tuning than to the nature of the algorithm used.\n\n**(2) Stability of the proposed method.**\n\nWe run our method on the Gaussian setup five times and provide the means and standard deviations in **Tables 1, 2, and 3 of the attached PDF file**. In Figure 1, we also provide the plot of $\\text{BW}_2^2\\text{-UVP}$ (\\%) between the ground truth EOT plan $\\pi^{*}$ and the learned plan $\\pi$ of ENOT and MLE-SB during training for $DIM=128$. Note that the steps on the plots represent different values due to the different nature of the algorithms. There are IPF steps for MLE-SB and outer problem steps for ENOT (ours). As we can see, both methods stably converge.\n\n**Concluding remarks**. \nWe would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns about our work. If so, we kindly ask that you consider increasing your rating. We are also open to discussing any other questions you may have.\n\n**Additional references.**\n\n[1] Vargas, Francisco, et al. \"Solving schrödinger bridges via maximum likelihood.\" Entropy 23.9 (2021): 1134."}, {"review_id": "xotDPmgMBR", "reviewer": "Reviewer_KyaP", "summary": "The paper proposes to solve dynamic entropic optimal tansport (EOT), also known as Schrödinger bridge problem, with nerual solver. Specifically, the authors propose a saddle-point, maximin, formulation of EOT, yielding a GAN-resemble algorithm that can be trained in an end-to-end fashion. Experiments are conducted on 2D toy datasets and images translation. ", "questions": "- Experiments seem to support that $\\epsilon$=0, where the dynamics reduce to flow, works much better than SDE. This seems to contradict recent insights from diffusion models [1], where SDE usually performs better than ODE on the same tasks. Can the authors comment on that?\n- why is there a \"1/|f|\" term in Alg 1, when computing the KL? And what's $f_{n,m}$? I understand tht $f_n$ is the drift output at time step $n$. \n\n[1] Elucidating the Design Space of Diffusion-Based Generative Models", "limitations": "Limitations are included in Sec6, where the authors mentioned the potential computational burden caused by simulation and back-prop through the SDE dynamics during training.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The saddle point reformulation of EOT is interesting. The proposed algorithm can be viewed as a stochastic control problem, where the terminal cost function $\\beta$ is learned so that the resulting policy approaches terminal distribution. \n- Writing is generally clear and easy to follow. Sufficient related works and preliminaries are included.\n- Experiment are extensive and include many related baselines.", "weaknesses": "- The proposed algorithm aims to minimize (12), which is not well-defined when $\\epsilon$=0, as the KL term will blow up. Even though $\\epsilon$=0 is algorithmically applicable, it makes the current algorithm disconnected from the mathematical framework. I'll be more convinced if the authors can provide additional justification (maybe connection to OT).\n\n- The proposed method is closely related to recent maximim OT [1], which consists of the same two networks (potential + policy) and the same training losses. From my understanding, the two methods coincide when $\\epsilon$=0 and $N=1$. Are the authors aware of [1]? Given that the proposed ENOT seems to work best when $\\epsilon$=0, can the author compare with [1]?\n\n- Given that DiffSB [13] was compared throughout most Sec 5, I suggest the authors to compare in Sec 5.4 to [2], which applies DiffSB to unpaired super-resolution image datasets.\n\n- While Sec 2 has introduced sufficient background and comparison between EOT and SB, which I do appreciate, I think their connection to Sec 4, which is the main content, is rather weak. Given that the proposed algorithm closely relates to dual formulation (e.g. (26) in Appendix), I suggest including those parts in Sec 2.\n\n[1] On amortizing convex conjugates for optimal transport  \n[2] Conditional Simulation Using Diffusion Schrödinger Bridges  \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer KyaP, thank you for your comments. Here are the answers to your questions.\n\n**(1) Connection to OT for unregulated case $\\epsilon=0$.**\n\n**We emphasize that our work focuses on developing a new algorithm for solving entropic OT and the equivalent SB problem. This implies that $\\epsilon>0$.** We present the empirical results of our algorithm for $\\epsilon=0$ only for completeness since it *technically* allows the use of $\\epsilon=0$ (see lines 226-227). Given your interest in studying this case in detail, we provide **the proof sketch** to show that for $\\epsilon=0$ our saddle-point objective\n$$\n\\sup_{\\beta} \\inf_{T_f} \\mathcal{L}(\\beta,T_f) = \\mathbb{E}_{T\\_f}( \\int\\_0^1 ||f(X_t, t)||^2 dt) - \\int\\_{\\mathcal{Y}} \\beta(y) d\\pi_1^{T_f}(y)+ \\int\\_{\\mathcal{Y}} \\beta(y) d\\mathbb{P}_1(y) ,\n$$\n\n$$\ns.t. \\quad X_0 \\sim \\mathbb{P}_0, \\quad T_f : dX_t = f(X_t, t)dt.\n$$\nis a max-min reformulation of the **Benamou-Brenier** problem (the dynamic OT with $\\ell^{2}$ cost):\n$$\n\\mathcal{L}^*=\\inf_f \\mathbb{E}\\_{T_f}[\\int_0^1 ||f(X\\_t, t)||^2 dt], \\quad s.t. \\quad X\\_0 \\sim \\mathbb{P}\\_0, \\quad  T\\_{f} : dX_t = f(X_t, t)dt, \\quad X_1 \\sim \\mathbb{P}_1.\n$$\nwhich searches for an ODE drift $f$ which moves the mass of $\\mathbb{P}_0$ to $\\mathbb{P}_1$.\n\n**SKETCH OF THE PROOF**:\n\n*Step 1 (Auxiliary functional, analog of Lemma B.3).*\n\nWe introduce an auxiliary functional:\n$$\\widetilde{\\mathcal{L}}(\\beta,H)=\\int\\_{\\mathcal{X}} \\|x-H(x)\\|^{2}d\\mathbb{P}\\_{0}(x)-\\int\\_{\\mathcal{X}} \\beta(H(x))d\\mathbb{P}\\_0(x)+\\int\\_{\\mathcal{Y}} \\beta(y)d\\mathbb{P}\\_1(y),$$\nand consider the following maximin reformulation of OT with the quadratic cost [1, Eq.4]:\n$$\n\\sup_{\\beta}\\inf_{H}\\widetilde{\\mathcal{L}}(\\beta,H) = \\inf\\_{T\\sharp \\mathbb{P}\\_0 = \\mathbb{P}\\_1} \\int\\_{\\mathcal{X}} ||x - T(x)||^2 d\\mathbb{P}\\_0(x) = \\mathcal{L}^{*}.\n$$\n\n*Step 2 (Solution of the inner problem is always an OT map).*\n\nIt can be shown that the minimizer of the inner problem exists ($H^{\\beta}$), i.e.:\n$$\nH^{\\beta}\\in \\text{arg}\\min\\_{H\\sharp\\mathbb{P}_0=\\mathbb{P}'}\\int\\_{\\mathcal{X}} \\big\\lbrace\\|x-H(x)\\|^{2}- \\beta(H(x))\\big\\rbrace d\\mathbb{P}\\_0(x)\n$$\nMoreover, $H^{\\beta}$ is an OT map between $\\mathbb{P}_0$ and $\\mathbb{P}'\\stackrel{def}{=}H^{\\beta}\\sharp \\mathbb{P}_0$.\n\n*Step 3 (Equivalence for inner objective values).*\n\nSince $H^{\\beta}$ is the OT map between $\\mathbb{P}\\_0,\\mathbb{P}' $, it can be represented as an ODE with zero acceleration ($\\frac{df^{\\beta}(x(t), t)}{dt}=0$) solution $T_{f^{\\beta}}$ to the Benamour Brenier problem between $\\mathbb{P}\\_0,\\mathbb{P}'$, for which $\\|x-H^{\\beta}(x)\\|^{2}=\\int_{0}^1 ||f^{\\beta}(X_t, t)||^2 dt$, i.e.:\n\n$$\\inf_{H}\\tilde{\\mathcal{L}}(\\beta,H) = \\inf_{T_f}\\mathcal{L}(\\beta, T_f).$$\n\n*Step 4 (Equivalence of the saddle point objective).*\n\nTake $\\sup$ over $\\beta\\in\\mathcal{C}\\_{b,2}(\\mathcal{Y})$ and get the final equivalence:\n$$\\sup\\_{\\beta}\\inf\\_{H}\\widetilde{\\mathcal{L}}(\\beta,H)= \\sup\\_{\\beta}\\inf\\_{H}\\mathcal{L}(\\beta,T_{f})= \\mathcal{L}^{*}.$$\n\nWe cannot give the full proof due to the length limit of the answer, but per request, we can provide it. **At the same time, our paper focuses on solving the problem with $\\epsilon > 0$.**\n\n**(2) Comparision with [1].**\n\nWe are aware of [1]. However, since they address the unregularized OT problem, we have not compared our results to theirs. As per your request, we have trained our method with $\\epsilon=0$ using the same image benchmark setup [2] and present our results alongside the results from Table 2 of [1] **in Table 4 of the attached pdf**. Since the method presented in [1] compares with the MMR method from [2], we also present its results.\n\nAs we can see, ENOT with $\\epsilon=0$ works better than the MM-R solver but slightly underperforms compared to [1]. \n\n**(3) Comparison with [3].**\n\nWe do not compare with DiffSB since the authors do not consider unpaired translation for the spaces of images larger than grayscale $32$x$32$. In their official GitHub repository there is no config for unpaired translation at all. Moreover, the authors themselves used not the Wiener prior for SB, see [3, Appendix J.3]. After trying to tune hyperparameters which we used for the Colored MNIST setup on our own and obtaining poor results, we decided not to scale this algorithm further to unpaired Celeba setup.\n\n**(4) Including dual form in Sec 2.**\n\nWe agree that dual form formulation of OT problems also could be discussed in Section 2 since our proofs are based on it. We will add a discussion about dual OT formulation and methods based on it (such as [1] and [2]) to the main text of the final version.\n\n**(5) Comparison of SDE and ODE.**\n\nOur proposed algorithm has a larger gradient variance when $\\epsilon$ is larger, which may affect the final quality. To improve the results, one can use more steps for sampling from SDE or adjust the learning rate. In the paper, we present results for different $\\epsilon$ while keeping all the other hyperparameters the same.\n\n**(6) why is there a \"1/|f|\" term in Alg 1, when computing the KL? And what's $f_{n,m}$? I understand tht $f_n$ is the drift output at time step $n$.**\n\nWe use $f_{n, m}$ to denote drift output at time step $n$ for the $m$-th object of the input sample batch. We use the average of the drifts as an estimation of $\\int_{0}^{1} ||f(X_t, t)||^2 dt$ in the training objective.\n\n**Concluding remarks.**\nWe would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns about our work. If so, we kindly ask that you consider increasing your rating. We are also open to discussing any other questions you may have.\n\n**Additional references.**\n\n[1] Amos,. \"On amortizing convex conjugates for optimal transport.\"\n\n[2] Korotin, et al. \"Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark.\"\n\n[3] De Bortoli, et al. \"Diffusion Schrödinger bridge with applications to score-based generative modeling.\""}]

Author rebuttal:
Dear reviewers, thank you for taking the time to review our paper.

Your valuable feedback and constructive comments are greatly appreciated. We are particularly pleased that all reviewers found our paper well-written and easy to read (TSgQ, LQN6, WYHi, KyaP). We are also pleased that you find our duality gap analysis and guarantee of the quality of the saddle point solution important (TSgQ, WYHi), that our experiments are extensive and include many related baselines (KyaP), and that the whole work is great and quite complete (WYHi).

Please, find the answers to your questions below. **Please note that we have added tables and figures in the attached pdf to support our responses to the reviewers WYHi, KyaP, and TSgQ.**

Meta-review: The paper proposes a novel algorithm for estimating the entropic optimal transport by
bridging the gap with the Schrödinger bridge problem. The authors propose a saddle-point, maxmin optimization problel
that allows to train the neural OT.

All reviewers agree that the paper proposes a novel approaches with sound theoretical arguments and
supported by empirical evidences.
Decision: Accept (oral)