Supervised Tree-Wasserstein Distance

Abstract

To measure the similarity of documents, the Wasserstein distance is a powerful tool, but it requires a high computational cost. Recently, for fast computation of the Wasserstein distance, methods for approximating the Wasserstein distance using a tree metric have been proposed. These tree-based methods allow fast comparisons of a large number of documents; however, they are unsupervised and do not learn task-specific distances. In this work, we propose the Supervised Tree-Wasserstein (STW) distance, a fast, supervised metric learning method based on the tree metric. Specifically, we rewrite the Wasserstein distance on the tree metric by the parent-child relationships of a tree and formulate it as a continuous optimization problem using a contrastive loss. Experimentally, we show that the STW distance can be computed fast, and improves the accuracy of document classification tasks. Furthermore, the STW distance is formulated by matrix multiplications, runs on a GPU, and is suitable for batch processing. Therefore, we show that the STW distance is extremely efficient when comparing a large number of documents.

References in corpus (5)

• Distributed Representations of Words and Phrases and their Compositionality
• Large Batch Training of Convolutional Networks
• From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification
• Hierarchical Optimal Transport for Document Representation
• From Trees to Continuous Embeddings and Back: Hyperbolic Hierarchical Clustering