Sliced Wasserstein Distance for Improving Class Separability: Experiments and Reflections

Can We Make Classes More "Clearly Separated" in Feature Space?

In machine learning tasks, feature separability often directly determines the upper bound of a classifier's performance. If two classes highly overlap in feature space, even the most powerful model will struggle to draw a clear decision boundary. Recently, a Reddit developer shared an inspiring experiment: can we learn a feature transformation that makes the distributions of two classes more separated, without destroying the underlying geometric structure of the data?

The core tool behind this idea is the Sliced Wasserstein Distance (SWD). The author attempted to learn a transformation T(x) via a neural network that maximizes the SWD between two class distributions after transformation. The experiment on the Breast Cancer dataset yielded interesting yet contradictory results — it performed well on decision trees but poorly on other algorithms.

Core Principles of Sliced Wasserstein Distance

From Optimal Transport to Distribution Metrics

The Wasserstein distance (also known as Earth Mover's Distance) originates from optimal transport theory and measures the "difference" between two probability distributions. Intuitively, it represents the minimum cost of moving one pile of dirt from one distribution's shape into another distribution's shape. Compared to metrics like KL divergence, the Wasserstein distance is more sensitive to the geometric positions of distributions and can provide meaningful measurements even when two distributions have no overlap.

More precisely from a mathematical perspective, the p-th order Wasserstein distance is defined as the infimum of the expected transport cost over all joint distributions (transport plans) that map distribution μ to distribution ν. This definition stems from the optimal transport problem proposed by the French mathematician Monge in the 18th century, later given a dual relaxation form by Kantorovich in the 1940s, transforming it from a purely mathematical problem into a computable optimization problem. The reason Wasserstein distance is so favored in machine learning is that it is a true metric (satisfying symmetry and the triangle inequality), whereas KL divergence is asymmetric and becomes infinite when distribution supports don't overlap, and Jensen-Shannon divergence, while symmetric, has gradient signals that approach zero when distributions are separated. It is precisely these desirable properties that enabled Arjovsky et al. to propose WGAN in 2017 to solve the training instability of traditional GANs, and that provide a solid theoretical foundation for the feature space optimization in this article.

Why "Slice" — Efficient Approximation for High-Dimensional Computation

Directly computing Wasserstein distance in high-dimensional spaces is extremely expensive — for n sample points, exact computation requires solving a linear programming problem with complexity O(n³log n), and faces the curse of dimensionality in high dimensions. Sliced Wasserstein Distance cleverly circumvents this problem: it randomly projects high-dimensional distributions onto many one-dimensional directions, computes the Wasserstein distance on each one-dimensional direction (which has a closed-form solution in 1D — just sorting, with complexity only O(n log n)), and then averages across all projection directions. This preserves the core properties of optimal transport while dramatically reducing computational complexity, making it ideal as a differentiable loss function in neural network training.

The theoretical foundation of Sliced Wasserstein Distance comes from the Radon transform — an integral transform that decomposes high-dimensional functions into one-dimensional projections along all directions. The existence of its inverse transform guarantees that the original high-dimensional distribution's information can be fully recovered from the set of one-dimensional projections. Bonneel et al. proved in 2015 the topological equivalence between SWD and the true Wasserstein distance: SWD equals zero if and only if the true Wasserstein distance equals zero, and both induce the same weak convergence topology. This theoretical guarantee means SWD is not merely an "approximation" but a complete metric in the sense of distribution comparison. In practice, SWD has been successfully applied in Sliced Wasserstein GAN (Deshpande et al., 2018), texture synthesis, domain adaptation, and many other scenarios, demonstrating its effectiveness as a training objective.

Because SWD is differentiable and computationally efficient, it becomes the ideal optimization objective for "pulling apart" two class distributions.

Experimental Design: Finding Balance Between Separation and Fidelity

Core Optimization Approach

The author's method can be summarized as a clear pipeline:

Raw data → Learn transformation → Maximize inter-class SWD under geometric constraints → Standard ML pipeline

Specifically, given two class distributions A and B, the goal is to learn a transformation T(x) such that:

max SWD(T(A), T(B))

The Necessity of Geometric Constraints

A key engineering insight emerges here. If the neural network merely maximizes SWD without constraints, it will "cheat" — distorting the data in arbitrary ways as long as it pushes the two distributions further apart. Such unconstrained optimization may cause SWD values to skyrocket, but it completely destroys the intrinsic structure of the data, leaving downstream classifiers helpless.

To address this, the author introduced geometric constraints, requiring the transformation to preserve the data's original geometric relationships as much as possible while increasing inter-class distance. This is essentially a trade-off between "separability" and "fidelity" — making classes more separable without rendering the data unrecognizable.

This approach has deep theoretical roots in metric learning and manifold learning. Classical isometric mapping methods (like ISOMAP) require low-dimensional embeddings to preserve geodesic distances in high-dimensional space; triplet loss in contrastive learning pulls same-class samples closer while pushing different-class samples apart, but implicitly prevents space collapse through margin constraints. A more direct constraint form is Lipschitz continuity — requiring transformation T to satisfy ||T(x₁) - T(x₂)|| ≤ L·||x₁ - x₂||, thereby limiting the maximum stretching of the transformation on space. This constraint is widely adopted in Wasserstein GAN discriminator design (implemented via spectral normalization or gradient penalty), and the geometric constraints in this experiment can be seen as a concrete application of the same philosophical idea in the feature preprocessing scenario.

Result Analysis: Why Does Only the Decision Tree Succeed?

Why Decision Trees Performed Well

Experimental results showed that on the Breast Cancer dataset, features transformed via SWD performed well on decision trees. There's a reasonable explanation: decision trees make decisions through axis-aligned splits along coordinate axes, and they are insensitive to monotonic transformations or local stretching of features — as long as the ordering relationship between classes on a feature is strengthened, decision trees benefit. SWD maximization likely strengthens class separation ordering on certain dimensions.

On a deeper level, when decision trees select the optimal splitting feature at each node, the criteria used (whether information gain, Gini impurity, or mean squared error) are essentially functions of order statistics — they only care about "on each side of a threshold, what are the proportions of each class," not the actual distances between data points. This means decision trees have a natural invariance to any monotonic transformation that preserves class ordering. For example, applying transformation f(x) = x³ to feature x, as long as f is monotonic, the tree's split results remain completely unchanged. SWD maximization on one-dimensional projection directions is essentially widening the ordering gap between two classes — when the one-dimensional Wasserstein distance increases, the quantile differences between samples of both classes in that direction increase, directly corresponding to decision trees finding purer split points on that feature. This explains why a method optimized at the distribution level happens to be perfectly suited for order-based learners like decision trees.

The Deeper Reasons Other Classifiers Failed

What truly deserves reflection is: why does this method work well for decision trees but perform terribly for logistic regression, SVM, KNN, and other algorithms?

Algorithms that rely on distance metrics (like KNN, SVM with RBF kernel) or those that assume linear separability (like logistic regression) are extremely sensitive to global geometric structure. When the transformation distorts the overall topology of feature space to increase SWD, the distance relationships and linear relationships these algorithms rely on are destroyed, and performance naturally collapses.

Specifically, KNN predictions are entirely based on the local topological relationship of "who are the nearest neighbors." If transformation T greatly stretches space in some regions while compressing it in others, originally neighboring same-class samples may be pushed apart while different-class samples are pulled closer, degrading neighborhood voting results. The RBF kernel of SVM, K(x,y) = exp(-γ||x-y||²), is highly sensitive to Euclidean distance — spatial distortion directly changes the eigenstructure of the kernel matrix, potentially making data that was linearly separable in kernel space become inseparable. Logistic regression assumes that the log-odds of class posterior probability is a linear function of features — a nonlinear transformation T may destroy this linear relationship. Even if the distribution centers of two classes are pushed further apart, if the decision boundary becomes highly nonlinear, the linear model of logistic regression becomes even harder to fit.

This demonstrates that "more separated in distribution" does not equal "more friendly to all classifiers." SWD measures overall differences at the distribution level, while different classifiers care about separability at different levels.

Key Insights and Future Exploration Directions

The Alignment Problem Between Optimization Objectives and Downstream Tasks

Perhaps the greatest value of this experiment lies not in whether it "succeeded," but in the universal problem it exposed: the misalignment between optimization objectives and final evaluation criteria. The author optimized inter-class SWD, but the true goal was to improve classification accuracy. The gap between the two causes the optimization to work for some models but fail for others.

This is actually a microcosm of the surrogate loss problem in machine learning. Since classification accuracy itself is a non-differentiable step function, we must use differentiable surrogate objectives as replacements — cross-entropy loss is a surrogate for classification accuracy, hinge loss is a convex relaxation of 0-1 loss, and SWD is a distribution-level surrogate for "class separability." The consistency between surrogate loss and true objective (Fisher consistency) is a core problem in statistical learning theory: only when minimizing the surrogate loss guarantees that the true loss also approaches optimality is the surrogate "safe." The lesson from this experiment indicates that SWD, as a surrogate for classification accuracy, lacks this consistency guarantee. The reason end-to-end learning paradigms (like backpropagating classification loss directly to feature extraction layers in deep learning) succeed is precisely because they eliminate intermediate surrogate objectives, letting feature representations serve the final task directly. This also reminds us that any two-stage approach of "first optimize representation, then train classifier" needs careful verification of the alignment between its surrogate objective and the downstream task.

Degeneration Risk of Unconstrained Optimization

The experiment also reaffirms a classic lesson in deep learning: given enough degrees of freedom, neural networks will always find shortcuts to "game the system." Simply maximizing some distance metric often yields degenerate solutions. Introducing geometric constraints and regularization terms essentially tells the model "what constitutes meaningful separation."

Promising Improvement Directions Worth Exploring

Starting from this experiment, there are several avenues for follow-up research:

• Refined constraint design: One could try preserving local neighborhood structure (similar to Locally Linear Embedding in manifold learning, or preserving KL divergence of conditional probabilities as in t-SNE) or introducing Lipschitz constraints to limit the degree of transformation distortion (e.g., constraining the spectral norm of each network layer to within 1 via spectral normalization).
• Joint optimization tailored to specific classifiers: Rather than optimizing a generic SWD, it's better to incorporate the downstream classifier's loss into end-to-end training. For example, one could construct a joint loss L = L_classifier - λ·SWD, letting the feature transformation simultaneously serve distribution separation and classification accuracy.
• Validation on more datasets: Testing only on the relatively simple Breast Cancer binary classification dataset (30 features, 569 samples) limits the generalizability of conclusions. Validation is needed on more complex, higher-dimensional data, such as deep feature spaces of image data or word embedding spaces in natural language processing.

Conclusion

This is a typical exploration where "negative results also have value." Using SWD to reshape feature space, while not yielding a one-size-fits-all solution for every algorithm, clearly reveals the subtle and complex relationship between distribution-level metrics and different classifiers. For researchers, experiments with this interweaving of failures and partial successes often inspire more thinking than a polished SOTA number — reminding us that in machine learning, a "better metric" must be tied to "specific tasks and models" to have real meaning.