# Constrained Sliced Wasserstein Embedding

arXiv: https://arxiv.org/abs/2506.02203 · HTML: https://mlsanity.com/paper/2506.02203/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 72 of its 117 references that are in the map.
- First posted on arXiv 2 Jun 2025.

## Builds on (most cited first, 30 of 72)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [PCT: Point cloud transformer](https://mlsanity.com/paper/2012.09688/) (2020)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Geometric Dataset Distances via Optimal Transport](https://mlsanity.com/paper/2002.02923/) (2020)
- [Primal Wasserstein Imitation Learning](https://mlsanity.com/paper/2006.04678/) (2020)
- [Max-Sliced Wasserstein Distance and its use for GANs](https://mlsanity.com/paper/1904.05877/) (2019)
- [Wasserstein Weisfeiler-Lehman Graph Kernels](https://mlsanity.com/paper/1906.01277/) (2019)
- [Wasserstein-2 Generative Networks](https://mlsanity.com/paper/1909.13082/) (2019)
- [Subspace Robust Wasserstein Distances](https://mlsanity.com/paper/1901.08949/) (2019)
- [Statistical and Topological Properties of Sliced Probability Divergences](https://mlsanity.com/paper/2003.05783/) (2020)
- [Generative Modeling with Optimal Transport Maps](https://mlsanity.com/paper/2110.02999/) (2021)
- [Lagrangian Duality for Constrained Deep Learning](https://mlsanity.com/paper/2001.09394/) (2020)
- [Distributional Sliced-Wasserstein and Applications to Generative Modeling](https://mlsanity.com/paper/2002.07367/) (2020)
- [SoftSort: A Continuous Relaxation for the argsort Operator](https://mlsanity.com/paper/2006.16038/) (2020)
- [Recent Advances in Optimal Transport for Machine Learning](https://mlsanity.com/paper/2306.16156/) (2023)
- [Sliced Wasserstein Kernels for Probability Distributions](https://mlsanity.com/paper/1511.03198/) (2015)
- [Wasserstein Adversarial Imitation Learning](https://mlsanity.com/paper/1906.08113/) (2019)
- [Sliced Wasserstein Distance for Learning Gaussian Mixture Models](https://mlsanity.com/paper/1711.05376/) (2017)
- [Wasserstein Embedding for Graph Learning](https://mlsanity.com/paper/2006.09430/) (2020)
- [Avoiding Overfitting: A Survey on Regularization Methods for Convolutional Neural Networks](https://mlsanity.com/paper/2201.03299/) (2022)
- [Achieving Zero Constraint Violation for Constrained Reinforcement Learning via Primal-Dual Approach](https://mlsanity.com/paper/2109.06332/) (2021)
- [Probably Approximately Correct Constrained Learning](https://mlsanity.com/paper/2006.05487/) (2020)
- [Constrained Learning with Non-Convex Losses](https://mlsanity.com/paper/2103.05134/) (2021)

## Built on by (most cited first, 2 of 2)

- [AL-CoLe: Augmented Lagrangian for Constrained Learning](https://mlsanity.com/paper/2510.20995/) (2025)
- [Sliced Inner Product Gromov-Wasserstein Distances](https://mlsanity.com/paper/2605.08546/) (2026)

## Abstract (arXiv)

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.