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.

翻译：本研究深入探讨了利用树结构上的1-瓦瑟斯坦距离（即树瓦瑟斯坦距离TWD）进行自监督学习的问题，其中TWD定义为两个树嵌入向量之间的L1距离。在自监督学习方法中，余弦相似度常被用作目标函数，但关于瓦瑟斯坦距离的应用尚未得到充分研究。训练瓦瑟斯坦距离在数值上具有挑战性，因此本研究通过实证方法探索了优化瓦瑟斯坦距离自监督学习的策略，并发现了稳定的训练流程。具体而言，我们评估了两种TWD（全变差和ClusterTree）与多种概率模型（包括softmax函数、ArcFace概率模型和单纯形嵌入）的组合。我们提出了一种简单而有效的基于杰弗里散度的正则化方法来稳定优化过程。通过在STL10、CIFAR10、CIFAR100和SVHN上的实证实验，我们发现softmax函数与TWD的简单组合相比标准SimCLR获得了显著更低的结果。此外，TWD与SimSiam的简单组合无法训练模型。我们发现模型性能取决于TWD与概率模型的组合方式，而杰弗里散度正则化有助于模型训练。最后，我们证明适当的TWD与概率模型组合优于基于余弦相似度的表征学习。