Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark

Abstract:
Despite the recent popularity of neural network-based solvers for optimal transport (OT), there is no standard quantitative way to evaluate their performance. In this paper, we address this issue for quadratic-cost transport—specifically, computation of the Wasserstein-2 distance, a commonly-used formulation of optimal transport in machine learning. To overcome the challenge of computing ground truth transport maps between continuous measures needed to assess these solvers, we use inputconvex neural networks (ICNN) to construct pairs of measures whose ground truth OT maps can be obtained analytically. This strategy yields pairs of continuous benchmark measures in high-dimensional spaces such as spaces of images. We thoroughly evaluate existing optimal transport solvers using these benchmark measures. Even though these solvers perform well in downstream tasks, many do not faithfully recover optimal transport maps. To investigate the cause of this discrepancy, we further test the solvers in a setting of image generation. Our study reveals crucial limitations of existing solvers and shows that increased OT accuracy does not necessarily correlate to better results downstream. Solving optimal transport (OT) with continuous methods has become widespread in machine learning, including methods for large-scale OT [11, 36] and the popular Wasserstein Generative Adversarial Network (W-GAN) [3, 12]. Rather than discretizing the problem [31], continuous OT algorithms use neural networks or kernel expansions to estimate transport maps or dual solutions. This helps scale OT to large-scale and higher-dimensional problems not handled by discrete methods. Notable successes of continuous OT are in generative modeling [42, 20, 19, 7] and domain adaptation [43, 37, 25]. In these applications, OT is typically incorporated as part of the loss terms for a neural network model. For example, in W-GANs, the OT cost is used as a loss function for the generator; the model incorporates a neural network-based OT solver to estimate the loss. Although recent W-GANs provide state-of-the-art generative performance, however, it remains unclear to which extent this 35th Conference on Neural Information Processing Systems (NeurIPS 2021). success is connected to OT. For example, [28, 32, 38] show that popular solvers for the Wasserstein-1 (W1) distance in GANs fail to estimate W1 accurately. While W-GANs were initially introduced with W1 in [3], state-of-the art solvers now use both W1 and W2 (the Wasserstein-2 distance, i.e., OT with the quadratic cost). While their experimental performance on GANs is similar, W2 solvers tend to converge faster (see [19, Table 4]) with better theoretical guarantees [19, 26, 16]. Contributions. In this paper, we develop a generic methodology for evaluating continuous quadraticcost OT solvers (W2). Our main contributions are as follows: • We use input-convex neural networks (ICNNs [2]) to construct pairs of continuous measures that we use as a benchmark with analytically-known solutions for quadratic-cost OT (M3, M4.1). • We use these benchmark measures to evaluate popular quadratic-cost OT solvers in highdimensional spaces (M4.3), including the image space of 64ˆ 64 CelebA faces (M4.4). • We evaluate the performance of these OT solvers as a loss in generative modeling of images (M4.5). Our experiments show that some OT solvers exhibit moderate error even in small dimensions (M4.3), performing similarly to trivial baselines (M4.2). The most successful solvers are those using parametrization via ICNNs. Surprisingly, however, solvers that faithfully recover W2 maps across dimensions struggle to achieve state-of-the-art performance in generative modeling. Our benchmark measures can be used to evaluate future W2 solvers in high-dimensional spaces, a crucial step to improve the transparency and replicability of continuous OT research. Note the benchmark from [35] does not fulfill this purpose, since it is designed to test discrete OT methods and uses discrete low-dimensional measures with limited support. Notation. We use P2pRq to denote the set of Borel probability measures on R with finite second moment and P2,acpRq to denote its subset of absolutely continuous probability measures. We denote by ΠpP,Qq the set of the set of probability measures on R ˆ R with marginals P and Q. For some measurable map T : R Ñ R, we denote by T 7 the associated push-forward operator. For φ : R Ñ R, we denote by φ its Legendre-Fenchel transform [10] defined by φpyq “ maxxPRD rxx, yy  ́ φpxqs. Recall that φ is a convex function, even when φ is not. 1 Background on Optimal Transport We start by stating the definition and some properties of optimal transport with quadratic cost. We refer the reader to [34, Chapter 1] for formal statements and proofs. Primal formulation. For P,Q P P2pRq, Monge’s primal formulation of the squared Wasserstein-2 distance, i.e., OT with quadratic cost, is given by W2pP,Qq def “ min T 7P“Q ż RD }x ́ T pxq} 2 dPpxq, (1) where the minimum is taken over measurable functions (transport maps) T : R Ñ R mapping P to Q. The optimal T ̊ is called the optimal transport map (OT map). Note that (1) is not symmetric, and this formulation does not allow for mass splitting, i.e., for some P,Q P P2pRq, there is no map T that satisfies T 7P “ Q. Thus, Kantorovich proposed the following relaxation [14]: W2pP,Qq def “ min πPΠpP,Qq ż RDˆRD }x ́ y} 2 dπpx, yq, (2) where the minimum is taken over all transport plans π, i.e., measures on R ˆ R whose marginals are P and Q. The optimal π ̊ P ΠpP,Qq is called the optimal transport plan (OT plan). If π ̊ is of the form ridRD , T ̊s7P P ΠpP,Qq for some T ̊, then T ̊ is the minimizer of (1). Dual formulation. For P,Q P P2pRq, the dual formulation of W2 is given by [40]: W2pP,Qq “ max f‘gď 12 } ̈}2 „ ż

text: Solving optimal transport (OT) with continuous methods has become widespread in machine learning, including methods for large-scale OT [11, 36] and the popular Wasserstein Generative Adversarial Network (W-GAN) [3, 12]. Rather than discretizing the problem [31], continuous OT algorithms use neural networks or kernel expansions to estimate transport maps or dual solutions. This helps scale OT to large-scale and higher-dimensional problems not handled by discrete methods. Notable successes of continuous OT are in generative modeling [42, 20, 19, 7] and domain adaptation [43, 37, 25].
In these applications, OT is typically incorporated as part of the loss terms for a neural network model. For example, in W-GANs, the OT cost is used as a loss function for the generator; the model incorporates a neural network-based OT solver to estimate the loss. Although recent W-GANs provide state-of-the-art generative performance, however, it remains unclear to which extent this
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
success is connected to OT. For example, [28, 32, 38] show that popular solvers for the Wasserstein-1 (W1) distance in GANs fail to estimate W1 accurately. While W-GANs were initially introduced with W1 in [3], state-of-the art solvers now use both W1 and W2 (the Wasserstein-2 distance, i.e., OT with the quadratic cost). While their experimental performance on GANs is similar, W2 solvers tend to converge faster (see [19, Table 4]) with better theoretical guarantees [19, 26, 16].
Contributions. In this paper, we develop a generic methodology for evaluating continuous quadraticcost OT solvers (W2). Our main contributions are as follows:
• We use input-convex neural networks (ICNNs [2]) to construct pairs of continuous measures that we use as a benchmark with analytically-known solutions for quadratic-cost OT (M3, M4.1).
• We use these benchmark measures to evaluate popular quadratic-cost OT solvers in highdimensional spaces (M4.3), including the image space of 64ˆ 64 CelebA faces (M4.4).
• We evaluate the performance of these OT solvers as a loss in generative modeling of images (M4.5).
Our experiments show that some OT solvers exhibit moderate error even in small dimensions (M4.3), performing similarly to trivial baselines (M4.2). The most successful solvers are those using parametrization via ICNNs. Surprisingly, however, solvers that faithfully recover W2 maps across dimensions struggle to achieve state-of-the-art performance in generative modeling.
Our benchmark measures can be used to evaluate future W2 solvers in high-dimensional spaces, a crucial step to improve the transparency and replicability of continuous OT research. Note the benchmark from [35] does not fulfill this purpose, since it is designed to test discrete OT methods and uses discrete low-dimensional measures with limited support.
Notation. We use P2pRDq to denote the set of Borel probability measures on RD with finite second moment and P2,acpRDq to denote its subset of absolutely continuous probability measures. We denote by ΠpP,Qq the set of the set of probability measures on RD ˆ RD with marginals P and Q. For some measurable map T : RD Ñ RD, we denote by T 7 the associated push-forward operator. For φ : RD Ñ R, we denote by φ its Legendre-Fenchel transform [10] defined by φpyq “ maxxPRD rxx, yy ´ φpxqs. Recall that φ is a convex function, even when φ is not.

heading: 1 Background on Optimal Transport
text: We start by stating the definition and some properties of optimal transport with quadratic cost. We refer the reader to [34, Chapter 1] for formal statements and proofs.
Primal formulation. For P,Q P P2pRDq, Monge’s primal formulation of the squared Wasserstein-2 distance, i.e., OT with quadratic cost, is given by
W22pP,Qq def“ min
T 7P“Q
ż
RD
}x´ T pxq}2
2 dPpxq, (1)
where the minimum is taken over measurable functions (transport maps) T : RD Ñ RD mapping P to Q. The optimal T˚ is called the optimal transport map (OT map). Note that (1) is not symmetric, and this formulation does not allow for mass splitting, i.e., for some P,Q P P2pRDq, there is no map T that satisfies T 7P “ Q. Thus, Kantorovich proposed the following relaxation [14]:
W22pP,Qq def“ min
πPΠpP,Qq
ż
RDˆRD
}x´ y}2
2 dπpx, yq, (2)
where the minimum is taken over all transport plans π, i.e., measures on RD ˆ RD whose marginals are P and Q. The optimal π˚ P ΠpP,Qq is called the optimal transport plan (OT plan). If π˚ is of the form ridRD , T˚s7P P ΠpP,Qq for some T˚, then T˚ is the minimizer of (1). Dual formulation. For P,Q P P2pRDq, the dual formulation of W22 is given by [40]:
W22pP,Qq “ max f‘gď 12 }¨}2
„ ż
RD fpxqdPpxq `
ż
RD gpyqdQpyq

, (3)
where the maximum is taken over all f P L1pP,RD Ñ Rq and g P L1pQ,RD Ñ Rq satisfying fpxq ` gpyq ď 12}x´ y}
2 for all x, y P RD. From the optimal dual potential f˚, we can recover the optimal transport plan T˚pxq “ x´∇f˚pxq [34, Theorem 1.17].
The optimal f˚, g˚ satisfy pf˚qc “ g˚ and pg˚qc “ f˚, where uc : RD Ñ R is the c´transform of u defined by ucpyq “ minxPRD “ 1{2}x´ y}2 ´ upxq ‰ . We can rewrite (3) as
W22pP,Qq “ max f
„ ż
RD fpxqdPpxq `
ż
RD f cpyqdQpyq

, (4)
where the maximum is taken over all f P L1pP,RD Ñ Rq. Since f˚ and g˚ are each other’s c-transforms, they are both c-concave [34, M1.6], which is equivalent to saying that functions ψ˚ : x ÞÑ 12}x} 2 ´ f˚pxq and φ˚ : x ÞÑ 12}x} 2 ´ g˚pxq are convex [34, Proposition 1.21]. In particular, ψ˚ “ φ˚ and φ˚ “ ψ˚. Since
T˚pxq “ x´∇f˚pxq “ ∇ ˆ }x}2
2 ´ f˚pxq
˙
“ ∇ψ˚, (5)
we see that the OT maps are gradients of convex functions, a fact known as Brenier’s theorem [6].
“Solving” optimal transport problems. In applications, for given P,Q P P2pRDq, the W2 optimal transport problem is typically considered in the following three similar but not equivalent tasks:
• Evaluating W22pP,Qq. The Wasserstein-2 distance is a geometrically meaningful way to compare probability measures, providing a metric on P2pRDq. • Computing the optimal map T˚ or plan π˚. The map T˚ provides an intuitive way to interpolate between measures. It is often used as a generative map between measures in problems like domain adaptation [36, 43] and image style transfer [16]. • Using the gradient BW22pPα,Qq{Bα to update generative models. Derivatives of W22 are used implicitly in generative modeling that incorporates W2 loss [19, 33], in which case P “ Pα is a parametric measure and Q is the data measure. Typically, Pα “ Gα7S is the measure generated from a fixed latent measure S by a parameterized function Gα, e.g., a neural network. The goal is to find parameters α that minimize W22pPα,Qq via gradient descent.
In the generative model setting, by definition of the pushforward Pα “ Gα7S, we have
W22pPα,Qq “ ż
z
f˚pGαpzqqdSpzq ` ż
RD g˚pyqdQpyq,
where f˚ and g˚ are the optimal dual potentials. At each generator training step, f˚ and g˚ are fixed so that when we take the gradient with respect to α, by applying the chain rule we have:
BW22pPα,Qq Bα “ ż
z
JαGαpzqT∇f˚ ` Gαpzq ˘ dSpzq, (6)
where JαGαpzqT is the transpose of the Jacobian matrix of Gαpzq w.r.t. parameters α. This result still holds without assuming the potentials are fixed by the envelope theorem [29]. To capture the gradient, we need a good estimate of ∇f˚ “ idRD ´T˚ by (5). This task is somewhat different from computing the OT map T˚: since the estimate of ∇f˚ is only involved in the gradient update for the generator, it is allowed to differ while still resulting in a good generative model.
We will use the generic phrase OT solver to refer to a method for solving any of the tasks above.
Quantitative evaluation of OT solvers. For discrete OT methods, a benchmark dataset [35] exists but the mechanism for producing the dataset does not extend to continuous OT. Existing continuous solvers are typically evaluated on a set of self-generated examples or tested in generative models without evaluating its actual OT performance. Two kinds of metrics are often used:
Direct metrics compare the computed transport map T̂ with the true one T˚, e.g., by using L2 Unexplained Variance Percentage (L2-UVP) metric [16, M5.1], [17, M5]. There are relatively few direct metrics available, since the number of examples of P,Q with known ground truth T˚ is small: it is known that T˚ can be analytically derived or explicitly computed in the discrete case [31, M3], 1-dimensional case [31, M2.6], and Gaussian/location-scatter cases [1].
Indirect metrics use an OT solver as a component in a larger pipeline, using end-to-end performance as a proxy for solver quality. For example, in generative modeling where OT is used as the generator loss [19, 27], the quality of the generator can be assessed through metrics for GANs, such as the Fréchet Inception distance (FID) [13]. Indirect metrics do not provide clear understanding about the quality of the solver itself, since they depend on components of the model that are not related to OT.

heading: 2 Continuous Dual Solvers for Quadratic Cost Transport
text: While our benchmark might be used to test any continuous solver which computes map T˚ or gradient ∇f˚, in this paper, we perform evaluation only on dual-form continuous solvers based on (3) or (4). Such solvers have straightforward optimization procedures and can be adapted to various datasets without extensive hyperparameter search. In contrast, primal-form solvers based on (1), e.g., [18, 43, 21, 23], typically parameterize T˚ using complicated generative modeling techniques that depend on careful hyperparameter search and complex optimization procedures [24].
We summarize existing continuous dual form solvers in Table 1. These fit a parametric function fθ (or ψθ) to approximate f˚ (or ψ˚ “ idRD ´ f˚). The resulting fθ produces an approximate OT map idRD´∇fθ“∇ψθ « T˚ and derivative ∇fθ“ idRD´∇ψθ needed to update generative models (6).
To our knowledge, none of these solvers has been quantitatively evaluated in a non-Gaussian setting. For tMMs, tMM-Bs, and tQCs, the quality of the recovered derivatives ∇f˚ for BW22pPα,Qq{Bα has only been evaluated implicitly through GAN metrics. Moreover, these three solvers have not been quantitatively evaluated on solving OT tasks. We now overview each solver from Table 1.
tLSs optimizes an unconstrained regularized dual form of (3) [36]:
max f,g
„ ż
RD fpxqdPpxq `
ż
RD gpyqdQpyq

´Rpf, gq. (7)
The entropic or quadratic regularizer R penalizes potentials f, g for violating the constraint f ‘ g ď 12} ¨ }
2 [36, M3]. In practice, f “ fθ and g “ gω are linear combinations of kernel functions [11] or neural networks [36]. The parameters θ, ω are obtained by applying stochastic gradient ascent (SGA) over random mini-batches sampled from P,Q. Most other solvers are based on an expansion of (4):
W22pP,Qq “ max f
ż
RD fpxqdPpxq `
ż
RD
“fcpyq hkkkkkkkkkkkkkkkikkkkkkkkkkkkkkkj
min xPRD
„
1 2 }x´ y}2 ´ fpxq

dQpyq. (8)
The challenge of (8) is the inner minimization over x P RD, i.e., evaluating f cpyq. The main difference between existing solvers is the procedure used to solve this inner problem.
tMM-Bs uses a neural network fθ as the potential trained using mini-batch SGA [27]. To solve the inner problem, the authors restrict the minimization of x to the current mini-batch from P instead of RD. The strategy is fast but leads to overestimation of the inner problem’s solution since the minimum is taken over a restricted subset.
tMM-v1s exploits the property that f˚ “ 12} ¨ } 2 ´ ψ˚, where ψ˚ is convex [39]. The authors parametrize fθ “ 12} ¨ } 2 ´ ψθ, where ψθ is an input convex neural network (ICNN) [2]. Hence, for every y P RD, the inner problem of (8) becomes convex in x. This problem can be solved using SGA to high precision, but doing so is computationally costly [16, MC.4].
tMMs uses a formulation equivalent to (8) [30]:
W22pP,Qq “ max f
ż
RD fpxqdPpxq `
ż
RD min H
„
1 2 }Hpyq ´ y}2 ´ fpHpyqq

dQpyq, (9)
where the minimization is performed over functions H : RD Ñ RD. The authors use neural networks fθ and Hω to parametrize the potential and the minimizer of the inner problem. To train θ, ω, the authors apply stochastic gradient ascent/descent (SGAD) over mini-batches from P,Q. tMMs is generic and can be modified to compute arbitrary transport costs and derivatives, not just W22, although the authors have tested only on the Wasserstein-1 (W1) distance. Similarly to tMMv1s, tMMv2s parametrizes fθ “ 12} ¨ }
2 ´ ψθ, where ψθ is an ICNN [26]. For a fixed fθ, the optimal solution H is given by H “ p∇ψθq´1 which is an inverse gradient of a convex function, so it is also a gradient of a convex function. Hence, the authors parametrize Hω “ ∇φω, where φω is an ICNN, and use tMMs to fit θ, ω. tW2s uses the same ICNN parametrization as [26] but introduces cycle-consistency regularization to avoid solving a maximin problem [16, M4].
Finally, we highlight the solver tQCs [19]. Similarly to tMM-Bs, a neural network fθ is used as the potential. When each pair of mini-batches txnu, tynu from P,Q is sampled, the authors solve a discrete OT problem to obtain dual variables tf˚n u, tg˚nu, which are used to regress fθpxnq onto f˚n . Gradient deviation. The solvers above optimize for potentials like fθ (or ψθ), but it is the gradient of fθ (or ψθ) that is used to recover the OT map via T “ x´∇fθ. Even if }f ´ f˚}2L2pPq is small, the difference }∇fθ ´∇f˚}2L2pPq may be arbitrarily large since ∇fθ is not directly involved in optimization process. We call this issue gradient deviation. This issue is only addressed formally for ICNN-based solvers tMMv1s, tMMv2s, tW2s [16, Theorem 4.1], [26, Theorem 3.6]. Reversed solvers. tMMs, tMMv2s, tW2s recover not only the forward OT map ∇ψθ « ∇ψ˚ “ T˚, but also the inverse, given by Hω « pT˚q´1 “ p∇ψ˚q´1 “ ∇ψ˚, see [26, M3] or [16, M4.1]. These solvers are asymmetric in P,Q and an alternative is to swap P and Q during training. We denote such reversed solvers by tMM:Rs, tMMv2:Rs, tW2:Rs. In M4 we show that surprisingly tMM:Rs works better in generative modeling than tMMs.

heading: 3 Benchmarking OT Solvers
text: In this section, we develop a generic method to produce benchmark pairs, i.e., measures pP,Qq such that Q “ T 7P with sample access and an analytically known OT solution T˚ between them. Key idea. Our method is based on the fact that for a differentiable convex function ψ : RD Ñ R, its gradient ∇ψ is an optimal transport map between any P P P2,acpRDq and its pushforward ∇ψ7P by ∇ψ : RD Ñ RD. This follows from Brenier’s theorem [6], [41, Theorem 2.12]. Thus, for a continuous measure P with sample access and a known convex ψ, pP,∇ψ7Pq can be used as a benchmark pair. We sample from ∇ψ7P by drawing samples from P and pushing forward by ∇ψ. Arbitrary pairs pP,Qq. It is difficult to compute the exact continuous OT solution for an arbitrary pair pP,Qq. As a compromise, we compute an approximate transport map as the gradient of an ICNN using tW2s. That is, we find ψθ parameterized as an ICNN such that ∇ψθ7P « Q. Then, the modified pair pP,∇ψθ7Pq can be used to benchmark OT solvers. We choose tW2s because it exhibits good performance in higher dimensions, but other solvers can also be used so long as ψθ is convex. Because of the choice of tW2s, subsequent evaluation might slightly favor ICNN-based methods. Extensions. Convex functions can be modified to produce more benchmark pairs. If ψ1, . . . , ψN are convex, then σpψ1, . . . , ψN q is convex when σ : RN Ñ R is convex and monotone. For example, c ¨ ψ1 (c ě 0q, ř
n ψn, maxn ψn are convex, and their gradients produce new benchmark pairs.
Inversion. If ∇ψθ is bijective, then the inverse transport map for pP,∇ψθ7Pq exists and is given by p∇ψθq´1. For each y P RD, the value p∇ψθq´1pyq can be obtained by solving a convex problem [39, M6], [16, M3]. All ICNNs ψθ we use have bijective gradients ∇ψθ, as detailed in Appendix B.1.

heading: 4 Benchmark Details and Results
text: We implement our benchmark in PyTorch and provide the pre-trained transport maps for all the benchmark pairs. The code is publicly available at
https://github.com/iamalexkorotin/Wasserstein2Benchmark The experiments are conducted on 4 GTX 1080ti GPUs and require about 100 hours of computation (per GPU). We provide implementation details in Appendix B.

heading: 4.1 Datasets
text: High-dimensional measures. We develop benchmark pairs to test whether the OT solvers can redistribute mass among modes of measures. For this purpose, we use Gaussian mixtures in dimensions D “ 21, 22, . . . , 28. In each dimension D, we consider a random mixture P of 3 Gaussians and two random mixtures Q1,Q2 of 10 Gaussians. We train approximate transport maps ∇ψi7P « Qi (i “ 1, 2) using the tW2s solver. Each potential is an ICNN with DenseICNN architecture [16, MB.2]. We create a benchmark pair via the half-sum of computed potentials pP, 12 p∇ψ1 `∇ψ2q7Pq. The first measure P is a mixture of 3 Gaussians and the second is obtained by averaging potentials, which transforms it to approximate mixtures of 10 Gaussians. See Appendix A.1 and Figure 1 for details.
Images. We use the aligned images of CelebA64 faces dataset1 [22] to produce additional benchmark pairs. First, we fit 3 generative models (WGAN-QC [19]) on the dataset and pick intermediate training checkpoints to produce continuous measures QkEarly,QkMid,QkLate for the first 2 models (k “ 1, 2) and the final checkpoint of the third model (k “ 3) to produce measure P3Final. To make measures absolutely continuous, we add small Gaussian noise to the generator’s output. Each checkpoint (Early, Mid, Late, Final) represents images of faces of a particular quality. Next, for k P t1, 2u and Cpkt P tEarly, Mid, Lateu, we use tW2s solver to fit an approximate transport map ∇ψkCpkt for the pair pP3Final,QkCpktq, i.e., ∇ψkCpkt7P3Final « QkCpkt. The potential ψkCpkt is a convolutional ICNN with ConvICNN64 architecture (MB.1). For each Cpkt, we define a benchmark pair pPCelebA,QCpktq def“pP3Final, rp∇ψ1Cpkt `∇ψ2Cpktq{2s7P3Finalq. See Appendix A.2 and Figure 2 for details.

heading: 4.2 Metrics and Baselines
text: Baselines. We propose three baseline methods: identity tIDs, constant tCs and linear tLs. The identity solver outputs T id “ idRD as the transport map. The constant solver outputs the mean value of Q, i.e., T 0 ” EQrys ” µQ. The linear solver outputs T 1pxq “ Σ ´ 12 P ` Σ 1 2 P ΣQΣ 1 2 P ˘ 1 2 Σ ´ 12 P px´ µPq ` µQ, i.e., the OT map between measures coarsened to Gaussians [1, Theorem 2.3].
Metrics. To assess the quality of the recovered transport map T̂ : RD Ñ RD from P to Q, we use unexplained variance percentage (UVP) [16]: L2-UVPpT̂ q def“ 100 ¨ }T̂ ´ T˚}2L2pPq{VarpQq%. Here T˚ is the OT map. For values « 0%, T̂ approximates T˚ well. For values ě 100%, map T̂ is far from optimal. The constant baseline provides L2-UVPpT 0q “ 100%.
1http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
pPCelebA,QCpktq def“ ` P3Final, 12 p∇ψ 1 Cpkt `∇ψ2Cpktq7P3Final ˘ . In the visualization, Cpkt is Early.
To measure the quality of approximation of the derivative of the potential ridRD ´ T̂ s « ∇f˚ that is used to update generative models (6), we use cosine similarity (cos):
cospid´ T̂ , id´ T˚q def“ xT̂ ´ id,∇ψ˚ ´ idyL2pPq
}T˚ ´ id}L2pPq ¨ }T̂ ´ id}L2pPq P r´1, 1s.
To estimate L2-UVP and cos metrics, we use 214 random samples from P.

heading: 4.3 Evaluation of Solvers on High-dimensional Benchmark Pairs
text: We evaluate the solvers on the benchmark and report the computed metric values for the fitted transport map. For fair comparison, in each method the potential f and the map H (where applicable) are parametrized as fθ “ 12} ¨ }
2 ´ ψθ and Hω “ ∇φω respectively, where ψθ, φω use DenseICNN architectures [16, MB.2]. In solvers tQCs, tLSs, tMM-Bs,tMMs we do not impose any restrictions on the weights θ, ω, i.e. ψθ, φω are usual fully connected nets with additional skip connections. We provide the computed metric values in Table 2 and visualize fitted maps (for D “ 64) in Figure 3. All the solvers perform well (L2-UVP« 0, cos « 1) in dimensionD “ 2. In higher dimensions, only tMMv1s, tMMs, tMMv2s, tW2s and their reversed versions produce reasonable results. However, tMMv1s solver is slow since each optimization step solves a hard subproblem for computing f c. Maximin solvers tMMs,tMMv2s,tMM:Rs are also hard to optimize: they either diverge from the start (Û) or diverge after converging to nearly-optimal saddle point (í). This behavior is typical for maximin optimization and possibly can be avoided by a more careful choice of hyperparameters.
For tQCs, tLSs,tMM-Bs, as the dimension increases, the L2-UVP drastically grows. Only tMM-Bs notably outperforms the trivial tLs baseline. The error of tMM-Bs is explained by the overestimation of the inner problem in (8), yielding biased optimal potentials. The error of tLSs comes from bias introduced by regularization [36]. In tQCs, error arises because a discrete OT problem solved on sampled mini-batches, which is typically biased [5, Theorem 1], is used to update fθ. Interestingly, although tQCs, tLSs are imprecise in terms of L2-UVP, they provide a high cos metric. Due to optimization issues and performance differences, wall-clock times for convergence are not representative. All solvers except tMMv1s converged in several hours. Among solvers that
substantially outperform the linear baseline, i.e. tMMs, tMMv1s, tMMv2s, tW2s, tMM-Bs, the fastest converging one is tMM-Bs, but it is biased. tMMs, tMMv2s, tW2s require more time.

heading: 4.4 Evaluation of Solvers in CelebA 64ˆ 64 Images Benchmark Pairs
text: For evaluation on the CelebA benchmark, we excluded tLSs and tMMv1s: the first is unstable in high dimensions [33], and the second takes too long to converge. ICNN-based solvers tMMv2s, tW2s and their reversed versions perform roughly the same in this experiment. For simplicity, we treat them as one solver tW2s. In tW2s, we parametrize fθ “ 12} ¨ }
2 ´ ψθ and Hω “ ∇φω, where ψθ, φω are input-convex neural nets with ConvexICNN64 architecture (MB.1). All the other solvers are designed in the generative modeling setting to work with convolutional architectures for images. Thus, in tMMs, tQCs, tMM-Bs we parametrize networks fθ as ResNet and Hω as U-Net (in tMMs). In turn, in tMM:Rs we parametrize Tθ by UNet and gω by ResNet. We compute the transport map QCpkt Ñ PCelebA for each solver on three image benchmarks. The results are in Figure 4 and Table 3 and echo patterns observed on high-dimensional problems (M4.3). tQCs, tMM-Bs suffer from extreme bias thanks to the high dimension of images, and the derivative of W22 computed by these solvers is almost orthogonal to the true derivative (cos « 0). This means that these solvers do not extract W22. tMMs, tMM:Rs, tW2s recover the transport maps well. tMMs’s map is slightly noisier than the one by tMM:Rs, a minor example of gradient deviation.

heading: 4.5 Evaluation of Solvers in Generative Modeling of CelebA 64ˆ 64 Faces
text: Based on our previous evaluation, many existing neural OT solvers are notably imprecise. This leads us to ask: To what extent does solver quality matter in real-world applications?
To address this question, we evaluate the most promising solvers in the task of generative modeling for CelebA 64 ˆ 64 images of faces. For comparison, we add tQCs, which has good generative performance [19]. For each solver, we train a generative network Gα with ResNet architecture from [19] to map a 128-dimensional normal distribution S to the data distribution Q. As the loss function for generator, we use W22pPα,Qq “W22pGα7S,Qq estimated by each solver. We perform GAN-style training, where gradient updates of the generator alternate with gradient steps of OT solver (discriminator) (MB.2.3). We show sample generated images in the top row of each subplot of Figure 5 and report FID [13]. On the bottom row, we show the pushforward of the OT map from Pα “ Gα7S to Q extracted from the OT solver. Since the model converged (Pα « Q), the map should be nearly equal to the identity. tW2s provides the least quality (Figure 5a). This can be explained by the use of ConvICNN: the other solvers use convolutional architectures and work better. In general, the applicability of ICNNs to image-based tasks is questionable [16, M5.3] which might be a serious practical limitation. tQCs has strong generative performance (Figure 5b). However, as in M4.3-4.4, the recovered map is far from the identity. We suspect this solver has decent generative performance because it approximates some non-W22 dissimilarity measure in practice. tMMs results in a generative model that produces blurry images (Figure 5c). The computed transport map idRD ´∇fθ is too far from the identity due to the gradient deviation. This leads to inaccurate gradient computation used to update the generator and explains why the generator struggles to improve. We emphasize that in M4.4 tMMs does not notably suffer from the gradient deviation. Probably, this is due to measures being absolutely continuous and supported on the entire RD. This is not the case in our generative modeling setup, where generated and data measures are supported on low-dimensional manifolds in RD.
Reversed tMM:Rs overcomes the problem of tMMs with the gradient deviation but still leads to blurry images (Figure 5d). Interestingly, the fitted transport map Tθ significantly improves the quality and images Tθ ˝Gαpzq are comparable to the ones with tQCs solver (Figure 5b). We emphasize that formulations from tMMs, tMM:Rs solvers are maximin: using them in GANs requires solving a challenging min-max-min optimization problem. To handle this, we use three nested loops and stochastic gradient descent-ascent-descent. In our experiments, the training was not stable and often diverged: the reported results use the best hyperparameters we found, although there may exist better ones. The difficulty in selecting hyperparameters and the unstable training process are limitations of these solvers that need to be addressed before using in practice.

heading: 5 Conclusion
text: Our methodology creates pairs of continuous measures with ground truth quadratic-cost optimal transport maps, filling the missing gap of benchmarking continuous OT solvers. This development allows us to evaluate the performance of quadratic-cost OT solvers in OT-related tasks. Beyond benchmarking the basic transport problem, our study of generative modeling reveals surprising patterns: bad OT solvers can yield good generative performance, and simply reversing asymmetric solvers can affect performance dramatically. Limitations. We rely on ICNN gradients as W2 optimal transport maps to generate pairs of benchmark measures. It is unclear whether analogous constructions can be used for other costs such as W1. We also limit our benchmark pairs to be absolutely continuous measures while limiting the ground truth transport maps to be gradients of ICNNs, which may not have enough representational power. While we reveal a discrepancy between performance in OT-related tasks and performance in generative modeling, in-depth study is needed to answer questions such as what exact dissimilarity metric tQCs implies that explains its generative performance while poorly approximating W2. Potential impact. We expect our benchmark to become a standard benchmark for continuous optimal transport as part of the ongoing effort of advancing computational OT, in particular, in its application to generative modeling. As a result, we hope our work can improve the quality and reusability of OT-related research. One potential negative is that our benchmark might narrow the evaluation of future OT solvers to the datasets of our benchmark. To avoid this, besides particular benchmark datasets, in M3 we describe a generic method to produce new benchmark pairs. ACKNOWLEDGEMENTS. The problem statement was developed in the framework of Skoltech-MIT NGP program. The work of Evgeny Burnaev was supported by the Ministry of Science and Higher Education of the Russian Federation grant No. 075-10-2021-068. The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grants W911NF2010168 and W911NF2110293, of Air Force Office of Scientific Research award FA9550-19-1-031, of National Science Foundation grants IIS-1838071 and CHS-1955697, from the CSAIL Systems that Learn program, from the MIT–IBM Watson AI Laboratory, from the Toyota–CSAIL Joint Research Center, from a gift from Adobe Systems, from an MIT.nano Immersion Lab/NCSOFT Gaming Program seed grant, and from the Skoltech–MIT Next Generation Program.

Review summaries:
The authors present a continuous Wasserstein-2 benchmark using input convex neural networks to obtain ground truth continuous OT maps. They evaluate continuous OT solvers based on how well they reproduce the ground truth transport map, how well they approximate the derivative of the potential, and how well they perform in a GAN framework. This provides a basis to evaluate current and future continuous 2-Wasserstein solvers.

This work introduces a continuous Wasserstein-2 benchmark to assess the qualities of different neural optimal transport solvers. The work uses input convex neural networks (ICNN) to construct pairs of measures whose ground truth OT maps can be obtained analytically. This strategy yields pairs of continuous benchmark measures in high-dimensional spaces such as spaces of images. The authors thoroughly evaluate existing optimal transport solvers. The study reveals crucial limitations of existing solvers.