Augmented Sliced Wasserstein Distances
Authors: Xiongjie Chen
Yongxin Yang
Yunpeng Li
Venue: ICLR 2022 Poster (2022)

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.

Published as a conference paper at ICLR 2022
AUGMENTED SLICED WASSERSTEIN DISTANCES
Xiongjie Chen', Yongxin Yang" and Yunpeng Li'
Universitty of Surrey, University of Edinburgh, Huaweii Noah's Ark Lab { [xiongjie, chen, yunpeng. 1i 1 (suurry ac. uk, ?yongxin. "yongxin.yangeenn yangged,
ABSTRACT
MXXiv:2006.0881277 [cs.LG] 17 Mar M 2022
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 metries, called augmented sliced Wasserstein distances (ASWDs), constructed by first mapping samples to higher-dimensional hypersurfaces parame- terized by neural networks. It is derived from a key observation that (random) linear projections of samples residing 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 an injective neural network architecture, Numerical results demonstrate that the ASWD significantly outperforms other Wasserstein variants for both synthetic and real-world problems.
INTRODUCTION
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., 2019b; Cuturi and Doucet, 2014; Courty et al., 2016).
Despite its favorable properties, such as robustness to disjoint supports and numerical stability (Arjovsky 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 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 focus on alternative OT-based metrics that are computationally efficient and 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 al., 2018) 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 et 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- 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 dinensiinnii space has a closed-fom solution based on sorting. (Deshpande et al., 2019) extends the sliced Wasserstein distance the max-sliced Wasserstein distance (Max-SWD). by finding single projection direction with the maximal distance between projected samples. The subspace robust Wasserstein distance extends the idea of slicing projecting distributions linear subspaces (Paty and

---

Published as a conference paper at ICLR 2022
Projected sample
sample
from Sample from
ASWD
SWD
80
ASWD
SWD
-60
-60
Frequency
Frequency
40
20
20
x
60
Distance between 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 embedded in the ASWD learns data patterns (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 evident in Figure (c)), as almost all random projection directions a -00-dimensional space lead to significant distances between 1-dimensional projections. In contrast, random linear mappings in the SWD often produce closer 1 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 (d)), leading to low projection efficiency in high- nigh-dimensional spaces.
Cuturi, 2019). However, the linear nature of these projections usually leads to low projection efficiency of the resulted metries in nigh-ddmennional spaces (Deshpande et al., 2019; Kolouri et al., 2019a).
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. 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 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 a task mask-dependent problem and requires domain knowledge, and the impact on performance from different defining functions is still unclear.
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 metrie; (iii) We provide a mechanism in which the hypersurface where high-dimensionall 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 aata-adaptive, i.c. 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- real-world 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. Related work are discussed in Section 4. Numerical experiment results are presented and discussed in Section 5. We conclude the paper in Section 6.
2

---

Published as a conference paper at ICLR 2022
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 Pk: (S2) be a set of Borel probability measures with finite k-th moment on Polish metric space (S2,d) (Villani, 2008). Given two probability measures f4, v Px (S2), the
Wasserstein distance of order k E [1,+ between js and v is defined as:
inf WE(ne)) = (ET ((
where df is the cost function, V(, represents the set all transportation plans 7, i.e. joint distributions whose marginals are μa and v, respectively.
While the Wasserstein distance is generally intractable for nigh-dimensional distributions, there are several favorable cases where the optimal transport problem can be efficiently solved. If μ and are continuous one-dimensional measures defined on a linear space equipped with the L* norm, the Wasserstein distance between y and has closed-fore foommmffoommoooooooo solution (Peyré and Cuturi, 2019):
Wu(uuV) (1)
(2)
where Fyr1 and Fr1 are inverse cumulative distribution functions (CDFs) of μu and v, respectively. In practice, Wasserstein distances Wk( D) between one-dimensional empirical distributions ū= N. > Oz,, and ū ū= > Sy., can be computed by sorting one-dimensional samples from empirical distributions, and evaluating the distances between sorted samples (Kolouri al., 2019b);
Wx(ūi,)=
[1. -M,tm|
where N the number of samples, Ig [n] and [n[n] are the indices of sorted samples satisfying Fr,(n) [n-][n-]] and yr, S (r,[+++1]. respectielly
Radon transform and generalized Radon transform: The Radon transform (Radon, 1917) maps a function f(+) E (RR)) to the space of functions defined over spaces of hyperplanes in IR.. The Radon transform of f()) is defined by line integrals of f()) along all possible hyperplanes in RRd.
RT((10)= Jzd ( ((=)(((-(xr)))) (xr
(4)
where / E IR and 0 E Sd-1 represent the parameters of hyperplanes ERd (x,0) -1, 8(-) is the one-dimensional Dirac delta function, and (--) refers the Euclidean inner product.
By replacing the inner product (rr,0) in Equation (4) with (xx,0), 0). specific family of functions named as defining function in (Kolouri al., 2019a), the generalized Radon transform (GRT) (Beylkin, 1984) is defined as integrals of f-) along hypersurfaces defined by {æ Rd ((x,0)=t])
gf(t.)) f(=)0(t-
(5)
where / E R, E and Pg is compact set of all feasible 0, e.g. Ωg = - for ((x,0) -) (x,0). f In particular, a function ((x,0) defined on X (IR*) {(0) with C Rd is called a defining function of GRTs if it satisfies conditions H.1 - H.4 given in (Kolouri al., 2019a).
For probability measures μ (IRd), the Radon transform and the GRT can be employed as push-forward operators, and the generated push-forward measures Ry = G,, = # are defined as follows (Bonneel et al., 2015):
(t,0) 5(1 - ((0-(x.0))m dd,, Jgd
(6)
(t,0)) 8(t 5(t-B(x,00 3(xx0)))d
3

---

Published as a conference paper at ICLR 2022
Notably, the Radon transform is linear bijection (Helgason, 1980), and the sufficient conditions for GRTs to be bijective are provided in (Homan and Zhou, 2017).
Sliced Wasserstein distance and generalized sliced Wasserstein distance: By applying the Radon transform to μ and v/ to obtain multiple projections, the sliced Wasserstein distance (SWD) decomposes the -isesiinenional Wasserstein distance into multiple one-dimensional Wasserstein distances which can be efficiently evaluated (Bonneel et al., 2015). The k-SW between and y is defined by:
( ((,
:
SWDx
where the Radon transform R defined by Equation (6) is adopted as the measure push pusw-forward operator. The GSWD generalizes the idea of SWD by slicing distributions with hypersurfaces rather than hyperplanes (Kolouri et al., 2019a). The GSWD is defined as:
GSWD&(HLV) (/o (.)
where the GRT 9 defined by Equation (7) is used as the measure push-forward operator. From Equation (3), with L random projections and N samples, the SWD and GSWD between j and can be approximated by:
N
SWD&().U)
Krggin .t) (1
(10)
NL 1= N
t
(11)
NT (=1n=1
where 1' and 1y are sequences consisting of the indices of sorted samples which satisfy (T (n] 01) < (Frg((++1): 1), 0i). (Urg(n) 01) < (1r 0z) in the SWD, and ((<rg(=)) 01) S 3( [n) 0j). [(urg[m) 0t) S [[urg(n+1)..)) 0t) in the GSWD. The approximation error in estimating SWDs using Equation (10) is derived in (Nadjahi et al., 2020). is proved in (Bonnotte, 2013) that the SWD a valid distance metric. The GSWD is valid metric except for its neural network variant (Kolouri et al., 2019a).
3 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
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 one-dimensional space. However, it is not straightforward functions (xx,0) for the GRT, since one needs to first check if (z2,0) () satisfies the conditions to be a defining function (Kolouri al., 2019a; 19; Beylkin, 1984), and then whether the corresponding GRT is bijective or not (Homan and Zhou, 2017). In practice, the assumption of the transform can be relaxed, as Theorem shows that as long as the transform is injective, the corresponding ASWD metric is valid distance metric.
design defining
To help us define the augmented sliced Wasserstein distance, we first introduce the spatial Radon transform which includes the Radon transform and the polynomial GRT special cases (See Remark 5). Definition 1. Given a measurable injective mapping g(-): IRdd d - Redo and a function f(() E T.- (TRd), the spatial Radon transform of f(-) is defined as
H.f (l.(t.)g)=
(12)
where LE R. and 0 Sdo--1 are the parameters of hypersurfaces ERE (g(x:)0) =t}.
4

---

Published as a conference paper at ICLR 2022
Similar to the Radon transform and the GRT, the spatial Radon transform also be used to generate push-forwar measure H₂ = H## for μ Ps (IR)) as in Equations (6) and (7):
Hy (t.00.)))
((t-(g(z),0)
(13)
Remark 1. Note that the spatial Radon transform can be interpreted as applying the vanilla Radon transform to Pg. where Bg refers to the push-forwer measure 94 ie given a measurable injective mapping g() (.):: :IR. + IRdo, the spatial Radon transform defined by Equation (13) can rewritten as:
H (t.t,.:g)=
((--)((tt((((((()) (g)))
E
| 8(t b(t- (1,0(
=R,,(t.0)
(14)
Hence the spatial Radon transform inherits the theoretical properties of the Radon transform and incorporates nonlinear projections through 9()).
In what follows, we use v to denote probability measures μ.ve (a(RR) that satisfy =(²²)= for VO'' CRd
Lemma 1. Given an injective mapping g(-): -Rdd and two probability measures μ.v E P(R.),, for all & R and 0 E gdo-1, (u, if and only fu i.e. the spatial Radon transform an injection
Ps (T.TR.). Moreover, the spatial Radon transform is injection Ps (TR.) and only if the mapping g(.) is an injection.
See Appendix A for the proof of Lemma
We now introduce the augmented sliced Wasserstein distance, by utilizing the spatial Radon transform the measure push-forward operator:
Definition 2. Given two probability measures μ.ve (IR())
( and an injective mapping g( IRd the augmented sliced Wasserstein distance (ASWD) of order e (1,+000) is defined as:
Rd
ASWD&(p.Vg)>
W (Hz
(15)
where E gde-1, W& is the Wasserstein distance defined by Equation (I), and H. refers to the spatial Radon transform defined by Equation (13).
Remark 2. Following the connection between the spatial Radon transform and the vanilla Radon transform as shown in Equation (14), the ASWD can be rewritten
ASWD&(H,1g) = ( ( W (Ra, (-0))
(16)
where jyy and iy are probability measures on Rd which satisfy g(x) Bog for x and g(y) Dg for y-uv.
Theorem 1. The augmented sliced Wasserstein distance (ASWD) order E [1,+ defined by +R RRd is a metric on P& (RR)) ) if and only if if.)()) is injective.
Equation (15) with mapping g(-): IR -
The proof of Theorem is provided in Appendix C. Theorem shows that the ASWD is metric given a fixed injective mapping 9(-). In practical applications, the mapping g(-) needs to be optimized to project samples onto discriminating hypersurfaces. We show Corollary 1.1 that the ASWD between and v with the optimized g(-) is also metric under mild conditions.
Corollary 1.1. The augmented sliced Wasserstein distance (ASWD) of order k E [1,+00) between two probability measures μ.v Pa (Rd) defined by Equation (15) with the optimal mapping
g" (-) argmax (ASD -Lp,pxx:g)}
(17)
metric on Ps (R.),, where T(u,v,A:g) X A (1) + 18) (
5

---

Published as a conference paper at ICLR 2022
The proof of Corollary 1.1 is provided in Appendix D.
Remark 3. Corollary 1.1 shows that given measures M1 Pu PR(R). the triangle inequality holds for the ASWD when g()) is optimized for each pair of measures, as shown in Appendix D. It is worth noting that >> 1 is sufficient conditiin for the ASWD to be metric- further discussed in Remark 6, 0< A 1 can also lead to finite g[((z))))))))))((()))) (0(z)ll2 in various scenarios, resulting in valid metrics. The discussion on the impact of X the performance of the ASWD in practice can be found in Appendix G.2.
3.2 NUMERICAL IMPLEMENTATION
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.
Injective neural networks: As stated in Lemma and Theorem 1, the injectivity of g()) 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 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:):
I- (x) [r)(w(())
(18)
It is trivial to show that 9u (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 the output of all preceding layers.
Optimization objective: We aim to slice distributions with maximally discriminating hypersurfaces between two distributions while avoiding the projected samples being arbitrarily large, so that the projected samples between the compared distributions are finite and most dissimilar regarding the ASWD, 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 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:
C(p.v:gWAA) ( W
(19)
where A >>0 O and the regularization term (lu,v.:g..) - = A(E [lla [B] + B (1) on the magnitude of the neural network's 's output is used, otherwise the projections may arbitrarily large.
Remark 4. The regularization coefficient À adjusts the introduced non-linearity in the evaluation of the ASWD by controlling the norm of ow( in Equation (18). In particular, when X - cc, the nonlinear term the (.) shrinks to O. The intrinsic dimension of the augmented space, i.e. the number of non-zero dimensions in the augmented space, is hence explicitly controlled by the flexible choice of (.())) and implicitly regularized by Lp..P,s..w).
By plugging the optimized x-) argmax(( V. 9w, A)) into Equation (15), we obtain the 9-.
empirical version of the ASWD. Pseudocode is provided in Appendix E.
RELATED WORK
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 it has been proved to be valid distance metric if and only if they adopt injective GRTs, which only include the circular functions and a 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 the aforementioned limited class of candidates. In addition, the

---

Published as a conference paper at ICLR 2022
selection of defining function is usually task-dependent and needs domain knowledge, and the impact on performance from different defining functions is still unclear.
To tackle those limitations, (Kolouri et al., 2019a) proposed the GSWD-NN, which directly takes the outputs of a 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 number of nodes in the neural network's '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 GSWD-NN, as the projection results are determined by the inputs and weights of the neural network. 3) The GSWD-ND is a pseudo-metric since uses a vanilla neural network, rather than Radon transform 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.
Another notable variant of the SWD is the distributional sliced Wasserstein distance (DSWD) (Nguyen et 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 to the distributional generalized sliced Wasserstein distance (DGSWD) (Nguyen 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.c. they are still not able to produce nonlinear projections adaptively.
Our contribution differs from previous work in three ways: 1) The ASWD is data- data-adaptive, i.e. the hyper- surfaces where high- high-dimensional distributions are projected onto can 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 a 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- ing nonlinear projections into Radon transform either is restricted to only a few candidates of defining functions (GRTs) or breaks the framework of Radon transforms (neural networks in GSWD-NN), in contrast, the spatial Radon transform provides novel way of defining nonlinear Radon-type transforms.
5 EXPERIMENTS
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 tfte-of-the-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 et al., 2019b) using the ASWD (Appendix I), image color transferring (Appendix J) and sliced Wasserstein barycenters (Appendix K).
To examine the robustness of the ASWD, throughout the experiments, we adopt the injective network architecture given in Equation (18) and set &w to be a single fully-connected layer neural network whose output dimension equals its input dimension, with a ReLU layer as its activation function. The order l is set to be 2 in all experiments.
5.1 SLICED WASSERSTEIN FLOWS
We first consider the problem of evolving source distribution μ to a target distribution v by minimizing slice-based Wasserstein distances between μ and in the sliced Wasserstein flow task reported in (Kolouri et al., 2019a).
0:4= VSWD(,.))..
(20)
where b4 refers to the updated source distribution at each iteration f. The SWD in Equation (20) 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 in this task. The set of hyperparameter values used in this experiment can be found in Appendix F.I.
Code to expermmetettttiiiemmeeteetttt is available at: https: / / thub. com/x x ongji echen/ ASWI..
7

---

Published as a conference paper at ICLR 2022
SWD GSWD-Poly
WWDD-Circular
GSWD-NN
DSWD
ASWD
8 Gaussian
Knot
Togia (W2)
500
1000
1500
2000
500
1000
1500
2000
Iteration
Iteration
Moon
Swiss Roll
1.0
log Seculte (W2)
0.5
0.7
0.0
0.0
0.5
1000
2000
1000 1500
2000
Iteration
Iteration
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- Wasserstein distances over 50 runs. A more extensive set of experimental results can be found in Appendix G.1.
Figure 2:
Without loss of generality, we initialize Ni to be the standard normal distribution (0.I)). We repeat each experiment times and record the 2- 2-Wasserstein distance between μ and at every iteration. In Figure 2, plot the 2- W-Wassesstein distances between the source and target distributions 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 best performing distance metrics, including the ASWD, the DSWD, the SWD, the GSWD-NN 1, which directly generates projections through of degree 3, circular defining functions, out of the 12 distance metrics we compared.
function
layer MIP,, as well the GSWD with the polynomial
We observe from Figure 2 that the ASWD 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 compared distance metrics and 8 target distributions are included in Appendix G. The ASWD outperforms the compared state-of-the-ate sliced-based based Wasserstein distance metrics with 7 out of the target distributions except for the 25-Gaussian. This is achieved through the simple injective network architecture given in Equation (18) and one layer fully-con colle-connected 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 regularization coefficient X, the choice of the dimensionality de of the augmented space, and the optimization of hypersurfaces in the ASWD. Details can be found in Appendix G.2.
only leads smaller 2- Wassersttei distances, but
5.2 GENERATIVE MODELING
In this experiment, use sliced-based Wasserstein distances for generative modeling task described in (Nguyen et al., 2021). The task is to generate images using generative adversarial networks (GANs) either the CIFAR10 dataset (64x 64 resolution) (Krizhevsky,
(Goodfellow et al., 2014) trained 2009) or the CelebA dataset (64 x64 resolution) (Liu et al., 2015). Denote the hidden layer and the output layer of the discriminator by hi and De. and the generator by Ge, we train GAN models with the following objectives:
min SWD(h (pr). (G+
(21)
max Exzo, [log( Du (h (z)))) +Bz-p, [log(1 - De (ho (ho V.
(22)
where Pz is the prior of latent variable z and p, is the distribution of real data. The SWD in Equation (21) is replaced by the ASWD and other variants of the SWD compare their performance. The
8

---

Published as a conference paper at ICLR 2022
310
300
DSWD
DSWD
260
GSWD
250
GSWD
score
SWD
SWD
210
200
160
150
110
100
60
50
100
150
200
50
20
40
60
80
100
Epoch
Epoch
Figure FID scores of generative models trained with different metrics on CIFARIO (left) and CelebA (right) datasets with L = 1000 projections. The error bar represents the standard deviation of the FID scores the specified training epoch among 10 simulation runs.
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 CIFARIO 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 RAM, and RTX 6000 graphic card with 22GB memories.
Table 1:
ASWD
L
(Bonneel al.., 2015)
(Kolouri et al., 201
(Nguyen et al., 2021)
FID
(SV(t)
FID
FID
(s/it)
(
10
121.4.1..0
0.34
108.3.15.6
0.41
3.1
0.55
65.7.132
0.58
100
104.6:5.2
0.35
105 2.2
0.74
66.5 3.9
0.57
62.5
0.60
502
0.36
98.2#55.1
2.22
5,73
1.30
39,3=3.22
1.38
CELEBA
94.8:12.5
0.35
95.1::4.2
0.40
86.0
0.53
81.2=113
0.59
100
88.7.15.7
0.36
86.7.13.55
0.75
76.1.1.3. 3.5
0.55
73.2 2.6
0.61
1000
86.51:4.1
85.2:16.3
2.19
71.3::4.7
1.38
GSWD with the polynomial defining function and the DGSWD is not included this experiment due to its excessively high computational cost in high-dimensional space. The Fréchet Inception Distance (FID score) (Heusel et al., 2017) 7) 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 F.2.
We run 200 and 100 training epochs to train the GAN models on the CIFARIO and the CclebA dataset, respectively. Each experiment repeated for 10 times and results reported in Table 1. With the same number of projections and a similar computation cost, the ASWD leads to significantly improved FII) scores among all evaluated distances metrics both datasets, which implies that images generated with the ASWD are of higher qualities. Figure plots the FID scores recorded during the training process. The GAN model trained with the ASWD exhibits faster convergence it reaches smaller FID scores with fewer epochs. Randomly selected samples of generated images are presented in Appendix H.I.
6 CONCLUSION
We proposed a novel variant of the sliced Wasserstein distance, namely the augmented sliced Wasserstein distance (ASWD), which 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 a valid distance metric and presented its numerical implementation. We reported empirical performance of the ASWD over state-ofthhe-s sliced Wasserstein metrics 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 in a sliced Wasserstein 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 includes the topological properties the ASWD. We leave this topic as interesting future research direction.

---

Published as a conference paper at ICLR 2022
REFERENCES
J. Altschuler, J. Niles-Weed, and P. Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 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), pages 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 (ICMI.) pages 573-582, Long Beach, California, USA, 2019.
E. Bernton, P. E. Jacob, M. Gerber, and C. Robert. On aarameter with the wasserstein distance. Information and Inference: A Journal of the IMA, 8(4):657-676, 2019.
G. Beylkin, The inversion problem and applications of the generalized Radon transform. Communications on Pure and Applied Mathematicx, 37(5):579-599, 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.
N. Courty, R. Flamary, D. Tuia, and A. Rakotomamoniy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence
2016.
M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Proc. Advances in Neu- 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 Machine Learning (ICML), pages 685-693, Beijing, China, 2014.
Deshpande, Hu, R. Sun, A. Pyrros, et al. Max-sliced Wasserstein distance and its use for GANs. In Proc. IEEE Conference Computer Vision and Patlern Recognition (CVPR), pages 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 Reseurch (JMLR), 19()-590--44,, 2018.
L. Ehrenpreis. The universality of the Radon transform, chapter pages 299-363, Oxford University Press, Oxford, UK. 2003.
S. Ferradans, N. Papadakis, G. Peyré, and J, Aujol. Regularized discrete optimal transport. SIAM Journal on Imaging Sciences, 7(3): 113::1853-182, 2014.
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), pages 1574-1583, Okinawa, Japan, 2019.
I. Goodfellow, A. Pouget, M. Mirza, B. Xu, F. D. Warde, al. Generative adversarial nets. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 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), pages 5767-5777, Long Beach, California, USA, 2017.
S. 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), pages 6626-6637, Long Beach, California, USA, 2017.
10

---

Published as a conference paper at ICLR 2022
A. Homan and H. Zhou. Injectivity and stability for a generic class of generalized Radon transforms. The Journal of Geometric Analysis, 27(2):1115-159, 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), pages 4700-4708, Hawaïi, USA, 2017.
M. Karami, D. Schuurmans, J. Sohl-Dickstein, L. Dinh, and D. Duckworth. Invertible convolutional flow. In Proc. Advances Neural Information Processing Systems (NeurIPS), pages 5636 5636-5646, Vancouver, Canada, 2019.
S. Kolouri, K. Nadjahi, U. Simsekli, R. Badeau, and G. Rohde. Generalized sliced Wasserstein distances. Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 261-272, Vancouver, Canada, 2019a.
Kolouri, P. Pope, C. E. Martin, and K. Rohde. ternational Conference on Learning Representations (ICLR), New Orleans, Louisiana, USA, 2019b.
autoencoders. In Proc. In-
A. Krizhevsky. Learning multiple layers of features from tiny images. Tech Report, 2009.
Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-basee learning applied to document recognition. Proceedings of the IEEE, 86(1 (6(11)):228-253 2324, 1998.
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proc. IEEE International Conference on Computer Vision (ICCV), pages 3730-3738, Las Condes, Chile, 2015.
Muzellec and M. Cuturi. Subspace detours: Building transport plans are on subspace pro- jections. Proc. Advances in Neural Information Processing Systems (NeurIPS), 32:6917-6628 2019.
K. Nadjahi al. Statistical and topological properties of sliced probability divergences. Proc. Advances in Neural Information Processing Systems (NeurIPS), 33, 2020.
K. Nguyen, N. Ho, T. Pham, and H. Bui. Distributional sliced-Wasserse Wasserstein and applications to generative modeling. In Proc. International Conference on Learning Representations (ICLR), Vienna, Austria, 2021.
F. Paty and M. Cuturi. Subspace robust Wasserstein distances. In Proc. International Conference on Machine Learning (ICML), pages 5072-5081, Long Beach, California, USA, 2019.
G. Peyré and M. Cuturi. Computational optimal transport. Foundations and Trends@ in Machine Learning, 1(5-6):355-60%) 2019.
J. Radon. Uber die bestimmung funktionen durch ihre integralwerte laengs gewisser mannigfaltigkeiten. Ber. Verh. Saechs. Akad. Wiss. Leipzig Math. Phys. KI., 69:262, 1917.
D. Rezende and S. Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning (ICMI.), pages 1530-1538, Lille, France, 2015.
Salimans, H. Zhang, A. Radford, and D. Metaxas. Improving GANs using optimal transport. In Proc. International Conference on Learning Representations (ICLR), Vancouver, Canada, 2018.
C. Villani. Optimal Transport: old and new, volume 338. Berlin, Germany: Springer Science & Business Media, 2008

---

Published as a conference paper at ICLR 2022
APPENDIX A PROOF OF THE LEMMA 1
We prove that the spatial Radon transform defined with a measurable mapping g()): ):::: + Rd an injection on Ps (R.)) if and only if (-) is injective. In the following contents, use Pk (R:)) to denote of Borel probability measures with finite k-th moment on IR and f1 = fz is used to denote functions fi())X + I and f2(-): X + R that satisfy fi (x) = [2(x) for Vz X, and fi $22 is used to denote
functions (-): (xx + R and (2(+)= X + I that satisfy [x(x) 4 (x) for certain I E X. In addition,
for probability measures μ and V, we use / = to denote f..e/ Pss (IE5) that satisfy ((x²)= =00 (x) for
V and μ # v/ to denote U.VE E P& (R") that satisfy ((x) + v (22) for certain X' CR.
Proof. By using proof by contradiction, we first prove that if g() is injective, the corresponding spatial Radon transform is injective. If the spatial Radon transform defined with injective mapping g()) R + de is not injective, there exist μ,, V E PA / v, such that Hz (t. = Hi (t,6:g)
for Vt E R and V0 e gde
From Equation (14), for V R and v0 = the spatial Radon transform of μ can be written as:
(t,(t.0:))= R,,, (t, (t,0g)))) = Rv, (t.0),
(23)
(24)
where ü, and Pg refer to the push-forward measures J and yg respectively, From the assumption Hy (t = (t,0gg) and Equations (23) and (24), we know Rr, (4,0) = Re, (1,0) for VI E IR and v0 E -1, which implies i, = by since the Radon transform is injective.
Since g(-) is injective, for all measurable X S Rd E X' if and only if a - g(x) E e((2), which implies ( Plyx ((((()(())))) ()). Therefore,
/ dpg- = [ du,
(25)
x(x)
[ dig [x du.
(26)
((x)
Since By &gg from Equations (25) and (26) we have: [x S dv for Hence, for Rd:
J =0,,
(27)
which implies f = v, contradicting with the assumption l E Therefore, if H₂, = Hor FE In addition, from the definition of the spatial Radon transform in Equation (13), is trivial to show that yf v, (t,0g) = H. (:g).. Therefore, Hy = H. if and only if μ = v, i.e. the spatial Radon transform H defined with an injective mapping g(-) :R. + Rdo is injective.
We now prove that if the spatial Radon transform defined with mapping g(.) Rdd - Rdo is injective, g(-) must be injective. Again, we use proof by contradiction. If g() is not injective, there exist 3o E 2d such that zo yo and g(zo) g(yo). For Dirac measures μ1 and v1 defined by (x(x))={((xx0)
and v1 v)())= = 55 yo) dy. where CR, we know #n #11 as zo 7 yo.
We define variables x and ~ y. Then for variables -= g() ~ u2 and و = gg)) ~ U2, where H22 and i/2 are forward measures g# ua and I# it trivial to derive for VX C Rd
/
p2(9(>))=
8(4 ((
(28)
g(x))
v2(g()))) =
8(0 - 8 )dj,
(29)
9(y)
which implies = as -(x)))
From Equations (23), (24), (28) and (29), for vl. ER and VO -10
Hm, (t,0:g) = Ry, (t,0), = Ry,, (t,8)
Hv, (t,8:))
(30)
12

---

Published as a conference paper at ICLR 2022
contradicting with the assumption that the spatial Radon transform is injective. Therefore, if the spatial Radon transform is injective, g(-) must be injective. We conclude that the spatial Radon transform is injective if and only if the mapping g(-) is an injection on P& (IR...
13

---

Published as a conference paper at ICLR 2022
APPENDIX B SPECIAL CASES OF SPATIAL RADON TRANSFORMS
5emark The spatial Radon transform degenerates to the vanilla Radon transform when the mapping
g(-) is an identity mapping. In addition, the spatial Radon transform is equivalent to the polynomial GRT (Ehrenpreis, 2003) when g(æ) ( = (x "...,?"*h ), where a is multi-indice Or (7)i p e. E
satisfying |o:i| ンド 2j=1 ]i.. m. m is the degree the polynomial function, x = 1 Thi5 given
impul x= (a.) E ER& and da is the number of all possible multi-indices oz; that satisfies (a:|=10 m
We provide proof for the claim in Remark 5.
Proof. Given probability measure μs P(R.), the spatial Radon transform of js defined as:
(t.0.g) 5(t- ((t-()(z):0))d
(31)
R=
where R and E Sdo-1 are the parameters of hypersurfaces in Rd When the mapping g(-) is an identity mapping, i.e. g(x) = x for Vo E IR,, the spatial Radon transform degenerates to the vanilla Radon transform:
(1,(g,0:)) / 8(t ((1-(x,0)))) (z,0))d
(t,t.
(32)
For GRT's, (Ehrenpreis, 2003) provides a class of injective GRT's named polynomial GRT's by adopting homogeneous polynomial functions with odd degree m the defining function:
da
(u(t,0)
8(t
0. ddp,
s.t. a4/l=mmm -m,
(33)
where a; = (T%,1,--.....))) EN4 |a,|= Dnll=Ej= x II 1 This for x (I1 ERC da is the number of all possible multi-indices ai that satisfics aay| = m, and 0= (01...,(a. e Sd. In spatial Radon transform, for Vx Rd, when the mapping g(-) is defined as: g(r) = (x ......"'., )
(34)
the spatial Radon transform is equivalent to the polynomial GRT defined in Equation (33):
(t.(g.))))=
8(t
(()x))()))))dl
do
=/ 8(1
Σ 0,0 )du. 1=1
(35)
14

---

Published as a conference paper at ICLR 2022
APPENDIX C PROOF OF THEOREM 1
We provide a proof that the ASWD defined with a mapping (()): -R is metric on Pa(RR), if and only if (-) is injective, In what follows, we denote set of Borel probability measures with finite k-th moment on I by P& (RR)), and use u. Pa (R*) to refer to two probability measures defined on R
Proof. Symmetry: Since the k- balassersteiimnissttneeee is a metric thus symmetric (Villani, 2008):
W, (H W& (4
(36)
Therefore,
(1 W (H) (1.g))()))))))),
(H,:)))
=
(000- - (H d0 AV.DD(
Triangle inequality: Given an injective mapping g(+) : R + RR and probability measures
M1 (22, H3 E Ph (R.). since the k- Wasserstein distance satisfies the triangle inequality (Villani, 2008).
the following inequality holds:
(4 (Oo (1 $/2
(Lor (400)))))
( (n
1gh0 (. #
where the second inequality is due to the kinkowski in I (
Identity of indiscernibles: Since (H, = for PR(
have
ASWD&()4)g)) = S ( (HHu(,:0:)))
=0,
(37)
for VuE (.R).. Conversely, for VU PA if = 0, from the definition of the ASWD:
ASWDx(M (1
1 =0,
(38)
Due to the non- nog-negayyvyty of ke-th Wasserstein distance as it is metric on Ps: (IRe)) and the continuity of (<.))) on (R2) (Villani, 2008), W& (-,0gg))))))))))))))))))(()))) (-00g)) - holds for v0 E @dde-11 -1 if and only if (-,0gg) (-,0g))))) Again, given the spatial Radon transform is injective when g(.) is injective (see the proof in Appendix A),
= (-:0g) implics == // if g()) is injective.
In addition, if g()) ()) is not injective, the spatial Radon transform is not injective (see the proof in Appendix A), then 3/4, // E Pa (Rd), μ v such that H,, (-f+t)== = H. (-G0:g). which implies ASWD; (H,v:)) = for 1. Therefore, the ASWD satisfies the identity of indiscernibles if and only if g(-) is injective.
Non-negativity: The three axioms of a distance metric, i.c symmetry, triangle inequality, and identity of indiscemibles imply the non-negativity of the ASWD. Since the Wasserstein distance is non-negative,
15

---

Published as a conference paper at ICLR 2022
for for ME PR((". it can also be straightforwardly proved the ASWD between μ and v non-negative:
ASWD = (1 gág (H) t
(1 0'd00 =0.
(39)
Therefore, the ASWD is a metric on (R( if and only if (() is injective.
16

---

Published as a conference paper at ICLR 2022
APPENDIX D PROOF OF COROLLARY 1.1
We first introduce Lemma 2 to support the proof of Corollary 1.1. Lemma For À (1, +-0), the optimal mapping g* (+) defined by Equation (17) satisfies llg" (x)ll2 for VI E Ed ~ and Vx E 1R" ~M
Proof. Recall that Equation (16) the ASWD can be rewritten as: ASWD ASWDx(p, = (gg
SWD (
k - -( W (.00),,) (+0))d))))
(40)
where transformed variables = g(x) ~ Po for x ~ μ and & = g(v) ~ Dy for y ~ v, respectively. Combining the equation above with Equation (2):
(1 W (-0)) (.8))
ASWD,
(1 |FPotR (z) (2)-FPprll F-
Po440 (=)
(=))
( S. Ro( Ro(o,) (2)1) dzdf
t
(41)
where # denotes the push forward operator, Po:zER:: + (z,0) E R, and Ro (Pg)) #o Ro (v,)) = Pottig refer to one-dimensional measures obtained by slicing jg. bg with unit vector 0, (o and (o, are inverse cumulative distribution functions (CDFs) of Ro and R (o())) respectively.
By repeatedly applying the Minkowski's inequality to Equation (41), we obtain the following inequalities:
(ggr- ((F (noo,))
ASWD (p.:g))
dzd0)
160-11(11 (
$ (f, IFFoote (z)) " d0 a0)²
[6 / (Filaa) (2) (z))
*[ -oonn (Fnder dzd0 7 d2))
(42)
Let s = (2.6), then = FRO(AN)))) (s), de = dF (FRo(A))()):
[ PR (
| dFFdF
((,,0) di
(o(x))
(43)
where the last two equations are obtained through the definitions of the push-forward operators. Therefore, the following inequalities hold:
1 + +Ven (' (
ASWDx(),:1)) < (Mmdo doao) (z)
(2) dzdo
/ (((())))) (
[ (o(u),o)) /
< <Biag [19(())))) +BJ (199())
(44)
17

---

Published as a conference paper at ICLR 2022
Then obtain the following inequalities for the optimization objective:
ASWD - (x,,,
<(Bing [11g(>)]13 -  (1) - [I + +E [II(())
-(1 (B [1lg + (1
(45)
When we set X (1,+x)), if Rr TR ~ or the optimization objective approaches negative infinity, implying g(-) is not the optimal mapping. Therefore, by adopting Equation (17) as the optimization objective, the optimal mapping g*(+) satisfies ll9" (x)ll2 < co for Vir E Ed ~ μ and Vz E IR ~v.
such that [|g(1))ll -
[lg(y) oc.
Remark 6. It is worth noting that X> in Corollary 1.1 is sufficient condition for the supremum the optimization objective to be non-positive and (1))2)) < ou for Vo Rd ~P and Va E li Tgdda v. 0 0<XS 1 can also lead to finite [lg(z)]12 in various scenarios. Specifically, the upper bound of the optimization objective given in Equation (44) is obtained by applying:
[g(z), .)
(46)
where is the angle between 0 and g(x)). In high- sigh--iimensional spaces, Equation (46) gives very loose bound since digh-dimensional spaces the majority of sampled would be nearly orthogonal to g(x) and cos(a) ) is nearly zero with high probability (Kolouri et al., 2019a). Empirically we found that across all the experiment results, À in a candidate set of
all lead to finite g" (-)-
We now prove Corollary 1.1, i.c (I.ri() is a metric on (R where g-) is the optimal mapping defined by Equation (17) for XE (1,+00).
Proof. Symmetry: Since the k d--aassesteein distance is a metric thus symmetric (Villani, 2008):
W&
= Wt (-,00")))) ),H, Ha((-00))))
(47)
Therefore,
ASWD)("VVg")) (Tor) (H
( Wh (- ()
=ASWD((V.)gg))))
Triangle inequality: From Lemma 2, when X E (1, +o0), the optimal mapping g" (.) satisfies IIg" (z)ll2<< oo for Vz E and E ~ v, hence Av. is finite due to Equation (16). We then prove that (v, satisfies the triangle inequality.
Denote by gi. 9h. and g optimal mappings that result in the supremum of Equation (19) between Ni and (2g (4 and μ3. P2 and μ3, respectively, since (P1 ASWD, (H1.330)))) and ASWD& (p2.)):3)) are finite, the following equations snaloiiionnuatinnss hold:
ASWD& S (AA.D.((...::))))AAAWW.u.....
(48)
(sup///SWD(((urrsss)))=aaaaAAAA0 (12,433))
(49)
(50)
where the first inequality are from the metric property of the ASWD.
Identity of indiscernibles: Since W& = for tuu Pa (R()), have
A (=(())))0--) Wh (10g0" (1000000)) ():00)))))) =0,
(51)
1
18

---

Published as a conference paper at ICLR 2022
for VuE PR Conversely, for E (AA.E.R.RRR,,AAA.....Y... =0, from Equation (15):
t =0.
(
(Hp( H,HH) (=:0g"))
(52)
Due to the non non-negativity of k-th Wasserstein distance as it is a metric on Ps (1RR) and the continuity of (:-):)) on Pk (RR) (Villani, 2008), Wa (HH,((.,0:"))) (- =0 holds for v0 ESde-1, which implies H-
H-(g) for A0 Sd. 11, Therefore, given the spatial Radon transform is injective when g (-) is injective, (,((,000" = H implies ==
Non-negativity: Since the Wasserstein distance is non non-negative, for Vp v (R(RR), the ASWD defined with optimal mappings g(-) between and is also non-negative:
(1
ASWD:
t
gu 0* =0.
(53)
Therefore, the ASWD defined with the optimal mapping 9 (-) is non-negative, symmetric, and satisfies the triangle inequality and the identity of indiscernibles, i.e. the ASWD defined with optimal mappings g" (-) is also a metric.
19

---

Published as a conference paper at ICLR 2022
APPENDIX E PSEUDOCODE FOR THE EMPIRICAL VERSION OF THE ASWD
Algorithm 1 The augmented sliced Wasserstein distance. All of the for loops can be parallelized.
Require: Sets of samples {zn, E E {yn +
Require: Randomly initialized injective neural network (.):. -RR
Require: Number of projections L, hyperparameter X, learning rate c, number of iterations M; Initialize x 0.m 1;
2: while has not converged and wn S M do
3: Draw a set of samples (0:)/=1 from E8d--1
4: for n= to N do
5:
Compute 9w (zn) and (yn);
Calculate the regularization term LA t L, + A[( (10( ) (لاا) y²]: N
6:
7:
end for
8:
for for/=1to/ L do
9:
Compute ((),,01) (gw (zn)..t)) 0t) (wn) (1)).01) for each
10:
Sort ((z,,)) 81) and ((y,. 01) in ascending order s.t. (11r//1].0)) 01) S (20 and
B(yrg,1)) S ((1r(t+1).888 81);
11:
Calculate the ASWD: D D+ D (001)((()))))))))
12:
end for
C+ D-D-Lxx Lx;
14: Update by gradient ascent w - we WuLi
Reset D= =0, update m m++;; 1;
16: end while
17: Draw a set of samples {0: from e 1do-1;
18: for no N do
Compute Iu (xr.) and (yn): 20: end for
19:
21: for forl=10 L do
22:
Compute (1),, - (1n)),000)) ((yn, (wn),,)) for each mi;
23:
Sort 8(2,, 81) and 3(yn, fr) in ascending order s.t. (( 01) S (( 01) and
B(yr,(m) 01) S0Ur)tn++))
24:
Calculate the ASWD: D (
25:
end for
26:
Output: Augmented sliced Wasserstein distance D.
In Algorithm 1, Equation (17) is used as the optimization objective, where the regularization term L,, > (x) [x) + -Ej. |I gow (y) [b] is used. This particular choice of the regu- larization term facilitates the proofs to Lemma and subsequently to Corollary 1.1 with details in Ap- pendix D. fact, have also examined other types of regularization terms such as the Lg norm of the output of g()), and empirically they produce similar numerical results as the current regularization term.
20

---

Published as a conference paper at ICLR 2022
APPENDIX F EXPERIMENTAL SETUPS
F.1 HYPERPARAMETERS IN THE SLICED WASSERSTEIN FLOW EXPERIMENT
We randomly generate 500 samples both for target distributions and source distributions. We initialize the source distributions (Ao as standard normal distributions (V, where 7 is a identity matrix. We update source distributions using Adam optimizer, and the learning rate=0.002. For all methods, we set the order k 2. When testing the ASWD, the number of iterations M in Algorithm to to 10.
a 2-dimensional
In the sliced Wasserstein flow experiment the mapping g²(,) optimized by maximizing Equation (17) was found to be finite for all values of À in the set of (0.01, 0.05, 0.1, 0.5, 1, 10, 100), which is sufficient condition for the ASWD to be valid metric shown in the proof of corollary 1.1 provided in Appendix D. In addition, numerical results presented in Appendix G.2 indicate that empirical errors in the experiment not sensitive to the choice of À in the candidate set (0.01, 0.05, 0.1, 0.5). The reported results in the main paper are produced with À x= 0.1.
F.2 NETWORK ARCHITECTURE IN THE GENERATIVE MODELING EXPERIMENT
Denote a convolutional layer whose kernel size s with kernels by Convo s), and fully-connected layer whose input and output layer have and $2 neurons by FC(s1 x $2). The network structure used the generative modeling experiment is configured to be the same as described in (Nguyen et al., 2021):
hy (6: x C (40L
Conv (4 BatchNormallizatoon BaccaNlommiiaattoo---LaLLLLLLLL
Conv256 (4 4) BatchNNormalzzttoon - LeakyeeLU((
Convs12(4 BatchNorraliiatioon Tanh
Outpur (512 4x44
:o:Conv x Sigmoid Ourput +(1=+++100000
(1x1x1)
: C (Cnn//rannggssseeg(((x
BatchNormalization -eeLnnnn-eeee---CC-aannaaaaaa
(4x4) -
BatchNormalization ReLU CenLU--ConTraansoseeee 128 (4x4) (4xx4)
BatchNormalizationn -
- ReL.UU-CCo//aaappsssss BaddNNomallaaiinn-RRRLU----
ReLU
4)
BatchNormalization -> Con/Transpose + Tanh
Ouput (64x64x3)
Ousput (8192)- vector
Φ: FCF99112228192
We train the models with the Adam optimizer, and set the batch size to 512. Following the setup in (Nguyen al., 2021), the learning rate is set 0.0005 and beta=(0.5, 0.999) for both CIFARIO dataset and CelebA dataset. For all methods, we set the order & to 2. For the ASWD, the number of iterations M in Algorithm 1 is set 5. The hyperparameter À set 1.01 to guarantee that the ASWD being a valid metric and introduce slightly larger regularization of the optimization objective due to the small output values from the feature layer hig.
21

---

Published as a conference paper at ICLR 2022
APPENDIX G ADDITIONAL
RESULTS IN THE SLICED WASSERSTEIN FLOW EXPERIMENT
G.1
FULL EXPERIMENTAL RESULTS ON THE SLICED WASSERSTHIN EXPERIMENT
Figure 4 shows the full experimental results on the sliced Wasserstein flow experiment.
SWD
GSWD-Circular
Max-GSWD-Poly 5
GSWD-NN 1 MaX-GSWD-NN ASWD
GSWD-Poly 3 GSWD-Poly
Max-SWD
Max-GSWD-Circular DSWD
Max-G5WD-Poly
Gaussian
Knot
0.5
0.0
1.00
500
1000
1500
2000
1000
1500
2000
Iteration
Iteration
Moon
Swiss Roll
1.0
0.2
toga
0.0
0.0
-0.2
500
1000
1500
2000
500
1000 Iteration
1500
2000
Iteration
25 Gaussian
Rectangle
ZMM (MM)oroo
0.0
0.5
1.00
500
1000
1500
2000
1000
1500
Iteration
Iteration
Heart
Circle
1.0
0.5
10giotW2
0.0
0.0
-055
-1.0
1.0
500
1000
1500
2000
500
1000
1500
2000
Iteration
Iteration
Figure 4: Full experimental results on the sliced Wasserstein flow example. The first and third columns are target distributions. The second and fourth columns are log 2- 2-Wasserstein distances between the target distributions and the source distributions. The horizontal axis shows the number of training iterations. Solid lines and shaded areas represent the average values and 95% confidence intervals of log 2- Wasserstein distances over 50 runs.
22

---

Published as a conference paper at ICLR 2022
G.2 ABLATION STUDY
Impact of the injectivity and optimization of the mapping
In this ablation study, we compare ASWDs constructed by different mappings to GSWDs with different predefined defining functions, and investigate the effects of the optimization and injectivity of the adopted mapping (-) used in the ASWDs. In what follows, "ASWD-vanills is used to denote ASWDs that employ randomly initialized neural network (-) to parameterize the injective mapping (.) = (.)] i.e. the mapping (-) is not optimized in the ASWD-vanill and the results of ASWD- D-nailla reported in Figure are obtained by slicing with random hypersurfaces. Furthermore, the "ASWD-non-injectvve refers to ASWDs that do not use the injectivity trick, i.e. the mapping g. (.) is not guaranteed to be injective. In addition, the both setups in the "ASWD-vanilla" and "ASWD-nO mon-injeective" resulting in random non-ing non-injective mapping g. (.). The reported experiment results in this ablation study is calculated over 50 runs, and the neural network (-) is reinitialized randomly in each run.
-aaiiaa--oo--innneeeeee ioniieeeeveee adopts
From Figure 5, it can be observed that the ASWD-vanilla shows comparable performance to GSWDs defined by polynomial and circular defining functions, which implies GSWDs with predefined defining functions are as uninformative as slicing distributions with random hypersurfaces constructed by the ASWD. In GSWDs, the hypersurfaces are predefined and cannot be optimized since they are determined by the functional forms of the defining functions. On the contrary, we found that the optimization of hypersurfaces in the ASWD framework can help improve the performance of the slice-based Wasserstein distance. As in Figure 5, the ASWD and the ASWD-non non-injective present significantly better performance than methods that do not optimize their hypersurfaces (ASWD-vanill)) ASWD- sanilla-oon-ineectvve and GSWDs). In terms of the impact of the injectivity of the mapping Iu. in this experiment, the ASWD- -vaailll exhibits smaller 2- Wasserstein distances than the injective in all tested distributions, and the ASWD leads to more stable training than the ASWD-non: injective. Therefore, the injectivity the mapping gw does not only guarantee the ASWD to be a valid distance metric as proved in Section 3, but also better empirical performance in this experiment setup.
Impact of the regularization coefficient
We also evaluated the sensitivity of performance of the ASWD with respect to the regularization coefficient A. The ASWD is evaluated with different values of and compared with other basce Wasserstein metrics in this ablation study. The numerical results presented in Figure 6 indicates that different values of À in the candidate set (0.01, 0.05, 0. 0.5 lead similar performance of the ASWD, i. the performance of the ASWD is not sensitive to A. Additionally, the ASWDs with different values of A in the candidate set outperform the other evaluated slice-based Wasserstein metrics.
We have also evaluated the performance of the ASWD when the range of X is much larger than in the candidate set. Specifically, as presented in Figure 7, when À is set to be large values, e.g 10 or 100, the resulted ASWD leads decreased performance on par with the SWD, This is consistent with our expectation that excessive regularization will eliminate nonlinearity as discussed in Remark 2, leading to similar performance with the SWD.
In addition, the effect of the regularization term on the performance of investigated in this ablation study. The performance of the MaX-GSWD-NN and the Max-GSWD-NI trained with the regularization term used in the ASWD are compared in Figure 8. From the numerical results presented in Figure 8, the Max- Bax-GSWD-NN 1 with regularization leads to performance similar to the MaX-GSWD-NN I without regularization, implying that the performance gap between the ASWD and the MAX-GSSD-NN is not due to the introduction of the regularization term.
was also
Choice of injective mapping
We reported in Figure the performance of the ASWD defined with other types of injective mappings other than Equation (18). In particular, examined two invertible mappings, including the planar flow and radial flow (Rezende and Mohamed, 2015), as alternatives to the injective mapping defined by Equation (18). The numerical results presented in Figure 9 show that the ASWD defined with planar flow and radial flow produced better performance than GSWD variants in most setups. They exhibit slightly worse performance compared with the ASWD with injective mapping defined in Equation (18), possibly due to the additional restriction in invertible mapping imposed by the planar flow and radial flow.
Choice of the dimensionality do of the augmented space
23

---

Published as a conference paper at ICLR 2022
To investigate how the dimensionality dg of the augmented space affects the performance of the ASWD, different choices of dg are employed in the ASWD. Specifically, the injective network architecture (x) [z, 0z (x)] E + Rd given in Equation (18) is adopted and ow is set to be single fully-connected neural networks whose output dimension equals (1, 2, 3, 4) times its input dimension, i.e dg (2d,3d, 3d,,4,,4dd) respectively, where d is the dimensionality of æ. 'he numerical results are presented in Figure 10, and it can be found that the ASWDs present similar results across different choices of dg. can also be observed in Figure 10 that the ASWDs with different choices of dg consistently produce better performance than the other evaluated slice-based Wasserstein metrics.
SWD
GSWD-Circular GSWD-NN
ASWD-vanilla
GSWD-Poly GSWD-Poly
ASWD--annllaa-non-njjeeiie ASWD-non-njeettvv
ASWD
Gaussian
Knot
0.5
0.0
193000
00.5
1.0
500
1000
1500
2000
500
1000
1500
2000
Iteration
Iteration
Moon
Swiss Roll
0.4
log1o
0.5
0.2
0.0
0.5
Iteration
Iteration
25 Gaussian
Rectangle
1.0
(2
0.5
-1
1.00
-1
500
1000
1500
2000
-1
500
1000
1500
2000
Iteration
Iteration
Heart
Circle
0.5
0.0
0.0
-0.5
1.0
1000
1500
2000
1000
1500
2000
Iteration
Iteration
Figure 5: Ablation study on the impact from injective neural networks and the optimization of hypersurfaces on the ASWD. ASWDs with different mappings are compared GSWDs with different defining functions. The first and third columns show target distributions. The second and fourth columns plot log 2- 2-Wasserstein distances between the target distributions and the source distributions. In the second and fourth columns, the horizontal axis shows the number of training iterations. Solid lines and shaded areas represent the average values and 95% confidence intervals of log 2- Wasserstein distances over 50 runs.
24

---

Published as a conference paper at ICLR 2022
SWD GSWD-Poly 3
GSWD-Circular DSWD
ASWD (A = 0.01) ASWD (A 0.05)
ASWD (A 0.1)
ASWD (1 = 0.5)
8 Gaussian
Knot
0.5
logietW22
logiotW2]
0.0
-1
1.0
-1
500
1000
1500
2000
500
1000
1500
2000
lteration
lteration
Moon
Swiss Roll
1.0
0.4
(W91
0.5
1og1001 ol
0.2
-2
0.0
0.0
2000
1500
lteration
Iteration
25 Gaussian
Rectangle
1.0
(ZM))160
0.5
0.0
0.5
1.0
11
500
1000
1500
2000
-1
500
1000
1500
2000
lteration
lteration
Heart
Circle
1.0
0.5
(091otW2)
0.5
10g10
0.0
0.0
-0.5
-1.0
1.0
lteration
Iteration
Figure 6: Ablation study on the impact the regularization coefficient X. The performance of the ASWDs with different values of À are compared with other slice-based Wasserstein metrics. The first and third columns show target distributions. The second and fourth columns plot log Wasserstein distances between the target distributions and the source distributions. In the second and fourth columns, horizontal axis shows the number training iterations. Solid lines and shaded areas represent the average values and 95% confidence intervals of log 2- Wasserstein distances over 50 runs.
25

---

Published as a conference paper at ICLR 2022
ASWD 100)
ASWD (1 10)
ASWD 0.1)
SWD
8 Gaussian
Knot
0.5
(og1o(W2)
1091o(W2)
0.0
0.5
2000
1000
2000
lteration
lteration
Moon
Swiss Roll
1.0
0.4
Logiol (ooo
0.2
0.0
0.5
0.2
0
0
500
1000
1500
2000
0
500
1000
1500
2000
Iteration
Iteration
25 Gaussian
Rectangle
1000
2000
1000
2000
lteration
Iteration
Heart
Circle
1.0
0.5
log
0.5
0.0
0.0
00.5
-1.0
1.0
500
1000
1500
2000
1000 1500 2000
Iteration
Iteration
Figure 7: Ablation study on the impact from large values of X. The performance of the ASWDs with large values of X, e.g 10 and 100, are compared with the SWD. The first and third columns show target distributions. The second and fourth columns plot log 2- Wasserstein distances between the target distributions and the source distributions. In the second and fourth columns, the horizontal axis shows the number of training iterations. Solid lines and shaded areas represent the average values and 95% confidence intervals of log 2- Wasserstein distances over 50 runs.
26

---

Published as a conference paper at ICLR 2022
ASWD Max-GSWD-NN 1
Max-GSWD-NN 1 Reg
GSWD-Circular SWD
DSWD
8 Gaussian
Knot
0.5
(ogio(W2)
0.0
O
-1
1.0
500
1000
1500
2000
500
1000
1500
2000
Iteration
Iteration
Moon
Swiss Roll
1.0
0.4
(0g10(W2)
0.5
0.0
0.0
-0.2
1000
2000
1000
2000
Iteration
Iteration
25 Gaussian
Rectangle
1.0
0.5
0
0.0
0.5
1.0
0
500
1000
1500
2000
500
1000
1500
2000
Iteration
Iteration
Heart
Circle
1.0
0.5
(WW)
0.5
Ogga Be (W(
0.0
0.55
-1.0
1.0
2000
1000
2000
Iteration
Iteration
Figure 8: Ablation study on the impact from the regularization term on the performance of the Max- GSW-NN. The first and third columns show target distributions. The second and fourth columns plot log Wasserstein distances between the target distributions and the source distributions. In the second and fourth columns, the horizontal axis shows the number of training iterations. Solid lines and shaded areas represent the average values and 95% confidence intervals of log 2- Wasserstein distances over 50 runs.
27

---

Published as a conference paper at ICLR 2022
SWD GSWD-Poly 3
GSWD-Circular DSWD
GSWD-NN ASWD
ASWD (with planar flow) ASWD (with radial flow)
Gaussian
Knot
0.5
0.0
0.5
2000
1000
lteration
Iteration
Moon
Swiss Roll
1.0
0.2
Logion
0.0
0.0
0.5
500
1000
1500
2000
500
1000
1500
2000
Iteration
lteration
25 Gaussian
Rectangle
1.0
(2M)
0.5
0.0
0.5
1.0
500
1000
1500
2000
500
1000
1500
2000
Iteration
lteration
Heart
Circle
1.0
0.5
WW2
0.5
0.0
0.0
log
-0.5
-1.0
1500 2000
1000
Iteration
Iteration
Figure Ablation study on the impact from the choice of injective networks. The performance of the ASWDs with different types of injective networks are compared with other slice-based Wasserstein metrics. The first and third columns show target distributions. The second and fourth columns plot log 2- Wasserstein distances between the target distributions and the source distributions. In the second and fourth columns, the horizontal axis shows the number of training iterations. Solid lines and shaded areas represent the average values and 95% confidence intervals of log 2- Wasserstein distances over 50 runs.
28

---

Published as a conference paper at ICLR 2022
ASWD (dg = d) ASWD (dg = 4 d)
ASWD (dg = 3 ASWD (dg = dd)
DSWD
GSWD-Poly 3 SWD
GSWD-Circular
8 Gaussian
Knot
0.5
0.0
0.5
2000
lteration
Iteration
Moon
Swiss Roll
1.0
2 1
10g10
0.2
0.0
0.0
0.5
500
1000
1500
2000
500
1000
1500
2000
lteration
Iteration
25 Gaussian
Rectangle
1.0
(ZM))160)
-2
500
1000
1500
2000
-1
1000
1500
2000
lteration
Iteration
Heart
Circle
1.0
0.5
0.5
0.0
0.0
-0.5
0.5
-1.0
1.0
500
1000
1500
2000
-1
500
1000
1500
2000
Iteration
Iteration
Figure 10: Ablation study on the choice of the dimensionality de of the augmented space. The perfor- mance of the ASWDs with different choices of dg are compared with other slice-based Wasserstein metrics. The first and third columns show target distributions. The second and fourth columns plot log 2- Wasserstein distances between the target distributions and the source distributions. In the second and fourth columns, the horizontal axis shows the number of training iterations. Solid lines and shaded areas represent the average values and 95% confidence intervals of log 2- Wasserstein distances over 50 runs.
29

---

Published as a conference paper at ICLR 2022
APPENDIX H ADDITIONAL
RESULTS IN THE GENERATIVE MODELING EXPERIMENT
H.1
SAMPLES OF GENERATED IMAGES OF CIFARIO AND CELEBA DATASETS
(a) CelebA (L= 10)
(b) CelebA (L 100)
(c) CelebA (L 1000)
(d) CIFARIO (L- = 10)
(e) CIFARIO (L 100)
(f) CIFARIO (L (L= 1000)
Figure 11: Visualized experimental results of the ASWD on CelebA and CIFARIO dataset with 10, 100, 1000 projections. The first row shows randomly selected samples of generated CelebA images, the second row shows randomly selected samples of generated CIFAR 10 images.
H.2 EXPERIMENT RESULTS ON MNIST DATASET
In the generative modelling experiment on the MNIST dataset, we train a generator by minimizing different blice-based Wasserstein metrics, including the ASWD, the DSWD, the GSWD (circular), and the SWD. Denote by G& the generator, the training objective of the experiment can be formulated (Bernton et al., 2019):
min Easpp...ppp SWD(2,G
(54)
where and are the prior of latent variable and the real data distribution, respectively. In other words, the SWD, or other slice-based Wasserstein metrics, can be considered as discriminator in this framework. By replacing the SWD with the ASWD, the DSWD, and the GSWD, we compare the perfor- mance learned generative models trained with different metrics. In this experiment, different methods are compared using different number of projections L= (10,1000). The 2- Wasserstein distance and the SWD between generated images and real images are used as metrics for evaluating performances of different generative models. The experiment results are presented in Figure 12. Quality of generated images and convergence rate: It can observed from Figure 12 that the ASWD outperforms all the other methods regarding both the 2- Wasserstein distance and the SWD between generated and real images. In particular, the generative model trained with the ASWD produces smaller Wasserstein distances within less iteration, which implies the generated images leads to higher convergence rates of generative models. In addition, the ASWD shows that is able to generate higher quality images than the SWD and the GSWD with 1000 projections using only less
of higher quality and the ASWD
30

---

Published as a conference paper at ICLR 2022
DSWD-10
DSWD-10
96
DSWD-1000
D5WD-1000
GSWD-10
28
GSWD-10
88
GSWD-1000
GSWD-1000
SWD-10
SWD-10
5WD-1000
60
W5(Oreal Ptake
ASWD-10
SW>(Presi Piate)
ASWD-10
ASWD-1000
20
ASWD-1000
72
64
58
12-
5000
10000
15000
20000
5000
10000
15000
20000
Iteration
Iteration
(a)
(b)
Figure 12: Visualized experimental results of different based Wasserstein metries the MNIST dataset with 10, 1000 projections. (a) Comparison between the SWD, the GSWD, the DSWD, and the ASWD using the 2- Wassersttei distance between fake and real images the evaluation metric. (b) Comparison between the SWD, the GSWD, the DSWD, and the ASWD using the SWD between fake and real images as the evaluation metric.
as 10 projections. other words, the ASWD has higher projection efficiency than the other dlce-bbsed Wasserstein metrics. The ASWD also has the smallest SWD distance shown in Figure 12. Although the SWD converges slightly faster than the ASWD in terms of the SWD between fake and real images, this is due to the training objective and the evaluated metric the same for the SWD. Randomly selected images generated by different blice Wasserstein metrics are presented in Figure 13.
Computation cost of the ASWD: The execution time per mini-batch (512 samples) of different methods are compared in Figure 14a. We evaluate the SWD by varying the number of projections in the set (10, 1000, 10000) and all the other methods in the set (10, 1000). We found that although the SWD requires much fewer computational time than the DSWD and the ASWD, the quality of generated data poor even when the number of projections L increases to 10000. The GSWD is also computationally efficient when using 10 projections, but it requires the highest execution time and generates the highest 2- Wasserstein distance among all compared methods when the number of projections increases 1000. The huge difference in the execution time of the GSWD with 10 and 1000 projections is due to the GSWD needs to calculate distance matrices of shape L, where N and L are the number of samples and projections respectively, which is more computationally expensive than calculating inner products when the number of projections increases. The DSWD requires similar computational time as the ASWD in this example, while the ASWD generates higher quality images in terms of Wasserstein distances.
Besides, have also evaluated the effect of batch size the computation cost of different slice-based Wasserstein metrics. Due to the out-of- -ottmo--mmory error caused by the excessively high computation cost of the GSWDs in high-dimensional space, the GSWD is not included in this comparison. Specifically, the ASWD, the DSWD, and the SWD are compared in this experiment, and the computation time of the evaluated methods with reported in Figure 14b. From the results presented in Figure 14b. it can be observed that, similar to the computational complexity of the DSWD and the SWD, the computational complexity of the ASWD empirically tends to scale in O (NlogN).
10, 1000 projections and N (3 214 2 21 samples are
31

---

Published as a conference paper at ICLR 2022
3
2
5
9
5
3
4
7
G
3
9
0
3
0
7
7
4
8
3
3
0
5
2
A 4 7
6
6
R9
2
7
9
1
6
9
2
8
3
8
0
9
3
4
9
7
4
9
0
9
1
3
3
7
3
9
3
9
5
0
9
7
1
4
A
/
8
9
0
0
0
0
3 3
7
9
2
7
8
4
5
8
8
6
3
0
6
6
7
0
6
5 7
7
7
3
2
4
4
0 5
3
7
/
0
7
3
0
5
0
7
2
9
9
5
R
3
O
7
3
7
5
7
0
0
J
9
5
2 5
14
3
0
6
6
3
4
8
6
O
2
2
ASWD-10 10
SWD-10
DSWD-10
-0
GSWD-10 10
5
J.
9
S
O
2
6
6
O
6
2
0
4
6
3
2
O
5
2
7
2
S
6
4
1
3
9
5
5
2
7
7
4
6
9
4
A
4
3
9
6 4 / 9 8
7
9
4
4
5
5
5 0
3
5
6
8
6
5
9
0
0
2
2
5
0
4
5
7
4
0
4
8
8
7 0 5 3
8
7
4
?
9
7
6
6
0
8 3
5
2
9
2
4
4
4
0
/
3
17
0 8
9
9
4
7
1 3
1
7
6
8
5
3
5
3
6
5
0
0
7
6
9
0
9
9
O
9
0
V g
A
9
5
7
7
ASWD-1000
SWD-1000
DSWD-1000
GSWD-1000
Figure Randomly selected samples generated by different metrics, and 1000 refer the number of projections.
ASAVD DSWD
ASWD-1000
ASWD-10
DSW0D-1000
DSWD-10
GSWD
SWD-1000
SWD-10
Computation time (second)
2-2
W2 pree
45
2-4
12
2.
2-
36
33
10
20
40
26
2
Execution time per minibatch ((s)
Batch
(a)
(b)
Figure 14: (a) The execution time of different methods and the 2- Wasserstein distance between the real images and the fake images generated by their corresponding models. Each dot of the curve of SWD corresponds to the performance of the SWD with the number of projections = {10, 1000, 10000), in sequence. Each dot of the other curves correspond to the performance of the other methods with the number of projections L= metrics, the horizontal axis is the number of samples from the compared distributions, and the horizontal axis is the averaged time consumption for calculating the distances. It can be found that the evaluated metrics tend to scale in ONNog.N)
[1110,10000 [10,1000), in sequence. (b) Computation cost calculating different
32

---

Published as a conference paper at ICLR 2022
APPENDIX 1 SLICED WASSERSTEIN AUTOENCODERS
We train an autoencoder using the framework proposed (Kolouri et al., 2019h), where an encoder and decoder are jointly trained by minimizing the following objective:
min BCE (CF 05.00
(55)
where ø is the encoder, & is the decoder, Ps is the prior distribution of latent variable, BCE(-.) is the binary cross entropy loss between reconstructed images and real images, and Li(:-) is the L1 loss between reconstructed images and real images. We train this model using different slice-based Wasserstein metrics, including the ASWD, the DSWD, the SWD, and the GSWD. Here we use the ring distribution as the prior distribution as shown in Figure 16. We report the binary cross entropy loss during test time and the Wasserstein distance between prior and the encoded latent variable ((x) in Figure 15. The dlce-based Wasserstein metrics used as the third term in Equation (55) is also recorded at each iteration and presented Figure in order to analyze the factors that causes the differences in the performance of models trained with different slice-based Wasserstein metrics.
can observed from the first two columns of Figure 15 that while the model trained with the ASWD, SWD, and DSWD lead to similar 2- Wassertein distance between the encoded latent variable distribution and the prior distribution, which implies that they have similar coverage of the prior distribution shown in Figure 16, the model trained with the ASWD converges slightly faster to smaller binary cross entropy loss than the others. As shown in Figure 15(b) and 15(c). since the obtained GSWDs trivially small compared with the other metrics, models trained with GSWDs only focuses on the reconstruction loss and therefore produce reconstructed images of higher qualities in terms of the binary cross entropy. However, although models trained with the GSWD polynomial and the GSWD circular lead to smaller binary cross entropy loss than the ASWD, their latent distributions present very different data structures from the specified prior distribution, which can be problematic in certain applications where the support of the latent distribution is required to be within particular range.
Some MNIST images randomly generated by SWAEs trained with different metrics are given in Figure 17.
B210
20
COWD
GOLD Fnjynemit1
CONDD-Moyremmr)
baseed Wasserstein Metrits
SAII
GOWIDC roler
ARNO
200
Cross Entrapy
-20
100
-30
-00
44
2000
2053
2015
1000
2500 2003
Iteration
Iteration
iteration
(a)
(b)
(C)
Figure 15: Convergence behavior of SWAES trained with different blsce-based Wasserstein metrics. (a) The 2- Wasserstein distance between the prior distribution p. and the distribution of encoded feature (x)). (b) The binary cross entropy loss between the reconstruction and real data. (c) The slice- based Wasserstein metric used to train the model.
33

---

Published as a conference paper at ICLR 2022
(a) ASWD latent space
(b) DSWD latent space
(c) SWD latent space
(d) GSWD (poly 3) latent space
(e) GSWD (circular) latent
(f) Prior distribution
Figure 16: Comparisons between the encoded latent space generated by different slice-based Wasser- stein metries.
9
8
9
0
9
9
9
2
3
9
8
5
3
5
4
8
7
9
2
5
/
9
9
9
7
9
9
5
9
9
8
8
7
0
9
0
2
4
8
2
/
4
2
6
9
2
9
9
9
6
8
5
9
5
9
9
9
9
2
9
8
2
2
0
4
7
9
2
4
9
8 2
0
3
4
3
3
0
7
1
9
2 2
7
0
/
5
9
5
9
2
2
7 6
0
8
8
2
O
9
2
6
O
6
5
8
5
6
8
2
0
5
6
2
6
3
0
5
9
2
9
2
7
3
(a)
ASWD samples
(b)
DSWD samples
(c)
SWD samples
3
4
5
0
9 3
0
0
9 3
6 7 8
2 8 0 0 7
0
3
9
7
2
6
8
3
9
7
9
7
9
9
9
9
6
O
5 9 8 8 7 6
0 8
0
2
2
5
0
9
3
6
9
/
9
6
6 4 3
7
5
2
9
9
9
3
9
9
8
6
9
/
0
1
2
0
5
8
2
1
7
6 2 0
3
6
9
6
3
0
3
8
3 3
5
9
0
7 4 6
9
/
8
2
3
0
3
5
5
O
6
3 9
9
2
4
8
0
5
(d) GSWD (poly 3) samples
(e) GSWD (cireular) samples
(f) MNIST samples
Figure 17: MNIST images randomly generated by SWAES trained with different metrics.
34

---

Published as a conference paper at ICLR 2022
APPENDIX COLOR TRANSFERRING
Color transferring can be formulated as an optimal transport problem (Bonneel et al., 2015; Radon, 1917). In this task, the color palette of source image is transferred to that of a target image, while keeping the content of source image unchanged. To achieve this, the optimal transport can be used to find the alignment between image pixels by calculating the optimal mapping of color palettes. In this experiment, instead of solving the optimal mapping in the original space, we first project the distribution onto one-dimensional spaces and average the alignment between one-dimensional samples as an approximation of the optimal mapping in the original space. After obtaining the approximation, we replace pixels of the souree image with the averaged corresponding pixels in the target image. To reduce the computational cost, utilize the approach proposed in (Muzellec and Cuturi, 2019), where the K-means algorithm used to cluster the pixels of both source and target images, and then implement color transfer for the quantized images whose pixels consist of the centers of 3000 clusters rather than the original source and target images.
We present the results of color transferring in Figure 18. It be observed that the ASWD and the DSWD produce sharper images than the SWD and the GSWD (polynomial), we conjecture that is because the ASWD and the DSWD generate better alignment of pixels. The GSWD (circular) tends to generate high contrast images, but sometimes produces images with low color diversity as in Figure 18(a).Th The Ma has the highest contrast among all methods, but this is due it only uses single projection to obtain the transport mapping, thus there is no need to average different pixels from the target image. A disadvantage of the Max-SWD that the transferred images generated by Max-SWD is not smooth enough and do not look realistic. The ASWD can generate smooth and realistic images than the SWD and the Max-SWD, even when the number of projections is as small as 10. Transferred images obtained by transferring colors from the source to target using standard optimal transport maps is also presented in Figure 18 for reference (Ferradans et al., 2014).
35

---

Published as a conference paper at ICLR 2022
Source
Target
SW n=100
SW n- 100
ASW n=10
ASW n=100
McxSW
DSW n=
DSW n=100
GSWicircular)
GSWccircular n=100
GSWipoly)
GSWrpoty) n=100
Standard map
(a)
Source
Target
SW n-10
SW n-00
ASW n-10
ASW n-100
MaxSW
DSW n-10
DSW n=100
GSWicirclaa n=10
GSWcroular n=100
GSWipol) n-10
GSWipoly) n-100
Standard OT
(b)
Source
Target
SW n-
SW n=100
ASW n-10
ASW n=100
MaxSW
DSW n-10
DSW m=100
CSWIcircula))nnl n=10
GSWccrcular) n=100
GSWipdly) n=10
GSWipoly) n=100
Standard OT
(c)
Source
Target
SW n=10
SW n=100
ASW n=10
ASW n=100
MaxSW
DSW n=10
DSW n=100
GSWicircular) n=10
GSWccircular n=100
GSWIpoly) n=10
GSWIpoly) 1=100
Standard OT map
(d)
Figure 18: Top rows are source images and target images, lower rows show transferred im- ages obtained by using different methods with different number of projections. Source and target images are from (Bonneel et al., 2015) and https: //i thub. com/chia56028/ Color- Transfer Transfer-between. between Images.
36

---

Published as a conference paper at ICLR 2022
APPENDIX K SLICED WASSERSTEIN BARYCENTER
Sliced Wasserstein distances can also be applied in the barycenter calculation and shape inter- polation (Bonneel al., 2015). Here we compare barycenters produced by different slice-based Wasserstein metrics, including the GSWD (circular and polynomial), the ASWD, the SWD, and the DSWD. Specifically, we compute barycenters of different shapes consisting of point clouds, as shown in Figure 19. Each object in Figure 19 corresponds to specific barycenter with different weights. The Wasserstein barycenters are also presented in Figure 19 for reference, which provide geometrically meaningful barycenters at the expense of significantly higher computational cost.
Formally, a sliced- Wasserstein barycenter of objects μ (MI,H2, "NN E Pa(RR)) assigned with is defined as:
weights w = 1W1,"02"""" WN ER)
Bar(u, argmin (E mc SWD(U.)))
(56)
In this experiment, set = and compute barycenters corresponding to different weights. The results are given in Figure 19.
From Figure 19, it can be observed that the ASWD produces similar barycenters that of the SWD, which are sharper than the DSWD, and more meaningful than the GSWD (polynomial). The flexibility of the injective neural networks g(-) and its optimization in the ASWD can be potentially combined with specific requirements in particular tasks to generate calibrated barycenters we leave this as a future research direction.
(a) ASWD barycenters
(b) GSWD (circular) barycenters
(c) GSWD (polynomial) barycenters
(d) DSWD barycenters
(e) SWD barycenters
(f) Wasserstein barycenter
Figure 19: Sliced Wasserstein barycenters generated by the ASWD, the GSWD, the DSWD, and the SWD, and the Wasserstein barycenter.
37


Reviews:
[{"review_id": "yiEdIh-SWLa", "reviewer": "Reviewer_ViRm", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "Strengths\n- The paper is well written and results are easy to follow.\n- The new metric is well justified with respect to existing literature, and many of its theoretical properties are analyzed.\n- The amount of numerical experiments and the variety of problems is very impressive!\n\nWeaknesses\n- Many of the points where mathematical formality is required is unfortunately omitted. For example at the beginning of page 4, I'm still wondering what are those \"certain conditions\". How large or small, or scalable is the approximation in (10) or (11)? What are the \"certain non-trivial requirements\" in Section 3.1. What are the \"certain conditions\" in Remark 1? What are the \"certain constraints \" in Section 3.2 Optimization objective?", "summary_of_the_paper": "This manuscript introduces the concept of augmented sliced Wasserstein distances. The main idea is to extend the sliced Wasserstein distance based on mapping samples to higher -dimensional hypersurfaces. The proposed distance is shown to be a metric. Moreover, given that the optimal choice of the nonlinear maps is rather computationally intensive to obtain, an approximation based on neural networks is proposed. Several experiments are shown where a better performance is obtained with respect to existing methods.", "summary_of_the_review": "Correctness: \n- I believe the statements are correct, but I did not checked ever single line of all the proofs. \n\nNovelty:\n- Many of the concepts are natural extensions of the sliced Wasserstein distances. However, there is sufficient novelty.\n\nEmpirical novelty:\n- A very complete numerical examples set is presented.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "nNj8RLJGnI", "reviewer": "Reviewer_1Q3z", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "- Overall, the paper is well written. It thoroughly  reviews the background and main works on SWD.  The rationale behind the proposed method is justified and aims to design an appropriate non-linear projection mapping that renders the induced sliced-Wasserstein distance (SWD) a valid metric. Theoretical analyses of conditions under which the metric property is preserved are provided.  In this regard the contribution is of interest and significative. Intensive empirical evaluations support these developments and show how the approach contributes efficient SWD computation.\n\n- Specifically,  instead of seeking a generalized Radon Transform (defined by non-linear defining functions that have to follow stringent constraints) or finding the distribution of (non-linear) random projections that maximize the expected 1D Wasserstein distances, the main idea is rather to learn a non-linear mapping function $g$ onto a high-dimension hyper-surfaces. The function $g$ acts as a measurable mapping that generates push-forward measures of the distributions to be compared. Hence, linear random projections onto 1D are performed  to obtain the related augmented SWD (ASWD) as expected 1D Wasserstein distances. The procedure results in defining a spatial Radon transform based on $g$ to compute the SWD which is clearly a clever idea. Therefore, establishing that ASWD is a metric simply derives from the injective property $g$. This facilitates the theoretical study of ASWD. \n\n- As an another interesting contribution, learning $g$ from data is achieved by maximizing a regularized version of ASWD. The range of the regularization parameter $\\lambda$ ensuring that the obtained $g$ leads to a valid metric  is stated. To guarantee the injection property, a neural network (NN) is considered and $g$ is get as the concatenation of the input and the output of the NN. The idea inspired from DenseNet is tricky and presents a clear advantage of avoiding learning complex NN. \n\n- Intensive empirical evaluations are conducted. They show impressive results compared to existing methods as SWD, distributional SWD (DSWD), generalized SWD (GSWD). In a synthetic problem of evolving source distribution to target one, ASWD achieves the smallest distance. The ASWD is also used to train a generative model for image generation. The reported performances (FID score) highlight the effectiveness of ASWD.\n\n- ASWD is shown to be a valid metric. What are the topological properties of ASWD compared to plain Wasserstein distance, SWD? Beyond empirical evaluations, can one theoretically relate ASWD to WD and SWD?\n\n- The main feature of ASWD is to project the samples onto a $d_\\theta$-high-dimensional hyper-surface using $g$.  Due to the design of $g$ in Eq. 19, $d_\\theta > d$. Does the choice of $d_\\theta$ impact the quality of the obtained ASWD?\n\n- How evolve the computation time according to the batch size? In Nguyen et al., 2021 the computation complexity tends to scale in $O(n \\lg n)$.\n\n- Unless mistaken, it seems that the FID scores in Table 1 for CIFAR 10 and CelebA are not competitive with the ones reported on Figure 3 in Nguyen et al, 2021. Indeed Nguyen et al. 2021 show smaller FID of order 60 for CIFAR 10 and 70 for CelebA. Can the authors explain the shift in the reported performances in Table 1? \n\n\nMinor comments\n- In remark 4, the expression 'rank of the augmented' is not clear and is undefined.\n\n*After rebuttal*\n\nI read authors rebuttal. They address most  raised points. Indeed they study the influence of $d_\\theta$ on  the computed ASWD as long as the computational complexity wich is similar to the one  of distributional SWD up to a constant. Also the mismatch in the reported FID scores is clarified and  fixed. One left out concern is  the topological property of ASWD.", "summary_of_the_paper": "The paper proposes a new slice-based approach to efficiently compute the Wasserstein distance between two distributions $\\nu$ and $\\mu$. The method termed ASWD (augmented sliced Wasserstein Distance) first  projects the samples from $\\nu$ and $\\mu$ onto a higher dimensional space using a non linear injective mapping function and then uses the classical random linear projections onto 1D to compute the sliced Wasserstein Distance. Overall, the procedure amounts to applied a spatial Radon Transform to perform the slicing. Theoretical results establish conditions  under which ASWD is a metric. A numerical algorithm is given along with the design of the injective mapping using NN. Empirical evaluations on simulation datasets and on generative modeling highlight the potential of the proposed method over existing approaches.", "summary_of_the_review": "The paper is well written and proposes a new sliced Wasserstein distance using the spatial Radon defined Transform defined as line integrals of a function along all hyperplances over a hyper-surface induced by a non-linear injective mapping. The new distance is supported by theoretical analysis and empirical experiments showing its effectiveness. Topological properties of the distance and some details on the empirical  evaluations remain to be clarified. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "eKgAPQdRJ1p", "reviewer": "Reviewer_RRsR", "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": "#### Strength: \nThe paper is well-written and the approach is mostly well presented. The authors introduce the ASWD and derive a set of theoretical properties and propose a numerical algorithm to approximate ASWD. They finally conduct several experiments to illustrate the advantages of ASWD over the state-of-the-art of sliced-based Wasserstein distance like SWD, GSWD, and DSWD.\nThe code is given, results seem reproducible.\n\n#### Weakness: \n\n- The authors argue that projecting the original data samples of the source and target distributions into higher-dimensional space leads to maximally discriminating projected samples, by constraining the $k$-th moments of the projected samples. In my opinion, the latter constraint could be well motivated and clarified (or at least well explained or discussed). \n\n- Since the projection is made by the same mapping $g$, It might be not guaranteed maximal discrimination over the projected samples. Intuitively, I think the mapping $g$ should be data-dependent and/or making a trade-off with different regularization parameters, i.e. the functional $L(\\mu, \\nu, \\lambda; g)$ becomes $L(\\mu, \\nu, \\lambda_1, \\lambda_2; g)$.\n\n- As was shown in Remark 3 that if the tuning parameter $\\lambda>1$ then ASWD is a proper metric, whereas the authors conducted their experiments with values $\\lambda<1$, because it leads to outperforming the other sliced-based Wasserstein metrics. In my opinion, this restricts the applications of ASWD, since one has to investigate a panel of tuning parameters $\\lambda<1$ to first check if ASWD is still a valid distance then if it gives good performance compared to the other SWDs. \n\n- The non-linearity of the injective mapping $g(\\cdot) = [Id, \\phi_\\omega(\\cdot)]$ seems a simple injective NN, allowing to gain efficient time complexity. Does it possible to construct $g$ by other injective NNs? It is probably that the fact for a small $lambda <1$ ASWD outperforms the sliced-based Wassrstein distance is linked to the choice of this special parameterization of the injective NN.\n\n- How is the dimensionality of the latent space $d_\\theta$ decided and is its impact? The paper does not seem to have included discussions on this seemingly important hyperparameter.\n\n#### Typo:\n- page 4: (Definition 1) \"Given an measurable\" -> \"Given a measurable\"", "summary_of_the_paper": "This paper introduces the augmented sliced Wasserstein distance (ASWD), a new variety of sliced Wasserstein distance (SWD), that allows comparing two probability distributions by combining a nonlinear embedding of the sample data points to a higher-dimensional space with a slicing scheme to calculate 1D Wasserstein distances between uniformly projection directions. The authors introduce the spatial Radon transform, which includes the standard Radon transform and the special case of polynomial generalized Random transform (introduced in Kolouri et al. 2019). They further prove that ASWD is a valid metric if and only if the mapping is injective. Several experiments are conducted on generative modelling (CIFAR10, CelebA,  MNIST, color transferring).", "summary_of_the_review": "The authors introduce the ASWD a new method of SWDs and derive a set of theoretical properties and propose a numerical algorithm to approximate ASWD. The ASWD is a proper metric if and only if the original data samples embedding in a higher-dimensional space through a nonlinear injective mapping. In my opinion, the optimizing ASWD with respect to the injective mapping is strongly depending on the choice of the tuning parameter which guarantees that ASWD is a valid metric if it is greater than $1$. However, in practice, it is better to conduct experiments with values less than $1$. This creates some \"broken pieces\" between theoretical fundings and experiments.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "PeGijZDvNpT", "reviewer": "Reviewer_6LWm", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "\nThis paper is overall well written, and ASWD is well motivated.\nThe proposed method seems to have some originality, though the technique of using non-linear functions to improve the efficiency of slicing has been proposed in GSWD-NN.\nThe author theoretically show that their proposed method is a valid distance metric even if they use neural networks for the non-linear functions.\nI did not check the theoretical results carefully, but they are correct as far as I can tell.\nThe significance of ASWD has been shown through various experiments such as flow, generalize model, barycenter, and color transfer.\n\nMy concern is about the difference between ASWD and GSWD-NN.\nTo my understanding, the main differences between them are\n(1) GSWD-NN is not a valid metric because they uses non-injective function, and (2) ASWD calculates SWD after the non-linear mappings, while max-GSWD-NN directly calculates 1-dimensional Wasserstein distance after the mapping.\n\nQuestions:\n- (a) As the authors claim, GSWD-NN is not a valid metric. However, if they use an injective function, would GSWD-NN be a valid metric? In other words, is the claim against the inappropriate choice of $g(x)$? Or is GSWD inherently difficult to combine with NNs?\n- (b) The second question is about the cause of the difference in performance. The experiment in Sec G.2 indicates that there is not much difference between ASWD and ASWD-non-injective in performances when optimizing NN. So is the difference in performance between ASWD and GSWD-NN due to the random projection (or some other reasons, e.g., regularization term $L_\\lambda$)?\n", "summary_of_the_paper": "This paper proposes a variant of sliced Wasserstein distance, named augmented sliced Wasserstein distance.\nASWD maps input data points to hypersurfaces using neural networks, then calculates SWD on the hypersurfaces.\nASWD alleviates the low efficiency problem of SWD for high-dimensional data.\nVarious tasks including flow, generative modeling, and barycenters show the advantage of AWSD against some existing methods.", "summary_of_the_review": "This paper is well written and well motivated. The proposed method seems to have some originality, though the technique of using non-linear functions has been proposed. Various experiments show its significance over existing methods.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a variant of sliced wasserstein distance , where the slicing operation is performed with a neural network. The resulting distance is studied and experiments on synthetic data and as cost in generative modeling are performed.

While the idea of the paper is not that novel, the work is overall well executed. Reviewers agreed that the paper is borderline weak accept. Accept as a poster.
Decision: Accept (Poster)