H IERARCHICAL S LICED W ASSERSTEIN DISTANCE
A BSTRACT
Sliced Wasserstein (SW) distance has been widely used in different application scenarios since it can be scaled to a large number of supports without suffering from the curse of dimensionality.
The value of sliced Wasserstein distance is the average of transportation cost between one-dimensional representations (projections) of original measures that are obtained by Radon Transform (RT).
Despite its efficiency in the number of supports, estimating the sliced Wasserstein requires a relatively large number of projections in high-dimensional settings.
Therefore, for applications where the number of supports is relatively small compared with the dimension, e g., several deep learning applications where the mini-batch approaches are utilized, the complexities from matrix multiplication of Radon Transform become the main computational bottleneck.
To address this issue we propose to derive projections by linearly and randomly combining a smaller number of projections which are named bottleneck projections .
We explain the usage of these projections by introducing Hierarchical Radon Transform (HRT) which is constructed by applying
Radon Transform variants recursively.
We then formulate the approach into a new metric between measures, named Hierarchical Sliced Wasserstein (HSW) distance.
By proving the injectivity of HRT, we derive the metricity of HSW.
Moreover, we investigate the theoretical properties of HSW including its connection to SW variants and its computational and sample complexities.
Finally, we compare the computational cost and generative quality of HSW with the conventional SW on the task of deep generative modeling using various benchmark datasets including
CIFAR10, CelebA, and Tiny ImageNet 1 .
I NTRODUCTION
Wasserstein distance (Villani, 2008; Peyr ´ e & Cuturi, 2020) has been widely used in applications, such as generative modeling on images (Arjovsky et al., 2017; Tolstikhin et al., 2018; litu Rout et al., 2022), domain adaptation to transfer knowledge from source to target domains (Courty et al., 2017; Bhushan Damodaran et al., 2018), clustering problems (Ho et al., 2017), and various other applications (Le et al., 2021; Xu et al., 2021; Yang et al., 2020).
Despite the increasing importance of
Wasserstein distance in applications, prior works have alluded to the concerns surrounding the high computational complexity of that distance.
When the probability measures have at most n supports, the computational complexity of Wasserstein distance scales with the order of O ( n 3 log n ) (Pele &
Werman, 2009).
Additionally, it suffers from the curse of dimensionality, i.e., its sample complexity (the bounding gap of the distance between a probability measure and the empirical measures from its random samples) is of the order of O ( n − 1 /d )
(Fournier & Guillin, 2015), where n is the sample size and d is the number of dimensions.
Over the years, numerous attempts have been made to improve the computational and sample complexities of the Wasserstein distance.
One primal line of research focuses on using entropic regularization (Cuturi, 2013).
This variant is known as entropic regularized optimal transport (or in short entropic regularized Wasserstein).
By using the entropic version, one can approximate the
Wasserstein distance with the computational complexities O ( n 2 ) (Altschuler et al., 2017; Lin et al., 2019b;a; 2020) (up to some polynomial orders of approximation errors).
Furthermore, the sample complexity of the entropic version had also been shown to be at the order of O ( n − 1 / 2 ) (
Mena &
Weed, 2019), which indicates that it does not suffer from the curse of dimensionality.
Another line of work builds upon the closed-form solution of optimal transport in one dimension.
A notable distance metric along this direction is sliced Wasserstein (SW) distance (Bonneel et al., 2015).
SW is defined between two probability measures that have supports belonging to a vector space, e.g, R d .
As defined in (Bonneel et al., 2015), SW is written as the expectation of one-dimensional
Wasserstein distance between two projected measures over the uniform distribution on the unit sphere.
Due to the intractability of the expectation, Monte Carlo samples from the uniform distribution over the unit sphere are used to approximate SW distance.
The number of samples is often called the number of projections that is denoted as L .
On the computational side, the projecting directions matrix of size d × L is sampled and then multiplied by the two data matrices of size n × d resulting in two matrices of size n × L that represent L one-dimensional projected probability measures.
Thereafter, L one-dimensional Wasserstein distances are computed between the two corresponding projected measures with the same projecting direction.
Finally, the average of those distances yields an approximation of the value of the sliced Wasserstein distance.
Prior works (Kolouri et al., 2018a; Deshpande et al., 2018; 2019; Nguyen et al., 2021a;b) show that the number of projections L should be large enough compared to the dimension d for a good performance of the SW.
Despite the large L , SW has lots of benefits in practice.
It can be computed in O ( n log 2 n ) time, with the statistical rate O ( n − 1 / 2 ) that does not suffer from the curse of dimensionality, while becoming more memory efficient 2 compared with the vanilla Wasserstein distance.
For these reasons, it has been successfully applied in several applications, such as (deep) generative modeling (Wu et al., 2019; Kolouri et al., 2018a; Nguyen & Ho, 2022a), domain adaptation (Lee et al., 2019), and clustering (Kolouri et al., 2018b).
Nevertheless, it also suffers from certain limitations in, e.g., deep learning applications where the mini-batch approaches (Fatras et al., 2020) are utilized.
Here, the number of supports n is often much smaller than the number of dimensions.
Therefore, the computational complexity of solving L one-dimensional Wasserstein distance, Θ( Ln log 2 n ) is small compared to the computational complexity of matrix multiplication Θ( Ldn ) .
This indicates that almost all computation is for the projection step.
The situation is ubiquitous since there are several deep learning applications involving processing high-dimensional data, including images (Genevay et al., 2018; Nguyen & Ho, 2022b), videos (Wu et al., 2019), and text (Schmitz et al., 2018).
Motivated by the low-rank decomposition of matrices, we propose a more efficient approach to project original measures to their one-dimensional projected measures.
In particular, two original measures are first projected into k one-dimensional projected measures via Radon transform where k < L .
For convenience, we call these projected measures as bottleneck projections .
Then, new L one-dimensional projected measures are created as random linear combinations of the bottleneck projections.
The linear mixing step can be seen as applying Radon transform on the joint distribution of k one-dimensional projected measures.
From the computational point of view, the projecting step consists of two consecutive matrix multiplications.
The first multiplication is between the data matrix of size n × d
and the bottleneck projecting directions matrix of size d × k , and the second multiplication is between the bottleneck projecting matrix and the linear mixing matrix of size k × L .
Columns of both the bottleneck projecting directions matrix and the linear mixing matrix are sampled randomly from the uniform distribution over the corresponding unit-hypersphere.
For the same value of L , we show that the bottleneck projection approach is faster than the conventional approach when the values of L and d are large.
Contribution:
In summary, our main contributions are two-fold:
We formulate the usage of bottleneck projections into a novel integral transform, Hierarchical Radon Transform (HRT) which is the composition of Partial Radon Transform (Liang & Munson, 1997) (PRT) and Overparameterized Radon Transform (ORT).
By using HRT, we derive a new sliced Wasserstein variant, named Hierarchical sliced Wasserstein distance (HSW).
By showing that HRT is injective, we prove that HSW is a valid metric in the space of probability measures.
Furthermore, we discuss the computational complexity and the projection complexity of the proposed HSW distance.
Finally, we derive the connection between HSW and other sliced Wasserstein variants, and the sample complexity of HSW.
We design experiments focusing on comparing HSW to the conventional SW in generative modeling on standard image datasets, including CIFAR10, CelebA, and Tiny ImageNet.
We show that for approximately the same amount of computation, HSW provides better generative performance than SW and helps generative models converge faster.
In addition, we compare generated images qualitatively to reinforce the favorable quality of HSW.
Organization.
The remainder of the paper is organized as follows.
We first provide background about
Radon Transform, Partial Radon Transform, and sliced Wasserstein distance in Section 2.
In Section 3, we propose Overparameterized Radon Transform, Hierarchical Radon Transform,
Hierarchical sliced Wasserstein distance, and analyze their relevant theoretical properties.
Sectioncontains the application of HSW to generative models, qualitative results, and quantitative results on standard benchmarks.
We then draw concluding remarks in Section 5.
Finally, we defer the proofs of key results, supplementary materials, and discussion on related works to Appendices.
Notation.
For n ∈ N , we denote by [ n ] the set { 1 , 2 , . . . , n } .
For any d ≥ 2 , S d − 1 := { θ ∈ R d | ∥ θ ∥ 2 = 1 } denotes the d dimensional unit sphere, and U ( S d − 1 ) is the uniform measure over S d − 1 .
We denote the set of absolutely integrable functions on R d as L 1 (cid:0) R d (cid:1) :
= (cid:8) f : R d → R | (cid:82) R d | f ( x ) | dx < ∞ (cid:9) .
For p ≥ 1 , we denote the set of all probability measures on
R d that have finite p -moments as P p ( R d ) .
For m ≥ 1 , we denote µ ⊗ m as the product measure of m random variables that follow µ while A ⊗
m indicates the Cartesian product of m sets A .
The
Dirac delta function is denoted by δ .
For a vector X ∈ R dm , X :=
( x 1 , . . .
, x m ) , P X denotes the empirical measures 1 m (cid:80) mi =1 δ ( x − x i ) .
For any two sequences a n and b n , the notation a n = O ( b n ) means that a n ≤ Cb n for all n ≥ 1 , where C is some universal constant.
B ACKGROUND
We first review the definition of the Radon Transform.
We then review the sliced Wasserstein distance, and discuss its limitation in high-dimensional setting with relatively small number of supports.
Definition 1 (Radon Transform (Helgason, 2011)) .
The Radon Transform R : L 1 ( R d ) →
L 1 (cid:0) R × S d − 1 (cid:1) is defined as: ( R f )( t, θ ) =
(cid:82) R d f ( x ) δ ( t −⟨ x, θ ⟩ ) dx.
Note that, the Radon Transform defines a linear bijection.
Sliced Wasserstein distance: From the definition of the Radon Transform, we can define the sliced
Wasserstein distance as follows.
Definition 2 (Sliced Wasserstein Distance (Bonneel et al., 2015)) .
For any p ≥ 1 and dimension d ≥ 1 , the sliced Wasserstein-
p distance between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is given by:
[Equation]
where f µ ( · ) , f ν ( · ) are the probability density functions of µ , ν respectively, and W p ( µ, ν ) :
=
[Equation]
1 p is the Wasserstein distance of order p (Villani, 2008; Peyr´e & Cuturi, 2019).
With slightly abuse of notation, we use W p ( µ, ν ) and W p ( f µ , f ν ) interchangeably.
[Figure]
The main benefit of sliced Wasserstein is that the one-dimensional Wasserstein distance
W
[Equation]
has a
[Equation]
where F − 1 ( R f µ )(
[Equation]
,θ ) ( z ) is the inverse cumulative distribution function of the random variable that has the density ( R f µ )( · , θ ) (similarly for F − 1 ( R f ν )( · ,θ ) ( z ) ).
We denote the one-dimensional pushforward measures with the density ( R f µ )( · , θ ) as θ♯µ .
Sliced Wasserstein distance between discrete measures: In several applications, such as deep generative modeling (Deshpande et al., 2018; Nguyen & Ho, 2022b), domain adaptation (Lee et al., 2019), point cloud reconstruction (Nguyen et al., 2021c; 2023a), sliced Wasserstein distance had been used for discrete measures.
Let two measures µ and ν that have the pdfs f µ ( x ) = 1 n (cid:80) ni
=
α i δ ( x − x i ) , f ν ( y ) =
1 m (cid:80) mj
=1 β j δ ( y − y j )
( α i , β j > 0 ∀
i, j ), and a given projecting direction θ ∈ S d − 1 , the corresponding Radon Transform are ( R f µ )( z, θ ) = 1 n (cid:80) ni
=
α i δ ( z − θ ⊤ x i ) and ( R f ν )( z, θ ) =m (cid:80) mj
=1 β j δ ( z − θ ⊤ y j ) .
Since the expectation in Definition 2 is intractable, Monte Carlo estimation is used with L projecting directions θ 1 , . . .
, θ L ∼ U ( S d −
[Equation]
We denote X as the matrix of size n × d which has rows being supports of µ , [ x 1 , . . .
, x n ] ⊤ , Y as the matrix of size m × d which has rows being supports of ν , [ y 1 , . . .
, y m ] ⊤ , and Θ as the matrix of size d × L which has columns being sampled projecting directions [ θ 1 , . .
, θ L ] .
The supports of the Radon Transform measures from µ and ν are the results of matrix multiplication
¯ X = X · Θ (with the shape n × L ) and
¯ Y = Y · Θ (with the shape m × L ).
Columns of ¯ X and ¯ Y are supports of projected measures.
Therefore, L one-dimensional Wasserstein distances are computed by evaluating the quantile functions which are based on sorting columns.
A visualization is given in Figure 1.
Computational and projection complexities of sliced Wasserstein distance: Without the loss of generality, we assume that n ≥ m , the time complexity of sliced Wasserstein is O ( Ldn +
Ln log 2 n ) where Ldn is because of matrix multiplication ( X · Θ ) and Ln log 2 n is because of the sorting algorithm.
The projection complexity of SW is the memory complexity for storing the projecting directions Θ which is O ( Ld ) .
We would like to remark that, the value of number of projections L should be comparable to the number of dimension d for a good performance in applications (Deshpande et al., 2018; Nguyen et al., 2021a;b).
Computational issuse of sliced Wasserstein when d ≫ n : In deep learning applications (Deshpande et al., 2018; Lee et al., 2019; Kolouri et al., 2018a) where mini-batch approaches are used, the number of dimension d is normally much larger than the number of supports n , e.g., d = 8192 and n = 128 .
Therefore, log 2 n ≪ d that leads to the fact that the main computation of sliced Wasserstein is for doing projecting measures O ( Ldn ) .
To the best of our knowledge, prior works have not adequately addressed this limitation of sliced Wasserstein.
H IERARCHICAL S LICED W ASSERSTEIN DISTANCE
In this section, we propose an efficient way to improve the projecting step of sliced Wasserstein distance.
In particular, we first project measures into a relatively small number of projections ( k ), named bottleneck projections ( k < L ) .
After that, L projections are created by a random linear combination of the bottleneck projections.
To explain the usage of the bottleneck projections, we first introduce Hierarchical Radon Transform (HRT) in Section 3.1.
We then define the Hierarchical
Sliced Wasserstein (HSW) distance and investigate its theoretical properties in Section 3.2.
We show that the usage of bottleneck projections appears in an efficient estimation of HSW.
3. H IERARCHICAL R ADON T RANSFORM
To define the Hierarchical Randon Transform, we first need to review an extension of Radon Transform which is Partial Radon Transform (PRT).
After that, we propose a novel extension of Radon Transform which is named Overparameterized Radon Transform (ORT).
Definition 3 (Partial Radon Transform (Liang & Munson, 1997)) .
The Partial Radon Transform (PRT)
PR : L 1 ( R d 1 × R d 2 ) →
L 1 (cid:0) R × S d 1 − 1 × R d 2 (cid:1) is defined as: ( PR f )( t, θ, y )
= (cid:82) R d 1 f (
x, y ) δ ( t − ⟨ x, θ ⟩ ) dx.
Given a fixed y , the Partial Radon Transform is the Radon Transform of f ( · , y ) .
Definition 4 (Overparameterized Radon Transform) .
The Overparameterized Radon Transform (ORT) OR : L 1 ( R d ) → L 1 (cid:0) R ⊗ k × S ( d − 1) ⊗ k (cid:1) is defined as:
[Equation]
where t 1: k :=
( t 1 , . . .
, t k ) ∈ R ⊗ k and θ 1: k :=
( θ 1 , . .
, θ k ) ∈ ( S d − 1 ) ⊗ k .
Definition 4 is called “overparameterized” since the dimension of the transformed function’s arguments is higher than the original dimension.
Our motivation for ORT comes from the success of overparametrization in deep neural networks (Allen-Zhu et al., 2019).
Proposition 1.
The Overparameterized Radon Transform (ORT) is injective, i.e., for any functions f,
g ∈ L 1 ( R d ) , ( OR f )( t 1: k , θ 1: k ) =
( OR g )( t 1: k , θ 1: k ) ∀ t 1: k , θ 1: k implies that f = g .
Since ORT is an extension of RT, the injectivity of ORT is derived from the injectivity of RT.
The proof of Proposition 1 is in Appendix C.1.
We now define the Hierarchical Radon Transform (HRT).
Definition 5 (Hierarchical Radon Transform) .
Hierarchical Radon Transform (HRT) HR : L 1 ( R d ) → L
[Equation]
R × S
[Equation]
d
[Equation]
k × S k
[Equation]
1 (cid:1) is defined as:
[Equation]
where v ∈ R , ψ = ( ψ 1 , . . . ,
ψ k ) ∈ S k − 1 , and θ 1: k
= ( θ 1 , . .
, θ k ) ∈ ( S d − 1 ) ⊗ k .
Definition 5 is called “hierarchical” since it is the composition of Partial Radon Transform and Overparameterized Radon Transform.
We can verify that ( HR f )( v, θ 1: k , ψ ) =
( PR ( OR f ))
( v, θ 1: k , ψ ) (we refer the reader to Appendix B for the derivation).
To the best of our knowledge, ORT and HRT have not been proposed in the literature.
Proposition 2.
The Hierarchical Radon Transform (HRT) is injective, i.e., for any functions f,
g ∈ L( R d ) , ( HR f )( v, θ 1: k , ψ ) =
( HR g )( v, θ 1: k , ψ ) ∀ v, θ 1: k , ψ implies that f = g .
Since HRT is the composition of PRT and ORT, the injectivity of HRT is derived from the injectivity of ORT and PRT.
The proof of Proposition 2 is in Appendix C.2.
Hierarchical Radon Transform of discrete measures: Let f ( x ) = 1 n (cid:80) nj
=1 α i δ ( x − x j ) , we have ( HR f )( v, θ 1: k , ψ ) = 1 n (cid:80) ni
=
α i δ (cid:16) v − (cid:80) kj =1 ⟨ x i , θ j ⟩ ψ j (cid:17) =
1 n (cid:80) ni
=
α i δ (cid:0) v − ψ ⊤ Θ ⊤ x
i (cid:1) , where Θ is the matrix columns of which are θ 1: k .
Here, θ 1: k are bottleneck projection directions, ψ is the mixing direction, ψ ⊤ Θ is the final projecting direction, θ ⊤ j x i , . . .
, θ ⊤ j x n for any j ∈ [ k ] are the bottleneck projections, and ψ ⊤ Θ ⊤ x i , . . .
, ψ ⊤ Θ ⊤ x n is the final projection.
We consider that the bottleneck projections are the spanning set of a subspace that has the rank at most k and the final projection belongs to that subspace.
Based on the linearity of Gaussian distributions, we provide the result of HRT on multivariate Gaussian and mixture of multivariate Gaussians in Appendix B.
Applications of Hierarchical Radon Transform: In this paper, we focus on showing the benefit of the HRT in the sliced Wasserstein settings.
However, similar to the Radon Transform, the HRT can also be applied to multiple other applications such as sliced Gromov Wasserstein (Vayer et al., 2019), sliced mutual information Goldfeld & Greenewald (2021), sliced score matching (sliced Fisher divergence)
Song et al.
, sliced Cramer distance (Kolouri et al., 2019a), and other tasks that need to project probability measures.
3. H IERARCHICAL S LICED W ASSERSTEIN DISTANCE
By using Hierarchical Radon Transform, we define a novel variant of sliced Wasserstein distance which is named Hierarchical Sliced Wasserstein (HSW) distance.
Definition 6.
For any p ≥ 1 , k ≥ 1 , and dimension d ≥ 1 , the hierarchical sliced Wasserstein distance of order p between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is given by:
[Equation] [Equation]
where θ 1 , . .
, θ k ∼ U ( S d − 1 ) and ψ ∼ U ( S k − 1 ) .
Properties of hierarchical sliced Wasserstein distance: First, we have the following result for the metricity of HSW.
Theorem 1.
For any p ≥ 1 and k ≥ 1 , the hierarchical sliced Wasserstein HSW p,k ( · , · ) is a metric on the space of probability measures on R d .
Proof of Theorem 1 is given in Appendix C.3.
Our next result establishes the connection between the HSW, max hierarchical sliced Wasserstein (Max-HSW)
(see Definition 8 in Appendix B), max sliced Wasserstein (Max-SW)
(see Definition 7 in Appendix B), and Wasserstein distance.
We refer the reader to Appendix B for more theoretical properties of the Max-HSW.
The role of Max-HSW is to connect HSW with Max-SW that further allows us to derive the sample complexity of HSW.
Proposition 3.
For any p ≥ 1 and k ≥ 1 , we find that
[Equation]
where we define
[Equation]
as max hierarchical sliced p -Wasserstein, and max sliced p -Wasserstein, respectively.
Proof of Proposition 3 is given in Appendix C.4.
Given the bounds in Proposition 3, we demonstrate that the hierarchical sliced Wasserstein does not suffer from the curse of dimensionality for the inference purpose, namely, the sample complexity for the empirical distribution from i.i.d.
samples to approximate their underlying distribution is at the order of O ( n − 1 / 2 ) where n is the sample size.
Proposition 4.
Assume that P is a probability measure supported on compact set of R d .
Let
X 1 , X 2 , . . .
, X n be i.i.d.
samples from P and we denote P n = 1 n (cid:80) ni
=
δ X i as the empirical measure of these data.
Then, for any p ≥ 1 , there exists a universal constant C > 0 such that
[Equation]
where the outer expectation is taken with respect to the data X 1 , X 2 , . .
X n .
Proof of Proposition 4 is given in Appendix C.5.
Monte Carlo estimation:
Similar to SW, the expectation in Definition 6 is intractable.
Therefore, Monte Carlo samples ψ 1 , . . .
, ψ L ∼
U ( S k − 1 ) (by abuse of notations) and θ 1 , 1 , . . .
, θ 1 ,H , . . .
, θ k, 1 , . .
, θ k,H ∼ U
( S d − 1 ) are used to approximate the HSW, which leads to the following approximation:
[Equation]
[Figure]
Computational and projection complexities on discrete measures: It is clear that the time complexity of the Monte Carlo estimation of HSW with discrete probability measures of n supports is O ( Hkdn + HLkn + HLn log 2 n ) .
The projection complexity of HSW is O ( Hdk + kL ) .
For a fast computational complexity and a low projection complexity, we simply choose H = 1 .
In this case, the computational complexity is O ( kdn + Lkn + Ln log 2 n ) and the projection complexity is
O ( dk + kL ) .
Recall that k < L , which implies the estimation of HSW is faster and more efficient in memory than the estimation of SW.
This fast estimator is unbiased, though its variance might be high.
Benefit of HSW when d ≫ n : For the same computation complexity, HSW can have a higher number of final projections L than SW.
For example, when d = 8192 and n = 128 (the setting that we will use in experiments), SW with L = 100 has the computational complexity proportion to 104 .
94 × 10 6 and the projection complexity proportion to 0 .
82 × 10 6 .
In the same setting of d and n ,
HSW with H = 1 , k = 50 and L = 500 has the computational complexity proportion to 89 .
45 × 10and the projection complexity proportion to 0 .
× 10 6 .
Implementation of HSW on discrete measures: The slicing process of HSW contains H matrices of size d × k in the first level and a matrix of size k × L in the second level.
The projected measures are obtained by carrying out matrix multiplication between the support matrices with the projection matrices in the two layers in turn.
When H = 1 , we observe the composition of projection matrices in the two layers as a two-layer neural network with linear activation.
However, the weights of the neural network have a constraint (spherical constraint) and are sampled instead of being optimized.
In the paper, since we focus on the efficiency of the estimation, we consider only settings where H = 1 .
A visualization of the process when H = 1 is given in Figure 2.
Beyond single hierarchy with standard Radon transform: In HRT, PRT is applied on the arguments ( t 1 , . . .
, t k ) .
By applying many PRTs on subsets of arguments e.g., ( t 1 , t 2 ) , . . .
, ( t k − 1 , t k ) and then applying PRT again on the arguments of the output function, we derive a more hierarchical version of HRT.
Similarly, we could also apply ORT on the arguments multiple times to make the transform hierarchical.
Despite the hierarchy, Hierarchical Sliced Wasserstein distance still uses linear projections.
By changing from Radon Transform variants to Generalized Radon Transform (Beylkin, 1984) variants, we easily extend Hierarchical Radon Transform to Hierarchical Generalized Radon
Transform (HGRT).
Furthermore, we derive the Hierarchical Generalized Sliced Wasserstein (HGSW) distance.
Since HGSW has more than one non-linear transform layer, it has the provision of using a more complex non-linear transform than the conventional Generalized sliced Wasserstein (GSW)
(Kolouri et al., 2019b).
Compared to the neural network defining function of GRT (Kolouri et al., 2019b), HGSW preserves the metricity, i.e., HGSW satisfies the identity property due to the injectivity of HGRT with defining functions that satisfy constraints H1-H4 in (Kolouri et al., 2019b).
Since the current set defining functions e.g., circular functions and homogeneous polynomials with an odd degree are not scalable in high-dimension, we defer the investigation of HGSW and finding a more scalable function for future work.
Distributions of final projecting directions in HRT: We recall that the final projecting directions of HSW are ψ ⊤ 1 Θ , . . .
, ψ ⊤ L Θ , where Θ = ( θ 1: k ) , θ 1: k ∼ U ( S d − 1 ) , and ψ 1 . . .
, ψ L ∼ U ( S k − 1 ) .
It is clear that the final projecting directions are distributed uniformly from the manifold S :
= { x ∈
R d − 1 | x = (cid:80) ki
=1 ψ ⊤ i θ i , ( ψ 1 , . . .
ψ k ) ∈ S k − 1 , θ i ∈ S
d − 1 , ∀ i = 1 , . . .
, k } .
Therefore, HSW can be considered as the sliced Wasserstein with the projecting directions on a special manifold.
This is different from the conventional unit hyper-sphere.
To our knowledge, the manifold S has not been explored sufficiently in previous works, which may be a potential direction of future research.
[Table]
[Figure]
On the choices of k and L in HSW: We consider the HSW as an alternative option for applications that use the SW.
For a given value of L in the SW, we can choose a faster setting of HSW by selecting k ≤ LdL + d based on the analysis of their computational complexities.
Similarly, the HSW with Lfinal projections can be faster than the SW with L 1 projections by choosing k ≤ L 1 d − ( L 2 − L 1 )log 2 n d + L.
Later, we use this rule for the choices of k in the experiments while comparing HSW with SW.
E XPERIMENTS
In this section, we compare HSW with SW on benchmark datasets: CIFAR10 (with image size 32x32)
(Krizhevsky et al., 2009), CelebA (with image size 64x64), and Tiny ImageNet (with image size 64x64)
(Le & Yang, 2015).
To this end, we consider deep generative modeling with the standard framework of the sliced Wasserstein generator (Deshpande et al., 2018;
Nguyen et al., 2021a;
Deshpande et al., 2019; Nguyen & Ho, 2022b; Nadjahi et al., 2021).
We provide a detailed discussion of this framework including training losses and their interpretation in Appendix D.1.
The SW variants are used in the feature space with dimension d = 8192 and the mini-batch size 128 for all datasets.
The main evaluation metrics are FID score (Heusel et al., 2017) and Inception score (IS) (Salimans et al., 2016).
We do not report the IS score on CelebA since it poorly captures the perceptual quality of face images (Heusel et al., 2017).
The detailed settings about architectures, hyper-parameters, and evaluation of FID and IS are provided in Appendix F.
Our experiments aim to answer the following questions: (1) For approximately the same computation, is the HSW comparable in terms of perceptual quality while achieving better convergence speed?
For the same number of final projections L and a relatively low number of bottleneck projections k < L , how is the performance of the HSW compared to the SW?
(3) For a fixed value of bottleneck projections k , does increasing the number of final projections L improve the performance of the
HSW?
(4) For approximately the same computation, does reducing the value of k and raising the value of L lead to a better result?
After running each experiment 5 times, we report the mean and the standard deviation of evaluation metrics.
The HSW is usually better than the SW with a lower computation: We report the FID scores,
IS scores, the computational complexity ( × 10 6 ), and the projection complexity ( × 10 6 ) for the SW with L ∈ { 100 , 1000 , 10000 } and the HSW with ( k, L ) =
{ (70 , 2000) , (400 , 6000) , (3000 , 18000) } respectively in Table 1.
The computational complexities and the projection complexities are computed based on the big O notation analysis of SW and HSW in previous sections.
According to the table, the
HSW usually yields comparable FID and IS scores while having a lower computational complexity
[Table]
and a lower projection complexity on benchmark datasets.
We show random generated images on CIFAR10 in Figure 5 in Appendix D.2, on CelebA in Figure 6 in Appendix D.
[Table]
and on Tiny
ImageNet in Figure 7 in Appendix D.2 as a qualitative comparison.
Those generated images are consistent with the quantitative scores in Table
[Table]
The HSW leads to faster convergence than the SW with a lower computation: We plot the FID scores over training epochs of the SW and the HSW with the same setting as in Table 1 in Figure 3.
We observe that FID scores from models trained by the HSW (with a better computation) reduce faster than ones trained from the SW.
The same phenomenon happens with the IS scores in Figurein Appendix D.2 where IS scores from the HSW increase earlier.
The reason is that the HSW has a higher number of final projections than the SW, hence, it is a more discriminative signal than the SW.
Ablation studies on k and L in the HSW:
In the Table 2, we report the additional FID scores,
IS scores, the computational complexity ( × 10 6 ), and the projection complexity ( × 10 6 ) for the
HSW with ( k, L ) = { (500 , 1000) , (500 , 4000) , (100 , 50000) } .
First, we see that given k = 5000 , increasing L from 1000 to 4000 improves the generative performance on all three datasets.
Moreover, we observe that the HSW with the same number of final projections as the SW ( L = 1000 ) gives comparable scores while the complexities are only about half of the complexities of the SW.
When decreasing k from 500 to 400 and increasing L from 4000 to 6000 , the generative performance of the HSW is enhanced and the complexities are also lighter.
However, choosing a too-small k and a too-large L does not lead to a better result.
For example, the HSW with k = 100 , L =has high complexities compared to k = 500 , L = 4000 and k = 400 , L = 6000 , however, its FID scores and IS scores are worse.
The reason is because of the linearity of the HRT.
In particular, the k bottleneck projections form a subspace that has a rank at most k and the L final projections still lie in that subspace.
This fact suggests that the value of k should be also chosen to be sufficiently large compared to the true rank of the supports.
In practice, data often lie on a low dimensional manifold, hence, k can be chosen to be much smaller than d and L can be chosen to be large for a good estimation of discrepancy.
This is an interpretable advantage of the HSW compared to the SW since it can separate between the assumption of the ranking of supports and the estimation of the discrepancy between measures.
Max hierarchical sliced Wasserstein: We also compare Max-HSW
(see Definition 7 in Appendix B) with the conventional Max-SW in generative modeling in Table 3 in Appendix D.2.
We observe that the overparameterization from HRT could also improve the optimization of finding good projections.
We would like to recall that the Max-HSW is the generalization of the Max-SW, namely, Max-HSW with k = 1 is equivalent to the Max-SW.
Therefore, the performance of the Max-HSW is at least the same as the Max-SW.
C ONCLUSION
In this paper, we proposed a hierarchical approach to efficiently estimate the Wasserstein distance with provable benefits in terms of computation and memory.
It formed final projections by combining linearly and randomly from a smaller set of bottleneck projections.
We justified the main idea by introducing Hierarchical Radon Transform (HRT) and hierarchical sliced Wasserstein distance (HSW).
We proved the injectivity of the HRT, the metricity of the HSW, and investigated its theoretical properties including computational complexity, sample complexity, projection complexity, and its connection to other sliced Wasserstein variants.
Finally, we conducted experiments on deep generative modeling where the main computational burden due to the projection was highlighted.
In this setting,
HSW performed favorably in both generative quality and computational efficiency.
A CKNOWLEDGEMENTS
NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of
Machine Learning.
Supplement to “Hierarchical Sliced Wasserstein Distance”
In this supplement, we first discuss some related works in Appendix A. We then present some additional materials including hierarchical Radon transform on multivariate Gaussian and mixture of multivariate Gaussians, max sliced Wasserstein distance, max hierarchical sliced Wasserstein (Max-HSW) distance, their computation, and their theoretical properties in Appendix B. After that, we collect proofs for key results in the paper in Appendix C. In Appendix D, we include the detailed training objectives of our generative modeling framework and some additional experiments including convergence in IS scores of generative models, generated images, and the comparison between the
Max-HSW and the Max-SW.
Moreover, we report experimental settings including parameter choices and neural network architectures in Appendix F.
A M ORE R ELATED W ORKS
On the theoretical side, statistical properties of sliced Wasserstein distance in learning generative models are investigated in (Nadjahi et al., 2019).
Distributional convergence of the Sliced Wasserstein process is derived in Xi & Niles-Weed (2022).
Statistical, robustness, and computational guarantees for sliced Wasserstein distances are shown in (Nietert et al., 2022).
A differential private version of the SW is proposed in (Rakotomamonjy & Liva, 2021).
On the methodological side, authors in (Nguyen et al., 2021a;b) replace uniform distribution on the projecting directions on the unit hyper-sphere in SW with an estimated distribution that puts high probabilities for discriminative directions.
Spherical sliced Wasserstein which is a sliced Wasserstein variant on the hyper-sphere is introduced in (Bonet et al., 2022).
Dependent projecting directions are utilized in (Nguyen et al., 2023b).
Sliced partial optimal transport is proposed in (Bonneel &
Coeurjolly, 2019; Bai et al., 2022)
A fast biased approximation of the SW is proposed in (Nadjahi et al., 2021).
The Augmented SW is introduced (Chen et al., 2022).
A sliced Wasserstein variant between measures over tensors is defined in (Nguyen & Ho, 2022b).
Estimating Wasserstein distance with one-dimensional transportation plans from orthogonal projecting directions is used in (Rowland et al., 2019).
The SW is used in generative frameworks such as sliced iterative normalizing flows (Dai & Seljak, 2021) and fine-tuning pre-trained model (Lezama et al., 2021).
The SW gradient flows are investigated in (Liutkus et al., 2019; Bonet et al., 2021).
The SW is used for set representations in (Naderializadeh et al., 2021).
Variational inference using the SW is carried out in (Yi & Liu, 2021).
Similarly, the SW is used for approximate Bayesian computation in (Nadjahi et al., 2020).
B A DDITIONAL M ATERIALS
Hierarchical Radon Transform as the composition of Partial Radon Transform and Overparametrized Radon Transform: From the definitions of ORT (Definition 4) and PRT (Definition 3), we have:
[Equation]
v −
[Equation]
(cid:88) i =
t
[Equation]
ψ
[Equation]
(cid:33) dtdx
[Equation]
v −
[Equation]
where the second and the third equality are due to the Fubini’s theorem.
Monte Carlo approximation error of HSW: We have the following proposition for the Monte
Carlo approximation error:
Proposition 5.
Given p > 1 , k > 1 , and µ, ν ∈ P p ( R d ) , we have the following error of Monte Carlo estimation for HSW:
[Equation]
The proof is given in Appendix C.6.
From the proposition, we see that the error reduces proportionally to the square root of H and the square root of L .
Therefore, increasing H and L could lead to better performance.
Hierarchical Radon Transform of Multivariate Gaussian and Mixture of Multivariate Gaussians: Similar to the result of Radon Transform for Multivariate Gaussian and Mixture of Multivariate Gaussians (Kolouri et al., 2018b), given projecting direction Θ = ( θ 1: k ) , ψ , the one dimensional projected distributions from a Gaussian and a mixture of Gaussians with the HRT are also a Gaussian and a mixture of Gaussians respectively.
These results are directly obtained from the linearity of Gaussian distributions.
In more detail, let f := N ( µ, Σ) , we have ( HR f )( · , θ 1: k , ψ ) =
N ( ψ ⊤ Θ ⊤ µ, ψ ⊤ Θ ⊤ ΣΘ ψ ) .
Similarly, let f := (cid:80) ki
=1 w i N ( µ i , Σ i ) , we have ( HR f )( · , θ 1: k , ψ ) =
(cid:80) k
i =
1 w i N ( ψ ⊤ Θ ⊤ µ i , ψ ⊤ Θ ⊤ Σ
i Θ ψ ) .
Max sliced Wasserstein distance: We first review the definition of max sliced Wasserstein distance (Max-SW)
(Deshpande et al., 2019).
Definition 7.
For any p ≥ 1 and dimension d ≥ 1 , the max sliced Wasserstein-
p distance between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is given by:
[Equation]
Max-SW is also a valid metric between probability measures (Deshpande et al., 2019).
The benefit of Max-SW is that it uses only one projecting direction instead of L projecting directions as SW.
Therefore, its projection complexity is only O ( d ) compared to O ( Ld ) of SW.
However, Max-SW requires an iterative optimization procedure to find the max projecting direction on the unit hypersphere e.g., projected gradient asscent (Kolouri et al., 2019b).
We summarize the projected gradient ascent of the Max-SW in Algorithm 1.
Algorithm
Input: Probability measures: µ, ν , learning rate η , max number of iterations T .
Initialize θ while θ not converge or reach T do
[Equation]
θ =
θ ||
θ
|| 2 end while
Return: θ , W p (( R f µ )( · , θ ) , ( R f ν )( · , θ ))
Max hierarchical sliced Wasserstein distance: We now define the max hierarchical sliced Wasserstein distance.
Definition 8.
For any p ≥ 1 , k ≥ 1 and dimension d ≥ 1 , the max hierarchical sliced Wasserstein-
p distance between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is given by:
[Equation]
Similar to the idea of max sliced Wasserstein, Max-HSW improves the projection complexity of
HSW by avoiding using Monte Carlo samples for projecting directions and mixing directions.
In particular, Max-HSW finds the best set of bottleneck projections θ 1: k and the best mixing direction ψ in terms of maximizing the discrepancy between two interested measures.
Remark 1.
We have that when k
= 1 , the max hierarchical sliced Wasserstein distance reverts into the max sliced Wasserstein distance.
[Equation]
This is due to the fact that ψ ∈ S
= { 1 } .
In Max-HSW, we do not enforce any linear independence constraint over k bottleneck projection directions.
However, that kind of constraint can be done by forcing the k vectors to be orthogonal.
Formally, we could project k vectors into the Stiefel manifold by performing the Gram-Schmidt process.
We can only do maximization over k bottleneck vectors and take the expectation over the linear mixing vectors.
We denote this variant as Semi Max hierarchical sliced Wasserstein distance.
Definition 9.
For any p ≥ 1 , k ≥ 1 and dimension d ≥ 1 , the max hierarchical sliced Wasserstein-
p distance between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is given by:
[Equation]
The optimization can be done by stochastic gradient ascent that utilizes Monte Carlo samples ψ 1 , . . .
, ψ L from U ( S k − 1 ) .
Computation of max hierarchical sliced Wasserstein distance: Similar to the Max-SW, Max-HSW can be computed via the projected gradient ascent algorithm.
In practice, the gradient of θ 1: k and ψ can be computed by using backpropagation (chain rule) since the hierarchical Radon Transform can be seen as a two-layer neural network.
We summarize the projected gradient ascent of the Max-HSW in Algorithm 2.
Input: Probability measures: µ, ν , learning rate η , max number of iterations T .
Initialize θ while θ 1: k , ψ not converge or reach T do ψ
[Equation]
) ψ = ψ
||
[Equation]
|| 2 for i = 1 to k do
[Equation]
for end while
Properties of max hierarchical sliced Wasserstein distance: We first have the following result for the metricity of Max-HSW.
Theorem 2.
For any p ≥ 1 and k ≥ 1 , the hierarchical sliced Wasserstein Max-HSW p,k ( · , · ) is a metric on the space of probability measures on R d .
Proof of Theorem 2 is given in Appendix C.7.
We establish the connection between the HSW, max hierarchical sliced Wasserstein (Max-HSW), max sliced Wasserstein (Max-SW), and Wasserstein distance in Proposition 3.
The proof is given in Appendix C.4.
We demonstrate that the max hierarchical sliced Wasserstein does not suffer from the curse of dimensionality for the inference purpose, namely, the sample complexity for the empirical distribution from i.i.d.
samples to approximate their underlying distribution is at the order of O ( n − 1 / 2 ) .
Proposition 6.
Assume that P is a probability measure that has supports on compact set of R d .
Let
X 1 , X 2 , . . .
, X n be i.i.d.
samples from P and we denote P n = 1 n (cid:80) ni
=
δ X i as the empirical measure on data samples.
Then, for any p ≥ 1 and k ≥ 1 , there exists a universal constant C > 0 such that
[Equation]
where the outer expectation is taken with respect to the data X 1 , X 2 , . .
X n .
Proof of Proposition 6 is similar to the proof of Proposition 4.
We refer the reader to Appendix C.5.
Hierarchical Generalized Sliced Wasserstein distance: We first define the Overparameterized
Generalized Radon Transform, Partial Generalized Radon Transform, and Hierarchical Generalized
Radon Transform.
We recall two injective defining functions in (Kolouri et al., 2019b) which are the circular function g ( x, θ ) =
∥ x − r ∗ θ ∥ 2 for x ∈ R d , θ ∈ S d − 1 , (hyper-parameter r ∈ R + ) and the homogeneous polynomials with an odd degree g ( x, θ ) =
(cid:80) | α | = m θ α x α where α = ( α 1 , . . .
, α d α ) ∈
N d α , | α | = (cid:80) d α i =
α i , x α = (cid:81)
d α i =
x α
i i .
The summation in the polynomial defining function iterates over all possible multi-indices α such that | α | = m with m denotes the degree of the polynomial and θ ∈ S d α − 1 .
Definition 10 (Partial Generalized Radon Transform) .
Given the defining function g ( x, θ ) : R d 1 × Ω θ → R , the Partial Generalized Radon Transform (PGRT) PR : L 1 ( R d 1 × R d 2 ) →
L 1 (cid:0) R × Ω θ × R d 2 (cid:1) is defined as:
[Equation]
Definition 11 (Overparameterized Generalized Radon Transform) .
Given the defining function g ( x, θ ) : R d × Ω θ → R , the Overparameterized Radon Transform (OGRT) OR : L 1 ( R d ) →
L 1 (cid:0) R ⊗ k × S ( d − 1) ⊗ k (cid:1) is defined as:
[Equation]
where t 1: k :=
( t 1 , . . .
, t k ) ∈ R ⊗ k and θ 1: k :=
( θ 1 , . .
, θ k ) ∈ ( S d − 1 ) ⊗ k .
Definition 12 (Hierarchical Generalized Radon Transform) .
Given the defining functions g 1 ( x, θ ) : R d × Ω θ → R and g 2 ( x, ψ ) : R k × Ω ψ → R , Hierarchical Generalized Radon Transform (HGRT)
HGR
[Equation]
L
[Equation]
R
[Equation]
→ L 1 (cid:0) R × Ω ⊗ k θ × Ω ψ (cid:1) is defined as:
[Equation]
where v ∈ R , ψ ∈ Ω ψ and θ 1: k
= ( θ 1 , . .
, θ k ) ∈ Ω ⊗ k θ .
The injectivity of PGRT, OGRT, and HGRT follow the injectivity of GRT, i.e., the proof technique is similar to the case of ORT and HRT.
We refer the reader to the previous work (Kolouri et al., 2019b) for more information.
The benefit of HGRT is that it can create a much stronger non-linearity by stacking multiple injective ones.
Definition 13.
For any p ≥ 1 , k ≥ 1 , and dimension d ≥ 1 , the hierarchical generalized sliced
Wasserstein distance of order p between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is given by:
[Equation]
where θ 1 , . .
, θ k ∼ U
(Ω θ ) and ψ ∼ U (Ω ψ ) .
The proof technique for the metricity of HGSW is the same as HSW which is based on the metricity of the Wasserstein distance, the injectivity of HGRT, and the Minkowski inequality.
Distributional Hierarchical Sliced Wasserstein distance: In HSW, projecting directions are drawn from the uniform distribution which might not be the best choice in some settings.
We could replace it with some other distributions, e.g., Von Mises Fisher (Nguyen et al., 2021b) or implicit distribution (Nguyen et al., 2021a).
Definition 14.
For any p ≥ 1 , k ≥ 1 , and dimension d ≥ 1 , the distributional hierarchical sliced
Wasserstein distance of order p between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is given by:
[Equation]
where Γ( S d − 1 ) ( Γ( S k − 1 ) ) is a family of distributions over the unit-hypersphere.
The metricity of DHSW requires the continuity condition of Γ( S d − 1 ) and Γ( S k − 1 ) , i.e., they have supports on all the unit-hypersphere to guarantee the injectivity of the HRT.
We refer the reader to previous works (Nguyen et al., 2021a;b) for more details.
Augmented approaches for HSW: Authors in (Chen et al., 2022) propose to augment measures to higher dimensions for better linear separation.
We extend the framework to our hierarchical transform approach.
We first define the Spatial Hierarchical Radon transform.
Definition 15 (Spatial Hierarchical Radon Transform) .
Given a injective mapping g : R d →
R
[Equation] [Equation]
[Equation]
, Spatial Hierarchical Radon Transform
[Equation]
L
[Equation]
R × S
[Equation]
− 1) ⊗
[Equation]
× S k
[Equation]
1 (cid:17) is defined as:
[Equation]
where v ∈ R , ψ = ( ψ 1 , . . . ,
ψ k ) ∈ S k − 1 , and θ 1: k
= ( θ 1 , . .
, θ k ) ∈ ( S d ′ − 1 ) ⊗ k .
We can further extend the SHRT to hierarchical spatial Radon Transform (HSRT) by using the augmenting mappings twice.
Definition 16 (Hierarchical Spatial Radon Transform) .
Given injective mapping g 1 : R d → R d ′ ( d ′ ≥ d ) and g 2 : R k → R k ′ ( k ′ ≥ k ) , Hierarchical Spatial Radon Transform (SHRT) SHR :
L
[Equation]
d ) → L
[Equation]
R × S ( d ′ − 1) ⊗ k × S
k ′ − 1 (cid:17) is defined as:
[Equation]
where v ∈ R , ψ = ( ψ 1 , . . . ,
ψ k ) ∈
S
k ′ − 1 , and θ 1: k = ( θ 1 , . .
, θ k ) ∈ ( S d ′ − 1 ) ⊗ k .
From the above two transforms, we could derive the Augmented Hierarchical sliced Wasserstein (AHSW) distance and Hierarchical Augmented sliced Wasserstein (HASW) distance respectively.
We skip the definitions of AHSW and HASW here since they are straightforward from the definition of HSW.
The injectivity of SHRT and HSRT and the metricity of augmented distances can be proven based on the injectivity of augmenting mappings g 1 and g 2 .
We refer the reader to the previous work of Chen et al.
for more details.
C P ROOFS
In this appendix, we provide proofs for key results in the main text and in Appendix B.
C. P ROOF OF P ROPOSITION
Let us consider functions f,
g ∈ L 1 ( R d )
such that
[Equation]
It is clear from the Definition 4 that
[Equation]
Taking the integral of both sides in the above equation with respect to ( k − 1) variables t 2 , . .
, t k , we get
[Equation]
Next, by applying the Fubini’s theorem, we have
[Equation]
Note that for any i = 2 , . . .
, k − 1 , by using change of variables s i
=
t i − ⟨ x, θ i ⟩ , we have (cid:82)
R δ ( t i − ⟨ x, θ i ⟩
[Equation]
dt i =
[Equation]
R δ
[Equation]
s i ) ds i = 1 .
As a result,
[Equation]
or equivalently, ( R ) f
( t 1 , θ 1 )
=
R g ( t 1 , θ 1 ) .
Recall that the Radon transform R is injective.
Thus, we obtain that f = g , which completes the proof.
C. P ROOF OF P ROPOSITION
Let us consider arbitrary functions f, g ∈ L 1 ( R d ) satisfying ( HR
[Equation] [Equation]
, θ
[Equation] [Equation]
, ψ ) =
( HR g )( v, θ 1: k , ψ ) , where v ∈ R , ψ ∈ S k − 1 and θ 1: k ∈ ( S d − 1 ) ⊗ k .
It can be seen from the beginning of Appendix B that ( HR f )( v, θ 1: k , ψ ) =
( PR ( OR f ))
( v, ψ, θ 1: k ) and ( HR g )( v, θ 1: k , ψ ) =
( PR ( OR g ))
( v, ψ, θ 1: k ) .
Therefore, we get ( PR ( OR f ))( v, ψ, θ 1: k ) =
( PR ( OR g ))
( v, ψ, θ 1: k ) .
Since Partial Radon Transform is injective, we obtain h f = h g , or equivalently, ( OR f )( t 1: k , θ 1: k ) =
( OR g )( t 1: k , θ 1: k ) , which leads to f = g due to the injectivity of Overparametrized Radon Transform.
Hence, we reach the conclusion of the proposition.
C. P ROOF OF T HEOREM
To prove that the hierarchical sliced Wasserstein HSW p,k ( · , · ) is a metric on the space of all probability measure on R d for any p, k ≥ 1 , we need to verify four following criteria:
Symmetry:
For any p, k ≥ 1 , it is obvious that HSW p,k ( µ, ν ) =
HSW p,k ( ν, µ ) for any probability measures µ and ν .
Non-negativity: The non-negativity of HSW p,k ( · , · ) comes directly from the non-negativity of the Wasserstein metric.
Identity: For any p, k ≥ 1 , it is clear that when µ = ν , we have HSW p,k ( µ, ν )
= 0 .
Now, assume that HSW p,k ( µ, ν ) = 0 , then W p (( HR f µ )( · , θ 1: k , ψ ) , ( HR f ν )( · , θ 1: k , ψ ))
=
for almost all ψ ∈ S k − 1 , θ 1: k ∈ ( S d − 1 ) ⊗ k .
By applying the identity property of the Wasserstein distance, we have ( HR f µ )( · , θ 1: k , ψ ) =
( HR f ν )( · , θ 1: k , ψ ) for almost all ψ ∈ S k − 1 , θ 1: k ∈ ( S d − 1 ) ⊗ k .
Since the Hierarchical Radon Transform is injective, we obtain f µ = f ν , which implies that µ = ν .
Triangle Inequality: For any probability measures µ 1 , µ 2 , µ 3 , we find that
[Equation]
where the first inequality is due to the triangle inequality of Wasserstein metric, namely, we have
W p (( HR f µ 1 )
[Equation]
, ψ )
[Equation]
while the second inequality is an application of the Minkowski inequality for integrals.
Hence, the hierarchical sliced Wasserstein HSW p,k ( · , · ) is a metric on the space of all probability measures on R d for any p, k ≥ 1 .
C. P ROOF OF P ROPOSITION
The proof of this proposition is direct from the definition of the hierarchical sliced Wasserstein distance, the sliced Wasserstein distance, the max hierarchical sliced Wasserstein distance, and the max sliced Wasserstein distance.
Here, we provide the proof for the completeness.
(a
We start with
[Equation]
Subsequently, we have
[Equation]
θ k ∈ S d − 1 ,ψ ∈ S k −  inf π ∈ Π( µ,ν ) (cid:90)
R d × R d
|| ψ || p (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) k (cid:88)
j = θ ⊤ j ( x − y )
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) p dπ ( x, y ) 
 1 p = max θ 1 ,...,θ k ∈ S d −  inf π ∈ Π( µ,ν ) (cid:90)
R d × R d (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) k
(cid:88)
j = θ ⊤ j ( x − y )
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) p dπ ( x, y ) 
 1 p ≤ max θ ∈ S d −(cid:18) inf π ∈ Π( µ,ν ) (cid:90)
R d × R d k p (cid:12)(cid:12) θ
⊤ x − θ ⊤ y (cid:12)(cid:12) p dπ ( x, y ) (cid:19) 1 p = k · max θ ∈ S d −(cid:18) inf π ∈ Π( µ,ν ) (cid:90)
R d × R d (cid:12)(cid:12) θ ⊤ x − θ ⊤ y (cid:12)(cid:12) p dπ ( x, y ) (cid:19) 1 p = k · max θ ∈ S
d − 1 W p ( θ♯µ, θ♯ν ) = k · max θ ∈ S
d − 1 W p (( R ) f µ ( · , θ ) , ( R ) f ν ( · ,
[Equation]
Finally, by applying the Cauchy-Schwartz inequality, we get
[Equation]
Putting the above results together
[Equation]
we obtain the conclusion of the proposition.
(b) Firstly, it is obvious that
[Equation]
Subsequently, let θ ′
= θ ∗ = arg max θ ∈ S
d − 1 W p (( R f µ )( · , θ ) , ( R f ν )( · , θ )) and ψ ′ =
. , 0) ∈
S k − 1 .
Then, for any θ ′ 1 ∈ S d − 1 , we have
[Equation]
Since the above result holds for all arbitrary θ ′ 1 ∈ S
d −
[Equation]
we obtain
[Equation]
which completes the proof
[Equation]
C. P ROOF OF P ROPOSITION
For the ease of the presentation
[Equation]
we denote Θ ⊂ R d as the compact set of the probability measure P .
The result of Proposition 3 indicates that we have
[Equation]
where we define
[Equation]
Therefore, the conclusion of the proposition follows as long as we can demonstrate that
[Equation]
where C > 0 is some universal constant and the outer expectation is taken with respect to the data.
We first start with the property of the closed-form expression of Wasserstein distance in one dimension, which leads to the following equations:
[Equation]
where we denote F n,θ and F θ as the cumulative distributions of the two push-forward distributions θ♯P n and θ♯P .
Direct calculation indicates that
[Equation]
where we denote B as the set of half-spaces { z ∈ R d : θ ⊤ z ≤ x } for all θ ∈ R d such that ∥ θ ∥ ≤ 1 .
From Wainwright (2019), it is known that the Vapnik-Chervonenkis (VC) dimension of B is at most d + 1 .
Therefore, arrive at
[Equation]
with probability at least 1 − δ .
Collecting the above results, we finally obtain that
[Equation]
where C > 0 is some universal constant.
As a consequence, the conclusion of the proposition follows.
C. P ROOF OF P ROPOSITION
Using the Holder’s inequality
[Equation]
we have:
[Equation]
which completes the proof
[Equation]
C. P ROOF OF T HEOREM
Symmetry:
For any p ≥ 1 , it is clear that Max-HSW p ( µ, ν ) =
Max-HSW p ( ν, µ ) for any probability measures µ and ν .
Non-negativity: The non-negativity of Max-HSW p,k ( · , · ) comes directly from the non-negativity of the Wasserstein metric.
Existence of the max directions:
The unit hyperspheres S d − 1 and S k − 1 are compact and we have
W p (( HR f µ )( · , θ 1: k , ψ ) , ( HR f ν )( · , θ 1: k , ψ )) is continuous in terms of θ 1: k and ψ .
Therefore, there exists
[Equation]
Identity: For any p ≥ 1 and k ≥ 1 , it is clear that when µ = ν , then Max-HSW p ( µ, ν ) = 0 .
When Max-HSW p ( µ, ν ) = 0 , we have W p (( HR f µ )( · , θ 1: k , ψ ) , ( HR f ν )( · , θ 1: k , ψ ))
=
for almost all ψ ∈ S k − 1 , θ 1: k ∈ ( S d − 1 ) ⊗ k .
Applying the identity property of the Wasserstein distance, we have ( HR f µ )( · , θ 1: k , ψ ) =
( HR f ν )( · , θ 1: k , ψ ) almost all ψ ∈ S k − 1 , θ 1: k ∈ ( S d − 1 ) ⊗ k .
Since the
Hierarchical Radon Transform is injective, we obtain µ = ν .
Triangle Inequality: For any probability measures µ 1 , µ 2 , µ 3 , we find that
[Equation]
where the first equality is due to the existence of the max directions, the first inequality is due to the triangle inequality with Wasserstein metric, and the second inequality is an application of the definition of maxima.
D D EEP G ENERATIVE M ODELS
In this section, we first discuss the way we train generative models in our experiments with sliced
Wasserstein variants.
We then provide additional experimental results including convergence speed of the SW and the HSW in IS score, generated images from the SW and the HSW, and comparison between the Max-SW and the Max-HSW.
D. T RAINING D ETAIL
In short, we follow the generative framework from Deshpande et al.
The model distribution is parameterized as µ ϕ ( x ) ∈ P ( R d ) where µ ϕ ( x ) =
G ϕ ♯µ 0 , µ 0 :
= N (0 , I 128 ) , and G ϕ is a Resnet
He et al.
(2016)-type neural network.
The detail of the neural networks is given in Appendix F.
Previous works suggest to use a discriminator as a type of ground metric learning since the true metric between images is not given (unknown).
We denote the discriminator as a composite function D β 2 ◦ D βwhere D β 1 : R d → R d ′ and D β 2 : R d ′ → R .
In particular, D β 1 transforms original images to their latent features.
After that, D β 2 maps their features maps to their corresponding discriminative scores.
We denote the data distribution is ν data , our objectives are:
[Equation]
where m ≥ 1 is the mini-batch size and D ( · , · ) is the (Max-)SW distance or (Max-)HSW distance.
Relation to m -mini-batch energy distance: The objective for training the generator seen as a type of m -mini-batch energy distance Klebanov et al.
; Salimans et al.
(2018) with sliced
Wasserstein variants kernel.
Moreover, it is also known as mini-batch optimal transport Fatras et al.
(2020); Nguyen et al.
(2022a;b).
D. A DDITIONAL R ESULTS
Convergence of generative models from SW and HSW in IS scores: We plot the IS scores over training epochs of the SW and the HSW with the same setting as in Table 1 in Figure 4.
We observe that IS scores from models trained by the HSW (with better computation) increase faster than ones trained from the SW.
We would like to recall that reason is that the HSW has a higher number of final projections than the SW, hence, it has a more discriminative signal than the SW.
Radom generated images: We show randomly generated images from CIFAR10, CelebA, and Tiny
Imagenet from the SW and the HSW generative models in Figure 5- 7.
From these images, we observe that increasing L enhance the performance of the SW, and increasing both L and k yield better images for the HSW.
Also, the qualitative comparison from generated images between the SW and the HSW is consistent with the FID scores and IS scores in Table 1.
Comparison between the Max-HSW and the Max-SW : We would like to recall that Max-HSW is the generalization of the Max-SW since Max-HSW with k = 1 is equivalent to the Max-SW.
Therefore, the performance of Max-HSW is a least the same as the Max-SW.
To find out the benefit of k > 1 in Max-HSW, we run Max-HSW with k ∈ { 1 , 10 , 100 , 1000) , the slice learning rate η ∈ { 0 .
001 , 0 . 01 , 0 . 1 } , and the maximum number of iteration T ∈ { 1 , 10 , 100 } .
We report the best
FID scores and IS scores in Table 3.
We observe that Max-HSW k > 1 gives a better FID score than the Max-SW on CelebA dataset.
The reason might be due to the overparametrization of the
Max-HSW that leads to a better final slice.
However, on CIFAR10 and Tiny ImageNet, the Max-HSW k > 1 does not show any improvement compared to the Max-HSW k
=
(Max-SW).
This might be because the projected gradient ascent does not work well in a multiplayer structure like in HRT.
We will leave the investigation about optimization of the Max-HSW in future works since the Max-HSW has the potential to explain partially adversarial training frameworks.
E C OLOR T RANSFER VIA G RADIENT F LOW
In the color transfer setting, we are interested in transferring the color palette of a source image to the color palette of a target image.
Formally, we denote the color palette of the source image
[Figure] [Figure]
[Table] [Table]
[Table]
and the target image as X =
( x 1 , . .
, x n ) and Y =
( y 1 , . . .
y n ) of size n × 3 (RGB) in turn with n > 1 .
We would like to note that, the orders of points in X and Y are corresponding to the pixels in images.
We denote P X
= 1 n (cid:80) ni
=
δ x i and P Y = 1 n (cid:80) ni =
δ
y
i as two empirical measures on X and
Y respectively.
We now move P X to P Y under the following gradient flow: ˙ X ( t ) = ∇ X D
( P X , P Y ) with D as a chosen distance, e.g., SW and HSW.
We simply use the Euler scheme to solve this gradient flow with T iterations.
Since color is represented by three values in the set { 0 , . .
, 255 } , we round values in X ( T ) to the closest values in { 0 , . . .
, 255 } to obtain the final color palette.
[Figure] [Figure]
In our experiments, we compress original images to 3000 colors by using K-Means before doing color transfer.
In this case, n = 3000 is much larger than d = 3 , hence, using HSW does not give any computational benefits compared to SW.
Namely, computing one-dimension Wasserstein distances becomes the main computation.
However, HSW could provide the dependency between final projection directions since they are created from the same set of bottleneck projection directions.
In Figure 8, we compare SW, generalized sliced Wasserstein (GSW) (Kolouri et al., 2019b), HSW, and hierarchical generalized sliced Wasserstein (HGSW) which is a straightforward extension of HSW by replacing the Radon Transform variants by Generalized Radon Transform variants.
All the distances are computed with the order 2 ( p = 2 ) and we use the number of iterations T = 10000 with the step size 0 .
1 for the Euler scheme.
For GSW and HGSW, we use the polynomial of degree 3 as the defining function.
For SW and GSW, we set the number of projections L = 3 .
For HSW and HGSW, we set the number of bottleneck projections k = 10 and the number of final projections L = 3 .
In the figure, we report the computational time and the Wasserstein-2 distance between the corresponding color palette and the color palette of the target image.
We observe that HSW gives better Wasserstein-2 scores than SW in both two shown images while the computation time is comparable.
Similarly, the same phenomenon happens in the case of GSW and HGSW.
Comparing non-linear projecting and linear projecting, non-linear variants cost about four times more computation while their Wassersteinscores are higher.
However, it is worth noting that a better Wasserstein-2 score is not necessarily better since a non-linear transform creates a different flavor.
Overall, the hierarchical approach in
HSW and HGSW induces different gradient flows compared to SW and GSW which may yield some potential benefits.
F A DDITIONAL E XPERIMENTAL S ETTINGS
Additional settings: For all datasets, the number of training iterations is set to 50000.
We update the generator G ϕ for each 5 iterations using (Max-)SW and (Max-)HSW.
Moreover, we update D βand D β 2 ( the discriminator) every iterations.
The mini-batch size m is set 128 in all datasets.
We set the learning rate for G ϕ , D β 1 , and D β 2 to 0 . 0002 .
The optimizer that we use is Adam (Kingma & Ba, 2014) with parameters ( β 1 , β 2 ) =
(0 , 0 . 9) (slightly abuse of notations).
For SW and HSW, we use p = 2 (the order of Wasserstein distance).
FID scores and IS scores: We use 50000 random samples from trained models for computing the FID scores and the Inception scores.
In evaluating FID scores, we use all training samples for computing statistics of datasets 3 .
Neural networks: We present the architectures of our generative networks and the discriminative networks on CIFAR10, CelebA, and Tiny ImageNet as follow:.
• CIFAR10 : – G ϕ : ϵ ∈ R 128 ( ∼ N (0 , 1))
→ 4 × 4 × 256( Dense,linear ) →
ResBlock up 256 → ResBlock up 256 → ResBlock up
→ BN, ReLU, → 3 × 3 conv, 3 Tanh .
– D β 1 : x ∈ [ − 1 , 1] 32 × 32 × 3 → ResBlock down
→
ResBlock down
→
ResBlock down
→
ResBlock 128 → ResBlock 128 .
– D β 2 : x ∈ R 128 × 8 ×→ ReLU → Global sum pooling (128) → 1( Spectral normalization ) .
• CelebA and Tiny ImageNet: – G ϕ : ϵ ∈ R 128 ( ∼ N (0 , 1))
→ 4 × 4 × 256( Dense,linear ) →
ResBlock up 256 → ResBlock up 256 → ResBlock up
→
ResBlock up 256 → BN, ReLU, → 3 × 3 conv, 3 Tanh .
– D β 1 : x ∈ [ − 1 , 1] 32 × 32 × 3 → ResBlock down
→
ResBlock down
→
ResBlock down
→
ResBlock 128 → ResBlock 128 .
– D β 2 : x ∈ R 128 × 8 ×→ ReLU → Global sum pooling (128) → 1( Spectral normalization ) .