Supervised Tree-Wasserstein Distance

TL;DR

The paper introduces Supervised Tree-Wasserstein (STW), a fast, task-specific metric learning method for document similarity that improves classification accuracy and computational efficiency using a supervised, tree-based approximation of the Wasserstein distance.

Contribution

It proposes a novel supervised metric learning approach for the Wasserstein distance using tree metrics, enhancing both accuracy and computational speed for document comparison.

Findings

STW improves document classification accuracy.

STW runs efficiently on GPU for large-scale comparisons.

STW enables fast, batch processing of document similarities.

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.