Pooling by Sliced-Wasserstein Embedding

Navid Naderializadeh · Joseph F Comer · Reed Andrews · Heiko Hoffmann · Soheil Kolouri

Keywords: [ Graph Learning ] [ Vision ] [ Machine Learning ] [ Optimal Transport ] [ Representation Learning ]


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

Chat is not available.