Original: A UGMENTED S LICED W ASSERSTEIN D ISTANCES
Revised: A UGMENTED S LICED W ASSERSTEIN D ISTANCES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: While theoretically appealing, the application of the Wasserstein distance to large-scale machine learning problems has been hampered by its prohibitive computational cost.
Revised: While theoretically appealing, the application of the Wasserstein distance to large-scale machine learning problems has been hampered by its prohibitive computational cost.

Original: The sliced Wasserstein distance and its variants improve the computational efficiency through random projection, yet they suffer from low projection efficiency because the majority of projections result in trivially small values.
Revised: The sliced Wasserstein distance and its variants improve the computational efficiency through random projection, yet they suffer from low projection efficiency because the majority of projections result in trivially small values.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We show that the hypersurfaces can be optimized by gradient ascent efficiently.
Revised: We show that the hypersurfaces can be optimized by gradient ascent efficiently.

Original: We provide the condition under which the
Revised: We provide the condition under which the

Original: ASWD is a valid metric and show that this can be obtained by an injective neural network architecture.
Revised: ASWD is a valid metric and show that this can be obtained by an injective neural network architecture.

Original: Numerical results demonstrate that the ASWD significantly outperforms other Wasserstein variants for both synthetic and real-world problems.
Revised: Numerical results demonstrate that the ASWD significantly outperforms other Wasserstein variants for both synthetic and real-world problems.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Comparing samples from two probability distributions is a fundamental problem in statistics and machine learning.
Revised: Comparing samples from two probability distributions is a fundamental problem in statistics and machine learning.

Original: 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.
Revised: 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.

Original: The
Revised: The

Original: 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; Lee et al., 2019; Cuturi & Doucet, 2014; Claici et al., 2018; Courty et al., 2016;
Revised: 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; Lee et al., 2019; Cuturi & Doucet, 2014; Claici et al., 2018; Courty et al., 2016;

Original: Shen et al., 2018; Patrini et al., 2018).
Revised: Shen et al., 2018; Patrini et al., 2018).

Original: Despite its favorable properties, such as robustness to disjoint supports and numerical stability (Arjovsky et al., 2017), the Wasserstein distance suffers from high computational complexity especially when the sample size is large.
Revised: Despite its favorable properties, such as robustness to disjoint supports and numerical stability (Arjovsky et al., 2017), the Wasserstein distance suffers from high computational complexity especially when the sample size is large.

Original: Besides, the Wasserstein distance itself is the result of an optimization problem — it is non-trivial to be integrated into an end-to-end training pipeline of deep neural networks, unless one can make the solver for the optimization problem differentiable.
Revised: Besides, the Wasserstein distance itself is the result of an optimization problem — it is non-trivial to be integrated into an end-to-end training pipeline of deep neural networks, unless one can make the solver for the optimization problem differentiable.

Original: Recent advances in computational optimal transport methods focus on alternative OT-based metrics that are computationally efficient and differentiably solvable (Peyr´e & Cuturi, 2019).
Revised: Recent advances in computational optimal transport methods focus on alternative OT-based metrics that are computationally efficient and differentiably solvable (Peyr´e & Cuturi, 2019).

Original: Entropy regularization is introduced in the Sinkhorn distance (Cuturi, 2013) and its variants (Altschuler et al., 2017; Dessein et al., 2018; Lin et al., 2019) to smooth the optimal transport problem; as a result, iterative matrix scaling algorithms can be applied to provide significantly faster solutions with improved sample complexity (Genevay et al., 2019).
Revised: Entropy regularization is introduced in the Sinkhorn distance (Cuturi, 2013) and its variants (Altschuler et al., 2017; Dessein et al., 2018; Lin et al., 2019) to smooth the optimal transport problem; as a result, iterative matrix scaling algorithms can be applied to provide significantly faster solutions with improved sample complexity (Genevay et al., 2019).

Original: An alternative approach is to approximate the Wasserstein distance through slicing , i.e. linearly projecting, the distributions to be compared.
Revised: An alternative approach is to approximate the Wasserstein distance through slicing , i.e. linearly projecting, the distributions to be compared.

Original: The sliced Wasserstein distance (SWD) (Bonneel et al., 2015) is defined as the expected value of Wasserstein distances between one-dimensional random projections of high-dimensional distributions.
Revised: The sliced Wasserstein distance (SWD) (Bonneel et al., 2015) is defined as the expected value of Wasserstein distances between one-dimensional random projections of high-dimensional distributions.

Original: The SWD shares similar theoretical properties with the
Revised: The SWD shares similar theoretical properties with the

Original: Wasserstein distance (Bonnotte, 2013) and is computationally efficient since the Wasserstein distance in one-dimensional space has a closed form solution based on sorting.
Revised: Wasserstein distance (Bonnotte, 2013) and is computationally efficient since the Wasserstein distance in one-dimensional space has a closed form solution based on sorting.

Original: Deshpande et al.
Revised: Deshpande et al.

Original: (2019) extends the sliced Wasserstein distance to the max-sliced Wasserstein distance (max-SWD), by finding a single projection direction with the maximal distance between projected samples.
Revised: (2019) extends the sliced Wasserstein distance to the max-sliced Wasserstein distance (max-SWD), by finding a single projection direction with the maximal distance between projected samples.

Original: In Nguyen et al.
Revised: In Nguyen et al.

Original: (2020), the distributional sliced Wasserstein distance (DSWD) finds a distribution of projections that maximizes the expected distances over these projections.
Revised: (2020), the distributional sliced Wasserstein distance (DSWD) finds a distribution of projections that maximizes the expected distances over these projections.

Original: The subspace robust Wasserstein distance extends the idea of slicing to projecting distributions on linear subspaces (Paty & Cuturi, 2019).
Revised: The subspace robust Wasserstein distance extends the idea of slicing to projecting distributions on linear subspaces (Paty & Cuturi, 2019).

Original: However, the linear nature
Revised: However, the linear nature

Original: [Figure]
Revised: [Figure]

Original: Y
Revised: Y

Original: [Figure]
Revised: [Figure]

Original: Z (
Revised: Z (

Original: [Figure]
Revised: [Figure]

Original: X
Revised: X

Original: Y
Revised: Y

Original: SWD
Revised: SWD

Original: Z (b)    80 Distance between 1-d projections
Revised: Z (b)    80 Distance between 1-d projections

Original: F r e q u e n c y (c)
Revised: F r e q u e n c y (c)

Original: ASWD    4 Distance between 1-d projections
Revised: ASWD    4 Distance between 1-d projections

Original: [Figure]
Revised: [Figure]

Original: F r e q u e n c y (d)
Revised: F r e q u e n c y (d)

Original: [Figure]
Revised: [Figure]

Original: network in the learns data patterns (in the X - Y plane) and produces well-separate projected values ( Z -axis) between distributions in a random projection direction.
Revised: network in the learns data patterns (in the X - Y plane) and produces well-separate projected values ( Z -axis) between distributions in a random projection direction.

Original: The high projection efficiency of the ASWD is evident in Figure 1(c), as almost all random projection directions in a 100-dimensional space lead to significant distances between 1-dimensional projections.
Revised: The high projection efficiency of the ASWD is evident in Figure 1(c), as almost all random projection directions in a 100-dimensional space lead to significant distances between 1-dimensional projections.

Original: In contrast, random linear mappings in the SWD often produce closer 1-d projections ( Z -axis)
Revised: In contrast, random linear mappings in the SWD often produce closer 1-d projections ( Z -axis)

Original: (Figure 1(b)); as a result, a large percentage of random projection directions in the 100 -d space result in trivially small distances (Figure 1(d)), leading to a low projection efficiency in high-dimensional spaces.
Revised: (Figure 1(b)); as a result, a large percentage of random projection directions in the 100 -d space result in trivially small distances (Figure 1(d)), leading to a low projection efficiency in high-dimensional spaces.

Original: of these projections usually leads to low projection efficiency of the resulted metrics in high-dimensional spaces (Deshpande et al., 2019; Liutkus et al., 2019; Kolouri et al., 2019a).
Revised: of these projections usually leads to low projection efficiency of the resulted metrics in high-dimensional spaces (Deshpande et al., 2019; Liutkus et al., 2019; Kolouri et al., 2019a).

Original: More recently, there are growing interests and evidences that slice-based Wasserstein distances with nonlinear projections can improve the projection efficiency, leading to a reduced number of projections needed to capture the structure of the data distribution (Kolouri et al., 2019a; Nguyen et al., 2020).
Revised: More recently, there are growing interests and evidences that slice-based Wasserstein distances with nonlinear projections can improve the projection efficiency, leading to a reduced number of projections needed to capture the structure of the data distribution (Kolouri et al., 2019a; Nguyen et al., 2020).

Original: (Kolouri et al., 2019a) extends the connection between the sliced Wasserstein distance and the Radon transform (Abraham et al., 2017) to define generalized sliced Wasserstein distances (GSWDs) by utilizing generalized Radon transforms (GRTs).
Revised: (Kolouri et al., 2019a) extends the connection between the sliced Wasserstein distance and the Radon transform (Abraham et al., 2017) to define generalized sliced Wasserstein distances (GSWDs) by utilizing generalized Radon transforms (GRTs).

Original: It is shown in (Kolouri et al., 2019a) that the GSWD is indeed a metric if and only if the adopted GRT is injective.
Revised: It is shown in (Kolouri et al., 2019a) that the GSWD is indeed a metric if and only if the adopted GRT is injective.

Original: Injective GRTs are also used to extend the DSWD to the distributional generalized sliced Wasserstein distance (DGSWD)
Revised: Injective GRTs are also used to extend the DSWD to the distributional generalized sliced Wasserstein distance (DGSWD)

Original: (Nguyen et al., 2020).
Revised: (Nguyen et al., 2020).

Original: However, both the GSWD and the DGSWD are restricted by the limited class of injective GRTs, which utilize the circular functions and a finite number of harmonic polynomial functions with odd degrees as their defining function (Kuchment, 2006; Ehrenpreis, 2003).
Revised: However, both the GSWD and the DGSWD are restricted by the limited class of injective GRTs, which utilize the circular functions and a finite number of harmonic polynomial functions with odd degrees as their defining function (Kuchment, 2006; Ehrenpreis, 2003).

Original: The results reported in (Kolouri et al., 2019a; Nguyen et al., 2020) show impressive performance from the GSWD and the DGSWD, yet they require one to specify a particular form of defining function from the aforementioned limited class of candidates.
Revised: The results reported in (Kolouri et al., 2019a; Nguyen et al., 2020) show impressive performance from the GSWD and the DGSWD, yet they require one to specify a particular form of defining function from the aforementioned limited class of candidates.

Original: However, the selection of defining function is usually task-dependent and needs domain knowledge.
Revised: However, the selection of defining function is usually task-dependent and needs domain knowledge.

Original: In addition, the impact on performance from different defining functions is still unclear.
Revised: In addition, the impact on performance from different defining functions is still unclear.

Original: One variant of the GSWD (Kolouri et al., 2019a) is the GSWD-NN, which generates projections directly with neural network outputs to remove the limitations of slicing distributions with predefined GRTs.
Revised: One variant of the GSWD (Kolouri et al., 2019a) is the GSWD-NN, which generates projections directly with neural network outputs to remove the limitations of slicing distributions with predefined GRTs.

Original: In the GSWD-NN, the number of projections, which equals the number of nodes in the neural network’s output layer, is fixed.
Revised: In the GSWD-NN, the number of projections, which equals the number of nodes in the neural network’s output layer, is fixed.

Original: Hence different neural networks are needed if one wants to change the number of projections.
Revised: Hence different neural networks are needed if one wants to change the number of projections.

Original: There is also no random projections involved in the resulted GSWD-NN, since the projection results are determined by the neural network’s weights.
Revised: There is also no random projections involved in the resulted GSWD-NN, since the projection results are determined by the neural network’s weights.

Original: Besides, the GSWD-NN is a pseudo-metric since it uses a vanilla neural network, rather than the Radon transform or GRTs, as its push-forward operator.
Revised: Besides, the GSWD-NN is a pseudo-metric since it uses a vanilla neural network, rather than the Radon transform or GRTs, as its push-forward operator.

Original: Therefore, the GSWD-NN does not fit into the theoretical framework of the GSWD and does not inherit its geometric properties.
Revised: Therefore, the GSWD-NN does not fit into the theoretical framework of the GSWD and does not inherit its geometric properties.

Original: In this paper, we present the augmented sliced Wasserstein distance (ASWD), a distance metric constructed by first mapping samples to hypersurfaces in an augmented space, which enables flexible nonlinear slicing of data distributions for improved projection efficiency (See Figure 1).
Revised: In this paper, we present the augmented sliced Wasserstein distance (ASWD), a distance metric constructed by first mapping samples to hypersurfaces in an augmented space, which enables flexible nonlinear slicing of data distributions for improved projection efficiency (See Figure 1).

Original: Our main contributions include: (i)
Revised: Our main contributions include: (i)

Original: 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 a valid distance metric; (iii) We provide a mechanism in which the hypersurface where high-dimensional distributions are projected onto can be optimized and show that the optimization of hypersurfaces can help slice-based Wasserstein distances improve their projection efficiency.
Revised: 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 a valid distance metric; (iii) We provide a mechanism in which the hypersurface where high-dimensional distributions are projected onto can be optimized and show that the optimization of hypersurfaces can help slice-based Wasserstein distances improve their projection efficiency.

Original: Hence, the ASWD is data-adaptive, i.e. the hypersurfaces can be learned from data.
Revised: Hence, the ASWD is data-adaptive, i.e. the hypersurfaces can be learned from data.

Original: This implies one does not need to specify a defining function from a limited class of candidates; (iv) We demonstrate superior performance of the ASWD in numerical experiments for both synthetic and real-world datasets.
Revised: This implies one does not need to specify a defining function from a limited class of candidates; (iv) We demonstrate superior performance of the ASWD in numerical experiments for both synthetic and real-world datasets.

Original: The remainder of the paper is organized as follows.
Revised: The remainder of the paper is organized as follows.

Original: Section 2 reviews the necessary background.
Revised: Section 2 reviews the necessary background.

Original: We present the proposed method and its numerical implementation in Section 3.
Revised: We present the proposed method and its numerical implementation in Section 3.

Original: Numerical experiment results are presented and discussed in Section 4.
Revised: Numerical experiment results are presented and discussed in Section 4.

Original: We conclude the paper in Section 5.
Revised: We conclude the paper in Section 5.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: In this section, we provide a brief review of concepts related to the proposed work, including the
Revised: In this section, we provide a brief review of concepts related to the proposed work, including the

Original: Wasserstein distance, (generalized) Radon transform and (generalized) sliced Wasserstein distances.
Revised: Wasserstein distance, (generalized) Radon transform and (generalized) sliced Wasserstein distances.

Original: Wasserstein distance: Let P k (Ω) be a set of Borel probability measures with finite k -th moment on a Polish metric space (Ω ,d ) (Villani, 2008).
Revised: Wasserstein distance: Let P k (Ω) be a set of Borel probability measures with finite k -th moment on a Polish metric space (Ω ,d ) (Villani, 2008).

Original: Given two probability measures µ , ν ∈ P k (Ω) , whose probability density functions (PDFs) are p µ and p ν , the Wasserstein distance of order k ∈ [1 , + ∞ ) between µ and ν is defined as:
Revised: Given two probability measures µ , ν ∈ P k (Ω) , whose probability density functions (PDFs) are p µ and p ν , the Wasserstein distance of order k ∈ [1 , + ∞ ) between µ and ν is defined as:

Original: [Equation]
Revised: [Equation]

Original: where d ( · , · ) k is the cost function, Γ( µ,ν ) represents the set of all transportation plans γ , i.e. joint distributions whose marginals are p µ and p ν , respectively.
Revised: where d ( · , · ) k is the cost function, Γ( µ,ν ) represents the set of all transportation plans γ , i.e. joint distributions whose marginals are p µ and p ν , respectively.

Original: With a slight abuse of notation, we interchangeably use W k ( µ,ν ) and W k ( p µ ,p ν ) .
Revised: With a slight abuse of notation, we interchangeably use W k ( µ,ν ) and W k ( p µ ,p ν ) .

Original: While the Wasserstein distance is generally intractable for high-dimensional distributions, there are several favorable cases where the optimal transport problem can be efficiently solved.
Revised: While the Wasserstein distance is generally intractable for high-dimensional distributions, there are several favorable cases where the optimal transport problem can be efficiently solved.

Original: In particular, if µ and ν are continuous one-dimensional measures, the Wasserstein distance between µ and ν has a closed form solution (Bonneel et al., 2015):
Revised: In particular, if µ and ν are continuous one-dimensional measures, the Wasserstein distance between µ and ν has a closed form solution (Bonneel et al., 2015):

Original: [Equation]
Revised: [Equation]

Original: where F − 1 µ and F − 1 ν are inverse cumulative distribution functions (CDFs) of µ and ν , respectively.
Revised: where F − 1 µ and F − 1 ν are inverse cumulative distribution functions (CDFs) of µ and ν , respectively.

Original: Radon transform and generalized Radon transform: The Radon transform (Radon, 1917) maps a function f ( · ) ∈ L 1 ( R d ) to the space of functions defined over spaces of lines in R d .
Revised: Radon transform and generalized Radon transform: The Radon transform (Radon, 1917) maps a function f ( · ) ∈ L 1 ( R d ) to the space of functions defined over spaces of lines in R d .

Original: The Radon transform of f ( · ) is defined by line integrals of f ( · ) along all possible hyperplanes in R d :
Revised: The Radon transform of f ( · ) is defined by line integrals of f ( · ) along all possible hyperplanes in R d :

Original: [Equation]
Revised: [Equation]

Original: where t ∈ R and θ ∈ S
Revised: where t ∈ R and θ ∈ S

Original: d − 1 represent the parameters of hyperplanes in R d , δ ( · ) is the Dirac delta function, and (cid:104)· , ·(cid:105) refers to the Euclidean inner product.
Revised: d − 1 represent the parameters of hyperplanes in R d , δ ( · ) is the Dirac delta function, and (cid:104)· , ·(cid:105) refers to the Euclidean inner product.

Original: By replacing the inner product (cid:104) x,θ (cid:105) in Equation (3) with β ( x,θ ) , a specific family of functions named as defining function in Kolouri et al. (2019a), the generalized Radon transform (GRT) (Beylkin, 1984) is defined as:
Revised: By replacing the inner product (cid:104) x,θ (cid:105) in Equation (3) with β ( x,θ ) , a specific family of functions named as defining function in Kolouri et al. (2019a), the generalized Radon transform (GRT) (Beylkin, 1984) is defined as:

Original: [Equation]
Revised: [Equation]

Original: (4) where t ∈ R , θ ∈ Ω θ while Ω θ is a compact set of all feasible θ , e.g. Ω θ = S
Revised: (4) where t ∈ R , θ ∈ Ω θ while Ω θ is a compact set of all feasible θ , e.g. Ω θ = S

Original: d − 1 for β ( x,θ )
Revised: d − 1 for β ( x,θ )

Original: = (cid:104) x,θ (cid:105)
Revised: = (cid:104) x,θ (cid:105)

Original: (Kolouri et al., 2019a).
Revised: (Kolouri et al., 2019a).

Original: In practice, we can empirically approximate the Radon transform and the GRT of a probability density function p µ via:
Revised: In practice, we can empirically approximate the Radon transform and the GRT of a probability density function p µ via:

Original: [Equation]
Revised: [Equation]

Original: where x n ∼ p µ
Revised: where x n ∼ p µ

Original: and N is the number of samples.
Revised: and N is the number of samples.

Original: Notably, the Radon transform is a linear bijection (Helgason, 1980), and the GRT is a bijection if the defining function
Revised: Notably, the Radon transform is a linear bijection (Helgason, 1980), and the GRT is a bijection if the defining function

Original: β satisfies certain conditions (Beylkin, 1984).
Revised: β satisfies certain conditions (Beylkin, 1984).

Original: Sliced Wasserstein distance and generalized sliced Wasserstein distance:
Revised: Sliced Wasserstein distance and generalized sliced Wasserstein distance:

Original: By applying the Radon transform to p µ and p ν to obtain multiple projections, the sliced Wasserstein distance (SWD) decomposes the high-dimensional Wasserstein distance into multiple one-dimensional Wasserstein distances which can be efficiently evaluated (Bonneel et al., 2015).
Revised: By applying the Radon transform to p µ and p ν to obtain multiple projections, the sliced Wasserstein distance (SWD) decomposes the high-dimensional Wasserstein distance into multiple one-dimensional Wasserstein distances which can be efficiently evaluated (Bonneel et al., 2015).

Original: The k -SWD between µ and ν is defined by:
Revised: The k -SWD between µ and ν is defined by:

Original: [Equation]
Revised: [Equation]

Original: where the Radon transform R defined by Equation (3) is adopted as the measure push-forward operator.
Revised: where the Radon transform R defined by Equation (3) is adopted as the measure push-forward operator.

Original: The GSWD generalizes the idea of SWD by projecting distributions onto hypersurfaces rather than hyperplanes (Kolouri et al., 2019a).
Revised: The GSWD generalizes the idea of SWD by projecting distributions onto hypersurfaces rather than hyperplanes (Kolouri et al., 2019a).

Original: The GSWD is defined as:
Revised: The GSWD is defined as:

Original: [Equation]
Revised: [Equation]

Original: where the GRT G is used as the measure push-forward operator.
Revised: where the GRT G is used as the measure push-forward operator.

Original: The Wasserstein distances between one-dimensional distributions can be obtained by sorting projected samples and calculating the distance between sorted samples (Kolouri et al., 2019b): with L random projections, the SWD and GSWD between µ and ν can be approximated by:
Revised: The Wasserstein distances between one-dimensional distributions can be obtained by sorting projected samples and calculating the distance between sorted samples (Kolouri et al., 2019b): with L random projections, the SWD and GSWD between µ and ν can be approximated by:

Original: [Equation]
Revised: [Equation]

Original: where I lx and I ly are sequences consist of the indices of sorted samples which satisfy (cid:104)
Revised: where I lx and I ly are sequences consist of the indices of sorted samples which satisfy (cid:104)

Original: x I lx [ n ] ,θ l (cid:105) ≤ (cid:104) x
Revised: x I lx [ n ] ,θ l (cid:105) ≤ (cid:104) x

Original: I lx [ n +1] ,θ l (cid:105) , (cid:104)
Revised: I lx [ n +1] ,θ l (cid:105) , (cid:104)

Original: y
Revised: y

Original: I ly [ n ] ,θ l (cid:105) ≤ (cid:104)
Revised: I ly [ n ] ,θ l (cid:105) ≤ (cid:104)

Original: y
Revised: y

Original: I ly [ n +1] ,θ l (cid:105) in the SWD, and β ( x I lx [ n ] ,θ l ) ≤ β ( x I lx [ n +1] ,θ l ) , β ( y I ly [ n ] ,θ l ) ≤ β
Revised: I ly [ n +1] ,θ l (cid:105) in the SWD, and β ( x I lx [ n ] ,θ l ) ≤ β ( x I lx [ n +1] ,θ l ) , β ( y I ly [ n ] ,θ l ) ≤ β

Original: ( y I ly [ n +1] ,θ l ) in the GSWD.
Revised: ( y I ly [ n +1] ,θ l ) in the GSWD.

Original: It is proved in Bonnotte (2013) that the SWD is a valid distance metric.
Revised: It is proved in Bonnotte (2013) that the SWD is a valid distance metric.

Original: The GSWD is a valid metric except for its neural network variant (Kolouri et al., 2019a).
Revised: The GSWD is a valid metric except for its neural network variant (Kolouri et al., 2019a).

Original:  A UGMENTED SLICED W ASSERSTEIN DISTANCES
Revised:  A UGMENTED SLICED W ASSERSTEIN DISTANCES

Original: In this section, we propose a new distance metric called the augmented sliced Wasserstein distance (ASWD), which embeds flexible nonlinear projections in its construction.
Revised: In this section, we propose a new distance metric called the augmented sliced Wasserstein distance (ASWD), which embeds flexible nonlinear projections in its construction.

Original: We also provide an implementation recipe for the ASWD.
Revised: We also provide an implementation recipe for the ASWD.

Original: 3. S PATIAL R ADON TRANSFORM AND AUGMENTED SLICED W ASSERSTEIN DISTANCE
Revised: 3. S PATIAL R ADON TRANSFORM AND AUGMENTED SLICED W ASSERSTEIN DISTANCE

Original: In the definitions of the SWD and GSWD, the Radon transform (Radon, 1917) and the generalized Radon transform (GRT) (Beylkin, 1984) are used as the push-forward operator for projecting distributions to a one-dimensional space.
Revised: In the definitions of the SWD and GSWD, the Radon transform (Radon, 1917) and the generalized Radon transform (GRT) (Beylkin, 1984) are used as the push-forward operator for projecting distributions to a one-dimensional space.

Original: However, it is not straightforward to design defining functions β ( x,θ ) (Kolouri et al., 2019a) for the GRT due to certain non-trivial requirements for the function (Beylkin, 1984).
Revised: However, it is not straightforward to design defining functions β ( x,θ ) (Kolouri et al., 2019a) for the GRT due to certain non-trivial requirements for the function (Beylkin, 1984).

Original: In practice, the assumption of the transform can be relaxed, as Theorem 1 shows that as long as the transform is injective, the corresponding ASWD metric is a valid distance metric.
Revised: In practice, the assumption of the transform can be relaxed, as Theorem 1 shows that as long as the transform is injective, the corresponding ASWD metric is a valid distance metric.

Original: To help us define the augmented sliced Wasserstein distance, we first introduce the spatial Radon transform which includes the vanilla Radon transform and the polynomial GRT as special cases (See Remark 2).
Revised: To help us define the augmented sliced Wasserstein distance, we first introduce the spatial Radon transform which includes the vanilla Radon transform and the polynomial GRT as special cases (See Remark 2).

Original: Definition 1.
Revised: Definition 1.

Original: Given an injective mapping g ( · ): R d → R d θ and a probability measure µ ∈ P ( R d ) which probability density function (PDF) is p µ , the spatial Radon transform of p µ is defined as
Revised: Given an injective mapping g ( · ): R d → R d θ and a probability measure µ ∈ P ( R d ) which probability density function (PDF) is p µ , the spatial Radon transform of p µ is defined as

Original: [Equation]
Revised: [Equation]

Original: (11) where t ∈ R and θ ∈ S
Revised: (11) where t ∈ R and θ ∈ S

Original: d θ − 1 are the parameters of hypersurfaces in R d .
Revised: d θ − 1 are the parameters of hypersurfaces in R d .

Original: Remark 1.
Revised: Remark 1.

Original: Note that the spatial Radon transform can be interpreted as applying the vanilla Radon transform to the PDF of ˆ x = g ( x ) , where x ∼ p µ .
Revised: Note that the spatial Radon transform can be interpreted as applying the vanilla Radon transform to the PDF of ˆ x = g ( x ) , where x ∼ p µ .

Original: Denote the PDF of ˆ x by p ˆ µ g , the spatial Radon transform defined by Equation (11) can be rewritten as:
Revised: Denote the PDF of ˆ x by p ˆ µ g , the spatial Radon transform defined by Equation (11) can be rewritten as:

Original: [Equation]
Revised: [Equation]

Original: Hence the spatial Radon transform inherits the theoretical properties of the Radon transform subject to certain conditions of g ( · ) and incorporates nonlinear projections through g ( · ) .
Revised: Hence the spatial Radon transform inherits the theoretical properties of the Radon transform subject to certain conditions of g ( · ) and incorporates nonlinear projections through g ( · ) .

Original: In what follows, we use f 1 ≡ f 2 to denote functions f 1 ( · ) :
Revised: In what follows, we use f 1 ≡ f 2 to denote functions f 1 ( · ) :

Original: X → R and f 2 ( · ) :
Revised: X → R and f 2 ( · ) :

Original: X → R that satisfy f 1 ( x )= f 2 ( x ) for almost ∀ x ∈ X .
Revised: X → R that satisfy f 1 ( x )= f 2 ( x ) for almost ∀ x ∈ X .

Original: Lemma 1.
Revised: Lemma 1.

Original: Given an injective mapping g ( · ) :
Revised: Given an injective mapping g ( · ) :

Original: R d → R d θ and two probability measures µ,ν ∈ P ( R d ) whose probability density functions are p µ and p ν , respectively, for all t ∈ R and θ ∈ S d θ − 1 ,
Revised: R d → R d θ and two probability measures µ,ν ∈ P ( R d ) whose probability density functions are p µ and p ν , respectively, for all t ∈ R and θ ∈ S d θ − 1 ,

Original: H p µ ( t,θ ; g ) ≡H p ν ( t,θ ; g )
Revised: H p µ ( t,θ ; g ) ≡H p ν ( t,θ ; g )

Original: if and only if p µ ≡ p ν , i.e. the spatial Radon transform is injective.
Revised: if and only if p µ ≡ p ν , i.e. the spatial Radon transform is injective.

Original: Moreover, the spatial Radon transform is injective if and only if the mapping g ( · ) is an injection.
Revised: Moreover, the spatial Radon transform is injective if and only if the mapping g ( · ) is an injection.

Original: See Appendix A for the proof of Lemma 1.
Revised: See Appendix A for the proof of Lemma 1.

Original: Remark 2.
Revised: Remark 2.

Original: The spatial Radon transform degenerates to the vanilla Radon transform when the mapping g ( · ) is an identity mapping.
Revised: The spatial Radon transform degenerates to the vanilla Radon transform when the mapping g ( · ) is an identity mapping.

Original: When g ( · ) is a homogeneous polynomial function with odd degrees, the spatial
Revised: When g ( · ) is a homogeneous polynomial function with odd degrees, the spatial

Original: Radon transform is equivalent to the polynomial GRT (Ehrenpreis, 2003).
Revised: Radon transform is equivalent to the polynomial GRT (Ehrenpreis, 2003).

Original: Appendix B provides the proof of Remark 2.
Revised: Appendix B provides the proof of Remark 2.

Original: We now introduce the augmented sliced Wasserstein distance, by utilizing the spatial Radon transform as the measure push-forward operator:
Revised: We now introduce the augmented sliced Wasserstein distance, by utilizing the spatial Radon transform as the measure push-forward operator:

Original: Definition 2.
Revised: Definition 2.

Original: Given two probability measures µ,ν ∈ P k
Revised: Given two probability measures µ,ν ∈ P k

Original: ( R d ) , whose probability density functions are p µ and p ν , respectively, and an injective mapping g ( · ): R d → R d θ , the augmented sliced Wasserstein distance (ASWD) of order k ∈ [1 , + ∞ ) is defined as:
Revised: ( R d ) , whose probability density functions are p µ and p ν , respectively, and an injective mapping g ( · ): R d → R d θ , the augmented sliced Wasserstein distance (ASWD) of order k ∈ [1 , + ∞ ) is defined as:

Original: [Equation]
Revised: [Equation]

Original: where θ ∈ S d θ − 1 , W k is the k -Wasserstein distance defined by Equation (1), and H refers to the spatial
Revised: where θ ∈ S d θ − 1 , W k is the k -Wasserstein distance defined by Equation (1), and H refers to the spatial

Original: Radon transform defined by Equation (11).
Revised: Radon transform defined by Equation (11).

Original: Remark 3.
Revised: Remark 3.

Original: Following the connection between the spatial Radon transform and the vanilla Radon transform as shown in Equation (12), the ASWD can be rewritten as:
Revised: Following the connection between the spatial Radon transform and the vanilla Radon transform as shown in Equation (12), the ASWD can be rewritten as:

Original: [Equation]
Revised: [Equation]

Original: where ˆ µ g and ˆ ν g are probability measures on R d θ which satisfy g ( x ) ∼ ˆ µ g for x ∼ µ and g ( y ) ∼ ˆ ν g for y ∼ ν .
Revised: where ˆ µ g and ˆ ν g are probability measures on R d θ which satisfy g ( x ) ∼ ˆ µ g for x ∼ µ and g ( y ) ∼ ˆ ν g for y ∼ ν .

Original: Theorem 1.
Revised: Theorem 1.

Original: The augmented sliced Wasserstein distance (ASWD) of order k ∈ [1 , + ∞ ) defined by
Revised: The augmented sliced Wasserstein distance (ASWD) of order k ∈ [1 , + ∞ ) defined by

Original: Equation (13) with a mapping g ( · ): R d → R d θ is a metric on P k
Revised: Equation (13) with a mapping g ( · ): R d → R d θ is a metric on P k

Original: ( R d ) if and only if g ( · ) is injective.
Revised: ( R d ) if and only if g ( · ) is injective.

Original: Proof of Theorem 1 is provided in Appendix C.
Revised: Proof of Theorem 1 is provided in Appendix C.

Original: 3. N UMERICAL IMPLEMENTATION
Revised: 3. N UMERICAL IMPLEMENTATION

Original: 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.
Revised: 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.

Original: Injective neural networks: As stated in Lemma 1 and Theorem 1, the injectivity of g ( · ) is the sufficient and necessary condition for the ASWD being a valid metric.
Revised: Injective neural networks: As stated in Lemma 1 and Theorem 1, the injectivity of g ( · ) is the sufficient and necessary condition for the ASWD being a valid metric.

Original: Thus we need specific architecture designs on implementing g ( · ) by neural networks.
Revised: Thus we need specific architecture designs on implementing g ( · ) by neural networks.

Original: One option is the family of invertible neural networks (Behrmann et al., 2019; Karami et al., 2019; Song et al., 2019), which are both injective and surjective.
Revised: One option is the family of invertible neural networks (Behrmann et al., 2019; Karami et al., 2019; Song et al., 2019), which are both injective and surjective.

Original: However, the running cost of those models is usually much higher than that of vanilla neural networks.
Revised: However, the running cost of those models is usually much higher than that of vanilla neural networks.

Original: We propose an alternative approach by concatenating the input x of an arbitrary neural network to its output φ ω ( x ) :
Revised: We propose an alternative approach by concatenating the input x of an arbitrary neural network to its output φ ω ( x ) :

Original: [Equation]
Revised: [Equation]

Original: It is trivial to show that g ω ( x ) is injective, since different inputs will lead to different outputs.
Revised: It is trivial to show that g ω ( x ) is injective, since different inputs will lead to different outputs.

Original: Although embarrassingly simple, this idea of concatenating the input and output of neural networks has found success in preserving information with dense blocks in the DenseNet (Huang et al., 2017), where the input of each layer is injective to the output of all preceding layers.
Revised: Although embarrassingly simple, this idea of concatenating the input and output of neural networks has found success in preserving information with dense blocks in the DenseNet (Huang et al., 2017), where the input of each layer is injective to the output of all preceding layers.

Original: Optimization objective: We aim to project samples to maximally discriminating hypersurfaces between two distributions, so that the projected samples between distributions are most dissimilar subject to certain constraints on the hypersurface, as shown in Figure 1.
Revised: Optimization objective: We aim to project samples to maximally discriminating hypersurfaces between two distributions, so that the projected samples between distributions are most dissimilar subject to certain constraints on the hypersurface, as shown in Figure 1.

Original: Similar ideas have been employed to identify important projection directions (Deshpande et al., 2019; Kolouri et al., 2019a; Paty & Cuturi, 2019) or a discriminative ground metric (Salimans et al., 2018) in optimal transport metrics.
Revised: Similar ideas have been employed to identify important projection directions (Deshpande et al., 2019; Kolouri et al., 2019a; Paty & Cuturi, 2019) or a discriminative ground metric (Salimans et al., 2018) in optimal transport metrics.

Original: For the ASWD, the parameterized injective neural network g ω ( · ) is optimized by maximizing the following objective:
Revised: For the ASWD, the parameterized injective neural network g ω ( · ) is optimized by maximizing the following objective:

Original: [Equation]
Revised: [Equation]

Original: where λ> 0 and the regularization term L λ =
Revised: where λ> 0 and the regularization term L λ =

Original: [Equation]
Revised: [Equation]

Original: is used to control the norm of the output of g ω ( · ) , otherwise the projections may be arbitrarily large.
Revised: is used to control the norm of the output of g ω ( · ) , otherwise the projections may be arbitrarily large.

Original: Remark 4.
Revised: Remark 4.

Original: The regularization coefficient λ adjusts the introduced non-linearity in the evaluation of the
Revised: The regularization coefficient λ adjusts the introduced non-linearity in the evaluation of the

Original: ASWD by controlling the norm of φ ω ( · ) in Equation (15).
Revised: ASWD by controlling the norm of φ ω ( · ) in Equation (15).

Original: In particular, when λ →∞ , the nonlinear term φ ω ( · ) shrinks to 0 .
Revised: In particular, when λ →∞ , the nonlinear term φ ω ( · ) shrinks to 0 .

Original: The rank of the augmented space is hence explicitly controlled by the flexible choice of φ ω ( · ) and implicitly regularized by L λ .
Revised: The rank of the augmented space is hence explicitly controlled by the flexible choice of φ ω ( · ) and implicitly regularized by L λ .

Original: By plugging the optimized g
Revised: By plugging the optimized g

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ( L
Revised: ( L

Original: [Equation]
Revised: [Equation]

Original: µ,ν
Revised: µ,ν

Original: [Equation]
Revised: [Equation]

Original: g ω ,λ )) into Equation (13), we obtain the empirical version of the ASWD.
Revised: g ω ,λ )) into Equation (13), we obtain the empirical version of the ASWD.

Original: Notably, the regularization term L λ is only used when maximizing the objective in
Revised: Notably, the regularization term L λ is only used when maximizing the objective in

Original: Equation (16), once the optimization is completed, L λ is not used in the calculation of the ASWD defined by Equation (13).
Revised: Equation (16), once the optimization is completed, L λ is not used in the calculation of the ASWD defined by Equation (13).

Original: Pseudocode is provided in Appendix D.
Revised: Pseudocode is provided in Appendix D.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we describe the experiments that we have conducted to evaluate performance of the proposed distance metric.
Revised: In this section, we describe the experiments that we have conducted to evaluate performance of the proposed distance metric.

Original: 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.
Revised: 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.

Original: Hence we compare performance of the ASWD with the state-of-the-art distance metrics in the same examples and report results as below 1 .
Revised: Hence we compare performance of the ASWD with the state-of-the-art distance metrics in the same examples and report results as below 1 .

Original: To examine the robustness of the ASWD, throughout the experiments, we adopt the injective network architecture given in Equation (15) and set φ ω to be a one layer fully-connected neural network whose outputs’ dimension equals its inputs’ dimension, with a ReLU layer as its activation function.
Revised: To examine the robustness of the ASWD, throughout the experiments, we adopt the injective network architecture given in Equation (15) and set φ ω to be a one layer fully-connected neural network whose outputs’ dimension equals its inputs’ dimension, with a ReLU layer as its activation function.

Original: 4. S LICED W ASSERSTEIN FLOWS
Revised: 4. S LICED W ASSERSTEIN FLOWS

Original: We first consider the problem of evolving a source distribution µ to a target distribution ν by minimizing
Revised: We first consider the problem of evolving a source distribution µ to a target distribution ν by minimizing

Original: Wasserstein distances between µ and ν in the sliced Wasserstein flow task reported in Kolouri et al. (2019a).
Revised: Wasserstein distances between µ and ν in the sliced Wasserstein flow task reported in Kolouri et al. (2019a).

Original: [Equation]
Revised: [Equation]

Original: where µ t refers to the updated source distribution at each iteration t .
Revised: where µ t refers to the updated source distribution at each iteration t .

Original: The SWD in Equation (17) can be replaced by other sliced-Wasserstein distances to be evaluated.
Revised: The SWD in Equation (17) can be replaced by other sliced-Wasserstein distances to be evaluated.

Original: As in Kolouri et al.
Revised: As in Kolouri et al.

Original: (2019a)
Revised: (2019a)

Original: , the 2-Wasserstein distance was used as the metric for evaluating performance of different distance metrics in this task.
Revised: , the 2-Wasserstein distance was used as the metric for evaluating performance of different distance metrics in this task.

Original: The set of hyperparameter values used in this experiment can be found in Appendix E.1.
Revised: The set of hyperparameter values used in this experiment can be found in Appendix E.1.

Original: [Figure]
Revised: [Figure]

Original: Without loss of generality
Revised: Without loss of generality

Original: [Figure]
Revised: [Figure]

Original: we initialize µ 0 to be the standard normal distribution N (0 ,I ) .
Revised: we initialize µ 0 to be the standard normal distribution N (0 ,I ) .

Original: We repeat each experiment 50 times and record the 2-Wasserstein distance between µ and ν at every iteration.
Revised: We repeat each experiment 50 times and record the 2-Wasserstein distance between µ and ν at every iteration.

Original: In Figure
Revised: In Figure

Original: [Figure]
Revised: [Figure]

Original: we plot the 2 -Wasserstein distances between the source and target distributions as a function of the training epochs and the 8-Gaussian, the Knot, the Moon, and the Swiss roll distributions are respective target distributions.
Revised: we plot the 2 -Wasserstein distances between the source and target distributions as a function of the training epochs and the 8-Gaussian, the Knot, the Moon, and the Swiss roll distributions are respective target distributions.

Original: For clarity, Figure 2 displays the experiment results from the 6 best performing distance metrics, including the ASWD, the DSWD, the SWD, the GSWD-NN 1, which directly generates projections through a one layer MLP, as well as the GSWD with the polynomial of degree 3, circular defining functions, out of the 12 distance metrics we compared.
Revised: For clarity, Figure 2 displays the experiment results from the 6 best performing distance metrics, including the ASWD, the DSWD, the SWD, the GSWD-NN 1, which directly generates projections through a one layer MLP, as well as the GSWD with the polynomial of degree 3, circular defining functions, out of the 12 distance metrics we compared.

Original: We observe from Figure 2 that the ASWD not only leads to lower 2-Wasserstein distances, but also converges faster by achieving better results with fewer iterations than the other methods in these four target distributions.
Revised: We observe from Figure 2 that the ASWD not only leads to lower 2-Wasserstein distances, but also converges faster by achieving better results with fewer iterations than the other methods in these four target distributions.

Original: A complete set of experimental results with 12 compared distance metrics andtarget distributions are included in Appendix F.1.
Revised: A complete set of experimental results with 12 compared distance metrics andtarget distributions are included in Appendix F.1.

Original: The ASWD outperforms the compared state-of-the-art sliced-based Wasserstein distance metrics with 7 out of the 8 target distributions except for the 25-Gaussian.
Revised: The ASWD outperforms the compared state-of-the-art sliced-based Wasserstein distance metrics with 7 out of the 8 target distributions except for the 25-Gaussian.

Original: This is achieved through the simple injective network architecture given in Equation (15) and a one layer fully-connected neural network with equal input and output dimensions throughout the experiments.
Revised: This is achieved through the simple injective network architecture given in Equation (15) and a one layer fully-connected neural network with equal input and output dimensions throughout the experiments.

Original: In addition, ablation study is conducted to study the effect of injective mappings, the regularization coefficient λ , and the optimization of hypersurfaces in the ASWD.
Revised: In addition, ablation study is conducted to study the effect of injective mappings, the regularization coefficient λ , and the optimization of hypersurfaces in the ASWD.

Original: Details can be found in Appendix F.2.
Revised: Details can be found in Appendix F.2.

Original: 4. G ENERATIVE MODELING
Revised: 4. G ENERATIVE MODELING

Original: In this experiment, we use the sliced-based Wasserstein distances for a generative modeling task described in
Revised: In this experiment, we use the sliced-based Wasserstein distances for a generative modeling task described in

Original: Nguyen et al.
Revised: Nguyen et al.

Original: The task is to generate images using generative adversarial networks (GANs) (Goodfellow et al., 2014) trained on either the CIFAR10 dataset (64 × 64 resolution)
Revised: The task is to generate images using generative adversarial networks (GANs) (Goodfellow et al., 2014) trained on either the CIFAR10 dataset (64 × 64 resolution)

Original: (Krizhevsky, 2009) or the
Revised: (Krizhevsky, 2009) or the

Original: CELEBA dataset (64 × 64 resolution) (Liu et al., 2015).
Revised: CELEBA dataset (64 × 64 resolution) (Liu et al., 2015).

Original: Denote the hidden layer and the output layer of the discriminator by h ψ and D Ψ , and the generator by G Φ , we train GAN models with the following objectives:
Revised: Denote the hidden layer and the output layer of the discriminator by h ψ and D Ψ , and the generator by G Φ , we train GAN models with the following objectives:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: where p z is the prior of latent variable z and p r is the distribution of real data.
Revised: where p z is the prior of latent variable z and p r is the distribution of real data.

Original: The SWD in Equation (18) is replaced by the ASWD and other variants of the SWD to compare their performance.
Revised: The SWD in Equation (18) is replaced by the ASWD and other variants of the SWD to compare their performance.

Original: The GSWD with the polynomial defining function and the DGSWD is not included in this experiment due to its excessively high computational cost in high-dimensional space.
Revised: The GSWD with the polynomial defining function and the DGSWD is not included in this experiment due to its excessively high computational cost in high-dimensional space.

Original: The Fr´echet Inception Distance (FID score) (Heusel et al., 2017) is used to assess the quality of generated images.
Revised: The Fr´echet Inception Distance (FID score) (Heusel et al., 2017) is used to assess the quality of generated images.

Original: More details on the network structures and the parameter setup used in this experiment are available in Appendix E.2.
Revised: More details on the network structures and the parameter setup used in this experiment are available in Appendix E.2.

Original: We run 200 and 100 training epochs to train the GAN models on the CIFAR10 and the CELEBA dataset, respectively.
Revised: We run 200 and 100 training epochs to train the GAN models on the CIFAR10 and the CELEBA dataset, respectively.

Original: Each experiment is repeated for 10 times.
Revised: Each experiment is repeated for 10 times.

Original: We report experimental results in Table 1.
Revised: We report experimental results in Table 1.

Original: With the same number of projections and a similar computation cost, the ASWD leads to significantly improved
Revised: With the same number of projections and a similar computation cost, the ASWD leads to significantly improved

Original: FID scores among all evaluated distances metrics on both datasets, which implies that images generated with the ASWD are of higher qualities.
Revised: FID scores among all evaluated distances metrics on both datasets, which implies that images generated with the ASWD are of higher qualities.

Original: Figure 3 plots the FID scores recorded during the training process.
Revised: Figure 3 plots the FID scores recorded during the training process.

Original: The GAN model trained with the ASWD exhibits a faster convergence as it reaches lower FID scores with fewer epochs.
Revised: The GAN model trained with the ASWD exhibits a faster convergence as it reaches lower FID scores with fewer epochs.

Original: Randomly selected samples of generated images are presented in Appendix G.
Revised: Randomly selected samples of generated images are presented in Appendix G.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We proposed a novel variant of the sliced Wasserstein distance, namely the augmented sliced Wasserstein distance (ASWD), which is flexible, has high projection efficiency, and generalizes well.
Revised: We proposed a novel variant of the sliced Wasserstein distance, namely the augmented sliced Wasserstein distance (ASWD), which is flexible, has high projection efficiency, and generalizes well.

Original: The ASWD adaptively updates the hypersurface where the samples are projected onto by learning from data.
Revised: The ASWD adaptively updates the hypersurface where the samples are projected onto by learning from data.

Original: We proved that the ASWD is a valid distance metric and presented its numerical implementation.
Revised: We proved that the ASWD is a valid distance metric and presented its numerical implementation.

Original: We reported empirical performance of the ASWD over state-of-the-art sliced Wasserstein metrics in numerical experiments.
Revised: We reported empirical performance of the ASWD over state-of-the-art sliced Wasserstein metrics in numerical experiments.

Original: The ASWD leads to the smallest distance errors over the majority of datasets in a sliced
Revised: The ASWD leads to the smallest distance errors over the majority of datasets in a sliced

Original: Wasserstein flow task and superior performance in a generative modeling task.
Revised: Wasserstein flow task and superior performance in a generative modeling task.

Original: As of our future research, a potential extension of the ASWD’s application is in computer graphics, for instance, calculating the sliced Wasserstein barycenters with the ASWD.
Revised: The ASWD can be extended in several directions. One future research topic is to incorporate the framework of the ASWD into the sliced Gromov-Wasserstein distance, such that data-adaptive, nonlinear projections can be learned to compare distributions whose supports do not necessary lie in the same metric space. It would be also interesting to explore the identification of the barycenter of objects via the ASWD, e.g. through an optimization scheme alternating between optimization of the nonlinear maps in the ASWD framework and the barycenter. Another interesting area is the incorporation of the ASWD with state-of-theart generative models to further exploit projection efficiency.

Revised: We anticipate that they will provide promising directions for enabling the generalization for sliced Wasserstein distance to be applied in wider domains.