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: The sliced Wasserstein distance (SWD) is a metric on the space of probability measures and serves as a practical alternative to the standard metric arising from the optimal transport (OT) problem, the Wasserstein distance. Indeed, while the Wasserstein distance suffers from computational and statistical limitations on large-scale and high-dimensional settings, SWD can alleviate these issues while offering similar theoretical properties. Hence, SWD has been increasingly popular within the statistical and machine learning (ML) community in recent years.  SWD is defined as an expectation that is intractable in general; in practice, it is then usually approximated with a simple Monte Carlo estimation, but prior studies reported that this method might induce an important approximation error on high-dimensional settings, leading to an underestimation of the "true" dissimilarity between two distributions. Therefore, one would need to compute a Monte Carlo approximation based a very large amount of samples to obtain an accurate approximation, which increases the computational complexity of SWD.   Several variants of SWD have recently been introduced to overcome this problem and include the "generalized sliced Wasserstein distances" (GSWD): this class of distances aim at offering a better computational efficiency, by comparing distributions from their nonlinear one-dimensional representations (instead of the linear ones used in SWD). This paper focuses on this line of work: the limitations of GSWD are identified to motivate the formulation of a novel class of distances, called "Augmented Sliced Wasserstein distances" (ASWD).  Similarly to SWD and GSWD, the definition of ASWD is connected to the Radon transform: the authors introduce a novel type of Radon transform, the "spatial Radon transform" (Definition 1), and use it to define the class of augmented sliced Wasserstein distances.  Then, they study the theoretical properties of the spatial Radon transform and ASWD. Their main result states that ASWD verify all metric axioms if and only if the mapping $g$ (that characterizes the underlying spatial Radon transform) (1) is injective (Theorem 1), or (2) maximizes ASWD over the space of bounded and injective functions (Corollary 1.1). The authors then introduce a specific class of mappings (Equation (15)) which meet the conditions of Corollary 1.1 thus guarantee that the resulting ASWD is a metric. Such mappings require training a neural network to make ASWD data-adaptive. Finally, the superior empirical performance of ASWD against existing variants of SWD (including GSWD) is illustrated on various problems, including gradient flows, generative modeling, color transferring and barycenters of measures.