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 author proposed a variant of sliced Wasserstein distance (SWD) called **augmented sliced Wasserstein distance** (ASWD) which replace the linear projections in SWD with a non-linear one. The construction of ASWD consists of two steps: (1) define spatial Radon transform (SRT) for probability measures (2) replace the Radon transform in the original SWD with this new SRT. The key intuition of SRT is to first transform the function $\pmb{x}$ to high dimensional hyperspace by transformation $g(\pmb{x})$ and then perform Radon transform on functions defined in this hyperspace.   Theoretically, the author proves the conditions required for SRT to be injective. Later, the author uses this result to show the corresponding ASWD is a distance metric. In addition, the author proposes a training objective for finding the transformation $g(\cdot)$.   Empirically, the author conducts both toy and practical experiments to demonstrate the advantages of ASWD in terms of higher projection efficiency and also provides some ablation studies in the appendix to investigate the effects of different $g(\cdot)$ choices.