Year: 2021

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 minmax 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.

heading: 1 INTRODUCTION
text: Since the discovery of Generative Adversarial Networks (GANs, Goodfellow et al. (2014)), there has been a surge in generative modeling (Radford et al., 2016; Arjovsky et al., 2017; Brock et 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., 2017), convergence (Sanjabi et al., 2018), and stability (Miyato et al., 2018; Kim et al., 2021) 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 as the loss for generative models, see Figure 1a. This is the most prevalent class of methods which includes WGAN (Arjovsky et al., 2017) and its modifications: WGAN-GP (Gulrajani et al., 2017), WGAN-LP (Petzka et al., 2018), and WGAN-QC (Liu et al., 2019). Second, the optimal transport map is used as a generative model itself, see Figure 1b. Such approaches include LSOT (Seguy et al., 2018), AE-OT (An et al., 2020a), ICNN-OT (Makkuva et al., 2020), W2GN (Korotin et al., 2021a). Models of the first class have been wellstudied, 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 of such generative models depends on the underlying AEs, in which decoding transformations are often not accurate; as a result this deficiency limits practical applications in high-dimensional ambient spaces. For this reason, using OT in the latent space does not necessarily guarantee superior performance in generative modeling.
The focus of our paper is the second class of OT-based 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 translation, the learner has to fit a map between two data distributions which preserves the image content. CycleGAN-based models (Zhu et al., 2017) are widely used for this purpose. However, they typically 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:
1. We propose an end-to-end algorithm (M4.3) to fit OT maps for the quadratic cost (Wasserstein-2 distance) between distributions located on the spaces of equal dimensions (M4.1) and extend the method to unequal dimensions as well (M4.2). We prove error bounds for the method (M4.4).
2. We demonstrate large-scale applications of OT maps in popular computer vision tasks. We consider image generation (M5.1) and unpaired image restoration (M5.2) tasks.
Our strict OT-based framework allows the theoretical analysis of the recovered transport map. The OT 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 on 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 as the loss.
Notations. In what follows, X and Y are two complete metric spaces, µ(x) and ν(y) are probability distributions on X and Y , respectively. For a measurable map T : X → Y , T#µ denotes the pushforward distribution of µ, i.e., the distribution for which any measurable set E ⊂ Y satisfies T#µ(E) = µ(T
−1(E)). For a vector x, ‖x‖ denotes its Euclidean norm. We use 〈x, y〉 to denote the inner product of vectors x and y. We use Π(µ, ν) to denote the set of joint probability distributions on X ×Y whose marginals are µ and ν, respectively (couplings). For a function f : RD → R∪{±∞} its Legendre–Fenchel transform (the convex conjugate) is f(y) = supx∈RD{〈x, y〉 − f (x)}. It is convex, even if f is not.

heading: 2 BACKGROUND ON OPTIMAL TRANSPORT
text: Consider a cost of transportation, c : X × Y → R defined over the product space of X and Y . Monge’s Formulation. The optimal transport cost between µ and ν for ground cost c(·, ·) is
Cost(µ, ν) def= inf T#µ=ν ∫ X c (x, T (x)) dµ(x), (1)
where the infimum is taken over all measurable maps T : X → Y pushing µ to ν, see Figure 2. The map T ∗ on which the infimum in (1) is attained is called the optimal transport map. Monge’s formulation does not allow splitting. For example, when µ is a Dirac distribution and ν is a non-Dirac distribution, the feasible set of equation (1) is empty.
Kantorovich’s Relaxation. Instead of asking to which particular point y ∈ Y should all the probability mass of x be moved, Kantorovich (1948) asks how the mass of x should be distributed among all y ∈ Y . Formally, a transport coupling replaces a transport map; the OT cost is given by:
Cost(µ, ν) def= inf π∈Π(µ,ν) ∫ X×Y c (x, y) dπ(x, y), (2)
where the infimum is taken over all couplings π ∈ Π(µ, ν) of µ and ν. The coupling π∗ attaining the infimum of (2) is called the optimal transport plan. Unlike the formulation of (1), the formulation of (2) 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 π∗ is deterministic, i.e., π∗ = [idX , T ∗]#µ for some T ∗ : X → Y , then T ∗ minimizes (1). Duality. The dual form of (2) is given by (Kantorovich, 1948):
Cost(µ, ν) = sup (u,v) {∫ X u(x)dµ(x) + ∫ Y v(y)dν(y) : u(x) + v(y) ≤ c(x, y) } , (3)
with u ∈ L1(µ), v ∈ L1(ν) called Kantorovich potentials. For u : X → R and v : Y → R define their c-transforms by uc(y) = infx∈X {c (x, y)− u (x)} and vc(x) = infy∈Y{c (x, y)− v (y)} respectively. Using c-transform, (3) is reformulated as (Villani, 2008, M5)
Cost(µ, ν)=sup v {∫ X vc(x)dµ(x)+ ∫ Y v(y)dν(y) } =sup u {∫ X u(x)dµ(x)+ ∫ Y uc(y)dν(y) } . (4)
Primal-dual relationship. For certain ground costs c(·, ·), the primal solution T ∗ of (1) can be recovered from the dual solution u∗ of (3). For example, if X = Y = RD, c(x, y) = h(x− y) with strictly convex h : RD → R and µ is absolutely continuous supported on the compact set, then
T ∗(x) = x− (∇h)−1 ( u∗(x) ) , (5)
see (Santambrogio, 2015, Theorem 1.17). For general costs, see (Villani, 2008, Theorem 10.28).

heading: 3 OPTIMAL TRANSPORT IN GENERATIVE MODELS
text: OPTIMAL TRANSPORT COST AS THE LOSS (Figure 1a). Starting with the works of Arjovsky & Bottou (2017); Arjovsky et al. (2017), 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 ν and fake distribution µθ, the goal is to minimize Cost(µθ, ν) w.r.t. the parameters θ. Typically, µθ is a pushforward distribution of some given distribution, e.g., N (0, I), via generator network Gθ. The Wasserstein-1 distance (W1), i.e., the transport cost for ground cost c(x, y) = ‖x − y‖, is the most practically prevalent example of such a loss. Models based on this loss are known as Wasserstein GANs (WGANs). They estimateW1(µθ, ν) based on the dual form as given by (4). For W1, the optimal potentials u∗, v∗ of (4) satisfy u∗ = −v∗ where u∗ is a 1-Lipschitz function (Villani, 2008, Case 5.16). As a result, to computeW1, one needs to optimize the following simplified form:
W1(µθ, ν) = sup ‖u‖L≤1 {∫ X u(x)dµθ(x)− ∫ Y u(y)dν(y) } . (6)
In WGANs, the potential u is called the discriminator. Optimization of (6) reduces constrained optimization of (4) with two potentials u, v to optimization of only one discriminator u. In practice, enforcing the Lipschitz constraint on u is challenging. Most methods to do this are regularizationbased, 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 certain hard restrictions on the discriminator’s architecture (Anil et al., 2019; Tanielian & Biau, 2021).
General transport costs (other than W1) can 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 similar to (6) for W1. Consequently, the challenging estimation of the c-transform uc is needed. To avoid this, Sanjabi et al. (2018) consider the dual form of (3) with two potentials u, v instead form (4) with one u and softly enforce the condition u(x) + v(y) ≤ c(x, y) via entropy or quadratic regularization. Nhan Dam et al. (2019) use the dual form of (4) and amortized optimization to compute uc via an additional neural network. Both methods work for general c(·, ·), though the authors test them for c(x, y) = ‖x − y‖ only, i.e., W1 distance. Mallasto et al. (2019) propose a fast way to approximate the c-transform and test the approach (WGAN-(q, p)) with several costs, in particular, the Wasserstein-2 distance (W2), i.e., the transport cost for the quadratic ground cost c(x, y) = 12‖x−y‖
2. Specifically forW2, Liu et al. (2019) approximate the c-transform via a 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 et 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 (3). In practice, these approaches usually learn the ground OT cost c(·, ·). The aforementioned methods use OT cost in the ambient space to train GANs. There also exist approaches using OT cost in the latent space. For example, Tolstikhin 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 1b). Methods to compute the OT map (plan) are less common in comparison to those computing the cost. Recovering the map from the primal form (1) or (2) usually yields complex optimization objectives containing several adversarial terms (Xie et al., 2019; Liu et 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 (M2) makes it possible to recover the OT map via solving the dual form (3). Dual-form based methods primarily considerW2 cost due to its nice theoretical properties and relation to convex functions (Brenier, 1991). In the semi-discrete case (µ is continuous, ν is 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, Seguy et al. (2018) use the entropy (quadratic) regularization to recover the dual potentials and extract OT map from them via the barycenteric projection. Taghvaei & Jalali (2019), Makkuva et al. (2020), Korotin et al. (2021a) employ input-convex neural networks (ICNNs, see Amos et al. (2017)) to parametrize potentials in the dual problem and recover OT maps by using their gradients.
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 forW2 (Korotin et al., 2021b) reveals their crucial limitations, which negatively affect their scalability, such as poor expressiveness of ICNN architectures or bias due to regularization.

heading: 4 END-TO-END SOLUTION TO LEARN OPTIMAL MAPS

heading: 4.1 EQUAL DIMENSIONS OF INPUT AND OUTPUT DISTRIBUTIONS
text: In this section, we use X = Y = RD and consider the Wasserstein-2 distance (W2), i.e., the optimal transport for the quadratic ground cost c(x, y) = 12‖x − y‖
2. 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 µ, ν in view there exists a unique OT plan π∗ minimizing (3) and it is deterministic, i.e., π∗ = [idRD , T ∗]#µ. Here T ∗ is an OT map which minimizes (1). Previous related works (Makkuva et al., 2020; Korotin et al., 2021a) assumed the absolute continuity of µ, which implied the existence and uniqueness of T ∗ (Brenier, 1991).
Let ψ(y) def= 12‖y‖ 2 − v(y), where v is the potential of (4). Note that
vc(x) = inf y∈RD
{ 1
2 ‖x− y‖2 − v(y)
} = 1
2 ‖x‖2 − sup y∈RD {〈x, y〉 − ψ(y)} = 1 2 ‖x‖2 − ψ(x). (7)
Therefore, (4) is equivalent to W22 (µ, ν) = ∫ X ‖x‖2 2 dµ(x) + ∫ Y ‖y‖2 2 dν(x) + sup ψ { − ∫ X ψ(x)dµ(x)− ∫ Y ψ(y)dν(y) } = (8)
Constant(µ, ν)− inf ψ {∫ X ψ(x)dµ(x) + ∫ Y ψ(y)dν(y) } = (9)
Constant(µ, ν)− inf ψ {∫ X sup y∈RD {〈x, y〉 − ψ(y)} dµ(x) + ∫ Y ψ(y)dν(y) } = (10)
Constant(µ, ν)− inf ψ { sup T ∫ X { 〈x, T (x)〉 − ψ ( T (x) )} dµ(x) + ∫ Y ψ(y)dν(y) } (11)
where between lines (10) and (11) we replace the optimization over y ∈ RD with the equivalent optimization over functions T : RD → RD. The equivalence follows from the interchange between the integral and the supremum (Rockafellar, 1976, Theorem 3A). Thanks to the following lemma, we may solve saddle point problem (11) and obtain the OT map T ∗ from its solution (ψ∗, T ∗). Lemma 4.1. Let T ∗ be the OT map from µ to ν. Then, for every optimal potential ψ∗,
T ∗ ∈ arg sup T ∫ X { 〈x, T (x)〉 − ψ∗ ( T (x) )} dµ(x). (12)
We prove Lemma 4.1 in Appendix A.1. For general µ, ν the arg supT set for optimal ψ ∗ might contain not only OT map T ∗, but other functions as well. Working with real-world data in experiments (M5.2), we observe that despite this issue, optimization (11) still recovers T ∗.
Relation to previous works. The use of the function T to approximate the c-transform was proposed 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. (2021a) who primarily assumed continuous µ, ν and reduced (11) to convex ψ and T = ∇φ for convex φ. Issues with nonuniqueness of solution of (12) were softened, but using ICNNs to parametrize ψ became necessary.
Korotin et al. (2021b) demonstrated that ICNNs negatively affect practical performance of OT and tested an unconstrained formulation similar to (11). As per the evaluation, it provided the best empirical performance (Korotin et al., 2021b, M4.5). The method bMM:Re they consider parametrizes 1 2‖ · ‖
2 − ψ(·) by a neural network, while we directly parametrize ψ(·) by a neural network (M4.3). Recent work by Fan et al. (2021) exploits formulation similar to (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 a 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.

heading: 4.2 UNEQUAL DIMENSIONS OF INPUT AND OUTPUT DISTRIBUTIONS
text: Consider the case when X = RH and Y = RD have different dimensions, i.e., H 6= D. In order to map the probability distribution µ to ν, a straightforward solution is to embed X to Y via some Q : X → Y and then to fit the OT map between Q#µ and ν for the quadratic cost on Y = RD. In this case, the optimization objective becomes
inf ψ sup T {∫ X { 〈Q(x), T ( Q(x) ) 〉 − ψ ( T ( Q(x) ))} dµ(x) + ∫ Y ψ(y)dν(y) } (13)
with the optimal T ∗ recovering the OT map from Q#µ to ν. For equal dimensions H = D and the identity embedding Q(x) ≡ x, expression (13) reduces to optimization (11) up to a constant. Instead of optimizing (13) over functions T : Q(X ) → Y , we propose to consider optimization directly over generative mappings G : X → Y:
L(ψ,G) def= inf ψ sup G {∫ X { 〈Q(x), G(x)〉 − ψ ( G(x) )} dµ(x) + ∫ Y ψ(y)dν(y) } (14)
Our following lemma establishes connections between (14) and OT with unequal dimensions: Lemma 4.2. Assume that exists a unique OT plan between Q#µ and ν and it is deterministic, i.e., [idRD , T ∗]#(Q#µ). Then G∗(x) = T ∗ ( Q(x) ) is the OT map between µ and ν for the Q-embedded quadratic cost c(x, y) = 12‖Q(x)− y‖ 2. Moreover, for every optimal potential ψ∗ of problem (14),
G∗ ∈ arg sup G ∫ X { 〈Q(x), G(x)〉 − ψ∗ ( G(x) )} dµ(x). (15)
We prove Lemma 4.2 in Appendix A.2 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 Qembedded quadratic cost between distributions µ and ν by solving the saddle point problem (14). Note the situation with nonuniqueness of arg supG is similar to M4.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., 2021, M5.2) but only on toy examples. We demonstrate that our method works well in large-scale generative modeling tasks (M5.1). Theoretical properties of OT maps for embedded costs are studied, e.g., in (Pass, 2010; McCann & Pass, 2020).

heading: 4.3 PRACTICAL ASPECTS AND OPTIMIZATION PROCEDURE
text: To optimize functional (14), we approximateG : RH → RD and ψ : RD → R with neural networks Gθ, ψω and optimize their parameters via stochastic gradient descent-ascent (SGDA) by using minibatches from µ, ν. The practical optimization procedure is given in Algorithm 1 below. Following the usual practice in GANs, 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.
Algorithm 1: Learning the optimal transport map between unequal dimensions. Input : Input distribution µ on X = RH ; output distribution ν on Y = RD;
generator network Gθ : RH → RD; potential network ψω : RD → R; number of iterations per network: KG, Kψ; embedding Q : X → Y;
Output: Trained generator Gθ representing OT map from µ to ν; 1 repeat 2 for kψ = 1 to Kψ do 3 Draw batch X ∼ µ and Y ∼ ν; 4 Lψ ← 1|Y | ∑ y∈Y ψω(y)− 1 |X| ∑ x∈X ψω ( Gθ(x) ) ; 5 Update ω by using ∂Lψ∂ω to minimize Lψ; 6 for kG = 1 to KG do 7 Draw batch X ∼ µ; 8 LG ← 1|X| ∑ x∈X [ ψ ( G(x) ) − 〈Q(x), Gθ(x)〉 ] ;
9 Update θ by using ∂LG∂θ to minimize LG; 10 until not converged; 11 return Gθ
Relation to previous works. WGAN by Arjovsky & Bottou (2017) usesW1 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 1 has similarities with WGAN’s training. The update of ψ (line 4) coincides with discriminator’s update in WGAN. The update of generator G (line 8) differs from WGAN’s update by the term −〈Q(·), Gθ(·)〉. Besides, in WGAN the optimization is infG supD. We have infψ supG, i.e., the generator in our case is the solution of the inner problem.

heading: 4.4 ERROR ANALYSIS
text: Given a pair (ψ̂, Ĝ) approximately solving (14), a natural question to ask is how good is the recovered OT map Ĝ. In this subsection, we provide a bound on the difference between G∗ and Ĝ based on the duality gaps for solving outer and inner optimization problems.
In (8), and, as the result, in (10), (11), (13), (14), it is enough to consider optimization over convex functions ψ, see (Villani, 2008, Case 5.17). Our theorem below assumes the convexity of ψ̂ although it might not hold in practice since in practice ψ̂ is a neural network. Theorem 4.3. Assume that exists a unique deterministic OT plan for Q-embedded quadratic cost between µ and ν, i.e., π∗ = [idRH , G∗]#µ for G∗ : RH → RD. Assume that ψ̂ is β-strongly convex (β > 0) and Ĝ : RH → RD. Define
1 = sup G L(ψ̂, G)− L(ψ̂, Ĝ) and 2 = sup G L(ψ̂, G)− inf ψ sup G L(ψ,G)
Then the following bound holds true for the OT map G∗ from µ to ν:
FID(Ĝ#µ, ν) L2 ≤ 2 · W22 (Ĝ#µ, ν) ≤ ∫ X ‖Ĝ(x)−G∗(x)‖2dµ(x) ≤ 2 β ( √ 1 + √ 2) 2, (16)
where FID is the Fré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., 2016).
We prove Theorem 4.3 in Appendix A.3. The duality gaps upper bound L2(µ) norm between computed Ĝ and true G∗ maps, and the W22 between true ν and generated (fake) distribution Ĝ#µ. Consequently, they upper bound FID between data ν and fake (generated) Ĝ#µ distributions.
Relation to previous works. Makkuva et al. (2020); Korotin et al. (2021a) prove related bounds for W2 with µ, ν located on the spaces of the same dimension. Our result holds for different dimensions.

heading: 5 EXPERIMENTS
text: We evaluate our algorithm in generative modeling of the data distribution from a noise (M5.1) and unpaired image restoration task (M5.2). Technical details are given in Appendix B. Additionally, in Appendix B.3 we test our method on toy 2D datasets and evaluate it on the Wasserstein-2 benchmark (Korotin et al., 2021b) in Appendix B.2. The code is in the supplementary material.

heading: 5.1 MODELING DATA DISTRIBUTION FROM NOISE DISTRIBUTION
text: In this subsection, µ is a 128-dimensional normal noise and ν the high-dimensional data distribution.
Let the images from ν be of size w × h with c channels. As the embedding Q : X → Y we use a naive upscaling of a noise. For x ∈ R128 we represent it as 1-channel 16× 8 image, stack c copies over channel dimension and bicubically upscale it to the size w × h of data images from ν. We test our method on MNIST 32× 32 (LeCun et al., 1998), CIFAR10 32× 32 (Krizhevsky et al., 2009), and CelebA 64× 64 (Liu et al., 2015) image datasets. In Figure 5, we show random samples generated by our approach, namely Optimal Transport Modeling (OTM). To quantify the results, in Tables 1 and 2 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 on 50K real and generated samples. Additionally, in Appendix B.3, we test our method on 128×128 CelebA faces. We provide qualitative results (images of generated faces) in Figure 11.
For comparison, we include the scores of existing generative models of three types: (1) OT map as the generative model; (2) OT cost as 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 or comparable to existing alternatives.

heading: 5.2 UNPAIRED IMAGE RESTORATION
text: In this subsection, we consider unpaired image restoration tasks on CelebA faces dataset. In this case, the input distribution µ consists of degraded images, while ν are clean images. In all the cases, embedding Q is a straightforward identity embedding Q(x) ≡ x. In image restoration, optimality of the restoration map is desired since the output (restored) image is expected to be close to the input (degraded) one minimizing the transport cost. Note that GANs
Table 2: Results on CelebA dataset.
Model Related Work FID ↓ DCGAN Radford et al. (2016) 52.0 DRAGAN Kodali et al. (2017) 42.3 BEGAN Berthelot et al. (2017) 38.9 NVAE Vahdat & Kautz (2020) 13.4 NCP-VAE Aneja et al. (2021) 5.2 WGAN Arjovsky et al. (2017) 41.3 WGAN-GP Gulrajani et al. (2017) 30.0 WGAN-QC Liu et al. (2019) 12.9 AE-OT An et al. (2020a) 28.6 W2GN+AE Korotin et al. (2021a) 17.2 AE-OT-GAN An et al. (2020b) 7.8 OTM Ours 7.7
do not seek for an optimal mapping. However, in practice, due to implicit inductive biases such as convolutional architectures, GANs still tend to fit low transport cost maps (Bézenac et al., 2021).
The experimental setup is shown in Figure 6. We split the dataset in 3 parts A, B, C containing 90K, 90K, 22K samples respectively. To each image we apply the degradation transform (decolorization, noising or occlusion) and obtain the degraded dataset containing of 3 respective parts A, B, C. For unpaired training we use part A of degraded and part B of clean images. For testing, we use parts C.
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 are given in Table 3, details of each experiment and qualitative results are given below.
As the baseline, we include WGAN-GP. For a fair comparison, we fit it using exactly the same hyperparameters as in our method. This is possible due to the similarities between our method and WGAN-GP’s training procedure, see discussion in M4.3.
Denoising. To create noisy images, we add white normal noise with σ = 0.3 to each pixel. Figure 7 illustrates image denoising using our OTM approach on the test part of the dataset. We show additional qualitative results for varying σ in Figure 16 of M(B.3).
Colorization. To create grayscale images, we average the RGB values of each pixel. Figure 8 illustrates image colorization using OTM on the test part of the dataset.
Inpainting. To create incomplete images, we replace the right half of each clean image with zeros. Figure 9 illustrates image inpainting using OTM on the test part of the dataset.

heading: 6 CONCLUSION
text: Our method fits OT maps for the embedded quadratic transport cost between probability distributions. Unlike predecessors, it scales well to high dimensions producing applications of OT maps directly in ambient spaces, such as spaces of images. The performance is comparable to other existing generative models while the complexity of training is similar to that of popular WGANs.
Limitations. For distributions µ, ν we 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 12‖x − y‖
2. Nevertheless, our approach extends to other costs c(·, ·), see Fan et al. (2021). When the dimensions are unequal, we restrict our analysis to embedded quadratic cost 12‖Q(x) − y‖
2 where Q equalizes dimensions. Choosing the embedding Q might not be straightforward in some practical problems, but our evaluation (M5.1) shows that even naive choices of Q work well. 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 problems, we expect our optimal transport approach to improve the performance of the existing models. However, biases in data might lead to biases in the pushforward samples. This should be taken into account when using our method in practical problems.
Reproducibility. The PyTorch source code is provided in the supplementary material. The instructions to use the code are included in the README.md file.

heading: A PROOFS

heading: A.1 PROOF OF LEMMA 4.1
text: Proof. It is enough to prove that ψ∗(x) = 〈T ∗(x), x〉 − ψ∗ ( T (x) ) holds µ-almost everywhere, i.e., T ∗(x) ∈ arg sup y∈RD {〈x, y〉 − ψ∗(y)}. Since ν = T ∗#µ, we use (9) with ψ ← ψ∗ to derive
W22 (µ, ν)− ∫ X 1 2 ‖x‖2dµ(x)− ∫ Y 1 2 ‖y‖2dν(y) =
− ∫ X ψ∗(x)dµ(x)− ∫ Y ψ∗(y)dν(y) = − ∫ X ψ∗(x)dµ(x)− ∫ Y ψ∗ ( T ∗(x) ) dµ(x) =
− ∫ X [ ψ∗(x) + ψ∗ ( T ∗(x)︸ ︷︷ ︸
≥〈T∗(x),x〉
)] dµ(x) ≤ − ∫ X 〈T ∗(x), x〉dµ(x) = (17)
∫ X 1 2 ‖x− T ∗(x)‖2dµ(x)− ∫ X 1 2 ‖x‖2dµ(x)− ∫ X 1 2 ‖T ∗(x)‖2dµ(x) =
W22 (µ, ν)− ∫ X 1 2 ‖x‖2dµ(x)− ∫ Y 1 2 ‖y‖2dν(y).
As a result, inequality (17) becomes the equality, in particular, ψ∗(x) + ψ∗ ( T ∗(x) ) = 〈T ∗(x), x〉 holds µ-almost everywhere.

heading: A.2 PROOF OF LEMMA 4.2
text: Proof. Let QW22 denote the Q-embedded quadratic cost. We use the change of variables formula to derive
QW22 (µ, ν) = inf π∈Π(µ,ν) ∫ X×Y 1 2 ‖Q(x)− y‖2dπ(x, y) =
inf π′∈Π(Q#µ,ν) ∫ X×Y 1 2 ‖x− y‖2dπ′(x, y) =W22 (Q#µ, ν), (18)
i.e., computing the OT plan for QW22 (µ, ν) boils down to computing the OT plan forW22 (Q#µ, ν). It follows that [idRH , T ∗ ( Q(x) ) ]#µ = [idRH , G∗]#µ is an OT plan for QW22 (µ, ν), and G∗ is the OT map. Inclusion (15) now follows from Lemma 4.1.

heading: A.3 PROOF OF THEOREM 4.3
text: Proof. Pick any G′ ∈ arg supG L(ψ̂, G) = arg supG ∫ X { 〈Q(x), G(x)〉 − ψ̂ ( G(x) )} dµ(x) or,
equivalently, for all x ∈ RH , G′(x) ∈ arg supy { 〈Q(x), y〉 − ψ̂(y) } . Consequently, for all y ∈ RD
〈Q(x), G′(x)〉 − ψ̂ ( G′(x) ) ≥ 〈Q(x), y〉 − ψ̂(y),
which after regrouping the terms yields ψ̂(y) ≥ ψ̂ ( G′(x) ) + 〈Q(x), y −G′(x)〉.
This means that Q(x) is contained in the subgradient ∂ψ̂ at G′(x) for a convex ψ̂. Since ψ̂ is β-strongly convex, for points G(x), G′(x) ∈ RD and Q(x) ∈ ∂ψ̂ ( G′(x) ) we derive
ψ̂ ( G(x) ) ≥ ψ̂ ( G′(x) ) + 〈Q(x), G(x)−G′(x)〉+ β
2 ‖G′(x)−G(x)‖2.
Regrouping the terms, this gives[ 〈Q(x), G′(x)〉 − ψ̂ ( G′(x) )] − [ 〈Q(x), G(x)〉 − ψ̂ ( G(x) )] ≥ β
2 ‖G′(x)−G(x)‖2.
Integrating w.r.t. x ∼ µ yields 1 = L(ψ̂, G′)− L(ψ̂, G′) ≥ β ∫ X 1 2 ‖G′(x)−G(x)‖2dµ(x) = β 2 · ‖G−G′‖2L2(µ). (19)
Let G∗ be the OT map from µ to ν. We use G∗#µ = ν to derive
L(ψ̂, G′) = ∫ X { 〈Q(x), G′(x)〉 − ψ̂ ( G′(x) )} dµ(x) + ∫ Y ψ̂(y)dν(y) =∫
X
{ 〈Q(x), G′(x)〉 − ψ̂ ( G′(x) )} dµ(x) + ∫ X ψ̂ ( G∗(x) ) dµ(x) =
∫ X 〈Q(x), G′(x)〉 − ψ̂ ( G′(x) ) + ψ̂ ( G∗(x) )︸ ︷︷ ︸ ≥〈Q(x),G∗(x)〉+β 12‖G′−G∗‖2  dµ(x) ≥∫ X 〈Q(x), G∗(x)〉dµ(x) + β ∫ X 1 2 ‖G′ −G∗‖2dµ(x). (20)
Let ψ∗ be an optimal potential in (14). Thanks to Lemma 4.2, we have
inf ψ sup G L(ψ,G) = L(ψ∗, G∗) =∫
X
{ 〈Q(x), G∗(x)〉 − ψ∗ ( G∗(x) )} dµ(x) + ∫ Y ψ∗(y)dν(y) =∫
X
{ 〈Q(x), G∗(x)〉 − ψ∗ ( G∗(x) )} dµ(x) + ∫ X ψ∗ ( G∗(x) ) dµ(x) =∫
X 〈Q(x), G∗(x)〉dµ(x) (21)
By combining (20) with (21), we obtain 2 = L(ψ̂, G′)− L(ψ∗, G∗) ≥ β ∫ X 1 2 ‖G′ −G∗‖2dµ(x) = β 2 · ‖G′ −G∗‖2L2(µ) (22)
The right-hand inequality of (16) follows from the triangle inequality combined with (19) and (22). The middle inequality of (16) follows from (Korotin et al., 2021a, Lemma A.2) and G∗#µ = ν.
Now we prove the left-hand inequality of (16). Let I be the feature extractor of the pre-trained InceptionV3 neural networks. FID score between generated (fake) Ĝ#µ and data distribution ν is
FID(Ĝ#µ, ν) = FD(I#Ĝ#µ, I#ν) ≤ 2 · W22 (I#Ĝ#µ, I#ν), (23)
where FD(·, ·) is the Fréchet distance which lower bounds 2 · W22 , see (Dowson & Landau, 1982). Finally, from (Korotin et al., 2021a, Lemma A.1) it follows that
W22 (I#Ĝ#µ, I#ν) ≤ L2 · W22 (Ĝ#µ, ν). (24)
Here L is the Lipschitz constant of I. We combine (23) and (24) to get the left-hand inequality in (16).

heading: B EXPERIMENTAL DETAILS
text: We use the PyTorch framework. All the experiments are conducted on 2×V100 GPUs. We compute inception and FID scores with the official implementation from OpenAI1 and TTUR2. The compared results are taken from the respective papers or publicly available source codes.
1IS: https://github.com/openai/improved-gan/tree/master/inception_score 2FID: https://github.com/bioinf-jku/TTUR

heading: B.1 GENERAL TRAINING DETAILS
text: MNIST (LeCun et al., 1998). On MNIST, we use x ∈ R128 and y ∈ R28×28. The batch size is 64, learning rate 2 · 10−4, optimizer Adam (Kingma & Ba, 2014) with betas (0.5, 0.999), gradient penalty coefficient λ = 10, and the number of training epochs T = 200. We observe stable training while updating G once in multiple ψ updates, i.e., kG = 1 and kψ = 5.
CIFAR10 (Krizhevsky et al., 2009). We use all 50000 samples while training. The latent vector x ∈ R128 and y ∈ R32×32×3, batch size 64, λ = 10, kG = 1, kψ = 5, T = 1000, Adam optimizer with betas (0, 0.99), and learning rate 2 · 10−4. CelebA (Liu et al., 2015). We use x ∈ R128 and y ∈ R64×64×3. The images are first cropped at the center with size 140 and then resized to 64 × 64. We consider all 202599 samples. We use Adam with betas (0.5, 0.999), T = 200, and learning rate 1 · 10−4. Other hyper-parameters are same as CIFAR10.
Image restoration. In the unpaired image restoration experiments, we use Adam optimizer with betas (0, 0.9), learning rate 1 ·10−3 and train for T = 200 epochs. Other hyper-parameters are same as the image generation experiment on CelebA.
CelebA128x128 (Liu et al., 2015). On this dataset, we resize the cropped images as in CelebA to 128 × 128, i.e. y ∈ R128×128×3. Here, KG = 10, Kψ = 1, λ = 0.01, learning rate 1 · 10−4 and betas=(0.5, 0.999). The batch size is 16 so as to fit in the GPU memory.
Toy datasets. The dimension is D = H = 2, total number of samples is 10000. We use the batch size 128, λ = 1,Kψ = 1,KG = 16, and T = 100. The optimizer is Adam with betas (0.5, 0.99) and learning rate 1 ·10−3. We use the following datasets: Gaussian to mixture of Gaussians3, two moons (sklearn.datasets.make_moons), circles (sklearn.datasets.make_circles), gaussian to S-curve (sklearn.datasets.make_s_curve), and gaussian to swiss roll (sklearn.datasets.make_swiss_roll).
Wasserstein-2 benchmark (Appendix B.2). The dimension is D = H = 64 × 64 × 3. We use batch size 64, λ = 0, Kψ = 1, KG = 5, learning rate 10−4, and Adam optimizer with default betas.

heading: B.2 EVALUATION ON THE CONTINUOUS WASSERSTEIN-2 BENCHMARK
text: 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 (µ, ν) of continuous probability distributions with analytically known OT map T ∗ between them.
For evaluation, we use the ”Early” images benchmark pair (D = 12288), see (Korotin et al., 2021b, M4.1) for details. We adopt the L2-unexplained percentage metric (Korotin et al., 2021a, M5.1) to quantify the recovered OT map T̂ : L2-UVP(T̂ ) = 100 · ‖T̂ − T ∗‖2/Var(ν)%. For our method the L2-UVP metric is only ≈ 1%, see Table 4. This is comparable to the best dMM:Re method which the authors evaluate on their benchmark. The qualitative results are given in Figure 10.

heading: B.3 ADDITIONAL QUALITATIVE RESULTS
text: CelebA128x128. Figure 11 shows the embedding Q(x), the recovered transport map Ĝ(x), and independently drawn real samples y ∼ ν. The extended qualitative results on MNIST, CIFAR10, and CelebA are shown in Figure 12, Figure 13, and Figure 14 respectively. The samples are generated randomly (uncurated) by fitted optimal transport maps between noise and ambient space, e.g., spaces of high-dimensional images.
3https://github.com/AmirTag/OT-ICNN
Figure 15 illustrates latent space interpolation between the generated samples. Figure 16 shows denoising of images with varying levels of σ = 0.1, 0.2, 0.3, 0.4 by the model trained with σ = 0.3.
Toy datasets. Figure 17 shows the results of our method and related approaches (M3) on a toy 2D dataset. Figure 18 shows the results of our method applied to other toy datasets.

heading: B.4 NEURAL NETWORK ARCHITECTURES
text: This section contains architectures on MNIST (Table 5), CIFAR10 (Table 6), CelebA (Table 7), CelebA 128× 128 (Figure 11) generation, image restoration tasks (Table 8), evaluation on the toy 2D datasets and Wasserstein-2 images benchmark (Table 4).
The transport map G on MNIST is the generator of DCGAN (Radford et al., 2016). In the ResNet architectures, we use the ResidualBlock of NCSN (Song & Ermon, 2019).
In the unpaired restoration tasks (M5.2), we use UNet architecture for transport map G and convolutional architecture for potential ψ. Similarly to (Song & Ermon, 2019), we use BatchNormaliation (BN) and InstanceNormalization+ (INorm+) layers.
In the toy 2D examples, we use a simple multi-layer perceptron with 3 hidden layers consisting of 128 neurons each and LeakyReLU activation. The final layer is linear without any activation.
The transport map G and potential ψ architectures on CelebA 128 × 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 Unet4 architecture for transport map T and WGAN-QC discriminator’s architecture for ψ (Liu et al., 2019). These neural network architectures are the same as the authors of the benchmark use.
4https://github.com/milesial/Pytorch-UNet

Review summaries:
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.  The 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.  The 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).  The 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.  The 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.  The method compares similarly to other generative modeling models (in particular similar to WGANs). 

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. Experiments show that the proposed method works well and learns the image distribution successfully.  

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. 

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. Contributions 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.