Light Unbalanced Optimal Transport
Authors: Milena Gazdieva
Arip Asadulaev
Evgeny Burnaev
Alexander Korotin
Venue: NeurIPS 2024 poster (2024)
Keywords: unbalanced optimal transport
light solver
entropy regularization
generative modeling
Primary area: generative_models

Abstract: While the continuous Entropic Optimal Transport (EOT) field has been actively developing in recent years, it became evident that the classic EOT problem is prone to different issues like the sensitivity to outliers and imbalance of classes in the source and target measures. This fact inspired the development of solvers that deal with the *unbalanced* EOT (UEOT) problem $-$ the generalization of EOT allowing for mitigating the mentioned issues by relaxing the marginal constraints. Surprisingly, it turns out that the existing solvers are either based on heuristic principles or heavy-weighted with complex optimization objectives involving several neural networks. We address this challenge and propose a novel theoretically-justified, lightweight, unbalanced EOT solver. Our advancement consists of developing a novel view on the optimization of the UEOT problem yielding tractable and a non-minimax optimization objective. We show that combined with a light parametrization recently proposed in the field our objective leads to a fast, simple, and effective solver which allows solving the continuous UEOT problem in minutes on CPU. We prove that our solver provides a universal approximation of UEOT solutions and obtain its generalization bounds. We give illustrative examples of the solver's performance.

Light Unbalanced Optimal Transport
Milena Gazdieva
Arip Asadulaev
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute
ITMO University Artificial Intelligence Research Institute Moscow, Russia asadulaev@airi.nm net
Moscow, Russia
mile milena. gazdieva@skoltech. ru
Evgeny Burnaev
Alexander Korotin
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute Moscow, Russia
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute
Moscow, Russia w.koottinskkoltehhh.h.
burraeev@skolleech.rrr
Abstract
While the continuous Entropic Optimal Transport (EOT) field has been actively de- veloping in recent years, it became evident that the classic EOT problem is prone to different issues like the sensitivity to outliers and imbalance of classes in the source and target measures. This fact inspired the development of solvers that deal with the unbalanced EOT (UEOT) problem the generalization of EOT allowing for mitigating the mentioned issues by relaxing the marginal constraints. Surprisingly, it turns out that the existing solvers are either based on heuristic principles or heavy- weighted with complex optimization objectives involving several neural networks. We address this challenge and propose a novel theoretically justifiedd lightweight, unbalanced EOT solver. Our advancement consists of developing novel view the optimization of the UEOT problem yielding tractable and a non-min non-minimax opti- mization objective. We show that combined with light parametrization recently proposed the field our objective leads to fast, simple, and effective solver which allows solving the continuous UEOT problem in minutes our solver provides a universal approximation of UEOT solutions and obtain its generalization bounds. We give illustrative examples of the solver's performance. The code is publicly available at
CPU. We prove that
Introduction
The computational optimal transport (OT) has proven to be powerful tool for solving various
7 and biological data transfer [5] 42 (66] Historically, the majority of early works in the field were built upon solving the OT problem between discrete probability measures [10] 53 Only recently the advances in the field of generative models have led to explosive interest from the MI. community in developing the continuous OT solvers, see [38] for a survey. The setup of this problem assumes that the learner needs to estimate the OT plan between continuous measures given only empirical samples of data from them. Due to senvexty-relatte issues of OT problem [40], many works consider the EOT problem, i.c., use entropy regularizers which guarantee, e.g., the uniqueness of learned plans.
popular tasks, e.g., image-to-ing ro-imaa translation [68 17 50 28 image generation 67
13
Mcanwhile, researches attract attention other shortcomings of the classic OT problem. It enforces hard constraints on the marginal measures and, thus, does not allow for mass variations. As a result, OT shows high sensitivity an imbalance of classes and outliers in the source and target
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

measures which are almost inevitable for large-scale datasets. To overcome these issues, it is common to consider extensions of the OT problem, e-g., unbalanced OT/EOT (UOT/UEOT) 8 The unbalanced OT/FOT formulations allow for variation of total mass by relaxing the marginal constraints through the use of divergences.
The scope of our paper is the continuous UEOT problem. It seems that in this field, a solver that is fast, light, and theoretically justified has not yet been developed. Indeed, many of the existing solvers follow a kind of heuristical principles and are based on the solutions of discrete OT. For example, [45] uses regression to interpolate the discrete solutions, and [6 36 build a flow matching upon them. Almost all of the other solvers (9] [70] employ several neural networks with many syper-parameters and require umeeccnsuming optimization procedures. We solve the aforementioned shortcomings by introducing novel lightweight solver that can play the role of a simple baseline for unbalanced EOT.
Contributions. We develop novel lighteeight solver estimate continuous unbalanced EOT couplings between probability measures ( Our solver has a non-minimax optimization objective and employs the Gaussian mixture parametrization for the UEOT plans. We provide the generalization
bounds for our solver ($4.4) and experimentally test it on several tasks ($5.1 55.2).
- II). We use Pz.ac (R2) denote the set of absolutely continuous (a.c.) Borel probability measures differential entropy. The set of nonnegative measures on I with finite second moment is denoted as M2,+ (IR We use (RR)) to denote the space of all continuous functions Ç R + E for which 3a alS), = b(() such that Vz E Rd [S(x)] S + b||=x|]² Its subspace of functions which are additionally bounded from above is denoted as (IR)). For a.c. measure
Notations. We work in the Euclidian space (R², |l
R' with finite second moment and
(R)) (or M2,+ (RR"), we use p(z) to denote its density at point E Rd.. For a given a.c. measure def
E M2,+ x ) we denote its total mass by l/ = far (xx y)dxdy. We use 7=(x), 7v (v) to denote the marginals of (3x, v). They satisfy the equality |/ = 1l =|//|1 We write 7-1ax) to denote the conditional probability measure. Each such measure has unit total mass. We use f to def
denote the Fenchel conjugate of a function f: [(t) = SuPuER {ul. [u)]. }. We use T to denote the convex indicator of set A, i.e., (x) = 0f E AA +00 if A.
2 Background
Here we give an overview of the relevant entropic optimal transport (EOT) concepts. For additional
details on balanced EOT, refer to [10 24 53), unbalanced EOT to g 43
f-divergences for positive measures. For positive measures 11. .222,(R)) and lower
semi-continuous function f: IR U (oo), the f-divergence between , U2 is defined by
[
M1 (1) (x)
Dy (HIll
p2(z)dx if #1 < #2 and + otherwise.
We consider ft) which are convex, non-negative and attain zero uniquely when t = 1. In this case, Dy is valid measure of dissimilarity between two positive measures (see Appendix [foo for details). This means that Dy (1 (lp2) 2 0 and (y (l = 0 if and only if μu 1u2- In our paper, we also assume that all f that we consider satisfy the property that f is non-decreasing function.
Hillback-Leibles divergence DRL 8
62 is particular case of such f-divergence for positive def
measures. It has a generator function (t) = tlogt-t+1. Its convex conjugate =x)==)(ttt) 1. Another example is the x2 -iveerggee xx² which is generated by (t) def (t 1)² when 2 0 and Do otherwise. The convex conjugate of this function is fxa(t) = - if << -2 and e
2 +1 when / 2 2.
Remark. By the definition, convex conjugates of fKL and fx² divergences proper, non-negative and non-decreasing. These properties are used in some of our theoretical results.
), its entropy Entropy for positive measures. For μ E M22+ (IR""
given by
((m) def I (((r)log(((x)d)) + || if is a.c. and - 0o otherwise.
When μ (Rd)), i.c., || =1, equation I is the usual differential entropy minus 1.

---

Classic EOT formulation (with the quadratic cost). Consider two probability measures P2.ac( (R.)) De (2. For > 0, the EOT problem between p and q is
min /... xell(p.g)
x
vll (2. yddrdy (H),
(2)
JR JRd
2
where [p,, q) is the set of probability measures E E ac x Rd) with marginals p. 4 (transport plans). Plan n² attaining the minimum exists, it unique and called the EOT plan.
Classic EOT imposes hard constraints on the marginals which leads to several issues, e.g., sensitivity to outliers [4], inability handle potential measure shifts such class imbalances in the measures q. The UEOT problem [70 []] overcomes these issues by relaxing the marginal constraints [62].
Unbalanced EOT formulation (with the quadratic cost). Let Dj. and D/2 be two f-divergences over Rd., For two probability measures E P2,ac(R E P2.cc and e >> 0, the unbalanced EOT problem between p and q consists of finding a minimizer of
[
I vl1 7(2, v)dady =+I(7)++ Dr (Tellp) + D/2 (7ull)),
inf
yEMo, (Bt (Bd)
JRd
see Fig. I Here the minimum is attained for unique " which is called the unbalanced optimal entropic (UEOT) plan. Typical choices of fi (i E [1) are (i( = (t) or [i(1) = (() (T; > 0) yielding the scaled DRL. and D22, respectively. In this case, the bigger T1 (72) is, the more Yee (7yy) penalized for not matching the corresponding marginal distribution (g).
inf UEOT())
Tel Yy (1y
Yx
P
q
22
Tn
Figure I: Unbalanced EOT problem.
Remark 1. While the set M2,+ (IRd x Rd) contains not only a.c. measures, infimum in problem B is attained for a.c. measure 7 since -cH(r") term turns to +oc otherwise. Thus, almost everywhere in the paper we are actually interested in the subset of a.c. measures in M2,+ (IR. x IRd).
Remark 2. The balanced EOT problem is special case of 33. Indeed, let fi and f2 be the convex indicators of {1}. ie., f1 = f2 = I==1- Then the f-divergences D,, (7allp) and D,, (7p q) become infinity if p # Yo or q Ty- and become zeros otherwise.
Dual form of unbalanced EOT problem is formulated as follows
sup { (A.9) J
expl (o0( | vl1 }]ddd J1 J (4)
It is known that there exist two measurable functions 6", v/ delivering maximum to the following connection with the solution of the primal unbalanced problem
4 They have
||
vl(2
1 ("())
>", =xxp
eexp
eep
(5)
2e
Remark. For some of our results, we will need to restrict potentials (Φ. vb) in problem (4) to the space (IR") C2
Since established variants of dual forms typically correspond to other functional spaces, we derive and theoretically justify variant of the dual problem 44) in Appendix
A.3 Note that it may hold that optimal potentials "²" & (R. i.e., the supremum is not achieved within our considered spaces. This is not principle for our subsequent derivations.
Computational UEOT setup. Analytical solution for the unbalanced EOT problem is, in general, not known Moreover, in real-world setups where unbalanced EOT is applicable, the measures p, q are typically not available explicitly but only through their empirical samples (datasets).
We assume that data measures p. (/ P2 (Rd) are unknown and accessible only by a limited number of i.idd. empirical samples {zo...., ~.p {yo.... .... ym } q. We aim to approximate the optimal UEOT plan 7 solving (3) between the entire measures p. q. The recovered plans should allow the out-of-sample estimation, i.e., generation of samples from 7 (- (-|2:eew) where new is new test point (not necessarily present in the train data). Optionally, one may require the ability to sample from 7
Analytical solutions are known only for some specific cases. For example, [3] consider Gaussian measures and UEOT problem with DKL divergence instead of the differential entropy. This setup differs from ours.
3

---

Selver
Problem
Principles
What recovrrs? Scaling " (())p(x) stochatiic OT map 7 (ar,s)
Limitation
Selves estransform hased semi-dual
Complex objective,
max-min reformulation UOT using neural noss
neural networks
Custem
Regression on top discrese BOT
between balanced measures
Scaling factors and QT naps ro-scaled
Heuristically uses
minitstch OT appeotimations
consbined with ICNN-based salver
UDT
reformulaion UOT using neural Flow Matching on lop discrete UEDT using neural nets
Selves semi-dual max-min
Stochastic UOT map 1 (z. 2)
Complex max-minn objecrive,
neural networks
T6
UEUT
[0.1] [0.1 to traspprrt the mass
Purametrized vector
minibuich DT approximations
Heuricically taes
Scaling factiors parametrized conditional
X5
UBITT
Conditional Flow Masching om kop of discrete EOT between re balanced meusures using neural nets
vector field (+) a [0.1] to transport
Heuristiaally uses minibusch DJ agproximations
the mass between re-scaled messrres
U U-LightOT
(ours)
UBOT
Solves zoe- minimax reformulation of dual UEOT using Cousiian Mixtures
procadure to sample a o 725 ( and y to T (-la)
Density of UROT plan 7 together with light
Restricted to Gaussian Mixtune puramestizatio
Table 1: Comparison of the principles of existing UOT/UEOT solvers and our proposed light solver.
The described setup is typically called the continuous EOT and should not be twisted up with the discrete EOT [53 [O]. There the aim is to recover the (unbalanced) EOT plan between the empirical measures p p= N Di=ddn Say, 9 to و3= Sy, and out-of-sample estimations are typically not needed.
3 Related Work
Nowadays, the sphere of continuous OT/EOT solvers actively developing. Some of the early works related to this topic utilize OT cost as the loss function [27 25 2] These approaches are not relevant to us as they do not learn OT/EOT maps (or plans). We refer to [39 for a detailed review.
At the same time, there exist large amount of works within the discrete OT/EOT setup 10 15
69 [5]) see (53) for a survey. We again emphasize that solvers of this type are not relevant to us as they construct discrete matching between the given (train) samples and typically do not provide a generalization the new unseen (test) data. Only recently ML community started developing out-
sample estimation procedures based on discrete/batched OT. For example, [19 56 32
14 48 57
mostly develop such estimators using the barycentric projections of the discrete EOT plans. Though these procedures have nice theoretical properties, their scalability remains unclear.
Balanced OT/EOT solvers. There exists vast amount of neural solvers for continuous OT problem. Most of them learn the OT maps (or plans) via solving saddle point optimization problems 3 [8] 37
22 60 50 Though the recent works (28 6
30 tackle the EOT problem (2). they consider
its balanced version. Hence they are not relevant to us. Among these works, only [4] 30 evade non-trivial training/inference procedures and are ideologically the closest to ours. The difference between them consists of the particular loss function used. In fact, our paper proposes the solver which subsumes these solvers and generalizes them for the unbalanced case. The derivation of our solver is non-trivial and requires solid mathematical apparatus, see $1
Unbalanced OT/EOT solvers. A vast majority of early works in this field tackle the discrete UOT/UEOT setup 6 20 (54] but the principles behind their construction are not easy to generalize to the continuous setting. Thus, many of the recent papers that tackle the continuous unbalanced OT/EOT setup employ discrete solutions in the construction of their solvers. For example, [45] regress neural network on top of scaling factors obtained using the discrete UEOT while simultancously learning the continuous OT maps using an ICNN method [47]. In [6 and (36), the authors implement Flow Matching 44 FM] and conditional FM on top of the discrete UEOT plans, respectively. The algorithm of the latter consists of regressing neural networks on top of scaling factors and simultaneously learning a conditional vector field to transport the mass between re- re-balaneed measures. Despite the promising practical performance of these solvers, it is still unclear to what extent such approximations of UEOT plans are theoretically justified.
The recent papers 70 9 are more related to our study as they do not rely on discrete OT approxima- tions of the transport plan. However, they have non-trivial minimax optimization objectives solved using complex GAN-style procedures. Thus, these GANs often lean on heavy neural parametrization, may incur instabilities during training, and require careful hyperparameter selection [46]
For completeness, we also mention other papers which are only slightly relevant to us. For example, 23 considers incomplete OT which relaxes only one of the OT marginal constraints and is less general than the unbalanced OT. Other works [22 4 incorporate unbalanced OT into the training objective of GANs aimed at generating samples from noise.
In contrast to the listed works, our paper proposes theoretically justified and lighweight solver to the UEOT problem, see Table I for the detailed comparison of solvers.

---

4
def Light Unbalanced OT Solver
In this section, derive the optimization objective ($4.1) of our U-LightOT solver, present practical aspects of training and inference procedures ($4.2) and derive the generalization bounds for our solver (94.4). We provide the proofs of all theoretical results in Appendix A
4.1 Derivation of the Optimization Objective
Following the learning setup described above, we aim to get a parametric approximation "a.,w of the UEOT plan 7" Here 0, are the model parameters to learn, and it will be clear later why we split them into two groups. To recover 700 R y², our aim is to learn 0, w by directly minimizing the DRI. divergence between Ne and y²:
DKL (1T,,)), min (0,w)
The main difficulty of this optimizing objective (6) is obvious: the UEOT plan 7" is unknown.
Fortunately, below we show that one still can optimize (6) without knowing 7"
Recall that the optimal UEOT plan y has the form
5 We first make some changes of the variables.
Specifically, we define v² (y)
del
exp(
2nh (w)-Ilvll""
Formula (5) now reads as
2c
2 (x) ||=|| Jexp(
(a, v) ju (v) - (wl=) expl (1,1) Ju (v)-
(1, exp(
(7)
2e
Since the conditional plan has the unit mass, may write
(I. v) exp(
o (3)
(v(a))
(8)
Co-
where C. (z) exp( (e.e)) Jee dy) is the normalization costant ensuring that fRx 7 (y|x)dy = 1.
Consider the decomposition >² = (z) (w/l))))))..)))))))))) It shows that to obtain parametrization for the entire plan (", v), it is sufficient to consider parametrizations for its left marginal Y and the conditional measure 7 (y/z)) Meanwhile, equation (8) shows that conditional measures 7" (-|z) are
entirely described by the variable v². We these observations to parametrize Ne We set
expl (2.01/e)ua()))
79.0 (x, v) u (1)70(())))) (x)
Co (o(x)
where uw and Ug parametrize marginal measure " and the variable v", respectively. In turn, the
constant cæ (z) = fgu exp( (m-22) ( (ve(y)dy is the parametrization of c.- (z). Next, we demonstrate
our main result which shows that the optimization of (6) can be done without the access y*
Theorem 4.1 (Tractable form of DKL minimization). Assume that " is parametrized using 9 Then the following bound holds: eDKL (7" (110,,) S L(,, w) - L", where
(x)
|
/ Ja(-s (oguv(() pd
11r12
C(O, w)
log
)(xx)d++
(olu)dy-ellue II, (10)
JRd
Co
and constant (-L") is the optimal value of the dual form 4 The bound is tight in the sence that it turns to 0 O when vo =xx) (20"(x) (w)- 111l1//2)) and uw (x) =co(x) cxp(20" (m) [00-1=1//2]]]]
In fact, 11) is the dual form (4) but with potentials (Ø, v expressed through u. vo co):
(x) cx(x)
+ |
||yl|
((x) + (x)=
(0y) + vo = log vo (v)
2
2
Our result can be interpreted the bound the quality of approximate solution New recovered from the approximate solution to the dual problem (4) It can be directly proved using the original (ò, v notation of the dual problem, but we use (M v6) instead with this change of variables the form of the recovered plan "o is more interpretable (uw defines the first marginal, 0i- conditionals).
Instead of optimizing (6) to get "i we may optimize the upper bound I0) which is more tractable. Indeed, (10) is sum the expectations w.r.t. the probability measures p.q. We can obtain Monte-Carlo estimation of (10) from random samples and optimize it with stochastic gradient descent procedure w.r.t. (0, w). The main challenge here the computation of the variable co and term ||u I Below propose smart parametrization by which both variables can be derived analytically.
5

---

4.2 Parametrization and the Optimization Procedure
Parametrization. Recall that uw parametrizes the density of the marginal 7 which is unnormalized. Setting r = 0 in equation (7), we get (yll 0) x v² (y) which means that e" also corresponds to an unnormalized density of measure. These motivate us to use the unnormalized Gaussian mixture parametrization for the potential vo and measure (1)):
vo(y)= = > CanN(vlra,sSs) (x) (.((ollue)..) 1=1
(11)
Rd (i, EiHtE with ork, 2 0, rh, Mr and 0 Sk, Eg Raxd The covariance matrices scaled by e just for convenience.
Here def {ax, rk, Sk ) k=1
For this type of parametrization, it holds that |/ Σi= St. Moreover, there exist closed-free expressions for the normalization constant (x(x)) and conditional plan no r), see [4] Proposition
3.2]. Specifically, define Tu (2) = тk + Se.t and (xx)
del
del
ag expl" == "2a j. holds that
28
K
co (z) āz (x);: (0(()>) co(x)
1
Σ
(12)
k=1
Using this result and get the expression for Yow!
E ān (()ra())))))))
70 (2r, y) = a (x) (o(())) Σ (o(clle
(13)
Σk= (x(x)
(x))
Training. We recall that the measures ps & are accessible only by a number of empirical samples (see the learning setup in 2 Thus, given samples {21,,.., and {yI,,.,,y yar). we optimize the empirical analog of (10):
Σ(0, w) def
Σ fi(-elog 1=1
N
(x,)) co (x;)
| 2
1 M
Σ
M
f2(-elog (o()))
1 11? 2
J (14)
N
using minibatch gradient descent procedure w.r.t. parameters (0, w). In the parametrization of oe and uw (TT). we utilize the diagonal matrices Sg, Σ, This allows decreasing the number of learnable parameters in and speeding up the computation of inverse matrices ST xT
Inference. Sampling from the conditional and marginal measures To (v(xx) ( u. R is easy
and lightweight since they are explicitly parametrized as Gaussian mixtures, see 12
4.3 Connection Related Prior Works
The idea of using Gaussian Mixture parametrization for dual potentials in EOT-related tasks first appeared in the EOT/SB benchmark (29]. There was used to obtain the benchmark pairs of probability measures with the known EOT solution between them. In [4] the authors utilized this type of parametrization to obtain a light solver (LightSB) for the balanced EOT.
Our solver for unbalanced EOT (00) subsumes their solver for balanced EOT as well as one problem subsumes the other for the special case of divergences, see the remark in g
Let fi, f2 be convex
indicators of (1}. Then (i( t and (2()) and objective TOpp becomes
C(0,w) w) (-e log (x)
|>>|| )p(x)dr= [ (-Elogro(w) | )a(r)dy +oldyy+elleee ellualll
Rdd
co(x)
-e -e( log (x)
pla)dox / log vo (()g(y)dy-llul)
||0|| dr
l|vll
(())dy=
Rd co(x)"
JRa 2
Constt(p.9)
El/
log logco()))(z)dz)) -
JRa log (og)o()))())dd)) (15)
Rd
log (z) + alluw/11 - Const(p, q). (16)
Rd

---

Here line (15) depends exclusively on and exactly coincides with the LightSB's objective, see [4 Proposition 8]. At the same time, line (16) depends only on w, and its minimum is attained when tw = p. Thus, this part not actually needed in the balanced case, see [41 Appendix C).
4.4 Generalization Bounds and Universal Approximation Property
Theoretically, to recover the UEOT plan y", one needs to solve the problem L(0,w) () -> ming as stated in our Theorem 4.1 In practice, the measures p and q are accessible via empirical samples X def (II,..., } and Y def (y1,-.,Bar>> thus, one needs to optimize the empirical counterpart [(0,w) of w), see
Besides, the available potentials uws Vy over which one optimizes the objective come from restricted class of functions. Specifically, we consider unnormalized Gaussian mixtures u. Do with K and I. components respectively. Thus, one may naturally wonder: how close is the recovered plan T0. to the UEOT plan > given that (0, arg ming,w L(9,w)) w))
To address this question, we study the generalization error EDrr. (7*1lna,a) i.e., the expectation of DKL between y and V0. taken w.r.t. the random realization the train datasets ~ p, q.
Let (8, ()) arg (f ())x L(),w) denote the best approximators of L(0, w) in the admissi- ble class. From Theorem 4.1 it follows that that DEDrı. (7*1noass) can be upper bounded using E(C(,, () - L*). The latter can be decomposed into the estimation and approximation errors: B(C(O,)) O) -²* -----"""==-- -000 -(000))))((000))EEEC.....
[C(0,))_C) -(00 [(0,
(17)
Estimatioon
Approxination
We establish the quantitative bound for the estimation error in the proposition below.
Proposition 4.2 (Bound for the estimation error). Let p, q be compactly supported and assume
that I1, 12 are Lipshitz. Assume that the considered paramelric classes O Ω (9 (0,w)) consist of unnormalized Gaussian mixtures with K and L, components respectively with bounded means || || < R (for some R >> 0), covariances sT 3 Sk Σn 3 ST (for some O < S S) and weights a S 00, 31 S (for some < a S. A). Then the following holds:
B[ ()) L(. [ S O
+ + VM
where O()) hides the constants depending only on K, R, s, S, p. but not on sizes M,
This proposition allows to conclude that the estimation tend to infinity at the usual parametric rate. It remains to clarify the question: can we make the approximation error arbitrarily small? We answer this question positively in our Theorem below.
converges zero when N and M
Theorem 4.3 (Gaussian mixture parameterization for the variables provides the universal approxi- mation of UEOT plans). Let and be compactly supported and assume that Fi, J, are Lipshitz. Then for all 8 there exist Gaussian mixtures Ug. uw (IT) with scalar covariances S >xIa
Σ, = a > 0 of their components that satisfy DRI. (7" [70..) S (f(0,w) =
L*) 6.
In summary, results of this section show that one can make the generalization error arbitrarily small given sufficiently large amount of samples and components in the Gaussian parametrization. It means that theoretically our solver can recover the UEOT plan arbitrarily well.
5 Experiments
In this section, we test our U-LightOT solver on several setups from the related works. The code is written using PyTorch framework and is publicly available at
https: :/ thub. cwwmcoimmm/iicceccccccc
The experiments are issued in the convenient form of ipynb notebooks. Each experiment requires several minutes of training on CPU with 4 cores. Technical training details are given in Appendix B
5.1 Example with the Mixture of Gaussians
We provide an illustrative "Gaussians Mixture' example in 2D to demonstrate the ability of our solver to deal with the imbalance of classes in the source and target measures. We follow the experimental
7

---

x y
X pu())
y Ya.u
(a) Input, target measures
(b) U-LightOT, 7 10
(c) U-LiihhOT,, 10* (d) U-LightOT r = 10
Figure 2: Conditional plans (0, learned by our solver in Gaussians Mixture experiment with unbalancedness parameter [10", 10 102). Here Pa denotes the normalized first marginal twy i.c., = /
setup proposed in [6 Figure 2] and define the probability measures p. 9 as follows (Fig. 2ap:
p(x) | (, I))+ 3
(cr 0.1 T2),
(vv) - 4 ([,](-,))),.111 (2)11-1))+)) 12)+ - Mwl(1, -1)).
We test our U-LightOT solver with scaled DkL divergences, i.c., fi(t), are defined by т (t) 7( log
+ 1) where > 0. We provide the learned plans for E [1, 10 10a]. The results in Fig. 2 show that parameter T can used to control the level of unbalancedness of the learned plans. For = 1, U-LightOT solver truly learns the UEOT plans, Fig. 2d When increases, the solver fails to transport the mass from the input Gaussians to the closest target ones. Actually, for 102, our solutions are similar the solutions of [41 LightSB] solver which approximates balanced EOT plans between the measures. Hereafter, we treat as the umbalancedness parameter. In Appendix C
test the performance of our solver with Dy= divergence.
Remark. Here conduct all experiments with the entropy regularization parameter 0.05. The parameter E is responsible for the stochasticity of the learned transport Ya (+(x)). Since we are mostly interested in the correct transport of the mass (controlled by f1, f2) rather than the stochasticity, do not pay much attention to è throughout the paper.
5.2 Unpaired Image-to-Image Translation
Another popular testbed which is usually considered in OT/EOT papers is the unpaired image-to- image translation 71 task. Since our solver uses the parametrization based on Gaussian mixture, may be hard to apply U-LightOT for learning translation directly in the image space. Fortunately, nowadays it is common to use autoencoders [59 for more efficient translation. We follow the setup of [4] Section 5.4] and use pre-trained ALAE autoencoder [55] for 1024 1024 FFHQ dataset [34 of human faces. We consider different subsets of FFHQ dataset (Adult, Young, Man, Woman) and all variants of translations between them: Adulr YduitsyYoung and Woman<+Man.
The main challenge of the described translations the imbalance of classes in the images from source and target subsets, see Table 2 Let us consider in more detail Adult Young translation. the FFHQ dataset, the amount of adult men significantly outnum-
Class
Man
Woman
Young
15K
23K
Adult
7K
3.5K
bers the adult women, while the amount of young men is smaller than that of young women. Thus, balanced OT/EOT solvers are expected to translate some of adult man representatives to young woman ones. At the same time, solvers based on unbalanced transport are supposed to alleviate this issue.
Table 2: Number of train FFHQ images for each subset.
Baselines. We perform comparison with the recent procedure [6 UOT-FM] which considers roughly the same setup and demonstrates good performance. This method interpolates the results of unbalanced discrete OT to the continuous setup using the flow matching 44 For completeness, we include the comparison with LightSB and balanced optimal transport-based flow matching (OT- FM) 44 to demonstrate the issues of the balanced solvers. We also consider neural network based solvers relying on the adversarial learning such as UOT-SD 9 and UOT-GAN[70]

---

Young - Adult
Adult + Young
Man - Woman
Woman Man
10
10
38
10
120
Acc. (keep)
Acc. (keep)
Acc. (keep)
Acc. (keep)
10'
U LightOT (ours)
U UigghT (ours)
U- U-JggheT [ours]
U-UinTTT (ours)
60
UOT
UOT FM
UCT
UCT
50
LOT--AN
U33--GN
UGIT-GAN OT FM
CT FM
CT FM Lightsi
OT FM
40
Lightsi
Uigetii
Lightsil
40
80
100
60
70
90
100
60
60
100
60 Acc. (target)
100
Acc. (target)
Acc. (target)
Acc. (ta (arget)
Figure3 3: Visualization of pairs of accuracies (keep-targe() for our U-LightOT solver and other OT/EOT methods in the image translation experiment. The values of unbalancedness parameters for our U-LightOT solver (7) and [16] UOT-FM] (A reg_m) are specified directly on the plots.
Experiment 16 UOT-FM]
9 UOT-SD] 770 UOT-GAN]
U-LightOT (ours)
Time (sec)
03:21
18:11
16:30
02:38
Table 3: Comparison of wall-clock running times of unbalanced OT/EOT solvers in Young- AAdult
translation. The best results are in bold, second best are underlined.
Metrics. We aim to assess the ability of the solvers to perform the translation latent codes keeping the class of the input images unchanged, e.g., keeping the gender in Adult Young translation. Thus, we train a 99% MLP classifier for gender using the latent codes of images. We use to compute the accuracy of preserving the gender during the translation. We denote this number by accuracy (keep).
Meanwhile, it is also important ensure that the generated images belong to the distribution of the target images rather than the source ones, e.g., belong to the distribution of young people in our example. To monitor this property, we use another 99% MLP classifier identify whether each of the generated images belong to the target domain or to the source one. Then we calculate the fraction of the generated images belonging to the target domain which we denote as accuracy (larget).
Results. Unbalanced OT/EOT approaches are equipped with some kind of unbalancedness parameter (like 7 in our solver) which influences the methods' results: with the increase of unbalancednesss, accuracy of keeping the class increases but the accuracy of mapping to the target decreases. The latter is because of the relaxation of the marginal constraints in UEOT. For fair comparison, aim to compute the trade-offs between the keep/target accuracies for different values of the unbalancednesss. We do this for our solver and UOT-FM. We found that GAN-based unbalanced approaches show unstable behaviour, so we report UOT-SD and UOT-GAN's results only for one parameter value.
For convenience, we visualize the accuracies' pairs for our solver and its competitors in Fig. 3 The evaluation shows that our U-LightOT method can effectively solve translation tasks in high dimensions (d 512) and outperforms its alternatives in dealing with class imbalance issues. Namely, we see that our solver allows for achieving the best accuracy of keeping the attributes of the input images than other methods while it provides good accuracy of mapping to the target class. While balanced methods and some of the unbalanced ones provide really high accuracy of mapping to the target, their corresponding accuracies of keeping the attributes are better in dealing with class imbalance issues. As expected, for large parameter т, the results of our U-LightOT solver coincide with those for LightSB which is balanced solver. Meanwhile, our solver has the lowest wall-clock running time among the existing unbalanced sol solvers, see Table 3 for comparison. We demonstrate qualitative results of our solver and baselines in Fig. 4 The choice of unbalancedness parameter for visualization of our method and UOT-FM is detailed in Appendix B
worse than ours meaning that
We present the results of the quantitative comparison in the form of tables in Appendix D.I In Appendix C we perform the ablation study of our solver focusing on the selection of parameters 7, and number of Gaussian mixtures' components.
6 Discussion
Potential impact. Our light and unbalanced solver has a lot of advantages in comparison with the other existing UEOT solvers. First, it does not require complex max-min optimization. Second, it provides the closed form of the conditional measures ya (ylax) se (y(z)) of the UEOT plan.

---

Ingatt
DOT-50
DOTESD
UCT-CAN
LOT-GAN
OT-FM
LightS8
UOT-FM
LOT--M
Qurs
OuTS
(a) Young + Adult
(b) Adult
Young
Figntt
Figit
LOT-DD
UOT-SD
DOT-CCA
UOT-GAK
Lighesd
DOT-FM
DOTE
OUTS
Man Woman
(d) Woman Man
Figure 4: Unpaired translation with LightSB, OT-FM, UOT-FM and our U-LightOT solvers applied
in the latent space of ALAE (55 for FFHQ images 34 (002xx1024)
Moreover, allows for sampling both from the conditional measure (o(y))r) and marginal measure ua (ar) Rs (x)) Besides, the decisive superiority of our lightweight and unbalanced solver is its simplicity and convenience of use. Indeed, it has straightforward and non-minn non-minimax optimization objective and avoids heavy neural parametrization. As result, our lightweight and unbalanced solver converges in minutes on CPU. We expect that these advantages could boost the usage of solver as standard and easy baseline for UEOT task with applications in different spheres.
The limitations and broader impact of our solver discussed in Appendix E
ACKNOWLEDGEMENT The work Skoltech was supported by the Analytical center under the RF Government (subsidy agreement (00000D73321P500000
Grant No. 70-2021-00145 02.11 2021). We thank Kirill Sokolov, Mikhail Persiianov and Petr Mokrov for providing valuable feedback and suggestions for improving the proofs and clarity of our paper.
10

---

References
[1] R. Agrawal and T. Horel. Optimal bounds between f-divergences and integral probability metrics. Journal of Machine Learning Research, 22(128))1-59, 2021.
[2] M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 214-223. PMLR, 2017.
[3] A. Asadulaev, A. Korotin, V. Egiazarian, and E. Burnaev. Neural optimal transport with general cost functionals. In The Twelfth International Conference on Learning Representations, 2024.
[4] Y. Balaji, R. Chellappa, and S. Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. Advances in Neural Information Processing Systems, 33: 1993-11944, 2020.
[5] C. Bunne, S. G. Stark, G. Gut, S. Del Castillo, M. Levesque, K.- V.. Lehmann, L. Pelkmans, A. Krause, and Rätsch. Learning single-cell perturbation responses using neural optimal transport. Nature Methods, 10711)1779-1768, 2023.
[6] L. Chapel, R. Flamary, Wu, C. Févotte, and G. Gasso. Unbalanced optimal transport through non-negative penalized linear regression. Advances in Neural Information Processing Systems, 34:23270-23282, 2021.
[7] Chen, G.-H. Liu, and Theodorou. Likelihood training of schrödinger bridge using forward- backward sdes theory. In International Conference Learning Representations, 2021.
[8] L. Chizat. Unbalanced optimal transport: Models, numerical methods, applications. PhD) thesis, Université Paris sciences et lettres, 2017.
[9] J. Choi, J. Choi, and M. Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport. In Advances Neural Information Processing Systems, volume 36, 2024.
[10] M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pages 2292-2300, 2013.
[11] M. Daniels, T. Maunu, and P. Hand. Score- generative neural networks for scale optimal transport. Advances neural information processing systems, 34:12955-12965, 2021.
[12] Q. Dao, B. Ta, T. Pham, and A. Tran. Robust diffusion gan using semi-unbalanced optimal transport. arXiv preprint arXiv:2311.1710.1 2023.
[13] V. De Bortoli, J. Thornton, Heng, and A. Doucet. Diffusion schrödinger bridge with applica- tions to score-based generative modeling. Advances in Neural Information Processing Systems, 34:1 14:17695-17709 2021.
[14] N. Deb, P. Ghosal, and B. Sen. Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections. Advances in Neural Information Processing Systems, 34: 29736-29753, 2021.
[15] P. Dvurechenski,, D. Dvinskikh, A. Gasnikov, C. Uribe, and A. Nedich. Decentralize and randomize: Faster algorithm for Wasserstein barycenters. In Advances in Neural Information Processing Systems, pages 10760-10770, 2018.
[16] L. Eyring, D. Klein, T. Uscidda, G. Palla. N. Kilbertus, Z. Akata, and F. Theis. Unbalancedness in neural monge maps improves unpaired domain translation. In The Twelfth International Conference Learning Representations, 2024.
[17] J. Fan, A. Taghvaei, and Y. Chen. Scalable computations of Wasserstein barycenter via input convex neural networks, arXiv preprint arXiv:2007.04466. 2020.
[18] J. Fan, S. Liu, S. Ma, H.-M Zhou, and Y. Chen. Neural monge map estimation and its applications. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https://openreviex: nwt 2ttp://eereeeea.eeooooooeaaa/ane
Featured Certification.

---

[19] K. Fatras, Y. Zine, R. Flamary, R. Gribonval, and N. Courty. Learning with minibatch wasser- stein: asymptotic and gradient properties. In AISTATS 2020-23nd International Conference on Artificial Intelligence and Statistics, volume 108, pages 1 1-20, 2020.
[20] K. Fatras, T. Séjoumé, R. Flamary, and N. Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In International Conference on Machine Learning, pages 3186-3197. PMLR, 2021.
[21] R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel
A. Corenflos, K. Fatras, N. Fournier, et Pot: Python optimal transport. The Journal of Machine Learning Research, 22(1):3571-3578, 2021.
[22] M. Gazdieva, I.. Rout, A. Korotin, A. Kravchenko, A. Filippov, and E. Burnaev. An optimal transport perspective on unpaired image super-resolution. arXiv preprint arXiv:2202.01 116, 2022.
[23] M. Gazdieva, A. Korotin, D. Selikhanovych, and F. Burnaev. Extremal domain translation with neural optimal transport. In Advances in Neural Information Processing Systems, volume 36, 2023.
[24] A. Genevay. University, Mar. 2019. URL httttt::/tttteeeseeeeseeeee hal www.eneneeeeeee2e213990333
Optimal Transport for Machine Learning. Theses, PSL
[25] A. Genevay, G. Peyré, and M. Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pages 1608-1617. PMLR, 2018.
[26] N. Gozlan, C. Roberto, P.-M. Samson, and P. Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, 173(11):3377-3405 2017.
[27] I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville. Improved training of Wasserstein GANs. In Advances Neural Information Processing Systems, pages 5767-5777,
2017.
[28] N. Gushchin, A. Kolesov, A. Korotin, D. Vetrov, and E. Burnaev. Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems, 2023.
[29] N. Gushchin, A. Kolesov, P. Mokrov, P. Karpikova, A. Spiridonov, E. Burnaev, and A. Korotin. Building the bridge of schr'," odinger: A continuous entropic optimal transport benchmark. arXiv preprint arXiv:2306. 10161, 2023.
[30] N. Gushchin, S. Kholkin, E. Burnaev, and A. Korotin. Light and optimal schr\ "odinger bridge matching. In arXiv preprint arXiv:2402.0322 2024.
[31] G. H. Hardy, J. E. Littlewood, and G. Pólya. Inequalities. Cambridge university press, 1952.
[32] J.-C. Hütter and P. Rigollet. Minimax estimation of smooth optimal transport maps. 2021.
[33] H. Janati, B. Muzellec, G. Peyré, and M. Cuturi. Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems, 33: 10468-10479, 2020.
[34] T. Karras, S. Laine, and T. Aila. A style-based generator architecture for generative adversar- ial networks. In Proceedings of the IEEE/CVF conference on computer vision and puttern recognition, pages 4401-4410, 2019.
[35] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.69p0 19800990, 2014.
[36] D. Klein, Uscidda, F. Theis, and M. Cuturi. Generative entropic neural optimal transport to map within and across spaces. arXiv preprint arXiv:2310.09254, 2023.
[37] A. Korotin, L. Li, A. Genevay, J. M. Solomon, A. Filippov, and E. Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34:14593-14605, 2021.
12

---

[38] A. Korotin, L. J. Solomon, and E. Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021.
URL htt.htttts::/pperevieei httpeprrwweae..eeffoouoddd/dd--
[39] A. Korotin, A. Kolesov, and E. Burnaev. Kantorovich strikes back! wasserstein gans optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022.
[40] A. Korotin, D. Selikhanovych, and E. Burnaev. Kernel neural optimal transport. In International Conference on Learning Representations, 2023. URL http://///peeeeeeeeeeeeeeeeeeeee id=Zuc MHtUma4]
[41] A. Korotin, N. Gushchin, and E. Burnaev. Light schr)," odinger bridge. In The Twelfth International Conference on Learning Representations, 2024.
[42] T. Koshizuka and I. Sato. Neural lagrangian schr\ "{o} dinger bridge: Diffusion modeling for population dynamics. The Eleventh International Conference on Learning Representations, 2022.
[43] M. Liero, A. Mielke, and G. Savaré. Optimal kantorovich distance between positive measures. Inventiones mathematicae, 211(3):99--111) 2018.
problems and new hellinger-
[44] Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747) 2022.
[45] Lübeec, C. Bunne, G. Gut, J. S. del Castillo, L. Pelkmans, and D. Alvarez-Melis. Neural unbal- anced optimal transport via cycle-consistent semi-couplings. arXiv preprint arXiv:2209.16621 2022.
[46] M. Lucic, K. Kurach, M. Michalski, S. Gelly, and O. Bousquet. Are GANs created equal? a large-scale study. In Advances in neural information processing systems, pages 700-709, 2018.
[47] A. Makkuva, A. Taghvaei, S. Oh, and J. Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pages 6672 6681. PMLR, 2020.
[48] T. Manole, S. Balakrishnan, J. Niles-Weed, and L. Wasserman. Plugin estimation of smooth optimal transport maps. arXiv preprint arXiv:2107.12364, 2021.
[49] M. Mohri, A. Rostamizadeh, and A. Talwalkar. Foundations of machine learning. MIT press, 2018.
[50] P. Mokrov, A. Korotin, and F. Burnaev. Energy-guided entropic neural optimal transport. In The Twelfth International Conference on Learning Representations, 2024.
[51] Q. M. Nguyen, H. H. Nguyen, Y. Zhou, and L. M. Nguyen. On unbalanced optimal transport: Gradient methods, sparsity and approximation error. arXiv preprint arXiv:2202.0366 2022.
[52] T. T. Nguyen, H. D. Nguyen, F. Chamroukhi, and G. J. McLachlan. Approximation by finite mixtures of continuous density functions that vanish at infinity. Cogent Mathematics & Statistics, 711)1750061, 2020.
[53] G. Peyré, M. Cuturi, et al. Computational optimal transport. Foundations and Trends@ in Machine Learning, 11(--6)335-6077 2019.
[54] K. Pham, K. Le, N. Ho, T. Pham, and H. Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In International Conference on Machine Learning, pages 7673-7682. PMLR, 2020.
[55] S. Pidhorskyi, D. A. Adjeroh, and G. Doretto. Adversarial latent autoencoders. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 14104-14113, 2020.
13

---

[56] A.-A. Pooladian and J. Niles-Weed. Entropic estimation of optimal transport maps. arXiv preprint v7Xiv:2109.12004 2021.
[57] Rigollet and A. J. Stromme, On the sample complexity of entropic optimal transport. arXiv preprint arXiv: 13472, 2022.
[58] R. Rockafellar. Duality and stability in extremum problems involving convex functions. Pacific Journal of Mathematics, 21(1):167-187, 1967.
[59] R. Rombach, A. Blattmann, D. D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[60] L. Rout, A. Korotin, and E. Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2022.
[61] V. Seguy, B. B. Damodaran, R. Flamary, N. Courty, A. Rolet, and M. Blondel. Large scale optimal transport and mapping estimation. In International Conference on Learning Represen- tations, 2018.
[62] T. Séjourné, G. Peyré, and F--. Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprini arXiv:2211.0877: 2022.
[63] S. Shalev-Shwartz and Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
[64] A. Taghvaei and A. Jalali. 2-Wasserstein approximation via restricted convex potentials with application to improved training for GANs. arXiv preprint vrXiv:/902.07/00 2019.
[65] A. Tong, J. Huang, G. Wolf, D. Van Dijk, and S. Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference machine learning, pages 9526-952 9536. PMLR, 2020.
[66] F. Vargas, P. Thodoroff, A. Lamacraft, and N. Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9):1134, 2021.
[67] G. Wang, Y. Jiao, Q. Xu, Y. Wang, and C. Yang. Deep generative learning via schrödinger PMLR, 2021.
bridge. In International Conference on Machine Learning, pages
[68] Y. Xie, M. Chen, H. Jiang. T. Zhao, and H. Zha. On scalable and efficient computation of large scale optimal transport. In International Conference on Machine Learning, pages 6892. PMLR, 2019.
[69] Y. Xie, Y. Luo, and X. Huo. An accelerated stochastic algorithm for solving the optimal transport problem. arXiv preprint arXiv:2203.00813 2022.
[70] K. D. Yang and C. Uhler. Scalable unbalanced optimal transport using generative adversarial networks. In International Conference Learning Representations, 2018.
[71] J.-Y. Zhu, Zhang, D. Pathak, Darrell, A. A. Efros, O. Wang. and E. Shechtman. Toward multimodal image-to-image translation. Advances in neural information processing systems, 30, 2017.
14

---

A Proofs
A.1 Proof of Theorem 411
Proof. For convenience, we split this proof into several steps.
Step I. To begin with, we recall that the primal UEOT problem (3) attains a minimizer y- Then by substituting y directly in (3), we get
- 311?
(I, y)d - - Dr (Gall)) + D/2 (75ll9)
12
RR
2
u11
(1,, wddady +e S (1, (z, waddddy - ||1 +
Rd
2
Si (5(zz) p(z)
p(z)do+ [ f2
(g((V)
g(v)dy -ell>" +
gd
((3))
11 1112
|0 -ut
(I,, ydddy + (xx,))) (6 + 2i (v)
drdy + (18)
Rd
2
yog (x,u)=
Jan
Ta(xx) p(x))
p(x)ddo + [ S2 (v)
a(y)dy= = -01>=")11"1 +
((v)
Jec ( /.- ( (v)
f1 (2((x)
(T())
p((2) p(xx)dx+
f2
(y)ddy=
p(x²)
9(v)
Rd
R2d
/ Rd
(x()
Vae (2r) p(x)
/
Tig (1)
Mr (1) ((v)
-elly" 111
fi
x5 (xx)
p(ax)dx+
f2
(v)
q(y)dy. (19)
p(x))
q(y)
Here in line 18
we use equation 5 which establishes the connection between y and the optimal
potentials o* of* delivering maximum to the dual UEOT problem Substituting these optimal
potentials to the dual problem, we get
1 y12 )} Jdddy // (1(-0"(x0)) (x)())))))
expl 10
(x)
RA Rd
2
=1h*
J Rd
-/ Jgd
-f [2((-
(20)
Note that (19) equals to (0) due to the equivalence between primal and dual UEOT problems, see We will use this fact in further derivations.
2
For every function f. its convex conjugate
F(t) supf ut f(u)) = WEB
11--ut-fle (005-f(uu f(u)] = inf fut + f(u)} = wE
ueR
-7(-t) ut ut f(u) Vu R.
(21)
Then for arbitrary a E Rd,, taking f f1, u
(x)
t- = 6² (x), we get
p(.)
-1(-0((x)) < (x)
(1((z) )
8² (z)
x²(x) + fi p(x)
(22)
plz)
-fe fi (1)))((()))d <
2x(x) Dd p(xx)
(x(x))
35 (x) + f1
p(x)) p(x)dd.
(23)
Similarly, for arbitrary y E gd
-
(v(y)
((y) 2 (v) + fz (6(0) (6(0)
q(v) )) a(v)dy.
(24)
15

---

Since 19
20
inequalities 22
23
24
actually turn to equalities.
Step 2. Since consider convex, lower semi-continuous functions f, it holds that 7 f= Then,
substituting f in formula (21 we derive
-((-t)=
inf {ut + f(u)} = =f(--t) S ut + F(u) Vu R.
(25)
HER
Thus, for arbitrary x E Rd taking f = fi, t - p(a) we get that -fi
[1(x))
I (x)
p(x)
plx)
[(--) Vu E R. At the same time, from step I, we know that - fi TT(- (x)), see
[2((a)
22(m)
0 (z)
+
p(x)
p(a)
22 Then
Y
O (x) + +T((- S 22(x)
p(x) + +
p(x)
(x)
(x) u) (Hil -1(
(26)
p(z)
Similarly, for arbitrary y E Rdd
7 (v) (v (y) u) s (fe(-u) (o) VUER.
(27)
(v)
Step 3. Now we ready prove the main result. We note that:
[ x (x. v) log (xx vddxdy
DKL (n ([ee.))
Je 7 (I. v) log (x, fddrdy + thowlll - 111"
(28)
While the ground-truth UEOT plan (x, y) and the optimal dual variables o= (x),, x (y) are con- nected via equation (5), our parametrized plan y0 (x. y) can be expressed using 6 (Ir), (v) as
ne (I, v) exp(s- (1o (x) + voo()) | vll /2)). see (9. Then
28 = = {
JA y (x, 4) (x) + (g)) dxdy
E-1 f
wll daddy + l1
11
JRaxed (14,))(00))+++ (x)
(x) the dw)) + ll ||>"
(x)
(6 (x) de (z) (x))
p(x)
1g(y)
- who (v)) I| - ||
1s ((1)
/ -
Rd
Rd
Rd
(C(0,w) - L').
16

---

A.2 Proof of Proposition 4.2
We begin with proving the auxiliary theoretical results (Propositions A.1 prove the main proposition of this section.
A.2)
which are needed to
Proposition A.1 (Rademacher bound on the estimation error). It holds that
B(C(0) [(0)] (Fi.P + 4R (F2,)
where Fi = -011-000 E x 22. F2 = (7((-v0) E Q for (o = log (u(a)
ca(z)
bo(y) = E log vy(y) + I!lll and Rx ()),)). (M q) denote the Rademacher complexity
63
626] of the functional classes U, V w.r.. to the sample sizes N, M of distributions
Proof of Proposition A.J
The derivation of this fact is absolutely analogous to
41 Proposition B.1],
50 Theorem 4] or 664 Theorem 3.4].
Proposition A.2 (Bound the Rademacher complexity of the considered classes). Let S A. let 0 < u & U. let 0 < w S W and > O. Consider the class of functions
||
F = (a+ F log uw (x) + log co(x)
2
F2 {y+ llg 1ul12
where (1) voly), co(x)) belong to the class
== for Σ on exp Uxr + o wk) with
(29)
ul Un = 3
V; wk S W; a ax SA..
Proposition B.2 we call the functions of the class V as constrained log-sum-exp quadratic functions. We assume that Fr, Jn are Lipshitz functions and measures p. are compactly supported with the supports lying in zero-centered ball radius R > O. Then
Following 41
Co VN'
G1
Rx S
RM S
VM
where the constants Co. C₁ do not depend sizes N, M of the empirical samples from p, q.
Proposition B.2], the Rademacher complexities of con- strained log-sum-exp quadratic functions + log (x), x + log co(x)) and y + log (y) are known to bounded by O VN
Proof of Proposition A.2 Thanks
O
VM respectively. According to the definition of Rademacher
z=x complexity, for single quadratic functions & + 2 (y + =²y) ) it is just equal to zero. Then, using the well-known scaling and additivity properties the Rademacher complexity [63], we get that
-E log (x) log co ||||| and y -5 log (v ||yll are bounded by O (x)) and
O M respectively. The remaining step is to recall that [1(x) (x) and f2(u) are Lipschitz. Therefore, according Talagrand's contraction principle [49] the Rademaher complexities of Fi and F2 also bounded by O and O respectively.
Proof of Proposition
4.2
The proof directly follows from Propositions A.I and A.2
A.3 Proof of Theorem
4.3
To begin with, we provide a quick reminder about the to derive the dual form of problem (3)
theorem which needed
Theorem A.3 (Fenchel-Rockaellle (58). Let (E, E') and (F, be two couples of topologically
paired spaces. Let A E + F be continuous linear operator and A F² + E'' be its adjoint. Let
17

---

f and be lower semi-continuous and proper convex fumctions defined on E and F respectively. If there exists dom s.t. g is continuous at Ar, then
-f--f))--gll-- zEE
sup
min J(A 7) + +
yEP"
and the min is attained. Moreover, if there exists maximizer a E E then there exists y E F satisfying Ar E Ög(T) and An E 0f(-x).
We note that in the below theorem C2 (2 does not denote the space of twice differentiable continuous functions. The exact definition of this space and C2,6(R.)) (R.4) is given in notations part.
Theorem A.4 (Dual form of problem (3)) The primal UEOT problem
3 has the dual counterpart
4 where the potentials (ø, 4p) belong to the space (2R)(RR)) x C2.b The minimum of is attained for a unique 7 E M22-(R2 x Rd).
Proof. We recall that in the primal form, the minimization is performed over functions belonging to M2,+ (R x Rd). In this proof, we suppose that this space is endowed with a coarsest topology
(M2,+ (I2 d R²)) which makes continuous the linear functionals y (dy, AÇ C2(IR x Rdd.. Then the topological space (M2,+ (R3 Rd), (M2,, (R Rd)) has topological dual (I x ), a (MM2+ (TRd x [) which, actually, is (linear) isomorphic to the space C2 x R 26 Lemma 9.9]. This fact opens opportunity apply the well-celebrated Fenchel- Rockafellar theorem. For this purpose, we will consider the following spaces: E C2(IR) x C2 C2(4),
F def (IR" x IRd) and their duals E" def M2,+ (I²") x M2 (I2) and fr def M2, (Rd x Rd).
Step J. Recall that the convex conjugate any function g M22+ (I RR for each Ç (M2,+ (R4 x R = C2 x IR as Ø(5) = (B f(7, g For the convenience of further derivations, we introduce additional functionals corresponding to the summands in the primal UEOT problem (3):
) ( defined
P(r) def Jge JRd
{
(1x y)dxdy - cH(())
2
(7(%)) def
Dr (Tellp);
F2(7w) del D5a (7vll4).
(30)
For our purposes, we need calculate the convex conjugates of these functionals. Fortunately, convex conjugates of f f-divergences Fi (7(7)) and F2(T)) are well-known, see [ Proposition 23], and equal to
[
F2(V) / Fa(b(w))(()dd.
(
F(((()))(())d)
To proceed, we calculate the convex conjugate of P(7):
[
| yl
P(()
7(x y) y)dady =11
JR Jzd
2
1/ ( ((1,0)((,, vddddy - (ff
| | (x, v)dzdy
sup (B
JRd
2
eH(r))
sup TEM2, xRd) Jee / (5(x, v)
|| - y11
(I, dddzyd +
2
(pe / -1(x, log (xx ydddy + l11111))
=H(1)
||0-v|2
8 (. / 7(2, v)
((xx y)
log (2x, v) dxdy + lmlll
sup (RU
18

---

sup rEM2, (R+ xRd) Rd J 7(2, (1og
(-E) { {
((x, v)
||
((x,)) (10gg(2;
2
ddxy - [mmlla]
(-8) { fer Jo 7(x,1) (log ((x,)))
((ar, y)
||2-111
sup 7EMs,- (Rt4
drdy 11 +
/ expl ((o, 1) ||=-
Jdady - fpc Jz expl ((x, v) ||> |x==v1
Jdady }
(31)
Ra Rdd
=0
sup MeM2, (R4x2) (-s) { DKL (all eIlexpf S(x, v)
{ DDa
liا-vl
2
8
((1, |ل
J J2 expl
Jxdy
(32)
Here in the transition from (11) to (32), we keep in mind our prior calculations of DKL in that DXL is non-negative and attains zero at the unique point
28
Recall
7(x v) = exp( S(,, v)
--
(33)
8
Thus, we get
P(() = / JRd
expl ()
Ia la-wwl wll
Jddy..
(34)
theorem in case. To begin with, we show that this thcorem is applicable to problem (3). i.c., that the functions under consideration satisfy the necessary conditions. Indeed, it is known that the convex conjugate of any functional (e.g., (.), F2(., P(-)) is lower semi-continuous and convex. Besides, the listed functionals are proper convex Indeed, the properness of Fi (.) and F-(-) follows from the fact that f f-divergences are known to be lower-semicontinuous and proper themselves, while properness of P(-) is evident from 332.
Step 2. Now ready to apply the
Now we consider the linear operator A : C2((R)) x C2(R)) + CIRI x Ied) which is defined as A(ø, v) (z, y) + <(x) + vyy). It is continuous, and its adjoint defined on M2,+ x Rd) as A(v) = (Ya, 7y). Indeed, (A(r), (u, u)) = (y, A(u, v)) = Jg= Jg- ((2 v)(()) + vvy)ddrdy= = F. 7 (x))(x)) + Jne Ty (y) (1)((())) ydy. Thus, the strong duality and the existence of minimizer for 3 follows from the Roccael-R theorem which states that problems
sup (R.)))C(((+xCCC(E))
(-P(A(6, Fi(-)) ( F2(-b)}
(35)
and
min NEM2 (R
(P()) + F1(7)) P2(T0))
(36)
equal. The uniqueness of the minimizer for comes from the strict convexity of P(-) (which holds thanks to the entropy term). Note that the conjugate of the sum of F1 and F2 is equal to the sum of their conjugates since they are defined for separate non-intersecting groups of parameters.
Next we prove that the supremum can be restricted to C2 (Rdd x IRd). Here we use "A" denote
the operation taking minimum between the function Rd R and real value k: (f A (x))) def min( ff (I), k). Then analogously to [26] Theorem 9.6], we get:
Proper convex function a real-valued convex function which has non-empty domain, never attains the value (-oo) and is not identically equal to (+oo). This property ensures that the minimization problem for this function has non-trivial solutions.
19

---

sup (0, ((,)))C(
{ P(A(O) Pi(-0) -2(-v))
sup (B.
lim 720000000
((A K1, 15 A 12)) Fi(-(AAk))) -2(-(1 k2))) S
(-P(A(Ó, -) -
(-(-)) - F2(-v)} (37)
sup (()V))C2
Since the another inequality is obvious, the two quantities equal which completes the proof.
Proof of Theorem44. 4.3 Our aim is to prove that for all 8 > 0 there exist unnormalized Gaussian mixtures u. and Up s.t. ((0,W) òc. We define
J 1)) det
J exp( (yy) | U112 Jrrdy- + {
J Jí( F((-6()))((adrr+ J J2(-v(u))ggt
Then from and Theorem A.44 it follows that L = inf (E))x))) J (6. Finally, using the definition of the infimum, we get that for all 8' there exist some functions (ō, 3)
(2,)(R (IR such that J [6, +) + For further derivations, set 8' ás and pick the corresponding (Ø, v
Step 1. We start with the derivation of some inequalities useful for future steps. (é, ) E xx2,((R))xxx11( and (vy) respectively. We recall that by the assumption of the theorem, measures p and g are compactly supported. Thus, there exist balls centered at = and = 0 and having some radius R > O which contain the supports of p and q respectively. Then we define
Since they have upper bounds à and 6 such that for all x, R4. x(x) S
x(x) ((x) max (0, max [||=|1 R² || Rt} < )(x) s ā;
((v) (1)) max 1lvl1 R² (2)) S 6.
We get that
((x) s ((x), (5v) S ())) =
f expl ((0) v()) || y/1
}dxdy z
Ra Jdd
2
{ Rd c(o) (())
|
1 drdy
(38)
2
Importantly, for all and within the supports of p and it holds that x(x)) 3(x) and ((v) E(y). respectively. Then
( J }
Rd
[ [2(- = - Ja(-)(v))a(vddd
(39)
Rd
Combining (38) and 339 we get that [(,, 3) < J [ v) < C* +
Before moving on, we note that functions exp( ((tt)// and (xp(vl))/e} are continuous and non- negative. Therefore, since measures p and q compactly supported, there exist some constants Cmin: Amin > O such that exp( @(w)/c}} > Cmin and exp(((s)/c} > hmin for all and y from the supports measures p and g respectively. We keep these constants for future steps.
Step 2. This step of our proof is similar to 41 Theorem 3.4]. We get that
max (0, ||>|| 1 R²)
à R²
exp (ōle)/2) S exp
I exp
exp((11=x1//)
(40)
20

---

max max |1yl1 - R2)
5 R
exp (V/n//) S exp
exp((llall/4)
exp
From this we can deduce that + exp( (W)))e) is normalizable density since it is bounded by the unnormalized Gaussian density. Moreover, we see that vanishes at the infinity. Thus, using the result [52] Theorem 5a], we get that for all 8 >> 0 there exists an unnormalized Gaussian mixture Vg = v3) such that
1lg -xxp(%/e))l sup (uz(y) exp exp(#(l//=) < 8". VERD
(41)
Following the mentioned theorem, we can set all the covariances in Ug to be scalar, i.e., define vg(y)= 1 (gN(ylFF..Eee for some K and ag R4. Fx Rd d R4 (k ......,KK For our future needs, we set
8" = òE 2
E hmin
<(TTE)* (1 (TT) exp a + R²
L
L2
emin
For simplicity, we consider the other mixture (v) def
llvl which is again unnormal- 1ig (2(()) exp( 2z
ized and has scalar covariances, see the proof of [4]] Theorem 3.4] for explanation. We denote the weights, means, and covariances of this mixture by ag R4, Th IRd and d E R4, respectively.
We derive that
J
Rd / exp (=((()) + +(y)
I - -yl?
) Jdxdy =
2
/ / expl ((T) expl
1
y11
expl {1))
Jdxdy
BB RRd
28
8
/ expl H(t)
|| - 1
exp
H(z(() - 88 dodyy
JRd Rd
2e
expl ((x)
|| -112
expl
}ra(y)
Rd Rd
expl (xx)
| y11²
8
Texp(
Jdrdy
Bd
Rdd
2x
expl ((x)
|
(1, v)
|19/12
texp
texp(
texpl
(oa(v) dxdy
3d
28
2s
=volv)
exp( ((x)
expp ||x
112
8"'
drdy =
Rd
E
25
((x)
|
f expf
(1
exp(
expf
Hoolyddy dr
Bdd
28
=co(2)
expl (((a)(()
||:0 - y11
8" E
expf
25 Idy do
gd
Rd
=(2rre
cf expl
((t)) E
||02|
JRc expf 0(x)
expl
tco (x)dx 8"<E(2Tt)"/a
gd
20
[(x)
||-2|
8
expl
expf 25
)co (2)
Rd
21

---

8" (2eme exp a+ R62
(==1t))
=(x=)
((x))
||-2|| 25
a+ R
8 expl
gd
expl
)co - a exp
(42)
Step 3. At this point, we will show that for every 8" > 0, there exists an unnormalized Gaussian mixture ug which is -lose to exp(((x)/))cg(x)) Using the closed-form expression for Cp(z) from 41 Proposition 3.2], we get that
|r
Ir + x>a1l1
(0(x)= Σ an exp(
exp(
Then
(xp()t)/))co((t) s
à - max (0, |>||* - 184) -1t")
(oo(2) S exp (a à @+
exp((1lt1)%/)colf)
exp
k
ā + Rtd
||r2|1
|0||| eexpf lrre + xd
0% exp
expf
expl
E
2=Ax
E
2=Ax
k=1
Ra
||rx11 } expl
|| + xAael' 1|2 2Ax |++|
Og exp
exp(
From this, we see that exp ((())/))co()) tends to zero while x approaches infinity. It means that I + exp (o corresponds to the normalizable density. Using [52] Theorem Sa], we get that for all 8" > 0 there exists an unnormalized Gaussian mixture ug such that
|/ exp((/e)coll) = sup (uz(x) exp(@x=l/)cc(())) 8". 8".
(43)
ZERO
Analogously with ug. we can set all the covariances in ua to scalar, i.e., define to Σl_ Nx eλila) for some L, By E Rd, & 14 ( (J). Moreover, consider Me (x) (1) exp ||]|]] which is again an unnormalized density with scalar covariances. We denote the weights, means, covariances of this mixture by Br Rt, kl Rd & R4 respectively.
Next we recall the equation (42) and get that
|0
+ Rt
42
exp(-
Mus(x) - 8"") 8"22/2 (d(())) exp {
25
Rd
||0||
||=||
+ R²
exp(
juz (x)) dx
88"
expl
)dk - 8"22/" 4 (d+1)
exp
Rd
25
Rd
23
(n)
=(2me)
El E||uw |11
=ō" {2TTE)
(1 [TE)
exp ā a+ R22
1)) (44)
8
Step 4. Now turn to other expressions. Using the property that function + logt emin Lipshitz on interval [fmin +-00 we get
(v) g })
(1())) }- -8") <
8"
log(exp(
log((epp (exp
hmin
((v)
}- 8n)) > F(y) 8"
log(expf
(45)
hmin
22

---

Similarly, we get that
lxg( Φ(v)
)x)(x)) 8") 2 Ō(v)
8"
+ log co (x)
(46)
Emi
We use this inequality, monotonicity of logarithm function, and 41 to derive
(1)
(1))
43
(b))
8"
(v( exp E }- - 88"- - log (g(()) > lex(e
8")
E
hmin
E8"
-v(y) - log (g(())
(47)
hmin
expl o(x)
(x)) -8" - log (x) > log(exp( x(x))
46
up (r)
Jco(x)
8"")
(vy) + log (z)
8"
(z(
E8"
o(y)
-Elog
(48)
E
Ein
ca
cx(x)
Emin
Recall that fi and f2 are non-deereasing functions. Moreover, they are Lipshitz with the constants T1, L2 respectively. Thus, we get
Eō"
l1u11
sō"
Fa(-((o)) 2 Fel log
= Fel (e(l-(log ve vv()) +
hmin
2E
hmin
e8"
2 /2 log vw)) |1yl
e8"
[2( - log (v)
L
(49)
2
Amin
2
hmin
uz
E8"
E8"
TT ((-)(())) 2 J1 -e log
11 log ( e log cg(xx)
(x(r)
mmin
Cmin
||011
E8"
the (x)
||00||
e8"
Fl(-eloguu(1) (x))
+ log (o)
) FT(-=llg
2
emin
co(x)
Gmin
Uw (wr) | 1
E8"
Ti(- log
(50)
co (x)
2
emin
Integrating these inequalities over all æ and y in supports of p and q respectively, we get
[ Rd F2(--(0)a / Fe(-e vo(y)
|1w12
E8"
)a(y)dy L1-
(51)
Jdd
hmin
(x(
||x:1|
3d {((Ō(()))(adda > fi(-6 log
(p(x)dt b
(52)
co co(xx)
2
emin
Finally, combining
44 and 52
we get
[(0,w) = J log uw(x)
|
|1y/12
+
log 0o (v) +
cp co(x)
2
2
Vo drf
$0.
a|luw|l 11l f Jil-s log cx(xx)
(x)
||0||
||y|1 )a(w))d
)(x)ddr
Fol-eloguo(e
Rd
expl (((o)) + x(y)
|| v11
)ddrdy
JR JRd
I J 7((-)(=))(())drr
Je( T2(--(u))a(uddd
88"
e8"
+ (2TE)) 1 (1 (mE) exp { a+ R² }) <
L
L2
6min
hmin
23

---

a R²
L + 0'+8"
L1
+
L2
e <(2TT)
(1+ (TE)
exp
Emin
hmin
8
L + de
2 + de
2
DEL (1o S ( L') <d
which completes the proof.
Remark. In fact, the assumption of the Lishitzness of (x), (2(x) can be omitted. Indeed, under the
"everything is compact" assumptions of Proposition 4.2 and Theorem 4.3 inputs to fi(.), J2(.) also always belong certain compact sets. The convex functions are known to be Lipshitz compact subsets of R, see 31 Chapter 3, $18]. and we actually do not need the Lipschitzness on the entire I.
B Experiments Details
B.1 General Details
To minimize objective we parametrize og- Sy of vy and 91, E. of uw in (11). Here we follow the standard practices in deep learning and parametrize logarithms log ag log 31 instead of directly parameterizing ay: 3r. In turn, variables Tk: la are parametrized directly as multi- dimensional vectors. We consider diagonal matrices Sk. Dn and parametrize them via their diagonal values (og(S))s, h,, and ( respectively. We initialize the parameters following the scheme in 41 In all our experiments, we use the Adam optimizer.
B.2 Details of the Experiment with Gaussian Mixtures
We use K L = 5, s = 0.05, = 3e 4 and batchsize 128. We do 2 10 gradient steps. For the LightSB algorithm, we use the parameters presented by the authors in the official repository.
B.3 Details of the Image Translation Experiment
We use the code and decoder model from
We download the data and neural network extracted attributes for the FFHQ dataset from
In the Adult class we include the images with the attribute Age 44; in the Young class with the AgeE 44]. We excluded the images with faces of children to increase the accuracy of classification per gender attribute. For the experiments with our solver, we use weighted DxL, divergence with parameters 5 specified in Appendix C. and set K = == 10, E = 0.05, Ir = 1. and batch size to 128. We do 5 10 gradient steps using Adam optimizer 355 and MultiStepLR scheduler with parameter 0.1 and milestones= [500, 1000). For testing [1 LightSB] solver, we use the official code (see
the link above) and instructions provided by the authors.
Baselines. For the OT-FM and UOT-FM methods, we parameterize the vector field (vr,)) for mass transport using 2-layer feed-forward network with 512 hidden neurons and ReLU activation. An additional sinusoidal embedding (65) was applied for the parameter t. The learning rate for the Adam optimizer was set le-4. To obtain optimal transport plan π (z, y) discrete OT solvers from the POT 21 package were used. These methods are built on the solutions (plans x' (x., y)) of discrete OT problems, to obtain them we use the POT (1 package. Especially for the UOT-FM, we use the ot unw.lanced. sinkhorn with the regularization equal to 0.05. We set the number of training and inference time steps equal to 100. To obtain results UOT-FM for Fig. 3 we run this method for 3K epochs with parameter reg_m E [5e- 4, 5e 3, 5e 2, 0.5, 1, 10, 10 103 10 10 109] and reported the mean values final metrics for independent launches with different seeds. In
Tables 20 21 22 for each translation we report the results for one chosen parameter specified in
24

---

(a) U-LightOT, 7 10
(b) U-LightOT,
(c) U-LightOT,
(d) U-LightOT, T
Figure 5: Conditional plans "ie. (y/æ² learned by our solver with scaled xx² divergences in Gaussians Mixture experiment E 2,, 10)).
Appendix D.I
We use the corresponding checkpoints of UOT-FM to visualize its performance in Fig.
4 For 9 UOT-SD), 70 UOT-GANJ we use the official code provided by the authors, see the links:
While both UOT-SB, UOT-GAN methods were not previously applied to the FFHQ dataset, we up a grid search for the parameters and followed the instructions provided by the authors for parameter settings. Both for UOT-SB, UOT-GAN, we used a 3-layer neural network with 512 hidden neurons, and ReLU activation was used for the generator networks and the potential and discriminator, respectively. Adam optimizer (35] with lr = 10-5 and Ir = 10-4 was used to train the networks in UOT-SB and UOT-GAN, respectively. We train the methods for 10K iterations and a batch size to 128. For UOT-SD, we used DRI. divergence with their unbalancedness parameter T = 0.002. For other parameters, used the default values provided by the authors for CIFAR-10 generation tasks. For all baseline models which use entropy regularization, we set = 0.05.
C Additional Discussion & Experiments with fdiivergences
Details about f-divergences between positive measures. In the classic form, f-divergences are defined measures of dissimilarity between two probability measures. This definition should be revised when dealing with measures of arbitrary masses. In the paragraph below we show that if the function f is convex, non-negative, and attains zero uniquely point {1} then Dy (pr llp2) is valid measure of dissimilarity between two positive measures.
Let f1, M2 E M2,+ (RR be two positive measures. The f-divergence satisfies Dy (un l|pa) 2 which is obvious from the non-negativity of f. From the definition of Dy (44 (12))))))))))) and the fact and only fi (z) u2 (z() holds up-everywhere. Actually, fi (ar) p2(x)) should hold for all as μn must be absolutely continuous w.r.t. μ42 (otherwise D, (Ma (l22) is assumed to equal +00).
that function f attains zero uniquely at a point {1}, we obtain that D, (ui (|
The usage of Dx² divergence. We tested the performance of our solver with scaled DKL divergences in the main text, see $5.1 $5.2 For completeness, here we evaluate our solver with Dx² divergence in Gaussian Mixture experiment. We use the same experimental setup as in $5.1 and present the qualitative results in Fig. 5 Interestingly, the solver's results differ from those which we obtain for DKL. divergence. For D divergence, supports of learned plans' marginals constitute only parts of source and target measures' supports when 5 = The issue disappears with slight increase of r, i.e. for 2. At the same time, further increase of 7 is useless, since the learned plans fail to deal with class imbalance issue. Thus, parameter 7 should be adjusted heuristically. In the case of DRI. divergence, supports coincide for all see Fig. This motivates to use DRi. divergences in our main experiments.
Parameter in Gaussian Mixture experiment. An ablation study unbalancedness parameter in Gaussian Mixture experiment conducted in $5.J and above in this section. For completeness, we also perform the quantitative assessment of our solver with DRL. divergence for different unbal- ancedness parameters т. We compute the normalized OT cost (R (a-v)² source and generated distributions, and the Wasserstein distance between the generated and target
between the
25

---

distributions (computed by discrete OT solver). For completeness, we additionally calculate the metries for the balanced [41 LightSB] approach. The results are presented in Table
Method
U U-LighOTT (r = 1) U-LighiOT (r = 10) U-LightTT (r = 50) U UiLighOT (F = 100) LightBB
OT cost (3)
2,023
2,913
3,874
3.931
3.952
Wardistance (J)
2,014
1.107
0.138
0.091
0.088
Table Normalized OT cost between the input and learned distributions, and Wasserstein distance between the learned and target distributions in Gaussian mixture experiment.
Results. Recall that the unbalanced nature of our solver leads to two important properties. Firstly, our solver better preserves the properties of the input objects than the balanced approaches indeed, allows for preserving object attributes (classes) even in the case of class imbalance. Secondly, due to the relaxed boundary condition for the target distribution, the distribution generated by our solver naturally less similar to the target distribution than for balanced methods.
The above intuitive reasoning is confirmed by the metries we obtained. Indeed, as the parameter increases, when our method becomes more and more similar to balanced approaches, the normalized OT cost between the source and generated distributions increases, and the Wasserstein distance between learned and target distributions decreases. LightSB [1] baseline, which is purely balanced approach, shows the best quality in terms of Wasserstein distance and the worst in terms of OT cost.
Parameters 7, & in image experiments. The effect of entropy regularization parameter & is well Namely, increasing the parameter stimulates the conditional distributions yo to become more dispersed. Still, below we provide an additional quantitative analysis of its influence on the learned translation. Besides, we address the question how does the parameter influence the performance of our solver in image translation experiments? To address this question, we learn the translations Young- Adult, Man-> Woman on FFHQ dataset varying the parameters т, e, see $5.2 for the experimental setup details. We test our solver with scaled DEL divergence training it for 5K iterations. Other hyperparameters are in Appendix B In Tables 5 8 14 we report the accuracy of keeping the attributes of the source images (e.g., gender in Young AAdll translation). In we report the accuracy of mapping to the correct target class (e.g., adult people in we report FD) metrics which is defined as Frechet distance between means and covariances of the learned and the target measures. For convenience, we
studied, see, e.g., 28 41
Tables 6
9 [5] [5]
Young Aduli translation). In Tables 7 TO 13 16
additionally illustrate the results of ablation studies on Fig. 6
10
20
50
10a
250
10a
10
0.01
96.37-10.07
96.12:10.25
93.4411.22
00.26:10.18
0.0
81.79
19.2
0.05
96.03-+0.08
96.55=0.13
93.1920.55
88.93++.94
1.64
80.50
0.77
0.1
95.20110.22
94.92.10.24
92.30
87.99
82.91
78.76
78.25
0.5
08.53+0.22
87.46:0.41
83.41
80.14+0.56
76.00
73.36
71.84+0.27
1.0
11.26=11099
77.31=0.33
74.33+0.54
71.01=1.13
67.13
66.71#0.13
Table 5: Test accuracy (t) of keeping the class in Young - Adult translation.
50
10a
10a
10a
0.01
38.9 98.9110.991
51.92:0.80
67.6821.06 06
75.49:00.00
1..
$3.6720..
15.27
0.05
00.09++0.16
03.19
60.01
77.41+0.67
31.78+0033
01.80=0.90
85.63+00776
0.1
4..18+1.50
56.74-0..8
70.77
78.31+0.91
03.7
87.07
88.21+0.23
0.5
50.51 10.34
65.61122.0
1.117
87.48.10.06
92.21 1.229
90.13
93.80110.55
1.0
40.61
83.82
80.56-10.50
93.78-10.35
95.04
95.50
Table 6: Test accuracy (1) of mapping to the target in Young + Adult translation.
10a
250
10a
10
0.01
26.77+1.599
17.3910.19
06.12-1.5
12.28 1.16
11.4610.70
10.85:E..12
0.05
42.05
31.47:10.03
23.90.11.49
18.31
0.116
17.15 1: 0.48
16.27
1..62
19.8915.551
0.1
20.05
0..6
31.15
0.54
25.86
24.07
26.30-10.15
0.5
114.49+1.06
104 0.33
98.2144.21
92.48=0.98 149.8140.86
89.42=0.10
88.74
1.
1.0
165.39+1.50
142.39
137.40
135.85=0.17
13.125=0.17
Table 7: Test FD of) of generated latent codes in Young - Adult translation.
Results show that increase of & negatively influences both accuracy of keeping the attributes of the source images and FD of generated latent codes which is caused by an increased dispersity of (0())z).
26

---

10
10a
10a
10"
0.01
96.1120.08
95.9640.03
03.60
89.88=0.74
81.1940.89
80.19+2.17
0.05
95.36=0.41
06.26+0.24
02.77=0.48
89.48=0.60
81.3440.35
0.
91.23=0.22
94.33-0.21
0..5
09
81.06+0.82
77.73=1.21
0.5
88.21
0..43
8.91
83.89=0.67
74.92
7..32+1.46
11.67
1..71.1.28
77.15
65.31
Table 8: Test accuracy (1)
keeping the class in Adult + Young translation.
10
10a
10a
10
0.01
0..45
76.38:10.36
0.28
$7.87=0.33
91.56+0.40
02.29:10.37
0.05
00.11
70 10.66+1..2
$4.43+0.40
87 0..39
91.98+0.00
02.29=#0771
65.29+1.54
76.58+1.10
81.29
80.65
92.33+0.00
92.87=00.11
0.5
69.1313.44
76.74:11.39
82.33 1.33316.55
86.86
94.0610..62
94.9110..45
13.13-11.80
0.58
91.00
95.66+40.41
96.37:10.52
Table 9: Test accuracy (f) of mapping to the target in Adult + Young translation.
10
20
50
10a
10a
10a
0.01
45.08
31.2410.62
22.44=0.32
18. 00.23
17.30
24.29
0.05
60.09:13.22 23
3.38
27. 0.007
30.79
21.59=4418
26.11 23
54.45+ 1.07
06.5240.70
5.22
34.02:00.25
44.25=9.33
40.164 7.51
0.5
128.91
121.00
116.78+10.. 10.57
111.49=5.02
102.51
102.31 0.111
1.7.33+1.64
173.11#1.67
171.32+13.12
156.25=44.40
152.73=0.38
Table 10: Test FD (4) of generated latent codes in Adult > Young translation.
10a
10a
10"
93.75
92.01
0.112
88.70100.66
87.62:10.34
0.05
93.40=0.117
0..10
91.91
90.30
88.05-10.20
87.77:10.28
92.99
02.78
91.24:0.46
80.57
87.45+0.62
87.21#0.11
0.5
89.96=0.10
88 88.37
87.21
86.75=1.06
84.09+0.31
83.54+0.41
1.0
84.9410.554
83.5410.41
81.80
80.72110.20
79.87:10.47
Table 11: Test accuracy (1) of keeping the class in Man + Woman translation.
10a
10
10c
61.3810.225 0.29
70.77
90.3210.70
93.7910..2
94.1911..40
0.05
1.11
0..39
80.01
90.23
03.87+0.22
94.4440.07
6..66=0.60
73.52+0.24
86.38
90.38
94.19+0.28
91.73+0.36
0.5
65.74
73.61+0.55
06.45
89.60
95.1420.58
95.59+0.61
1.0
08.1840.18
86.8640.74
91.50=0.74
95.63+0.14
96.93+0.117
Table 12: Test accuracy (1) of mapping to the target in Man + Woman translation.
10a
10
10c
37.8740..55
27.27 13.17
22.60
2..73
17.3241.29
0.05
61.61+2.58
11.25+11.78
34.28 1.04
27. 22.60
30.63
25.77=2.22
0.1
59.45
51.17
43.05
38.02
34.42.1.91
37.54 1.79
0.5
632045=6.007
1.0.82+
1.7.
6.77
103.24=2.32
100.71 1.66
1.0
182 2.2.0
161.51+1.50
156.1142.01
146.73=0.28
116.055#..10
Table 13: Test FD (4) of generated latent codes in Man
Woman translation.
10
10a
0.01
05.70:00.5
00.30
99.68
92.16.0.15
09.08
88.81100.05
0.05
95.55 0.008
05.05
93.67
92.32.0.22
89.66.10.17
88.39110.42
0.1
95.21 00226
04.83-10.19
93. 03.04
91.26
80.09
87.38-10.40
93.02=0.22
02.53+0.29
13.0011.01
80.43
1..87+1.20
86.24-0.41
1.0
00.00
88.83 83.7330.73
87.36
83.3221.08
83.1640.555
Table 14: Test accuracy (1) of keeping the class in Woman -> Man translation.
Interestingly, accuracy of mapping to the correct target class does not have evident dynamics w.r.t.
At the same time, when 7 increases, the learned plans provide worse accuracy for keeping the input latents' class but better FD of generated latent codes and accuracy of mapping to the target
27

---

10
50
10a
10a
10
0,01
51.5120.18
61.67+1009
07.59=0.39
02.52=0.19
88.61#0.48
88. 86.99=0.18
0.05
49. 29.26:0.72
64.09 1.08
76.88=0.17
03.47
88.59+0.65
89. 0.
79.74
63 03.99
0.55
82.15
89. 89.29+0.55
89.14+0.336
0.5
48.64
00.88
0.
12.10=1.17
89.82
90.63 10.633111008
62.62 2.
1.11
82.42 10.73
90.20
90.75.10...
Table 15: Test accuracy (1) of mapping to the target in Woman + Man translation.
10
20
50
10a
10a
0.01
1..00=1.771
33. 63.67
23. 23.85=1.155
20.83=1.27
16.18=0.099
18. 78.4443.44
0.05
52.30+1.200
39.79
29.66
2T.23
24.68+2.80
23.43+1.91
0.1
58
5..60
48.66 10.73
37.55
37.74
31.6310.42
32.83:1..02
0.5
131.17
120.63 0.776
0.664
104.85
101.29=0.11
102.26 1.58
1.0
186.46+0.92
169.64+0.63 63
160. 160.52+0.42
152.78=0.13
152.37=0.09
Table 16: Test FD (4) of generated latent codes in Woman -> Man translation.
Young - Adult
Adult Young
Worm
Man
100
o0
90
Acc. (keep)
80
Acc. (keep)
Acc. (keep)
Acc. (keep)
70-
600
00000
0.05
Cw0.55
=0.01
Cw0.5
0.01
0.05
50
40
Acc. (target)
Acc. (target)
Acc. (target)
Acc. (target)
Figure 6: Visualization of ablation studies on parameters т, E in image translation experiment.
class. It is an expected behavior since for bigger т, the constraints the marginals of the learned plans become more strict. That is, we enforce the marginals of the learned plans to closer to source and target measures which allows learning more accurate mappings target measure but does not allow keeping the source classes in the case of imbalance issues. Interestingly. in Adult Young translation, FD of learned latents and accuracy of mapping to the target do not change much for T 2 102 while accuracy of keeping the attributes exhibits significant drop between T = 10 and 103. Thus, we can treat т 10a as optimal since it provides the best trade-off between the quality of learned translations and their ability to keep the features of input latents. In the case of Addul translation, the values of accuracy and FD exhibit significant differences for considered Thus, we may consider more detailed scale and choose = 2.5 102 as the optimal one.
It important to note that method offers flexible way select domain translation configuration that allows for better preserving the properties of the original objects or generating a distribution closer to the target one. The final optimal configuration selection remains the discretion of the user. The highlighted values in the Tables are used for comparison with other approaches in DD
Remark. FD should be treated as relative measure of similarity between learned and target measures. The results obtained by balanced solver [41] LightSB] (cquivalent to ours for big T) are considered as gold standard.
Number of Gaussian components in potentials. For completeness, we perform an ablation study of our U-LightOT solver with different number on Gaussian components (K, L) in potentials vo and Uwe respectively. We run the solver in Young-AAduli translation with 5K steps, & = 0.05 and set 250. The quantitative results (accuracy of keeping the class, accuracy of mapping to the target, FD of generated latents vs target latents) are presented the Tables below.
The results show that in the considered task, our solver provides good performance even for small num- ber of Gaussian components. This can be explained by the smoothness of the latent representations of data in ALAE autoencoder.
28

---

1
2
16
64
128
01.17=0.45
81.13+122 27
81.25+1.08
80.08
81.1=0092
0..82
85.5220.71
83.70=0.48
2
84.99=0.59
84.5440.43
0..5
04.63
83.87=2.03
03.56+1.37
$4.66:0.44
86.74=0.54
0.86
81.1
01.37
01.65
83.75=101
81.00+1552
03.51
83,99=0,3
8
83.88
84.0810.881
83.7110.60
1..7641.98
84.69=0.38
84.30:10.99
85.05 1.55:1559
83.03=1.02
84.65 0.33
85.00
32
23.48=0.44
86.02+1.22
84.7940.60
04.44
85.24=0.84
84.06 1.00
44.73+0.51
=.24
8..21=0.27
0..23+0.59
81.12+0.12
81.6110.661
81.01=195
11.10=#1220
84.77
83.76=1.41
128
85.21=0.22
85.16:0.00
84.63+1.15
81.6120.49
84.12=0.57
84.1140.76
84.22
84.61=0.93
Table 17: Test accuracy (1) keeping the attributes in Young-+Adull AAdult translation for our U-LightOT
solver with different number of Gaussian components in potentials.
L
2
16
32
61
128
83.35 0.88
82.031 0.11
82.69 0.882
83.34 0.00
82.95 0.46
82.80
82.49 10.64
82.54
2
03.96=0.98
81.80+0.81
81.82
81
0.995
82.77=0.45
02.26:+0.98
82.36+0.377
82.81=0.38
4
0.17
0..47:0.229
8..3
6.66
82.70=0.64
82.48.10.37
83.23 1.22:.1.22
82.44=0.68
82.35
0.38
82.2331..52
82.19.10.72
0..8
82.55.10.118
82.85.10.28
82.87
82.43=0.61
1.
0.73
0..9
82.56=0.26
81.88
83.3440.57
82.90=0.39
32
00.17
01.39+0.62
81.76+0.227
81.28
82.47=1.30
82 02.06:40.41
02.35
82.56=0.88
02.3310.883
81.9740.61
00.22
0.76
50.06
128
82.39
82.1810..56
81.74110.19
82.22
81.5810.23
82.66.1.0.3
83.07 1.030
82.85
Table 18: Test accuracy (f) mapping to the target in -Adul- translation for our U-LightOT solver with different number of Gaussian components in potentials.
K
2
4
8
16
32
64
128
17.41 10.28
17.4240.36
18.10:1.411
17.8310.41
17.97
18.42
17.69
18.07
19.07
17.4710.83
16.87.0.225
17.22
4
10.03
16.66-1..14
17.02:10.84
18.01
16.71
17.07
16.59410.41
16.42:1..10
16.55=0.22
23.56+7.59
16.3740.23
16.81
17.21
18.37
17.13+0.77
16.96=0.86
16
17.88=1.58
18.06+1.61
18.25+2246
17.69=1.87
18.00=0.31
17.19
17.62:#0.93
18.97
10.1
17.1540.77
16.5010.110
2..9144.5
20.04122.6
19.31 3322
18.7944..03
22.5611.660
21.39 1.32
24.91
21.00:4..47
18.02
21.6313.14
19.56
17.53
20.26114.46
128
22.09
35.68421.07
33.80+21.67
20.46:2.33
25.30
22.73
22.09
22.72
Table 19: Test FD (4) of generated latent codes in Young- AAdult translation for our U-LightOT solver with different number of Gaussian components potentials.
D Additional Experimental Results
D.I Quantitative comparison with other methods in Image-to-Image translation experiment
In this section, we provide additional results of quantitative comparison of U-LightOT solver and other unbalanced and balanced OT/EOT solvers in image translation experiment. Tables 20 21 22 provide values used for plotting Fig. in the main text. The unbalancedness parameters used for our UightttOT solver and [16 UOT-FM] are specified in th Tables below. For obtaining the result of UOT-SD], we use their unbalancedness parameter T 0.002. For other details on methods parameters used to obtain the results below, see Appendix B Note that we do not include FID metric for assessing the quality generated images since we found that is not representative metric for assessing the performance of models performing the translation of ALAE latent codes.
LightSR]
U-LighiOT
Experiment
[66
LightSBM]
16 UOT-MMI
700
(ours)
Young +Aduli
67.71
78. 16
11.16 (reg_m n = 0.005)
45.71
73,85
11 = 250)
Adull Young
53.79
80.25
76.05 (reg_m m 0.005)
49.30
74.74
89.48 (r 102)
Mon Woman
76.05
87.82
84.42 (reg_m m 0.005)
75.50
84.04
90.30 (r 10")
Woman Man
88.10
86.10 (reg_m 0.05)
72.03
84.56
89.66
10º)
Table 20: Comparison of accuracies of keeping the attributes of the source images.
Experiment
OT-FM
41 LighiSB] or LightSBMI
Id UOT-PM]
UOT-SD
DOT-GAN
U-LighrOT
[6
(ours)
Young Aduut
93.28
85.97
74.10 (reg_m
0.005)
87.33
84.25
81.78
250)
Adult Young
96. 12
93.10
89.30 (reg_m 0.005)
97.39
95.88
87.79 (r 102)
Mon + Woman
93.33
94.37
92.35 (reg_n 0.005)
98.16
97.38
90.23 (r 102)
88.53 (reg__m 0.05)
88.50 10²)
Table 21: Comparison of accuracies of mapping to the target.
29

---

Experiment
OT-M
[16]
41
LightSBI or Ligh
UOT-FM]
UOFSD
UOIGANN
U-LighOT
(eurs)
Young AAdult
11.93
15.50
11.57 (reg_. 0.005)
13.25
11.23
17.15
250)
Adall Monan
14.10
21.41
17. (reg_m 0.005)
18.44
14.94
30.79 (r 10")
Mon- Woman
16.20
20.91
10.31 (reg_m 0.005)
16.13
22.41
27. (r 10")
Woman Man
11.42
30.87
6.99 (reg_m 0.05)
13.23
10.55
21.68 10")
Table 22: Comparison FD between the generated and learned latents.
D.2 Outlier robustness property of U-LightOT solver
To show the outlier robustness property our solver, we conduct the experiment on Gaussian Mixtures with added outliers and visualize the results in Fig. 7 The setup of the experiment, in general. follows the Gaussian mixtures experiment setup described in section 5.2 of our paper. The difference consists in outliers (small gaussians) added to the input and output measures.
(a) Measures p, q
(b) LightSB
(c) U-LightOT (ours)
Figure 7: Conditional plans "Yo "o (y)(z) learned by our solver (r 1) and LightSB in Gaussian Mixtures with ofliers experiment.
The resulty show that our U-LightOT solver successfully eliminates the outliers and manages to simultaneously handle the class imbalance issue. At the same time, the balanced LightSB [4]] solver fails to deal with either of these problems.
E Limitations and Broader Impact
Limitations. One limitation of our solver is the usage of the Gaussian Mixture parametrization which might restrict the scalability of our solver. This points the necessity for developing ways optimize objective with more general parametrization, e.g., neural networks.
Broader impact. Our work aims advance the field of Machine Learning. There many potential socictal consequences of our work, none of which we feel must be specifically highlighted here.
30

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: For each contribution listed in abstract and introduction we provide the links the sections about them.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of our solver in 6
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide the proofs and the assumptions in Appendix A
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide a full list of the experimental details in Appendix B available datasets. The code for the experiments is provided supplementary material.
We use publicly
Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The code for our solver is included in supplementary material and will be made public after acceptance of paper. We publicly available datasets.
Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We provide the experimental details in Appendix
7. Experiment Statistical Significance
Question: Does the paper report етror bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The focus of the papers is mostly theoretical. The experiments are provided just for illustration of the derived theoretical results.
Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
Answer: [Yes]
Justification: The computational resources discussed in Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The research conforms with NeurlPS Code of Ethics. We discuss the societal impact of our paper in S
31

---

10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: We discuss the societal impact of our paper in 6
11. Safeguards
Does the paper describe safeguards that have been put place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
The research does not release data or models that have a high risk for misuse and does not need safeguards.
Justification:
12. Licenses for existing assets
Question: Are creators or original owners assets (e.g., code, data, models), used in paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cited the creators all of the used assets.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code is included in supplementary material. The license for the code will be provided after the paper acceptance.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not include crowdsourcing experiments or research with human subjects.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Sub- jects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not include crowdsourcing experiments or research with human subjects.
32


Reviews:
[{"review_id": "gQFFbXQHe5", "reviewer": "Reviewer_nig3", "summary": "This paper presents a lightweight solver for Unbalanced Entropic Optimal Transport (UEOT) that does not rely on neural network parametrization. Instead, the authors parameterize the potential functions of UEOT using Gaussian Mixture Models (GMM). This parametrization enables the derivation of a tractable joint coupling. By incorporating the parameterized potential into the dual objective, the authors achieve a simple and tractable loss function. Additionally, the paper provides a universal approximation result for GMM parametrization. Experiments are conducted on toy data (GMM) and image-to-image (I2I) translation.", "questions": "- How does the performance change when parameterizing very high-dimensional and multi-modal GMM data with fewer or more $N,M$? \n- In the I2I experiments, how does the number of modes in the GMM affect performance?\n- In toy data experiments, does U-LightOT has lower transport plan costs and smaller Wasserstein distances between the target and generated distributions compared to other comparisons?", "limitations": "Discussed in Weakness section.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper proposes a simple and fast UEOT algorithm.\n- The paper justifies the GMM parametrization by presenting generalization bounds.\n- The paper demonstrates applicability to large-scale tasks such as I2I translation when combined with an autoencoder (AE).\n- The paper is well-written, clear, and easy to follow.", "weaknesses": "- The method of parameterizing the potential function using GMMs was already proposed in LightSB [1]. The only change here is the switch to a UOT objective, making the methodological contribution minimal. Aside from the universal approximation result, the theoretical contributions are also limited.\n\n- The experiments are not comprehensive. First, the experiments are conducted only on face-related data. More diverse datasets should be included. Second, the fairness of the comparisons is questionable. In the I2I experiments, the authors use the ALAE autoencoder, while some of comparison methods are implemented directly in the image space. All of the comparisons should be implemented in the latent space for fairness. Third, since U-LightOT are implemented on a latent space that captures attributes well, the attribute accuracy is expected to be high. Other than accuracy, more general metrics such as c-FID or FID should be used for comparison. Fourth, there is a lack of ablation studies on the number of Gaussian modals $N$ and $M$. This is very important and expected to be sensitive hyperparameter, thus, I believe authors should provide ablation studies on this parameter. Overall, the practical utility of the approach is questionable.\n\n[1] Light Schrodinger Bridge, ICLR, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough feedback. Please find the answers to your questions below.\n\n**(1) The method of parameterizing the potential function using GMMs was already proposed in LightSB [1]. The only change here is the switch to a UOT objective, making the methodological contribution minimal.**\n\nOur solver can be considered as the generalization of the one from [LightSB, 1] in the sense that it subsumes LightSB for the specific choice of $f$-divergences. However, this generalization is not straightforward or direct since our objective is built on the completely different principles:\n\n1. Our solver is derived from minimizing $D_{\\text{KL}}$ divergence (defined as a discrepancy between *positive measures*) between ground truth plan $\\gamma^*$ and its approximation $\\gamma_{\\theta,\\omega}$. This definition of divergence notably differs from the ordinary definition of $D_{\\text{KL}}$ for **probability measures** used in [1].\n\n2. We parametrize the entire plan using Gaussian mixtures while in [1] it is done only for conditional plans. It is an important difference, since in an unbalanced case the marginals of optimal plan do not coincide with source and target measures. Our parametrization allows sampling from the left marginal of UEOT plan and identifying potential outliers in the source measure.\n\n**(2) Aside from the universal approximation result, the theoretical contributions are also limited.**\n\nWe partially agree with the reviewer that the proof of our Universal Approximation Theorem (UAT) is the most difficult and tricky among the results obtained in our paper. However, the theoretical contributions of our paper are not limited to this theorem. Our other results include: \n(1) Theorem 4.1 $-$ the derivation of the tractable optimization objective in terms of the $D_{\\text{KL}}$ divergence between positive measures; (2) Proposition 4.2 $-$ the derivation of the bound for the estimation error of our solver; (3) Theorem A.4 $-$ derivation of the dual form of UEOT problem with the potentials belonging to the space $C_{2,b}(x)$ of continuous functions, bounded by the quadratic polynom (from the both sides) and additionally bounded by constant from above. Proof of each of these results is non-trivial and requires highly specialized knowledge in the diverse fields of mathematics and statistics.\n\n**(3) In the I2I experiments, the authors use the ALAE autoencoder, while some of comparison methods are implemented directly in the image space. All of the comparisons should be implemented in the latent space for fairness.**\n\nAll of the methods included in comparison in the image-to-image translation task were *implemented in the latent space* of the ALAE autoencoder which is mentioned in the paper, see Section 5.2, line 276. We agree that it might written more clearly and will additionally emphasize this aspect in the final version of our paper. \n\n**(4) Since U-LightOT are implemented on a latent space that captures attributes well, the attribute accuracy is expected to be high. Other than accuracy, more general metrics such as c-FID or FID should be used for comparison.**\n\nWe conduct the image translation experiment in the latent space of ALAE autoencoder. For this reason, we did not report the FID metrics assessing the quality of the generated images but rather focus on assessing the quality of the generated latents and focus on the Frechet distance (FD) defined as the difference in means and covariances of the generated and target latents.\n\nHowever, to fully address the raised question, we report the FID scores between the generated images (produced by ALAE decoder from the generated latent codes) and target images distributions in the Table below (Adult → Young translation). The results show that FID is nearly the same for all of the models under consideration. It supports our intuition that FID is indeed not a representative metric for assessing the performance of models performing the translation of latent codes.\n\n| 10k updates | ULight-OT (ours) | Light-SB [1]      | UOT-FM [2]           | Yang et al. [3]      | Choi et al. [4]   |\n|-------------|------------------|-------------------|----------------------|----------------------|-------------------|\n| FID         | $0.331 \\pm 0.03$ | $0.331 \\pm 0.03$  | $0.331 \\pm 0.04 $  | $0.344 \\pm 0.04 $  | $0.339 \\pm 0.03 $ |\n\n**(5a) Lack of ablation studies on the number of Gaussian modes N and M.***\n\nTo address the reviewer's concern, we perform additional experiments with our solver varying the number of Gaussian modes ($N$, $M$) in potentials. \n\n*Gaussians mixtures.* The setup of this experiment follows the setup introduced in our Section 5.1. We test our solver with diverse number of potentials $N\\in\\{1,3,5\\}$ and  $M\\in\\{1,2,3,4,5\\}$. The results are visualized in Fig. 1 of the **attached PDF file**. It can be seen that for insufficient number of modes in potentials, the solver exhibits issues with convergence and do not correctly solve the task. \n\n**(5b) How does the performance change when parameterizing very high-dimensional and multi-modal GMM data with fewer or more N and M?**\n\nUnfortunately, to assess the performance of our solver in such an experiment with multi-model GMM data, we need to have some kind of ground-truth solutions. However, for the multi-modal GMM data the solutions *are not available* making it hard to perform such an experiment. Following your comment, we qualitatively demonstrated the performance of our solver with varying number of potential modes $N,M$ for 2-dimensional Gaussian mixtures and quantitavely assess its performance in Image-to-Image translation task, see the answer above.\n\n**Concluding remarks**. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XY1hla21br", "reviewer": "Reviewer_t9a3", "summary": "The paper presents a novel approach to solving the continuous Unbalanced Entropic Optimal Transport (UEOT) problem. The authors introduce a lightweight, theoretically-justified solver that addresses the challenges of sensitivity to outliers and class imbalance in traditional Entropic Optimal Transport (EOT). The proposed method features a non-minimax optimization objective and employs Gaussian mixture parametrization for UEOT plans, resulting in a fast, simple, and effective solver. The authors provide theoretical guarantees for their solver's performance and apply it to simulated and image data.", "questions": "- The authors state that for OT-FM and UOT-FM in the FFHQ dataset, they use a 2-layer feed-forward network with 512 hidden neurons and ReLU activation. Where does this parameterization come from? It seems to be relatively small for a flow matching architecture on images, and does not seem to be the architecture used in the original papers.\n- Why are FID scores not reported for the image translation tasks?", "limitations": "The authors have considered the limitations and potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and easy to follow\n- The paper describes well related literature and clearly motivates the approach / why there is a need for this solver\n- The paper introduces a novel way to solve UEOT problems using Gaussian mixtures, even if the approach was previously used for balanced EOT problems as mentioned by the authors.\n- The authors thoroughly study generalization bounds.\n- The authors consider a wide range of competing methods.", "weaknesses": "- While the authors provide generalisation bounds, it would be helpful to assess the performance of the method on the UEOT plan between Gaussian distributions, see Janati et al. ,2020\n- As mentioned by the authors, the Gaussian mixture approach is likely to work only in low dimensions. It would be interesting to see when it fails, e.g. using the benchmark above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough feedback. Please find the answers to your questions below.\n\n**(1) Performance of the method on the UEOT plan between Gaussian distributions, see Janati et al.,2020 [5]**\n\nThank you for your suggestion. Unfortunately, a comparison of our method's solutions with the analytical solutions proposed in [5] is not relevant, since this paper consideres a different setup of the UEOT problem. Namely, this paper derives solutions for the UEOT problem (between Guassian measures) with $D_{\\text{KL}}$ as entropy regularization instead of the differential entropy used in our paper. We noted the difference between the problem we are considering and the one considered in [5] in our paper, see lines 91-92 and corresponding footnote.\n\n**(2) The Gaussian mixture approach is likely to work only in low dimensions. It would be interesting to see when it fails, e.g. using the benchmark above.**\n\nAs we explained in the previous answer the benchmark provided in [5] is not relevant for us as it considers another UEOT problem. \n\n**(3) The authors state that for OT-FM and UOT-FM in the FFHQ dataset, they use a 2-layer feed-forward network with 512 hidden neurons and ReLU activation. Where does this parameterization come from? It seems to be relatively small for a flow matching architecture on images, and does not seem to be the architecture used in the original papers.**\n\nIt is important to understand here that we run our experiments in the latent space of the ALAE autoencoder, and not on the images directly. Accordingly, we adapted the architectures of the neural networks used in OT-FM and UOT-FM to work with latent codes. In this case, architectures such as fully connected neural networks are relevant.\n\n**(4) Why are FID scores not reported for the image translation tasks?**\n\nWe conduct the image translation experiment in the latent space of ALAE autoencoder. For this reason, we did not report the FID metrics assessing the quality of the generated images but rather focus on assessing the quality of the generated latents and focus on the Frechet distance (FD) defined as the difference in means and covariances of the generated and target latents.\n\nHowever, to fully address the raised question, we report the FID scores between the generated images (produced by ALAE decoder from the generated latent codes) and target images distributions in the Table below. The results show that FID is nearly the same for all of the models under consideration. It supports our intuition that FID is indeed not a representative metric for assessing the performance of models performing the translation of latent codes.\n\n| 10k updates | ULight-OT (ours) | Light-SB [1]      | UOT-FM [2]           | Yang et al. [3]      | Choi et al. [4]   |\n|-------------|------------------|-------------------|----------------------|----------------------|-------------------|\n| FID         | $0.331 \\pm 0.03$ | $0.331 \\pm 0.03$  | $0.331 \\pm 0.04 $  | $0.344 \\pm 0.04 $  | $0.339 \\pm 0.03 $ |\n\n**Concluding remarks**. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.\n\n**References**\n\n[1] Korotin et al. \"Light schrodinger bridge\", ICLR 2024.\n\n[2] L. Eyring et al. Unbalancedness in neural monge maps improves unpaired domain translation. ICLR, 2024\n\n[3] J. Choi et al. Generative modeling through the semi-dual formulation of unbalanced optimal transport. arXiv preprint arXiv:2305.14777, 2023.\n\n[4] K. D. Yang et al. Scalable unbalanced optimal transport using generative adversarial networks. ICLR, 2018.\n\n[5] H. Janati et al. Entropic optimal transport between unbalanced gaussian measures has a closed form. NeurIPS, 2020.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CpTfGUH3Sm", "reviewer": "Reviewer_vYvs", "summary": "This work focuses on the largely computationally intractable efforts in unbalanced OT dual form where neural networks are used as a proxy (used as potentials) in order to approximate Wasserstein distances. In this work, the authors set out to significantly reduce this optimization procedure by decomposing the join optimal solution into conditionals which allows for both easier inference and a reduction in the number of parameters required. Experimental results are then carried out to show the success of this method beyond the improved efficiency.", "questions": "(1) The reduction you get appears to have (perhaps even superficially) some relationship to the way WAE decomposes the coupling into conditionals. More coincidentally, WAE also uses conditional Gaussians to parametrize the encoder distribution, although for different purposes. Do you have any comments if there is any deeper link here?\n\n(2) Do you have any intuition if one were to use other penalties beyond KL to enforce the mass constraint?", "limitations": "Yes", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "[+] The reduction in efficiency seems to be quite strong and effective way to reduce the overall parameters required to approximate OT distances\n[+] The theoretical results are sound and well motivated.\n[+] A generalization bound is also presented, attesting to the soundness one achieves with this light variation.", "weaknesses": "[-] Appears to be specific only to the case of having KL divergence penalties for the mass constraints.\n[-] Paper can appear a bit difficult and dense to read.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough feedback. Please find the answers to your questions below.\n\n**(1) Appears to be specific only to the case of having $D_{\\text{KL}}$ divergence penalties for the mass constraints. [...] Do you have any intuition if one were to use other penalties beyond $D_{\\text{KL}}$ to enforce the mass constraint?**\n\nOur solver admits different divergences except for the $D_{\\text{KL}}$ one. From the theoretical point of view, we describe the set of admissible divergences in our Appendix C (lines 549-672). Besides, we provide a numerical example illustrating the performance of our solver with $\\mathcal{D}_{\\chi^2}$ divergence, see Fig. 3 and description in lines 559-677.\n\n**(2) Paper can appear a bit difficult and dense to read.**\n\nWe are upset that you found our work difficult to read. We will try to improve this aspect if you indicate in more detail which points were difficult to understand.\n\n**(3) The reduction you get appears to have (perhaps even superficially) some relationship to the way WAE decomposes the coupling into conditionals. More coincidentally, WAE also uses conditional Gaussians to parametrize the encoder distribution, although for different purposes. Do you have any comments if there is any deeper link here?**\n\nThanks for asking. We think that there is no direct link. Indeed, in WAE, the encoder for each input $x$ outputs some Gaussian, while in our case, all the conditional Gaussians (more precisely, Gaussian mixtures) are tight together. This means that given one conditional distribution $\\gamma_{\\theta}(y|x=x_0)$, one can immediately express all the other $\\gamma_{\\theta}(y|x=x_{\\text{other}})$. In fact, the densities of all these conditional distributions are parameterized by a single scalar-valued function $v$; see eq. (8) in our paper. This is achieved because of the properties of the entropic optimal transport solutions which we exploited to construct our algorithm.\n\n**Concluding remarks**. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DYN1zjqAUa", "reviewer": "Reviewer_bWBu", "summary": "The proposal focuses on developing a fast solver for the unbalanced entropy-regularized optimal (EOT) transport between continuous Radon measures. The authors utilize the dual formulation of unbalanced EOT and use the relationship between the optimal potentials (i.e., the dual variables) and the primal transport plan. They then consider a parameterization of the transport plan and plan to minimize the KL divergence between this parameterized transport plan and the optimal one. Given that the optimal transport plan is unknown, the authors first use the relationship between the primal and dual solutions to reparameterize the dual variables and then derive a tight upper bound for the KL between the optimal plan and the parameterized one, which they propose to minimize. To deal with the normalization terms in their upper bound, the authors use a similar framework to that of Gushchin et al. [29] and assume the reparameterized dual variables are unnormalized Gaussian mixtures; this assumption enables analytic solutions to the otherwise difficult-to-calculate terms in the upper bound. Lastly, the authors provide a generalization error bound for their proposed framework. The paper provides two small-scale numerical examples to demonstrate their solver's efficiency: 1) two-dimensional Gaussian mixtures and 2) unpaired-image-to-image translation in the embedding space of an autoencoder, specifically ALAE, on an unbalanced subset of FFHQ dataset for Adult, Young, Man, Woman face.", "questions": "* How does the performance change as a function of $K$ (assuming $L=K$)?", "limitations": "Limitations are provided in the appendix.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "+ The paper is very well written and straightforward to follow.\n+ The clever parameterizations used in this paper (while they appear in some prior work), provide a unique approach for solving the UEOT problem between continuous measures.\n+ The provided generalization error bounds (while straightforward to derive), are important and certainly add value to the paper.\n+ The method is easy to implement and fast to train. Quick convergence on the CPU is a notable achievement unlocked by this work.", "weaknesses": "- One major weakness is that the paper does not discuss how $K$ and $L$, i.e., the number of Gaussians in the mixtures, affect the results. The generalization error bound mentions that $K$ and $L$ will appear as constants in the error bound, but the practical implications of the choice of $K$ and $L$ are missing from the paper. \n\n- Experiments are relatively modest: 2 experiments in low dimensions. It would be beneficial to have an experiment on robustness to outliers, as this is included in the main claim.\n\n- The paper claims to have a fast solver but lacks a detailed speed comparison for either experiment. It would be great to have a wall-clock comparison of competing methods.\n\n- The Gaussian mixture assumption limits the method's applicability to only low-dimensional problems, and it is not clear whether this limitation can be overcome.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough feedback. Please find the answers to your questions below.\n\n**(1) The paper does not discuss how K and L, i.e., the number of Gaussians in the mixtures, affect the results. The generalization error bound mentions that K and L will appear as constants in the error bound, but the practical implications of the choice of K and L are missing from the paper. [...] How does the performance change as a function of \n K(assuming K=L)?**\n\nTo address your question, we perform additional experiments (both on Gaussians and in the latent space of ALAE autoencoder) with our solver varying the number of Gaussian modes ($K$, $L$) in potentials. \n\nThe setup of this experiment follows the setup introduced in our Section 5.1. We test our solver with diverse number of potentials $K\\in\\{1,3,5\\}$ and  $L\\in\\{1,2,3,4,5\\}$. The results are visualized in Fig. 1 of the **attached PDF file**. It can be seen that for insufficient number of modes in potentials, the solver exhibits issues with convergence and do not correctly solve the task. \n\n**(2) It would be beneficial to have an experiment on robustness to outliers, as this is included in the main claim.**\n\nThank you for this valuable suggestion. We conduct the experiment on Gaussian Mixtures with added outliers and visualize the results in Fig. 2 of the **attached PDF file**. The setup of the experiment, in general, follows the *Gaussian mixtures* experiment setup described in section 5.2 of our paper. The difference consists in outliers (small gaussians) added to the input and output measures. The results show that our U-LightOT solver successfully eliminates the outliers and manages to simultaneously handle the class imbalance issue. At the same time, the balanced LightSB [4] solver fails to deal with either of these problems.\n\n**(3) [...] speed comparison for either experiment. It would be great to have a wall-clock comparison of competing methods.**\n\nThank you for the suggested idea. We compared the running time of our algorithm and unbalanced competitors on the image translation task (Adult$\\rightarrow$Young), its setup is described in section 5.2 of our paper. The results for all of the methods (wall-clock times for 10k updates) are presented in the Table below. We omit the results for other variants of translations since they are quite similar to the results in the Table.\n\nAs you can see, our proposed solver outperforms its competitors (unbalanced methods) in terms of convergence time.\n\n|      | ULight-OT | UOT-FM [1] | Yang et al. [2] | Choi et al. [3] |\n|---|---|---|---|---|\n| Time | **02:38**     | *03:21*   | 16:30      | 18:11 |\n\n**(4) The Gaussian mixture assumption limits the method's applicability to only low-dimensional problems, and it is not clear whether this limitation can be overcome.**\n\nIn general, methods based on Gaussian parametrization are usually not appropriate for tasks with complex data, e.g., images. We mention this limitation in our paper  (lines 256-257,702-707). Still, our aim was to develop a lightweight unbalanced solver which can serve as a simple and easy-to-use baseline in moderate-dimensional tasks. As expected, we get this ease in exchange for the rich parametrization required for large-dimensional tasks and vice versa.\n\n**Concluding remarks**. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.\n\n**References.**\n\n[1] L. Eyring et al. Unbalancedness in neural monge maps improves unpaired domain translation. ICLR, 2024\n\n[2] J. Choi et al. Generative modeling through the semi-dual formulation of unbalanced optimal transport. arXiv preprint arXiv:2305.14777, 2023.\n\n[3] K. D. Yang et al. Scalable unbalanced optimal transport using generative adversarial networks. ICLR, 2018.\n\n[4] Korotin et al. \"Light schrodinger bridge\", ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EHd1l8bTzw", "reviewer": "Reviewer_WAcu", "summary": "This paper proposes U-LightOT, a lightweight solver for the Unbalanced Entropic Optimal Transport (UEOT) problem. This method uses a Gaussian Mixture approximation for the potential $v_{\\theta}(y)$ and measure $u_{w}(x)$. This paper proves that under this approximation, the KL divergence to the ground truth UEOT plan has a tractable form. U-LightOT is evaluated on Gaussian Mixture and Unpaired Image-to-Image Translation tasks.", "questions": "1. In the Unpaired Image-to-Image Translation task, Table 2 only presents the accuracy of keeping the attributes of the source images. However, since the goal of this task is semantic translation, the accuracy of the target semantic is also required. For example, in the Young-to-Adult task, the accuracy of whether the generated image is indeed an adult image. Could you provide this target semantic accuracy results?\n2. In the Appendix, Tables 5 and 7 show the Frechet distance (FD) between the learned and target measures. I believe this FD metric evaluates whether the semantic translation is successful, at the marginal level. Generally, increasing τ\\tauτ decreases (improves) FD metrics in Table 5 and 7. Could you clarify how the optimal $\\tau$ is selected? I am curious because when $\\tau$ is overly large, U-LightOT achieves worse accuracy compared to other models in Table 2.\n3. Could you present the FD results (Tables 5 and 7) for the other models?\n4. For the optimal $\\tau=500$ in the Man-to-Woman task, Table 6 shows an accuracy of 83.85 at best, while Table shows an accuracy of 92.85. Could you clarify which result is correct?", "limitations": "The authors addressed the limitations and broader impact of their work.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. This paper provides a theoretical analysis of the generalization bounds and the universal approximation property for the Gaussian mixture parametrization.\n2. The proposed method is a lightweight solver for the UEOT problem, which requires several minutes of CPU training for the experiment in Sec 5.\n3. This paper is easy to follow.", "weaknesses": "1. The optimization objective and Gaussian Mixture approximation in Sec 4 are similar to [1].\n\n2. While this paper provides the universal approximation property for Gaussian Mixure approximation, I have concerns about whether this Gaussian Mixture parametrization can achieve decent results for more complex distributions, such as in generative modeling within the data space on CIFAR-10. \n\n[1] Korotin, Alexander, Nikita Gushchin, and Evgeny Burnaev. \"Light schr\\\" odinger bridge.\" ICLR 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough feedback. Please find the answers to your questions below.\n\n**(1) The optimization objective and Gaussian Mixture approximation in Sec 4 are similar to [1].**\n\nOur solver can be considered as the generalization of the one from [LightSB, 1] in the sense that it subsumes LightSB for the specific choice of $f$-divergences. However, this generalization is not straightforward or direct since our objective is built on the completely different principles:\n\n1. Our solver is derived from minimizing $D_{\\text{KL}}$ divergence (defined as a discrepancy between positive measures) between ground truth plan $\\gamma^*$ and its approximation $\\gamma_{\\theta,\\omega}$. This definition of divergence notably differs from the ordinary definition of $D_{\\text{KL}}$ for probability measures used in [1].\n\n2. We parametrize the entire plan using Gaussian mixtures while in [1] it is done only for conditional plans. It is an important difference, since in an unbalanced case the marginals of optimal plan do not coincide with source and target measures. Our parametrization allows sampling from the left marginal of UEOT plan and identifying potential outliers in the source measure.\n\n**(2) While this paper provides the universal approximation property for Gaussian Mixure approximation, I have concerns about whether this Gaussian Mixture parametrization can achieve decent results for more complex distributions, such as in generative modeling within the data space on CIFAR-10.**\n\nIn general, methods based on Gaussian parametrization are usually not appropriate for tasks with complex data, e.g., images. We mention this limitation in our paper  (lines 256-257,702-707). Still, our aim was to develop a lightweight unbalanced solver which can serve as a simple and easy-to-use baseline in moderate-dimensional tasks. As expected, we get this ease in exchange for the rich parametrization required for large-dimensional tasks and vice versa.\n\n**(3) In the Unpaired Image-to-Image Translation task, Table 2 only presents the accuracy of keeping the attributes of the source images. However, since the goal of this task is semantic translation, the accuracy of the target semantic is also required. For example, in the Young-to-Adult task, the accuracy of whether the generated image is indeed an adult image. Could you provide this target semantic accuracy results?**\n\nWe are working on this and aim to add the results soon.\n\n**(4) In the Appendix, Tables 5 and 7 show the Frechet distance (FD) between the learned and target measures. I believe this FD metric evaluates whether the semantic translation is successful, at the marginal level. Generally, increasing $\\tau$ decreases (improves) FD metrics in Table 5 and 7. Could you clarify how the optimal tau is selected? I am curious because when tau is overly large, U-LightOT achieves worse accuracy compared to other models in Table 2.**\n\nIn Appendix C (Tables 5, 7), we perform an ablation study of our method in order to show that it offers a flexible way to select a domain translation configuration (unbalancedness parameter $\\tau$) that either allows for very good level of preserving the properties of the input objects or generation of a distribution which is a very good approximation of the target distribution. In that section, we highlighted the parameter which is optimal in the sense that it provides the best tradeoff between the closeness of the learned translations and target ones (*Pareto-optimal*), and the ability of the learned latents to keep the features of input latents. However, the final selection of the optimal configuration remains at the discretion of the user. \n\n\n**(5) Could you present the FD results (Tables 5 and 7) for the other models?**\n\nWe will add the results soon.\n\n**(6) For the optimal tau=500 in the Man-to-Woman task, Table 6 shows an accuracy of 83.85 at best, while Table shows an accuracy of 92.85. Could you clarify which result is correct?**\n\nThe Tables provide results for different number of training steps and parameters $\\tau$. Table 2 shows the accuracy for $\\tau=100$ and 5K steps of the algorithm which is specified in Appendix B.3. Table 6 presents the results for only 3K steps and different parameters $\\tau$. (Note that the value for $\\tau=100$ in the Table 6 (88.59 $\\pm$ 0.40) is close to 92.85 - up to the difference in number of training steps.) \n\n**References.**\n\n[1] Korotin, Alexander, Nikita Gushchin, and Evgeny Burnaev. \"Light schr\" odinger bridge.\" ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear reviewers, 

thank you for your thorough and detailed reviews! We are highly inspired by the fact that you agree on the importance of our theoretical results (Reviewer bWBu, vYvs), clarity (Reviewer WAcu, bWBu, t9a3, nig3) of our paper and mark the efficiency of the our solver (Reviewer bWBu, vYvs). We hope that our U-LightOT algorithm would be easy to use in practical applications.

We will incorporate the changes suggested by the reviewers in the final version of our paper. We list the changes below:

(a) **Main text** $-$ addition of the Table with wall-clock comparison of our U-LightOT solver and its competitors (Reviewer ) plus minor requested clarifications,

(b) **Additional experiment in Appendix C** section $-$ ablation study of our solver with different number of gaussian modes in potentials(**Reviewers  nig3, bWBu**),

(c) **Additional experiment in Appendix E** section $-$ *Gaussian Mixtures with outliers* experiment showing the robustness of our solver towards potential outliers (**Reviewer bWBu**).

Please find Figures for experiments requested by the reviewers  nig3, bWBu in the **attached PDF file**.

Please find the answers to your questions below.

Meta-review: The paper introduces a new solver for Unbalanced Entropic Optimal Transport, using a novel parameterization of the dual potential. The reviewers acknowledged that the paper offers theoretical insights into the approximation capabilities of the proposed method and that the authors made a significant effort in the rebuttal to improve some of the numerical experiments. However, all reviewers expressed concerns that the contributions are too incremental, primarily combining existing methodological and theoretical techniques from [39] (and [30], which is the same paper) and relying on the same set of experiments.

The introduction of the paper does not clearly outline the connection to [39] or explain the differences, failing to adequately motivate the relevance of the contribution. As a result, the paper does not clearly convey what makes this generalization to an unbalanced solver both challenging and important. Furthermore, the numerical evaluation of the proposed method’s improvements is mostly qualitative, without new compelling experiments to demonstrate the practical benefits from an applied perspective. Instead of relying on the same set of experiments as in [39], there should have been more innovation in this area, with stronger benchmarking to better highlight the advantages of the method.

For these reasons, I recommend rejecting this paper.
Decision: Accept (poster)