Augmented Sliced Wasserstein Distances
Authors: Xiongjie Chen
Yongxin Yang
Yunpeng Li
Venue: NeurIPS 2021 Submitted (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 the random projection, yet they suffer from low accuracy if the number of projections is not sufficiently large, 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.

25 Despite its favorable properties, such robustness disjoint supports and numerical stability [Arjovsky et al., 20 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 Augmented Sliced Wasserstein Distances
Anonymous Author(s) Affiliation Address email
Abstract
While theoretically appealing. the application of the Wasserstein distance 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 the random projection, yet they suffer from low accuracy if the number of projections is not sufficiently large, because the majority of projections result in trivially small values. In this work, we propose new family of distance metrics, called augmented sliced Wasserstein distances (ASWDs), constructed by first mapping samples to terized by neural networks. 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 valid metric and show that this can be obtained by injective neural network architecture. Numerical results demonstrate that the ASWD significantly outperforms other Wasserstein variants for both synthetic and real-world problems.
hypersurfaces parame-
11
13
15
Introduction
17
18
Comparing samples from two probability distributions is a fundamental problem in statistics and machine learning. The optimal transport (OT) theory [Villani, 2008] provides a powerful and flexible theoretical tool to compare degenerative distributions by accounting for the metric in the underlying spaces. The Wasserstein distance, which arises from the optimal transport theory, has become an increasingly popular choice in various machine learning domains ranging from generative models to transfer learning [Gulrajani et al., 2017; Arjovsky et al., 2017; Kolouri et al., 201 9b; Cuturi and Doucet, 2014; Courty et al., 2016].
19
20
21
22
23
24
26
27
it is non-trivial to be integrated into an end-to-end training pipeline of deep neural networks, unless 28
29 one can make the solver for the optimization problem differentiable. Recent advances computational
30 optimal transport methods focus on alternative OT basdd metrics that are computationally efficient and
31 solvable via a differentiable optimizer [Peyré and Cuturi, 2019]. Entropy regularization is introduced in the Sinkhorn distance [Cuturi, 2013) and its variants [ [Altschuler et al., 2017; Dessein et al., 201 8]
33 to smooth the optimal transport problem; as result, iterative matrix scaling algorithms can be applied
34 to provide significantly faster solutions with improved sample complexity [Genevay et al., 2019].
35 An alternative approach is 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 the expected value of Wasserstein distances between mne-iimensional random projections of high-dimensional distributions. The SWD shares similar theoretical properties with the
36
37
38
Submitted to 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Do not distribute.

---

Projected sample from Projected sample from
Sample from
Sample from
ASWD
SWD
80
880
ASWD
SWD
60
Frequency
Frequency
20
20
x
40
Distance between 1-d projections Distance between 1-d projections
(a)
(b)
(c)
(d)
Figure 1: (a) and (b) are visualizations of projections for the ASWD 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 ((a) shows that the injective neural network emhedded in the ASWD learns data patterns (in the X-Y plane) and produces well sellasee projected values (Z-axis) between distributions in a random projection direction. The high projection efficiency of the ASWD evident in Figure (c)), as almost all random projection directions in a 100-dimensional space lead to significant distances between I-dimensional projections. In contrast, random linear mappings in the SWD often produce closer 1-d projections (Z-axis) (Figure 1(b)); as result, a large percentage of random projection directions in the 100-d space result in trivially small distances (Figure 1(d)), leading to low projection efficiency in high- diteemeinnii spaces.
39 Wasserstein distance [Bonnotte, 2013] and is computationally efficient since the Wasserstein distance
40 in one-dimensional space has a closed-form solution based sorting. [Deshpande et al., 2019] extends the sliced Wasserstein distance to the max-sliced Wasserstein distance (Max-SWD). by finding a single projection direction with the maximal distance between projected samples. The subspace robust Wasserstein distance extends the idea of slicing to projecting distributions on linear subspaces [Paty and Cuturi. 2019]. However, the linear nature of these projections usually leads low projection efficiency of the resulted metrics in high- nigh-dimensional spaces [Deshpande et al., 2019; Kolouri et al., 2019a].
41
42
43
44
45
Different variants of the SWD have been proposed improve the projection efficiency of the SWD, either by introducing nonlinear projections or by optimizing the distribution of random projections. Specifically, [Kolouri et al., 2019a] extends the connection between the sliced Wasserstein distance and the Radon transform [Radon, 1917] to introduce generalized sliced Wasserstein distances (GSWDs) by utilizing generalized Radon transforms (GRTs), which are defined by nonlinear defining functions and lead to nonlinear projections. A variant named the GSWD-NN was proposed in [Kolouri et al., 2019a] to generate nonlinear projections directly with neural network outputs, but it does not fit into the theoretical framework of the GSWD and does not guarantee a valid metric. In contrast, the distributional sliced Wasserstein distance (DSWD) and its nonlinear version, the distributional generalized sliced Wasserstein distance (DGSWD), improve their projection efficiency by finding a distribution of projections that maximizes the expected distances over these projections. The GSWD and the DGSWD exhibit higher projection efficiency than the SWD in the experiment evaluation, yet they require the specification of the particular form of defining functions from a limited class of candidates. However, the selection of defining functions is usually task-dependent problem and requires domain knowledge, and the impact on performance from different defining functions is still unclear.
47
49
50
51
52
53
54
55
56
57
58
60
61
In this paper, we present the augmented sliced Wasserstein distance (ASWD). a distance metric constructed by first mapping samples hypersurfaces in an augmented space, which enables flexible nonlinear slicing of data distributions for improved projection efficiency (See Figure 1). Our main contributions include: (i) We exploit the capacity of nonlinear projections employed in the ASWD by constructing injective mapping with arbitrary neural networks: (ii) We prove that the ASWD is valid distance metric; (iii) We provide a mechanism in which the hypersurface where nigh-dimensional distributions are projected onto can be optimized and show that the optimization of hypersurfaces can help improve the projection efficiency of slice-based Wasserstein distances. Hence, the ASWD is data-adaptive, i.e. the hypersurfaces can be learned from data. This implies one does not need to manually design a function from the limited class of candidates; (iv) We demonstrate superior performance of the ASWD in numerical experiments for both synthetic and real-world datasets.
82
63
64
65
66
67
68
70

---

72
The remainder of the paper is organized as follows. Section 2 reviews the necessary background. We present the proposed method and its numerical implementation Section Related work are discussed in Section 4. Numerical experiment results are presented and discussed in Section 5. We conclude the paper in Section
73
74
75
Background
76
77
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.
78
79
Wasserstein distance: Let Pk: (92) be a set of Borel probability measures with finite k-th moment on Polish metric space (22d) [Villani, 2008]. Given two probability measures μ., Ps((2),, whose probability density functions (PDFs) are Pμ and pv, the Wasserstein distance of order A: E (00)
81
82
between μ and is defined as:
Wu(m,))= ( d(xx,)) $
where d(::-)* is the cost function, T(,,v)) represents the of all transportation plans 7, i.e. joint distributions whose marginals are Pp and Pv. respectively. With a slight abuse of notation, we interchangeably use (1,,)) and Wa (Pp.Po)))
84
86 While the Wasserstein distance is generally intractable for high-dimensional distributions, there are several favorable cases where the optimal transport problem can be efficiently solved. If μ and 86 are continuous one-dimensional measures defined on linear space equipped with the L* norm, the Wasserstein distance between y and has closed- form solution [Peyré and Cuturi, 2019]:
(x(p,))= = ( - *
90
where F and F,11 are inverse cumulative distribution functions (CDFs) of μ and v. respectively.
91 Radon transform and generalized Radon transform: The Radon transform [Radon, 1917] maps ge
92 a function f(-) E L L²RR)) to the space of functions defined over spaces of lines in IR The Radon transform of f()) is defined by line integrals of f() along all possible hyperplanes in Rd:
93
R,((1,)= = gd
(x))(((- (x,0))
(3)
where t E IR and 0 @ represent the parameters of hyperplanes {ar E (z.6) t), 5(-) is the 95 Dirac delta function, and (.,) refers to the Euclidean inner product.
By replacing the inner product (xr,0) in Equation (3) with (xx,0), specific family of functions named 97 as defiming function in [Kolouri al., 2019a], the generalized Radon transform (GRT) [Beylkin, 1984] defined as integrals of f-) along hypersurfaces defined by f ((x,0)==t]
Gf(t, f f(z))((-((z.0))
(4)
99 where E R, 0 E
Qg while 2g is a compact set of all feasible 6, e.g. 2o Sd-1 for 8(x,8)=(z [Kolouri et al., 2019a].
101 In practice, we can empirically approximate the Radon transform and the GRT of probability density 102 function Py via:
N
(t,0
Σ
6(t
()),
(5)
N
==1
N
(p.(t
6(t ()xn
(6)
n=1
where vn ~ Pp and N is the number of samples. Notably, the Radon transform is a linear bijection 104 [Helgason, 1980], and the GRT a bijection if the defining function 9 satisfies certain conditions [Beylkin, 1984].
3

---

106
Sliced Wasserstein distance and generalized sliced Wasserstein distance: By applying the Radon transform to Pμ, and Pa to obtain multiple projections, the sliced Wasserstein distance (SWD) decom- poses the high-dimensional Wasserstein distance into multiple one-dimensional Wasserstein distances which can efficiently evaluated [Bonneel et al., 2015]. The k-SWD between μl and is defined by:
107
108
109
SWD&(u,D) = -fo gd-1
110 where the Radon transform R defined by Equation (3) is adopted as the measure forsaro operator.
!11 The GSWD generalizes the idea of SWD by slicing distributions with hypersurfaces rather than
112 hyperplanes [Kolouri et al., 2019a]. The GSWD is defined as:
(L (PP
(8)
113 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
115 distance between sorted samples [Kolouri et al., 2019b]: with random projections, the SWD and
116 GSWD between μ and be approximated by:
N Kezgm,.00) (w1)))))))))))))))) k
SWDx(u,)
(9)
NT. 1=
L 8 (1) -B
:
GSWD)(M,))
(10)
i=n==1
where l and r, are sequences consisting of the indices of sorted samples which satisfy (1a1 [v] (r) S
117
(11g(+11, 01), (915Im) B1) S (Ur)i=+1', 01) in the SWD, and ( B1) I ((x in+1]: 01).
118
(Vrgta) (1) S (Brg [1)] 0t) in the GSWD. is proved in [Bonnotte, 2013] that the SWD is valid
119
120
distance metric. The GSWD is valid metric except for its neural network variant [Kolouri et al., 2019a].
3
Augmented sliced Wasserstein distances
122
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.
124
125
3.1
Spatial Radon transform and augmented sliced Wasserstein distance
In the definitions of the SWD and GSWD, the Radon transform [Radon, 1917] and the generalized Radon transform (GRT) [Beylkin, 1984] are used as the push-forward operator for projecting distributions to one dimensiooii space. However, it not straightforward to design defining functions (1x #) [Kolouri et al., 2019a] for the GRT due to certain non- non-trivial requirements for the function [eylkin,, 1984]. In practice, the assumption of the transform be relaxed, as Theorem shows that as long as the transform is injective, the corresponding ASWD metric is a valid distance metric.
127
129
130
131
To help define the augmented sliced Wasserstein distance, first introduce the spatial Radon transform which includes the vanilla Radon transform and the polynomial GRT as special cases (See Remark 2).
133
Definition 1. Given injective mapping g()) Rdd - Rdd and probability measure μ E P(Rd))
135
136
whose probability density function (PDF) is Ppo the spatial Radon transform of Py is defined as
Hpp pp(z))((-- (
(11)
where EE R. and E are the parameters of hypersurfaces EI (g(x).0)=1)
137
138
Remark 1. Note that the spatial Radon transform can be interpreted as applying the vanilla Radon transform the PDF of f = g(z), where x ~ Pu- Denote the PDF of &t by Pp, the spatial Radon

---

140
transform defined by Equation (11) can be rewritten as:
Htp(t,0:g)) Lispp (0 ()()))))
Pampi, (6(t- (i.
| Pp., (0))(t (ā.6)
(Rpp,(t.))
(12)
Hence the spetial Radon transform inherits the theoretical properties of the Radon transform subject to certain conditions of g(-) and incorporates nonlinear projections through g(-).
142
143
In what follows, we use fi f2 to denote functions : and f2)) X + E that satisfy
144
fi (z) f2(x)) for VI eX
Lemma 1. Given an injective mapping g(-) :R. - Rdd and two probubility measures u. P(R4) whose probability density functions are Py and Pu, respectively, for all E R and E gd
148
147
(t. (pp(t.0:) if and only Pu i.e. the spatial Radon transform injective. Moreover, the spatial Radon transform is injective if and only if mapping 9(-) is an injection.
148
See Appendix A for the proof of Lemma 1.
150
Remark 2. The spatial Radon transform degenerates to the vanilla Radon transform when the mapping g(-) is an identity mapping. When g(-) is a homogeneous polynomial function with odd degrees, the spatial Radon transform is equivalent the polynomial GRT [Ehrenpreis, 2003].
151
152
153
Appendix provides the proof of Remark 2.
154
We now introduce the augmented sliced Wasserstein distance, by utilizing the spatial Radon transform the measure push-for forward operator:
Definition& 2. Given two probability measures Ju,2 E Pk (R.d), whose probability density functions are Pa and Pv. respectively, and an injective mapping 9()) :R= Wida the augmented sliced Wasserstein distance (ASWD) of order k E [1,+oc) is defined as:
156
157
158
(S
t
ASWD (1.,rg))
(13)
where E 5ds-1, & is the k- Wassersteei distance defined by Equation (I), and H. refers to the spatial Radon transform defined by Equation (11).
159
161
Remark 3. Following the connection between the spatial Radon transform and the vanilla Radon transform as shown in Equation (12), the ASWD can be rewritten
(
t
ASWD (H.)gg)=
=SWD&(igggg)
(14)
where by and bg are probability measures on RRdo which satisfy g(x) ~ fig for x ~ and gly) ~ ig for y-u
163
185
Theorem 1. The augmented sliced Wasserstein distance (ASWD) order E [1,+ defined by Equation (13) with a mapping g(-):TRd + RRd a metric on P& (TR.) if and only ifg)))) is injective.
186
187
The proof of Theorem is provided in Appendix C. Theorem shows that the ASWD is metric given a fixed injective mapping g(-). In practical applications, the mapping g(-) needs to be optimized. We show in Corollary 1.1 that the ASWD between μ and with the optimized g(-) -) is also metric under mild conditions.
188
169
170
Corollary 1.1. The augmented sliced Wasserstein distance (ASWD) of order k E [1, +00)
172
between two probability p, E (RRd) defined by Equation (13) with the optimal mapping
g" (-) Angmax(ASWDDiuuaa) is metric on (RR) when the optimization is confined to the set
173
of bounded and f {g(z) :R - ER (19()))))
The proof of Corollary 1.1 is provided in Appendix D.
Remark 4. Corollary 1.1 shows that given measures μt ,22,83 E (TRd)) the triangle inequality holds for the ASWD when g(-) is optimized for each pair of measures, shown in Appendix D.
176
5

---

178
3.2
Numerical implementation
179
We discuss in this section how to realize injective mapping g(+) with neural networks due to their expressiveness and optimize it with gradient based methods.
180
Injective neural networks: As stated in Lemma and Theorem 1, the injectivity of ()() is the sufficient and necessary condition for the ASWD being a valid metric. Thus we need specific architecture designs on implementing g()) by neural networks. One option is the family of invertible neural networks [Behrmann et al., 2019; Karami et al., 2019], which are both injective and surjectíve. 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 x of an arbitrary neural network to its output (z):
g()
182
183
184
185
186
(x) = (x))-
(15)
187
It is trivial to show that 9u. (x(x)) 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.
188
189
190
191
We aim to slice distributions with 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 and Cuturi, 2019] or a 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:
Optimization objective:
192
193
195
196
197
((p.trgw.A)
(gL) W (Hpu(.:0:g)) Hlp d0 - م2
(16)
where A 0 and the regularization term LA XE [[ (2)]12+ 11g (y)ila) used to control the norm of the output of (:),),, otherwise the projections may be arbitrarily large.
198
199
Remark 5. The regularization coefficient X adjusts the introduced non-linearity in the evaluation of
201
the ASWD by controlling the norm of ow in Equation (15). In particular, when X co, the nonlinear
term (-)()) shrinks to O. The rank of the augmented space is hence explicitly controlled by the flexible choice of (-) and implicitly regularized by Lx.
203
204
By plugging the optimized (-) arggmx( ((i,, U3 9w, into Equation (13), we obtain the
empirical version of the ASWD. Pseudocode is provided in Appendix E.
4
Related work
208
207
Recent work on slice-based Wasserstein distances mainly focused on improving their projection efficiency, leading to a reduced number of projections needed to capture the structure of data distributions [Kolouri et al., 2019a; Nguyen et al., 2021]. The GSWD proposes using nonlinear projections to achieve this goal, and has been proved to be valid distance metric if and only they adopt injective GRTs, which only include the circular functions and finite number of harmonic polynomial functions with odd degrees as their feasible defining functions [Ehrenpreis, 2003]. While the GSWD has shown impressive performance in various applications [Kolouri et al., 2019a], its defining function is restricted to the aforementioned limited class of candidates. In addition, the selection of defining function is usually task-dependent and needs domain knowledge, and the impact on performance from different defining functions is still unclear.
208
209
210
212
214
216
217
To tackle those limitations, [Kolouri et al., 2019a] proposed the GSWD-NN, which directly takes the outputs of neural network as its projection results without using the standard Radon transform or GRTs. However, this brings three side effects: 1) The number of projections, which equals the
218
219
220
number of nodes in the neural network's output layer, is fixed, thus new neural networks are needed
if one wants to change the number of projections. 2) There is no random projections involved in the
222
GSWD-NN, as the projection results are determined by the inputs and weights of the neural network. 3)
The GSWD-NNI is a pseudo-metric since it uses vanilla neural network, rather than Radon transform

---

224
or GRTs, as its push-forward operator. Therefore, the GSWD-NN does not fit into the theoretical framework of GSWD and does not inherit its geometric properties.
225
226
Another notable variant of the SWD) is the distributional sliced Wasserstein distance (DSWD) [Nguyen al., 2021]. By finding a distribution of projections that maximizes the expected distances over these projections, the DSWD can slice distributions from multiple directions while having high projection efficiency. Injective GRTs are also used to extend the DSWD the distributional generalized sliced Wasserstein distance (DGSWD) INguyen et al., 2021]. Experiment results show that the DSWD and the DGSWD have superior performance in generative modelling tasks Nguyen et al. [2021]. However, neither the DSWD nor the DGSWD have solved the problem with the GSWD, i.e. they are still not able to produce nonlinear projections adaptively.
227
228
229
230
231
232
233
234
Our contribution differs from previous work in three ways: 1) The ASWD is data-adaptive, i.e. the hyper- surfaces where high-dimensional distributions are projected onto be learned from data. This implies one does not need to specify a defining function from limited choices. 2) Unlike GSWD-NN, the ASWD takes novel direction to incorporate neural networks into the framework of sliced-based Wasserstein distances while maintaining the properties of sliced Wasserstein distances. 3) Previous work on introduc-
235
236
237
238
239
ing nonlinear projections into Radon transform either is restricted to only a few candidates of defining functions (GRT') or breaks the framework of Radon transforms (neural networks in GSWD-NN), in
240
241
contrast, the spatial Radon transform provides a novel way of defining nonlinear Radon-type transforms.
5
Experiments
242
243
In this section, we describe the experiments that we have conducted to 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., 2021] on CIFAR 10 [Krizhevsky, 2009], CelebA [Liu et al., 2015], and MNIST [LeCun et al., 1998] datasets (Appendix H.2). Hence, we compare performance of the ASWD with the oftt--of-hhe-ar distance metrics in the same examples and report results as below". 1 We provide additional experiment results in the appendices, including a sliced Wasserstein autoencoder (SWAE) [Kolouri al., 2019b] using the ASWD (Appendix I), image color transferring (Appendix J) and sliced Wasserstein barycenters (Appendix K).
245
246
247
248
249
250
251
252
To examine the robustness of the ASWD, throughout the experiments, adopt the injective network architecture given in Equation (15) and set &w to be a single fully-comnected layer neural network whose output dimension equals its input dimension, with ReLU layer as its activation function.
253
255
5.1
Sliced Wasserstein flows
256
We first consider the problem of evolving source distribution μ to a target distribution v by minimizing Wasserstein distances between and v in the sliced Wasserstein flow task reported in [Kolouri et al.,
257
258
201 19a].
0.14 = VSWD(
(17)
where i4 refers to the updated source distribution at each iteration 1. The SWD) in Equation (17) can be replaced by other sliced- Wasserstein distances to be evaluated. As in [Kolouri et al., 2019a], the 2- Wasserstein distance was used as the metric for evaluating performance of different distance metrics 261
259
260
262
in this task. The set of hyperparameter values used in this experiment can be found in Appendix F.
263
Without loss of generality, we initialize po to be the standard normal distribution N (0,T). We repeat each experiment 50 times and record the 2- Wasserstein distance between μ and at every iteration. In Figure 2, we plot the 2- Wassersstei distances between the source and target distributions as function of the training epochs and the 8-Gaussian, the Knot, the Moon, and the Swiss roll distributions are respective target distributions. For clarity, Figure 2 displays the experiment results from the 6 best performing distance metrics, including the ASWD, the DSWD, the SWD, the GSWD-NN 1, which directly generates projections through one layer MLP, as well as the GSWD with the polynomial of degree 3, circular defining functions, out of the 12 distance metrics we compared.
264
285
286
287
268
269
270
271
We observe from Figure 2 that the ASWD not only leads to smaller 2- Wasserstein distances, but also converges faster by achieving better results with fewer iterations than the other methods in these four
272
Code to reproduce experiment results is available at: https: betps://bbt... 1y /Y22200
7

---

ASWD
Polynomial 3
DSWD GSW-NN 1
Linear
Circular
Gaussian
Knot
logso (erooWZ
0.0
00.5
1.0
1000
1500
2000
500
1000
1500
2000
lteration
Iteration
Moon
Swiss Roll
0.4
log10
0.5
0.2
0.0
0.0
1000 1500
2000
0
500
1000
1500
2000
Iteration
Iteration
Figure 2: The first and third columns are target distributions. The second and fourth columns are log 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- 2-Wasserstein distances over 50 runs. A more extensive set of experimental results can be found in Appendix G.I.
310
ASWD
ASWD
DSWD
300
DSWD
260
GSWD
GSWD
FID score
SWD
FID score
250
SWD
210
200
160
150
110
60
50
100
150
200
20
40
60
80
100
Epoch
Epoch
Figure 3: FID scores of generative models trained with different metries CIFAR10 (left) and CelebA (right) datasets with T= 1000 projections. The error bar represents the standard deviation the FII) scores the specified training epoch among 10 simulation runs.
273
target distributions. A complete experimental results with 12 compared distance metrics and 8 target distributions are included in Appendix G.1. The ASWD outperforms the compared of the-art sliced-based Wasserstein distance metries with 7 out of the 8 target distributions except for
274
275
25- Gaussian. This is achieved through the simple injective network architecture given in Equation (15) neural network with equal input and output dimensions throughout the experiments. In addition, ablation study is conducted to study the effect of injective neural networks, the optimization of hypersurfaces in the ASWD. Details can be found in Appendix G.2.
277
and one layer fully-connecto
279
280
Generative modeling
In this experiment, we use the sliced-based Wasserstein distances for generative modeling task described [Nguyen et al., 2021]. The task is to generate images using generative adversarial
282
networks (GANs) [Goodfellow et al., 2014] trained on either the CIFAR 10 dataset (64 x64 resolution) [Krizhvssky, 2009] or the CelebA dataset (64 64 resolution) [Liu al., 2015]. Denote the hidden layer and the output layer of the discriminator by hy and De, and the generator by G4, train GAN
284

---

FID scores of generative models trained with different distance metrics. Smaller 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 CIFAR 10 dataset and CELEBA dataset. The running time per training iteration for one batch containing 512 samples is computed based on 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.
Table 1:
CIFARIO
SWD
GSWD
DSWD
ASWD
[Bonneel et al., 2015]
[Kolouri al., 2019a]
[Nguyen et al., 2021]
FID
(s/it)
FID
t (s/it)
FID
(s/it)
FID
{s/it)
10
192.645.7
0.32
189.516.00
0.35
79.014.2
0.48
73.243.1
0.55
100
155.0#2.99
0.32
155.9:332
0.70
7..2
0.51
66.7.13.2
0.57
1000
126.0:12.9
0.34
134.5:2.77
2.10
74.3
1.22
65.5.13.9
1.32
CELEBA
10
118.343.1
0.32
143.215..5
0.35
105.313.4
0.49
99.2.14.3
0.53
100
116.0=2288
0.33
1.0
0.69
103.1
0.51
94.3#22.2
0.56
1000
104.442.8
0.34
101.8+1.8
2.14
97.4#2.1
1.21
90.5+1.00
1.31
286
models with the following objectives:
min SWD( (po) (G+(p=)),
(18)
max Exap, [log( D (h V.P
(19)
where p. is the prior of latent variable and p. is the distribution of real data. The SWD in Equation (18) is replaced by the ASWD and other variants of the SWD to compare their performance. The
288
GSWD with the polynomial defining function and the DGSWD is not included in this experiment due
290
to its excessively high computational cost in high-dimensional space.The The Fréchet Inception Distance (FID score) [Heusel al., 2017] used assess the quality of generated images. More details on the network structures and the parameter setup used in this experiment are available in Appendix F.2.
291
292
We run 200 and 100 training epochs to train the GAN models on the CIFARIO and the CelebA dataset, respectively. Each experiment is repeated for times. We report experimental results in Table 1. With the same number of projections and similar computation cost, the ASWD leads to significantly improved FID scores among all evaluated distances metrics on both datasets, which implics that images generated with the ASWD of higher qualities. Figure 3 plots the FID scores recorded during the training process. The GAN model trained with the ASWD exhibits a faster convergence as it reaches smaller FIID scores with fewer epochs. Randomly selected samples of generated images presented in Appendix H.1.
294
296
298
300
6
Conclusion
302
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 hypersurfaces used to slice compared distributions by learning from data. We proved that the ASWD is a valid distance metric and presented its numerical implementation. We reported empirical performance of the ASWD over state-of-the-ae sliced Wasserstein metrics in various numerical experiments. We showed that ASWD with simple injective neural network architecture can lead to the smallest distance errors over the majority of datasets a sliced Wasserstein
304
306
308
flow task and superior performance in generative modeling tasks involving GANs and VAEs. We have also evaluated the applications of the ASWD in downstream tasks including color transferring and Wasserstein barycenters. What remains to be explored is the impact of the injective neural network
310
312
architecture used in the ASWD, e.g. the application of different types of invertible neural networks in the ASWD framework. We leave this topic potential research direction of our future work.

---

References
314
J. Altschuler.
and P. Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. In Pruc. Advances in Neural Information Processing Systems (NeurIPS), pages 1964-1974. Long Beach, California, USA, 2017.
316
317
318 M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In Proc. International Conference on Machine Learning (ICML). pages 214-223, Sydney, Australia, 2017.
319
320
J. Behrmann, W. Grathwohl, R. T. Q. Chen, D. Duvenaud, and Jacobsen. Invertible residual networks. In Proc. International Conference Machine Learning (ICML), pages 573-582, Long Beach, California, USA, 2019.
G. Beylkin. The inversion problem and applications of the generalized Radon transform. Communications on Pure and Applied Mathematics, 37(5):579 599, 1984.
323
324
N. Bonneel, Rabin, G. Peyré, and Plister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision. 51(1):22-455 4,.222245,222 2015.
325
326
N. Bonnotte. Unidimensional and evolution methods for optimal transportation. PhD thesis, Paris 11, 2013.
327
N. Courty, R. Flamary, D. Tuia, and A. Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions Pattern Anolysis and Machine Intelligence (IPAMI), 39(9):1853-1865 2016.
M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 2292-2300, Lake Tahoe, Nevada, USA, 2013.
M. Cuturi and A. Doucet. Fast computation of Wasserstein barycenters. In Proc. International Conference on Machine Learning (ICML), pages 685-693, Beijing, China, 2014.
332
335
1. Deshpande, Hu, R. Sun, A. Pyrros, et al. Max-sliced Wasserstein distance and its use for GANs. In Proc. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 10648-10656, Long Beach, California, 2019
334
325
336 A. Dessein, N. Papadakis, and J. Rouas. Regularized optimal transport and the rot mover's distance. The Journal of Machine Learning Research (JMIR), 19(1):590-642, 2018.
337
338
Ehrempreis. The universality of the Radon transform, chapter 5, pages 299-363. Oxford University Press, Oxford. UK,2003 2003.
339
340 A. Genevay, L. Chizat, Bach, M. Cuturi, and G. Peyré. Sample complexity of Sinkhorn divergences. Interna- tional Conference on Artificial Intelligence and Statistics (AISTATS), pages 1574-1583, Okinawa, Japan, 2019.
341
342 I. Goodfellow, J. Pouget, M. Mirza, B. Xu, F. D. Warde, ct al. Generative adversarial nets. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 2672-2680, Montréal, Canada, 2014.
L. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. Courville. Improved training of Wasserstein GANs. In Proc. Advances neural information processing systems (NeurIPS), pages 5767-5777, Long Beach, California, USA, 2017.
346
347
Helgason. The Radon transform, volume 2. Basel, Switzerland: Springer, 1980.
348
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), pages 6626-6637, Long Beach, California, USA, 2017.
349
350
G. Huang, 7. Liu, L. Van Der Maaten, and K. Q. Weinberger. Densely connected convolutional networks. In Proc. IEFE Conference on Computer Vision and Pattern Recognition (CVPR), pages 4700-4708, Hawaii, USA, 2017.
352
353
M. Karami, Schuurmans, Sohh-Dickstein, L. Dinh, and Duckworth. Invertible convolutional flow. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 5636-5646, Vancouver, Canada, 2019.
354
355
Kolouri, K. Nadjahi, Simsekli, R. Badeau, and G. Rohde. Generalized sliced Wasserstein distances. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 261 1222, Vancouver, Canada, 2019a.
356
357 S. Kolouri, P. E. Pope, C. Martin, and G. K. Rohde. Sliced- Soadd--assesstein autocncoders. 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.
10

---

380
Y. LeCun, L. Bottou. Y. Bengio, and P. Haffner. Gradient- learning applied to document recognition. Proceedings the IEEE, 86( 1):2278)22222 1998.
362
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proc. IEEE International Conference Computer Vision (ICCV), pages 3730-373 3738, Las Condes, Chile, 2015.
364 B. Muzellec and M. Cuturi. Subspace detours: Building transport plans that are optimal on subspace projections. Proc. Advances in Neural Information Processing Systems (NeurIPS), 32:6917- 2019.
K. Nguyen, N. Ho, T. Pham, and H. Bui. Distributional sliced- Wasserstein and applications to generative modeling. In Proc. International Conference on Learning Reprexentations (ICLR). Vienna, Austria, 2021.
368 F. Paty and M. Cuturi. Subspace robust Wasserstein distances. In Proc. International Conference on Machine Learning (ICML), pages 5072-5081, Long Beach, Califomia, USA, 2019.
G. Peyré and M. Cuturi. Computational optimal transport. Foundations and Trends@ in Machine Learning, 11 (5-6):355-607,2)) 2019.
J. Radon. Uber die hestimmung von funktionen durch ihre integralwerte laengs gewisser mannigfaltigkeiten Ber. Verh. Saechs. Akad. Wiss. Leipzig Math. Phys. K1., 69:262, 1917.
T. Salimans, H. Zhang, A. Radford, and D. Mctaxas. Improving GANs using optimal transport. In Proc. International Conference Learning Representations (ICLR), Vancouver, Canada, 2018.
C. Villani. Optimal Transport: old and new, volume 338. Berlin, Germany: Science 2008.
Science & Business Media,

---

Checklist
378
379
1. For all authors...
380
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
381
contributions and scope? [Yes] See abstract and Section 1.
382
(b) Did you describe the limitations of your work? [Yes] See Conclusion.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Appendix L.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
385
2. If you are including theoretical results...
(a) Did you state the full of assumptions of all theoretical results? [Yes] We state assumptions of our work Thenrem 1 and Corollary 1.1.
(b) Did you include complete proofs of all theoretical results? [Yes] Theoretical proofs are presented Appendices A, and
389
390
3. If you ran experiments...
391
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (cither in the supplemental material or as URL)? [Yes] The code to reproduce our results is given in https:/t///tt //bit. 1y/2Y23w0z.
392
393
394
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] Experiment details are provided in F.
395
396
(c) Did you report error bars (c.g., with respect to the random seed after running experiments multiple times)? [Yes] See Figure 2, Figure 3, Figure 7, Figure 10, and Table 1.
398
(d) Did you include the total amount of compute and the type of resources used (e.g.. type of GPUs, internal cluster, cloud provider)? [Yes] As is reported in the caption of Table our results computed based computer with an Intel (R) Xeon (R) Gold 5218 CPU 2.3 GHz and 16GB RAM, and a RTX 6000 graphic card with 22GB memories.
401
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
403
(a) If your work uses existing assets, did you cite the creators? [Yes] See Krizhevsky [2009]; LoCun et al. [1998]; Liu et al. [2015].
404
405
(b) Did you mention the license the assets? [N/A]
406
(c) Did you include any new assets either in the supplemental material or as URL? [N/A]
407
(d) Did you discuss whether and how consent was obtained from people whose data you're
408
using/curating" [N/A]
409
(e) Did you discuss whether data you are using/curating contains personally identifiable information or offensive content? [N/A]
410
5. If you used crowdsourcing conducted research with human subjects...
412
(a) Did you include the full text of instructions given to participants and screenshots, if applicaall??
413
[N/A]
414
(b) Did you descrihe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
415
416
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
417
12


Reviews:
[{"review_id": "lCPzv7ODvaH", "reviewer": "Reviewer_boZk", "summary": "The sliced Wasserstein distance (SWD) is a metric on the space of probability measures and serves as a practical alternative to the standard metric arising from the optimal transport (OT) problem, the Wasserstein distance. Indeed, while the Wasserstein distance suffers from computational and statistical limitations on large-scale and high-dimensional settings, SWD can alleviate these issues while offering similar theoretical properties. Hence, SWD has been increasingly popular within the statistical and machine learning (ML) community in recent years.\n\nSWD is defined as an expectation that is intractable in general; in practice, it is then usually approximated with a simple Monte Carlo estimation, but prior studies reported that this method might induce an important approximation error on high-dimensional settings, leading to an underestimation of the \"true\" dissimilarity between two distributions. Therefore, one would need to compute a Monte Carlo approximation based a very large amount of samples to obtain an accurate approximation, which increases the computational complexity of SWD. \n\nSeveral variants of SWD have recently been introduced to overcome this problem and include the \"generalized sliced Wasserstein distances\" (GSWD): this class of distances aim at offering a better computational efficiency, by comparing distributions from their nonlinear one-dimensional representations (instead of the linear ones used in SWD). This paper focuses on this line of work: the limitations of GSWD are identified to motivate the formulation of a novel class of distances, called \"Augmented Sliced Wasserstein distances\" (ASWD).\n\nSimilarly to SWD and GSWD, the definition of ASWD is connected to the Radon transform: the authors introduce a novel type of Radon transform, the \"spatial Radon transform\" (Definition 1), and use it to define the class of augmented sliced Wasserstein distances. \nThen, they study the theoretical properties of the spatial Radon transform and ASWD. Their main result states that ASWD verify all metric axioms if and only if the mapping $g$ (that characterizes the underlying spatial Radon transform) (1) is injective (Theorem 1), or (2) maximizes ASWD over the space of bounded and injective functions (Corollary 1.1).\nThe authors then introduce a specific class of mappings (Equation (15)) which meet the conditions of Corollary 1.1 thus guarantee that the resulting ASWD is a metric. Such mappings require training a neural network to make ASWD data-adaptive.\nFinally, the superior empirical performance of ASWD against existing variants of SWD (including GSWD) is illustrated on various problems, including gradient flows, generative modeling, color transferring and barycenters of measures. \n", "questions": "", "limitations": "", "rating": 5, "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": "This work proposes a new family of distances inspired by computational optimal transport, which can be easily applied to various classical machine learning applications. In that sense, it fits the interests of the NeurIPS community.\n\nI found the paper quite clear and well-written overall. The structure can be improved by incorporating the following changes,\n- l.46-47: \"Different variants of the SWD have been proposed to improve the projection efficiency of the SWD, either by introducing nonlinear projections or by optimizing the distribution of random projections.\" It seems more relevant to talk about the max-sliced Wasserstein distance (max-SWD) in this paragraph instead of in the previous one, since max-SWD was introduced in order to improve the \"projection complexity issue\" of SWD (see \"Max-Sliced Wasserstein Distance and its use for GANs\", Deshpande et al., 2019). \n- I suggest adding the equation that explains how to approximate equation (2) at the end of the paragraph \"Wasserstein distance\" (Section 2). This would clarify l.113-115, as well as equations (9) and (10)\n- l.101-105: I would remove this part since explaining how to approximate the Radon transform or GRT does not seem useful in this work. On the contrary, specifying the approximation for the regularization term in Section 3.2 (l.198) would be more relevant.\n- l.168: \"In practical applications, the mapping $g(\\cdot)$ needs to be optimized.\" At this stage, it is not obvious why this optimization is required. I suggest moving the explanation in Section 3.2 (e.g., l.191-196) to clarify this point.\n\nOverall, the theoretical claims look correct to me; my only concerns are detailed below.\n- Remark 2: \"$g(\\cdot)$ is a homogeneous polynomial function with odd degrees\" (l.151) contradicts equation (31) in Appendix B. On the other hand, it is not clear to me why the mapping defined in equation (31) is injective on $\\mathbb{R}^d$.\n- Appendix D: Why is it important to assume that the optimal mapping is bounded? Besides, equations (39), (40), (41) could be explained in more details.\n\nOther (minor) typos and inaccuracies are listed at the end of the review.\n\nIn my opinion, the main weakness of this work, which prevents me from giving a positive score, is the lack of originality and significance as compared to prior work. Indeed, ASWD build on GSWD and their definition differs from a simple trick (using the \"spatial Radon transform\" introduced by the authors, instead of the generalized Radon transform). The structure of this paper is the same as in \"Generalized Sliced Wasserstein Distances\" (Kolouri et al., 2019) [1]: the novel distance is defined, then its metric axioms are studied, finally its performance against the sliced Wasserstein distance (and its variants) is evaluated on practical applications. These two aspects are not necessarily concerning and help make the paper quite intuitive; besides, the empirical study provided by this paper is more extensive than the one in [1], since the compared methods and applications are more diverse. What is more questionable is that the advantages of ASWD as compared to SWD and GSWD are unclear to me in the end, given the contributions and discussion provided in this work. \n\nFirst, from a theoretical point of view, this paper only investigates the metric properties, and establishes that ASWD satisfy all axioms of metrics if and only if the mapping associated to its Radon transform is injective. While the authors argue that this is a significant advantage over GSWD, we don't know if this explains the superior performance of ASWD-based methods in practice. Besides, this result is similar to Proposition 1 in [1], which states that GSWD are metrics if and only if the corresponding generalized Radon transform is injective; consequently, the proofs in the appendix are largely inspired by [1].\n\nThen, the empirical performance of ASWD is not sufficiently discussed: while the results in Section 5 encourage its use against existing slice-based distances, the appendix illustrates that ASWD do not necessarily yield a superior performance:\n- Appendix I and K: ASWD and SWD achieve a similar performance, while ASWD is more computationally expensive.\n- Appendix I and J: the performance of GSWD should be studied on these experiments as well.\n\nThis shows that the superior performance of ASWD depends on the application, and the theoretical results do not help understand the behavior in practice. In particular, I am wondering whether the regularization term (equation (16)) might help improve the performance in the flows and generative modeling applications: it would be interesting to regularize the other methods to study the impact. \n\nHence, I encourage the authors to provide more significant evidence on why ASWD is a better alternative to SWD than GSWD, or to show a more nuanced picture in light of the empirical results in the appendix; for now, I am not convinced by the superiority of ASWD over SWD/GSWD given the motivation (in particular, l.234-241) nor the theoretical and empirical contributions.\n\nTypos:\n\nEquation (1): the integral should be over $\\Omega \\times \\Omega$\n\nl.92: \"over spaces of lines in $\\mathbb{R}^d$\" -> \"over spaces of hyperplanes in $\\mathbb{R}^d$\" \n\nl.99: \"while $\\Omega_\\theta$ is a compact set\" -> \"where $\\Omega_\\theta$ is a compact set\"\n\nl.110: \"where the Radon transform ... is adopted as the measure push-forward operator\" This can be further clarified by explaining the connection between the Radon transform and the push-forward operator that projects the two distributions (Proposition 6 in \"Sliced and Radon Wasserstein Barycenters of Measures\", Bonneel et al.).\n\nl.120: \"The GSWD is a valid metric except for its neural network variant [Kolouri et al., 2019a]\" This is not exactly true: it remains unknown whether neural networks-based GSWD (GSWD-NN) satisfy all axioms for a metric (see the explanation in Section 3.3 of \"Generalized Sliced Wasserstein Distances\", Kolouri et al.). Therefore, a correct statement would be, for example, \"The GSWD has been shown to be a valid metric, but the proof remains incomplete for its neural network variant [Kolouri et al., 2019a]\".\n\nDefinition 1: \n- The set $P(\\mathbb{R}^d)$ is not defined\n- What are the possible values for $d$ and $d_\\theta$?\n\nl.147: please define the operator \"\\equiv\".\n\nl.172: \"between two probability\" -> \"between two probability measures\"\n\nSection 5: The order $k$ of SWD/ASWD used in the experiments should be precised. \n\nl.256-257: \"by minimizing Wasserstein distances\" -> \"by minimizing slice-based Wasserstein distances\"\n\nl.425, 458, 463: \"for $\\forall x \\in \\mathbb{R}^d$\" -> \"for all $x \\in \\mathbb{R}^d$\". Same remark in l.428, 432, 434, 506...\n\nl.455: \"which probability density function is ...\" -> \"whose probability density function is ...\"\n\nl.498: \"that the the ASWD\"\n\nl.589: \"Sampels of generated images\" -> \"Samples of generated images\"\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations of this work are not clearly discussed, but the authors adequately addressed the societal impact (see Appendix L).\n", "review_text": "", "extra_scores": {"ethical_concerns": "I don't have any ethical concerns for this paper.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "hhEtTyRujMv", "reviewer": "Reviewer_7Tc3", "summary": "The author proposed a variant of sliced Wasserstein distance (SWD) called **augmented sliced Wasserstein distance** (ASWD) which replace the linear projections in SWD with a non-linear one. The construction of ASWD consists of two steps: (1) define spatial Radon transform (SRT) for probability measures (2) replace the Radon transform in the original SWD with this new SRT. The key intuition of SRT is to first transform the function $\\pmb{x}$ to high dimensional hyperspace by transformation $g(\\pmb{x})$ and then perform Radon transform on functions defined in this hyperspace. \n\nTheoretically, the author proves the conditions required for SRT to be injective. Later, the author uses this result to show the corresponding ASWD is a distance metric. In addition, the author proposes a training objective for finding the transformation $g(\\cdot)$. \n\nEmpirically, the author conducts both toy and practical experiments to demonstrate the advantages of ASWD in terms of higher projection efficiency and also provides some ablation studies in the appendix to investigate the effects of different $g(\\cdot)$ choices.", "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": "### Clarity\nThe main idea of ASWD is easy to understand and the author also gives a quite clear explanation on how to construct ASWD. Specifically, I appreciate the **Remarks**, which greatly improves the clarity of the paper. So I think the paper is clearly written and quite easy to follow. \n\n### Originality\nAs for the originality, the idea of using injective transformation to preserve some statistical properties is not new. For example, this is the foundation of normalizing flow, and it can also be found in the deep-kernel literature [1] to preserve the discrepancy validity. However, this seems to be new in the scope of SWD literature. In summary, I found the originality is ok but not significantly novel. \n\n### Quality and Significance\nIn terms of the research quality, the method is well-presented. I checked some of the proofs, they seem to be correct. Empirically, the author conducts several experiments to confirm the advantages of ASWD. However, I still have some concerns related to the experiments and the way the author presented the ASWD. In terms of significance, this work should be easy to implement and understand. So it should have some impact on the SWD community. \n\n### Questions\n1. You mentioned that the idea of SRT is to transform the original data $\\pmb{x}\\in\\mathbb{R}^d$ to $\\mathbb{R}^{d_\\theta}$ and $d_\\theta> d$, i.e. higher-dimensional space. However, I don't think this is necessary based on the theory you give. For structured data like images, the commonly used assumption is that they have a low-dimensional manifold. So I don't understand why you project it in higher-dimensional space. To me, it seems that this is due to the choice of the NN architecture (concatenation) and not from the theoretical arguments.\n2. The form of the transformation $g$ is quite simple. Have you evaluated the performance of using the invertible neural network for $g$? Despite the computational cost, does it give better performance?\n3. I wonder about the performance if you transform the data to low dimensional space by an arbitrary neural network. Indeed, strict injectivity cannot be guaranteed but one can roughly construct an injective network using the technique mentioned in [1]. \n4. What advantages of this strict injectivity can provide in practice? In the ablation studies, it seems that the non-injective NN can sometimes give even better performance compared to injective ASWD? Any real applications that actually require this strict injectivity?\n5. Despite the ASWD gives better performance compared to other baselines, the FID score is still quite low compared to the current state-of-the-art method. So what real advantages of this SWD can bring?\n\nMinor:\n1. Is SRT a special case of GRT in general? If so, I wonder about the expressiveness of SRT. Any transformations that SRT cannot represents but GRT can?\n\n\n\n[1] Li, Chun-Liang, et al. \"Mmd gan: Towards deeper understanding of moment matching network.\" arXiv preprint arXiv:1705.08584 (2017).\n\n---\nI appreciate the author's response. They managed to address most of my concerns. However, I agree with reviewer HUCr about the gap between the theory and the definition of $\\pmb{g}$. The search space of $\\pmb{g}$ and the optimal $\\pmb{g}^*$ should be re-defined to address this. Anyway, I agree that this work should have some potentials, I will raise my score. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "For limitations, see Main Review. \n\nAppendix L (societal impacts) addresses the potential negative social impact of this work. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "ZDmeX0e55-d", "reviewer": "Reviewer_xvVx", "summary": "The paper develops a computationally efficient metric between two probability distributions, which is called Augmented Sliced Wasserstein Distances (ASWD). The proposed metric applied neural networks to map two distributions from $\\mathbb{R}^d$ into high dimensional non-linear manifolds, then the mapped distributions are compared using the sliced Wasserstein distance. It is shown in the paper that ASWD is a well-defined metric, and an algorithm for empirically evaluating ASWD is provided. The empirical applicability of ASWD is demonstrated by numerical experiments including gradient flow, generative modeling, sliced Wasserstein autoencoder, image color transferring and \nsliced Wasserstein barycenter.", "questions": "", "limitations": "", "rating": 7, "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": "In general the paper is well-written, the proof is rigorous, and the numerical result is convincing. \n\n**major concerns:**\n\n1. How to select the regularization parameter $\\lambda$ is not discussed in the paper. Though it is discussed in Remark 5 that $\\lambda$ adjusts the introduced non-linearity and $\\lambda$ can not be too small (the output $g_{\\omega}$ might explode) or too large (totally kills the nonlinear term), how to select the lambda in practice is not clear to the readers. In appendix F, the authors report: empirical errors in the experiment are found to be not sensitive to the choice of $\\lambda$ in a candidate set of {0.01, 0.05, 0.1, 0.5}: does it means that the candidate set is not rich enough so that one can not observe variations?\n\n2. Empirically, the neural network $g$ is selected as $g^\\ast = argmax(\\mathcal L(\\mu,\\nu;g,\\lambda))$ instead of $g^\\ast = argmax(\\mathrm{ASWD}(\\mu,\\nu;g))$. With such $g^\\ast$, is the result in Corollary 1.1 still valid? \n\n**Minor suggestions:**\n\n1. Line 81: Is it necessary to assume the existence of PDFs for $\\mu$ and $\\nu$? Many interesting data-driven applications does not satisfy this assumptions. For instance, the empirical measure of a finite observations is not absolutely continuous w.r.t Lebesgue measure.\n\n2. Line 94: $\\delta(\\cdot)$ is the **1-dimensional** Dirac delta function.\n\n3. Line 434: need to restrict to measurable $\\mathcal X$, otherwise the integral is not well defined.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The author discussed societal impact in appendix L. An additional limitation of this work could be:\n\n- As ASWD applied neural networks to map the data, the resulting metrics might be hard to interpret.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "32xOsmrJYK_", "reviewer": "Reviewer_HUCr", "summary": "This paper introduces _Augmented Sliced Wasserstein Distances_ (ASWD) which consist, in a nutshell, in computing the usual SWD between pushed measures: $\\mathrm{ASWD}(\\mu,\\nu; g) = \\mathrm{SWD}(g \\circ \\mu, g \\circ \\nu)$ (note: the pushforward operator does not render properly, so I will use $\\circ$ instead). The idea is then to maximize on $g$ (under constraint/regularization) to obtain a pushforward map $g$ that provides the most discriminative projection directions. Authors provide an exhaustive condition on $g$ (namely, injectivity) to make the ASWD a metric between probability measures. Their approach is then showcased on a various set of benchmark experiments in computational optimal transport (gradient flows, generative models, auto-encoder, barycenters, color transfert...).", "questions": "", "limitations": "", "rating": 7, "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": "[Edit: increased grade 6-->7 after rebuttal]\n\n\nQuality:\nThe paper is overall interesting and supported by a large number of convincing experiments. Some of its theoretical aspects can be improved though:\n- When defining the ASWD, the map $g$ should be chosen to ensure that $g \\circ \\mu, \\nu$ still have finite $k$-th moments. \n- Corollary 1.1, in its current state, seems wrong/incomplete. It is not clear why the argmax should be reached. If my understanding is correct, it is actually likely that the supremum (even when restricted on bounded injective maps) is infinite (each instance is finite, as stated in line 496, but the supremum may not). For instance, let $\\mu = \\delta_x, \\nu = \\delta_y$ with $x,y \\in \\mathbb{R}^d$; then $\\mathrm{ASWD}(\\mu,\\nu; g) = \\int_\\theta |\\braket{\\theta, g(x) - g(y)}| \\mathrm{d} \\theta$. Taking $g = g_\\omega \\simeq \\omega \\arctan( \\cdot )$ may lead to arbitrarily large values (while such $g_\\omega$ is injective and bounded, at fixed $\\omega$). This is actually the point of line 199 (regularization in the NN); in my opinion, this should be the optimization problem used to define the ASWD metric (to be compared with the Vaniila-ASWD if one does not optimize on $g$), and it would be nice to have a proof that this supremum is reached, perhaps a statement about the (non-)convexity(?) of the problem, etc., as it is actually at the core of the work. \n\nClarity: \nThe paper is well-written and easy to read overall.\n\nOriginality/Significance:\nThis paper heavily builds on previous works (in particular those of Kolouri et al.) which it properly credits. It provides simple yet seemingly efficient variations that improve on previous results. A quick look at the code and at the variety of experiments run also suggest that, once made publicly available, these tools may be useful of the computational OT community.\n\n\nMinor remarks and suggestions:\n- The paper may benefits from being written in a measure-based formalism; the current definitions suggest that $\\mu,\\nu$ should have densities $p_\\mu, p_\\nu$ and thus that discrete measures are not encompassed, which is obviously not the case. \n- In Figure 14, it would be nice to include (Regularized) OT barycenters (available in PythonOptimalTransport for instance) to showcase the (potential) difference between sliced approaches and the \"classical one\". \n- (line 104) It may be worth briefly mentioning these \"certain conditions\" for the sake of completeness. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper mentions some of its limitations and potential negative social impact (which are quite general to the ML field); I did not identify specific impact this paper could have. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper received four reviews, with scores scattered around the acceptance threshold. All the reviewers basically evaluated the main idea of this paper interesting, and clarity of the presentation of the main idea positively, whereas evaluations on significance were somehow mixed. Most notably, the reviewers raised a significant number of criticisms on several aspects of the contents of this paper, and the authors seem to have responded to most of them with prospects of how they will revise the paper. Some comments were then followed by further discussion between the authors and the reviewers. All of these should be quite fruitful in further improving the paper, as evidenced by some reviewers raising their scores. At the same time, it suggests that the submitted version which we reviewed had a large room for significant improvements, and some reviewers, in the discussion among us, did show their concerns as to whether the planned revision will eventually yield a paper with sufficient quality. I therefore think that another round of reviewing should be necessary to judge the technical quality of this contribution after the planned extensive revision, so that I would not be able to recommend acceptance of this paper at this stage.

An additional minor point: Equation (39) in Appendix D seems incorrect, but equation (40) should be correct, as it has been proved in page 15, under the heading "Triangle inequality". Simply removing the right-hand side of equation (39) should fix this defect.
Decision: Reject