Revisiting Sliced Wasserstein on Images: From Vectorization to Convolution
Abstract
The conventional sliced Wasserstein is deﬁned between two probability measures that have realizations as vectors .
When comparing two probability measures over images, practitioners ﬁrst 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 ﬁrst limitation is that the spatial structure of images is not captured efﬁciently by the vectorization step; therefore, the later slicing process becomes harder to gather the discrepancy information.
The second limitation is memory inefﬁciency 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 can approximate the value 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 deﬁned between two probability measures that have supports belonging to a vector space, e.g, R d .
As deﬁned 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 ﬁrst step, L projecting directions are ﬁrst 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-deﬁned 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 deﬁned 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 efﬁcient 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 deﬁned 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 beneﬁts of the convolution operators on images, including their efﬁcient parameter sharing and memory saving as well as their superior performance in several tasks on images [28, 19], to introduce efﬁcient 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 ﬁrst 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 := { ✓ 2 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 ﬁnite p -moments.
For µ, ⌫ 2 P p ( R d ) , ⇧ ( µ, ⌫ ) := { ⇡ 2 P p ( R d ⇥ R d ) | R R d ⇡ ( x, y ) dx = ⌫ , R
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 2 R dm , X :=
( x 1 , . . .
, x m ) , P X denotes the empirical measures 1 m[-(cid:80) mi =
-]
P mi =1  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 ﬁrst review the deﬁnitions 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 0  1 , we ﬁrst deﬁne the Wasserstein- p distance [59, 50] between two probability measures µ 2 P p ( R d 0 ) and ⌫ 2 P p ( R d 0 ) , which is given by W p ( µ, ⌫ ) :
= ⇣ inf
[Equation]
⌘ 1 p .
When d 0 = 1 , the Wasserstein distance has a closed form which is W p ( µ, ⌫ ) = (
R 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:
[Equation]
]⌫ ) d ✓ , where ✓] µ is the push-forward probability measure of µ through the function T ✓ : R d 0 !
R with T ✓ ( x ) = ✓ > x .
For each ✓ 2 S d 0  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 0  1 ) are drawn uniformly from the unit sphere and the approximation of the sliced Wasserstein distance is given by:
[Equation]
i
[Equation]
⌫ ) . In practice, L should be chosen to be sufﬁciently large compared to the dimension d 0 , 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 deﬁned 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 µ 2 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 efﬁciency of the whole process.
[Figure]
Convolution operator: We now deﬁne 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 deﬁnition of the convolution operator with stride and dilation is as follows.
Deﬁnition 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 2 R c ⇥ d ⇥ d with a kernel size K 2 R c ⇥ k ⇥ k is X s,b ⇤ K =
Y,
[Equation] where d 0 = d  b ( k  1)  1 s + 1 .
For i = 1 , . . .
, d 0 and j = 1 , . . .
, d 0 , Y 1 ,i,j is deﬁned as: Y 1 , i
[Equation]
From its deﬁnition, we can check that the computational complexity of the convolution operator is
[Equation]
Convolution Sliced Wasserstein
In this section, we will deﬁne a convolution slicer that maps a tensor to a scalar by convolution operators.
Moreover, we discuss the convolution slicer and some of its speciﬁc 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 ﬁrst start with the deﬁnition of the convolution slicer, which plays an important role in deﬁning convolution sliced Wasserstein.
Deﬁnition 2 (Convolution Slicer)
For N  1 , given a sequence of kernels K (1) 2 R c (1) ⇥ d (1) ⇥ d (1) , . .
, K ( N ) 2 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 ) )
2 R 8 X 2 R c ⇥ d ⇥ d .
As indicated in Deﬁnition 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 deﬁnition of convolution-dilation slicers to Deﬁnition 5.
We ﬁrst start with the deﬁnition of the convolution-base slicer.
Deﬁnition 3 (Convolution-base Slicer)
Given X 2 R c ⇥ d ⇥ d ( d  2 ), 1.
When d is even, N is the biggest integer that satisﬁes d = 2 N  1 a with a is also an integer, sliced kernels are deﬁned as K (1) 2 R c ⇥ (2  1 d +1) ⇥ (2  1 d +1) and K ( h ) 2 R 1 ⇥ (2  h d +1) ⇥ (2  h d +1) for h = 2 , . . .
, N  1 , and K ( N ) 2 R 1 ⇥ a ⇥ a where a = d 2 N  1 .
Then, the convolution-base slicer
CS -b ( X | K (1) , . .
, K ( N ) ) is deﬁned as:
[Equation]
When d is odd, the convolution-base slicer CS -b ( X | K (1) , . .
, K ( N ) ) takes the form:
[Equation]
where K (1) 2 R c ⇥ 2 ⇥ 2 and K (2) , . . .
, K ( N ) are the corresponding sliced kernels that are deﬁned on the dimension d  1 .
The idea of the convolution-base slicer in Deﬁnition 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 ﬁrst 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) 2 R c ⇥ d ⇥ d that satisﬁes the constraint
[Equation]
We now discuss the second variant of the convolution slicer, named convolution-stride slicer, where we further incorporate stride into the convolution operators.
Its deﬁnition is as follows.
Deﬁnition 4 (Convolution-stride Slicer)
Given X 2 R c ⇥ d ⇥ d ( d  2 ), 1.
When d is even, N is the biggest integer that satisﬁes d = 2 N  1 a with a is also an integer, sliced kernels are deﬁned as K (1) 2 R c ⇥ 2 ⇥ 2 and K ( h ) 2 R 1 ⇥ 2 ⇥ 2 for h = 2 , . .
. , N  1 , and
K ( N ) 2 R 1 ⇥ a ⇥ a where a = d 2 N  1 .
Then, the convolution-stride slicer CS -s ( X | K (1) , . .
, K ( N ) ) is deﬁned as:
[Equation]
When d is odd, the convolution-stride slicer CS
-s ( X | K (1) ,
[Equation]
. . , K ( N ) ) takes the form:
[Equation]
where K (1) 2 R c ⇥ 2 ⇥ 2 and K (2) , . . .
, K ( N ) are the corresponding sliced kernels that are deﬁned on the dimension d  1 .
Similar to the convolution-base slicer in Deﬁnition 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 beneﬁt 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).
Deﬁnition 5 (Convolution-dilation Slicer)
Given X 2 R c ⇥ d ⇥ d ( d  2 ), 1.
When d is even, N is the biggest integer that satisﬁes d = 2 N  1 a with a is also an integer, sliced kernels are deﬁned as K (1) 2 R c ⇥ 2 ⇥ 2 and K ( h ) 2 R 1 ⇥ 2 ⇥ 2 for h = 2 , . .
. , N  1 , and K ( N ) 2 R 1 ⇥ a ⇥ a where a = d 2 N  1 .
Then, the convolution-dilation slicer CS -d ( X | K (1) , . .
, K ( N ) ) is deﬁned as:
[Equation]
When d is odd, the convolution-dilation slicer CS -d ( X | K (1) , . .
, K ( N ) ) takes the form:
[Equation]
where K ( [Equation] ) 2 R c ⇥ 2 ⇥ 2 and K (2) , . .
, K ( N ) are the corresponding sliced kernels that are deﬁned 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 ﬁeld 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 satisﬁes 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 deﬁne a non-linear slicer in Deﬁnition 7 in Appendix C.
The non-linear slicer can be seen as a deﬁning function in generalized Radon Transform [52] which was used in generalized sliced Wasserstein [24].
3.2 Convolution Sliced Wasserstein
Given the deﬁnition of convolution slicers, we now state general deﬁnition of convolution sliced
Wasserstein.
An illustration of the convolution sliced Wasserstein is given in Figure 2.
{+Deﬁnition 
+} For any p  1 , the convolution sliced Wasserstein (CSW) of order p > 0 between two given probability measures µ, ⌫ 2 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 ) 2 R c ( l ) ⇥ k ( l ) ⇥ k ( l ) for any l 2 [ N ] and U ( K ( l ) ) is the uniform distribution with the realizations being in the set K ( l ) which is deﬁned as K ( l ) :
= n
K
[Equation]
o , 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 speciﬁcally 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]
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 ` 2 [ N ] and i 2 [ 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 [Equation]
with a complexity of O ( m log 2 m ) for each i 2 [ 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 ﬂatten 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 ﬁrst 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 satisﬁes 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 ﬁnd that CSW p ( µ, ⌫ )  Max-SW p ( µ, ⌫ )  W p ( µ, ⌫ ) , where
[Equation]
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
P 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 deﬁnition of the convolution sliced
Wasserstein in Deﬁnition 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 modiﬁcations 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 ﬁrst 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 usually
[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]
qualitative comparison between SW and CSW variants.
From the ﬁgures, 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 ﬁgure 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 deﬁned 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 efﬁcient 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 reﬂect 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 identiﬁable 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?
[-1) ∈ R 1 × 2 − 1 d +1 × 2 − 1 d +
-]
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[-R 2 × 2 , K (
-]
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]