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

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 . 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 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 . 2 Continuous Dual Solvers for Quadratic Cost Transport . 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 .

Summary: 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.