# An Empirical Study of Self-supervised Learning with Wasserstein Distance

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 39 of its 61 references that are in the map.
- First posted on arXiv 16 Oct 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Improved Baselines with Momentum Contrastive Learning](https://mlsanity.com/paper/2003.04297/) (2020)
- [Barlow Twins: Self-Supervised Learning via Redundancy Reduction](https://mlsanity.com/paper/2103.03230/) (2021)
- [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://mlsanity.com/paper/1911.11763/) (2019)
- [MoverScore: Text Generation Evaluating with Contextualized Embeddings and Earth Mover Distance](https://mlsanity.com/paper/1909.02622/) (2019)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Max-Sliced Wasserstein Distance and its use for GANs](https://mlsanity.com/paper/1904.05877/) (2019)
- [Subspace Robust Wasserstein Distances](https://mlsanity.com/paper/1901.08949/) (2019)
- [Sliced Wasserstein Kernels for Probability Distributions](https://mlsanity.com/paper/1511.03198/) (2015)
- [Tree-Sliced Variants of Wasserstein Distances](https://mlsanity.com/paper/1902.00342/) (2019)
- [Understanding and Constructing Latent Modality Structures in Multi-modal Representation Learning](https://mlsanity.com/paper/2303.05952/) (2023)
- [Fair and Optimal Classification via Post-Processing](https://mlsanity.com/paper/2211.01528/) (2022)
- [Word Rotator's Distance](https://mlsanity.com/paper/2004.15003/) (2020)
- [A Note on Connecting Barlow Twins with Negative-Sample-Free Contrastive Learning](https://mlsanity.com/paper/2104.13712/) (2021)
- [Scalable Nearest Neighbor Search for Optimal Transport](https://mlsanity.com/paper/1910.04126/) (2019)
- [Simplicial Embeddings in Self-Supervised Learning and Downstream Classification](https://mlsanity.com/paper/2204.00616/) (2022)
- [Diffusion Earth Mover's Distance and Distribution Embeddings](https://mlsanity.com/paper/2102.12833/) (2021)
- [Sobolev Transport: A Scalable Metric for Probability Measures with Graph Metrics](https://mlsanity.com/paper/2202.10723/) (2022)
- [Fast Unbalanced Optimal Transport on a Tree](https://mlsanity.com/paper/2006.02703/) (2020)
- [Entropy Partial Transport with Tree Metrics: Theory and Practice](https://mlsanity.com/paper/2101.09756/) (2021)

## Abstract (arXiv)

In this study, we delve into the problem of self-supervised learning (SSL) utilizing the 1-Wasserstein distance on a tree structure (a.k.a., Tree-Wasserstein distance (TWD)), where TWD is defined as the L1 distance between two tree-embedded vectors. In SSL methods, the cosine similarity is often utilized as an objective function; however, it has not been well studied when utilizing the Wasserstein distance. Training the Wasserstein distance is numerically challenging. Thus, this study empirically investigates a strategy for optimizing the SSL with the Wasserstein distance and finds a stable training procedure. More specifically, we evaluate the combination of two types of TWD (total variation and ClusterTree) and several probability models, including the softmax function, the ArcFace probability model, and simplicial embedding. We propose a simple yet effective Jeffrey divergence-based regularization method to stabilize optimization. Through empirical experiments on STL10, CIFAR10, CIFAR100, and SVHN, we find that a simple combination of the softmax function and TWD can obtain significantly lower results than the standard SimCLR. Moreover, a simple combination of TWD and SimSiam fails to train the model. We find that the model performance depends on the combination of TWD and probability model, and that the Jeffrey divergence regularization helps in model training. Finally, we show that the appropriate combination of the TWD and probability model outperforms cosine similarity-based representation learning.