Augmented Sliced Wasserstein Distances
Authors: Xiongjie Chen
Yongxin Yang
Yunpeng Li
Venue:  (2021)

Abstract: While theoretically appealing, the application of the Wasserstein distance to large-scale machine learning problems has been hampered by its prohibitive computational cost. The sliced Wasserstein distance and its variants improve the computational efficiency through random projection, yet they suffer from low projection efficiency because the majority of projections result in trivially small values. In this work, we propose a new family of distance metrics, called augmented sliced Wasserstein distances (ASWDs), constructed by first mapping samples to higher-dimensional hypersurfaces parameterized by neural networks. It is derived from a key observation that (random) linear projections of samples residing on these hypersurfaces would translate to much more flexible nonlinear projections in the original sample space, so they can capture complex structures of the data distribution. We show that the hypersurfaces can be optimized by gradient ascent efficiently. We provide the condition under which the ASWD is a valid metric and show that this can be obtained by an injective neural network architecture. Numerical results demonstrate that the ASWD significantly outperforms other Wasserstein variants for both synthetic and real-world problems.

Under review as a conference paper at ICLR 2021
AUGMENTED SLICED WASSERSTEIN DISTANCES
Anonymous authors
Paper under double-blind review
ABSTRACT
While theoretically appealing, the application of the Wasserstein distance to large-scale machine learning problems has been hampered by its prohibitive computational cost. The sliced Wasserstein distance and its variants improve the computational efficiency through random projection, yet they suffer from low projection efficiency because the majority of projections result in trivially small values. In this work, we propose a new family of distance metrics, called augmented sliced Wasserstein distances (ASWDs), constructed by first mapping samples parameterized by neural networks. It is derived from key observation that (random) linear projections of samples residing on these hypersurfaces would translate to much more flexible nonlinear projections in the original sample space, so they can capture complex structures of the data distribution. We show that the hypersurfaces can be optimized by gradient ascent efficiently. We provide the condition under which the ASWD is a valid metric and show that this can be obtained by an injective neural network architecture. Numerical results demonstrate that the ASWD significantly outperforms other Wasserstein variants for both synthetic and real-world problems.
nigher-dimensional hypersurfaces
INTRODUCTION
Comparing samples from two probability distributions is fundamental problem in statistics and machine learning. The optimal transport (OT) theory (Villani, 2008) provides powerful and flexible theoretical tool to compare degenerative distributions by accounting for the metric in the underlying spaces. The Wasserstein distance, which arises from the optimal transport theory, has become an increasingly popular choice in various machine learning domains ranging from generative models to transfer learning (Gulrajani et al., 2017: Arjovsky et al., 2017; Kolouri al., 2019b; Lee et al., 2019; Cuturi & Doucet, 2014; Claici et al., 2018; Courty et al., 2016; Shen et al., 2018; Patrini et al., 2018).
Despite its favorable properties, such as robustness disjoint supports and numerical stability (Arjovsky et al., 2017), the Wasserstein distance suffers from high computational complexity especially when the sample size is large. Besides, the Wasserstein distance itself is the result of an optimization problem - it is non-trivial to be integrated into an end-to-end training pipeline of deep neural networks, unless one can make the solver for the optimization problem differentiable. Recent advances in computational optimal transport methods foeus on altemative -T-da metrics that are computationally efficient and differentiably solvable (Peyré & Cuturi, 2019). Entropy regularization is introduced in the Sinkhorn distance (Cuturi, 2013) and its variants (Altschuler et al., 2017: Dessein et al., 2018; Lin al., 2019) to smooth the optimal transport problem; as result, iterative matrix scaling algorithms can be applied to provide significantly faster solutions with improved sample complexity (Genevay al., 2019).
An alternative approach is to approximate the Wasserstein distance through slicing, i.e. linearly projecting, the distributions to be compared. The sliced Wasserstein distance (SWD) (Bonneel et al.. 2015) is defined as the expected value of Wasserstein distances between one-dimensional random projections of high-dimensional distributions. The SWD shares similar theoretical properties with the Wasserstein distance (Bonnotte, 2013) and is computationally efficient since the Wasserstein distance in one-dimensional space has closed form solution based on sorting. Deshpande et al. (2019) extends the sliced Wasserstein distance to the max-sliced Wasserstein distance (max-SWD). by finding single projection direction with the maximal distance between projected samples. In Nguyen et al. (2020), the distributional sliced Wasserstein distance (DSWD) finds distribution of projections that maximizes the expected distances over these projections. The subspace robust Wasserstein distance extends the idea of slicing to projecting distributions linear subspaces (Paty & Cuturi, 2019). However, the linear nature

---

Under review as a conference paper at ICLR 2021
Projected from Projected sample from
from Sample from
ASWD
SWD
-90
-90
ASWD
60
60
Freqqeency
40
40
20
20
x
Distanca betwoen projections Distance between projections
(a)
(b)
(c)
(d)
Figure 1: (a) and (b) are visualizations of projections for the ASWI) and the SWD between two 2-dimensional Gaussians. (c) and (d) are distance histograms for the ASWD and the SWD between two 100-dimensional Gaussians. Figure I(a) shows that the injective neural network embedded in the ASWD learns data patterss (in the X-Y plane) and produces well-separate projected values (Z-axis) between distributions in random projection direction. The high projection efficiency of the ASWD is evident in Figure ((c), as almost all random projection directions in 100-dimensional space lead to significant distances between 1- 1-dimensional projections. In contrast, random linear mappings in the SWD often produce closer 1-d projections (Z-axis) (Figure 1(b)); as a result, a large percentage of random projection directions in the 100-d space result in trivially small distances (Figure I(d),, leading a low projection efficiency in nigh-dimensional spaces.
of these projections usually leads to low projection efficiency of the resulted metrics in nigh-dimensional spaces (Deshpande al., 2019; I Liutkus et al., 2019; Kolouri et al., 2019a).
More recently, there are growing interests and evidences that slice-based Wasserstein distances with improve the projection efficiency, leading to reduced number of projections needed to capture the structure of the data distribution (Kolouri et al., 2019a; Nguyen et al., 2020). (Kolouri et al., 2019a) extends the connection between the sliced Wasserstein distance and the Radon transform (Abraham et al.. 2017) to define generalized sliced Wasserstein distances (GSWDs) by utilizing generalized Radon transforms (GRTs). It is shown in (Kolouri et al., 2019a) that the GSWD is indeed metric if and only if the adopted GRT is injective. Injective GRTs are also used to extend the DSWD to the distributional generalized sliced Wasserstein distance (DGSWD) (Nguyen et al., 2020). However, both the GSWD and the DGSWDD are restricted by the limited class of injective GRTs, which utilize the circular functions and a finite number of harmonic polynomial functions with odd degrees as their defining function (Kuchment, 2006; Ehrenpreis, 2003). The results reported in (Kolouri et al., 2019a; Nguyen et al., 2020) show impressive performance from the GSWD and the DGSWD, yet they require one to specify particular form of defining function from the aforementioned limited class of candidates. However, the selection of defining function is usually task-dependent and needs domain knowledge. In addition, the impact on performance from different defining functions is still unclear.
nonlinear projections
One variant of the GSWD (Kolouri et al., 2019a) is the GSWD-NN, which generates projections directly with neural network outputs remove the limitations of slicing distributions with predefined GRTs. In the GSWD-NN, the number of projections, which equals the number of nodes in the neural network's s output layer, is fixed. Hence different neural networks are needed if one wants change the number of projections. There is also no random projections involved in the resulted GSWD-NN, since the projection results are determined by the neural network's 's weights Besides, the GSWD-NN is a pseudo-ment -etrie since it uses a vanilla neural network, rather than the Radon transform or GRTs, as its push-forward operator. Therefore, the GSWD-NN does not into the theoretical framework of the GSWD and does not inherit. its geometric properties.
In this paper, we present the augmented sliced Wasserstein distance (ASWD), a distance metric constructed by first mapping samples to hypersurfaces in augmented space, which enables flexible nonlinear slicing of data distributions for improved projection efficiency (See Figure 1). Our main contributions include: (i)
We exploit the capacity of nonlinear projections employed in the ASWD by constructing injective mapping with arbitrary neural networks; (ii) We prove that the ASWD a valid distance metric; (ii) We provide mechanism in which the hypersurface where high-dimensional distributions are projected onto be
2

---

Under review as a conference paper at ICLR 2021
optimized and show that the optimization of hypersurfaces help slice-based Wasserstein distances improve their projection efficiency. Hence, the ASWD is data-adaptive, i.c. the hypersurfaces can be learned from data. This implies one does not need specify defining function from a limited class of candidates; (iv) We demonstrate superior performance of the ASWD numerical experiments for both synthetic and real- worlw datasets.
The remainder of the paper is organized as follows. Section reviews the necessary background. We present the proposed method and its numerical implementation in Section 3. Numerical experiment results are presented and discussed in Section 4. We conclude the paper in Section 5.
2 BACKGROUND
In this section, we provide a brief review of concepts related to the proposed work, including the Wasserstein distance, (generalized) Radon transform and (generalized) sliced Wasserstein distances.
Wasserstein distance: Let (S2) be set of Borel probability measures with finite k-th moment on a Polish metric space (S2,d) (Villani, 2008). Given two probability measures μ., ME P&()), whose probability density functions (PDFs) are Py and pa. the Wasserstein distance order h
E (L,+ between μ and v/ defined as:
Wk(Lv)= ( inf d(x..y)* dry(,))))))))))))
(TEI(a.))
where df..-) is the cost function, I(H,2) represents the set of all transportation plans % i... joint distributions whose marginals are Pa and Pes respectively. With a slight abuse of notation, we interchangeably use (i, and Wa
While the Wasserstein distance generally intractable for high-dimensione distributions, there are several favorable cases where the optimal transport problem can be efficiently solved. In particular, if and are continuous one-dimensional measures, the Wasserstein distance between μı and v has a closed form solution (Bonneel et al., 2015):
Wh(kp)) Walkep)= ((, K
where 1-11 and F,1 are inverse cumulative distribution functions (CDFs) of and v. respectively.
Radon transform and generalized Radon transform: The Radon transform (Radon, 1917) maps function f(-) (RR)) to the space of functions defined over spaces of lines in Rd. The Radon transform of f()) defined by line integrals of f(-) along possible hyperplanes in Rdd
R.f(1 = f
where E and 0 @egd- represent the parameters of hyperplanes in Rd, 8(-) is the Dirac delta function, and (-r-) refers to the Euclidean inner product.
By replacing the inner product (x,,0) in Equation (3) with B(x,0), specific family of functions named as defining function in Kolouri et al. (2019a), the generalized Radon transform (GRT) (Beylkin, 1984) is defined as:
Gf(t.0)=
f(x)))) Rd
(4)
where E P while 22p a compact set of all feasible 0, e.g. Ω4 Sd- for ((x.0) (x.0) (Kolouri al., 2019a).
In practice, we can empirically approximate the Radon transform and the GRT of a probability density function pp, via:
N 8(t (Zx,)))
(t)(+
(5)
N
n=1
1
(t.))
5(t 0))
(6)
N
3

---

Under review as a conference paper at ICLR 2021
where Iy, Pp and N is the number of samples. Notably, the Radon transform is linear bijection (Helga- son, 1980), and the GRT is a bijection the defining function 8 satisfies certain conditions (Beylkin, 1984).
Sliced Wasserstein distance and generalized sliced Wasserstein distance: By applying the Radon transform to p/. and pu to obtain multiple projections, the sliced Wasserstein distance (SWD) decomposes the nigh-dimensional Wasserstein distance into multiple one-dimensional Wasserstein distances which can be efficiently evaluated (Bonneel et al., 2015). The k-SWD between and v is defined by:
SWDR(HLP= (f (Ryal (pp x
where the Radon transform R defined by Equation (3) is adopted as the measure push-forward operator. The GSWD generalizes the idea of SWD by projecting distributions onto hypersurfaces rather than hyperplanes (Kolouri et al., 2019a). The GSWD is defined as:
GSWD = - W (CP.(.
where the GRT G is used as the measure push-forward operator. The Wasserstein distances between one-dimensional distributions can be obtained by sorting projected samples and calculating the distance between sorted samples (Kolouri et al., 2019b): with L random projections, the SWD and GSWD between f and can be approximated by:
I. N
SWDR(up) R3
T
(9)
(S)D)
(10)
where 1- and 1 are sequences consist of the indices of sorted samples which satisfy (Frg(n)0) O1) S (1rgb++1)) 01). (Wrgm)) S (111)++)) in the SWD, and (Zrr)1).0) S ((zrgp+11):0)) Blrry 0r) S B(u1p++1): (t) in the GSWD. It is proved in Bonnotte (2013) that the SWD is valid distance metric. The GSWD is valid metric except for its neural network variant (Kolouri et al., 2019a).
AUGMENTED SLICED WASSERSTEIN DISTANCES
In this section, we propose a new distance metric called the augmented sliced Wasserstein distance (ASWD). which embeds flexible nonlinear projections in its construction. We also provide an implementation recipe for the ASWD.
3.1 SPATIAL RADON TRANSFORM AND AUGMENTED SLICED WASSERSTEIN DISTANCE
the definitions of the SWD and GSWD, the Radon transform (Radon, 1917) and the generalized Radon transform (GRT) (Beylkin, 1984) are used as the furw-forward operator for projecting distributions to a one-dimensional space. However, is not straightforward to design defining functions (xx,0) (Kolouri al., 2019a) for the GRT due to certain non-trivial requirements for the function (Beylkin, 1984). In practice, the assumption the transform can be relaxed, as Theorem shows that as long as the transform injective, the corresponding ASWD metric is valid distance metric.
To help us define the augmented sliced Wasserstein distance, we first introduce the spatial Radon transform which includes the vanilla Radon transform and the polynomial GRT special cases (See Remark 2).
Definition 1. Given an injective mapping 9(:): :dd - Re and probability measure μ P(I))) which probability density function (PDF) is Py, the spatial Radon transform of Pp is defined as
Hpp = (
(11)
where E R and O E Sdo--1 the parameters hypersurfaces in TRd.
4

---

Under review as a conference paper at ICLR 2021
Note that the spatial Radon transform can be interpreted applying the vanilla Radon transform to the PDF of = where x Pmp Denote the PDF of & by Päg the spatial Radon transform defined by Equation (11) can be rewritten as:
Remark 1.
(t.(1.0:g)- FA 16(1 (g))
=Eappo (6(t-(2,01)
J Pa. (0))((t (i,
(Rpa,(t,0)
(12)
Hence the spatial Radon transform inherits the theoretical properties of the Radon transform subject to certain conditions of g()) and incorporates nonlinear projections through g(-)-
In what follows, we use fi = f2 to denote functions fi() X R and : that satisfy fi(x)= f2(x) for almost V.
Lemma 1. Given an injective mapping g(-): Rd + Rd and two probability measures μ.v P(Rd) whose probability density functions are Pp and Pu. respectively, for all E R and 0 E gdo-1, (pp(t, = (t.(t. (t,0:g) if and only if Pa =u, ie. the spatial Radon transform is injective. Moreover, the spatial Radon transform injective if and only if the mapping g(-) is an injection.
See Appendix A for the proof of Lemma 1.
Remark 2. The spatial Radon transform degenerates the vanilla Radon transform when the mapping g(.) is an identity mapping. When g(-) is homogeneous polynomial function with odd degrees, the spatial Radon trunsform is equivalent to the polynomial GRT (Ehrenpreis, 2003).
Appendix B provides the proof of Remark 2.
now introduce the augmented sliced Wasserstein distance, by utilizing the spatial Radon transform the measure push-forward operator:
Definition 2. Given two probability measures Jr, E (R.)),) whose probability density functions are and pv. respectively, and an injective mapping g(-) :RR - TRdo, the augmented sliced Wasserstein distance (ASWD) order k (1,+ is defined as:
(/ooari))-/ggoa W (H (Hp,(. (po())))
ASWDE(
(13)
where 0 S We is the k- t-Wasserstein distance defined by Equation (1), and H refers to the spatial Radon transform defined by Equation (11).
Remark 3. Following the connection between the spatial Radon transform and the vanilla Radon transform as shown Equation ( (12), the ASWD can be rewritten as:
ASWD((WRr)) (as (RPp, (.00))
(14)
where By and by are probability measures on Redo which satisfy g(x)) ~ Pig for x and gly) ~ Dg for yov.
Theorem 1. The augmented sliced Wasserstein distance (ASWD) of order k E [1, +oc) defined by
Equation (13) with a mapping g(.):Rdd RRdo a metric on Px(TR. (RR:) and only if g(:) is injective.
Proof of Theorem I is provided in Appendix C.
3.2 NUMERICAL IMPLEMENTATION
We discuss in this section how to realize injective mapping g(-) with neural networks due to their expressiveness and optimize with gradient based methods.
Injective neural networks: As stated in Lemma and Theorem 1, the injectivity of g(-) is the sufficient and necessary condition for the ASWD being valid metric. Thus we need specific architecture designs
5

---

Under review as a conference paper at ICLR 2021
implementing g()) by neural networks. One option is the family of invertible neural networks (Behrmann et al., 2019; Karami et al., 2019; Song et al., 2019), which are both injective and surjective. However, the running cost of those models is usually much higher than that of vanilla neural networks. We propose an alternative approach by concatenating the input & of an arbitrary neural network to its output (x):
In (x) - (x)).
(15)
It is trivial to show that (r) is injective, since different inputs will lead to different outputs. Although embarrassingly simple, this idea of concatenating the input and output of neural networks has found success in preserving information with dense blocks in the DenseNet (Huang et al., 2017), where the input of each layer is injective to the output of all preceding layers.
Optimization objectíve: We aim to project samples to maximally discriminating hypersurfaces between two distributions, so that the projected samples between distributions are most dissimilar subject to certain constraints on the hypersurface, as shown in Figure 1. Similar ideas have been employed to identify important projection directions (Deshpande et al., 2019; Kolouri et al., 2019a; Paty & Cuturi, 2019) or discriminative ground metric (Salimans et al., 2018) in optimal transport metrics. For the ASWD, the parameterized injective neural network (-) is optimized by maximizing the following objective:
(
Clu.ggw.)))=
(16)
where À >> and the regularization term LA = AEx [[llg (z) (19m used to control the norm of the output of gu (.), otherwise the projections may be arbitrarily large.
Remark 4. The regularization coefficient X adjusts the introduced non-linearity in the evaluation of the ASWD by controlling the norm of (u(-)) in Equation (15). In particular, when X A cc, the nonlinear term ow (.) shrinks to O. The rank of the augmented space is hence explicitly controlled by the flexible choice of (-) and implicitly regularized by Lx.
By plugging the optimized -h, - argmaaxggxxx44 (( into Equation (13), we obtain the empirical
9- version of the ASWD. Notably, the regularization term Lx is only used when maximizing the objective in Equation (16), once the optimization completed, L, is used the calculation of the ASWD defined by Equation (13). Pseudocode is provided in Appendix D.
4 EXPERIMENTS
In this section, we describe the experiments that we have conducted evaluate performance of the proposed distance metric. The GSWD leads to the best performance in a sliced Wasserstein flow problem reported in Kolouri et al. (2019a) and the DSWD outperforms the compared methods in the generative modeling task examined in Nguyen et al. (2020). Hence we compare performance of the ASWD with the state-of-the-art distance metrics in the same examples and report results as below
To examine the robustness of the ASWD, throughout the experiments, we adopt the injective network architecture given in Equation (15) and set ow to be a one layer fully-connected neural network whose outputs' dimension equals its inputs' dimension, with a ReLU layer as its activation function.
4.1 SLICED WASSERSTHIN FLOWS
We first consider the problem of evolving a source distribution μ target distribution by minimizing Wasserstein distances between and v in the sliced Wasserstein flow task reported in Kolouri et al. (2019a).
0,4== = VSWD(H....)))
(17)
where br refers to the updated source distribution at each iteration t. The SWD in Equation (17) can be replaced by other sliced- Wasserstein distances to evaluated. As Kolouri et al. (2019a), the 2- Wasserstein distance was used as the metric for evaluating performance of different distance metrics this task. The set of hyperparameter values used in this experiment can be found in Appendix E.
Code reproduce experiment results is available httph:ttt///...... //bitt//tttt./ 1y/211222222222

---

Under review as a conference paper at ICLR 2021
ASWD SWD
GSWD-Poly
DSWD GSWD- NN
GSWD-Circular
Gaussian
Knot
0.0
0.5
-5
500
1000
1500
1000
1500
2000
Iteration
Iteration
Moon
Swiss Roll
1.0
0.4
log
0.5
(0
0.2
22
0.0
0.5
-0.
500
1000 Iteration
1500
2000
0
500
1000
1500
2000
Iteration
Figure 2: The first and third columns are target distributions. The second and fourth columns are log 2- 2-Wasserstein distances between the target distribution and the source distribution. The horizontal axis show the number of training iterations. Solid lines and shaded areas represent the average values and 95% confidence intervals of log 2-sasserstein distances over 50 runs. A more extensive of experimental results can be found in Appendix F.I.
Without loss of generality, we initialize μi be the standard normal distribution (0.2). We repeat each experiment 50 times and record the 2-Wasserstein distance between and // at every iteration. In Figure we plot the 2- Wasserstein distances between the source and target distributions as a function of the training epochs and the 8- 8-Gaussian, the Knot, the Moon, and the Swiss roll distributions are respective target distributions. For clarity, Figure displays the experiment results from the best performing distance metrics, including the ASWD, the DSWD, the SWD, the GSWD-NN 1, which directly generates projections through one layer MI., as well as the GSWD with the polynomial of degree 3, circular defining functions, out of the 12 distance metrics compared.
We observe from Figure that the ASWD not only leads to lower 2- 2-Wasserstein distances, but also converges faster by achieving better results with fewer iterations than the other methods in these four target distributions. A complete set of experimental results with 12 compared distance metrics and target distributions are included in Appendix F.I. The ASWD outperforms the compared state-of-the-ar sliced-based Wasserstein distance metrics with out of the 8 target distributions except for the 25-Gaussian. This is achieved through the simple injective network architecture given in Equation (15) and layer fully-connected neural network with equal input and output dimensions throughout the experiments. In addition, ablation study is conducted to study the effect of injective mappings, the regularization coefficient X, and the optimization of hypersurfaces in the ASWD. Details can be found in Appendix F.2.
4.2 GENERATIVE MODELING
In this experiment, we the sliced-based Wasserstein distances for generative modeling task described in Nguyen et al. (2020). The task is to generate images using generative adversarial networks (GANs) (Good- fellow et al., 2014) trained on either the CIFARIO dataset (64x64 64 resolution) (Krizhevsky, 2009) the CELEBA dataset (64x64 resolution) (Liu et al., 2015), Denote the hidden layer and the output layer of the discriminator by he and Dv, and the generator by Gæ, we train GAN models with the following objectives:
min (p.
(18)
max Eg,p, (Iog(D((h + (logs 1 (hy V.6
(19)

---

Under review as a conference paper at ICLR 2021
Table 1: FII scores of generative models trained with different distance metrics. Lower scores indicate better image qualities. L is the number of projections, we run each experiment 10 times and report the average values and standard errors of FID scores for CIFARIO dataset and CELEBA dataset. The running time per training iteration for one batch containing 512 samples is computed based on a computer with an Intel (R)
Xeon (R) Gold 5218 CPU 2.3 GHz and 16GB of RAM, and RTX 6000 graphic card with 22GB memories.
SWD
GSWD
DSWD
ASWD
FID
I (s/it)
FID
(s/it)
FID
I (s/it)
FID
(s/it)
CIFARTO
10
192.6=t5.7
0.32
189.5:6.00
0.35
79.0+4.2
0.48
73.2:13.1
0.55
100
155.0+2.99
0.32
155.9=3.2
0.70
72.2+8.2
0.51
66.7+3322
0.57
1000
126.0=2.9
0.34
134.5#2.7
2.10
74.3+4.3
1.22
65.5+3.9
1.32
CELEBA
10
118.3+3.1
0.32
143.24:5.5
0.35
105.3+3.4
0.49
99.2+:433
0.53
100
116.0+2.88
0.33
1
0.69
103.1
0,51
94.3:22.2
0.56
1000
104.442.88
0.34
101.8
2.14
97.4
1.21
90.5
1.31
350
310
I
ASWD
ASWD
DSWD
300
DSWD
GSWD
GSWD
FID score
SWD
FID score
250
SWD
210
200
160
150
110
100
60
50
150
200
60
80
100
Epoch
Epoch
(a) FID scores CIFARIO (L= 1000)
(b) FID scores CELEBA (L= (L=1000)
Figure 3: FID scores of generative models trained with different metrics on CIFARIO and CELEBA datasets with L= 1000 projections. The error bar represents the standard deviation of the FID scores at the specified training epoch among 10 simulation runs.
where p= is the prior of latent variable and p, is the distribution of real data. The SWD in Equation (18) replaced by the ASWD and other variants of the SWD to compare their performance. The GSWD with the polynomial defining function and the DGSWD not included in this experiment due its excessively high computational cost in high-ddimnnsional space. The Fréchet Inception Distance (FID score) (Heusel al., 2017) is used to assess the quality of generated images. More details on the network structures and the parameter setup used in this experiment are available in Appendix E.2.
We run 200 and 100 training epochs train the GAN models on the CIFARIO and the CELEBA dataset, respectively. Each experiment is repeated for 10 times. We report experimental results in Table 1. With the same number of projections and a similar computation cost, the ASWD leads significantly improved FID scores among all evaluated distances metrics both datasets, which implies that images generated with the ASWD are higher qualities. Figure 3 plots the FID scores recorded during the training process. The GAN model trained with the ASWD exhibits faster convergence as it reaches lower FID scores with fewer epochs. Randomly selected samples of generated images are presented in Appendix
5 CONCLUSION
We proposed a novel variant of the sliced Wasserstein distance, namely the augmented sliced Wasserstein distance (ASWD), which is flexible, has high projection efficiency, and generalizes well. The ASWD adaptively updates the hypersurface where the samples are projected onto by learning from data. We proved that the ASWD valid distance metric and presented its numerical implementation. We
8

---

Under review as a conference paper at ICLR 2021
reported empirical performance of the ASWD over otate-of-the-ar sliced Wasserstein metries in numerical experiments. The ASWD leads to the smallest distance errors over the majority of datasets in sliced Wasserstein flow task and superior performance in generative modeling task.
The ASWD can be extended several directions. One future research topic is to incorporate the framework of the ASWD into the sliced Gromov- Wasserstein distance, such that data-adaptive, nonlinear projections can be learned to compare distributions whose supports do not necessary lie in the same metric space. would be also interesting to explore the identification of the barycenter of objects via the ASWD, e-g. through an optimization scheme alternating between optimization of the nonlinear maps in the ASWD framework and the barycenter. Another interesting area is the incorporation of the ASWD with state-of-the art generative models to further exploit projection efficiency. We anticipate that they will provide promising directions for enabling the generalization for sliced Wasserstein distance be applied wider domains.

---

Under review as a conference paper at ICLR 2021
REFERENCES
I. Abraham, R. Abraham, M. Bergounioux, and G. Carlier. Tomographic reconstruction from a few views: multi-marginal optimal transport approach. Applied Mathematics & Optimization, 75(1):55-73, 2017.
J. Altschuler, J. Kiles-Weed, and P. Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhom iteration. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pp. 1964-1974, Long Beach, California, USA, 2017.
M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In Proc. International Conference on Machine Learning (ICML), pp. 214-223, Sydney, Australia, 2017.
J. Behrmann, W. Grathwohl, R. T. Q. Chen, D. Duvenaud, and J. Jacobsen. Invertible residual networks. In Proc. International Conference on Machine Learning (ICML), pp. 573-582, Long Beach, Califomia, USA, 2019.
G. Beylkin. The inversion problem and applications of the generalized Radon transform. Communications on Pure and Applied Mathematics. 37(5):579-50) 5999,99999990 1984.
N. Bonneel, J. Rabin, G. Peyré, and H. Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51(1):22-45, 2015.
N. Bonnotte. Unidimensional and evolution methods for optimal transportation. PhD thesis, Paris 11, 2013.
S. Claici, E. Chien, and J. Solomon. Stochastic Wasserstein barycenters. In Proc. International Conference on Machine Learning (ICML). pp. 999-1008, Stockholm, Sweden, 2018.
N. Courty, R. Flamary, D. Tuia, and A. Rakotomamoniy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence (IPAMI), 185):18533 2016.
M. Cuturi. Sinkhom distances: Lightspeed computation of optimal transport. In Proc. Advances in Neural Information Processing Systems (NewrlPS), pp. 2292-2300, I Lake Tahoe, Nevada, USA, 2013.
M. Cuturi and A. Doucet. Fast computation of Wasserstein barycenters. In Proc. International Conference Machine Learning (ICML), pp. 685-693,-- Beijing, China, 2014.
I. Deshpande et al. Max-sliced Wasserstein distance and its use for GANs. In Proc. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10648-10656, Long Beach, California, USA, 2019.
A. Dessein, N. Papadakis, and J. Rouas. Regularized optimal transport and the rot mover's distance. The Journal of Machine Learning Research (JMLR), 19(1):90-6422 642, 2018.
L. Ehrenpreis. The universality of the Radon transform, chapter 5, pp. 299 363. Oxford University Press, Oxford, UK, 2003.
A. Genevay, L. Chizat, F. Bach, M. Cuturi, and G. Peyré. Sample complexity of Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 1574 1583, Okinawa, Japan, 2019.
I. Goodfellow et al. Generative adversarial nets. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pp. 2672-2680, Montréal, Canada, 2014.
I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. Courville. Improved training of Wasserstein GANs. In Proc. Advances in neural information processing systems (NeurIPS), pp. 5767-5777, Long Beach, Califomia, USA, 2017.
Helgason. The Radon transform, volume 2. Basel, Switzerland: Springer, 1980.
M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. GANs trained by two time-scale update rule converge to a local Nash equilibrium. In Proc. Advances in neural information processing systems (NeurIPS), pp. 6626-6633 Long Beach, California, USA, 2017.
G. Huang, Z. Liu, L. Van Der Maaten, and K. Q. Weinberger. Densely connected convolutional networks. In Proc. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 4700-4708, Hawaii, USA, 2017.
10

---

Under review as a conference paper at ICLR 2021
M. Karami, D. Schuurmans, J. Sohl-Dickstein L. Dinh, and D. Duckworth. Invertible convolutional flow. In Proc. Advances in Neural Information Processing Systems (NeurlPS), pp. 5636-5636660 5646, Vancouver, Canada, 2019.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In Proc. International Conference on Learning Representations (ICLR), San Diego, California, USA, 2015.
S. Kolouri, K. Nadjahi, U. Simsekli, R. Badeau, and G. Rohde. Generalized sliced Wasserstein distances. In Proc. Advances Neural Information Processing Systems (NeurIPS), pp. 261-272, Vancouver, Canada, 2019a.
S. Kolouri, P. E Pope, C. E. Martin, and G. K. Rohde. Bliced-Wasserse autoencoders. In Proc. International Conference on Learning Representations (ICLR), New Orleans, Louisiana, USA, 2019b.
A. Krizhevsky. Learning multiple layers of features from tiny images. Tech Report, 2009.
P. Kuchment, Generalized transforms of Radon type and their applications. In Proc. Symposia in Applied Mathematics, volume 63, pp. 67, San Antonio, Texas, USA, 2006.
C. Lee, T. Batra, M. B. Baig, and D. Ulbricht. Sliced Wasserstein discrepancy for unsupervised domain adaptation. In Proc. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp- 10285-10295 Long Beach, Califomia, USA, 2019.
T. Lin, N. Ho, and M. I. Jordan. On efficient optimal transport: an analysis of greedy and accelerated mirror descent algorithms. University of Califormia, Berkeley, 2019.
Liu, P. Luo, Wang, and Tang. Deep learning face attributes in the wild. Proc. IEEE International Conference on Computer Vision (ICCV), pp. 3730-3738, Las Condes, Chile, 2015.
A. Liutkus, U. Şimşekli, S. Majewski, Szymon, Alain D., and F. Stöter. Sliced Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. Proc. International Conference on Machine Learning (ICML), pp. 4104-4113, 2019.
K. Nguyen, N. Ho, T. Pham, and H. Bui. Distributional sliced- Wasserstein and applications to generative modeling. arXiv:2002.0736: 2020.
G. Patrini et al. Sinkhorn autoencoders. In Proc. Conference on Uncertainty in Artificial Intelligence (UAI), Tel Aviv, Israel, 2018.
Paty and M. Cuturi. Subspace robust Wasserstein distances. In Proc. International Conference on Machine Learning (ICML), pp. 5072-5081, Long Beach, California, USA, 2019.
G. Peyré and M. Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 115-66)355--07, 2019.
J. Radon. Uber die bestimmung von funktionen durch ihre integralwerte laengs gewisser manniffaltigkeiten. Ber. Verh. Saechs. Akad. Wiss. Leipzig Math. Phys. Kl., 69:262, 1917.
D. Rezende and S. Mohamed. Variational inference with normalizing flows. International Conference on Machine Learning (ICML), pp. 1530-1538, Lille, France, 2015.
T. Salimans, H. Zhang, A. Radford, and D. Metaxas. Improving GANs using optimal transport. In Proc. International Conference on Learning Representations (ICLR), Vancouver, Canada, 2018.
J. Shen, Y Qu, W Zhang, and Y Yu. Wasserstein distance guided representation learning for domain adaptation. Proc. AAAI Conference on Artificial Intelligence (AAAI), pp. 0058-4065 New Orleans, Louisiana, USA, 2018
Y. Song, C. Meng, and S. Ermon. Mintnet: Building invertible neural networks with masked convolutions. In Proc. Advances in Neural Information Processing Systems, pp. 11002-11012, Vancouver, Canada, 2019.
C. Villani. Optimal Transport: old and new, volume 338. Berlin, Germany: Springer Science & Business Media, 2008
11


Reviews:
[{"review_id": "qJV5jVTiOAr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Augmented Sliced Wasserstein distance between random vectors X&Y is defined as SWD between g(X) and g(Y) for an injective g. \n\nIt is claimed that this generalization solves the problem with SWD, i.e. a small distance for most of the projection directions.  Since the paper is purely experimental, then it is better to demonstrate that this issue is resolved (it is not obvious why the problem disappears).\n\nSince g_omega is also searched for by maximizing (16), then it becomes hard to interpret the value of metrics (also taking into account the role of lambda). What it measures now?\n\nThe code is given, results seem reproducible. Reported results with generative modeling compare ASWD with SWD, GSWD, DSWD. Since the paper is experimental, maybe it is natural to expect a comparison with SOTA generators that do not deal with generalizations of Wasserstein distance. Also, fake images generated by SWD, GSWD, DSWD are not given (only FID is shown).", "extra_scores": {}, "extra_text": {}}, {"review_id": "VHQZrxqjyLe", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper provides a notion of generalisation for sliced Wasserstein distances, that allows to explore nonlinear projections in arbitrary subspaces in a suitable and efficient way. The paper is well-written and provides a novel solution to the problem of exploiting nonlinear subspaces in computing distances. I feel this idea, although simple can be quite powerful, and can be extended to wide domains especially identifying objects based on arbitrary feature selections although I do not think the authors have explored that direction in this piece of work. \nComments:\n(1) Numerical results show the distance computed by ASWD is smaller than other measures. This can, however, be misleading in the sense that it might also be obtained by insufficient exploration of non-linear subspaces. For eg, if the sliced Wasserstein distances are computed along orthogonal directions to the primary features the Wasserstein distance obtained in that regard will also be very small, although this does not in anyway validate the superiority of the distance metric. I feel this argument therefore needs some clarification.\n(2) Can I use this method to construct arbitrary nonlinear projections of choice instead of the projection that gives the best distinction? For eg, suppose two pictures have several objects among which cars in the two pictures are the most distinguishing features. However, I also want to see what other features can help distinguish these two pictures from other pictures which do not have cars. How do I do that? In that broader sense I suppose my question deals with trying to identify barycenters of objects via sliced Wasserstein distances. Any thoughts in that regard could be useful for the reader.\n\nOverall I find the paper an interesting read although it requires some clarifications as outlined above.", "extra_scores": {}, "extra_text": {}}, {"review_id": "8nsKwlPIZNi", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces the augmented sliced Wasserstein distances (ASWD) to capture non-linear projections, as opposed to sliced Wasserstein distances. To this end, the architecture first maps inputs to higher dimensional hypersurfaces, before projecting. The designed ASWD is shown to be a metric as long as the initial mapping is injective. \n\nNumerically, the initial mapping is parameterized as the concatenation of the input and its output using a neural network. The loss aims at capturing a hypersurface that differentiates measures the most, and is regularized to control the initial mapping. The first task considers the minimization of sliced Wasserstein flows for different settings of synthetic data. The other task consists of generating images with GANs, using a sliced Wasserstein loss to learn the generator network.\n\nStrong points of the paper include:\n1.\tThe paper is well written. The literature review is very thorough and comparison with the proposed method is well explained.\n2.\tBoth theoretically and numerically, the method is well described, well compared to existing notions and yields convincing results.\n\nDrawbacks / questions:\n1.\tIt seems that a lot of projections are needed to retrieve visually satisfactory generated examples.\n2.\tHave you tried other types of injective maps than Eq 15?\n3.\tCan you comment on the impact of regularization strength lambda in practice?\n\nI recommend an accept for this paper, which, to the best of my knowledge, brings both theoretical and numerical valuable contributions to the literature.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Given two data measures in R^d, this paper proposes to use a NN to augment the representation of each data point found in these measures with additional coordinates. The measures are then compared using the sliced Wasserstein distance on these augmented representations. Because this augmentation is injective by design (the original vectors are part of the new representation) simple metric properties are kept. The authors propose to learn in a robust/adversarial way these augmentations. They propose simple experiments to illustrate that idea.

Although I found the idea interesting, I think it falls short of acceptance at ICLR. I agree with the sentiment of other reviewers 1 and 2 that defining another variant of robust/NN inspired variant of the W distance is interesting, but at this point the readership of the conference expects more than simple experiments on toy data and hard to interpret GAN results. I think there is value in the draft as it stands now, but that more efforts are needed to convince this variant is scalable / useful for other downstream tasks (e.g. W barycenters, or other easier to interpret W problems in lower dimensions).

minor comments
- as it stands, equation 2 is wrong if you do not add more conditions on the cost function d(.,.). 
- " the idea of SWD by projecting distributions onto hypersurfaces rather than hyperplane" -> this is wrong, the projection is done onto lines or curves, not hyperplanes or hypersurfaces.

Decision: Reject