Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This paper introduces _Augmented Sliced Wasserstein Distances_ (ASWD) which consist, in a nutshell, in computing the usual SWD between pushed measures: $\mathrm{ASWD}(\mu,\nu; g) = \mathrm{SWD}(g \circ \mu, g \circ \nu)$ (note: the pushforward operator does not render properly, so I will use $\circ$ instead). The idea is then to maximize on $g$ (under constraint/regularization) to obtain a pushforward map $g$ that provides the most discriminative projection directions. Authors provide an exhaustive condition on $g$ (namely, injectivity) to make the ASWD a metric between probability measures. Their approach is then showcased on a various set of benchmark experiments in computational optimal transport (gradient flows, generative models, auto-encoder, barycenters, color transfert...).
main_review: [Edit: increased grade 6-->7 after rebuttal]


Quality:
The paper is overall interesting and supported by a large number of convincing experiments. Some of its theoretical aspects can be improved though:
- When defining the ASWD, the map $g$ should be chosen to ensure that $g \circ \mu, \nu$ still have finite $k$-th moments. 
- Corollary 1.1, in its current state, seems wrong/incomplete. It is not clear why the argmax should be reached. If my understanding is correct, it is actually likely that the supremum (even when restricted on bounded injective maps) is infinite (each instance is finite, as stated in line 496, but the supremum may not). For instance, let $\mu = \delta_x, \nu = \delta_y$ with $x,y \in \mathbb{R}^d$; then $\mathrm{ASWD}(\mu,\nu; g) = \int_\theta |\braket{\theta, g(x) - g(y)}| \mathrm{d} \theta$. Taking $g = g_\omega \simeq \omega \arctan( \cdot )$ may lead to arbitrarily large values (while such $g_\omega$ is injective and bounded, at fixed $\omega$). This is actually the point of line 199 (regularization in the NN); in my opinion, this should be the optimization problem used to define the ASWD metric (to be compared with the Vaniila-ASWD if one does not optimize on $g$), and it would be nice to have a proof that this supremum is reached, perhaps a statement about the (non-)convexity(?) of the problem, etc., as it is actually at the core of the work. 

Clarity: 
The paper is well-written and easy to read overall.

Originality/Significance:
This paper heavily builds on previous works (in particular those of Kolouri et al.) which it properly credits. It provides simple yet seemingly efficient variations that improve on previous results. A quick look at the code and at the variety of experiments run also suggest that, once made publicly available, these tools may be useful of the computational OT community.


Minor remarks and suggestions:
- The paper may benefits from being written in a measure-based formalism; the current definitions suggest that $\mu,\nu$ should have densities $p_\mu, p_\nu$ and thus that discrete measures are not encompassed, which is obviously not the case. 
- In Figure 14, it would be nice to include (Regularized) OT barycenters (available in PythonOptimalTransport for instance) to showcase the (potential) difference between sliced approaches and the "classical one". 
- (line 104) It may be worth briefly mentioning these "certain conditions" for the sake of completeness. 
limitations_and_societal_impact: The paper mentions some of its limitations and potential negative social impact (which are quite general to the ML field); I did not identify specific impact this paper could have.
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: The sliced Wasserstein distance (SWD) is a metric on the space of probability measures and serves as a practical alternative to the standard metric arising from the optimal transport (OT) problem, the Wasserstein distance. Indeed, while the Wasserstein distance suffers from computational and statistical limitations on large-scale and high-dimensional settings, SWD can alleviate these issues while offering similar theoretical properties. Hence, SWD has been increasingly popular within the statistical and machine learning (ML) community in recent years.

SWD is defined as an expectation that is intractable in general; in practice, it is then usually approximated with a simple Monte Carlo estimation, but prior studies reported that this method might induce an important approximation error on high-dimensional settings, leading to an underestimation of the "true" dissimilarity between two distributions. Therefore, one would need to compute a Monte Carlo approximation based a very large amount of samples to obtain an accurate approximation, which increases the computational complexity of SWD. 

Several variants of SWD have recently been introduced to overcome this problem and include the "generalized sliced Wasserstein distances" (GSWD): this class of distances aim at offering a better computational efficiency, by comparing distributions from their nonlinear one-dimensional representations (instead of the linear ones used in SWD). This paper focuses on this line of work: the limitations of GSWD are identified to motivate the formulation of a novel class of distances, called "Augmented Sliced Wasserstein distances" (ASWD).

Similarly to SWD and GSWD, the definition of ASWD is connected to the Radon transform: the authors introduce a novel type of Radon transform, the "spatial Radon transform" (Definition 1), and use it to define the class of augmented sliced Wasserstein distances. 
Then, they study the theoretical properties of the spatial Radon transform and ASWD. Their main result states that ASWD verify all metric axioms if and only if the mapping $g$ (that characterizes the underlying spatial Radon transform) (1) is injective (Theorem 1), or (2) maximizes ASWD over the space of bounded and injective functions (Corollary 1.1).
The authors then introduce a specific class of mappings (Equation (15)) which meet the conditions of Corollary 1.1 thus guarantee that the resulting ASWD is a metric. Such mappings require training a neural network to make ASWD data-adaptive.
Finally, the superior empirical performance of ASWD against existing variants of SWD (including GSWD) is illustrated on various problems, including gradient flows, generative modeling, color transferring and barycenters of measures. 

main_review: This work proposes a new family of distances inspired by computational optimal transport, which can be easily applied to various classical machine learning applications. In that sense, it fits the interests of the NeurIPS community.

I found the paper quite clear and well-written overall. The structure can be improved by incorporating the following changes,
- l.46-47: "Different variants of the SWD have been proposed to improve the projection efficiency of the SWD, either by introducing nonlinear projections or by optimizing the distribution of random projections." It seems more relevant to talk about the max-sliced Wasserstein distance (max-SWD) in this paragraph instead of in the previous one, since max-SWD was introduced in order to improve the "projection complexity issue" of SWD (see "Max-Sliced Wasserstein Distance and its use for GANs", Deshpande et al., 2019). 
- I suggest adding the equation that explains how to approximate equation (2) at the end of the paragraph "Wasserstein distance" (Section 2). This would clarify l.113-115, as well as equations (9) and (10)
- l.101-105: I would remove this part since explaining how to approximate the Radon transform or GRT does not seem useful in this work. On the contrary, specifying the approximation for the regularization term in Section 3.2 (l.198) would be more relevant.
- l.168: "In practical applications, the mapping $g(\cdot)$ needs to be optimized." At this stage, it is not obvious why this optimization is required. I suggest moving the explanation in Section 3.2 (e.g., l.191-196) to clarify this point.

Overall, the theoretical claims look correct to me; my only concerns are detailed below.
- Remark 2: "$g(\cdot)$ is a homogeneous polynomial function with odd degrees" (l.151) contradicts equation (31) in Appendix B. On the other hand, it is not clear to me why the mapping defined in equation (31) is injective on $\mathbb{R}^d$.
- Appendix D: Why is it important to assume that the optimal mapping is bounded? Besides, equations (39), (40), (41) could be explained in more details.

Other (minor) typos and inaccuracies are listed at the end of the review.

In my opinion, the main weakness of this work, which prevents me from giving a positive score, is the lack of originality and significance as compared to prior work. Indeed, ASWD build on GSWD and their definition differs from a simple trick (using the "spatial Radon transform" introduced by the authors, instead of the generalized Radon transform). The structure of this paper is the same as in "Generalized Sliced Wasserstein Distances" (Kolouri et al., 2019) [1]: the novel distance is defined, then its metric axioms are studied, finally its performance against the sliced Wasserstein distance (and its variants) is evaluated on practical applications. These two aspects are not necessarily concerning and help make the paper quite intuitive; besides, the empirical study provided by this paper is more extensive than the one in [1], since the compared methods and applications are more diverse. What is more questionable is that the advantages of ASWD as compared to SWD and GSWD are unclear to me in the end, given the contributions and discussion provided in this work. 

First, from a theoretical point of view, this paper only investigates the metric properties, and establishes that ASWD satisfy all axioms of metrics if and only if the mapping associated to its Radon transform is injective. While the authors argue that this is a significant advantage over GSWD, we don't know if this explains the superior performance of ASWD-based methods in practice. Besides, this result is similar to Proposition 1 in [1], which states that GSWD are metrics if and only if the corresponding generalized Radon transform is injective; consequently, the proofs in the appendix are largely inspired by [1].

Then, the empirical performance of ASWD is not sufficiently discussed: while the results in Section 5 encourage its use against existing slice-based distances, the appendix illustrates that ASWD do not necessarily yield a superior performance:
- Appendix I and K: ASWD and SWD achieve a similar performance, while ASWD is more computationally expensive.
- Appendix I and J: the performance of GSWD should be studied on these experiments as well.

This shows that the superior performance of ASWD depends on the application, and the theoretical results do not help understand the behavior in practice. In particular, I am wondering whether the regularization term (equation (16)) might help improve the performance in the flows and generative modeling applications: it would be interesting to regularize the other methods to study the impact. 

Hence, I encourage the authors to provide more significant evidence on why ASWD is a better alternative to SWD than GSWD, or to show a more nuanced picture in light of the empirical results in the appendix; for now, I am not convinced by the superiority of ASWD over SWD/GSWD given the motivation (in particular, l.234-241) nor the theoretical and empirical contributions.

Typos:

Equation (1): the integral should be over $\Omega \times \Omega$

l.92: "over spaces of lines in $\mathbb{R}^d$" -> "over spaces of hyperplanes in $\mathbb{R}^d$" 

l.99: "while $\Omega_\theta$ is a compact set" -> "where $\Omega_\theta$ is a compact set"

l.110: "where the Radon transform ... is adopted as the measure push-forward operator" This can be further clarified by explaining the connection between the Radon transform and the push-forward operator that projects the two distributions (Proposition 6 in "Sliced and Radon Wasserstein Barycenters of Measures", Bonneel et al.).

l.120: "The GSWD is a valid metric except for its neural network variant [Kolouri et al., 2019a]" This is not exactly true: it remains unknown whether neural networks-based GSWD (GSWD-NN) satisfy all axioms for a metric (see the explanation in Section 3.3 of "Generalized Sliced Wasserstein Distances", Kolouri et al.). Therefore, a correct statement would be, for example, "The GSWD has been shown to be a valid metric, but the proof remains incomplete for its neural network variant [Kolouri et al., 2019a]".

Definition 1: 
- The set $P(\mathbb{R}^d)$ is not defined
- What are the possible values for $d$ and $d_\theta$?

l.147: please define the operator "\equiv".

l.172: "between two probability" -> "between two probability measures"

Section 5: The order $k$ of SWD/ASWD used in the experiments should be precised. 

l.256-257: "by minimizing Wasserstein distances" -> "by minimizing slice-based Wasserstein distances"

l.425, 458, 463: "for $\forall x \in \mathbb{R}^d$" -> "for all $x \in \mathbb{R}^d$". Same remark in l.428, 432, 434, 506...

l.455: "which probability density function is ..." -> "whose probability density function is ..."

l.498: "that the the ASWD"

l.589: "Sampels of generated images" -> "Samples of generated images"

limitations_and_societal_impact: The limitations of this work are not clearly discussed, but the authors adequately addressed the societal impact (see Appendix L).
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The author proposed a variant of sliced Wasserstein distance (SWD) called **augmented sliced Wasserstein distance** (ASWD) which replace the linear projections in SWD with a non-linear one. The construction of ASWD consists of two steps: (1) define spatial Radon transform (SRT) for probability measures (2) replace the Radon transform in the original SWD with this new SRT. The key intuition of SRT is to first transform the function $\pmb{x}$ to high dimensional hyperspace by transformation $g(\pmb{x})$ and then perform Radon transform on functions defined in this hyperspace. 

Theoretically, the author proves the conditions required for SRT to be injective. Later, the author uses this result to show the corresponding ASWD is a distance metric. In addition, the author proposes a training objective for finding the transformation $g(\cdot)$. 

Empirically, the author conducts both toy and practical experiments to demonstrate the advantages of ASWD in terms of higher projection efficiency and also provides some ablation studies in the appendix to investigate the effects of different $g(\cdot)$ choices.
main_review: ### Clarity
The main idea of ASWD is easy to understand and the author also gives a quite clear explanation on how to construct ASWD. Specifically, I appreciate the **Remarks**, which greatly improves the clarity of the paper. So I think the paper is clearly written and quite easy to follow. 

### Originality
As for the originality, the idea of using injective transformation to preserve some statistical properties is not new. For example, this is the foundation of normalizing flow, and it can also be found in the deep-kernel literature [1] to preserve the discrepancy validity. However, this seems to be new in the scope of SWD literature. In summary, I found the originality is ok but not significantly novel. 

### Quality and Significance
In terms of the research quality, the method is well-presented. I checked some of the proofs, they seem to be correct. Empirically, the author conducts several experiments to confirm the advantages of ASWD. However, I still have some concerns related to the experiments and the way the author presented the ASWD. In terms of significance, this work should be easy to implement and understand. So it should have some impact on the SWD community. 

### Questions
1. You mentioned that the idea of SRT is to transform the original data $\pmb{x}\in\mathbb{R}^d$ to $\mathbb{R}^{d_\theta}$ and $d_\theta> d$, i.e. higher-dimensional space. However, I don't think this is necessary based on the theory you give. For structured data like images, the commonly used assumption is that they have a low-dimensional manifold. So I don't understand why you project it in higher-dimensional space. To me, it seems that this is due to the choice of the NN architecture (concatenation) and not from the theoretical arguments.
2. The form of the transformation $g$ is quite simple. Have you evaluated the performance of using the invertible neural network for $g$? Despite the computational cost, does it give better performance?
3. I wonder about the performance if you transform the data to low dimensional space by an arbitrary neural network. Indeed, strict injectivity cannot be guaranteed but one can roughly construct an injective network using the technique mentioned in [1]. 
4. What advantages of this strict injectivity can provide in practice? In the ablation studies, it seems that the non-injective NN can sometimes give even better performance compared to injective ASWD? Any real applications that actually require this strict injectivity?
5. Despite the ASWD gives better performance compared to other baselines, the FID score is still quite low compared to the current state-of-the-art method. So what real advantages of this SWD can bring?

Minor:
1. Is SRT a special case of GRT in general? If so, I wonder about the expressiveness of SRT. Any transformations that SRT cannot represents but GRT can?



[1] Li, Chun-Liang, et al. "Mmd gan: Towards deeper understanding of moment matching network." arXiv preprint arXiv:1705.08584 (2017).

---
I appreciate the author's response. They managed to address most of my concerns. However, I agree with reviewer HUCr about the gap between the theory and the definition of $\pmb{g}$. The search space of $\pmb{g}$ and the optimal $\pmb{g}^*$ should be re-defined to address this. Anyway, I agree that this work should have some potentials, I will raise my score. 
limitations_and_societal_impact: For limitations, see Main Review. 

Appendix L (societal impacts) addresses the potential negative social impact of this work.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: The paper develops a computationally efficient metric between two probability distributions, which is called Augmented Sliced Wasserstein Distances (ASWD). The proposed metric applied neural networks to map two distributions from $\mathbb{R}^d$ into high dimensional non-linear manifolds, then the mapped distributions are compared using the sliced Wasserstein distance. It is shown in the paper that ASWD is a well-defined metric, and an algorithm for empirically evaluating ASWD is provided. The empirical applicability of ASWD is demonstrated by numerical experiments including gradient flow, generative modeling, sliced Wasserstein autoencoder, image color transferring and 
sliced Wasserstein barycenter.
main_review: In general the paper is well-written, the proof is rigorous, and the numerical result is convincing. 

**major concerns:**

1. How to select the regularization parameter $\lambda$ is not discussed in the paper. Though it is discussed in Remark 5 that $\lambda$ adjusts the introduced non-linearity and $\lambda$ can not be too small (the output $g_{\omega}$ might explode) or too large (totally kills the nonlinear term), how to select the lambda in practice is not clear to the readers. In appendix F, the authors report: empirical errors in the experiment are found to be not sensitive to the choice of $\lambda$ in a candidate set of {0.01, 0.05, 0.1, 0.5}: does it means that the candidate set is not rich enough so that one can not observe variations?

2. Empirically, the neural network $g$ is selected as $g^\ast = argmax(\mathcal L(\mu,\nu;g,\lambda))$ instead of $g^\ast = argmax(\mathrm{ASWD}(\mu,\nu;g))$. With such $g^\ast$, is the result in Corollary 1.1 still valid? 

**Minor suggestions:**

1. Line 81: Is it necessary to assume the existence of PDFs for $\mu$ and $\nu$? Many interesting data-driven applications does not satisfy this assumptions. For instance, the empirical measure of a finite observations is not absolutely continuous w.r.t Lebesgue measure.

2. Line 94: $\delta(\cdot)$ is the **1-dimensional** Dirac delta function.

3. Line 434: need to restrict to measurable $\mathcal X$, otherwise the integral is not well defined.
limitations_and_societal_impact: The author discussed societal impact in appendix L. An additional limitation of this work could be:

- As ASWD applied neural networks to map the data, the resulting metrics might be hard to interpret.
review_title: null

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

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