Neural Optimal Transport with General Cost Functionals
Authors: Arip Asadulaev
Alexander Korotin
Vage Egiazarian
Evgeny Burnaev
Venue: Submitted to ICLR 2023 (2023)
Keywords: Optimal Transport
Neural Networks
Generative Modelling
Unpaired Learning

Abstract: We present a novel neural-networks-based algorithm to compute optimal transport (OT) plans and maps for general cost functionals. The algorithm is based on a saddle point reformulation of the OT problem and generalizes prior OT methods for weak and strong cost functionals. As an application, we construct a functional to map data distributions with preserving the class-wise structure of data.

Under review as a conference paper at ICLR 2023
NEURAL OPTIMAL TRANSPORT WITH GENERAL COST FUNCTIONALS
Anonymous authors Paper under double-bling review
ABSTRACT
We present novel neural-networks wwwwnnes-wwwaseeeeeeeee based algorithm to compute optimal transport (OT) plans and maps for general cost functionals. The algorithm is based on a saddle point reformulation of the OT problem and generalizes prior OT methods for weak and strong cost functionals. As an application, we construct functional to map data distributions with preserving the class-wise structure of data.
INTRODUCTION
Optimal transport (OT) is powerful framework to solve mass- mass-moving and generative modeling problems for data distributions. Recent works (Korotin et al., 2022c Rout et al. 2022 Korotin et al. 2021b Fan al. 2021a Daniels et al. 2021) propose scalable neural methods to compute OT plans (or maps). They show that the learned transport plan (or map) can be used directly as the generative model in data synthesis (Rout et al. 2022 and unpaired learning (Korotin et 2022c Rout et al.) 2022 Daniels et al., [021] Gazdieva et al. 2022). Compared to WGANs (Arjovsky et al. (017)) which employ OT cost as the loss for generator (Rout et al.) 2022 $3), these methods provide better flexibility: the properties of the learned model can be controlled by the transport cost function.
Existing neural OT plan (or map) methods consider distance-based cost functions, c.g., weak or strong quadratic costs (Korotin et al. 2022c Fan et al., [021a Gazdieva et al. (022). Such costs
are suitable for the tasks of unpaired image- to to-image style translation (Zhu al. 2017 Figures 1,2) and image restoration (Lugmayr et al. 2020). However, they do not take into account the class-wise structure of data or available side information, e.g., the class labels. As result, such costs are hardly applicable certain tasks such as the dataset transfer where the preservation the class structure is needed (Figure [... We tackle this issue.
Contributions. We propose the extension of neural OT which allows to apply it to previously unreleased problems. For this, we develop a novel neural-netwooks--bare optimal transport plans for general cost functionals ($3). As example, we construct (#14 and test ($6) the functional for mapping data distributions with the preservation the class-wise structure.
algorithm to compute
T
Transport map
Q
2
IP
P
/
Unlabeled
Figure 1: The setup of class-guided dataset transfer. Input P= En ar distributions are mixtures of N classes. The task to learn a transport map T preserving the class. The learner has the access to labeled input data ~ IP and only partially labeled target data Q
target Q = B,
Notation.
The notation of our paper is based on that of PPty & Cuturi 2020 Korotin et al. 2022c).
For compact Hausdorf space S we use P(S) to denote the set of Borel probability distributions on S. We denote the space of continuous R-valued functions on S endowed with the supremum norm by C(S). Its dual space is the space M(S) P(S) of finite signed Borel measures over S. For

---

Under review as a conference paper at ICLR 2023
functional F: (U we use (h) def SuPHEM(S) [ss ( to denote its convex conjugate functional F C(S) + R foo}.
Let X be compact Hausdorf spaces and IP P(2), Q P()). We use IT(PP) C P(X x 3) to denote the subset of probability measures on A' x y which projection onto the first marginal P. We use [I, Q) C II(P) to denote the subset of probability measures (transport plans) on X' x ) with marginals IP, Q. For u,v ((X).C()) we write a & E C(X' x )) to denote the function u @ v (x. y) H u(x) + v(y). For a functional F MM x >)) + R we say that it is separably *-increasing if for all functions u. U E C(2). C()) and any function E C(X' x )) from u O e S e (point-wise) it follows (u + v) S (c). For a measurable map T" x + y. we
denote the associated push-forward operator by Tw. For Q1: Q2 E P())) with y C RD, the (square of) energy distance & (Rizzo & Székely 2016) between them is:
(Q1 = B Y2ll BJyy Y1 -
- -1112 Ygll-
(I)
where Y1 Q1, Y Q1, Q2 Y ~ Q2 independent random vectors. Energy distance !I a particular case of the Maximum Mean Discrepancy (Sejdinovic et al. (013). It equals zero only when Q1 Q2-
2 PRELIMINARIES
In this section, we provide key concepts of the optimal transport theory. Thought the paper, we consider compact X' = V c R and IP, Q E P(x), P()).
Strong OT. For cost function e C(X V), the optimal transport cost between P, Q is
Cost(IP def
I c(x, (d),)) v), TEl((..) xxy
inf
(2)
see (Villani 2008 11) Problem 2 admits a minimizer x.' E II(P, Q) which is called an OT plan (Santambrogio 2013, Theorem 1.4). It may be not unique Peerr al. 2019, Remark 2.3). Intuitively, the cost function c(2, y) measures how hard it is to move a mass piece between points x E X' and y E V. That is, x² shows how to optimally distribute the mass of P to Q, i..., with the minimal effort.
For cost functions c(z, y) = || - yll and c(z, y) = 1 || - 2 the OT cost 2 is called the Wasserstein-1 (WV, and the (square of) Wasserstein-2 (W2) distance, respectively, (Villani 2008
(1) or (Santambrogio) 2015 $ 2).
Weak OT. Consider weak cost function C: X' P()) + R. Its inputs are a point E X and distribution of y E V. The weak OT cost is (Gozlan et al. 2017 Fackhoff-Veraguas et al. 2019)
(Is
def
inf TEIT(P..)
/ C(a,, (+(x)))())
where 7 (-12x) denotes the conditional distribution. Weak formulation 3 reduces to strong formulation (2) when Clz, 14) = Sy c(x, y d)dp()) An other example of a weak cost function is the 7-weak quadratic cost C(z, pi) Jy æ - dll²du(y)
Var(u), where 7 2 0 and Var(u) is the variance of μu, (Korotin et al.) (222c Eq. 5), (Alibert al. 2019 $5.2). (Gozlan & Juillet 2020 $5.2) for details. For this cost, we denote the optimal value of 3 by W
2,7 and call it 7-weak Wasserstein-2. Regularized OT. The expression inside (2) is linear functional. is common to add lower-semi- continuous convex regularizer R MMXX )) R U foo} with weight y > 0:
Cost(P, Q def
( c(z, y) v) R(()) TEIl(P,0)
inf
(4)
Regularized formulation 4 typically provides several advantages over original formulation (2) For example, if (TT)) ) is strictly convex, the expression inside (4) is strictly convex functional in π and yields the unique OT plan x². Besides, regularized OT typically has better sample complexity (Genevay] 20019; Mena & Niles Weed, 2019 Genevay et al. 2019). Common regularizers are the entropic (Cuturi 2013b), quadratic (Essid & Solomon 2018 Lasso (Courty et al. 2016), etc.
2

---

Under review as a conference paper at ICLR 2023
General OT. Let F M(X )) + BRR(H (+00) be convex sewver--em-cooninuooi functional.
Assume that there exists и E [I, Q for which F(T)) < oc. Consider the problem:
Cos def
inf xElI(P,Q)
F(T).
(5)
The problem a generalization strong OT weak OT 3 regularized OT 4) following (Paty & Cuturi (2020), we call problem (5) general OT problem. Surprisingly, regularized OT 4 represents the same problem:
enough to put c(2:, y = 0, 7 and R(x) F(T) obtain (5) from
That is, regularized OT and general OT can be viewed equivalent formulations.
Existence and duality. With mild assumptions on F, problem (5) admits a minimizer x Paty &
Cuturi 2020 Lemma 1). If JF is separately *-increasing, the dual problem is
J ((x)dP(()) + / v(v)d((v) FF -
Cost(P P, Q) sup 0,.0
where optimization is performed over u, E C(X), C(C) which are called potentials, (Paty & Cuturi, 20200 Theorem 2) for details. Problem (6) also admits pair of minimizers u", v". The popular regularized functionals (4) are indeed separately *-increasing, including the OT regularized with entropy (Paty & Cuturi 2020, Example 7) or L (Paty & Cuturi 2020 Example 8). That is,
formulation (6) subsumes many known duality formulas for regularized OT.
ALGORITHM FOR LEARNING OPTIMAL TRANSPORT PLANS
In this section, we derive an algorithm to solve general OT problem with neural networks. We prove that (5) can be reformulated as saddle point optimization problem ( ($3.1) from the solution of which one may implicitly recover the OT plan (83.2). We give the proots in Appendix A
3.1 MAXIMIN REFORMULATION OF THE DUAL PROBLEM
In this subsection, we derive the dual form, which is alternative to (6 and can be used to get the OT plan x Our two following theorems constitute the main theoretical idea of our approach.
Theorem (Maximin reformulation of the dual problem). For *-separately increasing convex and lower-semi-comtinuout functional F :M(X' x y) + R U {+00} it holds
{ [F((T) - / oly)dd(w)]+ / olo)dQ(i) (7)
(Ps Q) suup inf [(v, =(upp)=sup inf =ETI(?)
v =ETT(F)
where the sup is taken over E()) and (y) is the marginal distribution over of the plan æ.
From 7 we also see that it is enough to consider values of F in E II(P) C MM(X' x x)). For convention, further derivations we always consider F(x) = for E M(X' x 3) ( \I((P..
Theorem (Optimal saddle points provide optimal plans). Let v* arg sup, inf ee(()) C(U, σ) be
any optimal potential. Then for every optimal transport plan x* II(P, Q) it holds:
π E arg inf L(v", *ETT(P)
(8)
If F is strictly convex in E I[(P), then L(v", T) strictly convex a functional of Consequently. it has a unique minimizer. As result, expression (8) is an equality. We have the following corollary.
Corollary (Every optimal saddle point provides the optimal transport plan). Assume additionally that F is strictly convex. Then the unique OT plan satisfies x arg arginfren(P)) L(v", .T.
Thanks to our results above, one may solve 7 obtain the OT plan π' from the solution (v",*") of the saddle point problem (7) We propose an algorithm to do this in the next subsections.
3.2 REPLACING MEASURES WITH STOCHASTIC MAPS
Formulation 7 requires optimization probability measures I[P). To make it practically feasible, we reformulate as the optimization over functions which generate these measures π.
We introduce latent space Z 7 and an atomless measure S P(Z) on it. e.g., S =O Iz). For every π E P(X x }), there exists a measurable function T, : X x + y which implicitly represents it Korotin et al. 2022c 64.1). Such Tz satisfies T (I, : m (r) for all
3

---

Under review as a conference paper at ICLR 2023
Algorithm I: Neural optimal transport (NOT) for general cost functionals
Input :distributions Q, accessible by samples;
mapping network To IR x R$ + R2;; potential network IR R: number of inner iterations Ky; empirical estimator F(x, ((X, Z)) for cost F(T):
Dutput:learnee stochastic OT map To representing an OT plan between distributions IP, Q: repeat
Sample batches Y ~ Q X ~ IP and for each X sample batch Z[x] S;
C.
wo (To(z,))) Y Σ (w):
X
=EZ/u]
yEY
Update w by using UC2.
dw for Ky = 1,2,.... ...., K do Sample batch X ~ IP and for each x E X sample batch Z|x] ~ S;
Cr F(x. (X, 2))
[X
Σ
Σ uw (To 2)):
Update 0 by using ºCr;
DEX
=E2[a]
00
until not converged;
a X. That is, given x X' and random latent vector z S, the function T produces sample (x z) e x (.] In particular, if a IP, the random vector [ar, Ta [x), is distributed as æ. Thus, every π E (I)P) can be implicitly represented as a function T2 : X x Z - y. Note that, in general, there might exist several suitable functions Tx,
Every measurable function T'; X x Z + y is an implicit representation of the measure тт which is the joint distribution random vector [xr, (z, =)] with æ ~ P, S. Consequently, the optimization over π E II( equivalent to the optimization over measurable functions T X' x Z + y. From Theorem | we have the following corollary.
Corollary 2. For s-separately increasing. lower-semiiconntinunn and convex F it holds
Cos Q) sup inf Σ(v, T) =sup inf ((TT) 1 xxx2
where the sup is taken over potentials v E C(V) and inf - over auue T x: + V
Here we identify F(T) del F(Fr) and Č(v,T) T) def [(U,TT)
Following the notation of (Korotin et al. 2022c), we say that functions 7 are stochastic maps (Figure
6). We say that T is stochastic OT map if represents some OT plan π², i.c., (x, -) = π (-1x) holds IP-almost surely for all o Y. From Theorem 2 and Corollary D we obtain the following result.
Corollary 3 (Optimal saddle points provide stochastic OT maps). Let v² E Eargsup, inf7β(rr T) be
any optimal potential. Then for every stochustic OT mup T it holds:
T E arginf inf E(u",
(10)
F is strictly convex in n, we have T E inf Č(u*, T)) + T is a stochastic OT map.
From our results it follows that by solving ( and obtaining an optimal saddle point T"*) one gets stochastic OT map T To guarantee that all the solutions are OT maps, one may consider adding strictly convex regularizers to F with some small weight, e.g., the conditional kernel variance Korotin et al. 2022b). Problem g replaces the constrained optimization over measures π E [I(P) in (7) with optimization over stochastic maps T in O making it practically feasible.
3.3
GENERIC PRACTICAL OPTIMIZATION PROCEDURE
To approach the problem 9 in practice, we use neural nets Tę R. x RR. - R and U. R - R to parametrize T and u.e respectively. We train them with the stochastic gradient descent (SGAD) by using random batches from P. Q S. The optimization procedure is given in Algorithm In the implementation, to update networks Ta and Uws use Adam optimizer (Kingma & Ba 2014
4

---

Under review as a conference paper at ICLR 2023
Algorithm 2: Neural optimal transport (NOT) with the class-guided cost functional Fo.
Input :distributions P = Σn, On n. Q = >n Bn Qn. S accessible by samples (unlabeled); weights on are known and samples from each F,., Qn are accessible (labeled); mapping network To IP" x R R potential network Vi RR + R:
number of inner iterations Kr;
Output : learned stochastic OT map To representing an OT plan between distributions IP, Q:
repeat
Sample (unlabeled) batches Y ~ ~, Q ~ IP and for each x E X sample batch Z[x] ~ S; Σ
C
+ X
Σ
Ew (To(z,2)) Y
Σ (1)); WEY
=EX
zEZ(u]
Update w by using UC.:
for bT = ......., Kr do Pick n f1, ...,2.....N.. N at random with probabilities (01,.... ........0.. aN):
Sample (labeled) batches Xn ~ P Yn ~ Q; for each E X sample batch Zn S;
LT + --EE (Xn, T(n,,
Σ Ex
Σ
ew (To(x,z)))
[xal
[
Update 0 by using &Cr. "9
2EZn(<]
until not converged;
Our Algorithm| | requires an empirical estimator F for F(T)). Providing such an estimator might be
non-trivial for general F. If F(T) = Sx C(z,, (-I1x))dlP i.e., the cost is weak (3). one may use the
following unbiased Monte-Care estimator: F(X, Z)) def XX- Erex Tz., Z[=))),
where C is the respective estimator for the weak cost C' and [[0]] denotes random batch of latent vectors S for given E X. For strong costs and the 7-weak quadratic cost, the estimator δ is given by (Korotin et al. 2022c Eq. 22 and 23) and our Algorithm [for general OT
5 reduces to
the NOT algorithm (Korotin et al. 2022c Algorithm 1) for weak (3) or strong 2 OT. Unlike the predecessor, our algorithm is suitable for general OT formulation (5) In the next section, we propose
cost functional FG (and provide an estimator for it) to solve the class-guides dataset transfer task.
CLASS-GUIDED DATASET TRANSFER WITH NEURAL OPTIMAL TRANSPORT
In this section, we show that general cost functionals (5) are useful, for example, for the class-guided dataset transfer. To begin with, we theoretically formalize the problem setup.
Let each input IP and output Q distributions be a mixture of N distributions (classes) (Pa) and (Qn J respectively. That is P E o, and Q >> Br where (n,, 3, 2 0 are the respective weights (class prior probabiliries) satisfying Σ ===1 an = 1 and Σx= 3, = 1. In this general setup, aim to find the transport plan л (2. 3) and y E y are the same for as much pairs (xr, y) ~ π as possible. That is, its respective stochastic map T should map each component P., (class) of P to the respective component Qn (class) of Q
IIP @) for which the classes of
X
The task above is related to domain adaptation or transfer learning problems. does not always have a solution with each F,, exactly mapped to Qn due to possible prior/posterior shift (Kouw & Loog 2018). We aim to find a stochastic map T between P and Q satisfying (PP,, x S) x Q. for all N. To solve the above discussed problem, we propose the following functional:
71
Fo(#) = Fo((T) def Σ an (TT Qn).
(11)
where & denotes the energy distance 11 Functional 11] is non-negative and attains zero value when the components of IP are correctly mapped to the respective components of Q (if this is possible).
Theorem 3. Functional Fa(T) is conve in क E [IP))) and s-separably increasing.
In practice, each of the terms 8² (T, 8), Qn) in admits estimation from samples from
'The functional Fo(T) is necessarily strictly convex.
5

---

Under review as a conference paper at ICLR 2023
Proposition 1 (Estimator for ). Let Xn Pn be a batch of Kx samples from class For each I E Xn, let [2] ~ S be a latent batch of size Kz. Consider a batch Yn ~ Qn of size Ky. Then
Ao (nn, T(X Yn) def
ΣΣ NEYn DE IEZn
Σ
ly T(x,
Kx
K2
1
Σ3 E 3 E (T( - T(x' 2')l
(12)
2. (K3 Kx) K2
mE =EZn(z) (ax) Z'EZ2'
is an estimator of E² (T#(P.x8), Qa) up to a constant T T-independent shift.
To estimate Fa(T)) one may separately estimate terms E² (TT((P,)))) Q.) for each n and sum them up with weights Ozg- We only estimate only n-th term with probability on at each iteration.
We highlight the two key details of the estimation of (IT) which are significantly different from estimation weak OT costs (3) appearing in related works (Korotin et al. 2022c; Fan et al. 2202la] Korotin al.) 2021b). First, one has sample not just from the input distribution IP,, but separately from each its component (class) P.... Moreover, one also has to be able to separately sample from the target distribution's ' Q components Qn. This is the part where the guidance (semi-supervision) happens. We note that to estimate costs such as strong weak (3), no target samples from Q are needed at all, i.e., they can be viewed as unsupervised.
In practice, we assume that the learner is given labeled empirical sample from IP for training. In contrast, we assume that the available samples from Q are only partially labeled (with 2 labeled data point per class). That is, we know the class label only for limited amount of data (Figure D In this case, all n cost terms (12 can still be stochastically estimated. These cost terms are used to learn the transport map To in Algorithm [ The remaining (unlabeled) samples will be used when training the potential fu, as labels are not needed to update the potential in (9). We provide the detailed procedure for learning with the functional FG (11) in Algorithm 2
5 RELATED WORK
Neural networks for OT. To the best of our knowledge, our method ($3 is the first to compute OT
plans for general cost functionals (5) Our duality formula (9) subsumes previously known formulas
for weak (3) and strong (2) functionals (Korotin et al.) [002aa] Eq. 7). (Korotin et al.]
[021b]] Eq. 9).
(Rout et al. 2022, Eq. 14), (Fan al., 2021a] Eq. I1), (Henry-Labordery 2019 Eq. II). (Gazdieva
et al. 2022 Eq. 10). For the strong quadratic cost, (Makkuva et al.]| (Korotin et al. 202 Eq.
10) consider analogous to (9) formulations restricted to convex potentials; they use Input Convex
Neural Networks (Amos al. (017) to approximate them. These nets are popular in OT (Korotin
ctt al. 2021c] Mokrov et al.] |021]] Fan et al. 202 Bunne et al. [02T] | Alvarez-Mels et al. 2021]
but OT algorithms based on them outperformed (Korotin al.
202 by the above-mentiones
unrestricted formulations. In (Genevay et al.) 2016 (Seguy et al.,
2017 Daniels et al., 2021 Fan
eet al., 2021b), the authors propose methods for /-divergence regularized functionals
The first
two methods recover biased plans which is notable issue in high dimensions (Korotin et al. 202 84.2). Method (Daniels et al.
202T is computationally heavy due using the Langevin dynamies.
Many approaches in generative learning use OT cost as the loss function to update generative models (WGANs (Arjovsky & Bottou| 2017 Petzka et al. 2017 Liu al., (019)). They are not related to our work as they do not compute OT plans (or maps). Importantly, saddle point problems such as (9)
significantly differ from GANs, see (Gazdieva et al.
2022 56.2).
Dataset transfer and domain adaptation. Deep distance based algorithms GGrettnn et al. 201 2
Long al. 2015 2017) or adversarial algorithms (Ganin & Lempitsky 2015 Long et al. 2018 1
are common solutions to the domain adaptation problem. Using neural networks, these methods align probability distributions while maintaining the discriminativity between classes (Wang & Deng 2018 Mostly they perform domain adaptation for image data at the feature level and are typically not used at the pixel level (data space). Pixel-level adaptation is typically performed by common
unsupervised image-to-ing -o-igge translation techniques such CycleGAN (Zhu et 2017 Hoffman
let al. 2018 Almahairi et al.) 2018) and UNIT (Huang et al. 2018 Liu et al. 2017)
OT and domain adaptation. Discrete OT solvers (EMD Nash
2000
Courty al. 2016).
Sinkhorn (Cuturi. 2013a), etc.) are usually employed to map labeled input samples to the unla-
beled partially labeled target samples. Combinations with neural feature extractors (Courty et al.

---

Under review as a conference paper at ICLR 2023
15
20
-35
-15 -20
20 05
15
-55 -10 -05 00 65 Lo 35 -10 -b 65
-s -Lo -ls ab
Le 15
(a) Input x. r Pn (circles),
(b) Discrete OT with the (c) NOT with the quadratic (d) Class-guided NOT learned cost (OT-SI)
target On (crosses).
cost
[Ours]
Figure 2: The results of mapping two moons using OT with different cost functionals.
2
(a) Input P.
(b) Target Q
(c) Guided NOT (Fo) [Ours]
Figure 3: Illustration of the mapping between two Gaussian mixtures learned by our Algorithm
2017 Redko et al. 2018 Bhushan Damodaran al. 2018 are popular in domain adaptation (Courty ct al. 2016). The available labels can be used to reconstruct the cost function and catch the data structure (Courty et al.) 2016 Stuart & Wolfram/ 2020 Liu et al. 2020 Li et al. 2019). Discrete OT performs a matching between the given empirical samples and does not provide out-of-sample estimates. In constrast, our method generalizes to unseen source data and generates
target data.
Conditional generative models. Conditional models, e.g., GAN (Mizaa Osindero 2014), Adver- sarial Autoencoder (Makhzani al.] 2015), the labels apply conditional generation. They not relevant to our work as we do not use any label information during the inference. Our learned mapping is based only the input content.
6 EXPERIMENTS
In this section, we test NOT with our cost functional FG on toy cases ($6.1) and image data ($6.2). For comparison, we consider NOT with Euclidean costs (Korotin et al.) 2022c Rout al. 2022
Korotin et al. 2b2 Fan et al., 2021a) and image-to-image translation models Isola et al. 2017
Huang et al. 2018 Liu et al. (017). The code written in PyTorch framework and will be made public along with the trained networks. On the image data, our method converges in 5-15 hours on Tesla V100 (16 GB). We give the training details (architectures, pre-processing. etc.) Appendix B Our Algorithm learns stochastic (one-to-many) transport maps 7(x, z). Following (Korotin et al. 2021b 85), we also test deterministic T(x, 2) = T(x). i.c., do not add a random noise & to input. This disables stochasticity and yiclds deterministic (one-to-one) transport maps + T(x). In $ (toy examples), we test only deterministic variant of our method. In 6.2 we test both cases.
6.1
TOY EXAMPLES
The moons. The task is to map two balanced classes of moons (red and green) between and Q (circles and crosses in Figure 2a respectively). The target distribution Q is P rotated by 90 degrees. The number of randomly picked labeled samples in each target moon 10. The maps with learned by
NOT with the quadratic cost (W2, Korotin et al. 2022c Rout et al. 2022 Fan et al. 202 lat) and our
functional Fa are given in Figures 2c and 2d respectively. For completeness, in Figure 2b we show
the matching performed by discrete OT-ST algorithm which learns the transport cost with neural
7

---

Under review as a conference paper at ICLR 2023
0
1
2
4
6
7
8
9
1 1
2 X
3 3
4 4
5 S
6
8
9
0
2
3
4
5
6
7
8
9
0
6
7
8
9
1
2
3
7
8
9
2
7
MUNT
1120
3
8
0
9
0
23
5667
8
9
CAN
/
3
4
6
9
OT00
223955
8
9
0
/
X
5
4
5
6
7
8
9
9
4
4
4
7
10
toursi W2.
2
4
4
4
Fo Fa tours tours!
1
223556787
9
Fo Pa
7
8
2 39 6
677 8
9
4
( MNIST + USPS
(b) MNIST + MNIST-M
Figure 4: The results mapping between two similar images datasets (related domains).
0
3
8
2
9
louns) War M OTDO 680 MUNI Fo K-
4
t
2
3
4
5
6
7
8
9
MUNIT y- x Pr.
$7
3
3
3
0
/
7
3
4
5
6
7
8
9
3
1
4
{
$2
7
8-5
8
5
7
1
6
9
6
2
0
5
K
1
4 1
1
$
3
GAN
9
/
3
7
8
4
3
9
5
O
Q
3
5
7 7 7
8
14
OTDO
1
T
0
O
21
& 3
A
8
3
7
O
6
2
4
2
8
7
1
6
(ouns((oura) WS,,
5
3
3
6
6
1
2
1
8
Fo Fa.
12
&
0
1
2
3
4
5
6
7
8
9
8 2
2
3
0
2
3
4
5
7
8
4
(a) MNIST + KMNIST
(b) FMNIST + MNIST
Figure 5: The results of mapping between two notably different datasets (unrelated domains).
net from known classes' correspondence (Liu al.
2020). The map for W2 does not preserve the
classes (Figure 2c), while our map solves the task (Figure 2d).
The Gaussians. Here both F. Q are balanced mixtures of 16 Gaussians, and each color denotes
unique class. The goal is to map Gaussians in P (Figure 3a) to respective Gaussians in Q (Figure (3b) which have the same color. The result of our method (10 known target labels per class) is given in Figure 3c] It correctly maps the classes. NOT for the quadratic cost is not shown as it results in the identity map (the same image as Figure 3a which completely mistaken in classes.
6.2
IMAGE DATA EXPERIMENTS
Datasets. We use MNIST (LeCun & Cortes 2010 MNIST-M (Ganin & Lempitsky 2015), Fashion- MNIST (Xiao et al. 2017), KMNIST (Clanuwat et al. 2018) datasets as P, Q Each dataset has 10 (balanced) classes and the pre-defined train-test split. We consider two cases. In the related domains case, source and target are close: MNIST + USPS, MNIST + MNIST-M. In the unrelated domains case, they notably differ: MNIST - KMNIST and FashionMNIST - MNIST. In all the cases, we use the default class correspondence between the datasets. For completeness, we provide an example with imbalanced classes and a non-defuult correspondence in Appendices and G respectively.
We compare our method to principal unsupervised translation models. We consider et al.) 2018). MUNIT(Huang al. 2018). We use the official implementations with the hyperparameters from the respective papers. Also we test NOT with Euclidean cost functions: the quadratic cost 1/1 || - y I| (W22 and the 7-weak (one-to-many) quadratic cost (W2,7 7 = to). The mentinee methods are unsupervised, i.c., they do not use the label information. For completeness, add (one-to-one) OTDD flow (Alvarez-Melis & Fusil
Baselines:
(one-to-many) AugCycleGANAAIlne
8

---

Under review as a conference paper at ICLR 2023
Image-to-Image Translation
Flows
Neural Optimal Transport
Datasels (32 x 32)
MUNIT
Aug CyclGAN
OTDD
Wa
W2,v
Ya. no
[Ours]
[Ours]
JG
MNIST + USPS
82,97
96.52
82.50
56.27
34.33
93.42
95.14
MNIST - MNIST
97.95
98.2
38.77
37.0
95.27
94.62
MNIST KMNNST
4,46
6.82
79.20
61.91
FMNIST MNIST
8.93
12.03
10.28
10.96
8.02
82.79
83.22
Table 1: Accuracyt of the maps learned by the translation methods in view.
Datasels (32 32)
MUNIT
Aug CycleGAN
OTDD
W2
W2.7
Fo. no Fa
[Ours]
[Ours]
MNIST + USPS
6.86
22.74
100
4.60
3.05
5.40
2,87
MNIST - M
11.68
26.87
19.43
17.48
18.56
6.67
MNIST KMNIST
100
9.46
17.26
9.69
FMNIST - MNIST
7.91
26.35
>. 100
7.51
7.02
7.14
5.26
Table 2: FID of the samples generated by the translation methods in view.
2021 2020). The method employs gradient flows to perform the transfer preserving the class label. Additionally. show the results of ICNN-based W2 OT method (Makkuva et al. 2019 Korotin
ct al. 202[a) in Appendix FF
Metrics. All the models are fitted on the train parts of datasets; all the provided qualitative and quantitative results are exclusively for test (unseen) data. To evaluate the visual quality, we compute FII (Heusel et al.] (017) of the mapped source test set w.r.t the target test set. To estimate the accuracy of the mapping we pre-train ResNet18 18 (He et al. (2016) classifier on the target dataset. We consider the mapping T' correct, if the predicted label for the mapped sample T(x, z) matches the corresponding label of x.
Results. Qualitative results are show in Figures 4 5 FID, accuracy in Tables 2 and respectively. To keep the figures simple, for all the models (one-to-one, one-to-many), we plot single output per an input. For completeness, we show multiple outputs per an input for our method in Appendix C Note that in our method and OTDD, we use 10 labeled samples per each class in the target distribution. The other methods under consideration do not use the label information.
In the related domains case (Figure 4. GAN-based methods and NOT with our guided cost Fo show high accuracy 2 90%. However, NOT with strong and weak quadratic costs provides low accuracy (35-50%). We presume that this is because for these dataset pairs the ground truth OT map for the (pixel-wise) quadratic cost simply does not preserve the class. This agrees with (Daniels et al. 2021 Figure 3) which test an entropy reguuayeee quadratic cost in similar MNIST USPS setup. For method with guided cost Fa. ablation study on 2 size presented in Appendix DD The OTDD gradient flows method provides reasonable accuracy on ANIIT->USH However, OTDD has much higher FID than the other methods. Visually, the OTDD results are comparable to (Alvarez-Melis & Fusi Figure 3).
1
In the unrelated domains case (Figure which is of our main interest, all the methods except our method do not preserve the class structure since they do not use label information. Consequently, their accuracy is around 10% (random guess). OTDD does not preserve the class structure and generates samples with worse FID. Only NOT with our cost FG preserves the class labels accurately.
Conclusion. Our approach provides FID scores which are better or comparable to principal image-to- image translation methods. In the related domain case (Figure 4 we provide comparable accuracy of class preservation. When the domains are unrelated, we notably outperform existing approaches.
6.3 DISCUSSION
Potential Impact. Our method is a generic tool to learn transport maps between data distributions. In contrast to many other generative models, it allows to control the properties of the learned map via choosing a task-specific cost functional F. Our method could be used for data generation and editing purposes, and analogously to GANs, have promising positive real-world applications, such as digital content creation and artistic expression. At the same time, generative models can also be used for negative data manipulation purposes such as deepfakes. In general, the impact of our work on society depends on the scope of its application and the task at the hand.
Limitations. To apply our method, one has to provide an estimator F(T) for the functional FF which may be non-trivial. Besides, the construction of a cost functional F for a particular downstream task may be not straightfoward. This should be taken into account when using the method in practice. Constructing task functionals J and estimators F is a promising future research avenue.

---

Under review as a conference paper at ICLR 2023
7 REPRODUCIBILITY
To ensure the reproducibility of our experiments, we provide the source code in the supplementary material. For toy experiments $ run twomoons_top -opppp and gaussian_toy. ipynb. teaeetttttanttee fer ipynb and dataset_ aetaase_transser_ no 2. ipynb. The detailed information about the data preprocessing and training hyperparameters is presented in an Appendix B
ipynb
For the dataset transfer experiments 6.2 run dataset_ dataset_trans
REFERENCES
J-J Alibert, Guy Bouchitté, and Thierry Champion. A new class of costs for optimal transport planning. European Journal of Applied Mathematics, 30(6):1229-12633 2019.
Amjad Almahairi, Sai Rajeshwar, Alessandro Sordoni, Philip Bachman, and Aaron Courville. Aug- mented eyclegan: Learning many-to-many mappings from unpaired data. International Confer- ence on Machine Learning, pp. 195-204. PMLR, 2018.
David Alvarez-Melis and Nicolo Fusi. Geometric dataset distances via optimal transport. Advances in Neural Information Processing Systems, 33:21 13:21428-21439, 2020.
David Alvarez-Melis and Nicolò Fusi. Dataset dynamics via gradient flows in probability space. In International Conference on Machine Learning, pp. 219-230. PMLR, 2021.
David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing functionals on the space of probabilities with input convex neural networks. arXiv preprint wxXv:22106.00774 2021.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In Proceedings of the 34th
International Conference Machine Learning-Voluum Volume 70, pp. 146-155. JMLR. org, 2017.
Martin Arjovsky and Léon Bottou. Towards principled methods for training generative adversarial networks. arXiv preprint arXiv:1701.04862 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214-223. PMLR, 2017.
Julio Backhoff- Veraguas, Mathias Beiglböck, and Gudmun Pammer. Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):1-2) 2019.
Backhoff-Veraguas
Bharath Bhushan Damodaran, Benjamin Kellenberger, Rémi Flamary, Devis Tuia, and Nicolas Courty. Deepjdot: Deep joint distribution optimal transport for unsupervised domain adaptation. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 447-463, 2018.
Charlotte Bunne, Laetitia optimal transport modeling of population dynamics, 2021.
Andreas Krause, and Marco Cuturi. Jkonet: Proximal
Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature. arXiv preprint arX1v:1812.01111, 2018.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9):1853-1865 2016.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. In Advances in Neural Information Processing Systems, pp. 3730-3739, 2017.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems 26: 27th Annual Conference on Neural Information Processing Systems 2013. Proceedings of meeting held December 5-8, 2013, Lake Tahoe Nevada, United States, pp. 2292 2300, 2013a. URtnttpp//paoossnnnss..o
4 htt.://e....ott.a..........a.aa.
10

---

Under review as a conference paper at ICLR 2023
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pp. 2292-2300, 2013b.
Grady Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for scage optimal transport. Advances in Neural Information Processing Systems, 34, 2021.
Montacer Essid and Justin Solomon. Quadratically regularized optimal transport on graphs. SIAM Journal on Scientific Computing, 4064):A1961-A1186 2018.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Yongxin Chen, and Haomin Zhou. Scalable computation of monge maps with general costs. arXiv preprint arXiv:2/06.038/1 2021a.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Variational wasserstein gradient flow. arXiv preprint urXiv:2112.02422) 2021b.
Yaroslav Ganin and Victor S. Lempitsky. Unsupervised domain adaptation by backpropagation. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning. ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 1180-1189 J....or 2015. URL http: htttp://proceedingg: mlr.
html
Milena Gazdieva, Litu Rout, Alexander Korotin, Alexander Filippov, and Evgeny Burnaev. Unpaired
image super-resolution with optimal transport maps. arXiv preprint 2rXiv:2202.0//66 2022.
Aude Genevay. retroay--eguulriized optimal transport for machine learning. Phl) thesis, Paris Sciences et Lettres (ComUR), 2019.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. In Advances in neural information processing systems, pp. 3440-3448, 2016.
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of sinkhorn divergences. In The 22nd international conference on artificial intelligence and statistics, pp. 1574-1583. PMLR, 2019.
Nathael Gozlan and Nicolas Juillet. On a mixture of brenier and strassen theorems. Proceedings of the London Mathematical Society, 120(3):334-46. 2020.
Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, 273(11):3327-340, 2017.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander J. Smola. A kernel two-sample test. J. Mach. Learn. Res., 13:723-773, 2012. URL http:
7/dl. //dl.acm.co wtt.aac.oo/ciaaiiiiooo./o
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Pierre Labrr-eae (martingale) optimal transport and anomaly detection with neural networks: A prim algorithm. Available at SSRN 3370910, 2019.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in neural information processing systems, pp. 6626-6637, 2017.
Judy Hoffman, Eric Tzeng, Taesung Park, Jun-Yan Zhu, Phillip Isola, Kate Saenko, Alexei A. Efros, and Trevor Darrell. Cycada: Cycle-consistent adversarial domain adaptation. In Jennifer G.
Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmäässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1994-2003. PMLR, 2018. URL http: 7/proceed ngs. ml prw.ssaaaaa httas///v000//ff/ fman 8a. html
11

---

Under review as a conference paper at ICLR 2023
Xun Huang, Ming- Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image-to-iman image translation. In Proceedings of the European conference on computer vision (ECCV), pp. 172-189, 2018.
Phillip Isola, Jun- Yan Zhu, Tinghui Zhou, and Alexei A Efros. image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 1 1125-1134, 2017.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprini 1412,6690, 2014.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations, 2021a. URL httphhttts////eeeeeeet
wwt.:e/ooooummmmoooooomm EXsa
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. arXiv preprint arXiv:2106.01954 2021b.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021c. URL
net httr//orummmiei 3tFAsEEE Pe
Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. arXiv preprint arXiv:2201.12245 12245, 2022a.
Alexander Korotin Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. arXiv preprint arXiv:2205. 2022b.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. arXiv preprint 12220 2022c.
Wouter M Kouw and Marco Loog. An introduction to domain adaptation and transfer learning. arXiv preprint / /806, 2018.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. MNIST, 2010. URL http: //yann. Lecun. cwm/exdb mn mnist/
Ruilin Li, Xiaojing Ye, Haomin Zhou, and Hongyuan Zha. Learning to match via inverse optimal transport. Journal of machine learning research, 20, 2019.
Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein GAN with quadratic transport cost.
In Proceedings of the IEEE International Conference on Computer Vision, pp. 4832 4841, 2019.
Ming-Yo Yu Liu, Thomas Breuel, and Jan Kautz. Unsupervised image-to-image translation networks. In Advances in neural information processing systems, pp. 700-708, 2017.
Ruishan Liu, Akshay Balsubramani, and James Zou. Learning transport cost from subset corre- spondence. In International Conference on Learning Representations, 2020. URL https //openreview.oo net.
Mingsheng Long, Yue Cao, Jianmin Wang, and Michael I. Jordan. Learning transferable features with deep adaptation networks. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 97-105. JMLR.org, 2015. URL html
http:://rrocedinngg mlr.
Mingsheng Long, Han Zhu, Jianmin Wang, and Michael I. Jordan. Deep transfer learning with joint adaptation networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney. NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 2208-2217. PMLR, 2017. html
BRL.htpp:d / /tte:://rocceeeee ngs wtwppssssw////ooo/op htt.://eeeeoooo........eeeoo//ao
12

---

Under review as a conference paper at ICLR 2023
Mingsheng Long, Zhangjie Cao, Jianmin Wang, and Michael I. Jordan. Conditional ad- versarial domain adaptation. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa- BisacBianchi, and Roman Garnett (eds.), Advances in Neu- ral Information Processing Systems 31: Anmual Conference on Neural Information Pro- cessing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 1647-1657, 2018. URL
neurips wwtpppper/2.///aaa....... Abstract html
Andreas Lugmayr, Martin Danelljan, and Radu Timofte. Ntire 2020 challenge on real-world image super-resolution: Methods and results. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition Workshops, pp. 494 495, 2020.
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoencoders. arXiv preprint arXiv:1511.056442 2015.
Ashok Vardhan Makkuva, Amirhossein Taghvaei. Sewoong Oh, and Jason D Lee. Optimal transport mapping via input convex neural networks. arXiv preprint arXiv: 1908. 10962, 2019.
Gonzalo Mena and Jonathan Niles-Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. Advances in Neural Information Processing Systems, 32, 2019.
Mehdi Mirza and Simon Osindero. Conditional generative adversarial nets. arXiv:1411.1785 2014.
arXiv preprint
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. arXiv preprint arXiv:2106.007.36 2021.
John C Nash. The (dantzig) simplex method for linear programming. Computing in Science & Engineering, 2(1):29-31, 2000.
François- Pierre Paty and Marco Cuturi. Regularized optimal transport is ground cost adversarial. In International Conference on Machine Learning, pp. 7532-7542. PMLR, 2020.
Henning Petzka, Asja Fischer, and Denis Lukovnicov. On the regularization of wasserstein gans. arXiv preprint 17XX::1709.8800 2017.
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport. Foundations and Trends@ in Machine Learning, 1 155-6))335-6677 2019.
levgen Redko, Nicolas Courty, Rémi Flamary, and Devis Tuia. Optimal transport for multi-source domain adaptation under target shift, arXiv preprint weXiv:1803.04890 2018.
Maria Rizzo and Gábor Székely. Energy distance, wiley interdisciplinary reviews: Computational statistics, 8(1):27-38, 2016.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net:
Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234 241. Springer, 2015.
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with transporr maps. In International Conference on Learning Representations, 2022. URL httpps://operreviee neet/Forum2id==SdLL22
transport
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94. 2015.
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv:1711.02283. 2017.
Dino Sejdinovic, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. Equivalence of distance-base and based statistics in hypothesis testing. The Annals of Statistics, pp. 2263- 2291, 2013.
13

---

Under review as a conference paper at ICLR 2023
Andrew M Stuart and Marie- Therese Wolfram. Inverse optimal transport. SIAM Journal on Applied Mathematics, 80(1):599-60 619, 2020.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Mei Wang and Weihong Deng. Deep visual domain adaptation: survey, Neurocomputingg 312: 135-153, 2018. doi: h00f//..eucom.2218.0555 URL neucom 2018 05.083
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 178.077708.07777, 2017.
Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223-2232, 2017.
14

---

Under review as a conference paper at ICLR 2023
A PROOFS
Proof of Theorem
We use the dual form 6 to derive
/ u(x)dP(x)- - F (u
Cost(P, Q) sup sup
vvv)dQ())
(13)
y
/ (v)d((())
sup
sup 1
u(x)dP(x) sup
(u (u@U)dT(EZ))) F(TT)
(14)
Xx)
/ (v)d((1)
sup sup
u(x)dP(2)) + inf
F(R)
(u (u (2) (z, 8n
(15)
12XY
sup
supinf inf F(T)) u(z)d(F -P)
v(y)dT(1)
(v)d(())
(16)
y
sup
sup inf (Ell(P)
F(T) u(z)d(a -))()) (()dd(())
(v)dQ(())
(17)
y y 1 v(v)dQ(v)
sup
sup inf TElI(P)
F(x) ((U)dT())
(v))((())
(18)
F(x)) - S u(v)dT(y)
sup
inf (ElI(P)
(19)
/,, (())d(())
sup F(TT)
v(V) dx (v)
=F(T)) = Cost [P, Q). (20)
(2(x)
In line (13) we group the terms involving the potential u. In line (14), we express the conjugate functional F by using definition. In transition to line (55). we replace inf operator with the equivalent sup, operator with the changed sign. In transition to (16). we put the term fx
under the inf, operator; we use definition (u 3) u(z) +w(y) to split the integral over (x, y) into two separate integrals over xt (xz and (ty) respectively. In transition [17), we restrict the inner inf, to probability measures
[I( which have IP as the first marginal, i.e. dr (z) = dP(x). This provides an upper bound on (T6), in particular, all u-dependent terms vanish, see (8]. As a result, we remove the sup,, operator in line (19). In transition to line (20) we substitute an optimal plan
rr" E [I(,, Q) C (1(Q) to upper bound (19). Since Cost(P, Q)) turns to be both an upper bound 200
and lower bound (13) for (19,, conclude that ( holds.
Proof of Theorem
Assume that n² & arrgint xen ,(v",z), i.e.
C(v >²)> inf [(v T) Cost(IP, Q). rEl(e))
We substitute v- and and see that
C(u", = [F(=)) I, v(y) dn (v)]+ v(y)dQ(v) = =(T²)) = Cost(PP Q),
dQli)
which is a contradiction. Thus, the assumption is wrong and g
holds.
Proof of Proposition First, prove that it is *-separately increasing. For n E MMX )) [IP)) it holds that F(T)) +oo Consequently,
f c(x, y)d v) F(T) / (((z))u()))d7((( F(T) -00 (21)
When II(P) it holds that is probability measure. We integrate u(x) + v(y) S c(z, v) w.r. substract F(T) and obtain
f ol, (o, FF()) //
+ ((()) F(x).
(22)
15

---

Under review as a conference paper at ICLR 2023
By taking the sup of (21) and 22 w.I.t MM )), obtain F F (u t
Next, we prove that F is convex. We prove that every term 8² (TT x 8), Qm) is convex in n.
First, show that + (T) def [P, S) is linear in (P).
Pick any Tı, T2, My E [I) which lie on the same line. Without loss of generatity we assume that Ts [T1, Tg]. i.., 7īy aTT (1 a2 for some E [0, 1]. We need to show that
(T() (m) + (1 fa
(23)
In what follows, for random variable U we denote its distribution by Law(U).
The first marginal distribution of each π,; is PP. From the gluing lemma (Villani 2008 51) it follows that there exists a triplet of (dependent) random variables (X, Y1, Y2) such that Law(X Y,) = x; for = 2, We define Y3 Y,, where r an independent random variable which takes values in {1,2] with probabilities {a, 1 a). From the construction of Y3 it follows that Law(X, Y2) is mixture of Law(X, = x and Law(X, Y2) T2 with weights and a. Thus, Law (X, Ys) = a + (1 TT = Tg. We conclude that Law (Ys) x=x) TT(-(a)) for P-almost all z E X' (recall
that (law(X)==)) On the other hand, again by the construction, the conditional Law(Yg) X =x)
mixture of Law(Yi =x) Ti (-1x)) and Law(Y2) X = =) x2(-12) with weights a and 1 Thus, ( = am, (-1ar) (1 (-1x)) holds true for P-almost all X.
Consider independent random variables X, ~ Pn and Z S. From the definition of T. we conclude that Law Law(T,, (zx, Z)) (-1=)>) for P-almost all E X' and, since Pn is a component of P., for Pllmmmosst all x E X' as well. As result, we define T; = Tn, (Xn, Z) and derive
Law X,, = x) = Ta (-12) am (+xx) + (1 ar ())(((12)))))) olaw T1 X,, = x)) + - aa) X,, =x)
for Pll-a all x E X. Thus, (X,, T3) is also a mixture of Law(Xn. .x,,Th Ti) and (X,, Te) with weights a and a. In particular, Law (T3) = oLaw(Ti) + (1 LT)). We note that Law(T,) = (TT) by the definition of fn and obtain (23)
Second, we highlight that for every P(V). the functional P(Y) 3 μ + E2 v)) is convex in μ.. Indeed, _ is particular case of (the square of) Maximum Mean Discrepancy (MMD, (Sejdinovic al. 2013). Therefore, there exists Hilbert space H and function Φ y - H (feature map), such that
111, ) ((v)du(v) J. (g) /2
Ed =
Since the kernel mean embedding az - Jy Φvv) du(y) is linear in μ and ||. 1-112 is convex, we conclude that 82 (n, v) is convex in J.. To finish the proof, it remains to combine the fact that + TTP x x8) linear and (2((, Qn.) is convex in the first argument.
Proof of Proposition
Direct calculation of the expectation of 12 yields the value
EJ - T(X,2)|
T(X, 2) T(X', 2(Xl
E Y - 2(X,
- T(X') 2'1l EI -
1 F Y'I + 2 F -Y
1
82 8 Qg.) + B
YYIl,
(24)
where Y, V Qn and (X, (X', (P,, x 8) are independent random variables. remains to note that 2 Y Y' is 7-independent constant.
"The proof is generic and works for any functional which equals too outside P(X y).
16

---

Under review as a conference paper at ICLR 2023
B EXPERIMENTS DETAILS
DATA PREPOSSESSING. We rescale the images to size 32x32 and normalize their channels to [-1, 1]. For the grayscale images, we repeat their channel times and work with 3-channel images. We do not apply any augmentations to data. We use the default train-test splits for all the datasets.
TRAINING DETAILS. In our Algorithm Ir = 10 for both Te and Vw- The number of inner iterations for To KT 10. Doing preliminary experiments. we noted that it is sufficient to use small mini-bater sizes Ky, K2 in (12). Therefore, we decided to average loss values over K small independent mini-batches (each from class n with probability an) rather than use a single large batch from one class. This is done parallel with tensor operations.
Adam (Kingma & Ba
2014) optimizer with
Two moons. We use 500 train and 150 test samples for each moon. We use the fully-connected net with 2 ReLU hidden layers size of 128 for both Ta and t.- We train the model for 10k iterations of % with Kg = 32, Kx Ky = (K2 plays no role as we do not use z here).
Gaussians. We use the fully-connected network with 2 ReI.U hidden layers size of 256 for both Te and U.e There are 10000 train and 500 test samples in each Gaussian. We train the model for 10k iterations of tw with KB 32, Kx Ky = 2 (Kz plays role here well).
Images. use WGAN-QC discriminator's ResNet architecture (He et al. (016)) for potential oa.
We use UNe (Ronneberger al. 2015) as the stochastic transport map To 2). To condition it on 2. we insert conditional instance normalization (CondIN) layers after each UNet's upscaling block(] We use CondIN from AugCycleGAN (Almahairi al. 2018). In experiments, is the 128-dimensional standard Gaussian noise.
The batch size is Kg 32, Kx Ky 2, K2 for training with When training without 2, we use the original UNet without conditioning; the batch parameters are the same (Kz does not matter). Our method converges in x 60k iterations of ow.
For comparison in image domain we use the official implementations with the hyperparameters from the respective papers: AugCycleGANDI (Almahairi et al. 2018). MUNIT Huang et al. 2018). For comparison with NOT 2W2, W2,x), we use the code shared by the authors of (Korotin et al. (022c)
OTDD flow details. As in our method, the number of labeled samples in each class is 10. We learn the OTDD flow between the labeled source datase[" and labeled target samples. Note the OTDD method does not use the unlabeled target samples. As the OTDD method does not produce out-of-samppe estimates, we train UNet to map the source data to the data produced by the OTDD flow via regression, Then we compute the metries on test (FID, accuracy) for this mapping network.
git BiLhhb/.cc//agggerr/////eetttttr
gi gtt.bb..c//aalllllaaaaoeeeeecoeo
github.ccm/Nv ccwm/N/1abs/MMNN11
We use only 15k source samples since OTDD is computationally heavy (the authors 2k samples).
17

---

Under review as a conference paper at ICLR 2023
C
ADDITIONAL EXAMPLES OF STOCHASTIC MAPS
In this subsection, we provide additional examples the learned stochastic map for Fa (with 2).
T(x1. Z)
X1
n(yl),)
n(())2)
P X2
m(((x,)
X
T(xm. z)
S
Noise
Figure 6: Implicit representation of π [IFP) via function T Tx Xx + V.
We consider all the image datasets from the main experiments (66). The results are shown in Figure and demonstrate that for fixed a and different 2, our model generates diverse samples.
Y-G 10.20 TO6,5) 19620 11x20 Hlc.0l X-P
0
2
3
4
$
6
7
8
9
O
I
2
3
4
5
6
7
8
9
01
2
3
4
56
7
9
2
3
5
7
9
0
0
1
2
3
5
0
3
8
9
01231
4 567
8
9
3
7
3
6
6
7
9
O
2
3 4
66
7
9
2
6
3
4
56
7
8
9
0
2
3
7
7
8
3
7
8
5
8
(a) MNIST to USPS
(b) MNIST MNIST-M
2
3
4
6
1
8
9
3
0
1
2
3
4
5
6
7
8
9
2
0
1
2
3
4
5
6
7
8
9
8
1
2
5
4
5
6
7
8
4
0
1
2
3
4
5
6
7
8
9
0
I
2
3
4
6
6
7
8
4
3
V-a
1
2
3
4
5
6
7
8
9
MNIST KMNIST
(d) PMNIST MNIST
Figure 7: Stochastic transport maps To (x, learned by our Algorithm 2
Additional examples.
18

---

Under review as a conference paper at ICLR 2023
D ABLATION STUDY OF THE LATENT SPACE DIMENSION
In this subsection, we study the structure of the learned stochastic map for Fa with different latent space dimensions Z. We consider MNIST USPS transfer task (10 classes). The results are shown As it can be seen, our model performs comparably for different Z.
in Figures 8 9 and Table 3
TIX TXX. TO.. TOT. TOTOO TO. 25)TH Tw..TH.. (5w... X-P
3
4
5
3
X
5
3
4
5
3
4
5
3
4
5
3
4
5
34
3
3
15
3
4
5
3
4
5
3
4
5
3
4
5
3
Y
5
1
4
5
3
4
5
3
4
5
3
4
6
3
4
5
7 1 7 4 44
3
3
4
5
3
4
5
3
5
3
1 4
6
3
7
5
3
5 5 5 5 5 5
3
4
5
3
9
5
3
4
5
3
6
8
4
5
3
3
4
5
3
4
5
3
9
5
3
4
8 5
3
4
5
3
5
3
4
5
3
4
5
3
4 4
3
4
5
3
5
3
4
5
3
4
5
3
5
3
4 1
5
3
4
5
3
C
5
3
5
3
1
5
3
5
3 4
11
(b) Z
(c)
Z
(d) 2 = 16
(e) 2 = 32
(1) 2 = 64
Figure 8: MNIST + USPS translation with functional FG and varying 2 = 1, 4,8, 32, 64.
X-Pr
1
2
3
4
5
6
7
8
9
O
2
3
4
5
6
7
8
9
2-8 z=4
1022
2
3
4
5
6
7
8
9
2
3
4
S
7
8
8
2- Z=32 2-16
2
3
4
5
7
8
9
2
3
4
5
6
7
8
9
2
3
4
5
6
7
8
9
Figure 9: Stochastic transport maps Te (z, 2)
learned by our Algorithm
2
with different sizes of Z.
Metrics
==1
Z 4
Z = 8
== 16
Z 32
Z 64
Accuracy
86.96
93.48
91.82
92.08
92.25
92.95
FID
4.90
5.88
4.63
3.80
4.34
4.61
Table 3: Accuracyt and FIDJ, of the stochastic maps MNIST + USPS learned by the our translation method with different noise size Z.
19

---

Under review as a conference paper at ICLR 2023
E IMBALANCED CLASSES
In this subsection, we study the behaviour of the optimal map for Fa when the classes are imbalanced in input and target domains. Since out method learns transport map from IP to Q, it should capture the class balance of the Q regardless of the class balance in F. We check this below.
classes in MNIST and n 3 classes in USPS. = 2 a3 and 31 = .2 = 333 3 That is, there no class 3 in the source dataset and not used anywhere during training. In turn, the target class 3 is not used when training To but is used when training fw. All the hyperparameters here are the same as in the previous MNIST + USPS experiments with 10 known labels target classes. The results are shown in Figure 10a and 11a We present deterministic (no 2) and stochastic (with 2) maps.
We consider MNIST USPS datasets with We assume that the class probabilities are
Our cost functional Fa stimulates the map to maximally preserve the input class. However, to transport Q model must change the class balance. We show the confusion matrix learned maps Ta in Figures [ [Tb] It illustrates that model maximally preserves the input classes 0, 1 and uniformly distributes the input classes O and into the class 2, as suggested by our cost functional.
n=
o= 1 1 n= 1
n=
1
3.1
1.6
y- On T(x) X
2
0.33
3.2
1.7
(a) Examples of transported digits x ++ (ar).
(b) Confusion matrix of (x
Figure 10: Imbalanced MNIST USPS translation with functional FG (deterministic, 2).
1
TIA.Al T(K, 23) TOL, 22) TUX,2)) Tix.2ol
2
2
2
2
3
1.7
7
0.16
3.3
1.9
(a) Examples transported digits x -> Ta(r, 2).
(b) Confusion matrix Ta 2).
Figure 11: Imbalanced MNIST + USPS translation with functional Fa (stochastic, with 2).
20

---

Under review as a conference paper at ICLR 2023
F ICNN-BASED DATASET TRANSFER
For completeness, we show the performance of ICNN-based method for the strong quadratic transport cost c(z, y) || - y| on the dataset transfer task. We use the non-minimax version (Korotin et al. 2021a a of the ICNN-based method by (Makkuva et al. 2019). We employ the publicly available code and dense ICNN architectures from the Wasserstein-2 benchmark repository The batch size is Kg = 32, total number of iterations is 100k, Ir = 3 10 and the Adam optimizer is used. The datasets are preprocessed as in the other experiments, see Appendix B
transfer are given in Figure 12 The results are reasonable in the first case (related domains). However, they are visually unpleasant in the second case (unrelated domains). This is expected as the second case is notably harder. More generally, as derived in the Wasserstein-2 benchmark (Korotin et al. 2021b), the ICNN models do not work well in the pixel space due to the poor expressiveness of ICNN architectures. The ICNN method transfer, and 35.6% and accuracy
The qualitative results for MNIST-+USPP and Fashion
achieved 18.8% accuracy and > 100 FID in the FMNIST-MMNI and 13.9 FID in the MNIST-USSE case. All the metrics are much worse than those achieved by general OT method with the class-guided functional Fa, see Table 2 for comparison.
2
3
4
5
6
7
2
1
2
3
4
5
6
7
8
9
ICNN y-
0
2
3
5
7
8
9
0
1
3
4
5
6
7
8
9
013
2
3
4
$
6
7
1
9
IONN
(a) -UIST transfer.
(b) FMNIST-MNIST transfer.
Figure 12: Results of ICNN-based method applied to the dataset transfer task.
G NON-DEFAULT CLASS CORRESPONDENCE
To show that method can work with any arbitrary correspondence between datasets, also consider FMNIST-MMNIT dataset transfer with the following non-default correspondence between the dataset classes:
0 9, 0, 2 1,3 2,4+ 3,5 4, 5, 6,8 7,9 7,9+ 8.
In this experiment, we use the same architectures and data preprocessing as in dataset transfer tasks; see Appendix [B] We use our Fa I as the cost functional and learn a deterministic transport map T (no =). In this setting, our method produce comparable results to the previously reported in Section accuracy equal to 83.1, and FID 6.69. The qualitative results are given in Figure 13
0
1
2
3
4
5
6
7
8
9
Go
3
5
6
9
0
4
7
8
y
Fo
(ours)
O
/
2
3
4
S
6
7
4
Figure 13: EMNIST mapping with FG no 2 cost, classes are permuted.
github.
21


Reviews:
[{"review_id": "X1LswAhqmpJ", "reviewer": "Reviewer_tFkj", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This present a novel neural-networks-based algorithm to compute optimal transport (OT) plans and maps for general cost functionals. \nThe algorithm generalizes prior OT methods for weak and strong cost functionals. \nAnd it constructs a functional to map data distributions with preserving the class-wise structure of data.", "summary_of_the_review": "Although the overall structure of this paper is clear, it lacks novelty. The experimental results do not show the superiority of the algorithm, and it is difficult to design a specific task cost function, lack of guidance.", "strength_and_weaknesses": "Strength: A functional is constructed to map data distributions with preserving the class-wise structure of data.\n\nWeakness: the general cost functionals lacks designed and selected in detail. \nLearned stochastic OT map via the algorithm that may be fake solution. \nIn Image-to-Image Translation task, Lack of comparison with other neural network-based algorithms for solving OT mappings, such as using Input Convex Neural Networks to estimate OT mapping.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "T-qBJf95wht", "reviewer": "Reviewer_7wv1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper is addressing the problem of finding optimal transport plans for general cost functionals. The algorithms is based on a saddle point reformulation of the problem where the marginal constraint is formulated as maximization. The plan on (x,y) is represented as a stochastic map from (x,z) to  (x,T(z,x)) where z is the latent random variable and T(z,x) represents the conditional distribution p_{y|x}, in a similar style as conditional GANs. The application of the proposed algorithm to a data transfer problem with class preservation is described and the results are illustrated. ", "summary_of_the_review": "I think the paper has enough content to be a good contribution but It should be presented in a way that highlights the main contributions and a more deeper discussion of the related works. ", "strength_and_weaknesses": "Strength: \n- the general ideas considered in the paper are novel and interesting. \n- the application that considered in the paper is interesting. \n- the mathematical results seem correct\n\nWeakness: \n- the discussion of related works does not go deep and it is merely a list. \n- there are many results presented that makes reading the paper a bit difficult. It might not be necessarily relevant to the main objectives. I think it is better to reduce them and expand the main ones, adding explanations and clarifications. \n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "QXN2SINhQ8f", "reviewer": "Reviewer_MBuL", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors first take a generalization to the strong and weak optimal transport problems by replacing the objective function from an expected cost wrt the transport plan to a general lower-semicontinuous functional. This formulation coincides with the regularized optimal transport, however, it is easier to adjust under different domains as one can directly design the functional rather than the additive regularizer.  It is shown that this problem admits a strong dual which is a max-min optimization problem, resembling the Kantorovich dual. The authors then show how to retrieve the optimal transportation plan from this dual, which coincides with the saddle points. By exploiting this result, they derive an NN-based solution algorithm, and afterward, they show how one can benefit from such a framework in a case where the underlying domain is concerned with having a \"cost\" representation that is beyond distances. They illustrate the performance of these methods over various numerical experiment settings.", "summary_of_the_review": "This is one of the best papers I have reviewed in a long time. From the accuracy of the literature review to the solid theoretical results, useful reformulations, interesting experiments, and honest but very efficient solution algorithms, the paper is written extremely well.", "strength_and_weaknesses": "**Strength:** Firstly, I would like to thank the author for such a clear motivation, literature review, theoretical analysis, numerical experiments, implementation, and discussions. I am **not** an expert in OT, but I used it in my research a couple of times, and as far as I am concerned the paper looks like a great fit for the community.\n\n**Weaknesses:** \n*I list some minor weaknesses. I am going to stay active during the discussion period in case the authors update the document or have questions about my review.*\n\n***Minor Weaknesses/Questions***\n- Section 1 starts very 'quick'. What is \"unpaired restoration and style translation\", for example?\n- Section 2: In the beginning we have $\\mathcal{X} = \\mathcal{Y}$. Is this a typo?\n- For Wasserstein-1 and Wasserstein-2 terminologies in Section 2, the given citations are specifically recent NN-based papers. However, the terminology definition comes from much earlier if I am not wrong.\n- \"General OT\" (5) requires citations. It may look like this is a new definition.\n- Section 3.1: \"which alternative to (6)\" -> typo\n- Theorem 1: \"minimax\" -> maximin ? (The rest says maximin)\n- Question: Is (11) used in any other paper before? If not, Proposition 1 can be a theorem.\n- Section 6: \"as well as the trained networks\" -> \"along with\" is maybe better?\n- Section 3.3 looks shorter than it has to be, in my view. There is a huge relevant literature in Stochastic Programming, but here it is a little hard to relate. The proposed unbiased estimator comes without discussion. \n\n**UPDATE**\nAfter reading (Korotin et al., 2022c) I am less enthusiastic about this work. I like the paper, but in the presence of the other one, I think the setting taken here can be derived relatively easily. I would like to update my score to \"5\" instead of \"8\". ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "D5t-Rsgwpf", "reviewer": "Reviewer_AhFA", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors proposed to estimate optimal transport maps that are parameterized as neural networks.\nThe proposed method is established based on reformulating the Kantorovich dual problem as maxmin problem which optimizes for a stochastic transport map and a potential jointly. Then, the authors also introduce a cost metric that encourages the transport map to preserve the class information, when mapping two datasets with conditional labels. \nThe authors evaluated their proposed method on a few datasets, including toy synthetic data and real-world image datasets. The proposed method is shown to be effective. ", "summary_of_the_review": "Overall this paper well-presented. Although the proposed method seems to be a little bit incremental since it is built upon existing formulation and algorithms. Since the presentation quality is high, I lean towards acceptance but I am still on the fence.", "strength_and_weaknesses": "Pro:\n-This manuscript is pretty well-written and easy to follow. \n-The proposed method is well motivated as the authors comprehensively introduced different OT problem settings step by step.\n-The experimental results are well-organized. The authors provided both qualitative and quantitative results to evaluate their proposed method.\n\nCons and questions:\n-There might be some concerns regarding the novelty. The proposed objective, as well as algorithm 1 are similar to that in one of the cited works [Korotin 2022c]. Although in this paper, the authors computed a class-guided cost.\n-In equation (11) and Proposition 1. It seems that the stochastic map is asked to map one mixture towards another. However, how to define the correspondence of classes between different domains? What will happen if the number of classes is different between two datasets? This may break down to the optimal transport between two mixture distributions. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject