Understanding Learning with Sliced-Wasserstein Requires Rethinking Informative Slices

TL;DR

This paper revisits classical Sliced-Wasserstein distances, proposing a simple rescaling method that makes all slices equally informative, leading to improved performance and easier integration into machine learning workflows.

Contribution

The paper introduces a data-dependent rescaling of the classical SWD, simplifying the approach to focus on a single global scaling factor, enhancing its effectiveness and theoretical properties.

Findings

Rescaling classical SWD improves performance in learning tasks.

A single global scaling factor suffices under certain data assumptions.

Classical SWD can match or outperform complex variants with proper configuration.

Abstract

The practical applications of Wasserstein distances (WDs) are constrained by their sample and computational complexities. Sliced-Wasserstein distances (SWDs) provide a workaround by projecting distributions onto one-dimensional subspaces, leveraging the more efficient, closed-form WDs for one-dimensional distributions. However, in high dimensions, most random projections become uninformative due to the concentration of measure phenomenon. Although several SWD variants have been proposed to focus on \textit{informative} slices, they often introduce additional complexity, numerical instability, and compromise desirable theoretical (metric) properties of SWD. Amidst the growing literature that focuses on directly modifying the slicing distribution, which often face challenges, we revisit the classical Sliced-Wasserstein and propose instead to rescale the 1D Wasserstein to make all slices equally informative. Importantly, we show that with an appropriate data assumption and notion of \textit{slice informativeness}, rescaling for all individual slices simplifies to \textbf{a single global scaling factor} on the SWD. This, in turn, translates to the standard learning rate search for gradient-based learning in common machine learning workflows. We perform extensive experiments across various machine learning tasks showing that the classical SWD, when properly configured, can often match or surpass the performance of more complex variants. We then answer the following question: "Is Sliced-Wasserstein all you need for common learning tasks?"