Set Representation Learning with Generalized Sliced-Wasserstein Embeddings
This addresses the need for effective set representation learning in machine learning, though it appears incremental as it builds on existing optimal transport concepts.
The paper tackles the problem of learning representations from set-structured data by proposing a geometrically-interpretable framework based on optimal mass transportation, specifically using Generalized Sliced Wasserstein distances for exact Euclidean embeddings, and demonstrates superiority over state-of-the-art methods in supervised and unsupervised tasks.
An increasing number of machine learning tasks deal with learning representations from set-structured data. Solutions to these problems involve the composition of permutation-equivariant modules (e.g., self-attention, or individual processing via feed-forward neural networks) and permutation-invariant modules (e.g., global average pooling, or pooling by multi-head attention). In this paper, we propose a geometrically-interpretable framework for learning representations from set-structured data, which is rooted in the optimal mass transportation problem. In particular, we treat elements of a set as samples from a probability measure and propose an exact Euclidean embedding for Generalized Sliced Wasserstein (GSW) distances to learn from set-structured data effectively. We evaluate our proposed framework on multiple supervised and unsupervised set learning tasks and demonstrate its superiority over state-of-the-art set representation learning approaches.