Efficient estimates of optimal transport via low-dimensional embeddings

1 Introduction . Optimal Transport metrics ( Kantorovich , 1960 ) or Wasserstein distances , have emerged successfully in the field of machine learning , as outlined in the review by Peyré et al . ( 2017 ) . They provide machinery to lift distances on X to distances over probability distributions in P ( X ) . They have found multiple applications in machine learning : domain adaptation Courty et al . ( 2017 ) , density estimation ( Bassetti et al. , 2006 ) and generative networks ( Genevay et al. , 2017 ; Patrini et al. , 2018 ) . However , it is prohibitively expensive to compute OT between distributions with support in a high-dimensional space and might not even be practically possible as the sample complexity can grow exponentially as shown by Dudley ( 1969 ) . Similarly , work by Weed et al . ( 2019 ) showed a theoretical improvement when the support of distributions is found in a low-dimensional space . Furthermore , picking the ground metric that one should use is not obvious when using high-dimensional data . One of the earlier ideas from Santambrogio ( 2015 ) showed that OT projections in a 1-D space may be sufficient enough to extract geometric information from high dimensional data . This further prompted Kolouri et al . ( 2018 ) to use this method to build generative models , namely the Sliced Wasserstein Autoencoder . Following a similar approach Paty & Cuturi ( 2019 ) and Muzellec & Cuturi ( 2019 ) project the measures into a linear subspace E of low-dimension k that maximizes the transport cost and show how this can be used in applications of color transfer and domain adaptation . This can be seen as an extension to earlier work by Cuturi & Doucet ( 2014 ) whereby the cost function is parameterized . One of the fundamental innovations that made OT appealing to the machine learning com- munity was the seminal paper by Cuturi ( 2013 ) that introduced the idea of entropic regularization of OT distances and the Sinkhorn algorithm . Since then , regularized OT has been successfully used as a loss function to construct generative models such as GANs ( Genevay et al. , 2017 ) or RBMs ( Montavon et al. , 2015 ) and computing Barycenters ( Cuturi & Doucet , 2014 ; Claici et al. , 2018 ) . More recently , the new class of Sinkhorn Divergences was shown by Feydy et al . ( 2018 ) to have good geometric properties , and interpolate between Maximum Mean Discrepancies ( MMD ) and OT . Building on this previous work , we introduce a general framework for approximating highdimensional OT using low-dimensional projections f by finding the subspace with the worst OT cost , i.e . the one maximizing the ground cost on the low-dimensional space . By taking a general family of parameterizable fφs that are 1-Lipschitz , we show that our method generates a pseudo-metric and is computationally efficient and robust . We start the paper in §2 with background on optimal transport and pseudo-metrics . In §3 we define the theoretical framework for approximating OT distances and show how both linear ( Paty & Cuturi , 2019 ) and non-linear projections can be seen as a special instance of our framework . In §4 we present an efficient algorithm for computing OT distances using Sinkhorn Divergences and fφs that are 1-Lipschitz under the L2 norm . We conclude in §5 with experiments illustrating the efficiency and robustness of our method . 2 Preliminaries . We start with a brief reminder of the basic notions needed for the rest of the paper . Let X be a set equipped with a map dX : X × X → R≥0 with non-negative real values . The pair ( X , dX ) is said to be a metric space and dX is said to be a metric on X if it satisfies the usual properties : • dX ( x , y ) = 0 if and only if x = y • dX ( x , y ) = dX ( y , x ) • dX ( x , z ) ≤ dX ( x , y ) + dX ( y , z ) If dX verifies the above except for the only if condition , it is called a pseudo-metric , and ( X , dX ) is said to be a pseudo-metric space . For a pseudo-metric , it may be that dX ( x , y ) = 0 while x 6= y . We write dX ≤ d′X if for all x , y dX ( x , y ) ≤ d′X ( x , y ) . It is easy to see that : 1 ) “ ≤ ” is a partial order on pseudo-metrics over X , 2 ) “ ≤ ” induces a complete lattice structure on the set of pseudo-metrics over X , where 3 ) suprema are computed pointwise ( but not infima ) . Consider X , Y , two metric spaces equipped with respective metrics dX , dY . A map f from X to Y is said to be α-Lipschitz continuous if dY ( f ( x ) , f ( x′ ) ) ≤ αdX ( x , x′ ) . A 1-Lipschitz map is also called non-expansive . Given a map f from X to Y one defines the pullback of dY along f as : f̂ ( dY ) ( x , x′ ) = dY ( f ( x ) , f ( x′ ) ) ( 1 ) It is easily seen that : 1 ) f̂ ( dY ) is a pseudo-metric on X , 2 ) f̂ ( dY ) is a metric iff f is injective , 3 ) f̂ ( dY ) ≤ dX iff f is non-expansive , 4 ) f̂ ( dY ) is the least pseudo-metric on the set X such that f is non-expansive from ( X , f̂ ( dY ) ) to ( X , dY ) . Thereafter , we assume that all metric spaces considered are complete and separable , i.e . have a dense countable subset . Let ( X , dX ) be a ( complete separable ) metric space . Let ΣX be the σ-algebra generated by the open sets of X ( aka the Borelian subsets ) . We write P ( X ) for the set of probability distributions on ( X , ΣX ) . Given a measurable map f : X → Y , and µ ∈P ( X ) one defines the push-forward of µ along f as : f # ( µ ) ( B ) = µ ( f−1 ( B ) ) ( 2 ) for B ∈ ΣY . It is easily seen that f # ( µ ) is a probability measure on ( Y , ΣY ) Given µ in P ( X ) , ν in P ( Y ) , a coupling of µ and ν is a probability measure γ over X × Y such that for all A in ΣX , B in ΣY , γ ( A×X ) = µ ( A ) , and γ ( X ×B ) = ν ( B ) . Equivalently , µ = π0 # ( π ) , and ν = π1 # ( π ) for π0 , π1 the respective projections . We write Γ ( µ , ν ) for the set of couplings of µ and ν . There are several ways to lift a given metric structure on dX to one on P ( X ) . We will be specifically interested in metrics on P ( X ) derived from optimal transport problems . The p-Wasserstein metric with p ∈ [ 1 , ∞ ) is defined by : Wp ( dX ) ( µ , ν ) p = inf γ∈Γ ( µ , ν ) ∫ X×X dpX dγ ( 3 ) Villani ( 2008 ) establishes that if dX is ( pseudo- ) metric so isWp ( dX ) . The natural ‘ Dirac ’ embedding of X into P ( X ) is isometric ( there is only one coupling ) . The idea behind the definition is that dpX is used as a measure of the cost of transporting units of mass in X , while a coupling γ specifies how to transport the µ distribution to the ν one . One can therefore compute the mean transportation cost under γ , and pick the optimal γ . Hence the name optimal transport . In most of the paper , we are concerned with the case X = Rd+ for some large d with a metric structure dX given by the Euclidean norm , and we wish to compute the W2 metric between distributions with finite support . Since OT metrics are costly to compute in high dimension , to estimate these efficiently , and mitigate the impact of dimension , we will use a well-chosen family of fs to push the data along a map with a low dimensional co-domain Y also equipped with the Euclidean metric . The reduction maps may be linear or not . They have to be non-expansive to guarantee that the associated pull-back metrics are always below the Euclidean one , and therefore we provide a lower estimate of W2 ( d2 ) . 3 Approximate OT with General Projections - GPW . With the ingredients from the above section in place , we can now construct a general framework for approximating Wasserstein-like metrics by low-dimensional mappings of X . We write simply W instead of Wp as the value of p plays no role in the development . Pick two metric spaces ( X , dX ) , ( Y , dY ) , and a family S = ( fφ : X → Y ; φ ∈ S ) of mappings from X to Y . Define a map from P ( X ) ×P ( X ) to non-negative reals as follows : dS ( µ , ν ) = sup S W ( dY ) ( fφ # ( µ ) , fφ # ( ν ) ) ( 4 ) Equivalently and more concisely dS can be defined as : dS ( µ , ν ) = sup φ W ( f̂φ ( dY ) ) ( µ , ν ) ( 5 ) It is easily seen that : 1. the two definitions are equivalent 2. dS is a pseudo-metric on P ( X ) 3. dS is a metric ( not just a pseudo one ) if the family fφ jointly separates points in X , and 4. if the fφs are non-expansive from ( X , dX ) to ( Y , dY ) , then dS ≤W ( dX ) The second point follows readily from the second definition . Each f̂φ ( dY ) is a pseudo-metric on X obtained by pulling back dY ( see preceding section ) , hence , so is W ( f̂φ ( dY ) ) on P ( X ) , and therefore dS being the supremum of this family ( in the lattice of pseudo-metrics over X ) is itself a pseudo-metric . The first definition is important because it allows one to perform the OT computation in the target space where it will be cheaper . Thus we have derived from S a pseudo-metric dS on the space of probability measures P ( X ) . We assume from now on that mappings in S are non-expansive . By point 4. above , we know that dS is bounded above by W ( dX ) . We call dS the generalized projected Wasserstein metric ( GPW ) associated to S. In good cases , it is both cheaper to compute and a good estimate .

Summary: This paper uses the fact that the Wasserstein distance is decreasing under 1-Lipschitz mappings from the ambient space to a feature space in order to propose more robust (to dimensionality ?) estimation of the Wasserstein distance between two probability distributions. A neural network with some sort of weight renormalization is used to produce 1-Lipschitz embeddings. The authors then maximise over the parametrized maps.