Neural Optimal Transport
Authors: Alexander Korotin
Daniil Selikhanovych
Evgeny Burnaev
Venue: ICLR 2023 notable top 25% (2023)
Keywords: weak optimal transport
neural networks

Abstract: We present a novel neural-networks-based algorithm to compute optimal transport maps and plans for strong and weak transport costs. To justify the usage of neural networks, we prove that they are universal approximators of transport plans between probability distributions. We evaluate the performance of our optimal transport algorithm on toy examples and on the unpaired image-to-image translation.

Published as a conference paper at ICLR 2023
NEURAL OPTIMAL TRANSPORT
Alexander Korotin
Daniil Selikhanovych
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute
Skolkovo Institute of Science and Technology Moscow, Russia
Moscow, Russia
se 1 ikkhaoovycchaani 1 100mmmmllll.... com
a. korotti@skolteechh..
Evgeny Burnaev
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute
Moscow, Russia
burnaev@@kklleech.c
ABSTRACT
We present novel neural-networks-tee based algorithm to compute optimal transport maps and plans for strong and weak transport costs. To justify the usage of neural networks, we prove that they are universal approximators of transport plans between probability distributions. We evaluate the performance of our optimal transport algorithm on toy examples and on the unpaired image-to-ime tonggge translation.
X X-3
x
T(x.2))
72x
T(x,2))
TXx
(a) Celeba (female) + anime, outdoor church.
(b) Handbags -> shoes, stochastic (one-to-many, W2.1).
deterministic (one-to-one, W2).
Figure 1: Unpaired translation with our Neural Optimal Transport (NOT) Algorithm
INTRODUCTION
Solving optimal transport (OT) problems with neural networks has become widespread in machine learning tentatively starting with the introduction of the large-scale OT (Seguy et al. 2017 and Wasserstein GANs (Arjovsky al. 2017). The majority of existing methods compute the OT cost and use it as the loss function to update the generator in generative models (Gulrajani a 2017 Liu al. 2019: Sanjabi et al. 2018 Petzka et al.) 7). Recently, (Rout et al. 2022 Daniels et al. 2021 have demonstrated that the OT plan itself can be used generative model providing comparable performance in practical tasks.
In this paper, focus the methods which compute the OT plan. Most recent methods (Korotin jet al., 2021b, Rout et al.) 2022) consider OT for the quadratic transport cost (the Wasserstein-2 distance, W2) and recover nonstochastic OT plan, i.e., deterministic OT map. In general, it may

---

Published as a conference paper at ICLR 2023
not exist. (Daniels al. 2021 recover the entropy- entropy-regularized stochastic plan, but the procedures for learning the plan and sampling from it are extremely cone-connuning due to using score- based models and the Langevin dynamics (Daniels et al. 2021 56).
Contributions. We propose novel algorithm to compute deterministic and stochastic OT plans with deep neural networks ($4.1 3/4.2). Our algorithm is designed for weak and strong optimal transport costs (#2) and generalizes previously known scalable approaches ($ $4.3). To reinforce the usage of neural nets, we prove that they are universal approximators of transport plans ($4.4). We show that our algorithm can be applied to large-scale computer vision tasks ($5).
Notations. We use X, y; sets of probability distributions on them. We denote the set of probability distributions on X' x y with marginals and Q by [I,, (). For measurable map T x Z y (or T X + }), we denote the associated nurh-forward operator by Ty..
denote Polish spaces and P(x), P()), P(Z) to denote the respective
2
PRELIMINARIES
In this section, provide key concepts of the OT theory (Villani 2008
Santambrogio
2015 Gozlan
et al. 2017 Backhoff- Veraguas et al. 2019) that we use in our paper.
Strong OT formulation. For P E P(X). e P()) and cost function c: X' x y + I., Monge's primal formulation of the OT cost is
((x))
T(xi)
X1.
c(x:
T(x2) Q
Co, P,
Q
def
/ ((x))()))PP)))) (1) T=
inf
P X2
T(Xa)
c(X,.T(Xn))
X
where the minimum is taken over measurable func- tions (transport maps) T X that map P to Q (Figure 2). The optimal T is called the OT map.
Figure 2: Monge's OT formulation.
Note that ( not symmetric and does not allow mass splitting, i.e. for some P. Q P(X), P()). there may be no T satisfying Ty P Q. Thus, (Kantorovitch) 1958) proposed the following relaxation:
Cost(PP,
Q) def
/ c(x, dr (x, inf =EI[(P.)) xxy
where the minimum taken all transport plans (Figure Baj. i.e., distributions X' x ) whose marginals are IP and Q. The optimal n E [I(P., Q) is called the optimal transport plan. If nr is of the form [id, T P E [1( Q) for some 7"* then T minimizes (T). In this case, the plan is called deterministic. Otherwise, it is called stochastic (nondeterministic)
C(x.T(y(x))
c(xiyy))
Y1 Y2
n(y|x,
X1
X1
P
X2
P
n(ylx,)
Ym
X
c(Xn.ym)
X C(xxx)ly)x,)
(a) Strong OT formulation
(b) Weak formulation
Figure 3: Strong (Kantorovich's) and weak (Gozlan et al. 2017) optimal transport fomulations.
example of OT cost for y DRD the (p-th power of) Wasserstein-p distance W,, i.e., formulation 2 with c(x, 3) || y/... Two its most popular cases are p 1,2 (W1, W2).
Weak OT formulation (Gozlan al. 2017 Let C:: X P(Y) + be a weak cost, i.c., function which takes a point x E X and distribution of y E y as input. The weak OT cost between P, Q is
Cost(IP def
inf =EIT(P.0)
/ C(I,
(3)
where <--12) denotes the conditional distribution (Figure 3b). Note that (3) generalization of Indeed, for cost C(., (1) Jy c(x, v)du(y), the weak formulation (3) becomes strong 22) An
2

---

Published as a conference paper at ICLR 2023
example of weak OT cost for X' = R2 is the y-weak (720) Wasserstein-2 (W2,)):
((2, / ----~ - dul Var(u)
(4)
2
Existence and duality. Throughout the paper, we consider weak costs Clz,, j.) which are lower bounded, convex μl and jointly lower semicontinuous in an appropriate sense. Under these assumptions, [Backhoff-Veraguas et al., 2019) prove that the minimizer x' of (3) always exists [] With mild assumptions on e, strong costs satisfy these assumptions. In particular, they are linear w.r.t μu, and, consequently, convex. The 7-weak quadratic cost (4) is lower-bounded (for S 1) and is also convex since the functional Var(μ) is concave in Ja. For the costs in view, the dual form of (3) is
Cost(P, Q) =su J, 1 (x)))((()))))) + + (())
sup
(5)
where f are the upper-bounded continuous functions with not very rapid growth (Backhofi- Veraguas ct al. 2019 Equation 1.2) and f is the weak C-transform of i.c.
inf {cla.n)- - f(u)du(u))
(6)
MEP())
Note that for strong costs C, the infimum is attained at any μ E P(V) supported on the arg inf yeylert(( ((v) set. Therefore, it suffices to use the strong c-transform:
(=) (x)
def
inf (c(z,, y) f())}
yEY
For strong costs formula (5) with 77 is the well known Kantorovich duality (Villani 2008 55).
Nonuniqueness. In general, an OT plan n is not unique, e.g., see (Peyré et al. 2019 Remark 2.3).
RELATED WORK
In lcale machine learning, OT costs are primarily used as the loss to learn generative models. Wasserstein GANs introduced by (Arjovsky al. 2017 Gulrajani al. 2017) are the most popular examples of this approach. We refer to (Korotin et al. 2022b 202 1b) for recent surveys of principles of WGANs. However, these models are out of scope of our paper since they only compute the OT cost but not OT plans or maps ($4.) To compute OT plans (or maps) is a more challenging problem, and only a limited number of scalable methods to solve have been developed.
We overview methods to compute OT plans (or maps) below. We emphasize that existing methods are designed only for strong OT formulation (2) Most of them search for deterministic solution (I. i.e. for map 1" rather than a stochastic plan π. although 7 might not always exist.
To compute the OT plan (map), (Lu et al. 2020 Xie et al. 2019) approach the primal formulation ① or (2). Their methods imply using generative models and yield complex optimization objectives with several adversarial regularizers, e.g., they used to enforce the boundary condition PT Q) As
a result, the methods are hard to setup since they require careful selection of hyperparameters.
In contrast, methods based on the dual formulation (5) have simpler optimization procedures. Most of such methods designed for OT with the quadratic cost, i.c., the Wasserstein-2 distance (W2). An evaluation of these methods is provided in (Korotin et al. (0211). Below we mention their issues.
Korotin et al. 2021a based on input- convex neural networks (ICNNs, see (Amos et al.] (017)) have solid theoretical justification, but do not provide sufficient performance in practical large-scale problems. Methods based on entropy regularized OT (Genevay et al. 2016; Seguy al. 2017 Daniels et al. 2222 recover regularized OT plan that biased from the true one, it is hard to sample from it or compute its density.
Methods by (Taghvaei & Jalali 2019 Makkuva et al. 2020
According to (Krottnn et al. 2021b), the best performing approach is [M::R], which is based on the maximin reformulation of (5) It recovers OT maps fairly well and has a good generative performance. The follow-up papers (Rout et al. 2022 Fan et al. 2022) test extensions of this approach for more
general strong transport costs c(., and apply to compute W2 barycenters (Korotin al. 2022a). Their key limitation is that it aims to recover a deterministic OT map T which might not exist.
Backhofff-Veree Veraguas et al. (211)) work with the subset Pp())) C P(Y)) whose p-th moment is finite. Hence- forth, we also work in P, ())) equipped with the Wasserstein-p topology. Since this detail is not principal for our subsequent analysis, to keep the exposition simple, we still write P(Y) but actually mean Pp ())).
3

---

Published as a conference paper at ICLR 2023
4 ALGORITHM FOR LEARNING OT PLANS
In this section, we develop a novel neural algorithm to recover solution 'T" of OT problem (3. The following lemma will play an important role our derivations.
Lemma (Existence of transport maps.). Let μ and be probability distributions on R&d and R.V. Assume that is atomless. Then there exists a measurable L: M -> +RR satisfying 1 v.
Proof. (Santambrogio 2015 Cor. 1.29) proves the fact for M N. The proof works for M N.L
Throughout the paper we assume that P, Q are supported on subsets X' C .R.. C RR respectively.
4.1 REFORMULATION OF THE DUAL PROBLEM
First, we reformulate the optimization in C-transform (6. For this, we introduce a subset C RRS with atomless distribution S on it, e.g., = Uniform 1)) or N0, 1).
Lemma 2 (Reformulation of the C-transform). The following equality holds:
(()) -nn=-inf/ Caz, ty f (((=))
(8)
where the infunum is taken over all measurable t: y.
Proof. For all x E X and t Z + y we have ((() < Clz, ((1, - d(t(2))dS(2)) The inequality is straightforward: we substitute of variables, Taking the infimum over I, we obtain
4 S to 66 to upper bound (x and use the change
(x) { ((2, (+3)
C
ft((z) (((2)))B((2)
(9)
Now let us turn (9) to an equality. We need to show that VE > there exists t : + y satisfying
(c)++ 2 Clx, (23 / (((())))S(2))
(10)
By (6 and the definition of inf, 3n E P()) such that f0(x) + 2 -(x, - Sy Sy (y)
Thanks to Lemma there exists - y such that μ² = S, i.e., (T)) holds true.
Now use Lemma to get analogous reformulation of the integral of fc in the dual form Lemma 3 (Reformulation of the integrated C-transform). The following equality holds:
J (()))P(()) inf (Cl,.Tee )+S) [ ( (T((z,))) dP(x), ' [
(11)
where the inner minimization is performed over all measurable functions T X x 2 - y
Proof. The lemma follows from the interchange between the infimum and integral provided by the Rockafellar's interchange theorem (Rockafellar 1976) Theorem 3A).
The theorem states that for function F A B R and distribution on A.
[ inf F(a, b)d((a) inf / F(a, (())d(a)) bEB
(12)
H:A-BBA
We apply (2). use A X, = P. and put B to be the space of measurable functions - Y, and F(a,b) b) = C(a, 8) fy f (y) (z). Consequently, we obtain that fxx f (x))d((x)) equals
inf / (((,, #
sw)d((())))))) ((GG-HO))))-(((((())))))))))))
dP(x)
(13)
Finally, we note that the optimization over functions H X (l: Z + }}} equals the optimization over functions : X' x 2 -. We put T(x, z) = (=)), use the change of variables for y T(x, 2) and derive 11 from
Lemma provides the way represent the dual form saddle point optimization problem.
4

---

Published as a conference paper at ICLR 2023
Corollary (Maximin reformulation of the dual problem). The following holds:
Cost(P, Q sup inf ((S,T),
(14)
where the functional L is defined by
T) def J (())dD())) / ((a,,T(
(T((, dP(x).
(15)
Proof. It suffices to substitute
into
We say that functions T X Z
+ y are stochastic maps. If a map I' is independent of 2, i.c., for all 2) E X Z we have T(x, T(), say the map is deterministic.
The idea behind the introduced notation is the fol- lowing. An optimal transport plan π² might be nondeterministic, i.e., there might not exist a de- terministic function T X y which satis- fies nT" [idx, 7] However, each transport plan æ E [I(P, Q) can be represented implicitly through a stochastic function 7' X' x Z - y. This fact is known as noise outsourcing (Kallenberg R'' and
S
Noise
T(xi Z)
T(x,2)
IP X2
T(x,.2)
1997 Theorem 5.10) for Z [0, 1] S Uniform((0 ((0 1)). Combined with Lemma the noise outsourcing also holds for general 2 R and atomless E 2(2). We visualize the idea in Figure 4 For plan æ, there might exist multiple maps 7" which represent it.
X
T(xm Z)
Figure 4: Stochastic function T(x, representing transport plan. The function's input is x E X' and S.
For pair of probability distributions IP, Q, we say that T is stochastic optimal transport map if realizes some optimal transport plan π². Such maps solve the inner problem in (14) for optimal f"
Lemma 4 (Optimal maps solve the maximin problem), For any maximizer f of and for any stochastic map T which realizes some optimal transport plan n", holds that T E arg inf L(F", T))
(16)
T
Proof. Let n-" be the OT plan realized by T We derive
Sx (w))d
f xxx
(1)ddx"'(, (a, -
(17)
where change the variables for y TT, 2) and the property dr di) Now
assume that T & arg infy L(F", In this case, from the definition (1)) we conclude that C(F T > (ost(IP @). However, derive substituting [7 into (15). we see that L(f" T*)
Ix C( T (2, + dP Cost(TP, Q), which is contradiction. Thus, (66) holds true.
Li
(g)m)
da (x))
5 indeed For the 7-weak quadratic cost which use in the experiments (55). maximizer exists, see (Alibert et al. 2019 55.22) or (Gozlan & Juillet 2020). Thanks to our Lemma 4 one may solve the saddle point problem (4) and extract an optimal stochastic transport map 7 from its solution (f*,T*). In general, the arg inf set for f may contain not only the optimal stochastic transport maps but other stochastic functions as well. In Appendix we show that for strictly convex (in μ1) costs C(x, p),, all the solutions of (14) provide stochastic OT maps.
4.2 PRACTICAL OPTIMIZATION PROCEDURE
To approach the problem (14 in practice, we neural networks Te RP x R - R and fu R + R to parameterize and f. respectively. We train their parameters with the stochastic gradient ascent-descent (SGAD) by using random batches from P. Q S, see Algorithm
5

---

Published as a conference paper at ICLR 2023
Algorithm 1: Neural optimal transport (NOT)
Input :dittributions P, Q, S accessible by samples: mapping network To RRP x R$ R
potential network fw IR + R; number of inner iterations Kr; (weak) cost C: A' x P()) - IR; empirical estimator Č(z, T(x, 2)) for the cost;
Output:learne stochastic OT map Ta representing an OT plan between distributions IP, repeat
Sample batches Y Y~Q, X ~ P: for each a E X sample batch 2₂ ~ S;
f [XT E [Z E=ez, [u(To( z)) M Eprr
Update w by using HCy. 8w
for KT = 1,2...., Kr do Sample batch IP; for each x E sample batch Zr S:
CT [X] ErEx (Ō(1, To( 7x)) [z-1 Σ=ez, Sw (z) &Cr: Update 0 by using 00
until not converged;
Our Algorithm requires an empirical estimator Ô for C(I,T(Z: T(x 4S). If the cost is strong, it is straightforward use the following unbiased Monte-Carlo estimator from a random batch Z S:
Cls, ((x,)))
c(x, (1x, =))
Σ cx2 T<x,))) Anf C(z, T(r, 2)).
(18)
Z
=EZ
For general costs C, providing an estimator might be nontrivial. For the 7-weak quadratic cost 4
such an unbiased Monte- Carlo estimator is straightforward derive:
C(z, T(x Z) def
Σ 3EZ
|| T(x, 2)1 7 2
(19)
2|2|
where 8² the (corrected) batch variance 6² = | ΣEEz (TT,, - 12) Erez T(x, 2)1 To
1
estimate strong costs (18), it is enough to sample single noise vector (12) 1). To estimate the y-weak quadratic cost one needs Z since the estimation of the variance 8² is needed.
4.3 RELATION TO PRIOR WORKS
novel approach to learn stochastic OT plans; it is not a GAN or WGAN-based solution endowed with additional losses such as the OT cost. WGANs (Arjovsky et al. 2017) do not learn an OT plan but use the (strong) OT cost as the loss to learn the generator network. Their problem is infz sup, f). The generator T solves the outer infy problem and is the first coordinate of an optimal saddle point (f*) In our algorithm [ problem (15) sup, inf C(f,T), the generator (transport map) T solves of the inner infr problem and is the second coordinate of an optimal saddle point (f*, ) Intuitively, in our case the generator T is adversarial to potential ff (discriminator), not vise-versa in GANs. Theoretically, the problem also significantly different swapping infyr and supy, in general, yields a different problem with different solutions, e.g., 1 inf, sup, sin(x + y) A sup, inf sin(x y) -1. Practically, we do Kr > updates of I" per one step of f. which again differs from common GAN practices, where multiple updates of f are done per step of T. Finally, in contrast to WGANs, we do not need to enforce any constraints on f, e.g., the 1-Lipschitz continuity.
Generative adversarial learning. Our algorithm
Stochastic generator parameterization. We add an additional noise input 2 to transport map T(x, 2) to make it stochastic. This approach is common technical instrument to parameterize one-to-many mappings in generative modeling, see (Almahairi et al.. 2018 $3.1) or (Zhu et al. 2017b] 93). the context of OT, Yanng& & Uhler 2019) employ stochastic generator learn transport plan in the unbalanced OT problem (Chizat 2017). Due to this, their optimization objective slightly resembles ours 115) However, this similanity is deceptive, Appendix G
Dual OT solvers. Our algorithm I recovers stochastic plans for weak costs It subsumes previously known approaches which learn deterministic OT maps for strong costs When the cost is strong (3) and transport map 7 is restricted to be deterministic T(x, 2) = (r), our Algorithm ] yields maximin method [MM:R, which was discussed in (Korotin al. 202 (2) for the quadratic cost || w| and further developed by (Rout et al. 2022) for the Q-embedded cost - -Q((x), y) and by [Fan al. 2022) for other strong costs clæ, y). These works are the most related our study.

---

Published as a conference paper at ICLR 2023
4.4 UNIVERSAL APPROXIMATION WITH NEURAL NETWORKS
In this section, we show that it is possible to approximate transport maps with neural nets.
Theorem 1 (Neural networks are universal approximators of stochastic transport maps). Assume that X, Z are compact and Q has finite second moment. Let TT' be a stochastic map from P to Q (not necessarity optimal). Then for any nonaffine continuous activation function which is continuously differentiable at at least one point (with nonzero derivative at that point) and for any > 0, there exists a neural network To
TR. IR.
RR satisfying and
Toe
T1 S E
(Px))
(20)
where = S, X 2 + RR) the space of quadratically integrable w.r.t. P & functions X' x 2 - R2 That is, the network To generates distribution which is c-close Q in W3
Proof. The squared norm ||T is equal to the second moment of Q since T pushes P S to Q. The distribution Q has finite second moment, and, consequently, T E L² Thanks to (Folland 1999 Proposition 7. 7.9). the continuous functions CO 2 - R) are dense in L According to (Kidger& & Lyons 2020 Theorem 3.2), the neural networks R' x R + ? with the above- mentioned activations
dense in C x Z + RQ) w.r.t. 1 norm and, consequently, W.r.t. 1,2 norm. Combining these results yields that neural nets are dense in I.2, and for every e >> 0 there necessarily exists network T3 satisfying the left inequality in (20). For Te, the right inequality follows from (Korotin et al. 2021a Lemma A.2).
Our Theorem 1 states that neural nets can approximate stochastic maps in L norm. It should be taken into account that such continuous nets To may be highly irregular and hard to learn in practice.
5 EVALUATION
We perform comparison with the weak discrete OT (considered as the ground truth) on toy 2D, 1D distributions in Appendices B C respectively. In this section, we test our algorithm on an unpaired image-to-image translation task. We perform comparison with popular existing translation methods in Appendix D The code is written in PyTorch Tramework and is publicly available at
Mti.nnneeuaalllla i ma 1T Transport.
Image datasets. We use the following publicly available datasets as IP, Q: aligned anime faces") celebrity faces (Liu et al. 2015). shoes (Yu & Grauman| 2014), Amazon handbags, churches from LSUN dataset (Yu al. 2015). outdoor images from the MIT places database (Zhou al.. 2014).
The size of datasets varies from 50K to 500K images.
Train-test split. We pick 90% of each dataset for unpaired training. The rest 10% are considered as the test set. All the results presented here are exclusively for test images, i.e., unseen data.
Transport costs. We experiment with the strong (? = 0) and 7-weak (y > 0) quadratic costs. Testing other costs, e.g., perceptual (Johnson et al. (0166) or semantic (Cherian & Sullivan| 2019), might be interesting practically, but these two quadratic costs already provide promising performance.
The other training details are given in Appendix
5.1 PRELIMINARY EVALUATION
In the preliminary experiments with strong cost (7 = 0), we noted that T(x, 22) becomes independent of z. For fixed potential f and a point x, the map T(x, .) learns to be the map pushing distribution S to some arg inf distribution j of (6). For strong costs, there are suitable degenerate distributions μt,
see the discussion around 7 Thus, for T it becomes unnecessary keep any dependence on z; simply learns a deterministic map T(x, = T(x). We call this behavior a conditional collapse.
Importantly, for the 7-weak cost (7 > 0), noted different behavior: the stochastic map T(x,,) 2) did not collapse conditionally. To explain this, we substitute (4) into (3) to obtain
[
1
f Var(=(ular) di(x))
W (P, Q
inf TEIl(.Q)
yll dr (I. 4)
xxy
(2)
The first term is analogous to the strong cost (W W2,0). while the additional second term stimulates the OT plan to be stochastic, i.c., to have high conditional variance.
"hhe proposition considers scalar-valued functions (Q = 1), but is analogous for vector-valued functions. kaggl e waaggle..cm//eeeeeeeeooe wwt.cc//eiitaaaaaaaaaaa igedan me faces
7

---

Published as a conference paper at ICLR 2023
X-
T(x)
T(x)
(a) Handbags + shoes, 128 x 128.
(b) Shoes + handbags, 128 x 128.
T(x)
T(x)
Celeba (female) anime, 64 x 64.
(d) Anime celeba (female), 64 64.
T(x) 1
T(x)
(c) Celeha (male) + celeba (female), 64 x 64.
(f) Anime -> shoes, 64 x 64.
Figure 5: Unpaired translation with deterministic OT maps (W2).
Taking into account our preliminary findings, perform two types of experiments. In $5.2 learn deterministic (one-to-one) translation maps Txx) for the strong cost (7 0), i.e. do not add z-channel. In $5.3 we learn stochastic cone-to-many) maps T(x 2) for the y-weak cost (? 0). For completeness, in Appendix A we study how varying affects the diversity of samples.
5.2
ONE-TO-ONE TRANSLATION WITH OPTIMAL MAPS
We learn deterministic OT maps between various pairs of datasets. We provide the results in Figures Extra results for all the dataset pairs that we consider are given in Appendix
la and 5
Being optimal, our translation map ((x) tries minimally change the image content in the L pixel space. This results in preserving certain features during translation. In shoes + handbags (Figures
56 pab. the image color and texture of the pushforward samples reflects those of imput samples. In celeba (female) <>> anime (Figures Ta 5c 5d). head forms, hairstyles are mostly similar for input and output images. The hair in anime usually bigger than that in celeba. Thus, when translating celeba (female) +> anime, the anime hair inherits the color from the celebrity image background. In outdoor + churches (Figure lap, the ground and the sky are preserved, celeba (male) - celeba (female) (Figure 5e) the face does not change. We also provide results for translation in the case when the input and output domains are significantly different, see anime + shoes (Figure 5f).
Related work. Existing unpaired translation models, e.g., CycleGAN (Zhu al. 2017a) UNIT (Liu et al. 2017]. typically have complex adversarial optimization objectives endowed with additional losses. These models require simultaneous optimization of several neural networks. Importantly, vanilla CycleGAN searches for a random translation map and is not capable of preserving certain attributes, c.g.. the color, see (Lu al. 2019, Figure 5b). To handle this issue, imposing extra losses is required (Benaim & Wolf [017:] Kim et al. 2017), which further complicates the hyperparameter straightforward objective (14); we use only networks (potential f. map T), see Table 2 for the comparison of hyperparameters. While the majority of
selection. contrast, our approach
existing unpaired translation models are based on GANs, recent work (Su et al. 2023) proposes diffusion model (DDIBs) and relates it to Schrödinger Bridge (Léonard 2014 i.e., entropic OT.
5.3
MANNY TRANSLATION WITH OPTIMAL PLANS
We learn stochastic OT maps between various pairs of datasets for the >-wekk quadratic cost. The parameter 7 equals 3 or 1 in the experiments. We provide the results in Figures 1b [bland In all the
8

---

Published as a conference paper at ICLR 2023
X
T(x, 2,)
T(x, 2z)
T(x,
T(x, 23)
T(x,2g)
(a) Celeba (female) 7 anime, 128 x 128 (W2.g) 3)
(b) Outdoor + church, 128 x 128 (W2.g) z)..
x-p
7zx X-P
Tix.z) T(x.2;) T(x,21) 21) X-P
Tix 211
Tox.z)
Txx,z)) 25x
f
(c) Anime >celeba (f), 64x64 (Wa.g) (d) Shoes + handbags, 64 (W2,1).
(e) Anime shoes, 64 64 (W2,1).
Figure 6: Unpaired translation with stochastic OT maps (W2,)).
cases, the random noise inputs 2 ~ are not synchronized for different inputs æ. The examples with the synchronized noise inputs are given in Appendix Extended results and examples of interpolation in the conditional latent space are given in Appendix H The stochastic map 7(x, =) preserves the attributes of the input image and produces multiple outputs.
Related work. Transforming a one-to-one learning pipeline to one-to-many nontrivial. Simply adding additional noise input leads to conditional collapse (Zhang) (011). This is resolved by AugCycleGAN (Almahairi et al. 2018) and M-UNIT (Huang et al. 2018). but their optimization objectives much more complicated then vanilla versions. Our method optimizes only 2 nets f. T in straightforward objective (14). It offers a single parameter r to control the amount of variability in the learned maps. We refer to Table for the comparison of hyperparameters of the methods.
DISCUSSION
Potential impact. Our method is novel generic tool to align probability distributions with determin- istic and stochastic transport maps. Beside unpaired translation, we expect our approach to be applied other one-to-one and one-to-many unpaired learning tasks well (image restoration, domain adaptation, etc.) and improve existing models in those fields. Compared to the popular models based on GANs (Goodfellow al.) 2014) diffusion models (Ho al. 2020), our method provides better interpretability of the learned map and allows to control the amount of diversity in generated samples (Appendix A). should be taken into account that OT maps we learn might be suitable not for all unpaired tasks. We mark designing task-specifie transport costs as a promising research direction.
Limitations. Our method searches for a solution (f", T*) of a saddle point problem [4 and extracts the stochastic OT map T from it. We highlight after Lemma and in $5.1 that not all T are optimal stochastic OT maps. For strong costs, the issue leads to the conditional collapse. Studying saddle points of (14) and arg inf sets (16) is important challenge address in the further research.

---

Published as a conference paper at ICLR 2023
Potential societal impact. Our developed method is at the junction of optimal transport and generative learning. In practice, generative models and optimal transport are widely used in entertainment (image-manipulation applications like adding masks to images, hair coloring, etc.), design, computer graphics, rendering, etc. Our method is potentially applicable to many problems appearing in mentioned industries. While the mentioned applications allow making image processing methods publicly available, a potential negative is that they might transform some jobs in the graphics industry.
Reproducibility. We provide the source code for all experiments and release the checkpoints for all models of S The details are given in README, .DD the official repository.
ACKNOWLEDDEMENTT. The work was supported by the Analytical center under the RF Government (subsidy agreement 0000000D7002 Grant No. 70-2021-00145 02.11.2021).
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-1260) 1263, 2019.
Amjad Almahairi, Sai Rajeshwar, Alessandro Sordoni, Philip Bachman, and Aaron Courville. Aug- mented cyclegan: Learning many-to-many mappings from unpaired data. In International Confer- ence on Machine Learning, pp. 195-204. PMLR, 2018
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning-Voi Volume 70, pp. 146-155. JMLR. org, 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- Backhoff-Veraguas Veraguas, Mathias Beiglböck, and Gudmun Pammer. Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):1-28, 2019.
Sagie Benaim and Lior Wolf. One-sided unsupervised domain mapping. Advances in Neural Information Processing Systems, 30, 2017.
Anoop Cherian and Alan Sullivan. Sem-gan: semantically-consises image-to-image translation. In 2019 ieee winter conference on applications of computer vision (wacv), pp- 1797-1806. IEEE, 2019.
Lenaic Chizat. Unbalanced optimal transport: Models, numerical methods, applications. PhD thesis, Université Paris sciences et lettres, 2017.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8188-88 1197, 2020.
Grady Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large-scale optimal transport. Advances in Neural Information Processing Systems, 34, 2021.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Yongxin Chen, and Hao-Min Zhou. Scalable computation of monge maps with general costs. In ICLR Workshop Deep Generative Models for Highly Structured Data, 2022. URL hothhtts:// /ttpss:///peeeeveiee net dorumm?dd-rBnc
Gerald B Folland. Real analysis: modern techniques and their applications, volume 40. John Wiley & Sons, 1999.
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.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural informa- fion processing systems, pp. 2672-2680. 2014.
10

---

Published as a conference paper at ICLR 2023
Nathael Gozlan and Nicolas Juillet. On mixture of brenier and strassen theorems. Proceedings of the London Mathematical Society, 120(3):434 46344344333,, 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-3405 2017.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of Wasserstein GANs. In Advances in Neural Information Processing Systems, pp. 5767-5777, 2017.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge local nash equilibrium. In Advances in neural information processing systems, pp. 6626-6637, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6855 2020.
Xun Huang, Ming- Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image-to-image translation. In Proceedings of the European conference on computer vision (ECCV), pp. 172-189, 2018.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In European conference on computer vision, pp. 694-711. Springer, 2016.
Olav Kallenberg. Foundations of modern probability, volume 2. Springer, 1997.
Leonid Kantorovitch. On the translocation of masses. Management Science, 5(1):1-4, 1958.
Patrick Kidger and Terry Lyons. Universal approximation with deep narrow networks. In Conference learning theory, pp. 2306-2327. PMLR, 2020.
Taeksoo Kim, Moonsu Cha, Hyunsoo Kim, Jung Kwon Lee, and Jiwon Kim. Learning to discover cross-domain relations with generative adversarial networks. In International Conference Machine Learning, pp. 1857-1865. PMLR, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference Learning Representations, 2021a. URL wwRp/htttss://peeeevviee ott.:/ee.eeet//oooooldoooooooo
EXsa
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34, 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 fent/forunt forumrid 3tFAsSE-PP
Alexander Korotin, Vage Egiazarian, Lingxiao Li. and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho Advances
Advances in Neural Information Processing Systems, 2022a.
//openreview eew net. forumiid=da TnaMN
Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein GANs are not optimal transport? In Tiirt Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022b.
1d-VtEEpi-deB dGlt
Christian Léonard. A survey of the schrödinger problem and some its connections with optimal transport. Discrete & Continuous Dynamical Systems, 34(4):1533, 2014.
11

---

Published as a conference paper at ICLR 2023
Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein GAN with quadratic transport cost. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1832-4841, 2019.
Ming-Yo Yu Liu, Thomas Breuel, and Jan Kautz. Unsupervised Advances in neural information processing systems, pp. 700-708, 2017.
translation networks. In
Yahui Liu, Marco De Nadai, Jian Yao, Nicu Sebe, Bruno Lepri, and Xavier Alameda-Pinede Gmm-unit: Unsupervised multi-domain and multi-modal image-to-image translation via attribute gaussian mixture modeling. arXiv preprint arXiv:2003.06775 2020.
Ziwei Liu, Ping Luo,, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Guansong Lu, Zhiming Zhou, Yuxuan Song, Kan Ren, and Yong Yu. Guiding the one-to-one mapping in cyclegan via optimal transport. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 4432 1432-4439, 2019.
Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Large-scale optimal transport via adversarial training with cycle-consistency. arXiv preprint arXiv:2003.06633, 2020.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. International Conference on Machine Learning, pp. 6672-6681. PMLR, 2020.
Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-GAN:
Training generative neural samplers using variational divergence minimization. In Advances in neural information processing systems, pp. 271-279, 2016.
Henning Petzka, Asja Fischer, and Denis Lukovnicov. On the regularization of wasserstein gans. arXiv preprint arXiv:
2017.
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport. Foundations and Trends@ in Machine Learning, 1 1(5-6))355-607, 2019.
R Tyrrell Rockafellar. Integral functionals, normal integrands and measurable selections. In Nonlinear operators and the calculus of variations, pp. 157-207. Springer, 1976.
Ralph Rockafellar. Characterization of the subdifferentials of convex functions. Pacific Journal of Mathematics, 17(3):97-5100, 1966.
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 transpor maps. In International Conference on Learning Representations, 2022. URL URL https: pttttp:://peerrviieeee net./ met./forum:s 2dd
transport
Kuniaki Saito, Kate Saenko, and Ming-Yu Yu Liu. Coco-funit: Few-shot unsupervised image translation with content conditioned style encoder. In European Conference Computer Vision, pp. 382-398, Springer, 2020.
Maziar Sanjabi, Jimmy Ba. Meisam Razaviyayn, and Jason D Lee. On the convergence and robustness of training gans with regularized optimal transport. Advances in Newral Information Processing Systems, 31, 2018.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(5-633) 2015.
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Math- ieu Blondel. Large scale optimal transport and mapping estimation. International Conference Learning Representations, 2017.
12

---

Published as a conference paper at ICLR 2023
Xuan Su, Jiaming Song, Chenlin Meng, and Stefano Ermon. Dual diffusion implicit bridges for image-to-image translation. In International Conference on Learning Representations, 2023. URL https: http///ppeeeeeeeeeeeeeeeeeeeeeep net http://ewwwe///oounrnteee TorumiiddSHHooveee..
Amirhossein Taghvaei and Amin Jalali. 2-Wasserstein approximation via restricted convex potentials with application to improved training for GANs. arXiv preprini urXiv:1902.077 2019.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Yujia Xie, Minshuo Chen, Haoming Jiang, Tuo Zhao, and Hongyuan Zha. On scalable and efficient computation of large scale optimal transport. volume 97 of Proceedings of Machine Learning Research, pp. 6882-6892 Long Beach, California, USA, 09-15 Jun 2019. PMLR. h
//proceedings. mlr press http////////e....ee htmlll111...llllll
Karren D. Yang and Caroline Uhler. Scalable unbalanced optimal transport using generative adversarial networks. In International Conference on Learning Representations, 2019. URL
httpps:/apenreviiew net. wttp://wwwwweeeeoouuuoel//elee H
Aron Yu and Kristen Grauman. Fine-grained visual comparisons with local learning. Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 192-199, 2014.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: /506 03366, 2015.
Yongqi Zhang. Xogan: One-to-many unsupervised image-to-image translation. arXiv preprint
arXiv:1805.07277, 2018.
Bolei Zhou, Agata Lapedriza, Jianxiong Xiao, Antonio Torralba, and Aude Oliva. Learning deep features for scene recognition using places database. 2014.
Jun-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, 2017a.
Jun-Yan Zhu, Richard Zhang. Deepak Pathak, Trevor Darrell, Alexei A Efros, Oliver Wang, and Eli Shechtman. Toward multimodal image-to-image translation. Advances in neural information processing systems, 30, 2017b.
13

---

Published as a conference paper at ICLR 2023
A HARIANCE-SIMILANE TRADE-OFF
In this section, we study the effect of the parameter 7 the structure of the learned stochastic map for the y-weak quadratic cost. We consider handbags + shoes translation (64 x 64) and test y {0, 3: 3: 1}. The results shown in Figure 7
Tox Tox,21 x-p
Tox. Zsl
Tux. T(x. 24)
Tix. Z4]
Tix 24)
Tix, Ze)
Tix.
(a) 7 0
(b) Y
(c)
(d) 7 =
Figure 7: Stochastic Handbags shoes translation with the w-eak quadratic cost for various y
Discussion. For y 0 there is no variety in produced samples (Figure 7ap. i.c.. the conditional collapse happens. With the increase of 7 (Figures 7b 7c the variety of samples increases and the style of the input images is mostly preserved. For y (Figure (d), the variety of samples is very high but many of them do not preserve the style of the input image. The parameter y can be viewed the trade-off parameter balancing the variance of samples and their similarity to the input.
B TOY 2D EXPERIMENTS
In this section, we test Algorithm I on toy 2D distributions P. i...e P=
Strong quadratic cost (y = 0). As we noted $5.1 and Appendix A for the strong quadratic cost, method tends to learn deterministic maps T(x, 2) T(x) which independent of the noise input z. For deterministic maps T(x), our method yields MM:RR method which has been evaluated
in the recent Wasserstein-2 benchmark by (Korotin al., 2021b). The authors show that the method recovers OT maps well on synthetic figh-dimensional pairs P. + with known ground truth OT maps. Thus, for brevity, we do not include toy experiments with our method for the strong quadratic cost.
Weak quadratic cost (r > there are no approaches to compare with. The analysis of computed transport plans for weak costs is challenging due to the lack of nontrivial pairs P., with known ground truth OT plan n". The
0). To our knowledge, our method is the first to solve weak OT, i.e.,
14

---

Published as a conference paper at ICLR 2023
(a) Input distribution P.
(b) Target distribution Q
(c) Fitted T (P x 8) = Q
x-
x-F T0r 21
X
Tix)
Vy (x)
(d) Map ((x) Jz T(x, a))8(2). (e) Learned stochastic map (ax, =).
(f) Map x fy ydn (ular).
Figure 8: Gaussian + Mixture of 8 Gaussians, learned stochastic map for the 1-weak quadratic cost.
situation is even worsened by the nonuniqueness of cope with this issue, we consider the weak quadratic cost with y = 1. For this cost, one may derive
C(I,1) / 211z v1l
1 -11 - - ydu(wi)? 2
Var(u) =
(21)
For cost (2T) and pair IP, Q. (Gozlan & Juillet 2020, Theorem 1.2) states that there exists PPuuique
(up constant) convex 2/ Rº R such that every OT plan n satisfies Vạ(x) Sy y dx (z/lz)
Besides, Ve Rº + R² is 1-Lipschitz. Let T(x, 2) be the stochastic map recovered by our Algorithm and let â be the corresponding plan. Let
T(x) y (d = ())(2))
(22)
Due to the above mentioned characterization of OT plans, T(x) should look like gradient V*)(2) of some convex function (x(z: and should nearly be contraction. Since here work in the 2D space, we are able to get sufficiently many samples from P and Q and obtain a fine approximation of an OT plan x² and Vib by discrete weak OT solver. We may sample random batches from X ~ P and Y ~ Q of size 20 and use ot weak from POT library to get some optimal x and Vú = Sy dx (y)z). We are going compare our recovered average map T with Vvh.
Datasets. We test 2 pairs P. Q: Gaussian > Mixture of Gaussians; Gaussian Swiss roll.
Neural Networks. We use multi-layer perceptrons as fw, with 3 hidden layers 100 neurons and ReLU nonlinearity. The input of the stochastic map To 2) is 2 4 dimensional. The two first dimensions represent the input x E TR.2 while the other dimensions represent the noise 2 ~ We employ Gaussian noise with = 0.1
Discussion. We provide qualitative results in Figures and 9 In both cases, the pushforward
distribution Th (P x S) matches the desired target distribution Q (Figures 8c and 9cb. Figures 8c
hhaps://yythoonagihubbi
15

---

Published as a conference paper at ICLR 2023
(a) Input distribution P.
(b) Target distribution Q
(c) Fitted T (P 8) =0 Q
x
7ix)
20x.2)
(d) Map (xx) J'z (ac =))8(z). (c) Learned stochastic map T(x, 2).
(f) Map x Iy ydr (ul=r).
Figure 9: Gaussian + Swiss Roll, learned stochastic OT map for the 1-weak quadratic cost.
and 9e show how the mass of points x ~ P is split by the stochastic map. The average maps T(x)
(Figures 8d 9d) indeed nearly match the ground truth V (Figures 8f 9fb obtained by POT. To quantify them, we compute -2-UVP (T) = 100% ||T - V Var( V the
//ar(//ar((
metric (Korotin et al.
2021a 15.1). Here we obtain small values 1% and x 3% for the Swiss Roll and 8 Gaussians examples which further indicates the similarity of the learned T' and the ground truth Vv(x).
Note that T indeed roughly equals a gradient of a convex function. The gradients of convex functions cycle monotone (Rockafellar 1966). Cycle monotonicity yields that for 21 * ata the segments [2i, (x)()) and I22. Vil Ø22 do not intersect in the inner points (Villani 2008 (8) Visually, we
see that in Figures 8d and 9d the segments [z, T(:) do not intersect for different x, which good.
C TOY 1D EXPERIMENTS
In this section, additionally test our Algorithm on toy ID distributions P. Q, i.c., P Q = In this case, transport plans are 2D distributions and can be conveniently visualized.
We experiment with the 1-weak quadratic cost 21). Following the discussion in the previ- ous section, we recall that an OT plan x" may be not unique. However, all OT plans satisfy Vel()) = Jyy y dnd (y|z) for some 1-smooth convex function i - R. This simply means that Va(aa= = Sy y dm" (ylx)) a monotone increasing 1-Lipschitz function V R + R. Below check that this necessary condition holds for T (22 where I is our learned stochastic map.
Datasets. We test 2 pairs P, Q Gaussian >> Mix of Gaussians; Gaussian + Mix of 3 Gaussians.
Neural Networks. We use the same networks in Appendix B This time, the input of the stochastic map To 2) is 1 + = dimensional, the input to fo I-dimensional.
Discussion. We provide qualitative results in Figures 10 and II For each case, plot the results of 3 random restarts of our method (# denotes our learned OT plan). Similarly to Appendix B we plot
For the sake of clarity, we slightly reformulated the property of the cycle monotone maps Villami
2008
16

---

Published as a conference paper at ICLR 2023
Input (1D)
Target (1D)
Learned it (2D),
Mapped T.(P 5) (1D)
0.7
x-P
T(x)
10x, 7.tex
D.. πix 21)
0.5
Denkty
0
0.3
-1
0.1
-2
6.0
0.2
0.6
0.6
Density
Donstty
(a) Input and output distributions.
(b) Learned plan स and marginal TT [P x 3), test 1.
Learned it (2D), ours
Mapped 7,(P x 5) (1D)
Learned (2D),
Mapped x,PPx
X TÚO
Tlx.2) T.fe x 5)
X TiX
Fix, f
Ux. Fix. -f 1
90
0.2
0.4
00.b
-2
0.0
0.2
0.4
Density
Density
(c) Learned plan क and marginal T (P x 8), test 2.
(d) Learned plan á and marginal Ta (P x 3), test
DOT plan it' (2D)
DOT plan n² (2D)
DOT plan (2D)
DOT plan n' (2D)
Vy (x)
(x(())
Ry'
[x,y)-1
(x.y)-11
(x.7)
(x,7)-n*
-1
-2
-1
-2
-2
-1
(e) Various optimal plans =² learned by discrete OT (considered here as the ground truth).
Figure 10: Stochastic plans between toy 1D distributions (Figure 10a) learned by NOT (Figures 10b 10c TOd and discrete OT (Figure 10ep with the I 1-weak quadratic cost. The figures with the 2D transport plans also demonstrate the average map x + fy yd (v|z) (conditional expectation).
the results obtained by discrete weak OT solver (ot weak from POT library). Namely, in Figures 10e| [Te we show its results obtained for restarts with differing seeds. Note that the average maps T computed by algorithm in both cases nearly match those computed by the discrete weak OT. This indicates that the transport cost of our computed plan Ω is since
[Cost of #] = 1 1
T(x) (dP(()) RV [ (2))()))
1 1 - ((P))())))))P() = Cost(IP, Q),
1
ie., it nearly equals the optimal cost. Here we use T(x) ss (+((x) observed from the experiments. To conclude, that the recovered plans are close the DOT considered the ground truth.
D COMPARISON WITH PRINCIPAL UNPAIRED TRANSLATION METHODS
We compare Algorithm with popular models for unpaired translation. We consider handbags + shoes (64 x 64), celeba male > female (64 x 64), ouidoor - church (128 128) translation. For quantitative comparison, compute Frechet Inception Distance (Heusel et al. 2017 FID) of the
i wtthh.ccom/meetteee///tttrttttte
17

---

Published as a conference paper at ICLR 2023
Input (1D)
Target (1D)
Learned It (2D), ours
Mapped T.(P 5) (1D)
0.7
T(x)
10x, 2) 7ip
[x. πux 21) A
0.5
Denkyy
0
0.3
1
0.1
-2
6.0
0.2
0.6
0.6
Density
Donstty
(a) Input and output distributions.
(b) Learned plan # and marginal TT P x 3), test 1.
Learned it (2D),
Mapped (1D)
Learned (2D), ours
Mapped T.(P (1D)
x T(X)
Tlx.2) T((xx)
x T(x)
TXx Tx
Fix.
Ix. Tix, f
0.0
0.2
0.4
0.6
1
0.0
0.2
0.4
0.5
Density
Denstty
(c) Learned plan क and marginal T (P 8). test 2.
(d) Learned plan á and marginal Ta (IP x 3), test 3.
DOT plan n²' (2D)
DOT plan n² (2D)
DOT plan (2D)
DOT plan n' (2D)
X Ty (x)
x++ Dp'(x)
X Ry
x Ry
Ix.y1
0x.y1
(x,))
(x,y)-
-1
(c) Various optimal plans x² Icarned by discrete OT (considered here as the ground truth).
Figure 11: Stochastic plans between toy 1D distributions (Figure 11a learned by NOT (Figures) 1b 11e 11d and discrete OT (Figure TTcp with the 1-weak quadratic cost. The figures with the 2D transport plans also demonstrate the average map x + Sy ydd (y/zz) (conditional expectation).
mapped test handbags subset w.r.t. the test shoes subset. The scores of our method and alternatives
given in Table [ The translated images shown in Figures 12 13 14
2017a). DiscoGAN B Methods. We compare our method with one-to-one CycleGAN Zhu et al.
(Kim al. 2017) and with one-to-many AugCycleGAN (Almahairi et al. 2018) and MUNIT (Huang et al. 2018). We use the official or community implementations with the hyperparameters from the respective papers. We choose the mbove-mentioned methods for comparison because they principal methods for one-to-one and one-to-many translation. Recent methods (GMM-UNIT (Liu al. 2020), COCO-FUNIT (Saito al. 2020). StarGAN (Choi et al. 2020)) are based on them and focus on specific details/setups such as style/content separation, few-shot learning, disentanglement, multi-domain transfer, which are out of scope of our paper.
Discussion. Existing one-to-one methods visually preserve the style during translation comparably to our method. Alternative one-to-many methods do not preserve the style at all. When the input
cycl legne igi giithb..comm discoqan qi htbbhbccow/allaal/auuuueee...ccl
g tithbb.ccm//w/aabs///Wvn
18

---

Published as a conference paper at ICLR 2023
and output domains similar (handbags- celeba male - female), the FID scores of all the models are comparable. However, most models are outperformed by NOT when the domains are distant (outdoor church), Figure [4 and the last in Table T For completeness, in Table we compare the number of hyperparameters of the translation methods view. Note that in contrast the other methods, optimize only neural networks transport map and potential.
Type
One-to-one
One-to-many
Disco GAN
Cycle GAN
NOT (ours)
AugCycle GAN
NOT
Method
MUNIT
(ours)
Handbags + shoes
18.84
15.76
13.44
22.42
16.00
13.77
x64x 64)
= 0.11
0.11
= 0.12
Celeba male + female x64x 64)
12.94
17.07
11.96
35.64
17.74
13.23
10.08
10.11
10.07
Outdoor + church (128x 128)
51,42
31.42
25.97
75.36
46.39
25.5
#0.12
#0.166
10.14
Table 1: Test FID, of the considered unpaired translation methods.
Type
One-to-one
One-to-many
Cycle GAN
AugCycle GAN
Method
MUNIT
GAN
(ours)
(ours)
Weights of cycle and
Weights of reconstruction Aus Ac: A
Hyperparameters
Weights eycle losses 71. 7/2
Diversity control parameter y
opimization
None
None
objectiives
identily
Acyer Asd
Total number hyperparameters
2
generators L.MM, 2 discriminators
2 generators, x 29.2M
2 generators 2x 11.MM
transport 9.7M,
generntees 2x 15.0M, discrimminatoes 2x8.3M
transport 9.7M,
Networks
discrimninaoors 2x0.7M
discriminators 2x2.8M
potential 22.9M (32.4M"]
potential 22.9M [32.4M*]
14M
Total number networks and parameters
networks 59.8M
networks 28.2M
2 networks
networks 7.0M
4 networks
networks 32.6M [42.1M*)
32.6M [42.
46.6M
Table 2: Comparison of the number hyperparameters of the optimization objectives, the number of networks and their parameters for the considered unpaired translation methods for 64 64 images.
For 128 128 images.
19

---

Published as a conference paper at ICLR 2023
X-
TIX)
(a) NOT (ours, W2), one-to-one.
TixI
(b) DiscoGAN, one-to-one.
T(x)
(c) CycleGAN, one-to-one.
X-P
X-P
K-P
Tourz
Ttx
Rox
Tix.z.l
(d) NOT (ours, W2, 3 one-to-many.
(e) MUNIT, one-to-many.
(I) AugCycleGAN
Figure 12: Handbags + shoes translation (64 x 64) by the methods in view.
20

---

Published as a conference paper at ICLR 2023
(a) NOT (ours, W2). one-to-one.
Fbd
(b) DiscoGAN, onc-to-one.
X
(c) CycleGAN, one-to-one,
e-p
A-P
Fix.t))
Tbx
fouin)
Tix.)))
Tir.
fic.
Plx zil
(d) NOT (ours, W2.3), one-to-many.
(c) MUNIT, one-to-many.
(f) AugCycleGAN
Figure 13: Celeba (male) - Celeba (female) translation (64 64) by the methods in view.
21

---

Published as a conference paper at ICLR 2023
T(x)
7IX
DiscoGAN, one-to-one.
(b) CycleGAN, one-to-one.
T(x)
(c) NOT (ours, W2), one-to-one.
Tix 21)
T(x,2))
T 2g1
T(x,z))
Txx.
7(x,
(d) AugCycleGAN, one-to-many.
(c) MUNIT, one-to-many.
Tix,21)
T(x
Tux, 2z)
(1) NOT (ours. W2,3) 3). one-to-neee one-to-many.
Figure 14: Outdoor church (128 128) translation with various methods.
22

---

Published as a conference paper at ICLR 2023
E EXPERIMENTAL DETAILS
Pre-processing. We beforehand rescale anime face images 512 x 512, and do 256 x 256 crop with the center located 14 pixels above the image center to get the face. Next, for all these datasets, we rescale RGB channels to [-1, and resize images to the required size (64 x or 128 128). We do not apply any augmentations to data.
Neural networks. We use WGAN-QC discriminator's ResNet architecture (Liu et al. 2019) for potential f. We use (Roneebrger et al., 2015) as the stochastic transport map T T(x 2). The noise 2 is simply additional 4th input channel (RGBZ), i.e., the dimension of the noise equals the image size (64 x 64 or 128 x 128). We use high-dimeene sigh-dimensional Gaussian noise with axis Ø = 0.1.
2014 with the default betas for both To and fa- The learning rate is Ir = 1- 10-1. The batch size is X 64. The number inner iterations is by 10. When training with the weak cost ( sample Zzr| x batch. In toy experiments, we do 10K total iterations of fw update. In the experiments with unpaired translation, our Algorithm[] converges in 40K iterations for most datasets.
Optimization. We the Adam optimizer (Kingma & Ba
noise samples per each image
Dynamic weak cost. In 55.3 we train the algorithm with the gradually changing Starting from y 0, we linearly increase to the desired value g or 1) during 25K first iterations of Sw-
Stability of training. In several cases, we noted that the optimization fluctuates around the saddle points or diverges. An analogous behavior of saddle point methods for OT has been observed in Korotin et al. 2021b). For the y-weak quadratic cost (y > 0), we sometimes experienced instabilities when the input P is notably less disperse than Q or when the parameter y high. Studying this behaviour and improving stability/convergence of the optimization is a promising research direction.
Computational complexity. The time and memory complexity of training deterministic OT maps T(x) is comparable to that of training usual generative models for unpaired translation. Our networks converge in 1-3 days on a Tesla V100 GPU (16 GB); wall-clock times depend on the datasets and the image sizes. Training stochastic T'z z) is harder since we sample multiple random per (we use [Z| = 4). Thus, we learn stochastic maps on 4 x Tesla V100 GPUs.
F OPTIMALITY OF SOLUTIONS FOR STRICTLY CONVEX COSTS
Our Lemma 4 proves that optimal maps T are contained in the arg inf sets of optimal potentials f but leaves the question what else may be contained in these arg inf sets open. Our following result shows that for strictly convex costs, nothing else beside OT maps is contained there.
Lemma 5 (Solutions of the maximin problem are OT maps). Let C(z, p)) be weak cost which is strictly convex p. Assume that there exists at least one potential f which maximizes dual form (5)
Consider any such optimal potential f E aag s inf C(J,T) It holds that
i E arginf inf L(F", + f is a stochastic OT map.
T
Proof of Lemma
15 By the definition of f*, we have
C(f f) = sup inf C(f, Cost(PP Q),
i.e., î attains the optimal cost. It remains to check that it satisfies πB(P x 8) Q i.e., î generates T (x, (xx 1 and ú. f(x .)3 for all x E X and define n = 1 (p= + ). Let T X' x - y be any stochastic map which satisfies T (ar, :)ss- = 1h for all x E X 4.1 By using the change of variables, we derive
Q from P. Let T be any true stochastic OT map. We denote μ²
Cost(IP 2 (T)) = / Cla, ((o,))))P())) Jr1 1/ ((wddd (o)
+ ((OQO)
(23)
g wtt.b.ccmmmileeiiaaaa lesi
23

---

Published as a conference paper at ICLR 2023
Since C' is convex in the second argument, we have
C(a, (f) = Cz., 2 (1 + (w)) 2 2 C( mC)) 0zz, Ax).
1
(24)
Since is strictly convex, the equality in 24 possible only when μ = Px. We also note that
J 5 = -/5 f (y) (12 + Pe)
(v) 2 / (())d + 2 / (()ddpa()).
1
2
We substitute these findings to ,(f,," and get
- Cost(P, +) C(", 2/2 + 2 1) - Cost((P Q)+ Cost(P, P. Q) Cost(IP, Q)
Thus, 23 is an equality P-almost surely for all a E X and μrz = fgg holds P-almost surely. This means that 7 and I generate the same distribution from IP x S, i.c., T is stochastic OT map. O
Our generic framework allows learning stochastic transport maps (Lemma 4). For strictly convex costs, all the solutions of our objective (14) are guaranteed to be stochastic OT maps (Lemma In the experiments, we focus on strong and weak quadratic costs, which are not strictly convex but still provide promising performance in the downstream task of unpaired image-to-image translation ($5) Developing strictly convex costs is promising research avenue for the future work.
G RELATION TO PRIOR WORKS IN UNBALANCED OPTIMAL TRANSPORT
In the context of OT, Yang & Uhler 2019) employ a stochastic generator to learn a transport plan
in the unbalanced OT problem (Chizat 2017). Due to this, their optimization objective slightly resembles our objective (5)). However, this similarity is deceptive. Unlike strong 2 or weak 3 [T(P,
OT, the unbalanced OT is unconstrained problem, i.e., there is need to satisfy
This makes unbalanced OT easier to handle: to optimize it one just has to parametrize the plan and backprop through the loss. The challenging part with which the authors deal is the estimation of the o-divergence terms in the unbalanced OT objective. These terms can be interpreted as a soft relaxation of the constraints π [IP,, Q), i.e., penalization for disobeying the constraints. The authors compute these terms by employing the variational (dual) formula from f-GAN (Nowozin al. 2016). This yields a GAN-style optimization problem minr max, which is similar to other problems in the generative adversarial framework. The problem we tackle is strong (2) and weak (3) OT which requires enforcing of the constraint æ E [I( Q). We reformulate the dual (weak) OT problem (5) into maximin problem (15) which can be used to recover the OT plan (via the stochastic map T)). Our approach can be viewed as a hard enforcement of the constraints. Our max, minr saddle point problem (15) is atypical for the traditional generative adversarial framework.
24

---

Published as a conference paper at ICLR 2023
H ADDITIONAL EXPERIMENTAL RESULTS
T(x)
(a) Celeba (female) > anime translation, 128 128.
T(x)
(b) Outdoor church, 128 128.
T(x)
(c) Handbags -> shoes, 128 x 128.
X-p
T(x)
(d) Shoes handbags, 128 128.
Figure 15: Unpaired translation with OT maps (W2). Additional examples (part 1).
25

---

Published as a conference paper at ICLR 2023
f(x)
(a) Celeba (female) + anime translation, 64 64.
f(x)
(b) Anime celeba (female) translation, 64 64.
T(x)
(c) Celeba (male) + celeba (female) translation, 64 x 64.
X- p
T(x)
(d) Anime + translation, x
Figure 16: Unpaired translation with OT maps (W2). Additional examples (part 2).
26

---

Published as a conference paper at ICLR 2023
x-P
T(x. 21)
7(x. 23)
7
Tix, 74)
(a) Input images and random translated examples T(r, z).
TiX, 2)
1,7
2/7
3/7
== 4/7
5/7
== 6/7
Tx)
(b) Interpolation in the conditional latent space, (1 0) 0222.
Figure 17: Celeba (female) + anime, 128 x 128, stochastic. Additional examples.
27

---

Published as a conference paper at ICLR 2023
T(x 21
7(x,
Tix, 7sl
Tix, Z4)
(a) Input images and random translated examples T(x, 2).
P
Tix
a 1/7
a 2/7
3/7
4/7
5/7
6J7
Tu.
(b) Interpolation in the conditional latent space, z = (1 01)2 0222
Figure 18: Outdoor + church, 128 128, stochastic. Additional examples.
28

---

Published as a conference paper at ICLR 2023
x-p P
Tix, 21)
T(x. zz)
T(x, 2g)
T(x, Za)
(a) Input images x and random translated examples T(x, 2).
x-P
T(x, z
a= 1/7
a 2/7
a = 3/7
a = 4/7
5/7
6/7
T(x. 22)
(b) Interpolation the conditional latent space, 2 (1 0.)a + 0.22.
Figure 19: Handbags - shoes translation, 128 128, stochastic. Additional examples.
29

---

Published as a conference paper at ICLR 2023
X-P
T(x, 21)
T(x, 22)
T(x, 23)
T(x, 24)
(a) Input images and random translated examples T'x z).
x-
Tox.z)
e- 1/11
2-111
a- 3/11
a 4/11
2-1//11
d 6/1
-- 7/11
a 8/11
a 9/11
a 10/11
Tix
(b) Interpolation in the conditional latent space, (1 0) 022.
Figure 20: Anime - celeba (female) translation, 64 64, stochastic. Additional examples.
30

---

Published as a conference paper at ICLR 2023
X-F
T(x.z1
Tix, Zal
ftx, 2,)
Tix, Tax,
(a) Input images or and random translated examples T(r, 2).
T(x.2.)
(b) Interpolation in the conditional latent space, z (1 a) a22.
Figure 21: Anime shoes translation, 64 64, stochastic. Additional examples.
31

---

Published as a conference paper at ICLR 2023
X-
T(x, 21)
Tix. Zal
Ttx. 2,)
T(x, Zal
(a) Input images ær and random translated examples ((x, 2).
2/1/10
(b) Interpolation in the conditional latent space, z = (1 ax)2 a22.
Figure 22: Shoes handbags, 64 x 64, stochastic. Additional examples.
32

---

Published as a conference paper at ICLR 2023
I EXAMPLES WITH THE SYNCHRONIZED NOISE
In this section, for handbagsssshoes sshoes (64x64) and outdoor (128x 128) datasets, we pick batch of input data .1...... N P and noise 21...... 2K S to plot the N K matrix of generated images To (In, =b). Our goal is to assess whether using the same o for different 2n leads to some shared effects such as the same form a generated shoe or church.
The images results are given in Figures 23 and 24 Qualitatively, we do not find any close relation between images produced with the same noise vectors for different input images In.
X- P
7(x,23) T(, 22) 7z),
Z(x.zs) T(X, Z4)
Za) Z,) T(x, Z6) Z6)
Z10) Z10) Tx)
Figure 23: Input images and random translated examples T'æ 2).
this example, noise inputs z are synchronized between different input images æ.
33

---

Published as a conference paper at ICLR 2023
X P
7(x.2A) T(x. T(x, T(x
T(x, Zs)
T(x, Z1o) T( T(X,2 Za) T(x,2,) T(x.zs)
Figure 24: Input images æ and random translated examples T(x, 2).
In this example, noise inputs 2 are synchronized between different input images x.
34


Reviews:
[{"review_id": "umkWRxxWGyj", "reviewer": "Reviewer_D9LF", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "This work presents a novel neural-networks-based algorithm to compute optimal transport maps and plans for both strong and weak transport costs. The problem is converted to a minmax optimization one using noise outsourcing idea. The work proves that the proposed neural-networks are universal approximators of transport plans between probability distributions. The algorithm is evaluated on both toy examples and on unpaired image-to-image translation.\n\nThe contributions of the work are\n1. proposes a novel framework to compute optimal transport maps and plans for both strong and weak transport costs based on noise outsourcing method;\n2. give rigorous proofs for the existence of the solution and show the proposed model is a universal approximator\n3. thorough experimental results to demonstrate the effectiveness of the proposed model", "summary_of_the_review": "This work presents a novel neural-networks-based algorithm to compute optimal transport maps and plans for both strong and weak transport costs. The problem is converted to a minmax optimization one using noise outsourcing idea. The work proves that the proposed neural-networks are universal approximators of transport plans between probability distributions. The algorithm is evaluated on both toy examples and on unpaired image-to-image translation.\n\nIt is well known that optimal transport maps/plans play fundamental roles in deep learning, but their computations are intrinsically complicated. An efficient and accurate computational method is highly desirable. This work offers a novel method for tackling this fundamental problem by converting it as a min-max optimization problem using the noise outsourcing idea. The algorithm can compute both deterministic and stochastic mappings with strong and weak transport costs. The various applications show the flexibility and effectiveness of the method. Although there are some theoretic ambiguity for the saddle points, the proposed method is inspiring the convincing.\n\n", "strength_and_weaknesses": "The work has the following strengths:\n\n1. The problem to tackle is to efficiently compute optimal transport maps or plans, which play fundamental roles in deep learning and are highly non-linear and notoriously hard to compute.\n2. By using the noise outsourcing idea, the problem is converted to a minmax optimization problem, which can be tackled efficiently\n3. The experimental results demonstrate the effectiveness of the proposed modal\n\nThe weakness is that the proposed algorithm searches for a solution of a saddle point problem and extract the stochastic OT map from it, but not all saddle points correspond to optimal stochastic OT maps, this type of ambiguity requires further theoretical exploration.", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Z4ADTqS80pq", "reviewer": "Reviewer_KZKB", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "This paper provides a scalable way to learn OT maps via a deep neural network. It is especially relevant due to the recent approaches demonstrating the use of OT maps for generative purposes, as opposed to the previous methods where OT was used as a loss when training generators.\n", "summary_of_the_review": "This is a very good paper.", "strength_and_weaknesses": "The paper is very well-written and pedagogical. It also mixes a good deal of theoretical and algorithmic approaches along with thorough empirical evaluations. Maybe the only weakness is that it is mostly motivated by the situation where no deterministic map approaches, and while it is clear that it might be theoretically the case, it is less clear when that situation occurs in practice. \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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "FOEu99AwDo", "reviewer": "Reviewer_NG9Z", "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": "This paper proposed a neural network based solver which estimates an explicit transport map rather than just using OT losses. Specifically, the authors introduced a weak OT formulation so as to find the potentially stochastic optimal transport map. The resulting objective becomes a minimax optimization which is learning a transport map and a potential simultaneously and both of them can be parameterized as neural networks. The proposed method is evaluated on a couple of datasets and compared against several baselines. The experimental results are quite encouraging as the estimated map can pushforwards image samples among different contexts but keep the style. ", "summary_of_the_review": "In this paper, the authors proposed a framework to estimate the optimal transport map parameterized as a neural network. The paper is well-written, and the contributions are clearly presented. The proposed weak OT cost extends the existing Kantorovich dual formulation, and yields a practical objective for training an OT map on large-scale datasets. The experimental result is quite encouraging. \nThus, I tend to recommend the acceptance of this work. ", "strength_and_weaknesses": "Pros:\n- Overall, the paper is well-written as the authors delivered their proposed approach clearly. \n- The organization of section 4 is encouraging and easy to follow. It is clear that the proposed method is motivated by the Kantorovich dual formulation with an additional C-transformation.\n- The experimental part is quite sufficient. The authors have conducted experiments on multiple datasets, including anime faces, CelebA datasets and so on. This is quite valuable as OT map estimation approaches are usually concerned with scalability.  \n\nCons and questions:\n- Is it possible to better justify the advantages of the optimal transport map formulation, rather than existing deep generative frameworks? An optimal transport map may allow for better interpretability and generalization. Is it possible to have so called out-of-sample transport results, like in figure 1 of[1], figure 4 of [2]? Also, given the ideally nice smoothness properties of OT maps, shall we expect some results like \"king - man + women - queen\" when we are transporting image samples?\n- While the problem formulation doesn't specify the ground metric. It seems that in the experiments, the authors are still using the quadratic cost. It is reasonable for the image datasets, however it could be nice if there are experimental results for other types of ground metrics, such as cosine similarity for language embeddings.\n- What are the values of z_1, z_2, and z_3 in figure 6, 7, 12, 17 and later?  Does that impose some kind of interpretability on the map?\n- From my understanding, it is true that under the given conditions, an optimal transport plan induced the transport map, which further characterizes the corresponding Wasserstein barycenters and the geodesic. However, given that the transport map is parameterized by a feedforward neural network (U-net), how does this U-net characterize the joint distribution of two distributions? In toy-1D experiments, the authors visualize the optimal plan but that is given by the ot.weak package.\nIn addition, the discriminator in a GAN can be viewed as a classifier. Are there any interpretations we can draw from the potential network f?\n\n[1] Perrot, Michaël, et al. \"Mapping estimation for discrete optimal transport.\" Advances in Neural Information Processing Systems 29 (2016).\n\n[2] Zhu, Jiacheng, et al. \"Functional optimal transport: map estimation and domain adaptation for functional data.\" arXiv preprint arXiv:2102.03895 (2021). \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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%