Understanding Learning with Sliced-Wasserstein Requires Re-thinking Informative Slices

Authors: Huy Tran, Ashkan Shahbazi, Yikun Bai, John R. Hershey, Soheil Kolouri
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Optimal Transport, Sliced Wasserstein, Concentration of Measure

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 classic Sliced-Wasserstein and propose instead to rescale the 1D Wasserstein to make all slices equally informative. Importantly, we show that with an appropriate 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 ML workflows. We perform extensive experiments across various machine learning tasks showing that the classic 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?

Decision: reject
Reviewer scores: 8, 6, 5

Reviews and discussion:
<thread>
creation date:
1730626503034

thread_type:
official_review

<thread_object>
creator: reviewer_tzmm

<content>
summary:
This paper proposes a simplified approach to the Sliced-Wasserstein Distance (SWD) by rescaling the Wasserstein distance for each one-dimensional slice, ensuring that all slices carry the same informativeness. This method demonstrates strong performance in tasks such as image generation and color transfer, rivaling or even surpassing more complex SWD variants. The core innovation lies in using a global scaling factor to balance the informativeness of different slices, thereby avoiding the complexity of directly modifying slice distributions. Overall, the idea is innovative, and the experimental results are convincing, showcasing the potential of the classic SWD with suitable parameter settings.

strengths:
1. Innovation: The rescaling method proposed in this paper offers a new approach for high-dimensional SWD applications. By using a global scaling factor to address the issue of selecting informative slices, it effectively simplifies the complex process of slice selection.  
2. Theoretical support: Under the assumption of an effective subspace, this paper provides a theoretical derivation of the rescaling strategy, offering mathematical support that explains the rationale behind this method.  
3. Experimental performance: In various machine learning tasks, the rescaled SWD method performs excellently, demonstrating that the performance gap between classic SWD and its complex variants can be bridged through appropriate hyperparameter tuning.

weaknesses:
1. Limitation of assumption: The theoretical analysis of this method relies on the assumption that data is distributed within a low-dimensional subspace. However, this assumption may not hold in certain high-dimensional data, particularly in datasets with complex or high-dimensional structures, such as text data or more intricate time series data, where the method's effectiveness remains unclear. Therefore, its applicability may be limited to datasets that satisfy this assumption.
2. Single metric of informativeness: The proposed global scaling factor only considers the informativeness of the slices in terms of their projection onto the effective subspace. However, different tasks may involve slices with more complex characteristics, and this single metric may not fully capture the contributions of all slices. Incorporating task-specific measures of informativeness could potentially enhance the method’s applicability and robustness.
3. Limited breadth of experimental validation: Although the paper validates the method’s effectiveness on multiple visual datasets, the types of experimental data are relatively narrow, lacking validation on other types of data, such as text or audio. Additionally, the paper does not deeply analyze performance on extremely high-dimensional data, which is a significant motivation behind designing SWD variants.
4. Limited generalization: The method performs well on classic machine learning tasks, but its applicability in more complex deep generative models, such as large-scale language models, remains unclear. Further research is needed to determine whether the method maintains its efficiency and stability on complex data distributions.
5. Sensitivity to hyperparameters in experiments: The proposed scaling method relies heavily on the choice of learning rate. Although the authors suggest that the learning rate can be adjusted as a hyperparameter, the specific impact on model convergence and adaptability lacks detailed analysis. Questions such as whether the results are limited to certain learning rate ranges and whether reliance on learning rate could lead to instability in practice require further exploration.

questions:
This paper provides new insights into optimizing Sliced-Wasserstein Distance, with a simplified method that is theoretically appealing and performs well across several tasks. However, due to limitations in assumptions, narrow experimental validation, and sensitivity to hyperparameters, the applicability of this method may be constrained. To enhance the comprehensiveness and generalizability of the work, it is recommended to validate the method on more complex tasks and datasets and further refine the theoretical analysis.

Suggested improvements based on the Weaknesses section:
1. Strengthen discussion and validation of assumptions: Further discuss the practicality and limitations of the effective subspace assumption in real-world applications, or explore ways to expand the method’s applicability without this assumption.
2. Expand experiments to diverse data types: Validate the method on a wider range of datasets, such as text and audio data, to verify its broad applicability. Greater diversity in experiments would enhance the generalizability and persuasiveness of the results.
3. More in-depth hyperparameter sensitivity analysis: Further analyze the effect of learning rate on the method's convergence. Testing across different learning rate ranges could ensure the robustness and stability of the method and clarify if it requires strict hyperparameter tuning.
4. Explore more flexible informativeness metrics: To adapt to a wider range of tasks, consider expanding the definition of the global scaling factor to incorporate more detailed task-specific information or introduce an adaptive informativeness metric.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699191245

thread_type:
official_review

<thread_object>
creator: reviewer_3tbe

<content>
summary:
This manuscript presents a novel solution of using sliced Wasserstein distance in general learning tasks. It has solid theoretical and empirical results to support their claim.

strengths:
1. There are a lot of derivations in the manuscript, provides a strong theoretical background for the discuss topic.

2. The ideas of 1D Wasserstein distance is quite novel.

3. Detailed experiments were conducted on some of the learning tasks.

weaknesses:
1. The experiment seems not sufficient to prove the claim "Sliced-Wasserstein is all you need for common learning tasks". Some common learning tasks like image classification and object detection is not tested.

2. Larger dataset for experiment is needed. Especially, toy datasets and MNIST images, which is a pretty small and easy learning task.

3. How SWD is used during the training is not very clear, a systematic model framework or a detailed algorithm might be helpful.

4. The compared methods in the experiment only has abbreviation with no explanation. 

5. How SWD is better is not clearly reflected in the experiments.

6. There are so many important things were discussed in the Appendix due to the page limit.

questions:
1. What is the computational complexity for the proposed SWD?

2. What's the purpose of finding LR bound for SW Gradient Flow? How it impact the results for the mentioned learning tasks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730704118340

thread_type:
official_review

<thread_object>
creator: reviewer_d93u

<content>
summary:
This paper revisits the classic Sliced-Wasserstein Distance (SWD) and proposes a new approach to enhance its effectiveness. The authors introduce a rescaling approach based on the "informativeness" of each slice, which they define as the slice's alignment with the data's effective subspace. Instead of modifying the slicing distribution to identify informative slices, this approach uses a single global scaling factor, which simplifies SWD calculations and also allows it to perform comparably with other variants.
The paper provides theoretical backing for the rescaling approach, showing how it aligns with different dimensionality settings and retains SWD's metric properties.
The authors validate this method through experiments on tasks like gradient flow, color transfer, and deep generative modeling, showing that a properly scaled SWD can match or even outperform other methods without added complexity.

strengths:
The paper offers a simpler and more computationally efficient approach to improving SWD by introducing a single global rescaling factor based on slice informativeness.

The rescaling approach can be absorbed into the learning rate tuning process, making it easy to implement without extensive adjustments to existing models.

It also reduces dependency on strict dimensionality assumptions, making SWD adaptable to more tasks without requiring the knowledge of the data's effective subspace dimension k.

The authors provide a solid theoretical foundation for their approach.

weaknesses:
I'm not an expert in this area, so I'll leave the weaknesses section blank for now.

questions:
Why the strict subspace condition can be relaxed given Assumption 4.1 provides a theoretical basis?

What other criteria could be used to define informativeness, and would they offer any additional benefits over alignment with the effective subspace?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732129542796

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to common questions and concerns

comment:
Dear AC and Reviewers,

We would like to truly thank all reviewers for their valuable time and dedicated efforts to provide us with these constructive feedbacks. We are encouraged by the general endorsement of novelty from all reviewers. In this common section, we attempt to clarify our paper and provide some additional results, which we may reference in our detailed response to each reviewer. We will update the paper accordingly throughout our discussions.

---

# C1. The paper's primary contributions

**Problem**: In machine learning applications using sliced-Wasserstein, the current dominating paradigm in the literature is to design a slicing distribution that emphasizes the difference between input measures to address the projection complexity of the original Sliced-Wasserstein distances [1][2][3][4]. This often results in additional computational/memory overheads (especially for optimization-based methods [1][3]), numerical instabilities, and the loss of desirable theoretical properties (e.g., being a proper metric) [2][4].

**Proposed**: We propose a shift in perspective via the novel reweighting formulation. Instead of explicitly modifying the slicing distributions to emphasize informative slices, we rescale the contributions for all slices (e.g., less informative slices should be rescaled more). This new paradigm requires rethinking what it means for a projection to be *informative*. We discovered that this generalized notion can be dependent on prior knowledge about the data distribution. We show that, under a mild data assumption and an appropriate notion of informativeness, we can design the reweighting function to make all slices *equally informative*. This recovers the *subspace Sliced-Wasserstein* variant that is exactly a scalar multiple of the classical Sliced-Wasserstein distance between the input measures. This reduces the original problem to learning rate search (without any knowledge of the subspace in practice), which is already standard in the machine learning workflows. In the end, after an extensive search for the next SW variant, what we obtain is exactly the well-studied classical SWD with all theoretical properties preserved and no additional computational/memory overheads.

This naturally raises an important question: "Is Sliced-Wasserstein all you need for common learning tasks?". Here, by "common learning tasks", we refer to vision tasks that are often used in the existing SW literature [1][2][3][4][5]. There may exist less common settings with other data modalities (e.g., text, audio) that will not be the focus of our paper, although there are no known limitations restricting our results to vision tasks. Our choice of tasks simply allows direct comparison with existing methods within our computational budget.

Lastly, our proposed formulation provides novel insights and opens up interesting research directions to explore different data assumptions as well as different notions of slice informativeness.

# C2. The subspace assumption is reasonable in practice

Following the proposed reweighing formulation, we provide the subspace assumption upon which rest many of our theoretical contributions, Here, we aim to clarify that the effective subspace assumption is reasonable in practice because real datasets (images, audios, etc.) naturally have low-dimensional linear structure in two key ways. First, in minibatch processing, d-dimensional data of batch size $B \ll d$ mathematically cannot span more than min(2B, d), which provides a natural bound on the data effective dimensionality. Second, real features typically show strong linear correlations, as evidenced by the rapid eigenvalue decay in PCA across diverse datasets - practical success of PCA and autoencoders demonstrates that high-dimensional data can often be approximated in much lower-dimensional spaces while preserving essential structure for downstream tasks.

It is worth noting that our reweighting formulation does **not** depend on this assumption. One may make different assumptions based on prior knowledge of the data. In our paper, we find that using the subspace assumption leads to a conveninent notion of informativeness. This allows us to define the reweighting function that simplifies the math to a scalar multiple of the classical Sliced-Wasserstein in expectation. We note that this paradigm open up future research venues to explore other data assumptions and notions of informative slices, beyond the scope of our current paper. 

Lastly, in section C4 below we show that our theoretical results remain relevant even when the data lives near a low-dimensional subspace (in parallel supported by extensive experimental results), addressing the comment of the reviewers about restrictive nature of our core assumption.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732548784132

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Engaging in Final Discussions and Requesting Reevaluations as the Discussion Period Closes

comment:
Dear Reviewers,

Thank you once again for your time and dedication to maintaining the high standards of the ICLR conference. As we near the end of the discussion period, we would like to check in and engage in constructive dialogue to address any remaining concerns you may have. 

We have provided detailed responses to your thoughtful feedback and revised the paper according to your constructive comments. We would be happy to engage with you further to clarify or expand on any points. Please let us know if there are additional questions we can address.

Best, The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper revisits Sliced-Wasserstein Distances (SWDs) and addresses their relative inefficiency in high dimensions by introducing a simple global rescaling factor to enhance slice informativeness while preserving theoretical properties. The method avoids the complexity of existing SWD variants and is aimed at improving their applicability in gradient-based learning workflows. Experiments demonstrate that their variant of SWD achieves competitive  performance compared to more complex alternatives in some machine learning tasks, mostly generative. 

Through the analysis of the reviews and the discussion, I find that the paper is still lacking decisive arguments both theoretical and experimental to make this variant of SWD stand out with respect to all the sliced variants of Wasserstein, and I will mostly follow reviewer  TZmm. In order to explain further my decision, I will detail some of those limitations:
  - limited theoretical insights: the subspace assumption is insufficiently described. Assumption 4.1 assumes that both target ad source distributions live both in the same subspace, which can be perceived as a strong limitation. Equivalences with other SWD variants are not very informative. 
  - the evaluation is somehow limited, and mostly focused on SWD variants. For equivalent batch sizes, why not considering directly Wasserstein ?  In the color transfer example, the idea of limited subspace of dimension < 3 is not discussed. In the case of FFHQ: the considered dimension is the one from the latent space (512). The metrics used for comparing the different variants of SWD are not not discussed, nor the results are analyzed. Results presented in the appendix should be chosen more carefully, and do not aim for exhaustivity (tables 5 to 18)

I encourage the authors to sharpen their arguments on why their variant of SWD should be preferred to the classical sliced Wasserstein, and to design experiments that clearly show benefits.