Revisiting Sliced Wasserstein on Images: From Vectorization to Convolution
Abstract
The conventional sliced Wasserstein is defined between two probability measures that have realizations as vectors .
When comparing two probability measures over images, practitioners first need to vectorize images and then project them to one-dimensional space by using matrix multiplication between the sample matrix and the projection matrix.
After that, the sliced Wasserstein is evaluated by averaging the two corresponding one-dimensional projected probability measures.
However, this approach has two limitations.
The first limitation is that the spatial structure of images is not captured efficiently by the vectorization step; therefore, the later slicing process becomes harder to gather the discrepancy information.
The second limitation is memory inefficiency since each slicing direction is a vector that has the same dimension as the images.
To address these limitations, we propose novel slicing methods for sliced Wasserstein between probability measures over images that are based on the convolution operators.
We derive convolution sliced Wasserstein (CSW) and its variants via incorporating stride, dilation, and non-linear activation function into the convolution operators.
We investigate the metricity of CSW as well as its sample complexity, its computational complexity, and its connection to conventional sliced Wasserstein distances.
Finally, we demonstrate the favorable performance of CSW over the conventional sliced
Wasserstein in comparing probability measures over images and in training deep generative modeling on images 1 .
Introduction
Optimal transport and Wasserstein distance [59, 51] have become popular tools in machine learning and data science.
For example, optimal transport has been utilized in generative modeling tasks to generate realistic images [2, 58], in domain adaptation applications to transfer knowledge from source to target domains [9, 3], in clustering applications to capture the heterogeneity of data [21], and in other applications [29, 62, 63].
Despite having appealing performance, Wasserstein distance has been known to suffer from high computational complexity, namely, its computational complexity is at the order of O ( m 3 log m ) [49] when the probability measures have at most m supports.
In addition, Wasserstein distance also suffers from the curse of dimensionality, namely, its sample complexity is at the order of O ( n − 1 /d )
[15] where n is the sample size.
A popular line of work to improve the speed of computation and the sample complexity of the Wasserstein distance is by adding an entropic regularization term to the Wasserstein distance [10].
This variant is known as entropic regularized optimal transport (or equivalently entropic regularized Wasserstein).
By using the entropic version, we the of Wasserstein distance with the computational complexities being at the order of O ( m 2 ) [1, 35, 36, 34] (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 ) [39], which indicates that it does not suffer from the curse of dimensionality.
Another useful line of work to improve both the computational and sample complexities of the
Wasserstein distance is based on the closed-form solution of optimal transport in one dimension.
A notable distance along this direction is sliced Wasserstein (SW) distance [6].
Due to the fast computational complexity O ( m log 2 m ) and no curse of dimensionality O ( n − 1 / 2 ) , the sliced Wasserstein has been applied successfully in several applications, such as generative modeling [61, 13, 25, 47], domain adaptation [31], and clustering [26].
The sliced Wasserstein is defined between two probability measures that have supports belonging to a vector space, e.g, R d .
As defined in [6], the sliced Wasserstein 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 the sliced Wasserstein distance.
The number of samples is often called the number of projections and it is denoted as L .
On the computational side, the computation of sliced Wasserstein can be decomposed into two steps.
In the first step, L projecting directions are first sampled and then stacked as a matrix (the projection matrix).
After that, the projection matrix is multiplied by the two data matrices resulting in two matrices that represent L one-dimensional projected probability measures.
In the second step, L one-dimensional Wasserstein distances are computed between the two corresponding projected measures with the same projecting direction.
Finally, the average of those distances is yielded as the value of the sliced Wasserstein.
Despite being applied widely in tasks that deal with probability measures over images [61, 13], the conventional formulation of sliced Wasserstein is not well-defined to the nature of images.
In particular, an image is not a vector but is a tensor.
Therefore, a probability measure over images should be defined over the space of tensors instead of vectors.
The conventional formulation leads to an extra step in using the sliced Wasserstein on the domain of images which is vectorization.
Namely, all images (supports of two probability measures) are transformed into vectors by a deterministic one-one mapping which is the "reshape" operator.
This extra step does not keep the spatial structures of the supports, which are crucial information of images.
Furthermore, the vectorization step also poses certain challenges to design efficient ways of projecting (slicing) samples to one dimension based on prior knowledge about the domain of samples.
Finally, prior empirical investigations indicate that there are several slices in the conventional Wasserstein collapsing the two probability measures to the Dirac Delta at zero [13, 12, 24].
Therefore, these slices do not contribute to the overall discrepancy.
These works suggest that the space of projecting directions in the conventional sliced Wasserstein (the unit hyper-sphere) is potentially not optimal, at least for images.
Contribution.
To address these issues of the sliced Wasserstein over images, we propose to replace the conventional formulation of the sliced Wasserstein with a new formulation that is defined on the space of probability measures over tensors.
Moreover, we also propose a novel slicing process by changing the conventional matrix multiplication to the convolution operators [16, 18].
In summary, our main contributions are two-fold:
We leverage the benefits of the convolution operators on images, including their efficient parameter sharing and memory saving as well as their superior performance in several tasks on images [28, 19], to introduce efficient slicing methods on sliced Wasserstein, named convolution slicers .
With those slicers, we derive a novel variant of sliced Wasserstein, named convolution sliced Wasserstein (CSW).
We investigate the metricity of CSW, its sample and computational complexities
and its connection to other variants of SW.
We illustrate the favorable performance of CSW in comparing probability measures over images.
In particular, we show that CSW provides an almost identical discrepancy between MNIST’s digits compared to that of the SW while having much less slicing memory.
Furthermore, we compare SW and CSW in training deep
generative models on standard benchmark image datasets, including CIFAR10, CelebA, STL10, and CelebA-HQ.
By considering the quality of the trained models, training speed, and training memory of CSW and SW, we observe that CSW has more favorable performance than the vanilla SW.
Organization.
The remainder of the paper is organized as follows.
We first provide background about Wasserstein distance, the conventional slicing process in the sliced Wasserstein distance, and the convolution operator in Section 2.
In Section 3, we propose the convolution slicing and the convolution sliced Wasserstein, and analyze some of its theoretical properties.
Section 4 contains the application of CSW to generative models, qualitative experimental results, and quantitative experimental results on standard benchmarks.
We conclude the paper In Section 5.
Finally, we defer the proofs of key results and extra materials in the Appendices.
Notation.
For any d ≥ 2 , S d − 1 := { θ ∈ R d | ||
θ ||
22 = 1 } denotes the d dimensional unit hyper-sphere in L 2 norm, and U ( S d − 1 ) is the uniform measure over S d − 1 .
Moreover, δ denotes the Dirac delta function.
For p ≥ 1 , P p ( R d ) is the set of all probability measures on R d that have finite p -moments.
For µ, ν ∈ P p ( R d ) , Π( µ, ν ) := { π ∈ P p ( R d × R d ) | (cid:82) R d π ( x, y ) dx =
ν, (cid:82) R d π ( x, y ) dy = µ } is the set of transportation plans between µ and ν .
For m ≥ 1 , we denotes µ ⊗ m as the product measure which has the supports are the joint vector of m random variables that follows µ .
For a vector X ∈ R dm , X :=
( x 1 , . . .
, x m ) , P X denotes the empirical measures 1 m (cid:80) mi =
δ 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.
Background
In this section, we first review the definitions of the Wasserstein distance, the conventional slicing, and the sliced Wasserstein distance, and discuss its limitation.
We then review the convolution and the padding operators on images.
Sliced Wasserstein: For any p ≥ 1 and dimension d ′ ≥ 1 , we first define the Wasserstein- p distance [59, 50] between two probability measures µ ∈ P p ( R d ′ ) and ν ∈ P p ( R d ′ ) , which is given by W p ( µ, ν ) :
=
(cid:16) inf
[Equation]
1 p .
When d ′ = 1 , the Wasserstein distance has a closed form which is W p ( µ, ν )
=
( (cid:82) 1 0 | F − 1 µ ( z ) − F − 1 ν ( z )
| p dz ) 1 /p where F µ and F ν are the cumulative distribution function (CDF) of µ and ν respectively.
Given this closed-form property of Wasserstein distance in one dimension, the sliced Wasserstein distance [6] between µ and ν had been introduced and admitted the following formulation:
SW pp ( µ, ν ) :
= (cid:82) S d − 1 W pp ( θ♯µ, θ♯ν ) dθ , where θ♯µ is the push-forward probability measure of µ through the function T θ : R d ′ → R with T θ ( x )
= θ
⊤ x .
For each θ ∈ S d ′ − 1 , W pp ( θ♯µ, θ♯ν ) can be computed in linear time O ( m log 2 m ) where m is the number of supports of µ and ν .
However, the integration over the unit sphere in the sliced Wasserstein distance is intractable to compute.
Therefore,
Monte Carlo scheme is employed to approximate the integration, namely, θ 1 , . . .
, θ L ∼ U ( S d ′ − 1 ) are drawn uniformly from the unit sphere and the approximation of the sliced Wasserstein distance is given by:
(cid:100) SW p p
[Equation]
µ, ν ) ≈
[Equation]
Li =1 W pp ( θ i ♯µ, θ i ♯ν ) .
In practice, L should be chosen to be sufficiently large compared to the dimension d ′ , which can be undesirable.
Sliced Wasserstein on Images: Now, we focus on two probability measures over images: µ, ν ∈
P p
( R c × d × d ) for number of channels c ≥ 1 and dimension d ≥ 1 .
In this case, the sliced Wasserstein between µ and ν is defined as:
[Equation]
where R : R c × d × d → R cd 2 is a deterministic one-to-one "reshape" mapping.
The slicing process: The slicing of sliced Wasserstein distance on probability measures over images consists of two steps: vectorization and projection.
Suppose that the probability measure µ ∈ P ( R c × d × d ) has n supports.
Then the supports of µ are transformed into vectors in R cd 2 and are stacked as a matrix of size n × cd 2 .
A projection matrix of size L × cd 2 is then sampled and has each column as a random vector following the uniform measure over the unit hyper-sphere.
Finally, the multiplication of those two matrices returns L projected probability measures of n supports in one dimension.
We illustrate this process in Figure 1.
Limitation of the conventional slicing: First of all, images contain spatial relations across channels and local information.
Therefore, transforming images into vectors makes it challenging to obtain that information.
Second, vectorization leads to the usage of projecting directions from the unit hyper-sphere, which can have several directions that do not have good discriminative power.
Finally, sampling projecting directions in high-dimension is also time-consuming and memory-consuming.
As a consequence, avoiding the vectorization step can improve the efficiency of the whole process.
[Figure]
Convolution operator: We now define the convolution operator on tensors [16], which will be used as an alternative way of projecting images to one dimension in the sliced Wasserstein.
The definition of the convolution operator with stride and dilation is as follows.
Definition 1 (Convolution)
Given the number of channels c ≥ 1 , the dimension d ≥ 1 , the stride size s ≥ 1 , the dilation size b ≥ 1 , the size of kernel k ≥ 1 , the convolution of a tensor X ∈ R c × d × d with a kernel size K ∈ R c × k × k is X s,b ∗ K =
Y,
Y ∈ R 1 × d ′ ×
d ′ where d ′ =
d − b ( k − 1) − 1 s
+ 1 .
For i = 1 , . . .
, d ′ and j = 1 , . . .
, d ′ , Y 1 ,i,j is defined as:
[Equation]
From its definition, we can check that the computational complexity of the convolution operator is
O
[Equation]
Convolution Sliced Wasserstein
In this section, we will define a convolution slicer that maps a tensor to a scalar by convolution operators.
Moreover, we discuss the convolution slicer and some of its specific forms including the convolution-base slicer, the convolution-stride slicer, the convolution-dilation slicer, and their non-linear extensions.
After that, we derive the convolution sliced Wasserstein (CSW), a family of variants of sliced Wasserstein, that utilizes a convolution slicer as the projecting method.
Finally, we discuss some theoretical properties of CSW, namely, its metricity, its computational complexity, its sample complexity, and its connection to other variants of sliced Wasserstein.
3.1 Convolution Slicer
We first start with the definition of the convolution slicer, which plays an important role in defining convolution sliced Wasserstein.
Definition 2 (Convolution Slicer)
For N ≥ 1 , given a sequence of kernels K (1) ∈
R c (1) × d (1) × d (1) , . . .
, K ( N ) ∈ R c
( N ) × d ( N ) × d ( N ) , a convolution slicer S ( ·| K (1) , . . .
, K ( N ) ) on
R c × d × d is a composition of N convolution functions with kernels K (1) , . .
, K ( N ) (with stride or dilation if needed) such that S ( X | K (1) , . .
, K ( N ) )
∈ R ∀
X ∈ R c
×
d × d .
As indicated in Definition 2, the idea of the convolution slicer is to progressively map a given data
X to a one-dimensional subspace through a sequence of convolution kernels, which capture spatial relations across channels as well as local information of the data.
It is starkly different from the vectorization step in standard sliced Wasserstein on images (1).
The illustration of the convolution slicer is given in Figure 2.
Sampling
Sampling
Sampling * * * * * * * * * * * *
[Figure]
We consider three particular types of convolution slicers based on using linear function on the convolution operator, named convolution-base, convolution-stride, and convolution-dilation slicers.
We defer the definition of convolution-dilation slicers to Definition 5.
We first start with the definition of the convolution-base slicer.
Definition 3 (Convolution-base Slicer)
Given X ∈ R c × d × d ( d ≥ 2 ), 1.
When d is even, N is the biggest integer that satisfies d = 2 N − 1 a with a is also an integer, sliced kernels are defined as K (1) ∈ R c × (2 − 1 d +1) ×
(2 − 1 d +1) and K ( h ) ∈ R 1 × (2 − h d
+1)
×
(2 − h d +1) for h = 2 , . . .
, N − 1 , and K ( N ) ∈ R 1 × a × a where a = d 2 N − 1 .
Then, the convolution-base slicer
CS -b ( X | K (1) , . .
, K ( N ) ) is defined as:
[Equation]
When d is odd, the convolution-base slicer CS -b ( X | K (1) , . .
, K ( N ) ) takes the form:
[Equation]
where K (1) ∈ R c × 2 × 2 and K (2) , . .
, K ( N ) are the corresponding sliced kernels that are defined on the dimension d − 1 .
The idea of the convolution-base slicer in Definition 3 is to reduce the width and the height of the image by half after each convolution operator.
If the width and the height of the image are odd, the first convolution operator is to reduce the size of the image by one via convolution with kernels of size 2 × 2 , and then the same procedure as that of the even case is applied.
We would like to remark that the conventional slicing of sliced Wasserstein in Section 2 is equivalent to a convolution-base slicer S ( ·| K (1) ) where K (1) ∈ R c × d × d that satisfies the constraint (cid:80) ch
=1 (cid:80) di
=1 (cid:80) dj
=1 K (1)2 h,i,j = 1 .
We now discuss the second variant of the convolution slicer, named convolution-stride slicer, where we further incorporate stride into the convolution operators.
Its definition is as follows.
Definition 4 (Convolution-stride Slicer)
Given X ∈ R c × d × d ( d ≥ 2 ), 1.
When d is even, N is the biggest integer that satisfies d = 2 N − 1 a with a is also an integer, sliced kernels are defined as K (1) ∈ R c × 2 × 2 and K ( h ) ∈ R 1 × 2 × 2 for h = 2 , . .
, N − 1 , and
K ( N ) ∈ R 1 × a × a where a = d 2 N − 1 .
Then, the convolution-stride slicer CS -s ( X | K (1) , . .
, K ( N ) ) is defined as:
[Equation]
. . , K ( N ) )
=
X ( N ) ,
X ( h ) =
 
X h =
X ( h − 1) 2 , 1 ∗ K ( h ) 1 ≤ h ≤ N − 1 ,
X ( h − 1) 1 , 1 ∗ K
( h ) h
= N,
When d is odd, the convolution-stride slicer CS
-s ( X | K (1) , . . .
, K ( N ) ) takes the form:
[Equation]
where K (1) ∈ R c × 2 × 2 and K (2) , . .
, K ( N ) are the corresponding sliced kernels that are defined on the dimension d − 1 .
Similar to the convolution-base slicer in Definition 3, the convolution-stride slicer reduces the width and the height of the image by half after each convolution operator.
We use the same procedure of reducing the height and the width of the image by one when the height and the width of the image are odd.
The benefit of the convolution-stride slicer is that the size of its kernels does not depend on the width and the height of images as that of the convolution-base slicer.
This difference improves the computational complexity and time complexity of the convolution-stride slicer over those of the convolution-base slicer (cf. Proposition 1). Definition 5 (Convolution-dilation Slicer)
Given X ∈ R c × d × d ( d ≥ 2 ), 1. When d is even, N is the biggest integer that satisfies d = 2 N − 1 a with a is also an integer, sliced kernels are defined as K (1) ∈ R c × 2 × 2 and K ( h ) ∈ R 1 × 2 × 2 for h = 2 , . . , N − 1 , and K ( N ) ∈ R 1 × a × a where a = d 2 N − 1 . Then, the convolution-dilation slicer CS -d ( X | K (1) , . .
, K ( N ) ) is defined as:
[Equation]
When d is odd, the convolution-dilation slicer CS -d ( X | K (1) , . .
, K ( N ) ) takes the form:
[Equation]
where K (1) ∈ R c × 2 × 2 and K (2) , . .
, K ( N ) are the corresponding sliced kernels that are defined on the dimension d − 1 .
As with the previous slicers, the convolution-dilation slicer also reduces the width and the height of the image by half after each convolution operator and it uses the same procedure for the odd dimension cases.
The design of kernels’ size of the convolution-dilation slicer is the same as that of the convolution-stride slicer.
However, the convolution-dilation slicer has a bigger receptive field in each convolution operator which might be appealing when the information of the image is presented by a big block of pixels.
Computational and projection memories complexities of the convolution slicers: We now establish the computational and projection memory complexities of convolution-base, convolutionstride, and convolution-dilation slicers in the following proposition.
We would like to recall that the projection memory complexity is the memory that is needed to store a slice (convolution kernels).
Proposition 1 (a) When d is even, N is the biggest integer that satisfies d = 2 N − 1 a with a is also an integer, and N = [log 2 d ] , the computational and projection memory complexities of convolution-base slicer are respectively at the order of O ( cd 4 ) and O ( cd 2 ) .
When d is odd, these complexities are at the order of O ( cd 2 + d 4 ) and O ( c + d 2 ) .
b) The computational and projection memory complexities of convolution-stride slicer are respectively at the order of O ( cd 2 ) and O ( c + [log 2 d ]) .
(c) The computational and projection memory complexities of convolution-dilation slicer are respectively at the order of O ( cd 2 ) and O ( c + [log 2 d ]) .
Proof of Proposition 1 is in Appendix B.4.
We recall that the computational complexity and the projection memory complexity of the conventional slicing in sliced Wasserstein are O ( cd 2 ) and
O ( cd 2 ) .
We can observe that the convolution-base slicer has a worse computational complexity than the conventional slicing while having the same projection memory complexity.
Since the size of kernels does not depend on the size of images, the convolution-stride slicer and the convolutiondilation slicer have the same computational complexity as the conventional slicing O ( cd 2 ) .
However, their projection memory complexities are cheaper than conventional slicing, namely, O ( c + [log 2 d ]) compared to O ( cd 2 ) .
Non-linear convolution-base slicer:
The composition of convolution functions in the linear convolution slicer and its linear variants is still a linear function, which may not be effective when the data lie in a complex and highly non-linear low-dimensional subspace.
A natural generalization of linear convolution slicers to enhance the ability of the slicers to capture the non-linearity of the data is to apply a non-linear activation function after convolution operators.
This enables us to define a non-linear slicer in Definition 7 in Appendix C.
The non-linear slicer can be seen as a defining function in generalized Radon Transform [52] which was used in generalized sliced Wasserstein [24].
3.2 Convolution Sliced Wasserstein
Given the definition of convolution slicers, we now state general definition of convolution sliced
Wasserstein.
An illustration of the convolution sliced Wasserstein is given in Figure 2.
Definition [-
-]6 For any p ≥ 1 , the convolution sliced Wasserstein (CSW) of order p > 0 between two given probability measures µ, ν ∈ P p ( R c × d × d ) is given by:
[Equation]
where the expectation is taken with respect to K (1) ∼ U
( K (1) ) , . . .
, K ( N ) ∼ U ( K ( N ) ) .
Here,
S ( ·| K (1) , . . .
, K ( N ) ) is a convolution slicer with K ( l ) ∈ R c
( l ) × k ( l ) × k
( l ) for any l ∈ [ N ] and U ( K ( l ) ) is the uniform distribution with the realizations being in the set K ( l ) which is defined as K ( l ) :
= (cid:110)
K ( l ) ∈ R c
( l ) × k ( l ) × k ( l ) | (cid:80) c ( l ) h =1 (cid:80) k
( l ) i ′
=
1 (cid:80) k
( l ) j ′
=
1 K ( i )2 h,i ′ ,j ′ = 1 (cid:111) , namely, the set K ( l ) consists of tensors
K ( l ) whose squared ℓ 2 norm is 1.
The constraint that ℓ 2 norms of K ( l ) is 1 is for guaranteeing the distances between projected supports are bounded.
When we specifically consider the convolution slicer as convolution-base slicer ( CS -b), convolution-stride slicer ( CS -s), and convolution-dilation slicer ( CS -d), we have the corresponding notions of convolution-base sliced Wasserstein (CSW-b), convolution-stride sliced Wasserstein (CSW-s), and convolution-dilation sliced Wasserstein (CSW-d).
Monte Carlo estimation and implementation: Similar to the conventional sliced Wasserstein, the expectation with respect to kernels K (1) , . .
, K ( N ) uniformly drawn from the sets K (1) , . .
, K ( N ) in the convolution sliced Wasserstein is intractable to compute.
Therefore, we also make use of Monte
Carlo method to approximate the expectation, which leads to the following approximation of the convolution sliced Wasserstein:
[Equation]
, (2) where K ( ℓ ) i are uniform samples from the sets K ( ℓ ) (which is equivalent to sample uniformly from S c ( l ) · k ( l )2 then applying the one-to-one reshape mapping) for any ℓ ∈ [ N ] and
i ∈
[ L ] .
Since each of the convolution slicer S ( ·| K (1) i , . . .
, K ( N ) i ) is in one dimension, we can utilize the closed-form expression of Wasserstein metric in one dimension to compute
W p (cid:16) S ( ·| K (1) i , . . .
, K ( N ) i ) ♯µ, S ( ·| K (1) i , . .
, K ( N ) i ) ♯ν (cid:17) with a complexity of O ( m log 2 m ) for each i ∈
[ L ] where m is the maximum number of supports of µ and ν .
Therefore, the total computational complexity of computing the Monte Carlo approximation (2) is O
( Lm log 2 m ) when the probability measures µ and ν have at most m supports.
It is comparable to the computational complexity of sliced
Wasserstein on images (1) where we directly vectorize the images and apply the Radon transform to these flatten images.
Finally, for the implementation, we would like to remark that L convolution slicers in equation (2) can be computed independently and parallelly using the group convolution implementation which is supported in almost all libraries.
Properties of convolution sliced Wasserstein: We first have the following result for the metricity of the convolution sliced Wasserstein.
Theorem 1 For any p ≥ 1 , the convolution sliced Wasserstein CSW p ( ., . ) is a pseudo-metric on the space of probability measures on R c × d × d , namely, it is symmetric, and satisfies the triangle inequality.
Proof of Theorem 1 is in Appendix B.1.
We would like to mention that CSW can might still be a metric since the convolution slicer might be injective.
Our next result establishes the connection between the convolution sliced Wasserstein and max-sliced Wasserstein and Wasserstein distances.
Proposition 2 For any p ≥ 1 , we find that CSW p ( µ, ν ) ≤ Max-SW p ( µ, ν ) ≤ W p ( µ, ν ) , where
Max-SW p ( µ, ν
[Equation]
:=
[Equation]
is max-sliced Wasserstein of order p .
Proof of Proposition 2 is in Appendix B.2.
Given the bounds in Proposition 2, we demonstrate that the convolution 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 3 Assume that P is a probability measure supported on compact set of R c × d × 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 3 is in Appendix B.3.
The result of Proposition 3 indicates that the sample complexity of the convolution sliced Wasserstein is comparable to that of the sliced Wasserstein on images (1), which is at the order of O ( n − 1 / 2 ) [4], and better than that of the Wasserstein metric, which is at the order of O ( n − 1 / (2 cd 2 ) )
Extension to non-linear convolution sliced Wasserstein:
In Appendix C, we provide a non-linear version of the convolution sliced Wasserstein, named non-linear convolution sliced Wasserstein.
The high-level idea of the non-linear version is to incorporate non-linear activation functions to the convolution-base, convolution-stride, and convolution-dilation slicers.
The inclusion of non-linear activation functions is to enhance the ability of slicers to capture the non-linearity of the data.
By plugging these non-linear convolution slicers into the general definition of the convolution sliced
Wasserstein in Definition 6, we obtain the non-linear variants of convolution sliced Wasserstein.
Experiments
In this section, we focus on comparing the sliced Wasserstein (SW) (with the conventional slicing), the convolution-base sliced Wasserstein (CSW-b), the convolution sliced Wasserstein with stride (CSW-s), and the convolution sliced Wassersstein with dilation (CSW-d) in training generative models on standard benchmark image datasets such as CIFAR10 (32x32) [27], STL10 (96x96) [8], CelebA (64x64), and CelebA-HQ (128x128) [37].
We recall that the number of projections in SW and CSW’s variants is denoted as L .
Finally, we also show the values of the SW and the CSW variants between probability measures over digits of the MNIST dataset [30] in Appendix D.1.
From experiments on
MNIST, we observe that values of CSW variants are similar to values of SW while having better projection complexities.
In generative modeling, we follow the framework of the sliced Wasserstein generator in [13] with some modifications of neural network architectures.
The details of the training are given in Appendix D.2.
We train the above model on standard benchmarks such as CIFAR10 (32x32) [27], STL10 (96x96) [8],
CelebA (64x64), and CelebAHQ (128x128) [37].
To compare models, we use the FID score [20] and the Inception score (IS) [55].
The detailed settings about architectures, hyperparameters, and evaluation of FID and IS are given in
Appendix E. We first show the FID scores and IS scores of generative models trained by SW and CSW’s variants with the number of projections L ∈ { 1 , 100 , 1000 } in Table 1.
In the table, we report the performance of models at the last training epoch.
We do not report the IS scores on CelebA and CelebA-HQ since the IS scores are not suitable for face images.
We then demonstrate the FID scores and IS scores across training epochs in Figurefor investigating the convergence of generative models trained by SW and CSW’s variants. After [Figure]
[Table] that, we report the training time and training memory of SW and CSW variants in Table 5.
Finally, we show randomly generated images from SW’s models and CSW-s’ models on CelebA dataset in
Figure 4.
Generated images of all models on all datasets are given in Figures 5-8 in Appendix D.2.
Summary of FID scores and IS scores: According to Table 1, on CIFAR10, CSW-d gives the lowest values of FID scores and IS scores when L = 1 while CSW-s gives the lowest FID scores
when
L = 100 and L = 1000 .
Compared to CSW-s, CSW-d and CSW-b yield higher FID scores and lower
IS scores.
However, CSW-d and CSW-b are still better than SW.
On CelebA, CSW-s performs the best in all settings.
On STL10, CSW’s variants are also better than the vanilla SW; however, it is unclear which is the best variant.
On CelebA-HQ, SW gives the lowest FID score when L = 1 .
In contrast, when L = 100 and L = 1000 , CSW-s is the best choice for training the generative model.
Since the FID scores of L = 1 are very high on CelebA-HQ and STL10, the scores are not very meaningful for comparing SW and CSW’s variants.
For all models, increasing L leads to better quality.
Overall, we observe that CSW’s variants enhance the performance of generative models.
FID scores and IS scores across epochs: From Figure 3, we observe that CSW’s variants help the generative models converge faster than SW
when L = 100 and L = 1000 .
Increasing the number of projections from 100 to 1000 , the generative models from both SW and CSW’s variants become better.
Overall, CSW-s is the best option for training generative models among CSW’s variants since its FID curves are the lowest and its IS curves are the highest.
Generated images: We show randomly generated images on CelebA dataset in Figure 4 and Figure(Appendix D), and generated images on CIFAR10, CelebA, STL10, and CelebA-HQ in Figures 5-8 as
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure] [Equation]
[Figure] [Equation] [Figure] qualitative comparison between SW and CSW variants.
From the figures, we can see that generated images of CSW-s is more realistic than ones of SW.
The difference is visually clear when the number of projections L is small e.g., L = 1 and L = 100 .
When L = 1000 , we can still figure out that
CSW-s is better than SW by looking at the sharpness of the generated images.
Also, we can visually observe the improvement of SW and CSW-s when increasing the number of projections.
In summary, the qualitative results are consistent with the quantitative results (FID scores and IS scores) in Table 1.
For the generated images of CSW-b and CSW-d, we also observe the improvement compared to the
SW which is consistent with the improvement of FID scores and IS scores.
Non-linear convolution sliced Wasserstein: We also compare non-linear extensions of SW and
CSW variants in training generative models on CIFAR10 in Appendix D. For details of non-linear extensions, we refer to Appendix C. From experiments, we observe that convolution can also improve the performance of sliced Wasserstein in non-linear projecting cases.
Compared to linear versions, non-linear versions can enhance the quality of the generative model or yield comparable results.
Conclusion
We have addressed the issue of the conventional slicing process of sliced Wasserstein when working with probability measures over images.
In particular, sliced Wasserstein is defined on probability measures over vectors which leads to the step of vectorization for images.
As a result, the conventional slicing process cannot exploit the spatial structure of data for designing the space of projecting directions and projecting operators.
To address the issue, we propose a new slicing process by using the convolution operator which has been shown to be efficient on images.
Moreover, we investigate the computational complexity and projection memory complexity of the new slicing technique.
We show that convolution slicing is comparable to conventional slicing in terms of computational complexity while being better in terms of projection memory complexity.
By utilizing the new slicing technique, we derive a novel family of sliced Wassersein variants, named convolution sliced Wasserstein.
We investigate the properties of the convolution sliced Wasserstein including its metricity, its computational and sample complexities, and its connection to other variants of sliced
Wasserstein in literature.
Finally, we carry out extensive experiments in comparing digits images and training generative models on standard benchmark datasets to demonstrate the favorable performance of the convolution sliced Wasserstein.
Acknowledgements
NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of
Machine Learning.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2. If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results? [Yes] (b) Did you include complete proofs of all theoretical results? [Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A] Supplement to "Revisiting Sliced Wasserstein on Images: From Vectorization to Convolution" In this supplement, we first discuss related works and the potential impacts and limitations of our works in Appendix A. In Appendix B, we provide proofs for key results in the paper.
In Appendix C, we introduce non-linear versions of the convolution sliced Wasserstein, max convolution sliced Wasserstein, and convolution projected robust Wasserstein.
In Appendix D, we include additional experiments for comparing measures over MNIST’s digits via sliced Wasserstein and convolution sliced Wasserstein.
Also, we further provide generated images for convolution sliced Wasserstein under generative model settings , and generative experiemnts on max convolution sliced Wasserstein and convolution projected robust Wasserstein.
Finally, in Appendix E, we include details of experimental settings in the paper.
A Related Works, Potential Impact, and Limitations
Sliced Wasserstein is used for the pooling mechanism for aggregating a set of features in [40].
Sliced Wasserstein gradient flows are investigated in [38, 5].
Variational inference based on sliced
Wasserstein is carried out in [64].
Similarly, sliced Wasserstein is used for approximate Bayesian computation in [41].
Statistical guarantees of training generative models with sliced Wasserstein is derived in [43].
Other frameworks for generative modeling using sliced Wasserstein are sliced iterative normalizing flows [11] and run-sort-rerun for fine-tuning pre-trained model [32].
Differentially private sliced Wasserstein is proposed in [53].
Approximating Wasserstein distance based on one-dimensional transportation plans from orthogonal projecting directions is introduced in [54].
To reduce the projection complexity of sliced Wasserstein, a biased approximation based on the concentration of
Gaussian projections is proposed in [42].
Augmenting probability measures to a higher-dimensional space for a better linear separation is used in augmented sliced Wasserstein [7].
Projected Robust
Wasserstein (PRW) metrics that find the best orthogonal linear projecting operator onto k >dimensional space and Riemannian optimization techniques for solving it are proposed in [48, 33, 22].
Sliced Gromov Wasserstein, a fast sliced version of Gromov Wasserstein, is proposed in [57].
The slicing technique is also be applied in approximating mutual information [17].
We would like to recall that all the above works assume working with vector spaces and need to use vectorization when dealing with images.
In [56], convolution is used for learning the ground cost metric of optimal transport while it is used to project measures to one-dimensional measures in our work.
Potential Impact: This work addresses a fundamental problem of designing a slicing process for sliced Wasserstein on images and it can be used in various applications that perform on images.
Therefore, it could create negative potential impacts if it is used in applications that do not have good purposes.
Limitations: One limitation of CSW is that it is a pseudo metric on the space of all distribution over tensors.
However, this is because we do not assume any structure on distribution over images.
In practice, many empirical investigations show that image datasets belong to some geometry group (symmetry, rotation invariant, translation invariant, and so on).
Therefore, the set of distributions over images might be a subset of the set of distributions over tensors.
If the convolutional transform can hold the injectivity on the set of distributions over images, CSW can be a metric on the space of distributions over images.
In our applications, we compare the value of sliced Wasserstein and convolution sliced Wasserstein on MNIST digits in Table 4 in Appendix D.1, we found that the values of CSW are closed to the value of SW that can be considered as a test for our hypothesis of metricity of CSW.
To our best knowledge, there is no formal definition of the space of distributions over images and its property.
Therefore, we will leave this for future work.
In deep learning applications, sliced Wasserstein is computed between empirical distributions over mini-batches of samples that are randomly drawn from the original distribution [1].
This is known as mini-batch optimal transport with sliced Wasserstein kernel that is used when dealing with very large scale distributions and implicit continuous distributions.
When using mini-batches, both Wasserstein distance, sliced Wasserstein distance, and convolutional sliced Wasserstein will lose its metricity to become a loss [14].
Therefore, metricity is not the deciding factor in some applications of sliced
Wasserstein such as deep generative model, deep domain adaptation, and so on.
This partially explains the better performance of CSW on our deep generative model experiments in Table 1.
B Proofs
In this appendix, we provide proofs for key results in the main text.
B.1 Proof of Theorem
For any p ≥
[Equation]
it is clear that when µ = ν , then CSW p ( µ, ν ) = 0 .
Furthermore, CSW p ( µ, ν )
[Equation]
CSW p ( ν, µ ) for any probability measures µ and ν .
Therefore, to obtain the conclusion of the theorem [Equation] it is sufficient to demonstrate that is satisfies the triangle inequality.
Indeed, for any probability measures µ 1 , µ 2 , µ
[Equation]
we find that
[Equation]
where the first inequality is due to the triangle inequality with Wasserstein metric
[Equation]
namely
[Equation]
we have
[Equation]
while the second inequality is an application of Minkowski inequality for integrals
[Equation]
As a consequence
[Equation]
we obtain the conclusion of the theorem.
B.2 Proof of Proposition
The proof of this proposition is direct from the definition of the convolution sliced Wasserstein.
Here, we provide the proof for the completeness.
Indeed, since the convolution slicer S ( ·| K (1) , . . .
, K ( N ) ) is a mapping from R c × d × d to R , it is clear that
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
where the second inequality is due to the inequality with ℓ 2 norm of convolution of matrices and the fact that the ℓ 2 norm of each tensor K ( i ) is 1 for all i ∈
[ N ] .
In addition, we find that
[Equation]
Putting the above results together, we obtain the conclusion of the proposition.
B.3 Proof of Proposition
From the assumption of Proposition 3, we denote Θ ⊂ R c ×
d × d as the compact set that the probability measure P is supported on.
Based on the result of Proposition 2, we have
[Equation]
where max-SW p ( P n , P )
=
max θ
[Equation]
R cd
[Equation]
: ∥
[Equation] ∥≤ [Equation] W p ( θ♯P n , θ♯P ) .
Therefore, to obtain the conclusion of the proposition, it is sufficient to demonstrate that E [ max-SW p ( P n , P )] ≤ C [Equation] cd 2 + 1) log 2 n/n for some universal constant C > 0 .
Indeed, from the closed-form expression of Wasserstein metric in one dimension, we have
[Equation]
where F n,θ and F θ are respectively the cumulative distributions of θ♯P n and θ♯P .
Furthermore, we have the following relation:
[Equation]
where A is the set of half-spaces { y ∈ R cd 2 : θ ⊤ y ≤ x } for all θ ∈ R cd 2 such that ∥ θ ∥ ≤ 1 .
The
Vapnik-Chervonenkis (VC) dimension of A is upper bounded by cd 2 + 1 (see the reference [60]).
Therefore, with probability at least 1 − δ we obtain that
[Equation]
Putting the above results together, we can conclude that E [ max-SW p ( P n , P )] ≤
C (cid:112) ( cd 2 + 1) log 2 n/n for some universal constant C > 0 .
As a consequence, we obtain the conclusion of the proposition.
B.4 Proof of Proposition
(a) We first consider the computational and projection memory complexities of the convolution-base slicer.
When d is even, we can write down d = 2 [log 2 d ] − 1 · d 2 [log2 d ] − 1 .
Direct calculation indicates that the computational complexity of convolution-base slicer is
O
[Equation]
Similarly [Equation] we can check that the projection memory complexity of convolution-base slicer is
[Equation]
Therefore
[Equation]
we obtain the conclusion of part (a) when d is even.
Moving to the case when d is odd, the computational complexity of convolution-base slicer becomes
O
[Equation]
Similarly, we can check that when d is odd, the projection memory complexity of convolution-base slicer is O (cid:16) [Equation] + [Equation]
1 −
[Equation]
(cid:17) d
[Equation]
(cid:17)
= O ( cd 2 ) .
As a consequence, we obtain our claims with the computational and projection memory complexities of convolution-base slicer.
(b) We now establish the computational and projection memory complexities of convolution-stride slicer.
When d is even, we can write down d = 2 [log 2 d ] − 1 · d 2 [log2 d ] − 1 .
Then, the computational complexity of convolution-stride slicer is
O
[Equation]
=
[Equation]
Similarly, the projection memory complexity of convolution-stride slicer is
O
[Equation]
When d is odd, the computational complexity of convolution-stride slicer is O
[Equation]
Similarly, we can check that when d is odd, the projection memory complexity of convolution-stride slicer is O (cid:16) 4 c +
[Equation]
d
[Equation]
[log
[Equation]
d − [Equation] [Equation] [log
[Equation] d ]) . As a consequence, we obtain the conclusion of part (b).
(c) Since the convolution-dilation slicer is designed in the same way as that of the convolution-stride slicer, its computational complexity and projection memory complexity can be derived in the same manner as those of the convolution-stride slicer.
As a consequence, we reach the conclusion of part (c).
C Non-linear Versions of Convolution Sliced Wasserstein, Max Convolution Sliced Wassestein, and Convolution Projected Robust Wasserstein
In this appendix, we consider an extension of convolution sliced Wasserstein to non-linear convolution sliced Wasserstein to enhance the ability of convolution sliced Wasserstein to capture the non-linearity of the data.
Moreover, we also propose the max sliced version of convolution sliced Wasserstein to overcome the projection complexity [12].
Non-linear convolution sliced Wasserstein: We first state the definition of non-linear convolutionbase slicer.
Definition 7 (Non-Linear Convolution-base Slicer) Given X ∈ R c × d × d ( d ≥ 2 ) and a non linear activation σ ( · ) , 1.
When d is even, N = [log 2 d ] , sliced kernels are defined as K (
1) ∈ R 1 × 2 − 1 d +1 × 2 − 1 d +
and
K ( [Equation] ) ∈ R 1 × 2 − h d +1 × 2 − h d
+1 for h = 2 , . . .
, N − 1 , and K ( N ) ∈ R 1 × a × a where a = d 2 N − 1 .
Then, the non-linear convolution-base slicer NCS -b
( X | K (1) , . . .
, K ( N ) ) is defined as:
[Equation]
When d is odd, the non-linear convolution-base slicer NCS -b
( X | K (1) , . . .
, K ( N ) ) takes the form:
[Equation]
where K (1) ∈ R c × 2 × 2 and K (2) , . .
, K ( N ) are the corresponding sliced kernels that are defined on the dimension d − 1 .
The main idea of non-linear convolution-based slicer is that we incorporate non-linear activation function σ ( . ) into the layers of the convolution-base slicer.
Using that idea, we also can extend the convolution-stride and convolution-dilation slicers to their nonlinear versions, named nonlinear convolution-stride and convolution-dilation slicers.
We respectively denote these slicers as NCS -s
( X | K (1) , . . .
, K ( N ) ) and NCS -d
( X | K (1) , . . .
, K ( N ) ) .
Using the non-linear convolution-base slicer for sliced Wasserstein, we obtain the corresponding non-linear convolution-base sliced Wasserstein as follows.
[Table]
Definition 8 For any p ≥ 1 , the non-linear convolution-base sliced Wasserstein (NCSW-b) of order p > 0 between two given probability measures µ, ν ∈ P p ( R c × d × d ) is given by:
[Equation]
where NCS -b ( ·| K (1) , . .
, K ( N ) ) is a non-linear convolution-base slicer with K ( i
) ∈
R c
( i ) × k
( i )
× k
( i ) for any i ∈
[ N ] and U ( K ( i ) ) is the uniform distribution with the realizations being in the set
K ( i ) =
[Equation]
By replacing the non-linear convolution-base slicer NCS -b ( ·| K (1) , . . .
, K ( N ) )
in Definition 5 by nonlinear convolution-stride slicer NCS -s ( ·| K (1) , . . .
, K ( N ) ) and non-linear convolution-dilation slicer NCS -d ( ·| K (1) , . . .
, K ( N ) ) , we respectively have the non-linear convolution-stride sliced Wasserstein (NCSW-s) and non-linear convolution-dilation sliced Wasserstein (NCSW-d).
In Appendix D, we provide experiment results with non-linear convolution sliced Wasserstein on generative models.
Max Convolution sliced Wasserstein: Similar to the definition of Max-SW [12], the definition of max convolution sliced Wasserstein (Max-CSW) is as follow:
Definition 9 For any p ≥ 1 , the max convolution sliced Wasserstein (Max-SW) of order p >between two given probability measures µ, ν ∈ P p ( R c × d × d ) is given by:
[Equation]
where S ( ·| K (1) , . . .
, K ( N ) ) is a convolution slicer and K 1 , . . .
, K N are defined as in Definition 6.
The constrained optimization in Max-CSW is solved by projected gradient ascent that is similar to
[Table]
Max-SW.
Similar to CSW, Max-CSW also has three variants that are corresponding to three types of proposed convolution slicer, namely, Max-CSW-b, Max-CSW-s, and Max-CSW-d.
Convolution projected robust Wasserstein:
As a generalization of Max-SW, projected robust
Wasserstein (PRW) [48] finds the best subspace of k > 1 dimension that can maximize the Wasserstein distance between projected measures.
Given two probability measures µ, ν ∈ P p ( R d ) , the projected robust Wasserstein distance between µ and ν is defined as:
[Equation]
where V k ( R d ) := { U ∈ R d × k | U ⊤ U
= I k } is the Stefel Manifold.
To define the convolution projected robust Wasserstein, we first define the k -convolution slicers:
Definition 10 ( k -Convolution Slicer)
For N ≥ 1 , given a sequence of kernels K (1) ∈
R c (1) × d (1) × d (1) , . . .
, K ( N ) ∈ R c
( N ) × d ( N ) × d ( N ) ,
a k -convolution slicer S k ( ·| K (1) , . . .
, K ( N ) ) on
R c × d × d is a composition of N convolution functions with kernels K (1) , . .
, K ( N ) (with stride or dilation if needed) such that S k
( X | K (1) , . . .
, K ( N ) )
∈ R k ∀
[Equation]
∈ R c × d × d .
From the above definition, we can define the convolution projected robust Wasserstein as follow:
Definition 11 For any p ≥ 1 , the convolution projected sliced Wasserstein (CPRW) of order p >between two given probability measures µ, ν ∈ P p
( R c × d × d ) is given by:
[Equation]
where S k ( ·| K (1) , . . .
, K ( N ) ) is a k -convolution slicer and K 1 , . . .
, K N are defined as in Definition 6.
We can obtain three instances of k -convolution slicers by modifying the number of channels from [Figure] to k in the convolution-base slicer,the convolution-stride slicer, and the convolution-dilation slicer.
As a result, we obtain three variants of CPRW which are CPRW-b, CPRW-s, and CPRW-d.
D Additional Experiments
In this section, we first present experiments on comparing probability measures over MNIST’s digits in Appendix D.1.
Then, we provide details of training generative models and additional experimental results in Appendix D.2.
D.1 Comparing Measures over MNIST’s digits
In the MNIST dataset, there are 60000 images of size 28 × 28 of digits from 0 to 9.
We compute SW between two empirical probability measures over images of every two digits, e.g., 1 and 2, 1 and 3, and so on.
To compare on the same digit, e.g, 1, we split images of the same digit into two disjoint sets and then compute the SW between the corresponding empirical probability measures [Table]
[Table] Meaningful measures of discrepancy: We approximate the SW and the CSW’s variants by a finite number of projections, namely, L = 1 , L = 10 , and L = 100 .
We show the mean of approximated values of L = 100 over 5 different runs and the corresponding standard deviation in
Table 4.
According to the table, we observe that SW and CSW’s variants can preserve discrepancy between digits well.
In particular, the discrepancies between probability measures of the same digit are relatively small compared to the discrepancies between probability measures of different digits.
Moreover, we see that the values of CSW-s and CSW-d are closed to the values of SW on the same pairs of digits.
We also show similar tables for L = 1 and L = 10 in Tables 2-3.
From these tables, we observe that the number of projections can affect the stability of both SW and CSW’s variants.
Furthermore, with a smaller value of L , the standard deviations of 5 different runs of both SW and
CSW’s variants are higher than values with L = 100 .
Projection memory for slicers: For SW, the conventional slicing requires L · 784 float variables for L projecting directions of 28 · 28 dimension.
On the other hand, CSW only needs L · 338 float variables since each projecting direction is represented as three kernels K
(1) ∈ R 15 × 15 , K (2) ∈ R 8 × 8 , and K (3) ∈
R 7 × 7 .
More importantly, CSW-s and CSW-d require only L · 57 float variables since they are represented by three kernels K (1) ∈
R 2 × 2 , K (
) ∈ R 2 × 2 , and K (3) ∈
R 7 × 7 .
From this experiment, we can see that using the whole unit-hypersphere as the space of projecting directions can be sub-optimal when dealing with images. [Figure] [Table]
D.2 Generative models
We parameterize the model distribution p ϕ ( x ) ∈ P ( R c × d × d ) and p ϕ ( x ) =
G ϕ ♯ϵ where ϵ is the standard multivariate Gaussian of 128 dimension and G ϕ is a neural network with Resnet architecture [19].
Since the ground truth metric between images is unknown, we need a discriminator as a type of ground metric learning.
We denote the discriminator as a function T β 2 ◦ T β 1 where
T β 1 : R c × d × d → R c ′ ×
d ′ × d ′ and T β 2 : R c ′ ×
d ′ × d ′ → R .
In greater detail, T β 1 maps the original images to their corresponding features maps and T β 2 maps their features maps to their corresponding discriminative scores.
Let the data distribution is µ , our training objectives are:
[Equation]
where m ≥ 1 is the mini-batch size and D ( · , · ) is the SW or CSW’s variants.
The above training procedure follows the papers [13, 44] that can be seen as an application of mini-batch optimal transport [14, 46, 45] with sliced Wasserstein kernels.
Training time and training memory: We report in Table 5 the training speed in the number of iterations per second and the training memory in megabytes (MBs).
We would like to recall that the time complexity and the projection memory complexity of CSW-s and CSW-d are the same.
Therefore, we measure the training time and the training memory of CSW-s as the result for both
CSW-s and CSW-d.
We can see that increasing the number of projections L costs more memory and [Figure] also slows down the training speed.
However, the rate of increasing memory of CSW is smaller than
SW.
For CSW-s and CSW-d, the extent of saving memory is even better.
As an example, L =in CSW-s and CSW-d costs less memory than SW with L = 100 while the performance is better (see Table 1).
In terms of training time, CSW-s and CSW-d are comparable to SW and they can be computed faster than CSW.
We refer the readers to Section 3 for a detailed discussion about the computational time and projection memory complexity of CSW’s variants.
Random generated images: We show some images that are drawn randomly from models trained by SW, CSW-b, CSW-s, and CSW-d on CIFAR10.
CelebA, STL10, and CelebA-HQ in Figure 5,
Figure 5, Figure 6, Figure 7, and Figure 8 in turn.
From these figures, we again observe the effect of changing the number of projections L , namely, a bigger value of L leads to better-generated images. [Figure]
Moreover, we observe that convolution sliced Wasserstein variants provide more realistic images than the conventional sliced Wasserstein.
These qualitative comparisons are consistent with the quantitative comparison via the FID scores and the IS scores in Table 1.
Results of Max Convolution sliced Wasserstein: We train generative models with Max-SW and
Max-CSW variants.
We search for the best learning rate in { 0 . 1 , 0 . 01 } and the number of update steps in {10,100}.
We report the best results on CIFAR10, CelebA, and CelebA-HQ for all models in
Table 6.
From this table, we observe that Max-CSW-s gives the best result on CIFAR10 and CelebA while Max-CSW-d is the best on CelebA-HQ.
This strengthens the claim that convolution slicers are better than conventional ones.
We also would like to recall that the computational time and memory of Max-CSW variants are better than Max-SW. [Table]
[Table] Results of Convolution projected sliced Wasserstein: As generalization of Max-SW and Max-CSW, we use PRW and CPRW-s with k ∈ { 2 , 4 , 16 } to train generative models.
We search for the best learning rate in { 0 . 1 , 0 . 01 } and the number of update steps in {10,100}.
The result on CIFAR is given in Table 6.
According to the table, CPRW-s is better than PRW with all choice of k which reinforces the favorable performance of convolution slicers.
Results of non-linear convolution sliced Wasserstein: We report FID scores and IS scores of generative models trained by non-linear sliced Wasserstein (NSW) [24] and non-linear convolution sliced Wasserstein (NCSW) variants including NCSW-b, NCSW-s, and NCSW-d on CIFAR10 in
Table 7.
The non-linear sliced Wasserstein is a variant of generalized sliced Wasserstein where we use a non-linear activation function after the linear projection, namely, g ( x, θ )
=
σ ( θ ⊤ x ) .
For NSW and NCSW variants, we choose σ () as the Sigmoid function.
Compared to linear versions in Table 1, we can see that including the non-linear activation function can improve the scores in some cases, e.g., NSW and NCSW-s. We also show FID scores and IS scores across training epochs in Figure 9.
Similar to the linear case, NCSW’s variants can help generative models converge faster than NSW.
E Experimental Settings
Architectures of neural network: We illustrate the detail of neural network architectures including the generative networks and the discriminative networks on CIFAR10 in Table 8, CelebA in Table 9,
STL10 in Table 10, and CelebA-HQ in Table 11.
Other settings: We set the number of training iterations to 50000 on CIFAR10, CelebA, and
CelebA-HQ and to 100000 on STL10.
For each 5 iterations, we update the generator G ϕ by the [Figure] [Table] Spectral normalization corresponding SW and CSW variants.
For the discriminator, we update T β 1 and T β 2 every iterations.
We set the mini-batch size m to 128 on CIFAR10 and CelebA, set to 32 on STL10, and set toon CelebA-HQ.
The learning rate of G ϕ , T β 1 , and T β 2 is set to 0 . 0002 .
We use Adam [23] for optimization problems with ( β 1 , β 2 ) =
Calculation of scores: For the FID scores and the Inception scores, we calculate them based on 50000 random samples from trained models.
For the FID scores, we calculate the statistics of datasets on all training samples.
[Table]
128 → 4 × 4 × 256 , dense
[Table]
[Table]
ResBlock up
ResBlock up
BN, ReLU,
[Table]
×
[Table]
conv
[Table]
[Table]
128 → 3 × 3 × 256 , dense , linear
ResBlock up
ResBlock up ResBlock up
ResBlock up ResBlock up
BN, ReLU, 3 × 3 conv, 3 Tanh
[Table]
ResBlock
[Table]
ResBlock
[Table]
[Table]
ϵ ∈ R
[Table]
∼ N
[Table]
) 128 → 4 × 4 × 256 , dense , linear
ResBlock up
ResBlock up
ResBlock up ResBlock up ResBlock up BN, ReLU, 3 × 3 conv, 3 Tanh
[Table]
ResBlock
[Table]
ResBlock
[Table]
[Table]
[Table]
Spectral normalization