An Empirical Study of Self-Supervised Learning with Wasserstein Distance

Area of Science:

• Machine Learning
• Computer Vision

Background:

• Self-supervised learning (SSL) commonly uses cosine similarity.
• Wasserstein distance, specifically Tree-Wasserstein distance (TWD), is less explored in SSL.
• Training Wasserstein distance can be numerically challenging.

Purpose of the Study:

• Investigate optimizing SSL with TWD.
• Identify stable training procedures for TWD in SSL.
• Evaluate TWD's effectiveness against cosine similarity in representation learning.

Main Methods:

• Employed two TWD types: total variation and ClusterTree.
• Tested various probability models: softmax, ArcFace, and simplicial embedding.
• Introduced Jeffrey divergence-based regularization for optimization stability.

Main Results:

• A softmax and TWD combination outperformed standard SimCLR on benchmark datasets (STL10, CIFAR10/100, SVHN).
• TWD with SimSiam failed to train, indicating sensitivity to model combinations.
• Jeffrey divergence regularization significantly aided model training.

Conclusions:

• Model performance in SSL with TWD is highly dependent on the chosen probability model.
• TWD, when appropriately combined with probability models and regularization, surpasses cosine similarity-based methods for representation learning.