Constrained Sliced Wasserstein Embedding

TL;DR

This paper proposes a constrained learning method to optimize slicing directions in Sliced Wasserstein distances, improving the efficiency and informativeness of high-dimensional probability measure comparisons.

Contribution

It introduces a gradient-based primal-dual approach to learn meaningful slicing directions, reducing the number of slices needed and enhancing performance in embedding comparisons.

Findings

Improved slicing directions lead to better embedding representations.

Numerical results show enhanced performance on images, point clouds, and protein sequences.

The method reduces computational complexity in high-dimensional Wasserstein distance calculations.

Abstract

Sliced Wasserstein (SW) distances offer an efficient method for comparing high-dimensional probability measures by projecting them onto multiple 1-dimensional probability distributions. However, identifying informative slicing directions has proven challenging, often necessitating a large number of slices to achieve desirable performance and thereby increasing computational complexity. We introduce a constrained learning approach to optimize the slicing directions for SW distances. Specifically, we constrain the 1D transport plans to approximate the optimal plan in the original space, ensuring meaningful slicing directions. By leveraging continuous relaxations of these transport plans, we enable a gradient-based primal-dual approach to train the slicer parameters, alongside the remaining model parameters. We demonstrate how this constrained slicing approach can be applied to pool high-dimensional embeddings into fixed-length permutation-invariant representations. Numerical results on foundation models trained on images, point clouds, and protein sequences showcase the efficacy of the proposed constrained learning approach in learning more informative slicing directions. Our implementation code can be found at https://github.com/Stranja572/constrainedswe.