Title: Pooling by Sliced-Wasserstein Embedding
Authors: Navid Naderializadeh
Joseph F. Comer
Reed W Andrews
Heiko Hoffmann
Soheil Kolouri
Venue: NeurIPS 2021
Keywords: Pooling
feature aggregation
permutation invariance
optimal transport
set learning
representation learning

Abstract: Learning representations from sets has become increasingly important with many applications in point cloud processing, graph learning, image/video recognition, and object detection. We introduce a geometrically-interpretable and generic pooling mechanism for aggregating a set of features into a fixed-dimensional representation. In particular, we treat elements of a set as samples from a probability distribution and propose an end-to-end trainable Euclidean embedding for sliced-Wasserstein distance to learn from set-structured data effectively. We evaluate our proposed pooling method on a wide variety of set-structured data, including point-cloud, graph, and image classification tasks, and demonstrate that our proposed method provides superior performance over existing set representation learning approaches. Our code is available at https://github.com/navid-naderi/PSWE.

TL;DR: We present a novel pooling method, called Pooling by Sliced-Wasserstein Embedding (PSWE), which leverages the (generalized) sliced-Wasserstein distances to map an input set of arbitrary cardinality to a fixed-size representation.

Review discussion:
Pros:
* Paper is relatively clear written and easy to follow. The content is pretty much self-contained with sufficient background information about wasserstein distance in Section 3.
* A straightforward modification of the sliced-wasserstein distance by introducing two trainable components: trainable slicers and trainable reference sets.
* The problem being addressed here is important and has many practical applications.
Cons:
* Motivation of the introduced modifications is not clear. In line 58, the authors motivate the proposed changes by claiming that “PSWE provides an exact Euclidean embedding for the (generalized) sliced-Wasserstein (SW) distance”.
However, without the two proposed modifications, the conventional generalized sliced-Wasserstein distance also provides Euclidean embedding during distance computation.
* Section 4.3 “On projection complexity of sliced distances” is not clear in its current form. The authors claimed that samples often live on a low-dimensional manifold and therefore a large number of projections L should be used (line 191).
However the exact projection complexity (compute/memory) analysis is not presented.
* Also in Section 4.3, experiments with different numbers of projections have been conducted in [A1]. What is the main message the authors try to convey here? Are trainable slicers more efficient than random projection slicers? From my point of view, random projection in practice is a very strong baseline and it is quite critical to compare the proposed trainable slicer and random projection slicer in terms of classification accuracy and computational cost.
* [A2] is one of the early works that proposed the concept of sliced-wasserstein distance and should be discussed.
[A1] Sliced Wasserstein Discrepancy for Unsupervised Domain Adaptation. CVPR 2019\
[A2] Wasserstein barycenter and its application to texture mixing. SSVM, 2011.
No
First we would like to genuinely thank you for your generous service to the community as a reviewer amidst these troubling times, and for your constructive feedback. Below please find our specific comments.
* **Without the two proposed modifications, the conventional generalized sliced-Wasserstein distance also provides Euclidean embedding during distance computation.** Here we clarify that this point is not true.
First, a Euclidean embedding for a given distance (e.g., SW), is a mapping from an input space (e.g., space of probability measures) to a fixed-dimensional vector space in which the Euclidean distance between two vectors corresponds to the given distance between the two inputs in the input space. Our work provides such a mapping to the embedding space leveraging the idea of linearized Wasserstein embedding proposed in the prior work, and extending it to sliced-Wasserstein distances.
* **The exact projection complexity (compute/memory) analysis is not presented.** In addition to the performance results reported in Tables 1 and 2 that show the sensitivity of approach to the choice of number of slices, $L$, we have added computational analysis of the method with respect to the number of slices. Below we include the per-epoch training and inference times of different pooling approaches on the ModelNet40 dataset, which should complete the projection complexity picture and address the reviewer's concern. All experiments were run in isolation on a $2.90$ GHz Intel&reg; Xeon&reg; Gold 6226R CPU and two $32$ GB NVIDIA&reg; Tesla&reg; V100S GPUs.
**Average train time per epoch (seconds)**
| Backbone |        GAP       |    PMA    |   FSPool   |    PSWE (L=1)   |    PSWE (L=4)   |   PSWE (L=16)   |   PSWE (L=64)   |   PSWE (L=256)  |  PSWE (L=1024)  |
    |----------|------------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|
    |   MLP    | 21.35 $\pm$ 0.09 | 26.58 $\pm$ 0.12 | 22.82 $\pm$ 0.06 | 21.89 $\pm$ 0.02 | 21.91 $\pm$ 0.02 | 21.95 $\pm$ 0.02 | 22.33 $\pm$ 0.02 | 23.79 $\pm$ 0.09 | 30.17 $\pm$ 0.24 |
    |   ISAB   | 26.27 $\pm$ 0.03 | 30.50 $\pm$ 0.11 | 27.31 $\pm$ 0.03 | 26.71 $\pm$ 0.02 | 26.74 $\pm$ 0.02 | 26.76 $\pm$ 0.02 | 27.03 $\pm$ 0.03 | 28.06 $\pm$ 0.04 | 33.99 $\pm$ 0.05 |
**Average test time per epoch (seconds)**
| Backbone |        GAP       |    PMA    |   FSPool   |    PSWE (L=1)   |    PSWE (L=4)   |   PSWE (L=16)   |   PSWE (L=64)   |   PSWE (L=256)  |  PSWE (L=1024)  |
    |----------|------------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|
    |   MLP    | 0.16 $\pm$ 0.00 | 0.60 $\pm$ 0.00 | 0.40 $\pm$ 0.00 | 0.17 $\pm$ 0.00 | 0.17 $\pm$ 0.00 | 0.19 $\pm$ 0.00 | 0.27 $\pm$ 0.00 | 0.54 $\pm$ 0.00 | 1.68 $\pm$ 0.00 |
    |   ISAB   | 0.51 $\pm$ 0.00 | 0.97 $\pm$ 0.00 | 0.76 $\pm$ 0.00 | 0.52 $\pm$ 0.00 | 0.53 $\pm$ 0.00 | 0.54 $\pm$ 0.00 | 0.62 $\pm$ 0.00 | 0.89 $\pm$ 0.00 | 2.04 $\pm$ 0.01 |
* **Are trainable slicers more efficient than random projection slicers?** We believe that a critical point is missing here. The existing work on SW distances (e.g., [A1]) show that random slices work reasonably well, when random slices are sampled for each batch, for instance in generative modelling or domain adaptation. The goal in these frameworks is to compare two distributions via their random projections. In our framework and for global pooling, we require a fixed set of slices (which are not resampled) that are used to provide a representation for the input sets, $X_i$s. Selecting this fixed set randomly has severe consequences, as one can completely miss the discriminative/descriptive variations in the input space (again note that the slices are not resampled).
Hence, while the difference between optimizing the slices or using random slices is not crucial for former applications, it becomes absolutely necessary in our general pooling mechanism to optimize the slices.
* **[A2] is one of the early works that proposed the concept of sliced-wasserstein distance and should be discussed.** We regret having missed this prominent work on sliced-Wasserstein distance in our references. We will add and discuss this reference in the camera-ready version of the manuscript.
We hope that we have addressed your concerns and clarified any confusions.
In particular, we think the notion of the Euclidean embedding and the big picture of the unique requirements for defining a global pooling framework were missing in the original set of reviews. We hope that you revise your decision in light of the provided information, and we would be happy to discuss any further questions you may have.
Thanks to the authors for preparing the answers.
Regarding the runtime analysis, we can observe that when L=256 / L =1024, the inference time grows significantly.
In fact, the inference time of PSWE at L=1024 is actually more than double of the inference time of the PMA baseline but with similar performance (PSWE 87.6 acc/2.04 sec vs PMA 87.6 acc/0.97 sec). It is hard to justify why the proposed method is more favorable. I am leaning toward my original rating at this point.
We thank the reviewer for the additional question regarding the head-to-head comparison between PSWE and PMA. We agree that in the case of the ModelNet40 dataset and for a large number of slices, the inference time of PSWE is higher than PMA (while their training times are comparable).
However, aside from the fact that PSWE is geometrically interpretable and theoretically grounded as opposed to PMA, which is not, PSWE has a clear performance advantage over PMA in other domains, such as graphs and images. As we mentioned in the paper, for image recognition, PMA provided such a poor performance that we decided not to include its results in Table 3.
Moreover, we have performed a new set of image recognition experiments on the Places-Extra69 dataset, where we resolved a few minor issues with the image augmentations used during training and testing. Below please find the updated results of our experiment. As the table demonstrates, our proposed PSWE pooling method, while significantly outperforming PMA with both backbones, considerably outperforms global average pooling (GAP) and FSPool with the "$16 \times 16$ Patches + MLP" and ResNet18 backbones, respectively. This shows the effectiveness of the proposed PSWE pooling method when used in conjunction with different backbone types and complexities. We will add the updated results to the camera-ready version of the manuscript.
| Backbone | GAP              | FSPool           | PSWE             |
| -------- | ---------------- | ---------------- | ---------------- |
| $16 \times 16$ Patches + MLP       | 29.10 $\pm$ 0.10   | 35.40 $\pm$ 1.80   | 35.20 $\pm$ 0.20   |
| ResNet18       | 58.63 $\pm$ 0.03 | 57.30 $\pm$ 1.80 | 58.30 $\pm$ 0.02 |
We would be happy to continue the discussion to address any remaining concerns you might have.
I have very little to add beyond the summary above. Some small questions for the authors:
1. Can the sets $X_i$ be multisets, which would give rise to different empirical distributions (i.e., weighted delta functions) and how would this affect the method?
2. Why use ${\cal X}$ and ${\cal X'}$ in Eqn (1) instead of ${\cal X}_i$ and ${\cal X}_j$?
3. Are the regularity conditions presented in [39] (L.131) satisfied when the inputs are learned from a neural network?
4. Exactly what is learned in the set $X_0$ (L.171)? My understanding is that it's the $x_n^0$ that are learned/perturbed during learning, so it's unclear why the underlying distribution $p_0$ is needed. These points are never resampled.
5. The paper discusses computational complexity in terms of computing PSWE. Can the authors provide running times as a function of L for the experiments?
No
We would like to genuinely thank you for your service to the community as a reviewer amidst these troubling times, and for your positive evaluation of our paper. Below please find our responses to your thoughtful comments.
* **Can the sets $X_i$ be multisets?** Yes, absolutely! This is an interesting question that can broaden the application of our proposed framework. One can imagine different strategies for assigning weights to the set elements. In this regime, we need to revert equation (11) and use equation (3) directly to obtain the Monge map. Computationally, the approach will remain to entail two one-dimensional interpolations.
* **Replacing $\mathcal{X}$ and $\mathcal{X}'$ with $\mathcal{X}_i$ and $\mathcal{X}_j$.** This is a great suggestion. We will update the manuscript accordingly.
* **Are the regularity conditions presented in [39] (L.131) satisfied when the inputs are learned from a neural network?** This is a good question. Here we clarify that we treat the backbone as a mapping from a generic input into a set, which is then treated as an empirical measure.
So the backbone does not need to satisfy the regularity conditions. The conditions, however, must be satisfied for the slicer, $g_\theta$. In our experiments, in order to avoid complications, we used a linear slicing procedure, which satisfies all regularity conditions presented in [39].
* **Exactly what is learned in the set (L.171)?** Your understanding is spot-on. We do not resample $p_0$ and what is learned is trully the $x^0_n$s.
That said, we can certainly learn a parametric probability distribution for the reference method, for instance using a Gaussian Mixture, and sample the reference.
* **Can the authors provide running times as a function of L for the experiments?** Following your suggestion, we ran extensive experiments to perform wall-clock run-time comparisons between our proposed method and the baselines. Below we include the per-epoch training and inference times of different pooling approaches on the ModelNet40 dataset. All experiments were run in isolation on a $2.90$ GHz Intel&reg; Xeon&reg; Gold 6226R CPU and two $32$ GB NVIDIA&reg; Tesla&reg; V100S GPUs.
**Average train time per epoch (seconds)**
| Backbone |        GAP       |    PMA    |   FSPool   |    PSWE (L=1)   |    PSWE (L=4)   |   PSWE (L=16)   |   PSWE (L=64)   |   PSWE (L=256)  |  PSWE (L=1024)  |
    |----------|------------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|
    |   MLP    | 21.35 $\pm$ 0.09 | 26.58 $\pm$ 0.12 | 22.82 $\pm$ 0.06 | 21.89 $\pm$ 0.02 | 21.91 $\pm$ 0.02 | 21.95 $\pm$ 0.02 | 22.33 $\pm$ 0.02 | 23.79 $\pm$ 0.09 | 30.17 $\pm$ 0.24 |
    |   ISAB   | 26.27 $\pm$ 0.03 | 30.50 $\pm$ 0.11 | 27.31 $\pm$ 0.03 | 26.71 $\pm$ 0.02 | 26.74 $\pm$ 0.02 | 26.76 $\pm$ 0.02 | 27.03 $\pm$ 0.03 | 28.06 $\pm$ 0.04 | 33.99 $\pm$ 0.05 |
**Average test time per epoch (seconds)**
| Backbone |        GAP       |    PMA    |   FSPool   |    PSWE (L=1)   |    PSWE (L=4)   |   PSWE (L=16)   |   PSWE (L=64)   |   PSWE (L=256)  |  PSWE (L=1024)  |
    |----------|------------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|-----------------|
    |   MLP    | 0.16 $\pm$ 0.00 | 0.60 $\pm$ 0.00 | 0.40 $\pm$ 0.00 | 0.17 $\pm$ 0.00 | 0.17 $\pm$ 0.00 | 0.19 $\pm$ 0.00 | 0.27 $\pm$ 0.00 | 0.54 $\pm$ 0.00 | 1.68 $\pm$ 0.00 |
    |   ISAB   | 0.51 $\pm$ 0.00 | 0.97 $\pm$ 0.00 | 0.76 $\pm$ 0.00 | 0.52 $\pm$ 0.00 | 0.53 $\pm$ 0.00 | 0.54 $\pm$ 0.00 | 0.62 $\pm$ 0.00 | 0.89 $\pm$ 0.00 | 2.04 $\pm$ 0.01 |
We will add these results to the camera-ready version of our paper.
Thank you again for your comments and your time.
Thank you for your response. I am still in support of this paper but other reviews do raise some concerns, mostly with respect to empirical evaluation. I'm particularly concerned with the comment that you were not able to evaluate on ImageNet due to its non-commercial license yet other datasets such as Places-Extra69, ModelNet40 and IMDB datasets appear to be released under similar terms.
We appreciate the reviewer for their positive evaluation of our work.
As for ImageNet, we agree that it would be invaluable to show the effectiveness of our method on this dataset. To date, our legal department has strongly advised us against using this dataset in our studies.
However, we will follow up with them to see if they believe we can use it for this work, in which case, we will add the experiments to the camera-ready version of the manuscript accordingly.
This paper introduces a general geometric-aware pooling function. The key idea is based on sliced Wasserstein distance. The input is transformed to a set of points no matter what the original format is.
Then, the Wasserstein distance between this input set and a trainable reference set is computed. The sliced Wasserstein distance can be computed by sorting the input feature in a straightforward fashion. This paper demonstrates the performance on set classification, graph classification, and image classification tasks. In some cases, the proposed method outperforms its counterparts.
Pros:
1. This paper presents a great motivation to design a geometric-aware pooling function. Compared to others, this method treats the input features as  a set of samples from a probability distribution. This novel interpretation makes it possible to incorporate tools in geometric optimization.
2. The idea of using a sliced Wasserstein distance is technically sound. In the 1D case, the Wasserstein distance can be easily computed by sorting.
So this significantly reduces the complexity of Wasserstein distance computation.
3. This paper introduces the idea of using a trainable set, which can be viewed as  a dataset prior.
Cons:
1. The proposed method seems only applicable to the global level and it is not able to replace the a general pooling (e.g., strided average pooling) with the proposed method. As a consequence, this pooling module can only work on classification tasks.
2. My another concern is the practical value of the proposed method.
Although this method has a good motivation, it’s still quite challenging to incorporate it into a general deep learning architecture and achieve comparable performance to state-of-the-arts.
3. In the point cloud classification case, the performance of the proposed method is much lower than existing point-based methods. E.g., PointNet without the spatial transformer can achieve more than 87% on this dataset.
In addition, ModelNet40 is a small scale dataset, so testing with the SOTA encoders (DGCNN or KPConv) will be more convincing.
4. In the graph classification case, the proposed method does not show consistent improvements over its counterparts, either.
5.
Similarly, in the image classification case, the method is only comparable to FSPool.
Also, I am confused why the method does not evaluate on the most common datasets such as CIFAR and/or ImangeNet.
No
We would like to genuinely thank you for your service to the community as a reviewer amidst these troubling times. Below we provide specific answers to your comments.
* **The proposed method is only applicable to global level pooling.** While in this work, we have focused on global representations for input samples (such as point-clouds, graphs, and images), we would like to clarify that our method is not necessarily limited to global pooling.
Indeed, our approach is a generic mechanism for embedding an input set into a fixed-dimensional representation and, therefore, one can also use it for local pooling. This setting, while interesting, is not the focus of the current work, but we will mention it in the camera-ready version of the manuscript as a promising avenue for future work.
* **The approach is not applicable to general deep learning architectures.** We respectfully disagree with the reviewer, as our proposed PSWE method is indeed a global pooling mechanism that can be used in any modern deep learning architecture as demonstrated in our experiments and the implementation code included in the Supplementary Material.
* **The perfomance on point-cloud datasets is significantly lower than SoA.** This is certainly true.
But again, we believe that the big picture is completely missed here. Since we are proposing a generic global pooling mechanism, it is only fair to have a head-to-head comparison with other global pooling approaches including GAP, and PMA, regardless of the type of backbone used for the experiment. This is precisely what we have done in our work.
* **In graph classification, PSWE does not provide a consistent improvement.** First, we clarify that FSPool can be considered as a special case of our proposed approach as pointed out in lines 98-99.
More precisely, FSPool is a special case of PSWE for which the reference is frozen to the uniform distribution and the slicers are frozen to identity. Between FSPool and PSWE, we can see that there is a consistent improvement over alternative methods like GAP and PMA.
In addition, our work provides a theoretical foundation for the otherwise heuristic pooling mechanisms, including FSPool.
* **Image Recognition results.** The reason we did not test our method on ImageNet is purely legalistic. As a commercial organization, we do not have the legal permission to use the dataset in our research.
Regarding the CIFAR dataset, given the small size of images in this dataset, the empirical distributions (i.e., the output of the convolutional layers before global pooling) have very few samples (e.g., $2\times 2$ in $256$-dimensional space), which is unsuitable for any distributional modeling approach considering higher orders of statistics. We will add this information in the camera-ready version of the manuscript.
We hope that we have addressed the reviewer's questions and clarified any confusions. Thank you again, we hope that you reconsider your rating, and we would be happy to discuss any further questions you may have.
Thank you for providing the response and additional discussion!
However, I still have three major questions: 1. more evidences that this work can be applied to local pooling should be presented. I encourage authors to conduct such an experiment and compare against max/average pooling. 2. This current result on ModelNet40 is way lower than SOTA; it's even lower than vanilla PointNet. Can authors provide more details on the architecture used for ModelNet40? 3. ModelNet40 is also under academic-used-only restriction. Missing results on CIFAR/ImageNet makes me worry about the empirical value of this work.
So I am leaning to rejection given the current concerns.
We thank the reviewer for the additional questions and would like to further clarify the raised points below:
* Regarding the local pooling, as we mentioned in our rebuttal, while this is doable, we do not find the suggested experiment relevant to the current paper. In our work, we focus on the problem of permutation-invariant pooling for sets of features.
Therefore, respectfully, we believe the application of our method to local pooling is beyond the scope of this paper, and we leave it for future work.
* We use two types of backbones for the ModelNet40 experiments, the complete details of which can be found in Section 2.3 (lines 27-46) in the Supplementary Material. As we mentioned in our initial response, our goal is not to provide state-of-the-art classification results on ModelNet40, but to demonstrate the effectiveness of our geometrically-interpretable and theoretically-grounded pooling mechanism in this dataset, alongside several other datasets with different data modalities. Based on our numerical results, we believe our pooling method can be readily paired with more sophisticated backbones if the primary goal is to push the SOTA performance on the ModelNet40 dataset.
* We reiterate that our method is most interesting in scenarios where the sets have numerous elements. In the vision domain, this translates to input images with higher resolution.
Therefore, low-resolution datasets such as CIFAR do not demonstrate the effectiveness of PSWE, and methods as simple as global average pooling (GAP) should work sufficiently well in those cases.
As for ImageNet, we agree that it would be invaluable to show the effectiveness of our method on this dataset. To date, our legal department has strongly advised us against using this dataset in our studies.
However, we will follow up with them to see if they believe we can use it for this work, in which case, we will add the experiments to the camera-ready version of the manuscript accordingly. We believe that our current experiments very well demonstrate the empirical value of our proposed method in a variety of scenarios with different backbones, modalities, and datasets.
We would be happy to continue the discussion to address any remaining concerns you might have.
I acknowledge the author's efforts to respond to my questions; I increase my rating slightly.
However, my concerns are not fully addressed: 1. Modelnet40 is not a large dataset and I think it's more convincing to demonstrate the practical value by reaching a reasonable performance on this dataset. 2. The scale of the experiments in the 2D vision domain is also too small.
That said I am mainly concerned the practical value of the this work.

Meta-review:
role: PC
data: comment: Reviewers have strong disagreements about this paper. After discussions, I concur with several
points adressed by the reviewers. Most importantly, we believe that the experimental studies need
to be made stronger, in term of comparison with SOTA methods (even if they do not reach SOTA), 
potential way for local pooling... 
In addition, I would like to mention the missing stronger discussion with the paper of Mialon 2021 et al.
on "trainable OT embedding ... ". Indeed, From my point of view, the proposed method is just a variant of 
the one of Mialon where the Wasserstein distance is replaced by a 1D Wasserstein and you consider multiple references
where one reference is obtained from each slicing (multiple references  is described in Section 3.5 in Mialon et al). 
An averaging over the row would give an embedding similar to those of Mialon. Main difference is that the input distribution is 
somehow pertubed as it is also projected through the slicing. Hence, I would have expected a comparison also
with this method.

Decision: Accept (Poster)