Augmented Sliced Wasserstein Distances

1 INTRODUCTION . Comparing samples from two probability distributions is a fundamental problem in statistics and machine learning . The optimal transport ( OT ) theory ( Villani , 2008 ) provides a powerful and flexible theoretical tool to compare degenerative distributions by accounting for the metric in the underlying spaces . The Wasserstein distance , which arises from the optimal transport theory , has become an increasingly popular choice in various machine learning domains ranging from generative models to transfer learning ( Gulrajani et al. , 2017 ; Arjovsky et al. , 2017 ; Kolouri et al. , 2019b ; Cuturi and Doucet , 2014 ; Courty et al. , 2016 ) . Despite its favorable properties , such as robustness to disjoint supports and numerical stability ( Arjovsky et al. , 2017 ) , the Wasserstein distance suffers from high computational complexity especially when the sample size is large . Besides , the Wasserstein distance itself is the result of an optimization problem — it is non-trivial to be integrated into an end-to-end training pipeline of deep neural networks , unless one can make the solver for the optimization problem differentiable . Recent advances in computational optimal transport methods focus on alternative OT-based metrics that are computationally efficient and solvable via a differentiable optimizer ( Peyré and Cuturi , 2019 ) . Entropy regularization is introduced in the Sinkhorn distance ( Cuturi , 2013 ) and its variants ( Altschuler et al. , 2017 ; Dessein et al. , 2018 ) to smooth the optimal transport problem ; as a result , iterative matrix scaling algorithms can be applied to provide significantly faster solutions with improved sample complexity ( Genevay et al. , 2019 ) . An alternative approach is to approximate the Wasserstein distance through slicing , i.e . linearly projecting , the distributions to be compared . The sliced Wasserstein distance ( SWD ) ( Bonneel et al. , 2015 ) is defined as the expected value of Wasserstein distances between one-dimensional random projections of high-dimensional distributions . The SWD shares similar theoretical properties with the Wasserstein distance ( Bonnotte , 2013 ) and is computationally efficient since the Wasserstein distance in one-dimensional space has a closed-form solution based on sorting . ( Deshpande et al. , 2019 ) extends the sliced Wasserstein distance to the max-sliced Wasserstein distance ( Max-SWD ) , by finding a single projection direction with the maximal distance between projected samples . The subspace robust Wasserstein distance extends the idea of slicing to projecting distributions on linear subspaces ( Paty and Cuturi , 2019 ) . However , the linear nature of these projections usually leads to low projection efficiency of the resulted metrics in high-dimensional spaces ( Deshpande et al. , 2019 ; Kolouri et al. , 2019a ) . Different variants of the SWD have been proposed to improve the projection efficiency of the SWD , either by introducing nonlinear projections or by optimizing the distribution of random projections . Specifically , ( Kolouri et al. , 2019a ) extends the connection between the sliced Wasserstein distance and the Radon transform ( Radon , 1917 ) to introduce generalized sliced Wasserstein distances ( GSWDs ) by utilizing generalized Radon transforms ( GRTs ) , which are defined by nonlinear defining functions and lead to nonlinear projections . A variant named the GSWD-NN was proposed in ( Kolouri et al. , 2019a ) to generate nonlinear projections directly with neural network outputs , but it does not fit into the theoretical framework of the GSWD and does not guarantee a valid metric . In contrast , the distributional sliced Wasserstein distance ( DSWD ) and its nonlinear version , the distributional generalized sliced Wasserstein distance ( DGSWD ) , improve their projection efficiency by finding a distribution of projections that maximizes the expected distances over these projections . The GSWD and the DGSWD exhibit higher projection efficiency than the SWD in the experiment evaluation , yet they require the specification of the particular form of defining functions from a limited class of candidates . However , the selection of defining functions is usually a task-dependent problem and requires domain knowledge , and the impact on performance from different defining functions is still unclear . In this paper , we present the augmented sliced Wasserstein distance ( ASWD ) , a distance metric constructed by first mapping samples to hypersurfaces in an augmented space , which enables flexible nonlinear slicing of data distributions for improved projection efficiency ( See Figure 1 ) . Our main contributions include : ( i ) We exploit the capacity of nonlinear projections employed in the ASWD by constructing injective mapping with arbitrary neural networks ; ( ii ) We prove that the ASWD is a valid distance metric ; ( iii ) We provide a mechanism in which the hypersurface where high-dimensional distributions are projected onto can be optimized and show that the optimization of hypersurfaces can help improve the projection efficiency of slice-based Wasserstein distances . Hence , the ASWD is data-adaptive , i.e . the hypersurfaces can be learned from data . This implies one does not need to manually design a function from the limited class of candidates ; ( iv ) We demonstrate superior performance of the ASWD in numerical experiments for both synthetic and real-world datasets . The remainder of the paper is organized as follows . Section 2 reviews the necessary background . We present the proposed method and its numerical implementation in Section 3 . Related work are discussed in Section 4 . Numerical experiment results are presented and discussed in Section 5 . We conclude the paper in Section 6 . 2 BACKGROUND . In this section , we provide a brief review of concepts related to the proposed work , including the Wasserstein distance , ( generalized ) Radon transform and ( generalized ) sliced Wasserstein distances . Wasserstein distance : Let Pk ( Ω ) be a set of Borel probability measures with finite k-th moment on a Polish metric space ( Ω , d ) ( Villani , 2008 ) . Given two probability measures µ , ν ∈ Pk ( Ω ) , the Wasserstein distance of order k∈ [ 1 , +∞ ) between µ and ν is defined as : Wk ( µ , ν ) = ( inf γ∈Γ ( µ , ν ) ∫ Ω×Ω d ( x , y ) kdγ ( x , y ) ) 1 k , ( 1 ) where d ( · , · ) k is the cost function , Γ ( µ , ν ) represents the set of all transportation plans γ , i.e . joint distributions whose marginals are µ and ν , respectively . While the Wasserstein distance is generally intractable for high-dimensional distributions , there are several favorable cases where the optimal transport problem can be efficiently solved . If µ and ν are continuous one-dimensional measures defined on a linear space equipped with theLk norm , the Wasserstein distance between µ and ν has a closed-form solution ( Peyré and Cuturi , 2019 ) : Wk ( µ , ν ) = ( ∫ 1 0 |F−1µ ( z ) −F−1ν ( z ) |kdz ) 1 k , ( 2 ) where F−1µ and F −1 ν are inverse cumulative distribution functions ( CDFs ) of µ and ν , respectively . In practice , Wasserstein distances Wk ( µ̃ , ν̃ ) between one-dimensional empirical distributions µ̃= 1N ∑N n=1δxn and ν̃ = 1 N ∑N n=1δyn can be computed by sorting one-dimensional samples from empirical distributions , and evaluating the distances between sorted samples ( Kolouri et al. , 2019b ) : Wk ( µ̃ , ν̃ ) = ( 1 N N∑ n=1 |xIx [ n ] −yIy [ n ] | k ) 1 k , ( 3 ) where N is the number of samples , Ix [ n ] and Iy [ n ] are the indices of sorted samples satisfying xIx [ n ] ≤xIx [ n+1 ] and yIy [ n ] ≤yIy [ n+1 ] , respectively . Radon transform and generalized Radon transform : The Radon transform ( Radon , 1917 ) maps a function f ( · ) ∈L1 ( Rd ) to the space of functions defined over spaces of hyperplanes in Rd . The Radon transform of f ( · ) is defined by line integrals of f ( · ) along all possible hyperplanes in Rd : Rf ( t , θ ) = ∫ Rd f ( x ) δ ( t−〈x , θ〉 ) dx , ( 4 ) where t∈R and θ ∈ Sd−1 represent the parameters of hyperplanes { x∈Rd | 〈x , θ〉= t } , δ ( · ) is the one-dimensional Dirac delta function , and 〈· , ·〉 refers to the Euclidean inner product . By replacing the inner product 〈x , θ〉 in Equation ( 4 ) with β ( x , θ ) , a specific family of functions named as defining function in ( Kolouri et al. , 2019a ) , the generalized Radon transform ( GRT ) ( Beylkin , 1984 ) is defined as integrals of f ( · ) along hypersurfaces defined by { x∈Rd | β ( x , θ ) = t } : Gf ( t , θ ) = ∫ Rd f ( x ) δ ( t−β ( x , θ ) ) dx , ( 5 ) where t∈R , θ∈Ωθ and Ωθ is a compact set of all feasible θ , e.g . Ωθ =Sd−1 for β ( x , θ ) = 〈x , θ〉 . In particular , a function β ( x , θ ) defined on X × ( Rd\ { 0 } ) with X ⊆Rd is called a defining function of GRTs if it satisfies conditions H.1 – H.4 given in ( Kolouri et al. , 2019a ) . For probability measures µ ∈ Pk ( Rd ) , the Radon transform and the GRT can be employed as push-forward operators , and the generated push-forward measuresRµ=R # µ , Gµ=G # µ are defined as follows ( Bonneel et al. , 2015 ) : Rµ ( t , θ ) = ∫ Rd δ ( t−〈x , θ〉 ) dµ , ( 6 ) Gµ ( t , θ ) = ∫ Rd δ ( t−β ( x , θ ) ) dµ . ( 7 ) Notably , the Radon transform is a linear bijection ( Helgason , 1980 ) , and the sufficient conditions for GRTs to be bijective are provided in ( Homan and Zhou , 2017 ) . Sliced Wasserstein distance and generalized sliced Wasserstein distance : By applying the Radon transform to µ and ν to obtain multiple projections , the sliced Wasserstein distance ( SWD ) decomposes the high-dimensional Wasserstein distance into multiple one-dimensional Wasserstein distances which can be efficiently evaluated ( Bonneel et al. , 2015 ) . The k-SWD between µ and ν is defined by : SWDk ( µ , ν ) = ( ∫ Sd−1 W kk ( Rµ ( · , θ ) , Rν ( · , θ ) ) dθ ) 1 k , ( 8 ) where the Radon transformR defined by Equation ( 6 ) is adopted as the measure push-forward operator . The GSWD generalizes the idea of SWD by slicing distributions with hypersurfaces rather than hyperplanes ( Kolouri et al. , 2019a ) . The GSWD is defined as : GSWDk ( µ , ν ) = ( ∫ Ωθ W kk ( Gµ ( · , θ ) , Gν ( · , θ ) ) dθ ) 1 k , ( 9 ) where the GRT G defined by Equation ( 7 ) is used as the measure push-forward operator . From Equation ( 3 ) , withL random projections andN samples , the SWD and GSWD between µ and ν can be approximated by : SWDk ( µ , ν ) ≈ ( 1 NL L∑ l=1 N∑ n=1 |〈xIlx [ n ] , θl〉−〈yIly [ n ] , θl〉| k ) 1 k , ( 10 ) GSWDk ( µ , ν ) ≈ ( 1 NL L∑ l=1 N∑ n=1 |β ( xIlx [ n ] , θl ) −β ( yIly [ n ] , θl ) | k ) 1 k , ( 11 ) where I lx and I l y are sequences consisting of the indices of sorted samples which satisfy 〈xIlx [ n ] , θl〉≤ 〈xIlx [ n+1 ] , θl〉 , 〈yIly [ n ] , θl〉 ≤ 〈yIly [ n+1 ] , θl〉 in the SWD , and β ( xIlx [ n ] , θl ) ≤ β ( xIlx [ n+1 ] , θl ) , β ( yIly [ n ] , θl ) ≤ β ( yIly [ n+1 ] , θl ) in the GSWD . The approximation error in estimating SWDs using Equation ( 10 ) is derived in ( Nadjahi et al. , 2020 ) . It is proved in ( Bonnotte , 2013 ) that the SWD is a valid distance metric . The GSWD is a valid metric except for its neural network variant ( Kolouri et al. , 2019a ) .

Summary: This paper introduces _Augmented Sliced Wasserstein Distances_ (ASWD) which consist, in a nutshell, in computing the usual SWD between pushed measures: $\mathrm{ASWD}(\mu,\nu; g) = \mathrm{SWD}(g \circ \mu, g \circ \nu)$ (note: the pushforward operator does not render properly, so I will use $\circ$ instead). The idea is then to maximize on $g$ (under constraint/regularization) to obtain a pushforward map $g$ that provides the most discriminative projection directions. Authors provide an exhaustive condition on $g$ (namely, injectivity) to make the ASWD a metric between probability measures. Their approach is then showcased on a various set of benchmark experiments in computational optimal transport (gradient flows, generative models, auto-encoder, barycenters, color transfert...).