Understanding Learning with Sliced-Wasserstein Requires Rethinking Informative Slices (2411.10651v1)

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?"

• The paper proposes a novel φ-weighting formulation to adjust 1D sliced-Wasserstein distances by enhancing slice informativeness.
• The authors introduce a global rescaling factor that simplifies SWD computations for high-dimensional data with low-dimensional support.
• Empirical results demonstrate that the modified classical SWD achieves performance comparable to or surpassing advanced variants without extra computational cost.

Understanding Learning with Sliced-Wasserstein Requires Rethinking Informative Slices

Introduction

The paper of Wasserstein distances (WDs) within optimal transport theory has been imperative in various machine learning applications, particularly for comparing data distributions. Due to the computational intensity and sample complexity of traditional WDs, the introduction of Sliced-Wasserstein Distances (SWDs) has offered an efficient proxy, leveraging projections onto one-dimensional (1D) subspaces. Despite this, the concentration of measure phenomenon poses a significant challenge in high-dimensional spaces, leading to most random projections being uninformative. This paper revisits the approach of leveraging these slices and investigates a novel method of rescaling them to ensure informativeness.

Main Contributions

The authors propose a unified formulation for rescaling SWDs that rethinks the conventional modification of slicing distributions. In high-dimensional spaces with data supported on lower-dimensional subspaces, a single global scaling factor can adjust the informativeness of slices, simplifying the SWD computations. The key insights are:

1. The $\phi$-Weighting Formulation: The authors introduce a novel approach that involves weighting the contribution of each 1D sliced Wasserstein distance based on a predefined informativeness function, rather than modifying the slicing distribution directly.
2. Global Rescaling Factor: The paper illustrates that with a sound assumption of the low-dimensional support of data, rescaling 1D Wasserstein slices simplifies to applying a single scaling factor to the SWD. This universal constant acts as an implicit reweighting mechanism, correcting the contribution based on informativeness.
3. Implications for ML Workflows: By aligning the learning rates with this notion of slice informativeness, the standard SWD can achieve comparable or superior performance to more advanced methods without compromising computational efficiency or stability. This approach readily integrates with existing workflows due to the use of conventional learning rate optimization practices.

Theoretical Analysis

Under the assumption that $d$-dimensional data inherently lies within much lower $k$-dimensional subspaces, the paper demonstrates analytically that every random projection’s contribution is implicitly down-weighted by its informativeness relative to the subspace. Moreover, the section establishes the Effective Subspace Scaling Factor (ESSF), which correlates the standard SWD to the intrinsic dimensions of data distribution.

Experimental Validation

Extensive experimentation across various machine learning tasks—ranging from gradient flow tasks on synthetic datasets to color transfer and generative modeling—substantiates the theoretical findings. The results illustrate that a well-tuned classical SWD, incorporating the proposed scaling modifications, can rival—or even surpass—the performance of more sophisticated SWD variants while retaining computational tractability.

Conclusion and Future Work

This paper challenges the pursuit of more intricate SWD variants by highlighting the potential of a properly configured classical SWD approach to attain high performance. The flexibility in defining informativeness functions opens avenues for future research to explore alternative reweighting functions tailored to specific data structures or applications. Future research could focus on optimizing these insights across different domains, potentially with empirical parameter tuning and broader applicability, to ensure maximal utilization of SWD's computational efficiency and rich theoretical properties.

By revisiting the foundations of slice selection and re-weighting strategies, the paper contributes a distinct perspective on the ongoing development and application of sliced optimal transport methodologies in machine learning.