Generative Modeling with Optimal Transport Maps
Authors: Litu Rout
Alexander Korotin
Evgeny Burnaev
Venue: ICLR 2022 Poster (2022)
Keywords: Optimal Transport Map
Generative Modeling
Unpaired Image Restoration

Abstract: With the discovery of Wasserstein GANs, Optimal Transport (OT) has become a powerful tool for large-scale generative modeling tasks. In these tasks, OT cost is typically used as the loss for training GANs. In contrast to this approach, we show that the OT map itself can be used as a generative model, providing comparable performance. Previous analogous approaches consider OT maps as generative models only in the latent spaces due to their poor performance in the original high-dimensional ambient space. In contrast, we apply OT maps directly in the ambient space, e.g., a space of high-dimensional images. First, we derive a min-max optimization algorithm to efficiently compute OT maps for the quadratic cost (Wasserstein-2 distance). Next, we extend the approach to the case when the input and output distributions are located in the spaces of different dimensions and derive error bounds for the computed OT map. We evaluate the algorithm on image generation and unpaired image restoration tasks. In particular, we consider denoising, colorization, and inpainting, where the optimality of the restoration map is a desired attribute, since the output (restored) image is expected to be close to the input (degraded) one.

Published as a conference paper at ICLR 2022
GENERATIVE MODELING WITH OPTIMAL TRANSPORT MAPS
Litu Rout
Alexander Korotin
Space Applications Centre Indian Space Research Organisation Ir@sac. isro. gov. in
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute (AIRT)
a. wworotin@kooleee wetsskkoleeh...... ru
Evgeny Burnaev
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute (AIRI) e. w..brnaev@skoll tech. ru
ABSTRACT
With the discovery of Wasserstein GANs, Optimal Transport (OT) has become a powerful tool for large-scale generative modeling tasks. In these tasks, OT cost is typically used as the loss for training GANs. In contrast to this approach, we show that the OT map itself can be used generative model, providing comparable performance. Previous analogous approaches consider OT maps as generative models only in the latent spaces due their poor performance in the original nigh-dimensional ambient space. In contrast, we apply OT maps directly in the ambient space, e-g., a space of nigh-dimensional images. First, we derive a min- max optimization algorithm to efficiently compute OT maps for the quadratic cost (Wasserstein-2 distance). Next, we extend the approach to the case when the input and output distributions are located in the spaces of different dimensions and derive error bounds for the computed OT map. We evaluate the algorithm on consider
image generation and unpaired image restoration tasks. In particular, denoising, colorization, and inpainting, where the optimality of the restoration map is desired attribute, since the output (restored) image is expected to be close to the input (degraded) one.
INTRODUCTION
Since the discovery of Generative Adversarial Networks (GANs, Goodfellow et al. (0114)) there has been surge in generative modeling (Radford et al. 2016 Arjovsky et al. [20 7 Brock al. 2019 Karras et al.) 2019). In the past few years, Optimal Transport (OT, Villani (2008) theory has been pivotal in addressing important issues of generative models. In particular, the usage of Wasserstein distance has improved diversity (Arjovsky et al. 2017] Gulrajani et al.) (017)) convergence (Sanjabi al. 2018), and stability (Miyato al. 2018 Kim et al. (021) of GANs.
Generative models based on OT can be split into two classes depending on what OT is used for. First, the optimal transport cost serves the loss for generative models, see Figure Ta This is the most prevalent class of methods which includes WGAN (Arjovsky et al. 2017) and its modifi- cations: WGAN-GP (Gulrajani al. (017). WGAN-LP (Petzka al. 2018). and WGAN-QC (Liu et al. 2019). Second, the optimal transport map is used as a generative model itself, see Fig- ure Tb Such approaches include LSOT (Seguy et al. 2018), AE-OT (An al. 2020a). ICNN-OT Makkuva et al. 2020), W2GN (Korotin et al.) 2021a). Models of the first class have been well- studied, but limited attention has been paid to the second class. Existing approaches of the second class primarily consider OT maps in latent spaces of pre-trained autoencoders (AE), see Figure 3 The performance such generative models depends on the underlying AEs, in which decoding transformations often not accurate; high-dimensions ambient spaces. For this reason, using OT in the latent space does not necessarily guarantee superior performance in generative modeling.
result this deficiency limits practical applications in

---

Published as a conference paper at ICLR 2022
Input distribution
generator network Ba
Fake distribution
OT map as generative map
OT cost as the loss function
Input distribution
Output distribution
True distribution
(a) OT cost the loss for the generative model.
OT map the generative model.
Figure 1: Two existing approaches to use optimal transport in generative models.
The focus of our paper is the second class of OT-bbaed models using OT map as the generative map. Finding an optimal mapping is motivated by its ability to preserve specific attributes of the input samples, a desired property in unpaired learning. For example, in unpaired image-to-image transla- tion. the learner has fit map between CycleGAN-brie
data distributions which preserves the image content. models (Zhu et al., (017)) are widely used for this purpose. However, they typi- cally have complex optimization objectives consisting of several losses (Amodio & Krishnaswamy 2019
Lu et al. 2019) in order to make the fitted map preserve the required attributes.
The main contributions of this paper are as follows:
We propose an end-to-end algorithm (84.3) fit OT maps for the quadratic cost (Wasserstein-2 distance) between distributions located on the spaces of equal dimensions ($4.I) and extend the method to unequal dimensions as well ($4.2). We prove error bounds for the method ($4.4). We demonstrate large-scale applications of OT maps in popular computer vision tasks. We con- sider image generation ($5.1) and unpaired image restoration ($5.2) tasks.
Our strict OT-based framework allows the theoretical analysis of the recovered transport map. The map obtained by our method can be directly used in large-scale computer vision problems which is in high contrast to previous related methods relying autoencoders and OT maps in the latent space. Importantly, the performance and computational complexity of our method is comparable to OT-based generative models using OT cost the loss.
Notations. In what follows, X' and y are two complete metric spaces, (u(z)) and (vy are probability distributions X' and V/ respectively. For measurable map T X' > V. T T+ denotes the pushforward distribution of i.e., the distribution for which any measurable set E C y satisfies Ty p(E)= (E)) For vector ||a/|| denotes its Euclidean norm. We use (ar, y) to denote the inner product of vectors x and y. We use [T(,,, v) to denote the set of joint probability distributions on X x y whose marginals are μ and v. respectively (couplings). For function f R ER U = its Legendre-Fenchee convex, even if is not.
transform (the convex conjugate) fly) suPzER" {(x, v) (x)) It is
2 BACKGROUND ON OPTIMAL TRANSPORT
Consider cost of transportation, c X IR defined over the product space of X' and V.
Monge's Formulation. The optimal transport cost between and for ground cost e( is
c(x.T(x.))
T(xi)
X
c(xo.T(x)))
T(xa)
Cost(u, v) def inf Twn=v
/ c (x, cxx,T(x) du(x),
(1)
/1
T(x.)
c(x.T(x))
X,
where the infimum is taken over all measurable maps T X + pushing y to v, see Figure The map T on which the infimum in (I) attained is called the optimal transport
Figure 2: Monge's OT.
2

---

Published as a conference paper at ICLR 2022
map. Monge's formulation does not allow splitting. For example, when j is Dirac distribution and V/ is a non-Dirac distribution, the feasible set of equation (T) is empty.
Kantorovich's Relaxation. Instead of asking which particular point y 3 should all the proba- bility mass o be moved, [aatorovich (1948) asks how the mass of æ should be distributed among all y E V. Formally, transport coupling replaces a transport map; the OT cost is given by:
Cost(j, v) def inf =ETT(())))
/ e c(x, v) dr(
where the infimum is taken over all couplings т 11(,, 1/) of μ and v. The coupling x¹ attaining the infimum of (2) is called the optimal transport plan. Unlike the formulation of ( the formu- lation of ( is well-posed, and with mild assumptions on spaces X.Y' and ground cost c(.,),, the minimizer of (2) always exists (Villani 2008 Theorem 4.1). In particular, if x* is deterministic, i.., πं' = 1 7 # for some
T" +y, then 7 minimizes
Duality. The dual form of (2 is given by (Kantorovich) 1948):
(// u(x))u(x)+ + I v(y)du(v): u(x) + v(y) s c(x,w) mn (a.0)
Cost(u, v) = sup
(3)
with u L' ('()), (v) called Kantorovich potentials. For X + R and y R de- fine their c-transforms by u (y) = infzex(e (c (2, y) u (2)) and v (x)) = inf inf v (v)}
respectively. Using c-transform, is reformulated as (Villani 2008 55)
Cost(H, v) = U (x)dd(())++ /, o(w)dr(() sup / u(x)du(x)+ / u (v)dd(u)) (4)
=sup
Primal-dual relationship. For certain ground costs c).,,), -), the primal solution T of can be recovered from the dual solution u" of (3). For example, if X' ) R c(x,)) y) h( y) with strictly convex h : RD + IR and μ is absolutely continuous supported on the compact set, then
T (x)
(Vh)"'(V'" (x)),
(5)
see (Santambrogio 2015 Theorem 1.17). For general costs, see (Villani 2008 Theorem 10.28).
OPTIMAL TRANSPORT IN GENERATIVE MODELS
OPTIMAL TRANSPORT COST AS THE LOSS (Figure Ta). Starting with the works Arjovsyy & Bottou (0177): Arjovsky et al. (017), the usage of OT cost as the loss has become a major way to apply OT for generative modeling. In this setting, given data distribution v and fake distribution Ma.. the goal is to minimize Cost v) w.r.t. the parameters 0. Typically, μg is pushforward distribution some given distribution, e.g., (0, I), via generator network Gg.
The Wasserstein- / distance (W1). i.e., the transport cost for ground cost c(x, v) = ||x yll. is the most practically prevalent example of such a loss. Models based on this loss are known as Wasserstein GANs (WGANs). They estimate W; (He,, based on the dual form as given by (4). For W1, the optimal potentials u", e" of (4) satisfy u" -v" where u² is a 1-Lipschitz function (Villani 2008 Case 5.16). As result, compute Wi, one needs optimize the following simplified form:
{ u(x))po((r) - u(g)ddu()) Wi(na, = sup [1
(6)
In WGANs, the potential u is called the discriminator. Optimization of 66) reduces constrained optimization of (4) with two potentials u, e to optimization of only one discriminator u. In practice, enforcing the Lipschitz constraint u is challenging. Most methods to do this are regularization- based, e.g., they use gradient penalty (Gulrajani et al., 2017 WGAN-GP) and Lipschitz penalty (Petzka et al. 2018, WGAN-LP). Other methods enforce Lipschitz property via incorporating cer- tain hard restrictions on the discriminator's 's architecture (Anil et al. 2019 Tanielian & Biau 2021).
General transport costs (other than W1) also be used as the loss for generative models. They are less popular since they do not have a dual form reducing to a single potential function simi- lar to (6) for Wi. Consequently, the challenging estimation of the c-transform of is needed. To
3

---

Published as a conference paper at ICLR 2022
avoid this, Sanjabi al. (2018) consider the dual form of (3) with two potentials u, instead form 4 with one u and softly enforce the condition u(x) + bly) I c(z,3) 3) via entropy or quadratic regularization. Nhan Dam et al. (2019) use the dual form of 4 and amortized optimization to compute wf via an additional neural network. Both methods work for general c(..)), though the authors test them for c(z; y) = || - y/l only, i.e., Wi distance. Mallasto et al. (2019) propose fast way to approximate the c-transform and test the approach (WGAN-(q, P)) with several costs, in partieular, the Wasserstein-2 distance (W2),, i..., the transport cost for the quadratie ground cost c(z, y) 3 || vll². Specifically for W2, Liu et al. (2019) approximate the c-transform via linear program (WGAN-QC).
A fruitful branch of OT-based losses for generative models comes from modified versions of OT cost, such as Sinkhorn (Genevay et al. 2018), sliced (Deshpande al., 2018) and minibatch (Fatras) et al. 2019) OT distances. They typically have lower sample complexity than usual OT and can be accurately estimated from random mini-batches without using dual forms such as In practice, these approaches usually learn the ground OT cost c)..))
The aforementioned methods use OT cost the ambient space to train GANs. There also exist ap- proaches using OT cost in the latent space. For example. Tolstikhhin et al. (2017); Patrini et al. (2020) use OT cost between encoded data and a given distribution as an additional term to reconstruction loss for training an AE, As the result, AE's latent distribution becomes close to the given one.
OPTIMAL TRANSPORT MAP AS THE GENERATIVE MAP (Figure lb Methods to compute the OT map (plan) are less common in comparison to those computing the cost. Recovering the map from the primal form (T) or (2) usually yields complex optimization objectives containing several adversarial terms (Xie al. 2019 Liu al. 2021 Lu et al., 2020).. Such procedures require careful hyperparameter choice. This needs to be addressed before using these methods in practice.
Primal-dual relationship (:2) makes it possible recover the OT map via solving the dual form (3) Dual-form based methods primarily consider W2 cost due to its nice theoretical properties and rela- tion convex functions (Brenier [191). In the semi-discrete case (μ is continuous,
discrete), An et al. 2020a) and Lei et al. 2019) compute the dual potential and the OT map by using the Alexandrov theory and convex geometry. For the continuous case, (Seuuy al. (2018) use the en- tropy (quadratic) regularization to recover the dual potentials and extract OT map from them via the barycenteric projection. Taghvaei Jalaii (2019), Makkuva al. (2020), Korotin et al. (2021a) employ input-convex neural networks (ICNNs, see Amos er al. (2017)) to parametrize potentials in the dual problem and recover OT maps by using their gradients.
Latent space
OT map
|Enc
Dee
distrbbtion
space
Figure 3: The existing most prevalent approach to use OT maps in generative models.
The aforementioned dual form methods compute OT maps in LATENT SPACES for problems such as domain adaptation and latent space mass transport, see Figure 3 OT maps in high-dimensional ambient spaces, e.g., natural images, are usually not considered. Recent evaluation of continuous OT methods for W2 (Korotin et al., (021b) reveals their crucial limitations, which negatively affect their scalability, such as poor expressiveness of ICNN architectures or bias due to regularization.
END-TO-END SOLUTION TO LEARN OPTIMAL MAPS
4.1 EQUAL DIMENSIONS OF INPUT AND OUTPUT DISTRIBUTIONS
In this section, we use X'
P and consider the Wasserstein-2 distance (W2),, i.e. the optimal
transport for the quadratic ground cost c(x, y) 1 || 3. We use the dual form (4) to derive a saddle point problem the solution of which yields the OT map T We consider distributions μ. with finite second moments. We assume that for distributions μ4, v in view there exists a unique OT
plan x" minimizing (3) and it is deterministic, i.e., x² [idgo, p, T*) T' Here T is an OT map
4

---

Published as a conference paper at ICLR 2022
which minimizes I Previous related works (Makkuva et al.) 2020 Korotin et al. 202 assumed the absolute continuity of μ., which implied the existence and uniqueness of T (Brenier 1991).
Let (v) def ||ybll2 v(y), where is the potential of (4 Note that
of (1)
inf yER
vll v(1)) I||=||
1
sup ((x,)) - +20)3 -51=011 +(x)) (7)
VER
Therefore, 4 is equivalent
W (H, = / | dp(x) + |1ll dv(x)) + sup {-/ ((z)dp(z)- - v(v)du(u) } (8)
Constant(Ju, v) inf { (())dp(()) / ((v)dv()) } (9)
Constant(p,,)) v) inf {/r VER"
sup {(x, - {1)) du(x)) + / (())d(()) (10)
Constant(yu, v) inf sup / x -((()))))(())) dp(x)) / ((o)du()) (11)
where between lines 110 and we replace the optimization over y E R.D with the equivalent RD. The equivalence follows from the interchange between
optimization over functions T integral and the supremum (Rockafellar 1976 Theorem 3A). We also provide an independent proof of equivalence specializing Rockafellar's interchange theorem in Appendix [. Thanks to the following lemma, we may solve saddle point problem and obtain the OT map T from its solution ( ).
Lemma 4.1. Let T be the OT map from μ to v. Then, for every optimal potential v6'"
[ (x. (())) - (T(x)))) du(x))
T sug
(12)
T
We prove Lemma 4.1 in Appendix A.2 For general μr, the supT set for optimal el* might con- tain not only OT map T"' but other functions as well. Working with real-world data in experiments (.5.), we observe that despite this issue, optimization (II) still recovers T".
Relation previous works. The use of function T to approximate the c-transform was pro- posed by Nhan Dam et al. (2019) to estimate the Wasserstein loss in WGANs. For W2, the fact that T is an OT map was used by Makkuva et al. (2020); Korotin et al. (202a)) who primarily assumed continuous fr, and reduced (11) to convex d and T Vø for convex Φ. Issues with non- uniqueness of solution of were softened, but using ICNNS to parametrize v became necessary.
Korotin al. 2021b) demonstrated that ICNNs negatively affect practical performance of OT and tested an unconstrained formulation similar to (TT). As per the evaluation, provided the best em- pirical performance (Korotin al. 2021b 84.5). The method [MM:R] they consider parametrizes -11-112 vb-)) by a neural network, while we directly parametrize v-) by neural network ($4.3).
Recent work by Fan et al. 2021) exploits formulation similar 11 for general costs c(::)) While their formulation leads to a max-min scheme with general costs Fan et al.) 2021 Theorem 3), our approach gives rise to min-max method for quadratic cost. In particular, we extend the formulation to learn OT maps between distributions in spaces with unequal dimensions, see the next subsection.
4.2 UNEQUAL DIMENSIONS OF INPUT AND OUTPUT DISTRIBUTIONS
Consider the case when X Rf and y D have different dimensions, i.c., H * D. In order to map the probability distribution μ to straightforward solution is to embed X' to y via some Q: + y and then to fit the OT map between Q and v for the quadratic cost on y = PR_ In this case, the optimization objective becomes
inf sup { ((Q()) - 4u(x)+ ((o)dl())) T
(13)
5

---

Published as a conference paper at ICLR 2022
with the optimal 7* recovering the OT map from 1# /s to v. For equal dimensions H D and the identity embedding Q(x) = 2, expression (33 reduces to optimization
[11 up to a constant.
Instead of optimizing 13 over functions T Q(x) + V, propose to consider optimization
directly over generative mappings G : + +.
G) def inf sup U (1Q(z). - (G(())) du(x)) / (o))d(())
[
(14)
Our following lemma establishes connections between (14 and OT with unequal dimensions:
Lemma 4.2. Assume that exists unique OT plan between Q# n and v and it is deterministic, i.e., [idgo, + (Q # p.) Then GO (x) T (Q(z) is the OT mup between and v for the Q-embedded quadratic cost c( v) [1) yll². Moreover, for every optimal potential 8 of problem 14
G E srg / ((Q(z)))((()) v= (G(x))) dp(x).
(15)
G
We prove Lemma 4.2 in Appendix A.3 and schematically present its idea in Figure 4 Analogously to Lemma 4.1 it provides a way to compute the OT map G for the Q- embedded quadratic cost between distribu- tions lr and by solving the saddle point problem (14). Note the situation with non- uniqueness of arg supg is similar to $4.1
T*
Q#1
G*
Embedding
Output
Output Space
Input space
clx.y)=1|00x)y/1
μ
Relation to previous works. In practice, learning OT maps directly between spaces of unequal dimensions was considered in the work by (Fan et al. [02T] $5.2) but only on toy examples. We demonstrate that our method works well in large-scale generative modeling tasks ($5.1). Thcoretical properties of OT maps for embedded costs are studied, e.g., in (Pass 2010 McCann & Pass 2020
Input
Figure 4: The scheme of our approach for learning OT maps between unequal dimensions. the fig- ure, the setup of $5.1 is shown: l is a noise, Q is the bicubic upscaling, v a distribution images.
4.3 PRACTICAL. ASPECTS AND OPTIMIZATION PROCEDURE
To optimize functional (4). we approximate G R + RD and sb R - R with neural networks Ge. tha and optimize their parameters via stochastic gradient descent-ascent (SGDA) by using mini- batches from v. The practical optimization procedure is given in Algorithm | below. Following the usual practice in GANs, we add a small penalty (8I.3) on potential th. for better stability. The penalty is not included in Algorithm D keep it simple.
Bottou 2017) uses Wi as the loss to update the generator while we solve a diferent task - we fit the generator G to be the OT map for Q-embedded quadratic cost. Despite this, our Algorithm | has similarities with WGAN's training. The update of v (line 4) coincides with discriminator's update in WGAN. The update of generator G (line 8) differs from WGAN's update by the term - -QQ(,,, G+)-)). Besides, in WGAN the optimization is infe supp. We have inf supg, i.e., the generator in case is the solution of the inner problem.
Relation to previous works. WGAN Arjoovsk
4.4 ERROR ANALYSIS
Given a pair (v) G) approximately solving 14). natural question to ask is how good is the recov- ered OT map G. In this subsection, we provide bound on the difference between G* and Ō based on the duality gaps for solving outer and inner optimization problems.
In and, as the result. in 100.
it is enough to consider optimization over convex Our theorem below assumes the convexity of & although might not hold in practice since in practice & is neural network.
14
functions th, see (Villani 2008 Case 5.17).

---

Published as a conference paper at ICLR 2022
Algorithm 1: Learning the optimal transport map between unequal dimensions.
Input Input distribution on X' :R: output distribution // on y = IR D;
generator network Go IR. + IR.:; potential network t/ : RRD R;
number of iterations per network: Kg. K embedding Q A' - )/
Output: Trained generator Gg representing OT map from p to vi; repeat
for k2 1 K do
Draw batch X μ and
Lv t |テ| Eyer (y) 지& E (Go(z));
Update by using
&Ce aw
to minimize Lus
for ka to Kg do
Draw batch X ~ μú
LG - Xi Ezex [)(G(=)) Go())
Update 0 by using
&Cc 00
to minimize Lo;
until not converged:
return Ge
Theorem 4.3. Assume that there exists a unique deterministic OT plan for Q-embedded quadratic cost between μ and u, i.e. π² = lidæ", G* 41 for G* R + RDP. Assume that vb is 8-strongly convex (B > 0) and G R2 + P. Define
= [(], G) L(B. G)
and
= sup C(ù, G) inf sup L(t, G)
Then the following bound holds true for the OT map G* from μ
FID(G L2
f ((((() - G*)))) & glvā ve2)?
2 C2
(16)
where FID the Fréchet inception distance (Heusel et al. 2017) and L is the Lipschitz constant of the feature extractor of the pre-trained InceptionV3 neural network (Szegedy et al. (016).
We prove Theorem 4.3 in Appendix A.4 The duality gaps upper bound (²) norm between com- puted & and true G* maps, and the W2 between true and generated (fake) distribution G Consequently, they upper bound FID between data and fake (generated) & distributions.
Relation to previous works. Makkuva et al. (2020) Korotin al. (202 la) prove related bounds for W2 with μ, located the spaces of the same dimension. Our result holds for different dimensions.
5 EXPERIMENTS
We evaluate our algorithm in generative modeling of the data distribution from noise (3..1 and unpaired image restoration task (5(5.2). Technical details are given Appendix B Additionally, in Appendix B.4 we test our method on toy 2D datasets and evaluate it on the Wasserstein-2 benchmark
(Korotin al. 2021b) in Appendix B.2 The code is in the supplementary material.
5.1
MODELING DATA DISTRIBUTION FROM NOISE DISTRIBUTION
In this subsection, y is 192-dimensional normal noise and the high-dimensional data distribution.
Let the images from be of size w x with channels. As the embedding Q X + y we use naive upscaling of noise. For R represent it as 3-channel 8 image and bicubically upscale it to the size w x h of data images from v. For grayscale images drawn from v, we stack c copies over channel dimension.
We test our method on MNIST 32 32 LeCun et al. 1998), CIFAR1O 32 x 32 (Krizhvssy et al.) 2009), and CelebA 64 64 (Liu et al. 2015) image datasets. In Figure 5 show random samples
7

---

Published as a conference paper at ICLR 2022
Y- Gon o00
Gon o00
109 00
2
8
5
5
4
2
1-1
(a) MNIST, 32 x 32, grayscale
(b) CIFARIO, 32 x 32, RGB
(c) CelebA, 64 x 64, RGB
Figure 5: Randomly generated MNIST, CIFAR1O, and CelebA samples by our method (OTM).
Table 1: Results on CIFARIO dataset.
Table 2: Results on CelebA dataset.
Related Work
Inception
Model
Related Work
FID
NVAE
Vahdal & Kaut
2020
51.71
Radioed
EBM
PixelIQN
Ostrovsia et al
2018
6.02
5.29
40.58
49.46
DRAGAN
Kodali of al
2017
42.3
NCSN
NCP-VAE
DCGAN
Song Ermon
Kadiord et al
et al 2021
2116
6.64-10.14
8.8740.12
37.70
24.08
25.32
NCP-VAB
NVAR
BEGAN
Ancja et 202
Beniei
Vahdat Kanfz
2020
5.2
38.9
13.4
WGAN-GP
WGAN
3P-WGAN
Guirajant
Nian Dian et
Arjovsky of
2017
2019
6.49:t0.09
7.38 +: 0.08
39.40
28.8
55.2
WGAN-QC
WGAN-GP
WGAN
AE-OT
Gulrajani et 2017
Arjosky et al
Liu ct al
2019
21
28.6
30.0
12.9
AF.OT-GAN
AE-OT
el al.
20200
2012k
28.5
17.1
AE-OT-AN
An al.
20.0m
7,8
OTM
Ours
7.42:10.06
21.78
OIM
Ours
6.5
generated by our approach, namely Optimal Transport Modeling (OTM). To quantify the results,
in Tables [ [nd we give the inception (Salimans et al. 2016) and FID (Heusel et al.] 2017) scores of generated samples. Similar to (Song & Ermon] 2019 Appendix B.2), we compute them 50K real and generated samples. Additionally, in Appendix [B.4 we test our method on 128 x 128 CelebA
faces. We provide qualitative results (images of generated faces) in Figure 11
For comparison, we include the scores existing generative models three types: (1) OT map as the generative model; (2) OT cost the loss; (3) not OT-based. Note that models of the first type compute OT in the latent space of an autoencoder in contrast to our approach. According to our evaluation, the performance of our method is better comparable to existing alternatives.
5.2 UNPAIRED IMAGE RESTORATION
In this subsection, we consider unpaired image restoration tasks on CelebA faces dataset. In this case, the input distribution y consists of degraded images, while v are clean images. In all the cases, embedding Q is straightforward identity embedding (Ix) x.
In image restoration, optimality of the restoration map desired since the output (restored) image expected to be close to the input (degraded) one minimizing the transport cost. Note that GANs do not seek for an optimal mapping. However, practice, due to implicit inductive biases such as convolutional architectures, GANs still tend to fit low transport cost maps (Bézenac et al., 2021).
The experimental setup is shown in Figure We split the dataset in parts A, B, C containing 90K, 90K, 22K samples respectively. To each image apply the degradation transform (decolorization, noising or occlusion) and obtain the degraded dataset containing of respective parts A, B, C. For unpaired training part A of degraded and part B of clean images. For testing,
parts C.
Train
Test
Clean
images
C
Degraded
images
B
C
Train
Test
Figure 6: The training/testing scheme that we use for unpaired restoration tasks.
To quantify the results we compute FID of restored images w.r.t. clean images of part C. The scores for denoising, inpainting and colorization given in Table 3 details of each experiment and qual- itative results are given below.
As baseline, we include WGAN-GP, For fair comparison, fit it using exactly the same hyper- parameters as in our method OTM-GP. This is possible due to the similarities between our method and WGAN-GP's training procedure, discussion in $4.3 In OTM, there is no GP (B3.).
8

---

Published as a conference paper at ICLR 2022
Model
Denoising
Colorization
Inpainting
Input
166.59
32.12
47.65
WGAN-GP
25.49
7.75
16.51
OTM-GP (ours)
10.95
5.66
9.96
OTM (ours) 5.92
5.65
8.13
Table 3: FID on test part C in image restoration experiments.
Denoising. To create noisy images, add white normal noise with 0.3 to each pixel. Fig- ure 7 illustrates image denoising using our OTM approach on the test part of the dataset. We show additional qualitative results for varying in Figure 15 S(B.4).
(a) Noisy samples. Figure 7: OTM for image denoising test C part of CelebA, 64 x 64.
(b) Pushforward samples.
(c) Original samples.
Colorization. To create grayscale images, average the RGB values of each pixel. Figure illustrates image colorization using OTM on the test part of the dataset.
(a) Grayscale samples.
(b) Pushforward samples.
(c) Original samples.
Figure 8: OTM for image colorization on test C part of CelebA, 64 x 64.
Inpainting. To create incomplete images, we replace the right half of each clean image with zeros. Figure illustrates image inpainting using OTM on the test part of the dataset.
3
(a) Occluded samples.
(c) Original samples. Figure 9: OTM for image inpainting test part of CelebA, 64 x 64.
(b) Pushforward samples.
6 CONCLUSION
Our method fits OT maps for the embedded quadratic transport cost between probability distribu- tions. Unlike predecessors, scales well to high dimensions producing applications of OT maps directly in ambient spaces, such spaces of images. The performance is comparable to other exist- ing generative models while the complexity of training is similar to that of popular WGANs.
assume the existence of the OT map between them. In practice, this might not hold for all real-world μ,, Working with equal dimensions, we focus on the quadratic ground cost 111 2||². Nevertheless, our approach extends other costs c(-,),, see Fan et al. (202T). When the dimensions are unequal, we restrict our analysis to embedded quadratic cost : QT)) y/ where Q equalizes dimensions. Choosing the embedding Q might not be straightforward in some practical problems, but our evaluation (.51) shows that even naive choices of Q work well.
Limitations. For distributions (a,
Potential impact and ethics. Real-world image restoration problems often do not have paired datasets limiting the application of supervised techniques. In these practical unpaired learning prob- lems, we expect our optimal transport approach to improve the performance of the existing models. However, biases data might lead to biases the pushforward samples. This should be taken into account when using our method in practical problems.

---

Published as a conference paper at ICLR 2022
Reproducibility. The PyTorch source code is provided at
The instructions to use the code are included in the README. md file.
7 ACKNOWLEDGMENT
This research supported by the computational resources provided by Space Applica- tions Centre (SAC), ISRO. The first author acknowledges the funding by HRD Grant No. 0303750FA700/SA//SROO Skoltech RAIC center was supported by the RF Government (subsidy agreement 0000000733332155000000 Grant No. 70-2021-00145 02. 01.2222...
REFERENCES
Matthew Amodio and Smita Krishnaswamy. Travelgan: Image-to-image translation by transfor- mation vector learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8983-8992, 2019.
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pp. 146-155, PMLR, 2017.
Dongsheng An, Yang Guo, Na Lei, Zhongxuan Luo, Shing-Tung Yau, and Xianfeng Gu. Ac-ot: A new generative model based on extended semi-discrete optimal transport. In International Conference on Learning Representations, 2020a. URL htttps::/ net/
forum?1d=#kIdyINNe
Dongsheng An, Yang Guo, Min Zhang, Xin Qi, Na Lei, and Xianfang Gu. Ae-ot-gan: Training gans from data specific latent distribution. In European Conference on Computer Vision, pp. 548-564. Springer, 2020b.
Jyoti Aneja, Alexander Schwing, Jan Kautz, and Arash Vahdat, Nep-vae: Variational autoencoders with noise contrastive priors. In Advances in Neural Information Processing Systems Conference, 2021.
Cem Anil, James Lucas, and Roger Grosse. Sorting out lipschitz function approximation. In International Conference on Machine Learning, pp. 291-301. PMLR, 2019.
Martin Arjovsky and Léon Bottou. Towards principled methods for training generative adversarial networks. arXiv preprint arXiv: 770100462, 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214 22432.3.722 PMLR, 2017.
David Berthelot, Thomas Schumm, and Luke Metz. Began: Boundary equilibrium generative ad- versarial networks. arXiv preprint arXiv: 703. 10717, 2017.
Emmanuel de Bézenac, Ibrahim Ayed, and Patrick Gallinari. Cyclegan through the lens of (dy- namical) optimal transport. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 132-147. Springer, 2021.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4)3775-477, 1991.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2019. URL
.ne/ htr:://rummmmm www.mmidd-ixssiiiii.
Biwei Dai and Uros Seljak. Sliced iterative normalizing flows. In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2021.
10

---

Published as a conference paper at ICLR 2022
Ishan Deshpande, Ziyu Zhang, and Alexander G Schwing. Generative modeling using the sliced wasserstein distance. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3483-34949 2018.
DC Dowson and BV Landau. The fréchet distance between multivariate normal distributions. Journal of multivariate analysis. 12(3):450-455 1982.
Yilun Du and Igor Mordatch. Implicit generation and generalization in based models. arXiv preprint 111:::19000066000 2019.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Yongxin Chen, and Haomin Zhou. Scalable computation of monge maps with general costs. arXiv preprint arXiv:2106.038/ 03.812 2021.
Kilian Fatras, Younes Zine, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein: asymptotic and gradient properties. arXiv preprint arXiv:191004009,,
2019.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn di- vergences. In International Conference Artificial Intelligence and Statistics, pp. 1608-1617. PMLR, 2018.
lan J Goodfellow, Jean Pouget-Abadis Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C Courville, and Yoshua Bengio. Generative adversarial nets. Advances in Neural Information Processing Systems Conference, 2674-2680, 2014.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Im- proved training of wasserstein gans. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 5769-5779, 2017.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Leygonie Jacob, Jennifer She, Amjad Almahairi, Sai Rajeswar, and Aaron Courville. W2gan: Re- covering an optimal transport map with a gan. 2018.
Leonid Vitalevich Kantorovich. On problem of monge. Uspekhi Mat. Nauk, pp. 2252222-222 225-226, 1948.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4401-4410, 2019.
Cheolhyeong Kim, Seungtae Park, and Hyung Ju Hwang. Local stability of wasserstein gans with abstract gradient penalty. IEEE Transactions Neural Networks and Learning Systems. 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1112.6980, 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.6114, 2013.
arXiv preprint
Naveen Kodali, Jacob Abernethy, James Hays, and Zsolt Kira. On convergence and stability of gans. arXiv preprint 1wX.v:1005.0222 2017.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations, EXsa
2021a. URL BtRL:tttst////eeeeeiiii
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. arXiv preprint arXiv:2106.01954, 2021b.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
11

---

Published as a conference paper at ICLR 2022
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2) 2324,, 1998.
Na Lei, Kehua Su, Li Cui, Shing-Tung Yau, and Xianfeng David Gu. A geometric view of optimal transportation and generative model. Computer Aided Geometric Design, 68:1-21, 2019.
Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein gan with quadratic transport cost. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), October 2019.
Shu Liu, Shaojun Ma, Yongxin Chen, Hongyuan Zha, and Haomin Zhou. Learning high dimensional wasserstein geodesics. arXiv preprint arXiv:2102.02992, 2021.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference computer vision, pp. 3730-3738, 2015.
Guansong Lu, Zhiming Zhou, Yuxuan Song, Kan Ren, and Yong Yu. Guiding the one-to-one map- ping in cyclegan via optimal transport. In Proceedings the AAAI Conference on Artificial Intelligence, volume 33, pp. 4432 4439, 2019.
Guansong Lu, Zhiming Zhou, Jian Shen, Cheng Chen, Weinan Zhang, and Yong Yu. Large-scale op- timal transport via adversarial training with cycle-consistency. arXiv preprint arXiv:2003.06635, 2020.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672- 6681. PMLR, 2020.
Anton Mallasto, Jes Frellsen, Wouter Boomsma, and Aasa Feragen. (q. p)- wasserstein gans: Com- paring ground metrics for wasserstein gans. arXiv preprint 1wX.iii:102..06444 2019.
Xudong Mao, Qing Li, Haoran Xie, Raymond YK Lau, Zhen Wang. and Stephen Paul Smol- ley. Least squares generative adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2794-2802, 2017.
Robert McCann and Brendan Pass. Optimal transportation between unequal dimensions. Archive for Rational Mechanics and Analysis, 238()):475-1520 2020.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018.
Quan Hoang Nhan Dam, Trung Le, Tu Dinh Nguyen, Hung Bui, and Dinh Phung. Threeplayer wasserstein gan via amortised duality. In Proc. of the 28th Int. Joint Conf. on Artificial Intelligence (IJCAI), 2019.
Georg Ostrovski, Will Dabney, and Rémi Munos. Autoregressive quantile networks for generative modeling. In International Conference Machine Learning, pp. 3936-3945. PMLR, 2018.
Brendan Pass. Regularity of optimal transportation between spaces with different dimensions. arXiv preprint arXiv: 10..0008..444 2010.
Giorgio Patrini, Rianne van den Berg, Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In Uncertainty in Artificial Intelligence, pp. 733-743. PMLR, 2020.
Henning Petzka, Asja Fischer, and Denis Lukovnikov. On the regularization of wasserstein gans. In International Conference Learning Representations, 2018.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. In 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL wtt.::ttp///aaaeo org/ wtgggba/////// aabs/1111. 06434
Tyrrell Rockafellar. Integral functionals, normal integrands and measurable selections. In Nonlinear operators and the calculus of variations, pp. 157-207. Springer, 1976.
12

---

Published as a conference paper at ICLR 2022
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems. 29: 2234-2242, 2016.
Maziar Sanjabi, Meisam Razaviyayn, Jimmy Ba, and Jason D Lee. On the convergence and bustness of training gans with regularized optimal transport. Advances in Neural Information Processing Systems, 2018:7091-7101, 2018.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Conference Learning Representations, 2018
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Proceedings of the 33rd Annual Conference on Neural Information Processing Systems, 2019.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818-881220 8826, 2016.
Amirhossein Taghvaei and Amin Jalali. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv preprint 1xxiv:1902.07111 2019.
Ugo Tanielian and Gerard Biau. Approximating lipschitz continuous functions with groupsort neu- ral networks. International Conference on Artificial Intelligence and Statistics, pp. 442-450.
PMLR, 2021.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. arXiv preprint 1.XXvv1711.015555 2017.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. In Advances in Neural Information Processing Systems Conference, 2020.
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. International Conference on Machine Learning, pp. 6882-6892. PMLR, 2019.
Jun Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image trans- lation using cycle-consistent adversarial networks. In Computer Vision (ICCV), 2017 IEEE International Conference on, 2017.
13

---

Published as a conference paper at ICLR 2022
A PROOFS
A.1 PROOF OF EQUIVALENCE: EQUATION 10 AND
Proof. Pick any T X - y. For every point X by the definition of the supremum have
(a,T(x)) - 4 (T(x)) sup [(x, v) v(4)).
vey
Integrating the expression w.r.t x ~ μ yields
J - (T(=) & / sup (a. (2)) du(x)) = LT.
yey
Since the inequality holds for all T + V, we conclude that
C2 = sup / (<,T(()) (T)))))du(x) - / T:X-V
sup (( v) - (v)) dp(x) = Lı. yE)
ie. L2 S L1. Now let us prove that the sup on the left side actually equals L1. To do this, we need show that for every e there exists 7 X' + y satisfying
/ ((x + (T (z))ddp(()) 2 -I-
First note that for every X by the definition of the supremum there exists yf (x) which provides
(x,y + (v)())) sup yEy
(2,,)) ((y)} c
We take (xx (x) for all x E X and integrate the previous inequality w.r.t. x ~ μ.. We obtain
/ - (T' (T 2
/
sup (10,8) v(1)) du(x) - L1
xx
which is the desired inequality.
A.2 PROOF OF LEMMA 4.1
Proof. It is enough to prove that (x) = (T), x) - 30 (T(x)) holds p-almost everywhere, i.c.,
T (x) arg sup ((x, v) - 40 (v)). Since v uu, we use b with &b + *0 to derive yERo
Wz(1,v) - / 1 =11113 / (1ulPaduo) =
2n (x)dp(x)) - -/ V (y))dd(v)) / 1b ()())d((z) I 10"
(x) 2i (I (d)))= S- / (T (x),
(17)
(x)x)
/ 11110--- TT) / / [1x1[[2pp(aa / (TT
w (H, / 1 - / 1 [.lliPPduud))
As result, inequality 117 becomes the equality, in particular, ()(x) + +0 (T) = (x))))
holds y-almost everywhere.
14

---

Published as a conference paper at ICLR 2022
A.3 PROOF OF LEMMA 4.2
Let QW2 denote the Q-embedded quadratic cost. We use the change of variables formula to derive
Proof.
J 1 (1Q(())- (z)))
QW2(A,))
inf nEll(N) Xxx /
T'El(upp.)) xxy 2 || dn (x, v) @²
inf
(18)
ie., computing the OT plan for QW v) boils down to computing the OT plan for ( μi, v). follows that fidga, T (Q(x)) (1 [idg", # is an OT plan for QW (1. v), and G* the
OT map. Inclusion (15) now follows from Lemma
A.4 PROOF OF THEOREM 4.3
Proof. Pick any G' argsupe [(d, G) argsupe fx (Q(z), G(z)) -(z)) } dp(x) or, equivalently, for æ e RH (R""G(()) E arg sup,
{ (((()).)) (() Consequently, for all y e RD
(Q()), (
{y))
which after regrouping the terms yields
v(v) 2 (C' + (Q(x),
This means that Q(x) is contained in the subgradient ai (' for a convex d. Since & is G-strongly convex, for points G(z)) Gx(x) E RD and Q(x) E 06 (x))) we derive
(G(x)) 2 (G'(x)) + (Q(x), G(x) G'(x)) + (G'(x) - G(x)]| 2
Regrouping the terms, this gives
((G()))] [(Q(z), (('()) (Q4) &(C()) 8 [ G(=)
Integrating w.r.t. x μ yields
xx (()(()))(x)][ddb)))) 3 || G'llia(a) (19)
= C( G*) - G) 3
Let G be the OT map from
We use G = v to derive
C) / - dp(x) J, M())d)())
-
1 (1)))) du(x)) (()))=))d(())=
(((x), -(C)) du(z))
x
/ (())) + |I' G)((-d((2).
1
(20)
Let 1/* be optimal potential in 14
Thanks Lemma 4.2 we have
inl sup L(v, G) f(
f (10(s), - 80 (a (x)) du(x) + / 10' (((()de())=
15

---

Published as a conference paper at ICLR 2022
/ - du(z)) / 3 (G(x))d)(===
J (((x),
(21)
By combining 20 with 21p. we obtain
1 11 Gr -
B
E2 = L(u, G*) C(v G 2
---II' -'G'-G"llenn (2
(22)
The right-hand inequality of (6) follows from the triangle inequality combined with (19) and 22
The middle inequality of (16) follows from (Krottin et al. 202 Lemma A.2) and Giyl
Now we prove the left-hand inequality of (6]. Let I be the feature extractor of the pre- traierd Inception V3 neural networks. FID score between generated (fake) G# and data distribution v is
FID = FD #V. I+V) S 2 2.W, (I G4 I+v),
(23)
where FD(:,) is the Fréchet distance which lower bounds 2 W22 see DDwsson & Landau 1982
Finally, from (Korotin et al. 2021a Lemma A.I) it follows that
(I G Iw s I.2 WW (2
(24)
Here L is the Lipschitz constant of We combine (23) and (24) to get the left-hand inequality in (1)).
B EXPERIMENTAL DETAILS
We use the PyTorch framework. All the experiments are conducted on 2x V100 GPUs. We compute inception and FID scores with the official implementation from OpenA and results are taken from the respective papers or publicly available source codes.
and TTUR? The compared
B. GENERAL TRAINING DETAILS
MNIST (LeCun et al., 1998). On MNIST, we use 192 and y E 32 The batch size 64, learning rate 2- 10 optimizer Adam (Kingma & Ba 2014 with betas (0, 0.9), gradient optimality coefficient À = 10, and the number training epochs
30. We observe stable training while updating & once in multiple G updates, i.e., kg and l 1.
CIFAR10 Krizhevsky et al. 2009). We use all 50000 samples while training. The latent vector R and y E R22232233 batch size 64, X 10, he 1, kg with betas (0, 0.9). and learning rate 2 10 for G and 1 10 for th.
1, T 1000, Adam optimizer
CelebA (Liu et al. 2015). We use E 192 and y R61x61x3 The images are first cropped the center with size 140 and then resized to 64 64. We consider all 202599 samples. We use Adam with betas (0, 0.9). T 200, K = 2, K = 1 and learning rate 2 10
Image restoration. In the unpaired image restoration experiments, we use Adam optimizer with betas (0, 0.9), KG = 5, K2
0, learning rate 1 10-4 and train for T 300 epochs.
CelebA128x128 (Liu et al., 2015). On this dataset, we resize the cropped images as in CelebA to 128 128, i.c. (0.as=(0.5,, 0.999). The batch size is reduced to 16
R 128x 28x3 Here. KG = 3, K = 1, X = 0.01, learning rate 1 10 1 and to fit in the GPU memory.
Anime128x125 This dataset consists of 500000 high resolution images. We resize the cropped images as in CelebA 128 128, i.e. y R 128x 1283 Here, Ke 5, K2 1, X 0.01, learning rate 2 10-1, batch size 16, and (0,
Toy datasets. The dimension is D 10000. We use the batch size 400, X
H
2, total number of samples is 0.1. K = 1, Kg 16. and
=
'IS: http:p//tth..ccceeeeee/////t/aet "FID: Anime: http
score
16

---

Published as a conference paper at ICLR 2022
T = 100. The optimizer is Adam with betas (0.5, 0.99) and learning rate 10-3. 3 We use the following datasets: Gaussian to mixture of Gaussian D two moons (sklearn. datasets (sklearn.datasets..tm www. saake_moons), circles (sklearn, datasets (aklearn.datasets.u.. make circles), gaussian to S-curve (sklearn. datasets. peatttaetssmaak___.ee. www.emakke_aaaeeeeee currve,, and gaussian to swiss roll (sklearn. waw.setaatsss..aao.m.mmmm
W benchmark (Appendix B.2).
The dimension is D H 64 x 64 x 3. We use batch size 64, À 0, K 1, K learning rate 10-4, and Adam optimizer with default betas.
B.2 EVALUATION ON THE CONTINUOUS WASSERSTEIN BENCIIARK
To empirically show that the method recovers the optimal transport maps well on equal dimensions, we evaluate it on the recent continuous Wasserstein-2 benchmark by Korotin et al. (2021b). The benchmark provides a number of artificial test pairs (11, v) of continuous probability distributions with analytically known OT map 7 between them.
we use the "Early" images benchmark pair (Korotin et al.) 2021b $4.1) for details. We adopt the L'-unexplained percentage metric (Koottin et al. 2021a 85.1) to quantify the recovered OT map T: For our method the f2-UVP metrie is only 1%, see Table 4 This is comparable to the best [MM:R] method which the authors evaluate on their benchmark. The qualitative results are given Figure I0
For evaluation, (D 12288),
Method
UV-U
[MMRR 1.4%
OTM (ours)
1.32%
CU-UV)( = 100 | - T T²////Vr(vv)
Table 4: 23-UVP metric of the recovered transport map on the "Early" images benchmark pair.
X-1
fix
TXK)-v
Figure 10: Qualitative results of OTM on the "Early" images benehmark pair (μ, v) by Korotin et al. 2021I). The 1st line shows samples x ~ μ. the 2nd line shows fitted OT map (z)), and the 3rd line shows the corresponding optimal map T (x)
B.3 FURTHER DISCUSSION AND EVALUATION
Generative modeling. In the experiments, we use the gradient penalty on & for better stability of optimization. The penalty is intended to make the gradient norm of the optimal WGAN critic equal to (Gulrajani et al., 2017] Corollary 1). This condition does not necessarily hold for optimal ai" in our case and consequently might introduce bias optimization.
To address this issue, we additionally tested an alternative regularization which we call the gradient optimality. For every optimal potential v/" and map G* of problem 14), we get from Lemma 4.2
Va Bewn (Q((),, (x)))) 10" (G (G*(I)]I Evw [Q(x)) -Eaw V (z)))) = 0. (25)
the above expression sim- plifies to Exep Vaji" (G* (x))) = O. Based on this property, we establish the following regularizer Al|Ez~m V (G(z))) for À > 0 and add this to Ey our Algorithm
Since l is normal noise distribution and Q(x) is naive upscaling ($5.1)
While gradient penalty considers expectation of norm, gradient optimality considers norm of expec- tation. The gradient optimality is always non-negative and vanishes the optimal point.
https: thuithtubbooo Ami TTag ICNN
17

---

Published as a conference paper at ICLR 2022
X
FID4
We conduct additional experiments with the gradient optimality and compare FID scores for different À in T It leads to an improve- ment of FID score from the earlier 7.7 with the gradient penalty to the current 6.5 with the gradient optimality on CelebA (Ta
0.001
16.91
0.01
16.22
0.1
16.70
1.0
10.01
10
6.50
Unpaired restoration. In the unpaired restoration experiments ($5.2), we test OTM with the gradient penalty to make a fair com- parison with the baseline WGAN-GP. We find OTM without regular- ization, i.e., À = 0 works better than OTM-GP (Table 3). In practice, more G updates for single update works fairly well (BB.I.
Table 5: Ablation study of
gradient optimality in OTM.
B.4 ADDITIONAL QUALITATIVE RESULTS
OTM works with both the grayscale and color embeddings of noise in the ambient space.
CelebA 128x128. Figure 11 shows the grayscale embedding (zz), the recovered transport map ((r), and independently drawn real samples y
Q(X)
G(X)
Y-v V
Figure 11: OTM between 128-dimensional noise and CelebA, 128 128. The 1st line shows the grayscale embedding Q (repeating bicubic upscaling of a noise, 16 x 8), the 2nd line shows corresponding generated samples, and the 3rd line shows random samples from the dataset.
Anime128x128, Figure shows the color embedding Q(z), the recovered transport map G(x),
and independently drawn real samples
Q(X)
G(X)
Figure 12: OTM between 192-dimensional noise and Anime, 128 x 128. The 1st line shows the color embedding Q (bicubic upscaling of a noise, x 8), the 2nd line shows corresponding generated samples, and the 3rd line shows random samples from the dataset.
The extended qualitative results with color embedding MNIST, CIFAR10, and CelebA shown in Figure 13a Figure [36 and Figure [3c respectively. Table shows quantiative results MNIST. The color embedding Q is bicubic upscaling of noise in l 3x888 The samples
18

---

Published as a conference paper at ICLR 2022
Table 6: Results on MNIST dataset.
Model
Related Work
FID
VAE
Kingra Welling
ISGAN
Mão et ml
7.8:0.66
BEGAN
Benthelot 22
13.1:110
WGAN
Arjovsky
2017
6.7.0.44
SIG
Dar & Seqpk
2021
4.5
AE-OT
20208
6.2=0.2
AE-OFGAAN
20208
3.2
OTM
Ours
2,4
generated randomly (uncurated) by fitted optimal transport maps between noise and ambient space, e.g., spaces of high-dimensional images. Figure illustratss latent space interpolation between the generated samples. [igure[5]shoon denoising images with varying levels of 0.1, 0.2, 0.3, 0.4 by the model trained with 0.3.
3
9
6
3
4
3
7
7
0
1
3
5
8
9
1 0 7 0 7 8
d
0
4
1
3
3
0 7 2
9
6
5 8 8 -
7
_
8
3
2
3
0 4
y 0
3
7
/
5
/ /
6
6
/
O
6
4
7
7
9
6
(a) MNIST, 32 x 32, grayscale
(b) CIFARIO, 32 x 32, RGB
(c) CelebA, 64 x 64, RGB
Figure 13: Randomly generated MNIST, CIFARIO, and CelebA samples by our method (OTM).
Figure 14: OTM for latent space interpolation on CelebA, 64. Extended samples.
Toy datasets. Figure 16 shows the results of our method and related approaches ($ on toy 2D dataset. Figure 17 shows the results of our method applied to other toy datasets.
B.5 NEURAL NETWORK ARCHITECTURES
This section contains architectures on CIFARIO (Table 7 CelebA (Table CelebA 128 128
(Figure II] generation, image restoration tasks (Table evaluation on the toy 2D datasets and Wasserstein-2 images benchmark (Table
In the unpaired restoration tasks (:5.2), we use UNet architecture for transport map and convolu- tional architecture for potential 15. Similarly to (Song & Ermon 2019), Normalaotoooooiioo (BN) and DastaneeNormalizattinn (INorm+) layers. In the ResNet architectures, we use the Resid- ualBlock of NCSN (Song & Ermon 2019).
In the toy 2D examples, we use a simple multi-layer perceptron with hidden layers consisting of 128 neurons each and LeakyReI.U activation. The final layer is linear without any activation.
19

---

Published as a conference paper at ICLR 2022
0.1 0.2 -= 0.3
0.4
0.1
0.2
0.3
0.4
(a) Noisy samples.
(b) Pushforward samples.
(c) Clean.
Figure 15: OTM for image denoising for varying levels of noise test C part of CelebA, 64 x 64.
10
10
-5
-10
-10
10
-10
-5
10
-10
-5
10
(a) LSOT (Seguy et al.
2018
(b) W2GAN Jacob al. 2018
(c) WGAN-LP Petzka et al.
2018
10
10
-5
-10
-10
-10
10
-10
-5
5
10
(d) ICNN-CT (aakkuva et al.
2020
(c) W2GN (Korotin et al. 2021a)
(I) OTM (ours)
Figure 16: Mapping between a Gaussian and a Mixture of 8 Gaussians in 2D by various methods.
The colors green, blue, and peru represent input, pushforward, and output samples respectively.
The transport map G and potential & architectures on MNIST 32 x 32, CelebA 128 x 128, and Anime 128 x 128 are the generator and discriminator architectures of WGAN-QC Liu et al. (2019) respectively.
In the evaluation on the Wasserstein-2 benchmark, we use publicly available Une architecture for transport map T and WGAN-QU discriminator's architecture for / (Liu et al. 2019). These neural network architectures are the same as the authors of the benchmark use.
20

---

Published as a conference paper at ICLR 2022
1.0
0.8
0.4
0.2
0.0
0.0
-0.
-0.4
-1.0
2
-1.0
-0.5
0.0
1.0
(a) Two moons
(b) Circles
100
6
4
2
2
-2
-4
-4
-10
-5
10
-5.0
225
0.0
2.5
5.0
7.5
(c) Curve
(d) Swiss Roll
Figure 17: OTM on toy datasets, D = 2. Here, the colors green, blue, and peru represent input, pushforward, and output samples respectively.
Table 7: Architectures for generation task on CIFAR10, 10, 32 x 32.
GO.)
Noise: x E R 128
Linear, Reshape, output shape: [128
4]
ResidualBlock Up. output shape: [128 x 8 x
ResidualBlock Up, output shape: [18 16 16
ResidualBlock Up, output shape: [128 x 32 x 32]
Conv, Tanh, output shape: [3 32 32]
Target: y E I. 32332232
ResidualBlock Down, output shape: [128 x 16 16)
ResidualBlock Down, output shape: [128 x 8)
ResidualB lock, output shape: 128 8 x 8)
ResidualBlock, output shape: [128 x x 8)
ReLU, Global sum pooling, output shape: [128 x x
Linear, output shape: [1]
21

---

Published as a conference paper at ICLR 2022
Table 8: Architectures for generation task on Celeba, 64 64.
G(.)
Noise: x E R 128
KonyTraanspose BN, LeakyReLU, output shape: [256 X 1
ConvTraasse Transpose, BN, LeakyReI.U, output shape: [512 x 4 4
Cony, PixelShuffle, BN, LeakyReL.U, output shape: [512 x 8
Conv, PixelShuffle, BN, LeakyReLU, output shape: [512 x 16 16]
Conv, PixelShuffle, BN, LeakyReLU, output shape: [512 32 32]
ConvTraassa Transpose, Tanh, output shape: [3 x 64 x 64]
(-)))
Target: y E IR 3X64X64
Conv, output shape: [128 x x 64
ResidualBlock Down, output shape: [256 x 32 32]
ResidualBlock Down, output shape: 2256 x 16 16]
ResidualBlock Down, output shape: [256 x 8 8
ResidualBlock Down, output shape: [28 x
Cony, output shape:
Table Architectures for restoration tasks on CelebA, 64 x 64.
Input: R 33x4
Conv, BN, LeakyReL.U, output shape: [256 x 64 x 64
Conv, LeakyReLUU, AvgPool, output shape: [256 32 32) x1
Conv, LeakyReL.U, AvgPool, output shape: 256 x 16 x 16) x2
Conv, LeakyReeu,,.. AvgPool, output shape: [256 8 x 8). x3
Cony, LeakyReLU, AvgPool, output shape: 256 x 4 x 4].x4
Nearest Neighbour Upsample, Conv, BN, ReLU, output shape: [256 x 8), y3
Add (y3, x3), output shape: [256 x 8 8). y3
Nearest Neighbour Upsample, Conv, BN, ReLU, output shape: [256 x 16 16,, y2
Add (y2, x2), output shape: [256 16 16), y2
Nearest Neighbour Upsample, Cony, BN, ReLU, output shape: [256 x 32 x 32), yl
Add (y1, x1), output shape: [256 32 32), yl
Nearest Neighbour Upsample, Conv, BN, ReL.U, output shape: [256 x 64 x 64).
Add (y, x), output shape: 2256 64 64].y y
Conv TonTTsanp, Tanh, output shape: [3 x 64 x 64]
((())
Target: y E IR 1xxxxx0
Conv, LeakyReLU, AvgPool. output shape: [256 x 32 x 32
Cony, LeakyRel.U, AvgPool, output shape:
256 x 16 x 16
Conv, LeakyRel.U,, AvgPool, output shape: [256
8)
Conv, Lea,yReLU, AvgPool, output shape: [256 X x 4
Linear, output shape: [1]
22


Reviews:
[{"review_id": "SOfipqV0a4", "reviewer": "Reviewer_vEgs", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "Pros:\n- The writing is clear and easy to follow.\n- The experimental results are solid and convincing.\n\nCons:\n- I am not sure if equation (10) and equation (11) is really equivalent to each other. To apply the Brenier theorem, namely replace $y$ with $T(x)$ in equation (10), we need to add the measure preserving constraint $T_\\#\\mu = \\nu$ in equation (11). Without this constraint, which requires $T$ to be a measure preserving map, equation (11) may goes to $\\infty$. For example, $T(x)\\equiv 0$ and $\\psi(y)=-\\infty$. The optimal transport map should be included in the feasible set of equation (11) if it is large enough, but without the measure preserving constraint, it's hard to prove that the learned map $T$ is measure preserving. Therefore, the authors need to make some effort to show the equivalence. \n\n- Furthermore, to make the conclusion more convincing, I also recommend the authors to experiment on more 2-dimensional toy sets. If the learned map $T$ is an optimal transport map, the samples of the Gaussian distribution should be evenly transported to the Mixture of 8 Gaussians in Fig. 14(f). But it is really difficult to find the 'even' transportation (especially for the southeast samples), which weakens the conclusion that the learned is an optimal transport map.", "summary_of_the_paper": "The paper proposes to learn the optimal transport map from the latent distribution to the data distribution directly by optimizing the W2 distance. To find the OT map, the authors replace $y$ with $T(x)$ in the Kantorovich dual problem.\nExperiments show that the proposed method works well and learns the image distribution successfully.\n\n", "summary_of_the_review": "My main concern comes from the deduction from equation (10) to equation (11). The authors may add more explanation to show the equivalence.", "strength_and_weaknesses": "", "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": "Pwe3MtchbSO", "reviewer": "Reviewer_qNeK", "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": "This work addresses an important problem and proposes an interesting solution using the OT map between input and output distributions.\nExperimental results show that the method performs similar to other existing methods.\n\n**Paper organization/presentation**\n\n- We need to wait till the middle of Section 4 to understand what the paper proposes. The introduction to the proposed idea is too long (even if  some of the concepts need to be provided,  a lot of that content can probably go to an Appendix).\n\n- Lemma 4.1 seems the relevant one. Maybe this could be better introduced in the paper (there's no  connection to the previous paragraphs, so it's hard to follow the story). \n\n- The paragraph following Lemma 4.1, that tries to connect the Lemma to the previous work is not particularly clear. The ultimate question is: have people tried this idea of estimating an OT map before? and if so, how is the proposed solution different from previous work (technical difference, algorithmic difference,....). \n\n- The technical extension to unequal dimensions seems straightforward if you have Q.\n\n- The paper discusses in multiple places the relation to previous work. This makes it hard to follow (the presentation is a little not linear, one needs to jump from one place to the other to get all the pieces). \n\n**Technical comments**\n- Most image restoration problems are ill-posed: there are multiple clean images that could lead to the same low-quality observation. This implies that it won't exist a deterministic OT plan (so no 1-1- map). This is particularly obvious in the case of image colorization: from a gray image estimate a possible RGB image. This raises the question of what is the formulation doing in this case? There seems to be a mismatch between the mathematical formulation and the practical use in inverse problems. \n\n- The paper claims optimality of the restoration map, but in most inverse problems it is unclear wether this mapping exists (in general there's a loss of information so no 1-1 mapping between low-quality - high-quality image).\n\n- One of the claims is that previous work using OT map didn't address the problem in the input space dimension. The paper does not provide intuition or a mathematical justification why the proposed formulation should work better than the other methods when addressing the problem in the original dimension.\n\n**Experiments**\n- The paper claims to be able to work on high-dimensional data, but in terms of image size the experiments consist of images of at most 64x64 RGB pixels. It's true that this could represent a high-dimensional dataset, but in general most of the problems with generated/restored images are visible at larger resolutions (this is particularly the case for image restoration). \n\n\n**Typos**\n- pag 3. Duality paragraph. $u^c(x)$ → $u^c(y)$.\n- pag 4. Alexandroff theory (?)\n- pag 4. Figure 3. This figure is not particularly illustrative.\n- pag4 - pag5. Indentation and presentation of equations (8) to (11) is confusing.\n- pag 5.  const($\\mu$,$\\nu$) → cost($\\mu$,$\\nu$)\n-pag 6. \"large-scale generative models\" doesn't seem particularly accurate for the shown experiments (32x32 to 64x64 images).\n\n \n", "summary_of_the_paper": "This work introduces a method for doing generative modeling, and conditional generative modeling (e.g., image restoration) using the optimal transport map between two distributions.\n\nThe paper first discusses the different uses of Optimal Transport (OT) on generative models, and shows that most of the works are on using the OT cost and not the map/plan itself. The paper also discusses some of the work that does show the OT map but only into latent spaces due to the intrinsic complexity of computing the OT map on the original space.\n\nThe proposed method uses the Wasserstein-2  distance as a cost measure and assumes the existence of a unique OT deterministic plan, implying the existence of a map.  The formulation of the OT plan estimation problem (Lemma 4.1) ends up addressing a saddle point problem (max-min).\n\nThe analysis requires that the dimensions of input and output spaces be the same. The method also proposes to use a deterministic and non-learnable mapping to embed the input signals on a space that has the same dimension of the output space, so the framework can be applied. This embedding needs to be manually defined depending on the application.\n\nThe paper presents experiments on image generation and image restoration on classical  datasets: MNIST, CIFRAR, CelebA (64x64). Regarding image restoration three different applications are analyzed: denoising, colorization inpainting.\n\nThe method compares similarly to other generative modeling models (in particular similar to WGANs).\n", "summary_of_the_review": "This paper introduces a mathematical formulation to recover the OT map that maps two distributions for doing generative modeling and (unpaired) image restoration tasks. The main weaknesses of the paper are the presentation and experimental validating the approach.  It's hard to understand what is different with respect to previous work also doing OT map estimation. What's particularly interesting about this mathematical formulation that enables it to work on the original dimension?  Experiments are only on small images which implies that the claim of working on high-dimensions and large-scale datasets is a little too strong.\n\n--\nAfter more discussion with the authors and after considering the updated manuscript I'm raising my score (marginally above the acceptance threshold).", "strength_and_weaknesses": "", "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": "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": "DF_pYrJRhbU", "reviewer": "Reviewer_fFFr", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "This paper sheds new light on using OT maps as generators by exploiting the OT map from latent space to data space with unequal dimensions. This is in contrast with previous approaches that manipulate distributions in the latent space. \n\nThe theoretical results of the error analysis of the proposed method are solid and novel in computer vision field. The proposed method does not add a sophisticated mechanism to compute the OT map compared to WGAN. Instead, the structure is very similar to that of WGAN model, which is a strength in terms of computational complexity. \n\nExtensive evaluation of the proposed method with previous methods is thoroughly conducted on both synthetic and real-world datasets. The proposed method performs comparable or better in terms of visual quality of the generated images as well as FID scores than competitive methods. \n\nA minor concern is that, as pointed out by previous studies (e.g., An 2020(a)), the OT map is generally highly discontinuous which makes approximating this map by DNNs tricky. I think this is probably one of the reasons why some previous methods avoid using DNNs to compute OT/generator maps (in contrast with the proposed method). A discussion on the singularities of the generator/OT map would be helpful for a deeper understanding of the proposed method. ", "summary_of_the_paper": "In this paper, a novel optimal transport (OT) map based generative model is proposed. The OT map is computed as the generator map from noise distribution to data distribution. Different from previous works that mainly focus on manipulating distributions in latent space, the proposed method directly computes the OT map in the data/ambient space. One of the benefits of this approach is that the FID is controlled by the W-2 loss in the data space (i.e., theorem 4.3). Experiment results show that the proposed method has comparable or better performance than competitive methods on various tasks. ", "summary_of_the_review": "Overall the paper is well-presented, with novel theoretical and empirical results. I would like to recommend accepting this paper. ", "strength_and_weaknesses": "", "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": "02mcD7GRtnD", "reviewer": "Reviewer_4T3P", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "This paper has some strengths. 1) the proposed method can apply OT maps directly in ambient space, and fit OT maps for Q-embedded quadratic cost between located on the spaces of unequal and equal dimensions. 2) The boundary error of OT maps is analyzed theoretically.\nThis paper has some weaknesses. 1) Compared with OT maps as loss function and OT maps as generation model, what is computational complexity of this paper. 2) what is the detailed model diagram of this paper, and how do you choose appropriate Q-embedding, what exactly is Q-embedding ?\n", "summary_of_the_paper": "This paper proposes a min-max optimization algorithm to apply OT maps directly in ambient space. And that extends the method the case when the input and output distributions are located in the spaces of different dimensions and derive error bounds for the computed OT map.\nContributions of this paper have three aspects: 1) the end-to-end algorithm is given to fit OT maps for Q-embedded quadratic cost, i.e., if two distributions located on the spaces of equal dimensions, the identity embedding Q(x)≡x (Wasserstein-2 distance); if not, choosing Q-embedded quadratic cost. 2) The theoretical analysis the error bounds is proved. 3) this paper demonstrates large-scale applications of OT maps in CV tasks, image generation and image restoration.\n", "summary_of_the_review": "The idea of this paper is very clear, the theoretical analysis is rigorous, and the experimental results can fully illustrate the performance and effectiveness of the method. But there are a few problems, The c-transform $u^{c}(x)$ on page3 should be $u^{c}(y)$, The computational complexity of the algorithm is not reflected in the text, what is Q-embedding? what is the relationship between it and the generated mapping?", "strength_and_weaknesses": "", "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": "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": {}}]

Author rebuttal:

Meta-review: The paper proposes a new method to learn OT maps, and reframes it in the GAN literature. The initial method works when computing maps between equal dimensions, through duality and an identity (10 - 11, amply discussed in the reviewing process). Lemma 4.1 provides the main result. While the discussion right below on the fact that several functions (non-OT maps) might maximize that criterion is not completely satisfactory, the result provides an interesting characterization. The second contribution adds a method to compute OT maps between spaces of unequal dimensions. Overall the contribution sounds a bit ad-hoc, and one wonders whether this does really work (comments such as "we add small gradient penalty (Gulrajani et al., 2017) on potential ψω for better stability. The penalty in not included in Algorithm 1 to keep it simple." are strange and point to instability) but the overall creativity and new ideas in the paper seem to have garnered enough support from reviewers to push for an accept.
Decision: Accept (Poster)